// 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 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 uint *packed [[buffer(5)]],
        device const ushort *scales [[buffer(6)]],
        device const ushort *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_value(packed, scales, biases, 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 uint *packed [[buffer(5)]],
        device const ushort *scales [[buffer(6)]],
        device const ushort *biases [[buffer(7)]],
        uint column [[thread_position_in_grid]]) {
    if (column >= args.u[0]) return;
    out[column] = qwen_quant_value(packed, scales, biases, args.u[4], column,
                                   args.u[0], 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 ushort *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(weights[(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 ushort *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(weight[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_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 ushort *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(weight[(ulong)channel * 4u]),
                  fma(qwen_bf16(history[1]), qwen_bf16(weight[(ulong)channel * 4u + 1u]),
                  fma(qwen_bf16(history[2]), qwen_bf16(weight[(ulong)channel * 4u + 2u]),
                      x[channel] * qwen_bf16(weight[(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 ushort *a_log [[buffer(5)]],
        device const ushort *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(dt_bias[head]);
    const float softplus = max(step, 0.0f) + log(1.0f + exp(-abs(step)));
    const float decay = exp(-exp(qwen_bf16(a_log[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 ushort *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(weight[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 ushort *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(weight[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(weight[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_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 [[thread_position_in_grid]]) {
    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);
    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);
        max_score = max(max_score, score * rsqrt((float)dim));
    }
    float denominator = 0.0f;
    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);
        denominator += exp(score * rsqrt((float)dim) - max_score);
    }
    for (uint column = 0; column < dim; column++) {
        float value = 0.0f;
        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);
            const float probability = exp(score * rsqrt((float)dim) - max_score) / denominator;
            value = fma(probability, 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]));
}