// 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]));
    }
}