3553 lines
157 KiB
Metal
3553 lines
157 KiB
Metal
// Qwen3.8 Flash Next primitives. Rust owns the graph and all state lifetimes;
|
|
// this file contains only the data-parallel kernels executed by Metal.
|
|
|
|
#include <metal_stdlib>
|
|
#ifdef DS4_METAL_HAS_TENSOR
|
|
#include <metal_tensor>
|
|
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
|
|
#endif
|
|
using namespace metal;
|
|
|
|
struct qwen_kernel_args {
|
|
uint u[16];
|
|
float f[8];
|
|
};
|
|
|
|
static inline float qwen_bf16(ushort value) {
|
|
return as_type<float>((uint)value << 16);
|
|
}
|
|
|
|
static inline ushort qwen_to_bf16(float value) {
|
|
uint bits = as_type<uint>(value);
|
|
bits += 0x7fffu + ((bits >> 16) & 1u);
|
|
return (ushort)(bits >> 16);
|
|
}
|
|
|
|
static inline float qwen_round_bf16(float value) {
|
|
return qwen_bf16(qwen_to_bf16(value));
|
|
}
|
|
|
|
static inline bfloat qwen_simd_sum_bf16(bfloat value) {
|
|
float sum = (float)value;
|
|
for (ushort offset = 16u; offset >= 1u; offset >>= 1u) {
|
|
sum += simd_shuffle_down(sum, offset);
|
|
}
|
|
return (bfloat)sum;
|
|
}
|
|
|
|
static inline float qwen_sigmoid_bf16(float value) {
|
|
const float input = qwen_round_bf16(value);
|
|
const float magnitude = qwen_round_bf16(input < 0.0f ? -input : input);
|
|
const float exponential = qwen_round_bf16(metal::precise::exp(magnitude));
|
|
const float denominator = qwen_round_bf16(1.0f + exponential);
|
|
const float tail = qwen_round_bf16(1.0f / denominator);
|
|
return input < 0.0f ? tail : qwen_round_bf16(1.0f - tail);
|
|
}
|
|
|
|
static inline float qwen_silu_sigmoid_bf16(float value) {
|
|
const bfloat input = (bfloat)value;
|
|
const bfloat magnitude = input < bfloat(0) ? -input : input;
|
|
const bfloat exponential = (bfloat)exp((float)magnitude);
|
|
const bfloat denominator = (bfloat)(1.0f + (float)exponential);
|
|
const bfloat tail = (bfloat)(1.0f / (float)denominator);
|
|
return (float)(input < bfloat(0) ? tail : (bfloat)(1.0f - (float)tail));
|
|
}
|
|
|
|
static inline float qwen_apply_rope(float value, float rotated, float theta) {
|
|
volatile float cosine_term = value * cos(theta);
|
|
volatile float sine_term = rotated * sin(theta);
|
|
return cosine_term + sine_term;
|
|
}
|
|
|
|
static inline ushort qwen_weight_u16(
|
|
device const uchar *data,
|
|
uint byte_offset,
|
|
ulong index) {
|
|
const ulong byte = (ulong)byte_offset + index * 2u;
|
|
if ((byte & 1u) == 0u) return *((device const ushort *)(data + byte));
|
|
return (ushort)data[byte] | ((ushort)data[byte + 1u] << 8u);
|
|
}
|
|
|
|
static inline uint qwen_weight_u32(
|
|
device const uchar *data,
|
|
uint byte_offset,
|
|
ulong index) {
|
|
const ulong byte = (ulong)byte_offset + index * 4u;
|
|
if ((byte & 3u) == 0u) return *((device const uint *)(data + byte));
|
|
return (uint)data[byte] |
|
|
((uint)data[byte + 1u] << 8u) |
|
|
((uint)data[byte + 2u] << 16u) |
|
|
((uint)data[byte + 3u] << 24u);
|
|
}
|
|
|
|
kernel void kernel_qwen_probe_affine_table(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device uint *out [[buffer(1)]],
|
|
device const uint *selected [[buffer(3)]],
|
|
device const uchar *packed [[buffer(5)]],
|
|
device const uchar *scales [[buffer(6)]],
|
|
device const uchar *biases [[buffer(7)]],
|
|
uint slot [[thread_position_in_grid]]) {
|
|
if (slot >= args.u[4]) return;
|
|
const uint expert = selected[slot];
|
|
const ulong row = (ulong)expert * args.u[1];
|
|
const ulong packed_index = row * (args.u[0] / (32u / args.u[2]));
|
|
const ulong parameter = row * (args.u[0] / args.u[3]);
|
|
out[slot * 3u] = qwen_weight_u32(packed, args.u[13], packed_index);
|
|
out[slot * 3u + 1u] = qwen_weight_u16(scales, args.u[14], parameter);
|
|
out[slot * 3u + 2u] = qwen_weight_u16(biases, args.u[15], parameter);
|
|
}
|
|
|
|
static inline float qwen_quant_weight(
|
|
device const uchar *packed,
|
|
device const uchar *scales,
|
|
device const uchar *biases,
|
|
uint packed_offset,
|
|
uint scales_offset,
|
|
uint biases_offset,
|
|
uint row,
|
|
uint column,
|
|
uint in_dim,
|
|
uint bits,
|
|
uint group_size) {
|
|
const uint per_word = 32u / bits;
|
|
const uint packed_columns = in_dim / per_word;
|
|
const uint groups = in_dim / group_size;
|
|
const uint word = qwen_weight_u32(packed, packed_offset,
|
|
(ulong)row * packed_columns + column / per_word);
|
|
const uint mask = (1u << bits) - 1u;
|
|
const uint quant = (word >> ((column % per_word) * bits)) & mask;
|
|
const uint group = row * groups + column / group_size;
|
|
return fma((float)quant,
|
|
qwen_bf16(qwen_weight_u16(scales, scales_offset, group)),
|
|
qwen_bf16(qwen_weight_u16(biases, biases_offset, group)));
|
|
}
|
|
|
|
static inline float qwen_quant_value(
|
|
device const uint *packed,
|
|
device const ushort *scales,
|
|
device const ushort *biases,
|
|
uint row,
|
|
uint column,
|
|
uint in_dim,
|
|
uint bits,
|
|
uint group_size) {
|
|
const uint per_word = 32u / bits;
|
|
const uint packed_columns = in_dim / per_word;
|
|
const uint groups = in_dim / group_size;
|
|
const uint word = packed[(ulong)row * packed_columns + column / per_word];
|
|
const uint mask = (1u << bits) - 1u;
|
|
const uint quant = (word >> ((column % per_word) * bits)) & mask;
|
|
const uint group = row * groups + column / group_size;
|
|
return fma((float)quant, qwen_bf16(scales[group]), qwen_bf16(biases[group]));
|
|
}
|
|
|
|
kernel void kernel_qwen_affine_mv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *packed [[buffer(5)]],
|
|
device const uchar *scales [[buffer(6)]],
|
|
device const uchar *biases [[buffer(7)]],
|
|
uint row [[thread_position_in_grid]]) {
|
|
const uint in_dim = args.u[0];
|
|
const uint out_dim = args.u[1];
|
|
if (row >= out_dim) return;
|
|
float sum = 0.0f;
|
|
for (uint column = 0; column < in_dim; column++) {
|
|
sum = fma(qwen_quant_weight(packed, scales, biases,
|
|
args.u[13], args.u[14], args.u[15], row, column,
|
|
in_dim, args.u[2], args.u[3]), x[column], sum);
|
|
}
|
|
out[row] = args.u[11] != 0u ? qwen_round_bf16(sum) : sum;
|
|
}
|
|
|
|
// Decode uses one activation row. Each SIMD group owns four output rows so
|
|
// the packed-weight decode shares every activation load across those rows,
|
|
// matching the affine qmv shape used by the reference MLX backend.
|
|
kernel void kernel_qwen_affine_qmv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const 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 simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint in_dim = args.u[0];
|
|
const uint out_dim = args.u[1];
|
|
const uint bits = args.u[2];
|
|
const uint group_size = args.u[3];
|
|
const uint values_per_word = 32u / bits;
|
|
const uint packed_columns = in_dim / values_per_word;
|
|
const uint words_per_group = group_size / values_per_word;
|
|
const uint groups_per_row = in_dim / group_size;
|
|
const uint first_row = group.x * 8u + simd_group * 4u;
|
|
const uint mask = (1u << bits) - 1u;
|
|
float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
|
|
for (uint word_index = lane; word_index < packed_columns; word_index += 32u) {
|
|
const uint first_column = word_index * values_per_word;
|
|
const uint quant_group = word_index / words_per_group;
|
|
uint words[4] = {0u, 0u, 0u, 0u};
|
|
float scale[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
float bias[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
for (uint output = 0; output < 4u; output++) {
|
|
const uint row = first_row + output;
|
|
if (row < out_dim) {
|
|
words[output] = qwen_weight_u32(
|
|
packed, args.u[13], (ulong)row * packed_columns + word_index);
|
|
const ulong parameter = (ulong)row * groups_per_row + quant_group;
|
|
scale[output] = qwen_bf16(qwen_weight_u16(scales, args.u[14], parameter));
|
|
bias[output] = qwen_bf16(qwen_weight_u16(biases, args.u[15], parameter));
|
|
}
|
|
}
|
|
for (uint packed_value = 0; packed_value < values_per_word; packed_value++) {
|
|
const float activation = x[first_column + packed_value];
|
|
const uint shift = packed_value * bits;
|
|
for (uint output = 0; output < 4u; output++) {
|
|
const float weight = fma((float)((words[output] >> shift) & mask),
|
|
scale[output], bias[output]);
|
|
sums[output] = fma(weight, activation, sums[output]);
|
|
}
|
|
}
|
|
}
|
|
for (uint output = 0; output < 4u; output++) {
|
|
const float sum = simd_sum(sums[output]);
|
|
if (lane == 0u && first_row + output < out_dim) {
|
|
out[first_row + output] = args.u[11] != 0u ? qwen_round_bf16(sum) : sum;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <ushort bits, ushort group_size>
|
|
static inline void qwen_affine_qmv_fast_impl(
|
|
constant qwen_kernel_args &args,
|
|
device float *out,
|
|
device const float *x,
|
|
device const uchar *packed,
|
|
device const uchar *scale_bytes,
|
|
device const uchar *bias_bytes,
|
|
uint packed_offset,
|
|
uint scale_offset,
|
|
uint bias_offset,
|
|
uint2 group,
|
|
uint simd_group,
|
|
uint lane,
|
|
uint expert,
|
|
uint input_slot,
|
|
uint output_slot) {
|
|
static_assert(bits == 4 || bits == 8, "Qwen affine QMV supports 4/8-bit packs");
|
|
static_assert(group_size == 32 || group_size == 64,
|
|
"Qwen affine QMV supports group sizes 32/64");
|
|
constexpr ushort packs_per_thread = bits == 2 ? 1 : 2;
|
|
constexpr ushort pack_factor = 32 / bits;
|
|
constexpr ushort values_per_thread = pack_factor * packs_per_thread;
|
|
constexpr ushort block_size = values_per_thread * 32;
|
|
constexpr ushort scale_step = group_size / values_per_thread;
|
|
const uint in_dim = args.u[0];
|
|
const uint out_dim = args.u[1];
|
|
const uint input_bytes_per_row = in_dim * bits / 8;
|
|
const uint groups_per_row = in_dim / group_size;
|
|
const uint first_row = group.x * 8u + simd_group * 4u;
|
|
const ulong table_row = (ulong)expert * out_dim + first_row;
|
|
device const uchar *weights = packed + packed_offset +
|
|
table_row * input_bytes_per_row + lane * packs_per_thread * 4u;
|
|
device const ushort *scales =
|
|
(device const ushort *)(scale_bytes + scale_offset);
|
|
device const ushort *biases =
|
|
(device const ushort *)(bias_bytes + bias_offset);
|
|
ulong parameter = table_row * groups_per_row + lane / scale_step;
|
|
device const float *input = x +
|
|
(ulong)input_slot * in_dim + lane * values_per_thread;
|
|
float result[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
|
|
for (uint column = 0; column < in_dim; column += block_size) {
|
|
float values[values_per_thread];
|
|
float input_sum = 0.0f;
|
|
for (ushort i = 0; i < values_per_thread; i++) {
|
|
const float value = input[i];
|
|
if constexpr (bits == 4) {
|
|
constexpr float divisors[4] = {1.0f, 16.0f, 256.0f, 4096.0f};
|
|
values[i] = value / divisors[i & 3u];
|
|
} else {
|
|
values[i] = value;
|
|
}
|
|
}
|
|
if constexpr (bits == 4) {
|
|
for (ushort i = 0; i < values_per_thread; i += 4u) {
|
|
const bfloat group_sum =
|
|
(bfloat)input[i] + (bfloat)input[i + 1u] +
|
|
(bfloat)input[i + 2u] + (bfloat)input[i + 3u];
|
|
input_sum += (float)group_sum;
|
|
}
|
|
} else {
|
|
for (ushort i = 0; i < values_per_thread; i++) input_sum += input[i];
|
|
}
|
|
for (ushort row = 0; row < 4; row++) {
|
|
device const uchar *row_weights = weights + row * input_bytes_per_row;
|
|
float dot = 0.0f;
|
|
if constexpr (bits == 4) {
|
|
for (ushort i = 0; i < values_per_thread / 4; i++) {
|
|
const ushort word = ((device const ushort *)row_weights)[i];
|
|
dot += values[4 * i] * (word & 0x000fu) +
|
|
values[4 * i + 1] * (word & 0x00f0u) +
|
|
values[4 * i + 2] * (word & 0x0f00u) +
|
|
values[4 * i + 3] * (word & 0xf000u);
|
|
}
|
|
} else {
|
|
for (ushort i = 0; i < values_per_thread; i++) {
|
|
dot += values[i] * row_weights[packed_offset + i];
|
|
}
|
|
}
|
|
const ulong row_parameter = parameter + (ulong)row * groups_per_row;
|
|
result[row] +=
|
|
qwen_bf16(scales[row_parameter]) * dot +
|
|
input_sum *
|
|
qwen_bf16(biases[row_parameter]);
|
|
}
|
|
weights += block_size * bits / 8;
|
|
parameter += block_size / group_size;
|
|
input += block_size;
|
|
}
|
|
|
|
for (ushort row = 0; row < 4; row++) {
|
|
const float sum = simd_sum(result[row]);
|
|
if (lane == 0u) {
|
|
out[(ulong)output_slot * out_dim + first_row + row] =
|
|
args.u[11] != 0u ? (float)(bfloat)sum : sum;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <ushort bits, ushort group_size>
|
|
kernel void kernel_qwen_affine_qmv_fast(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *packed [[buffer(5)]],
|
|
device const uchar *scales [[buffer(6)]],
|
|
device const uchar *biases [[buffer(7)]],
|
|
uint group [[threadgroup_position_in_grid]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
qwen_affine_qmv_fast_impl<bits, group_size>(
|
|
args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15],
|
|
group, simd_group, lane,
|
|
0u, 0u, 0u);
|
|
}
|
|
|
|
// 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(
|
|
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) {
|
|
constexpr uint k_lanes = 8u;
|
|
constexpr uint outputs_per_simdgroup = 4u;
|
|
constexpr uint outputs_per_group = 8u;
|
|
constexpr uint sub = 8u;
|
|
const uint k_lane = lane % k_lanes;
|
|
const uint simd_row = lane / k_lanes;
|
|
const uint output_row = group * outputs_per_group +
|
|
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};
|
|
|
|
for (uint quant_group = k_lane; quant_group < groups_per_row;
|
|
quant_group += k_lanes) {
|
|
const float scale = qwen_bf16(qwen_weight_u16(
|
|
scales, args.u[14], parameter_row + quant_group));
|
|
const float bias = qwen_bf16(qwen_weight_u16(
|
|
biases, args.u[15], parameter_row + quant_group));
|
|
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];
|
|
for (uint index = 0u; index < sub; index++) {
|
|
uint quantized;
|
|
if constexpr (bits == 4) {
|
|
const uchar byte = packed[weight_byte + index / 2u];
|
|
quantized = (byte >> ((index & 1u) * 4u)) & 15u;
|
|
} else {
|
|
quantized = packed[weight_byte + index];
|
|
}
|
|
weights[index] = scale * (float)quantized + bias;
|
|
}
|
|
for (uint vector = 0u; vector < rows; vector++) {
|
|
float sum = 0.0f;
|
|
for (uint index = 0u; index < sub; index++) {
|
|
sum += x[(ulong)vector * in_dim + column + index] * weights[index];
|
|
}
|
|
result[vector] += sum;
|
|
}
|
|
}
|
|
}
|
|
|
|
for (uint vector = 0u; vector < rows; vector++) {
|
|
result[vector] += simd_shuffle_down(result[vector], 4u);
|
|
result[vector] += simd_shuffle_down(result[vector], 2u);
|
|
result[vector] += simd_shuffle_down(result[vector], 1u);
|
|
if (k_lane == 0u && output_row < out_dim) {
|
|
out[(ulong)vector * out_dim + output_row] = args.u[11] != 0u
|
|
? qwen_round_bf16(result[vector])
|
|
: result[vector];
|
|
}
|
|
}
|
|
}
|
|
|
|
template <ushort bits, ushort group_size>
|
|
kernel void kernel_qwen_affine_qmv_batch_fast(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const 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 simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
qwen_affine_qmv_wide_impl<bits, group_size>(
|
|
args, out, x, packed, scales, biases, group.x, simd_group, lane);
|
|
}
|
|
|
|
template <ushort bits, ushort group_size>
|
|
kernel void kernel_qwen_affine_pair_qmv_wide(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *gate_out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device float *up_out [[buffer(4)]],
|
|
device const uchar *gate_packed [[buffer(5)]],
|
|
device const uchar *gate_scales [[buffer(6)]],
|
|
device const uchar *gate_biases [[buffer(7)]],
|
|
device const uchar *up_packed [[buffer(9)]],
|
|
device const uchar *up_scales [[buffer(10)]],
|
|
device const uchar *up_biases [[buffer(11)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint projection_groups = (args.u[1] + 7u) / 8u;
|
|
const bool up = group.x >= projection_groups;
|
|
qwen_affine_qmv_wide_impl<bits, group_size>(
|
|
args,
|
|
up ? up_out : gate_out,
|
|
x,
|
|
up ? up_packed : gate_packed,
|
|
up ? up_scales : gate_scales,
|
|
up ? up_biases : gate_biases,
|
|
group.x - (up ? projection_groups : 0u),
|
|
simd_group,
|
|
lane);
|
|
}
|
|
|
|
template <ushort bits, ushort group_size>
|
|
kernel void kernel_qwen_affine_gather_qmv_fast(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uint *selected [[buffer(3)]],
|
|
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 simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint slot = group.y;
|
|
qwen_affine_qmv_fast_impl<bits, group_size>(
|
|
args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15],
|
|
group, simd_group, lane,
|
|
selected[slot], args.u[5] != 0u ? slot : 0u, slot);
|
|
}
|
|
|
|
template <ushort bits, ushort group_size>
|
|
kernel void kernel_qwen_affine_gather_qmv_batch_fast(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uint *selected [[buffer(3)]],
|
|
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 simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint slot = group.y;
|
|
qwen_affine_qmv_fast_impl<bits, group_size>(
|
|
args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15],
|
|
group, simd_group, lane, selected[slot],
|
|
args.u[5] != 0u ? slot : slot / args.u[8], slot);
|
|
}
|
|
|
|
template <ushort bits, ushort group_size, bool gather, bool batch = false>
|
|
kernel void kernel_qwen_affine_pair_qmv_fast(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *gate_out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uint *selected [[buffer(3)]],
|
|
device float *up_out [[buffer(4)]],
|
|
device const uchar *gate_packed [[buffer(5)]],
|
|
device const uchar *gate_scales [[buffer(6)]],
|
|
device const uchar *gate_biases [[buffer(7)]],
|
|
device const uchar *up_packed [[buffer(9)]],
|
|
device const uchar *up_scales [[buffer(10)]],
|
|
device const uchar *up_biases [[buffer(11)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint projection_groups = (args.u[1] + 7u) / 8u;
|
|
const bool up = group.x >= projection_groups;
|
|
const uint2 projection_group = uint2(group.x - (up ? projection_groups : 0u), group.y);
|
|
const uint slot = group.y;
|
|
qwen_affine_qmv_fast_impl<bits, group_size>(
|
|
args,
|
|
up ? up_out : gate_out,
|
|
x,
|
|
up ? up_packed : gate_packed,
|
|
up ? up_scales : gate_scales,
|
|
up ? up_biases : gate_biases,
|
|
up ? args.u[9] : args.u[13],
|
|
up ? args.u[6] : args.u[14],
|
|
up ? args.u[7] : args.u[15],
|
|
projection_group,
|
|
simd_group,
|
|
lane,
|
|
gather ? selected[slot] : 0u,
|
|
batch
|
|
? (gather && args.u[5] == 0u ? slot / args.u[8] : slot)
|
|
: (gather && args.u[5] != 0u ? slot : 0u),
|
|
slot);
|
|
}
|
|
|
|
typedef decltype(kernel_qwen_affine_qmv_fast<4, 32>) qwen_affine_qmv_fast_b4g32;
|
|
typedef decltype(kernel_qwen_affine_qmv_fast<4, 64>) qwen_affine_qmv_fast_b4g64;
|
|
typedef decltype(kernel_qwen_affine_qmv_fast<8, 64>) qwen_affine_qmv_fast_b8g64;
|
|
typedef decltype(kernel_qwen_affine_qmv_batch_fast<4, 32>) qwen_affine_qmv_batch_fast_b4g32;
|
|
typedef decltype(kernel_qwen_affine_qmv_batch_fast<4, 64>) qwen_affine_qmv_batch_fast_b4g64;
|
|
typedef decltype(kernel_qwen_affine_qmv_batch_fast<8, 64>) qwen_affine_qmv_batch_fast_b8g64;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_fast<4, 32>) qwen_affine_gather_qmv_fast_b4g32;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_fast<4, 64>) qwen_affine_gather_qmv_fast_b4g64;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_fast<8, 64>) qwen_affine_gather_qmv_fast_b8g64;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_batch_fast<4, 32>) qwen_affine_gather_qmv_batch_fast_b4g32;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_batch_fast<4, 64>) qwen_affine_gather_qmv_batch_fast_b4g64;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_batch_fast<8, 64>) qwen_affine_gather_qmv_batch_fast_b8g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 32, false>) qwen_affine_pair_qmv_fast_b4g32;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 64, false>) qwen_affine_pair_qmv_fast_b4g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<8, 64, false>) qwen_affine_pair_qmv_fast_b8g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 32, true>) qwen_affine_gather_pair_qmv_fast_b4g32;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 64, true>) qwen_affine_gather_pair_qmv_fast_b4g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<8, 64, true>) qwen_affine_gather_pair_qmv_fast_b8g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 32, false, true>) qwen_affine_pair_qmv_batch_fast_b4g32;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 64, false, true>) qwen_affine_pair_qmv_batch_fast_b4g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<8, 64, false, true>) qwen_affine_pair_qmv_batch_fast_b8g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 32, true, true>) qwen_affine_gather_pair_qmv_batch_fast_b4g32;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<4, 64, true, true>) qwen_affine_gather_pair_qmv_batch_fast_b4g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_fast<8, 64, true, true>) qwen_affine_gather_pair_qmv_batch_fast_b8g64;
|
|
|
|
template [[host_name("kernel_qwen_affine_qmv_fast_b4g32")]]
|
|
kernel qwen_affine_qmv_fast_b4g32 kernel_qwen_affine_qmv_fast<4, 32>;
|
|
template [[host_name("kernel_qwen_affine_qmv_fast_b4g64")]]
|
|
kernel qwen_affine_qmv_fast_b4g64 kernel_qwen_affine_qmv_fast<4, 64>;
|
|
template [[host_name("kernel_qwen_affine_qmv_fast_b8g64")]]
|
|
kernel qwen_affine_qmv_fast_b8g64 kernel_qwen_affine_qmv_fast<8, 64>;
|
|
template [[host_name("kernel_qwen_affine_qmv_batch_fast_b4g32")]]
|
|
kernel qwen_affine_qmv_batch_fast_b4g32 kernel_qwen_affine_qmv_batch_fast<4, 32>;
|
|
template [[host_name("kernel_qwen_affine_qmv_batch_fast_b4g64")]]
|
|
kernel qwen_affine_qmv_batch_fast_b4g64 kernel_qwen_affine_qmv_batch_fast<4, 64>;
|
|
template [[host_name("kernel_qwen_affine_qmv_batch_fast_b8g64")]]
|
|
kernel qwen_affine_qmv_batch_fast_b8g64 kernel_qwen_affine_qmv_batch_fast<8, 64>;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_wide<4, 32>) qwen_affine_pair_qmv_wide_b4g32;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_wide<4, 64>) qwen_affine_pair_qmv_wide_b4g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_wide<8, 64>) qwen_affine_pair_qmv_wide_b8g64;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_wide_b4g32")]]
|
|
kernel qwen_affine_pair_qmv_wide_b4g32 kernel_qwen_affine_pair_qmv_wide<4, 32>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_wide_b4g64")]]
|
|
kernel qwen_affine_pair_qmv_wide_b4g64 kernel_qwen_affine_pair_qmv_wide<4, 64>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_wide_b8g64")]]
|
|
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(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const 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 = 64u;
|
|
constexpr uint threads = 128u;
|
|
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);
|
|
matmul2d<descriptor, execution_simdgroups<4>> multiply;
|
|
auto accum = multiply.template get_destination_cooperative_tensor<
|
|
tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>,
|
|
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 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]
|
|
: 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;
|
|
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,
|
|
args.u[0], bits, group_size)
|
|
: bfloat(0.0f);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
auto a = tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>(
|
|
xs, dextents<int32_t, 2>(tile_k, tile_m));
|
|
auto b = tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>(
|
|
ws, dextents<int32_t, 2>(tile_k, tile_n));
|
|
multiply.run(a, b, accum);
|
|
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 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;
|
|
}
|
|
}
|
|
}
|
|
|
|
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;
|
|
typedef decltype(kernel_qwen_affine_qmm_mpp<4, 32, 64>) qwen_affine_qmm_mpp_b4g32_bm64;
|
|
typedef decltype(kernel_qwen_affine_qmm_mpp<4, 64, 64>) qwen_affine_qmm_mpp_b4g64_bm64;
|
|
typedef decltype(kernel_qwen_affine_qmm_mpp<8, 64, 64>) qwen_affine_qmm_mpp_b8g64_bm64;
|
|
|
|
template [[host_name("kernel_qwen_affine_qmm_mpp_b4g32_bm32")]]
|
|
kernel qwen_affine_qmm_mpp_b4g32_bm32 kernel_qwen_affine_qmm_mpp<4, 32, 32>;
|
|
template [[host_name("kernel_qwen_affine_qmm_mpp_b4g64_bm32")]]
|
|
kernel qwen_affine_qmm_mpp_b4g64_bm32 kernel_qwen_affine_qmm_mpp<4, 64, 32>;
|
|
template [[host_name("kernel_qwen_affine_qmm_mpp_b8g64_bm32")]]
|
|
kernel qwen_affine_qmm_mpp_b8g64_bm32 kernel_qwen_affine_qmm_mpp<8, 64, 32>;
|
|
template [[host_name("kernel_qwen_affine_qmm_mpp_b4g32_bm64")]]
|
|
kernel qwen_affine_qmm_mpp_b4g32_bm64 kernel_qwen_affine_qmm_mpp<4, 32, 64>;
|
|
template [[host_name("kernel_qwen_affine_qmm_mpp_b4g64_bm64")]]
|
|
kernel qwen_affine_qmm_mpp_b4g64_bm64 kernel_qwen_affine_qmm_mpp<4, 64, 64>;
|
|
template [[host_name("kernel_qwen_affine_qmm_mpp_b8g64_bm64")]]
|
|
kernel qwen_affine_qmm_mpp_b8g64_bm64 kernel_qwen_affine_qmm_mpp<8, 64, 64>;
|
|
|
|
template <ushort tile_m>
|
|
kernel void kernel_qwen_bf16_qmm_mpp(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *weights [[buffer(5)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint tid [[thread_index_in_threadgroup]]) {
|
|
constexpr uint tile_n = 64u;
|
|
constexpr uint tile_k = 64u;
|
|
constexpr uint threads = 128u;
|
|
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);
|
|
matmul2d<descriptor, execution_simdgroups<4>> multiply;
|
|
auto accum = multiply.template get_destination_cooperative_tensor<
|
|
tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>,
|
|
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 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]
|
|
: 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;
|
|
ws[index] = weight_n < args.u[1]
|
|
? (bfloat)qwen_bf16(qwen_weight_u16(
|
|
weights, args.u[13],
|
|
(ulong)weight_n * args.u[0] + first_k + column))
|
|
: bfloat(0.0f);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
auto a = tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>(
|
|
xs, dextents<int32_t, 2>(tile_k, tile_m));
|
|
auto b = tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>(
|
|
ws, dextents<int32_t, 2>(tile_k, tile_n));
|
|
multiply.run(a, b, accum);
|
|
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 output_m = first_m + index[1];
|
|
if (output_n < args.u[1] && output_m < args.u[4]) {
|
|
out[(ulong)output_m * args.u[1] + output_n] =
|
|
qwen_round_bf16(accum[i]);
|
|
}
|
|
}
|
|
}
|
|
|
|
typedef decltype(kernel_qwen_bf16_qmm_mpp<32>) qwen_bf16_qmm_mpp_bm32;
|
|
typedef decltype(kernel_qwen_bf16_qmm_mpp<64>) qwen_bf16_qmm_mpp_bm64;
|
|
|
|
template [[host_name("kernel_qwen_bf16_qmm_mpp_bm32")]]
|
|
kernel qwen_bf16_qmm_mpp_bm32 kernel_qwen_bf16_qmm_mpp<32>;
|
|
template [[host_name("kernel_qwen_bf16_qmm_mpp_bm64")]]
|
|
kernel qwen_bf16_qmm_mpp_bm64 kernel_qwen_bf16_qmm_mpp<64>;
|
|
#endif
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_fast_b4g32")]]
|
|
kernel qwen_affine_gather_qmv_fast_b4g32 kernel_qwen_affine_gather_qmv_fast<4, 32>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_fast_b4g64")]]
|
|
kernel qwen_affine_gather_qmv_fast_b4g64 kernel_qwen_affine_gather_qmv_fast<4, 64>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_fast_b8g64")]]
|
|
kernel qwen_affine_gather_qmv_fast_b8g64 kernel_qwen_affine_gather_qmv_fast<8, 64>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_batch_fast_b4g32")]]
|
|
kernel qwen_affine_gather_qmv_batch_fast_b4g32 kernel_qwen_affine_gather_qmv_batch_fast<4, 32>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_batch_fast_b4g64")]]
|
|
kernel qwen_affine_gather_qmv_batch_fast_b4g64 kernel_qwen_affine_gather_qmv_batch_fast<4, 64>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_batch_fast_b8g64")]]
|
|
kernel qwen_affine_gather_qmv_batch_fast_b8g64 kernel_qwen_affine_gather_qmv_batch_fast<8, 64>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_fast_b4g32")]]
|
|
kernel qwen_affine_pair_qmv_fast_b4g32 kernel_qwen_affine_pair_qmv_fast<4, 32, false>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_fast_b4g64")]]
|
|
kernel qwen_affine_pair_qmv_fast_b4g64 kernel_qwen_affine_pair_qmv_fast<4, 64, false>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_fast_b8g64")]]
|
|
kernel qwen_affine_pair_qmv_fast_b8g64 kernel_qwen_affine_pair_qmv_fast<8, 64, false>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_fast_b4g32")]]
|
|
kernel qwen_affine_gather_pair_qmv_fast_b4g32 kernel_qwen_affine_pair_qmv_fast<4, 32, true>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_fast_b4g64")]]
|
|
kernel qwen_affine_gather_pair_qmv_fast_b4g64 kernel_qwen_affine_pair_qmv_fast<4, 64, true>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_fast_b8g64")]]
|
|
kernel qwen_affine_gather_pair_qmv_fast_b8g64 kernel_qwen_affine_pair_qmv_fast<8, 64, true>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_batch_fast_b4g32")]]
|
|
kernel qwen_affine_pair_qmv_batch_fast_b4g32 kernel_qwen_affine_pair_qmv_fast<4, 32, false, true>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_batch_fast_b4g64")]]
|
|
kernel qwen_affine_pair_qmv_batch_fast_b4g64 kernel_qwen_affine_pair_qmv_fast<4, 64, false, true>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_batch_fast_b8g64")]]
|
|
kernel qwen_affine_pair_qmv_batch_fast_b8g64 kernel_qwen_affine_pair_qmv_fast<8, 64, false, true>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_fast_b4g32")]]
|
|
kernel qwen_affine_gather_pair_qmv_batch_fast_b4g32 kernel_qwen_affine_pair_qmv_fast<4, 32, true, true>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_fast_b4g64")]]
|
|
kernel qwen_affine_gather_pair_qmv_batch_fast_b4g64 kernel_qwen_affine_pair_qmv_fast<4, 64, true, true>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_fast_b8g64")]]
|
|
kernel qwen_affine_gather_pair_qmv_batch_fast_b8g64 kernel_qwen_affine_pair_qmv_fast<8, 64, true, true>;
|
|
|
|
template <ushort bits, ushort group_size>
|
|
static inline void qwen_affine_qmv_safe_impl(
|
|
constant qwen_kernel_args &args,
|
|
device float *out,
|
|
device const float *x,
|
|
device const uchar *packed,
|
|
device const uchar *scale_bytes,
|
|
device const uchar *bias_bytes,
|
|
uint packed_offset,
|
|
uint scale_offset,
|
|
uint bias_offset,
|
|
uint2 group,
|
|
uint simd_group,
|
|
uint lane,
|
|
uint expert,
|
|
uint input_slot,
|
|
uint output_slot) {
|
|
constexpr ushort pack_factor = 32 / bits;
|
|
constexpr ushort values_per_thread = pack_factor;
|
|
constexpr ushort block_size = values_per_thread * 32;
|
|
constexpr ushort scale_step = group_size / values_per_thread;
|
|
const uint in_dim = args.u[0];
|
|
const uint out_dim = args.u[1];
|
|
const uint input_bytes_per_row = in_dim * bits / 8;
|
|
const uint groups_per_row = in_dim / group_size;
|
|
const uint first_row = group.x * 8u + simd_group * 4u;
|
|
if (first_row >= out_dim) return;
|
|
const uint used_row = out_dim < 8u ? first_row : min(out_dim - 4u, first_row);
|
|
const ulong table_row = (ulong)expert * out_dim + used_row;
|
|
device const uchar *weights = packed +
|
|
table_row * input_bytes_per_row + lane * 4u;
|
|
ulong parameter = table_row * groups_per_row + lane / scale_step;
|
|
device const float *input = x +
|
|
(ulong)input_slot * in_dim + lane * values_per_thread;
|
|
float result[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
|
|
uint column = 0;
|
|
for (; column + block_size < in_dim; column += block_size) {
|
|
float values[values_per_thread];
|
|
float input_sum = 0.0f;
|
|
for (ushort i = 0; i < values_per_thread; i++) {
|
|
if constexpr (bits == 4) {
|
|
constexpr float divisors[4] = {1.0f, 16.0f, 256.0f, 4096.0f};
|
|
values[i] = input[i] / divisors[i & 3u];
|
|
} else {
|
|
values[i] = input[i];
|
|
}
|
|
}
|
|
if constexpr (bits == 4) {
|
|
for (ushort i = 0; i < values_per_thread; i += 4u) {
|
|
const bfloat input_group =
|
|
(bfloat)input[i] + (bfloat)input[i + 1u] +
|
|
(bfloat)input[i + 2u] + (bfloat)input[i + 3u];
|
|
input_sum += (float)input_group;
|
|
}
|
|
} else {
|
|
for (ushort i = 0; i < values_per_thread; i++) input_sum += input[i];
|
|
}
|
|
for (ushort row = 0; row < 4u && used_row + row < out_dim; row++) {
|
|
device const uchar *row_weights = weights + row * input_bytes_per_row;
|
|
float dot = 0.0f;
|
|
if constexpr (bits == 4) {
|
|
for (ushort i = 0; i < values_per_thread / 4u; i++) {
|
|
const ushort word = qwen_weight_u16(row_weights, packed_offset, i);
|
|
dot += values[4u * i] * (word & 0x000fu) +
|
|
values[4u * i + 1u] * (word & 0x00f0u) +
|
|
values[4u * i + 2u] * (word & 0x0f00u) +
|
|
values[4u * i + 3u] * (word & 0xf000u);
|
|
}
|
|
} else {
|
|
for (ushort i = 0; i < values_per_thread; i++) {
|
|
dot += values[i] * row_weights[packed_offset + i];
|
|
}
|
|
}
|
|
const ulong row_parameter = parameter + (ulong)row * groups_per_row;
|
|
result[row] +=
|
|
qwen_bf16(qwen_weight_u16(scale_bytes, scale_offset, row_parameter)) * dot +
|
|
input_sum *
|
|
qwen_bf16(qwen_weight_u16(bias_bytes, bias_offset, row_parameter));
|
|
}
|
|
weights += block_size * bits / 8;
|
|
parameter += block_size / group_size;
|
|
input += block_size;
|
|
}
|
|
|
|
const int remaining = clamp(
|
|
(int)in_dim - (int)column - (int)lane * (int)values_per_thread,
|
|
0,
|
|
(int)values_per_thread);
|
|
if (remaining > 0) {
|
|
float values[values_per_thread];
|
|
float input_sum = 0.0f;
|
|
for (ushort i = 0; i < values_per_thread; i++) {
|
|
if ((int)i < remaining) {
|
|
if constexpr (bits == 4) {
|
|
constexpr float divisors[4] = {1.0f, 16.0f, 256.0f, 4096.0f};
|
|
values[i] = input[i] / divisors[i & 3u];
|
|
} else {
|
|
values[i] = input[i];
|
|
}
|
|
} else {
|
|
values[i] = 0.0f;
|
|
}
|
|
}
|
|
if constexpr (bits == 4) {
|
|
for (ushort i = 0; i < (ushort)remaining; i += 4u) {
|
|
const bfloat input_group =
|
|
(bfloat)input[i] + (bfloat)input[i + 1u] +
|
|
(bfloat)input[i + 2u] + (bfloat)input[i + 3u];
|
|
input_sum += (float)input_group;
|
|
}
|
|
} else {
|
|
for (ushort i = 0; i < (ushort)remaining; i++) input_sum += input[i];
|
|
}
|
|
for (ushort row = 0; row < 4u && used_row + row < out_dim; row++) {
|
|
device const uchar *row_weights = weights + row * input_bytes_per_row;
|
|
float dot = 0.0f;
|
|
if constexpr (bits == 4) {
|
|
for (ushort i = 0; i < (ushort)remaining / 4u; i++) {
|
|
const ushort word = qwen_weight_u16(row_weights, packed_offset, i);
|
|
dot += values[4u * i] * (word & 0x000fu) +
|
|
values[4u * i + 1u] * (word & 0x00f0u) +
|
|
values[4u * i + 2u] * (word & 0x0f00u) +
|
|
values[4u * i + 3u] * (word & 0xf000u);
|
|
}
|
|
} else {
|
|
for (ushort i = 0; i < (ushort)remaining; i++) {
|
|
dot += values[i] * row_weights[packed_offset + i];
|
|
}
|
|
}
|
|
const ulong row_parameter = parameter + (ulong)row * groups_per_row;
|
|
result[row] +=
|
|
qwen_bf16(qwen_weight_u16(scale_bytes, scale_offset, row_parameter)) * dot +
|
|
input_sum *
|
|
qwen_bf16(qwen_weight_u16(bias_bytes, bias_offset, row_parameter));
|
|
}
|
|
}
|
|
|
|
for (ushort row = 0; row < 4u && used_row + row < out_dim; row++) {
|
|
const float sum = simd_sum(result[row]);
|
|
if (lane == 0u) {
|
|
out[(ulong)output_slot * out_dim + used_row + row] =
|
|
args.u[11] != 0u ? (float)(bfloat)sum : sum;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <ushort bits, ushort group_size>
|
|
kernel void kernel_qwen_affine_qmv_safe(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const 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 simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
qwen_affine_qmv_safe_impl<bits, group_size>(
|
|
args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15],
|
|
group, simd_group, lane,
|
|
0u, 0u, 0u);
|
|
}
|
|
|
|
template <ushort bits, ushort group_size>
|
|
kernel void kernel_qwen_affine_qmv_batch_safe(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const 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 simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
qwen_affine_qmv_wide_impl<bits, group_size>(
|
|
args, out, x, packed, scales, biases, group.x, simd_group, lane);
|
|
}
|
|
|
|
template <ushort bits, ushort group_size>
|
|
kernel void kernel_qwen_affine_gather_qmv_safe(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uint *selected [[buffer(3)]],
|
|
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 simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint slot = group.y;
|
|
qwen_affine_qmv_safe_impl<bits, group_size>(
|
|
args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15],
|
|
group, simd_group, lane,
|
|
selected[slot], args.u[5] != 0u ? slot : 0u, slot);
|
|
}
|
|
|
|
template <ushort bits, ushort group_size>
|
|
kernel void kernel_qwen_affine_gather_qmv_batch_safe(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uint *selected [[buffer(3)]],
|
|
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 simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint slot = group.y;
|
|
qwen_affine_qmv_safe_impl<bits, group_size>(
|
|
args, out, x, packed, scales, biases, args.u[13], args.u[14], args.u[15],
|
|
group, simd_group, lane, selected[slot],
|
|
args.u[5] != 0u ? slot : slot / args.u[8], slot);
|
|
}
|
|
|
|
template <ushort bits, ushort group_size, bool gather, bool batch = false>
|
|
kernel void kernel_qwen_affine_pair_qmv_safe(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *gate_out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uint *selected [[buffer(3)]],
|
|
device float *up_out [[buffer(4)]],
|
|
device const uchar *gate_packed [[buffer(5)]],
|
|
device const uchar *gate_scales [[buffer(6)]],
|
|
device const uchar *gate_biases [[buffer(7)]],
|
|
device const uchar *up_packed [[buffer(9)]],
|
|
device const uchar *up_scales [[buffer(10)]],
|
|
device const uchar *up_biases [[buffer(11)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint projection_groups = (args.u[1] + 7u) / 8u;
|
|
const bool up = group.x >= projection_groups;
|
|
const uint2 projection_group = uint2(group.x - (up ? projection_groups : 0u), group.y);
|
|
const uint slot = group.y;
|
|
qwen_affine_qmv_safe_impl<bits, group_size>(
|
|
args,
|
|
up ? up_out : gate_out,
|
|
x,
|
|
up ? up_packed : gate_packed,
|
|
up ? up_scales : gate_scales,
|
|
up ? up_biases : gate_biases,
|
|
up ? args.u[9] : args.u[13],
|
|
up ? args.u[6] : args.u[14],
|
|
up ? args.u[7] : args.u[15],
|
|
projection_group,
|
|
simd_group,
|
|
lane,
|
|
gather ? selected[slot] : 0u,
|
|
batch
|
|
? (gather && args.u[5] == 0u ? slot / args.u[8] : slot)
|
|
: (gather && args.u[5] != 0u ? slot : 0u),
|
|
slot);
|
|
}
|
|
|
|
typedef decltype(kernel_qwen_affine_qmv_safe<4, 32>) qwen_affine_qmv_safe_b4g32;
|
|
typedef decltype(kernel_qwen_affine_qmv_safe<4, 64>) qwen_affine_qmv_safe_b4g64;
|
|
typedef decltype(kernel_qwen_affine_qmv_safe<8, 64>) qwen_affine_qmv_safe_b8g64;
|
|
typedef decltype(kernel_qwen_affine_qmv_batch_safe<4, 32>) qwen_affine_qmv_batch_safe_b4g32;
|
|
typedef decltype(kernel_qwen_affine_qmv_batch_safe<4, 64>) qwen_affine_qmv_batch_safe_b4g64;
|
|
typedef decltype(kernel_qwen_affine_qmv_batch_safe<8, 64>) qwen_affine_qmv_batch_safe_b8g64;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_safe<4, 32>) qwen_affine_gather_qmv_safe_b4g32;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_safe<4, 64>) qwen_affine_gather_qmv_safe_b4g64;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_safe<8, 64>) qwen_affine_gather_qmv_safe_b8g64;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_batch_safe<4, 32>) qwen_affine_gather_qmv_batch_safe_b4g32;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_batch_safe<4, 64>) qwen_affine_gather_qmv_batch_safe_b4g64;
|
|
typedef decltype(kernel_qwen_affine_gather_qmv_batch_safe<8, 64>) qwen_affine_gather_qmv_batch_safe_b8g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 32, false>) qwen_affine_pair_qmv_safe_b4g32;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 64, false>) qwen_affine_pair_qmv_safe_b4g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<8, 64, false>) qwen_affine_pair_qmv_safe_b8g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 32, true>) qwen_affine_gather_pair_qmv_safe_b4g32;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 64, true>) qwen_affine_gather_pair_qmv_safe_b4g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<8, 64, true>) qwen_affine_gather_pair_qmv_safe_b8g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 32, false, true>) qwen_affine_pair_qmv_batch_safe_b4g32;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 64, false, true>) qwen_affine_pair_qmv_batch_safe_b4g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<8, 64, false, true>) qwen_affine_pair_qmv_batch_safe_b8g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 32, true, true>) qwen_affine_gather_pair_qmv_batch_safe_b4g32;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<4, 64, true, true>) qwen_affine_gather_pair_qmv_batch_safe_b4g64;
|
|
typedef decltype(kernel_qwen_affine_pair_qmv_safe<8, 64, true, true>) qwen_affine_gather_pair_qmv_batch_safe_b8g64;
|
|
|
|
template [[host_name("kernel_qwen_affine_qmv_safe_b4g32")]]
|
|
kernel qwen_affine_qmv_safe_b4g32 kernel_qwen_affine_qmv_safe<4, 32>;
|
|
template [[host_name("kernel_qwen_affine_qmv_safe_b4g64")]]
|
|
kernel qwen_affine_qmv_safe_b4g64 kernel_qwen_affine_qmv_safe<4, 64>;
|
|
template [[host_name("kernel_qwen_affine_qmv_safe_b8g64")]]
|
|
kernel qwen_affine_qmv_safe_b8g64 kernel_qwen_affine_qmv_safe<8, 64>;
|
|
template [[host_name("kernel_qwen_affine_qmv_batch_safe_b4g32")]]
|
|
kernel qwen_affine_qmv_batch_safe_b4g32 kernel_qwen_affine_qmv_batch_safe<4, 32>;
|
|
template [[host_name("kernel_qwen_affine_qmv_batch_safe_b4g64")]]
|
|
kernel qwen_affine_qmv_batch_safe_b4g64 kernel_qwen_affine_qmv_batch_safe<4, 64>;
|
|
template [[host_name("kernel_qwen_affine_qmv_batch_safe_b8g64")]]
|
|
kernel qwen_affine_qmv_batch_safe_b8g64 kernel_qwen_affine_qmv_batch_safe<8, 64>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_safe_b4g32")]]
|
|
kernel qwen_affine_gather_qmv_safe_b4g32 kernel_qwen_affine_gather_qmv_safe<4, 32>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_safe_b4g64")]]
|
|
kernel qwen_affine_gather_qmv_safe_b4g64 kernel_qwen_affine_gather_qmv_safe<4, 64>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_safe_b8g64")]]
|
|
kernel qwen_affine_gather_qmv_safe_b8g64 kernel_qwen_affine_gather_qmv_safe<8, 64>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_batch_safe_b4g32")]]
|
|
kernel qwen_affine_gather_qmv_batch_safe_b4g32 kernel_qwen_affine_gather_qmv_batch_safe<4, 32>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_batch_safe_b4g64")]]
|
|
kernel qwen_affine_gather_qmv_batch_safe_b4g64 kernel_qwen_affine_gather_qmv_batch_safe<4, 64>;
|
|
template [[host_name("kernel_qwen_affine_gather_qmv_batch_safe_b8g64")]]
|
|
kernel qwen_affine_gather_qmv_batch_safe_b8g64 kernel_qwen_affine_gather_qmv_batch_safe<8, 64>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_safe_b4g32")]]
|
|
kernel qwen_affine_pair_qmv_safe_b4g32 kernel_qwen_affine_pair_qmv_safe<4, 32, false>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_safe_b4g64")]]
|
|
kernel qwen_affine_pair_qmv_safe_b4g64 kernel_qwen_affine_pair_qmv_safe<4, 64, false>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_safe_b8g64")]]
|
|
kernel qwen_affine_pair_qmv_safe_b8g64 kernel_qwen_affine_pair_qmv_safe<8, 64, false>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_safe_b4g32")]]
|
|
kernel qwen_affine_gather_pair_qmv_safe_b4g32 kernel_qwen_affine_pair_qmv_safe<4, 32, true>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_safe_b4g64")]]
|
|
kernel qwen_affine_gather_pair_qmv_safe_b4g64 kernel_qwen_affine_pair_qmv_safe<4, 64, true>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_safe_b8g64")]]
|
|
kernel qwen_affine_gather_pair_qmv_safe_b8g64 kernel_qwen_affine_pair_qmv_safe<8, 64, true>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_batch_safe_b4g32")]]
|
|
kernel qwen_affine_pair_qmv_batch_safe_b4g32 kernel_qwen_affine_pair_qmv_safe<4, 32, false, true>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_batch_safe_b4g64")]]
|
|
kernel qwen_affine_pair_qmv_batch_safe_b4g64 kernel_qwen_affine_pair_qmv_safe<4, 64, false, true>;
|
|
template [[host_name("kernel_qwen_affine_pair_qmv_batch_safe_b8g64")]]
|
|
kernel qwen_affine_pair_qmv_batch_safe_b8g64 kernel_qwen_affine_pair_qmv_safe<8, 64, false, true>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_safe_b4g32")]]
|
|
kernel qwen_affine_gather_pair_qmv_batch_safe_b4g32 kernel_qwen_affine_pair_qmv_safe<4, 32, true, true>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_safe_b4g64")]]
|
|
kernel qwen_affine_gather_pair_qmv_batch_safe_b4g64 kernel_qwen_affine_pair_qmv_safe<4, 64, true, true>;
|
|
template [[host_name("kernel_qwen_affine_gather_pair_qmv_batch_safe_b8g64")]]
|
|
kernel qwen_affine_gather_pair_qmv_batch_safe_b8g64 kernel_qwen_affine_pair_qmv_safe<8, 64, true, true>;
|
|
|
|
kernel void kernel_qwen_affine_gather_qmv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const int *selected [[buffer(3)]],
|
|
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 simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint in_dim = args.u[0];
|
|
const uint out_dim = args.u[1];
|
|
const uint bits = args.u[2];
|
|
const uint group_size = args.u[3];
|
|
const uint slot = group.y;
|
|
const uint expert = (uint)selected[slot];
|
|
const uint values_per_word = 32u / bits;
|
|
const uint packed_columns = in_dim / values_per_word;
|
|
const uint words_per_group = group_size / values_per_word;
|
|
const uint groups_per_row = in_dim / group_size;
|
|
const uint first_row = group.x * 8u + simd_group * 4u;
|
|
const uint mask = (1u << bits) - 1u;
|
|
float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
|
|
for (uint word_index = lane; word_index < packed_columns; word_index += 32u) {
|
|
const uint first_column = word_index * values_per_word;
|
|
const uint quant_group = word_index / words_per_group;
|
|
uint words[4] = {0u, 0u, 0u, 0u};
|
|
float scale[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
float bias[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
for (uint output = 0; output < 4u; output++) {
|
|
const uint row = first_row + output;
|
|
if (row < out_dim) {
|
|
const ulong table_row = (ulong)expert * out_dim + row;
|
|
words[output] = qwen_weight_u32(
|
|
packed, args.u[13], table_row * packed_columns + word_index);
|
|
const ulong parameter = table_row * groups_per_row + quant_group;
|
|
scale[output] = qwen_bf16(qwen_weight_u16(scales, args.u[14], parameter));
|
|
bias[output] = qwen_bf16(qwen_weight_u16(biases, args.u[15], parameter));
|
|
}
|
|
}
|
|
for (uint packed_value = 0; packed_value < values_per_word; packed_value++) {
|
|
const float activation =
|
|
x[(args.u[5] != 0u ? (ulong)slot * in_dim : 0u) +
|
|
first_column + packed_value];
|
|
const uint shift = packed_value * bits;
|
|
for (uint output = 0; output < 4u; output++) {
|
|
const float weight = fma((float)((words[output] >> shift) & mask),
|
|
scale[output], bias[output]);
|
|
sums[output] = fma(weight, activation, sums[output]);
|
|
}
|
|
}
|
|
}
|
|
for (uint output = 0; output < 4u; output++) {
|
|
const float sum = simd_sum(sums[output]);
|
|
if (lane == 0u && first_row + output < out_dim) {
|
|
out[(ulong)slot * out_dim + first_row + output] = qwen_round_bf16(sum);
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_weighted_sum10(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *experts [[buffer(2)]],
|
|
device const float *weights [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index >= args.u[0]) return;
|
|
const uint hidden = args.u[1] == 0u ? args.u[0] : args.u[1];
|
|
const uint row = index / hidden;
|
|
const uint column = index % hidden;
|
|
const ulong experts_base = (ulong)row * 10u * hidden;
|
|
const ulong weights_base = (ulong)row * 10u;
|
|
float weighted[10];
|
|
for (uint slot = 0; slot < 10u; slot++) {
|
|
weighted[slot] = qwen_round_bf16(
|
|
experts[experts_base + (ulong)slot * hidden + column] *
|
|
weights[weights_base + slot]);
|
|
}
|
|
// MLX's 10xhidden strided reduction uses eight rows per threadgroup:
|
|
// lanes 0 and 1 first fold rows 8 and 9, then lane 0 folds lanes 1..7.
|
|
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)]],
|
|
device const uchar *packed [[buffer(5)]],
|
|
device const uchar *scales [[buffer(6)]],
|
|
device const uchar *biases [[buffer(7)]],
|
|
uint column [[thread_position_in_grid]]) {
|
|
if (column >= args.u[0]) return;
|
|
out[column] = qwen_round_bf16(qwen_quant_weight(
|
|
packed, scales, biases, args.u[13], args.u[14], args.u[15],
|
|
args.u[4], column, args.u[0], args.u[2], args.u[3]));
|
|
}
|
|
|
|
kernel void kernel_qwen_affine_embedding_batch(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const int *tokens [[buffer(2)]],
|
|
device const uchar *packed [[buffer(5)]],
|
|
device const uchar *scales [[buffer(6)]],
|
|
device const uchar *biases [[buffer(7)]],
|
|
uint2 gid [[thread_position_in_grid]]) {
|
|
const uint column = gid.x;
|
|
const uint row = gid.y;
|
|
if (column >= args.u[0] || row >= args.u[4]) return;
|
|
out[(ulong)row * args.u[0] + column] = qwen_round_bf16(qwen_quant_weight(
|
|
packed, scales, biases, args.u[13], args.u[14], args.u[15],
|
|
(uint)tokens[row], column, args.u[0], args.u[2], args.u[3]));
|
|
}
|
|
|
|
kernel void kernel_qwen_ple_dequant(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const uint *packed [[buffer(2)]],
|
|
device const ushort *scales [[buffer(3)]],
|
|
device const ushort *biases [[buffer(4)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint dim = args.u[0];
|
|
if (index >= dim * args.u[1]) return;
|
|
const uint row = index / dim;
|
|
const uint column = index % dim;
|
|
out[index] = qwen_round_bf16(qwen_quant_value(
|
|
packed, scales, biases, row, column, dim, args.u[2], args.u[3]));
|
|
}
|
|
|
|
kernel void kernel_qwen_bf16_mv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *weights [[buffer(5)]],
|
|
uint row [[thread_position_in_grid]]) {
|
|
if (row >= args.u[1]) return;
|
|
float sum = 0.0f;
|
|
for (uint column = 0; column < args.u[0]; column++) {
|
|
sum = fma(qwen_bf16(qwen_weight_u16(weights, args.u[13],
|
|
(ulong)row * args.u[0] + column)),
|
|
x[column], sum);
|
|
}
|
|
out[row] = qwen_round_bf16(sum);
|
|
}
|
|
|
|
kernel void kernel_qwen_bf16_qmv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *weights [[buffer(5)]],
|
|
uint group [[threadgroup_position_in_grid]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint in_dim = args.u[0];
|
|
const uint out_dim = args.u[1];
|
|
const uint first_row = group * 8u + simd_group * 4u;
|
|
float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
for (uint column = lane; column < in_dim; column += 32u) {
|
|
const float activation = x[column];
|
|
for (uint output = 0; output < 4u; output++) {
|
|
const uint row = first_row + output;
|
|
if (row < out_dim) {
|
|
const float weight = qwen_bf16(qwen_weight_u16(
|
|
weights, args.u[13], (ulong)row * in_dim + column));
|
|
sums[output] = fma(weight, activation, sums[output]);
|
|
}
|
|
}
|
|
}
|
|
for (uint output = 0; output < 4u; output++) {
|
|
const float sum = simd_sum(sums[output]);
|
|
if (lane == 0u && first_row + output < out_dim) {
|
|
out[first_row + output] = qwen_round_bf16(sum);
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_bf16_gemv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *weights [[buffer(5)]],
|
|
threadgroup float *partials [[threadgroup(0)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
ushort lane [[thread_index_in_simdgroup]],
|
|
ushort simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
const uint in_dim = args.u[0];
|
|
const uint out_dim = args.u[1];
|
|
const uint mode = args.u[4];
|
|
device const float *input = x + (ulong)group.y * in_dim;
|
|
device float *output = out + (ulong)group.y * out_dim;
|
|
float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
|
|
if (mode == 2u) {
|
|
const uint first_row = group.x * 4u;
|
|
uint column = (uint)simd_group * 128u + (uint)lane * 4u;
|
|
for (; column < in_dim; column += 1024u) {
|
|
for (uint row = 0; row < 4u; row++) {
|
|
if (first_row + row < out_dim) {
|
|
const ulong base = (ulong)(first_row + row) * in_dim + column;
|
|
sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base)) * input[column];
|
|
sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base + 1u)) * input[column + 1u];
|
|
sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base + 2u)) * input[column + 2u];
|
|
sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base + 3u)) * input[column + 3u];
|
|
}
|
|
}
|
|
}
|
|
for (uint row = 0; row < 4u; row++) {
|
|
for (ushort offset = 16u; offset >= 1u; offset >>= 1u) {
|
|
sums[row] += simd_shuffle_down(sums[row], offset);
|
|
}
|
|
if (lane == 0u) partials[(uint)simd_group * 8u + row] = sums[row];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (simd_group == 0u && lane == 0u) {
|
|
for (uint row = 0; row < 4u && first_row + row < out_dim; row++) {
|
|
float value = partials[row];
|
|
for (uint part = 1u; part < 8u; part++) value += partials[part * 8u + row];
|
|
output[first_row + row] = qwen_round_bf16(value);
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
if (mode == 3u) {
|
|
const uint row_group = lane / 4u;
|
|
const uint column_lane = lane % 4u;
|
|
const uint first_row = group.x * 32u + row_group * 4u;
|
|
for (uint column = column_lane * 4u; column < in_dim; column += 16u) {
|
|
const uint valid = min(4u, in_dim - column);
|
|
for (uint row = 0; row < 4u && first_row + row < out_dim; row++) {
|
|
const ulong base = (ulong)(first_row + row) * in_dim + column;
|
|
for (uint i = 0; i < valid; i++) {
|
|
sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base + i)) * input[column + i];
|
|
}
|
|
}
|
|
}
|
|
for (uint row = 0; row < 4u; row++) {
|
|
sums[row] += simd_shuffle_down(sums[row], 2u);
|
|
sums[row] += simd_shuffle_down(sums[row], 1u);
|
|
if (column_lane == 0u && first_row + row < out_dim) {
|
|
output[first_row + row] = qwen_round_bf16(sums[row]);
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
const uint rows_per_group = mode == 1u ? 32u : 16u;
|
|
const uint first_row = group.x * rows_per_group + (uint)simd_group * 4u;
|
|
for (uint column = (uint)lane * 4u; column < in_dim; column += 128u) {
|
|
const uint valid = min(4u, in_dim - column);
|
|
for (uint row = 0; row < 4u && first_row + row < out_dim; row++) {
|
|
const ulong base = (ulong)(first_row + row) * in_dim + column;
|
|
for (uint i = 0; i < valid; i++) {
|
|
sums[row] += qwen_bf16(qwen_weight_u16(weights, args.u[13], base + i)) * input[column + i];
|
|
}
|
|
}
|
|
}
|
|
for (uint row = 0; row < 4u; row++) {
|
|
for (ushort offset = 16u; offset >= 1u; offset >>= 1u) {
|
|
sums[row] += simd_shuffle_down(sums[row], offset);
|
|
}
|
|
if (lane == 0u && first_row + row < out_dim) {
|
|
output[first_row + row] = qwen_round_bf16(sums[row]);
|
|
}
|
|
}
|
|
}
|
|
|
|
// out[4, N] = x[4, K] @ weight[N, K]^T. This is the decode shape used by
|
|
// the MTP hidden projection: stream each BF16 weight row once for all four
|
|
// hyper streams and preserve MLX's vec4 accumulation and reduction order.
|
|
kernel void kernel_qwen_bf16_gemv4(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *weights [[buffer(5)]],
|
|
uint group [[threadgroup_position_in_grid]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint in_dim = args.u[0];
|
|
const uint out_dim = args.u[1];
|
|
const uint row = group * 4u + simd_group;
|
|
if (row >= out_dim) return;
|
|
|
|
constexpr uint k_lanes = 32u;
|
|
constexpr uint unroll = 8u;
|
|
const uint vectors = in_dim / 4u;
|
|
const uint main = vectors - vectors % (k_lanes * unroll);
|
|
const ulong weight_row = (ulong)row * in_dim;
|
|
float result[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
|
|
for (uint base = 0u; base < main; base += k_lanes * unroll) {
|
|
float accumulated[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
for (uint i = 0u; i < unroll; i++) {
|
|
const uint column = (base + i * k_lanes + lane) * 4u;
|
|
const float4 weight = float4(
|
|
qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column)),
|
|
qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 1u)),
|
|
qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 2u)),
|
|
qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 3u)));
|
|
for (uint stream = 0u; stream < 4u; stream++) {
|
|
accumulated[stream] += dot(weight, *((device const float4 *)(
|
|
x + (ulong)stream * in_dim + column)));
|
|
}
|
|
}
|
|
for (uint stream = 0u; stream < 4u; stream++) {
|
|
result[stream] += accumulated[stream];
|
|
}
|
|
}
|
|
for (uint vector = main + lane; vector < vectors; vector += k_lanes) {
|
|
const uint column = vector * 4u;
|
|
const float4 weight = float4(
|
|
qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column)),
|
|
qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 1u)),
|
|
qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 2u)),
|
|
qwen_bf16(qwen_weight_u16(weights, args.u[13], weight_row + column + 3u)));
|
|
for (uint stream = 0u; stream < 4u; stream++) {
|
|
result[stream] += dot(weight, *((device const float4 *)(
|
|
x + (ulong)stream * in_dim + column)));
|
|
}
|
|
}
|
|
for (uint stream = 0u; stream < 4u; stream++) {
|
|
for (ushort offset = 16u; offset >= 1u; offset >>= 1u) {
|
|
result[stream] += simd_shuffle_down(result[stream], offset);
|
|
}
|
|
if (lane == 0u) {
|
|
out[(ulong)stream * out_dim + row] = qwen_round_bf16(result[stream]);
|
|
}
|
|
}
|
|
}
|
|
|
|
// MLX affine quantization, specialized for the private 8-bit/group-64 hyper
|
|
// packs. Rust owns pack creation and lifetime; this kernel only transforms the
|
|
// source BF16 matrices already admitted by the model loader.
|
|
kernel void kernel_qwen_quantize_hyper_q8(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device uchar *pack [[buffer(1)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
constexpr uint quant_group = 64u;
|
|
constexpr float bins = 255.0f;
|
|
const uint columns = args.u[0];
|
|
const uint source_rows = args.u[1];
|
|
const uint destination_row = args.u[2] + group.y;
|
|
const uint destination_rows = args.u[3];
|
|
if (group.y >= source_rows || group.x >= columns / quant_group) return;
|
|
|
|
float first = qwen_bf16(qwen_weight_u16(
|
|
weight, args.u[13], (ulong)group.y * columns + group.x * quant_group + lane));
|
|
float second = qwen_bf16(qwen_weight_u16(
|
|
weight, args.u[13], (ulong)group.y * columns + group.x * quant_group + lane + 32u));
|
|
const float minimum = simd_min(min(first, second));
|
|
const float maximum = simd_max(max(first, second));
|
|
float scale = max((maximum - minimum) / bins, 1e-7f);
|
|
const bool negative_side = abs(minimum) > abs(maximum);
|
|
scale = negative_side ? scale : -scale;
|
|
const float edge = negative_side ? minimum : maximum;
|
|
const float zero = round(edge / scale);
|
|
const bool at_zero = zero == 0.0f;
|
|
scale = at_zero ? scale : edge / zero;
|
|
const float bias = at_zero ? 0.0f : edge;
|
|
|
|
const ulong row = destination_row;
|
|
const ulong q_index = row * columns + (ulong)group.x * quant_group + lane;
|
|
pack[q_index] = (uchar)clamp(round((first - bias) / scale), 0.0f, bins);
|
|
pack[q_index + 32u] = (uchar)clamp(round((second - bias) / scale), 0.0f, bins);
|
|
if (lane == 0u) {
|
|
const ulong values = (ulong)destination_rows * columns;
|
|
const ulong groups_per_row = columns / quant_group;
|
|
const ulong parameter = row * groups_per_row + group.x;
|
|
device ushort *scales = (device ushort *)(pack + values);
|
|
device ushort *biases = scales + destination_rows * groups_per_row;
|
|
scales[parameter] = qwen_to_bf16(scale);
|
|
biases[parameter] = qwen_to_bf16(bias);
|
|
}
|
|
}
|
|
|
|
// Faithful port of MTPLX hyper_connection_v3 R1: four grouped RMS reductions,
|
|
// q8 down+inject projection, then the family's activation boundaries.
|
|
kernel void kernel_qwen_hyper_v3_r1(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *mix_out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *pack [[buffer(3)]],
|
|
device float *inject_out [[buffer(4)]],
|
|
device const uchar *norm_weight [[buffer(5)]],
|
|
device float *rms_out [[buffer(8)]],
|
|
uint3 group [[threadgroup_position_in_grid]],
|
|
uint3 thread_position [[thread_position_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr uint hidden = 2560u;
|
|
constexpr uint width = 10240u;
|
|
constexpr uint rank = 320u;
|
|
constexpr uint rows = 324u;
|
|
constexpr uint quant_group = 64u;
|
|
constexpr uint groups_per_row = width / quant_group;
|
|
const uint tid = thread_position.x;
|
|
const ulong batch_width = (ulong)group.y * width;
|
|
const ulong batch_rank = (ulong)group.y * rank;
|
|
const ulong batch_streams = (ulong)group.y * 4u;
|
|
threadgroup float group_sums[4];
|
|
threadgroup float partial[32];
|
|
|
|
float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
for (uint index = tid; index < width; index += 1024u) {
|
|
const float value = x[batch_width + index];
|
|
sums[index / hidden] += value * value;
|
|
}
|
|
for (uint stream = 0u; stream < 4u; stream++) {
|
|
const float value = simd_sum(sums[stream]);
|
|
if (lane == 0u) partial[simd_group] = value;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (simd_group == 0u) {
|
|
float total = partial[lane];
|
|
total = simd_sum(total);
|
|
if (lane == 0u) group_sums[stream] = total;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
float rms[4];
|
|
for (uint stream = 0u; stream < 4u; stream++) {
|
|
rms[stream] = rsqrt(group_sums[stream] / (float)hidden + 1e-6f);
|
|
}
|
|
if (group.x == 0u && tid < 4u) rms_out[batch_streams + tid] = rms[tid];
|
|
|
|
const uint row = group.x * 32u + simd_group;
|
|
if (row >= rows) return;
|
|
constexpr ulong q_bytes = (ulong)rows * width;
|
|
constexpr ulong parameter_count = (ulong)rows * groups_per_row;
|
|
device const uint *quantized = (device const uint *)pack;
|
|
device const ushort *scales = (device const ushort *)(pack + q_bytes);
|
|
device const ushort *biases = scales + parameter_count;
|
|
device const ushort *norm =
|
|
(device const ushort *)(norm_weight + args.u[13]);
|
|
float accumulator = 0.0f;
|
|
for (uint quant = lane; quant < groups_per_row; quant += 32u) {
|
|
const float scale = qwen_bf16(scales[(ulong)row * groups_per_row + quant]);
|
|
const float bias = qwen_bf16(biases[(ulong)row * groups_per_row + quant]);
|
|
const ulong base = (ulong)row * width + (ulong)quant * quant_group;
|
|
const uint activation_base = quant * quant_group;
|
|
float quantized_sum = 0.0f;
|
|
float normalized_sum = 0.0f;
|
|
for (uint word_index = 0u; word_index < 16u; word_index++) {
|
|
const uint word = quantized[(base >> 2u) + word_index];
|
|
for (uint byte = 0u; byte < 4u; byte++) {
|
|
const uint index = activation_base + word_index * 4u + byte;
|
|
const float normalized = x[batch_width + index] *
|
|
qwen_bf16(norm[index]) *
|
|
rms[index / hidden];
|
|
quantized_sum += (float)((word >> (byte * 8u)) & 255u) * normalized;
|
|
normalized_sum += normalized;
|
|
}
|
|
}
|
|
accumulator += scale * quantized_sum + bias * normalized_sum;
|
|
}
|
|
accumulator = simd_sum(accumulator);
|
|
if (lane == 0u) {
|
|
const float value = accumulator * 0.25f;
|
|
if (row < rank) {
|
|
mix_out[batch_rank + row] = qwen_round_bf16(value / (1.0f + exp(-value)));
|
|
} else {
|
|
inject_out[batch_streams + row - rank] =
|
|
qwen_round_bf16(2.0f / (1.0f + exp(-value)));
|
|
}
|
|
}
|
|
}
|
|
|
|
// Faithful port of MTPLX hyper_connection_v3 R2: q8 up projection followed by
|
|
// the four-stream gated mean. One SIMD group owns each output dimension.
|
|
kernel void kernel_qwen_hyper_v3_r2(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *pack [[buffer(3)]],
|
|
device const float *mix [[buffer(4)]],
|
|
device const uchar *norm_weight [[buffer(5)]],
|
|
device const float *rms [[buffer(8)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint lane [[thread_index_in_simdgroup]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr uint rank = 320u;
|
|
constexpr uint hidden = 2560u;
|
|
constexpr uint rows = 10240u;
|
|
constexpr uint quant_group = 64u;
|
|
constexpr uint groups_per_row = rank / quant_group;
|
|
const uint dimension = group.x * 32u + simd_group;
|
|
if (dimension >= hidden) return;
|
|
constexpr ulong q_bytes = (ulong)rows * rank;
|
|
constexpr ulong parameter_count = (ulong)rows * groups_per_row;
|
|
device const uint *quantized = (device const uint *)pack;
|
|
device const ushort *scales = (device const ushort *)(pack + q_bytes);
|
|
device const ushort *biases = scales + parameter_count;
|
|
device const ushort *norm =
|
|
(device const ushort *)(norm_weight + args.u[13]);
|
|
float mixed = 0.0f;
|
|
for (uint stream = 0u; stream < 4u; stream++) {
|
|
const uint row = stream * hidden + dimension;
|
|
float dot = 0.0f;
|
|
for (uint quant = 0u; quant < groups_per_row; quant++) {
|
|
const float scale = qwen_bf16(scales[(ulong)row * groups_per_row + quant]);
|
|
const float bias = qwen_bf16(biases[(ulong)row * groups_per_row + quant]);
|
|
const ulong base = (ulong)row * rank + (ulong)quant * quant_group;
|
|
const uint activation_base = quant * quant_group;
|
|
float quantized_sum = 0.0f;
|
|
float activation_sum = 0.0f;
|
|
for (uint word_index = lane; word_index < 16u; word_index += 32u) {
|
|
const uint word = quantized[(base >> 2u) + word_index];
|
|
for (uint byte = 0u; byte < 4u; byte++) {
|
|
const float value = mix[(ulong)group.y * rank + activation_base + word_index * 4u + byte];
|
|
quantized_sum += (float)((word >> (byte * 8u)) & 255u) * value;
|
|
activation_sum += value;
|
|
}
|
|
}
|
|
dot += scale * quantized_sum + bias * activation_sum;
|
|
}
|
|
dot = simd_sum(dot);
|
|
const uint index = stream * hidden + dimension;
|
|
const float normalized = x[(ulong)group.y * rows + index] *
|
|
qwen_bf16(norm[index]) * rms[(ulong)group.y * 4u + stream];
|
|
mixed += (1.0f / (1.0f + exp(-dot))) * normalized;
|
|
}
|
|
if (lane == 0u) {
|
|
out[(ulong)group.y * hidden + dimension] = qwen_round_bf16(mixed * 0.25f);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_repeat4(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint width = args.u[0];
|
|
const uint row_width = width * 4u;
|
|
const uint row = index / row_width;
|
|
if (index < max(args.u[4], 1u) * row_width) {
|
|
out[index] = x[(ulong)row * width + index % width];
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_copy_bf16(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *out [[buffer(1)]],
|
|
device const ushort *input [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] = input[index];
|
|
}
|
|
|
|
kernel void kernel_qwen_zero_rms(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint group [[threadgroup_position_in_grid]],
|
|
uint tid [[thread_index_in_threadgroup]],
|
|
uint threads [[threads_per_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
const uint width = args.u[0];
|
|
const uint group_size = args.u[1];
|
|
if (group >= (width / group_size) * max(args.u[4], 1u)) return;
|
|
const uint start = group * group_size;
|
|
threadgroup float *scale = scratch;
|
|
threadgroup float *partial = scratch + 1u;
|
|
float sum = 0.0f;
|
|
for (uint base = tid * 4u; base < group_size; base += threads * 4u) {
|
|
for (uint i = 0; i < 4u && base + i < group_size; i++) {
|
|
const float value = x[start + base + i];
|
|
sum += value * value;
|
|
}
|
|
}
|
|
sum = simd_sum(sum);
|
|
if (simd_group == 0u) partial[lane] = 0.0f;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (lane == 0u) partial[simd_group] = sum;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (simd_group == 0u) {
|
|
sum = simd_sum(partial[lane]);
|
|
if (lane == 0u) scale[0] = precise::rsqrt(sum / (float)group_size + args.f[0]);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (uint base = tid * 4u; base < group_size; base += threads * 4u) {
|
|
for (uint i = 0; i < 4u && base + i < group_size; i++) {
|
|
const uint index = start + base + i;
|
|
const float normalized = qwen_round_bf16(x[index] * scale[0]);
|
|
float multiplier = qwen_bf16(qwen_weight_u16(weight, args.u[13], index % width));
|
|
if (args.u[11] != 0u) multiplier = qwen_round_bf16(multiplier + 1.0f);
|
|
out[index] = qwen_round_bf16(multiplier * normalized);
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_silu_div4(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index >= args.u[0]) return;
|
|
const float value = qwen_round_bf16(x[index] * 0.25f);
|
|
out[index] = qwen_round_bf16(value * qwen_silu_sigmoid_bf16(value));
|
|
}
|
|
|
|
kernel void kernel_qwen_sigmoid(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] = qwen_sigmoid_bf16(x[index]);
|
|
}
|
|
|
|
kernel void kernel_qwen_sigmoid2_div4(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) {
|
|
const float value = qwen_round_bf16(x[index] * 0.25f);
|
|
out[index] = qwen_round_bf16(2.0f * qwen_sigmoid_bf16(value));
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_hyper_mix(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *normalized [[buffer(2)]],
|
|
device const float *mix [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint hidden = args.u[0];
|
|
if (index >= hidden * max(args.u[4], 1u)) return;
|
|
const uint row = index / hidden;
|
|
const uint column = index % hidden;
|
|
const ulong row_base = (ulong)row * hidden * 4u;
|
|
float value = 0.0f;
|
|
for (uint stream = 0; stream < 4u; stream++) {
|
|
const ulong offset = row_base + (ulong)stream * hidden + column;
|
|
const float product = qwen_round_bf16(normalized[offset] * mix[offset]);
|
|
value = qwen_round_bf16(value + product);
|
|
}
|
|
out[index] = qwen_round_bf16(value * 0.25f);
|
|
}
|
|
|
|
kernel void kernel_qwen_hyper_inject(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *residual [[buffer(2)]],
|
|
device const float *block [[buffer(3)]],
|
|
device const float *gate [[buffer(4)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint hidden = args.u[0];
|
|
const uint row_width = hidden * 4u;
|
|
if (index >= row_width * max(args.u[4], 1u)) return;
|
|
const uint row = index / row_width;
|
|
const uint local = index % row_width;
|
|
const float injected = qwen_round_bf16(
|
|
block[(ulong)row * hidden + local % hidden] *
|
|
gate[(ulong)row * 4u + local / hidden]);
|
|
out[index] = qwen_round_bf16(residual[index] + injected);
|
|
}
|
|
|
|
kernel void kernel_qwen_ple_gate(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *key [[buffer(2)]],
|
|
device const float *query [[buffer(3)]],
|
|
device const float *value [[buffer(4)]],
|
|
threadgroup bfloat *partials [[threadgroup(0)]],
|
|
uint3 group [[threadgroup_position_in_grid]],
|
|
uint3 thread_position [[thread_position_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
const uint hidden = args.u[0];
|
|
const uint tid = thread_position.x;
|
|
const uint stream = group.x;
|
|
if (stream >= 4u) return;
|
|
const ulong base = ((ulong)group.y * 4u + stream) * hidden;
|
|
const ulong value_base = (ulong)group.y * hidden;
|
|
bfloat score = 0.0bf;
|
|
for (uint i = 0; i < 4u; i++) {
|
|
const uint index = tid * 4u + i;
|
|
const bfloat product = (bfloat)((bfloat)key[base + index] * (bfloat)query[base + index]);
|
|
score = product + score;
|
|
}
|
|
score = qwen_simd_sum_bf16(score);
|
|
if (lane == 0u) partials[simd_group] = score;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (simd_group == 0u) {
|
|
score = lane < 20u ? partials[lane] : bfloat(0);
|
|
score = qwen_simd_sum_bf16(score);
|
|
if (lane == 0u) partials[0] = score;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
const bfloat scale = (bfloat)sqrt((float)hidden);
|
|
const float scaled = (float)(bfloat)((float)partials[0] / (float)scale);
|
|
const float transformed = scaled == 0.0f
|
|
? 0.0f
|
|
: qwen_round_bf16(copysign(
|
|
qwen_round_bf16(sqrt(max(abs(scaled), qwen_round_bf16(1.0e-6f)))),
|
|
scaled));
|
|
const float gate = qwen_sigmoid_bf16(transformed);
|
|
for (uint i = tid; i < hidden; i += 640u) {
|
|
out[base + i] = qwen_round_bf16(value[value_base + i] * gate);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_ple_conv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *gated [[buffer(2)]],
|
|
device const float *normalized [[buffer(3)]],
|
|
device const ushort *state [[buffer(4)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
device ushort *state_out [[buffer(8)]],
|
|
uint channel [[thread_position_in_grid]]) {
|
|
if (channel >= args.u[0]) return;
|
|
const ulong width = args.u[0];
|
|
for (uint row = 0; row < max(args.u[4], 1u); row++) {
|
|
const ulong index = (ulong)row * width + channel;
|
|
float value = 0.0f;
|
|
for (uint tap = 0; tap < 4u; tap++) {
|
|
const uint sequence = row + tap * 3u;
|
|
const float input = tap == 3u
|
|
? normalized[index]
|
|
: sequence < 9u
|
|
? qwen_bf16(state[(ulong)sequence * width + channel])
|
|
: normalized[(ulong)(sequence - 9u) * width + channel];
|
|
value += input * qwen_bf16(qwen_weight_u16(
|
|
weight, args.u[13], (ulong)channel * 4u + tap));
|
|
}
|
|
value = qwen_round_bf16(value);
|
|
const float activated = qwen_round_bf16(value * qwen_silu_sigmoid_bf16(value));
|
|
out[index] = qwen_round_bf16(gated[index] + activated);
|
|
}
|
|
for (uint i = 0; i < 9u; i++) {
|
|
const uint sequence = args.u[4] + i;
|
|
state_out[(ulong)i * width + channel] = sequence < 9u
|
|
? state[(ulong)sequence * width + channel]
|
|
: qwen_to_bf16(normalized[(ulong)(sequence - 9u) * width + channel]);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_add(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *a [[buffer(2)]],
|
|
device const float *b [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] = qwen_round_bf16(a[index] + b[index]);
|
|
}
|
|
|
|
kernel void kernel_qwen_zero_words(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device uint *out [[buffer(1)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[(ulong)args.u[1] + index] = 0u;
|
|
}
|
|
|
|
kernel void kernel_qwen_conv_silu(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device ushort *state [[buffer(3)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint channel [[thread_position_in_grid]]) {
|
|
if (channel >= args.u[0]) return;
|
|
const ulong c = channel;
|
|
const ulong width = args.u[0];
|
|
const float value =
|
|
qwen_bf16(state[c]) *
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], c * 4u)) +
|
|
qwen_bf16(state[width + c]) *
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], c * 4u + 1u)) +
|
|
qwen_bf16(state[2u * width + c]) *
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], c * 4u + 2u)) +
|
|
x[c] * qwen_bf16(qwen_weight_u16(weight, args.u[13], c * 4u + 3u));
|
|
state[c] = state[width + c];
|
|
state[width + c] = state[2u * width + c];
|
|
state[2u * width + c] = qwen_to_bf16(x[c]);
|
|
const float rounded = qwen_round_bf16(value);
|
|
out[c] = qwen_round_bf16(rounded * qwen_silu_sigmoid_bf16(rounded));
|
|
}
|
|
|
|
// Exact MTPLX empty-cache decode path: fused depthwise conv, SiLU and
|
|
// per-head q/k normalization. Later rows use kernel_qwen_gdn_step_fused.
|
|
kernel void kernel_qwen_gdn_conv_norm(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device ushort *state [[buffer(3)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
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]]) {
|
|
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;
|
|
if (channel >= width) return;
|
|
|
|
threadgroup float values[1024];
|
|
threadgroup float partial[32];
|
|
const float convolved =
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)) *
|
|
qwen_bf16(state[channel]) +
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 1u)) *
|
|
qwen_bf16(state[width + channel]) +
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 2u)) *
|
|
qwen_bf16(state[2u * width + channel]) +
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 3u)) *
|
|
x[channel];
|
|
const float activated = convolved / (1.0f + exp(-convolved));
|
|
state[channel] = state[width + channel];
|
|
state[width + channel] = state[2u * width + channel];
|
|
state[2u * width + channel] = qwen_to_bf16(x[channel]);
|
|
|
|
if (channel >= 2u * key_width) {
|
|
out[channel] = qwen_round_bf16(activated);
|
|
return;
|
|
}
|
|
|
|
values[tid] = activated;
|
|
float sum = simd_sum(activated * activated);
|
|
if (lane == 0u) partial[simd_group] = sum;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
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);
|
|
out[channel] = channel < key_width
|
|
? qwen_round_bf16(normalized * inverse_scale)
|
|
: qwen_round_bf16(normalized);
|
|
}
|
|
|
|
// MTPLX fused_gdn_conv_norm_rows, including its in-window convolution tail.
|
|
kernel void kernel_qwen_gdn_conv_norm_rows(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *q_out [[buffer(1)]],
|
|
device const float *qkv [[buffer(2)]],
|
|
device float *k_out [[buffer(3)]],
|
|
device float *v_out [[buffer(4)]],
|
|
device const ushort *state [[buffer(5)]],
|
|
device const uchar *weight [[buffer(6)]],
|
|
device ushort *state_out [[buffer(8)]],
|
|
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]]) {
|
|
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;
|
|
if (channel >= width) return;
|
|
|
|
threadgroup float values[1024];
|
|
threadgroup float partial[32];
|
|
const float w0 = qwen_bf16(qwen_weight_u16(weight, args.u[14], (ulong)channel * 4u));
|
|
const float w1 = qwen_bf16(qwen_weight_u16(weight, args.u[14], (ulong)channel * 4u + 1u));
|
|
const float w2 = qwen_bf16(qwen_weight_u16(weight, args.u[14], (ulong)channel * 4u + 2u));
|
|
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 float x0 = row < 3u
|
|
? qwen_bf16(state[(ulong)row * width + channel])
|
|
: qkv[(ulong)(row - 3u) * width + channel];
|
|
const float x1 = row + 1u < 3u
|
|
? qwen_bf16(state[(ulong)(row + 1u) * width + channel])
|
|
: qkv[(ulong)(row - 2u) * width + channel];
|
|
const float x2 = row + 2u < 3u
|
|
? qwen_bf16(state[(ulong)(row + 2u) * width + channel])
|
|
: 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));
|
|
if (value_channel) {
|
|
v_out[(ulong)row * (width - 2u * key_width) + channel - 2u * key_width] =
|
|
qwen_round_bf16(activated);
|
|
continue;
|
|
}
|
|
values[tid] = activated;
|
|
float sum = simd_sum(activated * activated);
|
|
if (lane == 0u) partial[simd_group] = sum;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
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);
|
|
if (channel < key_width) {
|
|
q_out[(ulong)row * key_width + channel] =
|
|
qwen_round_bf16(normalized * inverse_scale);
|
|
} else {
|
|
k_out[(ulong)row * key_width + channel - key_width] =
|
|
qwen_round_bf16(normalized);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
|
|
for (uint tail = 0u; tail < 3u; tail++) {
|
|
const uint sequence = args.u[4] + tail;
|
|
const float value = sequence < 3u
|
|
? qwen_bf16(state[(ulong)sequence * width + channel])
|
|
: qkv[(ulong)(sequence - 3u) * width + channel];
|
|
state_out[(ulong)tail * width + channel] = qwen_to_bf16(value);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_gdn_l2norm(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint head [[threadgroup_position_in_grid]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
constexpr uint dim = 128u;
|
|
constexpr uint key_width = 2048u;
|
|
if (head >= 16u) return;
|
|
const ulong qbase = (ulong)head * dim;
|
|
const ulong kbase = key_width + qbase;
|
|
float qv[4];
|
|
float kv[4];
|
|
float qsum = 0.0f;
|
|
float ksum = 0.0f;
|
|
for (uint i = 0; i < 4u; i++) {
|
|
const uint element = lane * 4u + i;
|
|
qv[i] = x[qbase + element];
|
|
kv[i] = x[kbase + element];
|
|
qsum += qv[i] * qv[i];
|
|
ksum += kv[i] * kv[i];
|
|
}
|
|
qsum = simd_sum(qsum);
|
|
ksum = simd_sum(ksum);
|
|
const float qinv = precise::rsqrt(qsum + 1.0e-6f);
|
|
const float kinv = precise::rsqrt(ksum + 1.0e-6f);
|
|
for (uint i = 0; i < 4u; i++) {
|
|
const uint element = lane * 4u + i;
|
|
const float normalized_q = qwen_round_bf16(qv[i] * qinv);
|
|
out[qbase + element] = qwen_round_bf16(normalized_q * 0.08837890625f);
|
|
out[kbase + element] = qwen_round_bf16(kv[i] * kinv);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_gdn_controls(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *controls [[buffer(1)]],
|
|
device const uchar *a_log [[buffer(5)]],
|
|
device const uchar *dt_bias [[buffer(6)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
constexpr uint heads = 48u;
|
|
constexpr uint controls_width = 6240u;
|
|
const uint rows = max(args.u[4], 1u);
|
|
if (index >= heads * rows) return;
|
|
const uint row = index / heads;
|
|
const uint head = index % heads;
|
|
const ulong base = (ulong)row * controls_width;
|
|
const float step = qwen_round_bf16(
|
|
controls[base + args.u[1] + head] +
|
|
qwen_bf16(qwen_weight_u16(dt_bias, args.u[14], head)));
|
|
const float maximum = max(step, 0.0f);
|
|
const float minimum = min(step, 0.0f);
|
|
const float exponential = qwen_round_bf16(
|
|
precise::exp(qwen_round_bf16(minimum - maximum)));
|
|
const float plus_one = 1.0f + exponential;
|
|
const float logarithm = plus_one == 1.0f
|
|
? exponential
|
|
: exponential * (precise::log(plus_one) / (plus_one - 1.0f));
|
|
const float softplus = qwen_round_bf16(
|
|
maximum + qwen_round_bf16(logarithm));
|
|
const float rate = precise::exp(
|
|
qwen_bf16(qwen_weight_u16(a_log, args.u[13], head)));
|
|
controls[base + args.u[1] + head] = precise::exp(-rate * softplus);
|
|
controls[base + args.u[0] + head] =
|
|
qwen_sigmoid_bf16(controls[base + args.u[0] + head]);
|
|
}
|
|
|
|
kernel void kernel_qwen_pack_gdn_controls(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *controls [[buffer(1)]],
|
|
device const float *z [[buffer(2)]],
|
|
device const float *b [[buffer(3)]],
|
|
device const float *a [[buffer(4)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
constexpr uint value_width = 6144u;
|
|
constexpr uint heads = 48u;
|
|
constexpr uint controls_width = value_width + 2u * heads;
|
|
if (index >= args.u[4] * controls_width) return;
|
|
const uint row = index / controls_width;
|
|
const uint column = index % controls_width;
|
|
if (column < value_width) {
|
|
controls[index] = z[(ulong)row * value_width + column];
|
|
} else if (column < value_width + heads) {
|
|
controls[index] = b[(ulong)row * heads + column - value_width];
|
|
} else {
|
|
controls[index] = a[(ulong)row * heads + column - value_width - heads];
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_gdn_delta_packed(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *qkv [[buffer(2)]],
|
|
device const float *controls [[buffer(3)]],
|
|
device float *state [[buffer(4)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint lane [[thread_index_in_simdgroup]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr uint dim = 128u;
|
|
constexpr uint key_width = 2048u;
|
|
constexpr uint rows_per_simdgroup = 8u;
|
|
constexpr uint values_per_lane = 32u;
|
|
constexpr uint partials_per_lane = 8u;
|
|
const uint head = group.y;
|
|
const uint row_group = group.x * 2u + simd_group;
|
|
const uint row_in_simdgroup = lane / 4u;
|
|
const uint lane_in_row = lane & 3u;
|
|
const uint value_index = row_group * rows_per_simdgroup + row_in_simdgroup;
|
|
const uint key_head = head / 3u;
|
|
const uint key_start = lane_in_row * values_per_lane;
|
|
device const float *q = qkv + (ulong)key_head * dim + key_start;
|
|
device const float *k = qkv + key_width + (ulong)key_head * dim + key_start;
|
|
device const float *v = qkv + 2u * key_width + (ulong)head * dim;
|
|
device float *row = state + ((ulong)head * dim + value_index) * dim + key_start;
|
|
float values[values_per_lane];
|
|
float partial[partials_per_lane];
|
|
const float decay = controls[args.u[1] + head];
|
|
for (uint i = 0; i < values_per_lane; i++) values[i] = row[i];
|
|
for (uint block = 0; block < partials_per_lane; block++) {
|
|
float sum = 0.0f;
|
|
for (uint i = 0; i < 4u; i++) {
|
|
const uint element = block * 4u + i;
|
|
values[element] *= decay;
|
|
sum += values[element] * k[element];
|
|
}
|
|
partial[block] = sum;
|
|
}
|
|
float prediction =
|
|
((partial[0] + partial[1]) + (partial[2] + partial[3])) +
|
|
((partial[4] + partial[5]) + (partial[6] + partial[7]));
|
|
prediction += simd_shuffle_xor(prediction, 1u);
|
|
prediction += simd_shuffle_xor(prediction, 2u);
|
|
const float delta =
|
|
(v[value_index] - prediction) * controls[args.u[0] + head];
|
|
for (uint block = 0; block < partials_per_lane; block++) {
|
|
float sum = 0.0f;
|
|
for (uint i = 0; i < 4u; i++) {
|
|
const uint element = block * 4u + i;
|
|
values[element] += k[element] * delta;
|
|
sum += values[element] * q[element];
|
|
}
|
|
partial[block] = sum;
|
|
}
|
|
float result =
|
|
((partial[0] + partial[1]) + (partial[2] + partial[3])) +
|
|
((partial[4] + partial[5]) + (partial[6] + partial[7]));
|
|
result += simd_shuffle_xor(result, 1u);
|
|
result += simd_shuffle_xor(result, 2u);
|
|
for (uint i = 0; i < values_per_lane; i++) row[i] = values[i];
|
|
if (lane_in_row == 0u) {
|
|
out[(ulong)head * dim + value_index] = qwen_round_bf16(result);
|
|
}
|
|
}
|
|
|
|
// mlx-lm gated_delta_step_packed_btree for S=2..4 verifier rows.
|
|
kernel void kernel_qwen_gdn_delta_rows(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *values_out [[buffer(1)]],
|
|
device const float *q_rows [[buffer(2)]],
|
|
device const float *k_rows [[buffer(3)]],
|
|
device const float *controls [[buffer(4)]],
|
|
device const float *state [[buffer(5)]],
|
|
device float *state_out [[buffer(8)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint lane [[thread_index_in_simdgroup]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr uint dim = 128u;
|
|
constexpr uint key_width = 2048u;
|
|
constexpr uint value_width = 6144u;
|
|
constexpr uint controls_width = 6240u;
|
|
constexpr uint rows_per_simdgroup = 8u;
|
|
constexpr uint values_per_lane = 32u;
|
|
constexpr uint partials_per_lane = 8u;
|
|
const uint head = group.y;
|
|
const uint row_group = group.x * 2u + simd_group;
|
|
const uint row_in_simdgroup = lane / 4u;
|
|
const uint lane_in_row = lane & 3u;
|
|
const uint value_index = row_group * rows_per_simdgroup + row_in_simdgroup;
|
|
const uint key_head = head / 3u;
|
|
const uint key_start = lane_in_row * values_per_lane;
|
|
device const float *q = q_rows + (ulong)key_head * dim + key_start;
|
|
device const float *k = k_rows + (ulong)key_head * dim + key_start;
|
|
device const float *state_row =
|
|
state + ((ulong)head * dim + value_index) * dim + key_start;
|
|
device float *state_out_row =
|
|
state_out + ((ulong)head * dim + value_index) * dim + key_start;
|
|
float values[values_per_lane];
|
|
float partial[partials_per_lane];
|
|
for (uint index = 0u; index < values_per_lane; index++) {
|
|
values[index] = state_row[index];
|
|
}
|
|
|
|
for (uint token = 0u; token < args.u[4]; token++) {
|
|
const ulong control_base = (ulong)token * controls_width;
|
|
const float decay = controls[control_base + args.u[1] + head];
|
|
for (uint block = 0u; block < partials_per_lane; block++) {
|
|
float sum = 0.0f;
|
|
for (uint index = 0u; index < 4u; index++) {
|
|
const uint element = block * 4u + index;
|
|
values[element] *= decay;
|
|
sum += values[element] * k[element];
|
|
}
|
|
partial[block] = sum;
|
|
}
|
|
float prediction =
|
|
((partial[0] + partial[1]) + (partial[2] + partial[3])) +
|
|
((partial[4] + partial[5]) + (partial[6] + partial[7]));
|
|
prediction += simd_shuffle_xor(prediction, 1u);
|
|
prediction += simd_shuffle_xor(prediction, 2u);
|
|
const ulong value_offset =
|
|
(ulong)token * value_width + (ulong)head * dim + value_index;
|
|
const float delta = (values_out[value_offset] - prediction) *
|
|
controls[control_base + args.u[0] + head];
|
|
for (uint block = 0u; block < partials_per_lane; block++) {
|
|
float sum = 0.0f;
|
|
for (uint index = 0u; index < 4u; index++) {
|
|
const uint element = block * 4u + index;
|
|
values[element] += k[element] * delta;
|
|
sum += values[element] * q[element];
|
|
}
|
|
partial[block] = sum;
|
|
}
|
|
float result =
|
|
((partial[0] + partial[1]) + (partial[2] + partial[3])) +
|
|
((partial[4] + partial[5]) + (partial[6] + partial[7]));
|
|
result += simd_shuffle_xor(result, 1u);
|
|
result += simd_shuffle_xor(result, 2u);
|
|
if (lane_in_row == 0u) {
|
|
values_out[value_offset] = qwen_round_bf16(result);
|
|
}
|
|
q += key_width;
|
|
k += key_width;
|
|
}
|
|
for (uint index = 0u; index < values_per_lane; index++) {
|
|
state_out_row[index] = values[index];
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_gdn_step_fused(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *qkv [[buffer(2)]],
|
|
device const float *controls [[buffer(3)]],
|
|
device float *state [[buffer(4)]],
|
|
device const uchar *conv_weight [[buffer(5)]],
|
|
device const uchar *a_log [[buffer(6)]],
|
|
device const uchar *dt_bias [[buffer(7)]],
|
|
device ushort *conv_state [[buffer(8)]],
|
|
device const uchar *norm_weight [[buffer(9)]],
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
ushort tid [[thread_index_in_threadgroup]],
|
|
ushort lane [[thread_index_in_simdgroup]],
|
|
ushort simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr uint dim = 128u;
|
|
constexpr uint key_width = 2048u;
|
|
constexpr uint conv_width = 10240u;
|
|
constexpr uint values_per_key = 3u;
|
|
constexpr float inverse_scale = 0.08838834764831845f;
|
|
const uint head = group.y;
|
|
const uint key_head = head / values_per_key;
|
|
threadgroup float *q = scratch;
|
|
threadgroup float *k = q + dim;
|
|
threadgroup float *v = k + dim;
|
|
threadgroup float *result = v + dim;
|
|
threadgroup float *partial = result + dim;
|
|
device ushort *conv_next = (device ushort *)(out + 6144u);
|
|
|
|
for (uint local = tid; local < 3u * dim; local += 256u) {
|
|
uint channel;
|
|
if (local < dim) {
|
|
channel = key_head * dim + local;
|
|
} else if (local < 2u * dim) {
|
|
channel = key_width + key_head * dim + local - dim;
|
|
} else {
|
|
channel = 2u * key_width + head * dim + local - 2u * dim;
|
|
}
|
|
const float value =
|
|
qwen_bf16(qwen_weight_u16(conv_weight, args.u[13], (ulong)channel * 4u)) *
|
|
qwen_bf16(conv_state[channel]) +
|
|
qwen_bf16(qwen_weight_u16(conv_weight, args.u[13], (ulong)channel * 4u + 1u)) *
|
|
qwen_bf16(conv_state[conv_width + channel]) +
|
|
qwen_bf16(qwen_weight_u16(conv_weight, args.u[13], (ulong)channel * 4u + 2u)) *
|
|
qwen_bf16(conv_state[2u * conv_width + channel]) +
|
|
qwen_bf16(qwen_weight_u16(conv_weight, args.u[13], (ulong)channel * 4u + 3u)) *
|
|
qkv[channel];
|
|
const float activated = value / (1.0f + exp(-value));
|
|
if (local >= 2u * dim || head % values_per_key == 0u) {
|
|
conv_next[channel] = conv_state[conv_width + channel];
|
|
conv_next[conv_width + channel] = conv_state[2u * conv_width + channel];
|
|
conv_next[2u * conv_width + channel] = qwen_to_bf16(qkv[channel]);
|
|
}
|
|
if (local < dim) {
|
|
q[local] = activated;
|
|
} else if (local < 2u * dim) {
|
|
k[local - dim] = activated;
|
|
} else {
|
|
v[local - 2u * dim] = qwen_round_bf16(activated);
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float squared = 0.0f;
|
|
if (tid < dim) {
|
|
squared = q[tid] * q[tid];
|
|
} else if (tid < 2u * dim) {
|
|
squared = k[tid - dim] * k[tid - dim];
|
|
}
|
|
squared = simd_sum(squared);
|
|
if (lane == 0u) partial[simd_group] = squared;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
const float q_scale =
|
|
rsqrt(partial[0] + partial[1] + partial[2] + partial[3] + 1.0e-6f) *
|
|
inverse_scale;
|
|
const float k_scale = rsqrt(partial[4] + partial[5] + partial[6] + partial[7] + 1.0e-6f);
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (tid < dim) {
|
|
q[tid] = qwen_round_bf16(q[tid] * q_scale);
|
|
} else if (tid < 2u * dim) {
|
|
k[tid - dim] = qwen_round_bf16(k[tid - dim] * k_scale);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
const float step = controls[args.u[1] + head] +
|
|
qwen_bf16(qwen_weight_u16(dt_bias, args.u[15], head));
|
|
const float softplus = max(step, 0.0f) + log(1.0f + exp(-abs(step)));
|
|
const float decay = exp(
|
|
-exp(qwen_bf16(qwen_weight_u16(a_log, args.u[14], head))) * softplus);
|
|
const float beta = 1.0f / (1.0f + exp(-controls[args.u[0] + head]));
|
|
|
|
for (uint value_index = simd_group; value_index < dim; value_index += 8u) {
|
|
const ulong state_row = ((ulong)head * dim + value_index) * dim;
|
|
const uint column = 4u * lane;
|
|
float s0 = state[state_row + column] * decay;
|
|
float s1 = state[state_row + column + 1u] * decay;
|
|
float s2 = state[state_row + column + 2u] * decay;
|
|
float s3 = state[state_row + column + 3u] * decay;
|
|
const float k0 = k[column];
|
|
const float k1 = k[column + 1u];
|
|
const float k2 = k[column + 2u];
|
|
const float k3 = k[column + 3u];
|
|
float prediction = s0 * k0 + s1 * k1 + s2 * k2 + s3 * k3;
|
|
prediction = simd_sum(prediction);
|
|
const float delta = (v[value_index] - prediction) * beta;
|
|
s0 += k0 * delta;
|
|
s1 += k1 * delta;
|
|
s2 += k2 * delta;
|
|
s3 += k3 * delta;
|
|
float value = s0 * q[column] + s1 * q[column + 1u] +
|
|
s2 * q[column + 2u] + s3 * q[column + 3u];
|
|
value = simd_sum(value);
|
|
state[state_row + column] = s0;
|
|
state[state_row + column + 1u] = s1;
|
|
state[state_row + column + 2u] = s2;
|
|
state[state_row + column + 3u] = s3;
|
|
if (lane == 0u) {
|
|
result[value_index] = qwen_round_bf16(value);
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float output_squared = tid < dim ? result[tid] * result[tid] : 0.0f;
|
|
output_squared = simd_sum(output_squared);
|
|
if (lane == 0u && simd_group < 4u) partial[simd_group] = output_squared;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
const float output_scale = rsqrt(
|
|
(partial[0] + partial[1] + partial[2] + partial[3]) / (float)dim + 1.0e-6f);
|
|
if (tid < dim) {
|
|
const float normalized = qwen_round_bf16(
|
|
result[tid] * output_scale *
|
|
qwen_bf16(qwen_weight_u16(norm_weight, args.u[9], tid)));
|
|
const float z = controls[(ulong)head * dim + tid];
|
|
const float gate = 1.0f / (1.0f + exp(-z));
|
|
out[(ulong)head * dim + tid] = qwen_round_bf16(gate * normalized);
|
|
}
|
|
}
|
|
|
|
// MTPLX verify-width GDN: convolution is evaluated against the immutable
|
|
// pre-window tail plus all earlier rows, while gated-delta recurrence advances
|
|
// sequentially through S=2..4 in one head-owned threadgroup.
|
|
kernel void kernel_qwen_gdn_verify(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *qkv [[buffer(2)]],
|
|
device const float *controls [[buffer(3)]],
|
|
device float *state [[buffer(4)]],
|
|
device const uchar *conv_weight [[buffer(5)]],
|
|
device const uchar *a_log [[buffer(6)]],
|
|
device const uchar *dt_bias [[buffer(7)]],
|
|
device const ushort *conv_state [[buffer(8)]],
|
|
device const uchar *norm_weight [[buffer(9)]],
|
|
threadgroup float *scratch [[threadgroup(0)]],
|
|
uint head [[threadgroup_position_in_grid]],
|
|
ushort tid [[thread_index_in_threadgroup]],
|
|
ushort lane [[thread_index_in_simdgroup]],
|
|
ushort simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr uint dim = 128u;
|
|
constexpr uint key_width = 2048u;
|
|
constexpr uint conv_width = 10240u;
|
|
constexpr uint values_per_key = 3u;
|
|
constexpr uint controls_width = 6240u;
|
|
constexpr uint value_width = 6144u;
|
|
constexpr float inverse_scale = 0.08838834764831845f;
|
|
const uint key_head = head / values_per_key;
|
|
threadgroup float *q = scratch;
|
|
threadgroup float *k = q + dim;
|
|
threadgroup float *v = k + dim;
|
|
threadgroup float *result = v + dim;
|
|
threadgroup float *partial = result + dim;
|
|
|
|
for (uint token = 0u; token < args.u[4]; token++) {
|
|
for (uint local = tid; local < 3u * dim; local += 256u) {
|
|
uint channel;
|
|
if (local < dim) {
|
|
channel = key_head * dim + local;
|
|
} else if (local < 2u * dim) {
|
|
channel = key_width + key_head * dim + local - dim;
|
|
} else {
|
|
channel = 2u * key_width + head * dim + local - 2u * dim;
|
|
}
|
|
float convolved = 0.0f;
|
|
for (uint tap = 0u; tap < 4u; tap++) {
|
|
const uint sequence = token + tap;
|
|
const float input = sequence < 3u
|
|
? qwen_bf16(conv_state[(ulong)sequence * conv_width + channel])
|
|
: qkv[(ulong)(sequence - 3u) * conv_width + channel];
|
|
convolved += input * qwen_bf16(qwen_weight_u16(
|
|
conv_weight, args.u[13], (ulong)channel * 4u + tap));
|
|
}
|
|
const float activated = convolved / (1.0f + exp(-convolved));
|
|
if (local < dim) {
|
|
q[local] = activated;
|
|
} else if (local < 2u * dim) {
|
|
k[local - dim] = activated;
|
|
} else {
|
|
v[local - 2u * dim] = qwen_round_bf16(activated);
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float squared = 0.0f;
|
|
if (tid < dim) {
|
|
squared = q[tid] * q[tid];
|
|
} else if (tid < 2u * dim) {
|
|
squared = k[tid - dim] * k[tid - dim];
|
|
}
|
|
squared = simd_sum(squared);
|
|
if (lane == 0u) partial[simd_group] = squared;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
const float q_scale =
|
|
rsqrt(partial[0] + partial[1] + partial[2] + partial[3] + 1.0e-6f) *
|
|
inverse_scale;
|
|
const float k_scale =
|
|
rsqrt(partial[4] + partial[5] + partial[6] + partial[7] + 1.0e-6f);
|
|
if (tid < dim) {
|
|
q[tid] = qwen_round_bf16(q[tid] * q_scale);
|
|
} else if (tid < 2u * dim) {
|
|
k[tid - dim] = qwen_round_bf16(k[tid - dim] * k_scale);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
const ulong control_base = (ulong)token * controls_width;
|
|
const float step = controls[control_base + args.u[1] + head] +
|
|
qwen_bf16(qwen_weight_u16(dt_bias, args.u[15], head));
|
|
const float softplus = max(step, 0.0f) + log(1.0f + exp(-abs(step)));
|
|
const float decay = exp(
|
|
-exp(qwen_bf16(qwen_weight_u16(a_log, args.u[14], head))) * softplus);
|
|
const float beta = 1.0f /
|
|
(1.0f + exp(-controls[control_base + args.u[0] + head]));
|
|
|
|
for (uint value_index = simd_group; value_index < dim; value_index += 8u) {
|
|
const ulong state_row = ((ulong)head * dim + value_index) * dim;
|
|
const uint column = 4u * lane;
|
|
float s0 = state[state_row + column] * decay;
|
|
float s1 = state[state_row + column + 1u] * decay;
|
|
float s2 = state[state_row + column + 2u] * decay;
|
|
float s3 = state[state_row + column + 3u] * decay;
|
|
const float k0 = k[column];
|
|
const float k1 = k[column + 1u];
|
|
const float k2 = k[column + 2u];
|
|
const float k3 = k[column + 3u];
|
|
float prediction = s0 * k0 + s1 * k1 + s2 * k2 + s3 * k3;
|
|
prediction = simd_sum(prediction);
|
|
const float delta = (v[value_index] - prediction) * beta;
|
|
s0 += k0 * delta;
|
|
s1 += k1 * delta;
|
|
s2 += k2 * delta;
|
|
s3 += k3 * delta;
|
|
float value = s0 * q[column] + s1 * q[column + 1u] +
|
|
s2 * q[column + 2u] + s3 * q[column + 3u];
|
|
value = simd_sum(value);
|
|
state[state_row + column] = s0;
|
|
state[state_row + column + 1u] = s1;
|
|
state[state_row + column + 2u] = s2;
|
|
state[state_row + column + 3u] = s3;
|
|
if (lane == 0u) result[value_index] = qwen_round_bf16(value);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float output_squared = tid < dim ? result[tid] * result[tid] : 0.0f;
|
|
output_squared = simd_sum(output_squared);
|
|
if (lane == 0u && simd_group < 4u) partial[simd_group] = output_squared;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
const float output_scale = rsqrt(
|
|
(partial[0] + partial[1] + partial[2] + partial[3]) /
|
|
(float)dim + 1.0e-6f);
|
|
if (tid < dim) {
|
|
const float normalized = qwen_round_bf16(
|
|
result[tid] * output_scale *
|
|
qwen_bf16(qwen_weight_u16(norm_weight, args.u[9], tid)));
|
|
const float z = controls[control_base + (ulong)head * dim + tid];
|
|
const float gate = 1.0f / (1.0f + exp(-z));
|
|
out[(ulong)token * value_width + (ulong)head * dim + tid] =
|
|
qwen_round_bf16(gate * normalized);
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_gdn_commit_conv_rows(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *state [[buffer(1)]],
|
|
device const float *qkv [[buffer(2)]],
|
|
uint channel [[thread_position_in_grid]]) {
|
|
const uint width = args.u[0];
|
|
if (channel >= width) return;
|
|
float tail[3];
|
|
for (uint tap = 0u; tap < 3u; tap++) {
|
|
const uint sequence = args.u[4] + tap;
|
|
tail[tap] = sequence < 3u
|
|
? qwen_bf16(state[(ulong)sequence * width + channel])
|
|
: qkv[(ulong)(sequence - 3u) * width + channel];
|
|
}
|
|
for (uint tap = 0u; tap < 3u; tap++) {
|
|
state[(ulong)tap * width + channel] = qwen_to_bf16(tail[tap]);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_gdn_norm_gate(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const float *controls [[buffer(3)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
constexpr uint dim = 128u;
|
|
constexpr uint value_width = 6144u;
|
|
constexpr uint controls_width = 6240u;
|
|
const uint head = group.x;
|
|
const uint row = group.y;
|
|
if (head >= 48u) return;
|
|
const ulong base = (ulong)row * value_width + (ulong)head * dim;
|
|
const ulong control_base = (ulong)row * controls_width + (ulong)head * dim;
|
|
float values[4];
|
|
float sum = 0.0f;
|
|
for (uint i = 0; i < 4u; i++) {
|
|
const uint element = lane * 4u + i;
|
|
values[i] = x[base + element];
|
|
sum += values[i] * values[i];
|
|
}
|
|
sum = simd_sum(sum);
|
|
const float scale = precise::rsqrt(sum / (float)dim + 1.0e-6f);
|
|
for (uint i = 0; i < 4u; i++) {
|
|
const uint element = lane * 4u + i;
|
|
const float normalized = qwen_round_bf16(values[i] * scale);
|
|
const float weighted = qwen_round_bf16(
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], element)) * normalized);
|
|
const float z = controls[control_base + element];
|
|
const float exponential = precise::exp(abs(z));
|
|
const float tail = 1.0f / (1.0f + exponential);
|
|
const float gate = z < 0.0f ? tail : 1.0f - tail;
|
|
out[base + element] = qwen_round_bf16(gate * weighted);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_swiglu(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *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]);
|
|
}
|
|
|
|
kernel void kernel_qwen_unpack_gdn_inputs(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *qkv [[buffer(1)]],
|
|
device const float *fused [[buffer(2)]],
|
|
device float *controls [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
constexpr uint qkv_width = 10240u;
|
|
constexpr uint controls_width = 6240u;
|
|
constexpr uint fused_width = qkv_width + controls_width;
|
|
if (index >= args.u[4] * fused_width) return;
|
|
const uint row = index / fused_width;
|
|
const uint column = index - row * fused_width;
|
|
if (column < qkv_width) {
|
|
qkv[(ulong)row * qkv_width + column] = fused[index];
|
|
} else {
|
|
controls[(ulong)row * controls_width + column - qkv_width] = fused[index];
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_softmax_precise_512(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *input [[buffer(2)]],
|
|
uint row [[threadgroup_position_in_grid]],
|
|
uint lid [[thread_position_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr uint reads = 4u;
|
|
input += (ulong)row * args.u[0];
|
|
out += (ulong)row * args.u[0];
|
|
threadgroup float local_max[32];
|
|
threadgroup float local_normalizer[32];
|
|
float values[reads];
|
|
for (uint i = 0; i < reads; i++) values[i] = input[lid * reads + i];
|
|
if (simd_group == 0u) {
|
|
local_max[lane] = -INFINITY;
|
|
local_normalizer[lane] = 0.0f;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
float maximum = -FLT_MAX;
|
|
for (uint i = 0; i < reads; i++) maximum = maximum < values[i] ? values[i] : maximum;
|
|
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 i = 0; i < reads; i++) {
|
|
values[i] = fast::exp(values[i] - maximum);
|
|
normalizer += values[i];
|
|
}
|
|
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 i = 0; i < reads; i++) {
|
|
out[lid * reads + i] = (float)(bfloat)(values[i] * normalizer);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_sparse_probabilities(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *logits [[buffer(2)]],
|
|
device const int *selected [[buffer(3)]],
|
|
uint row [[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 partial_max[32];
|
|
threadgroup float partial_sum[32];
|
|
const uint width = args.u[0];
|
|
const uint top_k = args.u[1];
|
|
const uint threads = args.u[12];
|
|
logits += (ulong)row * width;
|
|
selected += (ulong)row * top_k;
|
|
out += (ulong)row * top_k;
|
|
|
|
if (simd_group == 0u) {
|
|
partial_max[lane] = -INFINITY;
|
|
partial_sum[lane] = 0.0f;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
float maximum = -INFINITY;
|
|
for (uint token = lid; token < width; token += threads) {
|
|
maximum = max(maximum, logits[token] * args.f[0]);
|
|
}
|
|
maximum = simd_max(maximum);
|
|
if (lane == 0u) partial_max[simd_group] = maximum;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (simd_group == 0u) {
|
|
maximum = simd_max(partial_max[lane]);
|
|
if (lane == 0u) partial_max[0] = maximum;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
maximum = partial_max[0];
|
|
|
|
float sum = 0.0f;
|
|
for (uint token = lid; token < width; token += threads) {
|
|
sum += fast::exp(logits[token] * args.f[0] - maximum);
|
|
}
|
|
sum = simd_sum(sum);
|
|
if (lane == 0u) partial_sum[simd_group] = sum;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (simd_group == 0u) {
|
|
sum = simd_sum(partial_sum[lane]);
|
|
if (lane == 0u) partial_sum[0] = sum;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (lid < top_k) {
|
|
const uint token = (uint)selected[lid];
|
|
const float log_total = maximum + fast::log(partial_sum[0]);
|
|
out[lid] = fast::exp(logits[token] * args.f[0] - log_total);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_route_top10(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device int *ids [[buffer(1)]],
|
|
device float *weights [[buffer(2)]],
|
|
device const float *probabilities [[buffer(3)]],
|
|
uint row [[threadgroup_position_in_grid]],
|
|
uint lane [[thread_index_in_simdgroup]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
constexpr uint top_k = 10u;
|
|
constexpr uint simd_groups = 16u;
|
|
constexpr uint local_candidates = simd_groups * top_k;
|
|
constexpr uint candidates_per_lane = local_candidates / 32u;
|
|
threadgroup float local_probabilities[local_candidates];
|
|
threadgroup int local_indices[local_candidates];
|
|
threadgroup float merged_probabilities[top_k];
|
|
threadgroup int merged_indices[top_k];
|
|
|
|
ids += (ulong)row * 10u;
|
|
weights += (ulong)row * 10u;
|
|
probabilities += (ulong)row * args.u[0];
|
|
|
|
const int expert = (int)(simd_group * 32u + lane);
|
|
float candidate_probability = probabilities[expert];
|
|
for (uint rank = 0; rank < top_k; rank++) {
|
|
const float winner_probability = simd_max(candidate_probability);
|
|
const int winner_index = (int)simd_max(
|
|
candidate_probability == winner_probability ? (float)expert : -1.0f);
|
|
if (lane == 0u) {
|
|
const uint destination = simd_group * top_k + rank;
|
|
local_probabilities[destination] = winner_probability;
|
|
local_indices[destination] = winner_index;
|
|
}
|
|
if (expert == winner_index) candidate_probability = -INFINITY;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
if (simd_group == 0u) {
|
|
float candidates[candidates_per_lane];
|
|
int candidate_indices[candidates_per_lane];
|
|
for (uint slot = 0; slot < candidates_per_lane; slot++) {
|
|
const uint source = lane + slot * 32u;
|
|
candidates[slot] = local_probabilities[source];
|
|
candidate_indices[slot] = local_indices[source];
|
|
}
|
|
|
|
for (uint rank = 0; rank < top_k; rank++) {
|
|
float lane_probability = candidates[0];
|
|
int lane_index = candidate_indices[0];
|
|
uint lane_slot = 0u;
|
|
for (uint slot = 1; slot < candidates_per_lane; slot++) {
|
|
if (candidates[slot] > lane_probability ||
|
|
(candidates[slot] == lane_probability &&
|
|
candidate_indices[slot] > lane_index)) {
|
|
lane_probability = candidates[slot];
|
|
lane_index = candidate_indices[slot];
|
|
lane_slot = slot;
|
|
}
|
|
}
|
|
const float winner_probability = simd_max(lane_probability);
|
|
const int winner_index = (int)simd_max(
|
|
lane_probability == winner_probability ? (float)lane_index : -1.0f);
|
|
if (lane == 0u) {
|
|
merged_probabilities[rank] = winner_probability;
|
|
merged_indices[rank] = winner_index;
|
|
}
|
|
if (lane_index == winner_index) candidates[lane_slot] = -INFINITY;
|
|
}
|
|
|
|
if (lane == 0u) {
|
|
float denominator = 0.0f;
|
|
for (uint slot = 0; slot < top_k; slot++) {
|
|
denominator = qwen_round_bf16(
|
|
denominator + merged_probabilities[top_k - 1u - slot]);
|
|
}
|
|
for (uint slot = 0; slot < top_k; slot++) {
|
|
const uint source = top_k - 1u - slot;
|
|
ids[slot] = merged_indices[source];
|
|
weights[slot] = qwen_round_bf16(
|
|
merged_probabilities[source] / denominator);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_accumulate(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] += x[index] * args.f[0];
|
|
}
|
|
|
|
kernel void kernel_qwen_accumulate_sigmoid_scalar(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const float *gate [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint hidden = args.u[0];
|
|
if (index < hidden * max(args.u[4], 1u)) {
|
|
const uint row = index / hidden;
|
|
const float shared = qwen_round_bf16(
|
|
x[index] * qwen_sigmoid_bf16(gate[row]));
|
|
out[index] = qwen_round_bf16(out[index] + shared);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_split_q_gate(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *q [[buffer(1)]],
|
|
device const float *packed [[buffer(2)]],
|
|
device float *gate [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
const uint heads = args.u[0];
|
|
const uint dim = args.u[1];
|
|
const uint row_width = heads * dim;
|
|
if (index >= row_width * max(args.u[4], 1u)) return;
|
|
const uint row = index / row_width;
|
|
const uint local = index % row_width;
|
|
const uint head = local / dim;
|
|
const uint column = local % dim;
|
|
const ulong packed_base = (ulong)row * row_width * 2u;
|
|
q[index] = packed[packed_base + (ulong)head * dim * 2u + column];
|
|
gate[index] = packed[packed_base + (ulong)head * dim * 2u + dim + column];
|
|
}
|
|
|
|
kernel void kernel_qwen_head_norm_rope(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *x [[buffer(2)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint3 group [[threadgroup_position_in_grid]],
|
|
uint3 thread_position [[thread_position_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]]) {
|
|
const uint head = group.x;
|
|
const uint lid = thread_position.x;
|
|
const uint dim = args.u[0];
|
|
const uint rotary = args.u[1];
|
|
if (head >= args.u[2]) return;
|
|
const ulong batch_base = (ulong)group.y * args.u[2] * dim;
|
|
device const float *row = x + batch_base + (ulong)head * dim;
|
|
|
|
threadgroup float scale[1];
|
|
threadgroup float partial[32];
|
|
float variance = 0.0f;
|
|
float values[4];
|
|
for (uint i = 0; i < 4u; i++) {
|
|
const uint column = lid * 4u + i;
|
|
values[i] = column < dim ? row[column] : 0.0f;
|
|
variance += values[i] * values[i];
|
|
}
|
|
variance = simd_sum(variance);
|
|
if (simd_group == 0u) partial[lane] = 0.0f;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (lane == 0u) partial[simd_group] = variance;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (simd_group == 0u) {
|
|
variance = simd_sum(partial[lane]);
|
|
if (lane == 0u) {
|
|
scale[0] = metal::precise::rsqrt(variance / (float)dim + args.f[0]);
|
|
}
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (uint i = 0; i < 4u; i++) {
|
|
const uint column = lid * 4u + i;
|
|
if (column >= dim) continue;
|
|
const float normalized = qwen_round_bf16(values[i] * scale[0]);
|
|
float value = args.u[11] != 0u ? normalized : qwen_round_bf16(
|
|
normalized * qwen_bf16(qwen_weight_u16(weight, args.u[13], column)));
|
|
if (column < rotary) {
|
|
const uint rotary_half = rotary / 2u;
|
|
const uint pair = column < rotary_half ? column + rotary_half : column - rotary_half;
|
|
const float paired_normalized = qwen_round_bf16(row[pair] * scale[0]);
|
|
const float paired = args.u[11] != 0u ? paired_normalized : qwen_round_bf16(
|
|
paired_normalized * qwen_bf16(qwen_weight_u16(weight, args.u[13], pair)));
|
|
const float exponent = 2.0f * (float)(column % rotary_half) / (float)rotary;
|
|
const float inv_frequency = 1.0f / pow(args.f[1], exponent);
|
|
const float theta = (float)(args.u[3] + group.y) * inv_frequency;
|
|
value = qwen_apply_rope(
|
|
value, column < rotary_half ? -paired : paired, theta);
|
|
}
|
|
out[batch_base + (ulong)head * dim + column] = qwen_round_bf16(value);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_store_kv_bf16(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *cache [[buffer(1)]],
|
|
device const float *key [[buffer(2)]],
|
|
device const float *value [[buffer(3)]],
|
|
uint2 gid [[thread_position_in_grid]]) {
|
|
const uint width = args.u[0];
|
|
const uint index = gid.x;
|
|
const uint row = gid.y;
|
|
if (index >= width) return;
|
|
const ulong base = (ulong)(args.u[1] + row) * width * 2u;
|
|
const ulong input = (ulong)row * width + index;
|
|
cache[base + index] = qwen_to_bf16(key[input]);
|
|
cache[base + width + index] = qwen_to_bf16(value[input]);
|
|
}
|
|
|
|
kernel void kernel_qwen_qsa_store_raw(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *raw [[buffer(1)]],
|
|
device const float *projected [[buffer(2)]],
|
|
device float *query [[buffer(3)]],
|
|
uint2 gid [[thread_position_in_grid]]) {
|
|
const uint dim = args.u[0];
|
|
const uint query_width = args.u[1] * dim;
|
|
const uint index = gid.x;
|
|
const uint row = gid.y;
|
|
const ulong projected_base = (ulong)row * (query_width + dim);
|
|
if (index < query_width) {
|
|
query[(ulong)row * query_width + index] = projected[projected_base + index];
|
|
} else if (index < query_width + dim) {
|
|
raw[(ulong)(args.u[2] + row) * dim + index - query_width] =
|
|
qwen_to_bf16(projected[projected_base + index]);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_qsa_pool_key(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *pooled [[buffer(1)]],
|
|
device const ushort *raw [[buffer(2)]],
|
|
device const uchar *weight [[buffer(5)]],
|
|
uint column [[thread_position_in_grid]]) {
|
|
const uint dim = args.u[0];
|
|
if (column >= dim) return;
|
|
const ulong raw_start = (ulong)args.u[3] * dim;
|
|
float mean = 0.0f;
|
|
for (uint token = 0; token < args.u[2]; token++) {
|
|
mean += qwen_bf16(raw[raw_start + (ulong)token * dim + column]);
|
|
}
|
|
const ushort mean_bf16 = qwen_to_bf16(mean / (float)args.u[2]);
|
|
float variance = 0.0f;
|
|
for (uint i = 0; i < dim; i++) {
|
|
float item = 0.0f;
|
|
for (uint token = 0; token < args.u[2]; token++) {
|
|
item += qwen_bf16(raw[raw_start + (ulong)token * dim + i]);
|
|
}
|
|
item = qwen_bf16(qwen_to_bf16(item / (float)args.u[2]));
|
|
variance = fma(item, item, variance);
|
|
}
|
|
const float scale = rsqrt(variance / (float)dim + args.f[0]);
|
|
float value = qwen_bf16(mean_bf16) * scale *
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], column));
|
|
const uint rotary = 64u;
|
|
if (column < rotary) {
|
|
const uint rotary_half = rotary / 2u;
|
|
const uint pair = column < rotary_half ? column + rotary_half : column - rotary_half;
|
|
float paired_mean = 0.0f;
|
|
for (uint token = 0; token < args.u[2]; token++) {
|
|
paired_mean += qwen_bf16(raw[raw_start + (ulong)token * dim + pair]);
|
|
}
|
|
paired_mean = qwen_bf16(qwen_to_bf16(paired_mean / (float)args.u[2]));
|
|
const float paired = paired_mean * scale *
|
|
qwen_bf16(qwen_weight_u16(weight, args.u[13], pair));
|
|
const float theta = (float)args.u[3] *
|
|
pow(args.f[1], -2.0f * (float)(column % rotary_half) / (float)rotary);
|
|
value = qwen_apply_rope(
|
|
value, column < rotary_half ? -paired : paired, theta);
|
|
}
|
|
pooled[(ulong)args.u[1] * dim + column] = qwen_to_bf16(value);
|
|
}
|
|
|
|
kernel void kernel_qwen_qsa_scores(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *scores [[buffer(1)]],
|
|
device const float *query [[buffer(2)]],
|
|
device const ushort *pooled [[buffer(3)]],
|
|
uint block [[thread_position_in_grid]]) {
|
|
const uint dim = args.u[0];
|
|
if (block >= args.u[2]) return;
|
|
float score = 0.0f;
|
|
for (uint head = 0; head < args.u[1]; head++) {
|
|
float head_score = 0.0f;
|
|
for (uint i = 0; i < dim; i++) {
|
|
head_score = fma(query[(ulong)head * dim + i],
|
|
qwen_bf16(pooled[(ulong)block * dim + i]),
|
|
head_score);
|
|
}
|
|
score += max(head_score, 0.0f);
|
|
}
|
|
scores[block] = 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;
|
|
for (uint i = 1; i < args.u[0]; i++) {
|
|
const int value = selected[i];
|
|
uint j = i;
|
|
while (j > 0u && selected[j - 1u] > value) {
|
|
selected[j] = selected[j - 1u];
|
|
j--;
|
|
}
|
|
selected[j] = value;
|
|
}
|
|
}
|
|
|
|
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]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint heads = args.u[0];
|
|
const uint kv_heads = args.u[1];
|
|
const uint dim = args.u[2];
|
|
if (head >= heads) return;
|
|
const uint simd_group = tid / 32u;
|
|
const uint kv_head = head / (heads / kv_heads);
|
|
constexpr uint simdgroups = 8u;
|
|
constexpr uint values_per_lane = 8u;
|
|
threadgroup float q[256];
|
|
threadgroup float max_scores[simdgroups];
|
|
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;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
float result[values_per_lane] = {0.0f};
|
|
float max_score = -FLT_MAX;
|
|
float sum_exp = 0.0f;
|
|
const uint column = lane * values_per_lane;
|
|
for (uint block = simd_group; block < args.u[4]; block += simdgroups) {
|
|
const uint token_start = (uint)selected[block] * args.u[5];
|
|
for (uint inner = 0u; inner < args.u[5]; inner++) {
|
|
const uint token = token_start + inner;
|
|
const ulong base = (ulong)token * kv_heads * dim * 2u +
|
|
(ulong)kv_head * dim + column;
|
|
float score = 0.0f;
|
|
for (uint i = 0u; i < values_per_lane; i++) {
|
|
score += q[column + i] * qwen_bf16(cache[base + i]);
|
|
}
|
|
score = simd_sum(score);
|
|
const float new_max = max(max_score, score);
|
|
const float factor = exp(max_score - new_max);
|
|
const float exp_score = exp(score - new_max);
|
|
sum_exp = fma(sum_exp, factor, exp_score);
|
|
for (uint i = 0u; i < values_per_lane; i++) {
|
|
result[i] = fma(
|
|
result[i], factor,
|
|
exp_score * qwen_bf16(cache[base + kv_heads * dim + i]));
|
|
}
|
|
max_score = new_max;
|
|
}
|
|
}
|
|
for (uint token = args.u[6] + simd_group;
|
|
token < args.u[3];
|
|
token += simdgroups) {
|
|
const ulong base = (ulong)token * kv_heads * dim * 2u +
|
|
(ulong)kv_head * dim + column;
|
|
float score = 0.0f;
|
|
for (uint i = 0u; i < values_per_lane; i++) {
|
|
score += q[column + i] * qwen_bf16(cache[base + i]);
|
|
}
|
|
score = simd_sum(score);
|
|
const float new_max = max(max_score, score);
|
|
const float factor = exp(max_score - new_max);
|
|
const float exp_score = exp(score - new_max);
|
|
sum_exp = fma(sum_exp, factor, exp_score);
|
|
for (uint i = 0u; i < values_per_lane; i++) {
|
|
result[i] = fma(
|
|
result[i], factor,
|
|
exp_score * qwen_bf16(cache[base + kv_heads * dim + i]));
|
|
}
|
|
max_score = new_max;
|
|
}
|
|
|
|
if (lane == 0u) {
|
|
max_scores[simd_group] = max_score;
|
|
sum_exp_scores[simd_group] = sum_exp;
|
|
}
|
|
for (uint i = 0u; i < values_per_lane; i++) {
|
|
outputs[simd_group * dim + column + i] = result[i];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
float merged_max = max_scores[0];
|
|
for (uint group = 1u; group < simdgroups; group++) {
|
|
merged_max = max(merged_max, max_scores[group]);
|
|
}
|
|
float merged_sum = 0.0f;
|
|
float merged_value = 0.0f;
|
|
for (uint group = 0u; group < simdgroups; group++) {
|
|
const float weight = exp(max_scores[group] - merged_max);
|
|
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);
|
|
}
|
|
|
|
kernel void kernel_qwen_dense_attention_masked(
|
|
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 uchar *mask [[buffer(4)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint simd_group [[simdgroup_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];
|
|
const uint tokens = args.u[3];
|
|
if (head >= heads) return;
|
|
const uint kv_head = head / (heads / kv_heads);
|
|
constexpr uint simdgroups = 32u;
|
|
constexpr uint values_per_lane = 8u;
|
|
threadgroup float outputs[simdgroups * simdgroups];
|
|
threadgroup float max_scores[simdgroups];
|
|
threadgroup float sum_exp_scores[simdgroups];
|
|
float q[values_per_lane];
|
|
float k[values_per_lane];
|
|
float result[values_per_lane] = {0.0f};
|
|
const ulong query_base =
|
|
((ulong)row * heads + head) * dim + lane * values_per_lane;
|
|
const float attention_scale = precise::rsqrt((float)dim);
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
q[i] = attention_scale * query[query_base + i];
|
|
}
|
|
|
|
float max_score = -FLT_MAX;
|
|
float sum_exp = 0.0f;
|
|
for (uint token = simd_group; token < tokens; token += simdgroups) {
|
|
const bool use_key = mask[(ulong)row * tokens + token] != 0u;
|
|
if (use_key) {
|
|
const ulong base = (ulong)token * kv_heads * dim * 2u +
|
|
(ulong)kv_head * dim + lane * values_per_lane;
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
k[i] = qwen_bf16(cache[base + i]);
|
|
}
|
|
float score = 0.0f;
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
score = fma(q[i], k[i], score);
|
|
}
|
|
score = simd_sum(score);
|
|
const float new_max = max(max_score, score);
|
|
const float factor = fast::exp(max_score - new_max);
|
|
const float exp_score = fast::exp(score - new_max);
|
|
max_score = new_max;
|
|
sum_exp = fma(sum_exp, factor, exp_score);
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
result[i] = fma(
|
|
result[i], factor,
|
|
exp_score * qwen_bf16(cache[base + kv_heads * dim + i]));
|
|
}
|
|
}
|
|
}
|
|
if (lane == 0u) {
|
|
max_scores[simd_group] = max_score;
|
|
sum_exp_scores[simd_group] = sum_exp;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
max_score = max_scores[lane];
|
|
const float new_max = simd_max(max_score);
|
|
const float factor = fast::exp(max_score - new_max);
|
|
sum_exp = simd_sum(sum_exp_scores[lane] * factor);
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
outputs[lane * simdgroups + simd_group] = result[i];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
result[i] = simd_sum(outputs[simd_group * simdgroups + lane] * factor);
|
|
if (sum_exp != 0.0f) result[i] /= sum_exp;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
if (lane == 0u) {
|
|
const ulong output_base =
|
|
((ulong)row * heads + head) * dim + simd_group * values_per_lane;
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
out[output_base + i] = qwen_round_bf16(result[i]);
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_dense_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)]],
|
|
uint head [[threadgroup_position_in_grid]],
|
|
uint simd_group [[simdgroup_index_in_threadgroup]],
|
|
uint lane [[thread_index_in_simdgroup]]) {
|
|
const uint heads = args.u[0];
|
|
const uint kv_heads = args.u[1];
|
|
const uint dim = args.u[2];
|
|
const uint tokens = args.u[3];
|
|
if (head >= heads) return;
|
|
const uint kv_head = head / (heads / kv_heads);
|
|
constexpr uint simdgroups = 32u;
|
|
constexpr uint values_per_lane = 8u;
|
|
threadgroup float outputs[simdgroups * simdgroups];
|
|
threadgroup float max_scores[simdgroups];
|
|
threadgroup float sum_exp_scores[simdgroups];
|
|
float q[values_per_lane];
|
|
float result[values_per_lane] = {0.0f};
|
|
const ulong query_base = (ulong)head * dim + lane * values_per_lane;
|
|
const float attention_scale = precise::rsqrt((float)dim);
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
q[i] = attention_scale * query[query_base + i];
|
|
}
|
|
|
|
float max_score = -FLT_MAX;
|
|
float sum_exp = 0.0f;
|
|
for (uint token = simd_group; token < tokens; token += simdgroups) {
|
|
const ulong base = (ulong)token * kv_heads * dim * 2u +
|
|
(ulong)kv_head * dim + lane * values_per_lane;
|
|
float score = 0.0f;
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
score += q[i] * qwen_bf16(cache[base + i]);
|
|
}
|
|
score = simd_sum(score);
|
|
const float new_max = max(max_score, score);
|
|
const float factor = fast::exp(max_score - new_max);
|
|
const float exp_score = fast::exp(score - new_max);
|
|
max_score = new_max;
|
|
sum_exp = fma(sum_exp, factor, exp_score);
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
result[i] = fma(
|
|
result[i], factor,
|
|
exp_score * qwen_bf16(cache[base + kv_heads * dim + i]));
|
|
}
|
|
}
|
|
if (lane == 0u) {
|
|
max_scores[simd_group] = max_score;
|
|
sum_exp_scores[simd_group] = sum_exp;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
max_score = max_scores[lane];
|
|
const float new_max = simd_max(max_score);
|
|
const float factor = fast::exp(max_score - new_max);
|
|
sum_exp = simd_sum(sum_exp_scores[lane] * factor);
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
outputs[lane * simdgroups + simd_group] = result[i];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
result[i] = simd_sum(outputs[simd_group * simdgroups + lane] * factor);
|
|
if (sum_exp != 0.0f) result[i] /= sum_exp;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
if (lane == 0u) {
|
|
const ulong output_base = (ulong)head * dim + simd_group * values_per_lane;
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
out[output_base + i] = qwen_round_bf16(result[i]);
|
|
}
|
|
}
|
|
}
|
|
|
|
// MLX deliberately leaves its fused SDPA path when query rows times the GQA
|
|
// factor exceeds 32. These three kernels preserve the fallback's BF16
|
|
// boundaries: scaled Q @ K.t, precise softmax, then probabilities @ V.
|
|
kernel void kernel_qwen_attention_fallback_scores(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *scores [[buffer(1)]],
|
|
device const float *query [[buffer(2)]],
|
|
device const ushort *cache [[buffer(3)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
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];
|
|
const uint tokens = args.u[3];
|
|
if (head >= heads || row >= args.u[4]) return;
|
|
const uint kv_head = head / (heads / kv_heads);
|
|
const bfloat scale = (bfloat)precise::rsqrt((float)dim);
|
|
const ulong query_base = ((ulong)row * heads + head) * dim;
|
|
const ulong score_base = ((ulong)head * args.u[4] + row) * tokens;
|
|
for (uint token = 0u; token < tokens; token++) {
|
|
if (token > args.u[5] + row) {
|
|
if (lane == 0u) scores[score_base + token] = 0xff7fu;
|
|
continue;
|
|
}
|
|
const ulong key_base = (ulong)token * kv_heads * dim * 2u +
|
|
(ulong)kv_head * dim;
|
|
float sum = 0.0f;
|
|
for (uint column = lane; column < dim; column += 32u) {
|
|
const float q = (float)(bfloat)((bfloat)query[query_base + column] * scale);
|
|
sum = fma(q, qwen_bf16(cache[key_base + column]), sum);
|
|
}
|
|
sum = simd_sum(sum);
|
|
if (lane == 0u) scores[score_base + token] = qwen_to_bf16(sum);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_attention_fallback_softmax(
|
|
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]]) {
|
|
constexpr uint reads = 4u;
|
|
threadgroup float local_max[32];
|
|
threadgroup float local_normalizer[32];
|
|
const uint tokens = args.u[3];
|
|
const ulong base = (ulong)group * tokens;
|
|
float values[reads];
|
|
for (uint i = 0u; i < reads; i++) {
|
|
const uint token = lid * reads + i;
|
|
values[i] = token < tokens ? qwen_bf16(input[base + token]) : -INFINITY;
|
|
}
|
|
if (simd_group == 0u) {
|
|
local_max[lane] = -INFINITY;
|
|
local_normalizer[lane] = 0.0f;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
float maximum = -FLT_MAX;
|
|
for (uint i = 0u; i < reads; i++) maximum = max(maximum, values[i]);
|
|
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 i = 0u; i < reads; i++) {
|
|
values[i] = fast::exp(values[i] - maximum);
|
|
normalizer += values[i];
|
|
}
|
|
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 i = 0u; i < reads; i++) {
|
|
const uint token = lid * reads + i;
|
|
if (token < tokens) out[base + token] = qwen_to_bf16(values[i] * normalizer);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_attention_fallback_output(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const ushort *probabilities [[buffer(2)]],
|
|
device const ushort *cache [[buffer(3)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
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];
|
|
const uint tokens = args.u[3];
|
|
if (head >= heads || row >= args.u[4]) return;
|
|
const uint kv_head = head / (heads / kv_heads);
|
|
const ulong probability_base = ((ulong)head * args.u[4] + row) * tokens;
|
|
const ulong output_base = ((ulong)row * heads + head) * dim;
|
|
for (uint column = lane; column < dim; column += 32u) {
|
|
float sum = 0.0f;
|
|
for (uint token = 0u; token < tokens; token++) {
|
|
const ulong value = (ulong)token * kv_heads * dim * 2u +
|
|
(ulong)kv_head * dim + kv_heads * dim + column;
|
|
sum = fma(qwen_bf16(probabilities[probability_base + token]),
|
|
qwen_bf16(cache[value]), sum);
|
|
}
|
|
out[output_base + column] = qwen_round_bf16(sum);
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_causal_mask(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device uchar *mask [[buffer(1)]],
|
|
uint2 gid [[thread_position_in_grid]]) {
|
|
if (gid.x < args.u[0] && gid.y < args.u[4]) {
|
|
mask[(ulong)gid.y * args.u[0] + gid.x] = gid.x <= args.u[5] + gid.y;
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_prepare_dense_kv(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *keys [[buffer(1)]],
|
|
device const ushort *cache [[buffer(2)]],
|
|
device ushort *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] = cache[source];
|
|
values[target] = cache[source + width];
|
|
}
|
|
|
|
kernel void kernel_qwen_float_to_bf16(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device ushort *out [[buffer(1)]],
|
|
device const float *input [[buffer(2)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) out[index] = qwen_to_bf16(input[index]);
|
|
}
|
|
|
|
kernel void kernel_qwen_dense_attention_rows(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const ushort *queries [[buffer(2)]],
|
|
device const ushort *keys [[buffer(3)]],
|
|
device const ushort *values [[buffer(4)]],
|
|
device const uchar *mask [[buffer(8)]],
|
|
uint2 group [[threadgroup_position_in_grid]],
|
|
uint simd_group [[simdgroup_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];
|
|
const uint tokens = args.u[3];
|
|
if (head >= heads) return;
|
|
const uint kv_head = head / (heads / kv_heads);
|
|
constexpr uint simdgroups = 32u;
|
|
constexpr uint values_per_lane = 8u;
|
|
threadgroup float outputs[simdgroups * simdgroups];
|
|
threadgroup float max_scores[simdgroups];
|
|
threadgroup float sum_exp_scores[simdgroups];
|
|
float q[values_per_lane];
|
|
float k[values_per_lane];
|
|
float result[values_per_lane] = {0.0f};
|
|
queries += ((ulong)row * heads + head) * dim + lane * values_per_lane;
|
|
keys += ((ulong)kv_head * tokens + simd_group) * dim + lane * values_per_lane;
|
|
values += ((ulong)kv_head * tokens + simd_group) * dim + lane * values_per_lane;
|
|
mask += (ulong)row * tokens + simd_group;
|
|
const float attention_scale = precise::rsqrt((float)dim);
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
q[i] = attention_scale * qwen_bf16(queries[i]);
|
|
}
|
|
|
|
float max_score = -FLT_MAX;
|
|
float sum_exp = 0.0f;
|
|
for (uint token = simd_group; token < tokens; token += simdgroups) {
|
|
const bool use_key = mask[0] != 0u;
|
|
if (use_key) {
|
|
for (uint i = 0; i < values_per_lane; i++) k[i] = qwen_bf16(keys[i]);
|
|
float score = 0.0f;
|
|
for (uint i = 0; i < values_per_lane; i++) score += q[i] * k[i];
|
|
score = simd_sum(score);
|
|
const float new_max = max(max_score, score);
|
|
const float factor = fast::exp(max_score - new_max);
|
|
const float exp_score = fast::exp(score - new_max);
|
|
max_score = new_max;
|
|
sum_exp = sum_exp * factor + exp_score;
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
result[i] = result[i] * factor + exp_score * qwen_bf16(values[i]);
|
|
}
|
|
}
|
|
keys += simdgroups * dim;
|
|
values += simdgroups * dim;
|
|
mask += simdgroups;
|
|
}
|
|
if (lane == 0u) {
|
|
max_scores[simd_group] = max_score;
|
|
sum_exp_scores[simd_group] = sum_exp;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
max_score = max_scores[lane];
|
|
const float new_max = simd_max(max_score);
|
|
const float factor = fast::exp(max_score - new_max);
|
|
sum_exp = simd_sum(sum_exp_scores[lane] * factor);
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
outputs[lane * simdgroups + simd_group] = result[i];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
result[i] = simd_sum(outputs[simd_group * simdgroups + lane] * factor);
|
|
if (sum_exp != 0.0f) result[i] /= sum_exp;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
if (lane == 0u) {
|
|
const ulong output_base =
|
|
((ulong)row * heads + head) * dim + simd_group * values_per_lane;
|
|
for (uint i = 0; i < values_per_lane; i++) {
|
|
out[output_base + i] = qwen_round_bf16(result[i]);
|
|
}
|
|
}
|
|
}
|
|
|
|
kernel void kernel_qwen_gate_attention(
|
|
constant qwen_kernel_args &args [[buffer(0)]],
|
|
device float *out [[buffer(1)]],
|
|
device const float *attention [[buffer(2)]],
|
|
device const float *gate [[buffer(3)]],
|
|
uint index [[thread_position_in_grid]]) {
|
|
if (index < args.u[0]) {
|
|
out[index] = qwen_round_bf16(
|
|
attention[index] * qwen_sigmoid_bf16(gate[index]));
|
|
}
|
|
}
|