// Qwen3.8 Flash Next primitives. Rust owns the graph and all state lifetimes;
// this file contains only the data-parallel kernels executed by 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 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);
}

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] = sum;
}

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_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_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_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] = sum;
}

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]]) {
    if (index < args.u[0] * 4u) out[index] = x[index % args.u[0]];
}

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)]],
        uint group [[thread_position_in_grid]]) {
    const uint width = args.u[0];
    const uint group_size = args.u[1];
    if (group >= width / group_size) return;
    const uint start = group * group_size;
    float variance = 0.0f;
    for (uint i = 0; i < group_size; i++) variance = fma(x[start + i], x[start + i], variance);
    const float scale = rsqrt(variance / (float)group_size + args.f[0]);
    for (uint i = 0; i < group_size; i++) {
        const uint index = start + i;
        out[index] = x[index] * scale *
                     (1.0f + qwen_bf16(qwen_weight_u16(weight, args.u[13], index)));
    }
}

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 = x[index] * 0.25f;
    out[index] = value / (1.0f + exp(-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] = 1.0f / (1.0f + exp(-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]) out[index] = 2.0f / (1.0f + exp(-x[index] * 0.25f));
}

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]]) {
    if (index >= args.u[0]) return;
    float value = 0.0f;
    for (uint stream = 0; stream < 4u; stream++) {
        const uint offset = stream * args.u[0] + index;
        value = fma(normalized[offset], mix[offset], value);
    }
    out[index] = 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];
    if (index >= hidden * 4u) return;
    out[index] = residual[index] + block[index % hidden] * gate[index / hidden];
}

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)]],
        uint stream [[thread_position_in_grid]]) {
    const uint hidden = args.u[0];
    if (stream >= 4u) return;
    const ulong base = (ulong)stream * hidden;
    float score = 0.0f;
    for (uint i = 0; i < hidden; i++) score = fma(key[base + i], query[base + i], score);
    score *= rsqrt((float)hidden);
    const float transformed = copysign(sqrt(max(abs(score), 1.0e-6f)), score);
    const float gate = 1.0f / (1.0f + exp(-transformed));
    for (uint i = 0; i < hidden; i++) out[base + i] = value[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 ushort *state [[buffer(4)]],
        device const uchar *weight [[buffer(5)]],
        uint channel [[thread_position_in_grid]]) {
    if (channel >= args.u[0]) return;
    device ushort *history = state + (ulong)channel * 9u;
    float value = fma(qwen_bf16(history[0]),
                      qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)),
                  fma(qwen_bf16(history[3]),
                      qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 1u)),
                  fma(qwen_bf16(history[6]),
                      qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 2u)),
                      normalized[channel] *
                      qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 3u)))));
    for (uint i = 0; i < 8u; i++) history[i] = history[i + 1u];
    history[8] = qwen_to_bf16(normalized[channel]);
    out[channel] = gated[channel] + value / (1.0f + exp(-value));
}

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] = a[index] + b[index];
}

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;
    device ushort *history = state + (ulong)channel * 3u;
    float value = fma(qwen_bf16(history[0]),
                      qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u)),
                  fma(qwen_bf16(history[1]),
                      qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 1u)),
                  fma(qwen_bf16(history[2]),
                      qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 2u)),
                      x[channel] *
                      qwen_bf16(qwen_weight_u16(weight, args.u[13], (ulong)channel * 4u + 3u)))));
    history[0] = history[1];
    history[1] = history[2];
    history[2] = qwen_to_bf16(x[channel]);
    out[channel] = value / (1.0f + exp(-value));
}

kernel void kernel_qwen_gdn_step(
        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 *a_log [[buffer(5)]],
        device const uchar *dt_bias [[buffer(6)]],
        uint2 gid [[thread_position_in_grid]]) {
    const uint value_index = gid.x;
    const uint head = gid.y;
    const uint dim = args.u[0];
    const uint key_heads = args.u[1];
    const uint value_heads = args.u[2];
    if (value_index >= dim || head >= value_heads) return;
    const uint key_head = head / (value_heads / key_heads);
    device const float *q_raw = qkv + (ulong)key_head * dim;
    device const float *k_raw = qkv + (ulong)key_heads * dim + (ulong)key_head * dim;
    device const float *value = qkv + (ulong)key_heads * dim * 2u + (ulong)head * dim;
    float qsum = 0.0f;
    float ksum = 0.0f;
    for (uint i = 0; i < dim; i++) {
        qsum = fma(q_raw[i], q_raw[i], qsum);
        ksum = fma(k_raw[i], k_raw[i], ksum);
    }
    const float qscale = rsqrt(qsum + args.f[0]) * rsqrt((float)dim);
    const float kscale = rsqrt(ksum + args.f[0]);
    const float beta = 1.0f / (1.0f + exp(-controls[args.u[3] + head]));
    const float step = controls[args.u[4] + head] +
                       qwen_bf16(qwen_weight_u16(dt_bias, args.u[14], 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[13], head))) * softplus);
    device float *column = state + ((ulong)head * dim * dim) + value_index;
    float prediction = 0.0f;
    for (uint i = 0; i < dim; i++) {
        prediction = fma(column[(ulong)i * dim] * decay, k_raw[i] * kscale, prediction);
    }
    const float delta = (value[value_index] - prediction) * beta;
    float result = 0.0f;
    for (uint i = 0; i < dim; i++) {
        const ulong offset = (ulong)i * dim;
        const float updated = column[offset] * decay + k_raw[i] * kscale * delta;
        column[offset] = updated;
        result = fma(updated, q_raw[i] * qscale, result);
    }
    device float *head_out = out + (ulong)head * dim;
    head_out[value_index] = result;
}

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)]],
        uint head [[thread_position_in_grid]]) {
    const uint dim = args.u[0];
    if (head >= args.u[1]) return;
    device const float *row = x + (ulong)head * dim;
    float variance = 0.0f;
    for (uint i = 0; i < dim; i++) variance = fma(row[i], row[i], variance);
    const float scale = rsqrt(variance / (float)dim + args.f[0]);
    for (uint i = 0; i < dim; i++) {
        const ulong index = (ulong)head * dim + i;
        out[index] = row[i] * scale *
                     qwen_bf16(qwen_weight_u16(weight, args.u[13], i)) /
                     (1.0f + exp(-controls[index]));
    }
}

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;
    out[index] = gate[index] / (1.0f + exp(-gate[index])) * up[index];
}

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 *logits [[buffer(3)]],
        uint gid [[thread_position_in_grid]]) {
    if (gid != 0u) return;
    float max_value = -INFINITY;
    for (uint i = 0; i < args.u[0]; i++) max_value = max(max_value, logits[i]);
    float sum = 0.0f;
    for (uint i = 0; i < args.u[0]; i++) sum += exp(logits[i] - max_value);
    float selected_sum = 0.0f;
    for (uint slot = 0; slot < 10u; slot++) {
        float best = -1.0f;
        int best_id = -1;
        for (uint i = 0; i < args.u[0]; i++) {
            bool used = false;
            for (uint j = 0; j < slot; j++) used = used || ids[j] == (int)i;
            const float probability = exp(logits[i] - max_value) / sum;
            if (!used && probability > best) {
                best = probability;
                best_id = (int)i;
            }
        }
        ids[slot] = best_id;
        weights[slot] = best;
        selected_sum += best;
    }
    for (uint slot = 0; slot < 10u; slot++) weights[slot] /= selected_sum;
}

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]]) {
    if (index < args.u[0]) out[index] += x[index] / (1.0f + exp(-gate[0]));
}

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];
    if (index >= heads * dim) return;
    const uint head = index / dim;
    const uint column = index % dim;
    q[index] = packed[(ulong)head * dim * 2u + column];
    gate[index] = packed[(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)]],
        uint2 gid [[thread_position_in_grid]]) {
    const uint column = gid.x;
    const uint head = gid.y;
    const uint dim = args.u[0];
    const uint rotary = args.u[1];
    if (column >= dim || head >= args.u[2]) return;
    device const float *row = x + (ulong)head * dim;
    float variance = 0.0f;
    for (uint i = 0; i < dim; i++) variance = fma(row[i], row[i], variance);
    const float scale = rsqrt(variance / (float)dim + args.f[0]);
    float value = row[column] * scale *
                  (1.0f + 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 = row[pair] * scale *
                             (1.0f + 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 = value * cos(theta) + (column < rotary_half ? -paired : paired) * sin(theta);
    }
    out[(ulong)head * dim + column] = 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)]],
        uint index [[thread_position_in_grid]]) {
    const uint width = args.u[0];
    if (index >= width) return;
    const ulong base = (ulong)args.u[1] * width * 2u;
    cache[base + index] = qwen_to_bf16(key[index]);
    cache[base + width + index] = qwen_to_bf16(value[index]);
}

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)]],
        uint index [[thread_position_in_grid]]) {
    const uint dim = args.u[0];
    const uint query_width = args.u[1] * dim;
    if (index < query_width) {
        query[index] = projected[index];
    } else if (index < query_width + dim) {
        raw[(ulong)args.u[2] * dim + index - query_width] = qwen_to_bf16(projected[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 *
                  (1.0f + 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 *
                             (1.0f + 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 = value * cos(theta) + (column < rotary_half ? -paired : paired) * sin(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;
    }
}

static inline uint qwen_qsa_token(
        device const int *selected,
        uint ordinal,
        uint selected_count,
        uint ratio,
        uint tail_start) {
    const uint selected_tokens = selected_count * ratio;
    return ordinal < selected_tokens
        ? (uint)selected[ordinal / ratio] * ratio + ordinal % ratio
        : tail_start + ordinal - selected_tokens;
}

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)]],
        uint2 gid [[thread_position_in_grid]],
        uint lane [[thread_index_in_threadgroup]]) {
    const uint heads = args.u[0];
    const uint kv_heads = args.u[1];
    const uint dim = args.u[2];
    const uint column = gid.x;
    const uint head = gid.y;
    if (column >= dim || head >= heads) return;
    const uint kv_head = head / (heads / kv_heads);
    const uint tokens = args.u[4] * args.u[5] + args.u[7];
    const float attention_scale = rsqrt((float)dim);
    threadgroup float probabilities[2051];
    if (lane == 0u) {
        float max_score = -INFINITY;
        for (uint ordinal = 0; ordinal < tokens; ordinal++) {
            const uint token = qwen_qsa_token(selected, ordinal, args.u[4], args.u[5], args.u[6]);
            const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim;
            float score = 0.0f;
            for (uint i = 0; i < dim; i++) {
                score = fma(query[(ulong)head * dim + i], qwen_bf16(cache[base + i]), score);
            }
            probabilities[ordinal] = score * attention_scale;
            max_score = max(max_score, probabilities[ordinal]);
        }
        float denominator = 0.0f;
        for (uint ordinal = 0; ordinal < tokens; ordinal++) {
            probabilities[ordinal] = exp(probabilities[ordinal] - max_score);
            denominator += probabilities[ordinal];
        }
        for (uint ordinal = 0; ordinal < tokens; ordinal++) {
            probabilities[ordinal] /= denominator;
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    float value = 0.0f;
    for (uint ordinal = 0; ordinal < tokens; ordinal++) {
        const uint token = qwen_qsa_token(selected, ordinal, args.u[4], args.u[5], args.u[6]);
        const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim;
        value = fma(probabilities[ordinal],
                    qwen_bf16(cache[base + kv_heads * dim + column]),
                    value);
    }
    out[(ulong)head * dim + column] = value;
}

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)]],
        uint2 gid [[thread_position_in_grid]],
        uint lane [[thread_index_in_threadgroup]]) {
    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];
    const uint column = gid.x;
    const uint head = gid.y;
    if (column >= dim || head >= heads) return;
    const uint kv_head = head / (heads / kv_heads);
    const float attention_scale = rsqrt((float)dim);
    threadgroup float probabilities[2048];
    if (lane == 0u) {
        float max_score = -INFINITY;
        for (uint token = 0; token < tokens; token++) {
            const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim;
            float score = 0.0f;
            for (uint i = 0; i < dim; i++) score = fma(query[(ulong)head * dim + i], qwen_bf16(cache[base + i]), score);
            probabilities[token] = score * attention_scale;
            max_score = max(max_score, probabilities[token]);
        }
        float denominator = 0.0f;
        for (uint token = 0; token < tokens; token++) {
            probabilities[token] = exp(probabilities[token] - max_score);
            denominator += probabilities[token];
        }
        for (uint token = 0; token < tokens; token++) probabilities[token] /= denominator;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    float value = 0.0f;
    for (uint token = 0; token < tokens; token++) {
        const ulong base = (ulong)token * kv_heads * dim * 2u + (ulong)kv_head * dim;
        value = fma(probabilities[token], qwen_bf16(cache[base + kv_heads * dim + column]), value);
    }
    out[(ulong)head * dim + column] = value;
}

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] = attention[index] / (1.0f + exp(-gate[index]));
}