Qwen¶

pub const RopeScaling = struct {
    type: RopeScalingType,
    factor: f32,

    pub const RopeScalingType = enum {
        linear,   // Simple frequency division
        dynamic,  // NTK-Aware dynamic scaling
        yarn,     // Yet Another RoPE extensioN
        ntk,      // NTK-Aware interpolation
    };
};

// Qwen 1.x config enables dynamic NTK
.use_dynamic_ntk = true,
.use_logn_attn = true,

// Qwen2 config uses YARN scaling
.rope_scaling = RopeScaling{ .type = .yarn, .factor = 4.0 },

if (config.use_logn_attn) {
    logn_list = try allocator.alloc(f32, config.num_hidden_layers);
    for (logn_list.?, 0..) |*logn, i| {
        const layer_idx = @as(f32, @floatFromInt(i));
        logn.* = std.math.log(@as(f32, 512.0)) / std.math.log(2.0 + layer_idx);
    }
}

pub const QwenConfig = struct {
    variant: QwenVariant,
    vocab_size: u32,
    hidden_size: u32,
    num_attention_heads: u32,
    num_key_value_heads: u32,      // GQA support
    num_hidden_layers: u32,
    intermediate_size: u32,
    max_position_embeddings: u32,
    rms_norm_eps: f32,
    rope_theta: f32,               // 1,000,000 for all Qwen variants
    rope_scaling: ?RopeScaling,    // null for Qwen1, YARN for Qwen2
    use_bias: bool,
    use_sliding_window: bool,
    sliding_window: u32,
    use_flash_attn: bool,
    attention_dropout: f32,
    use_cache: bool,
    use_dynamic_ntk: bool,         // Qwen1 feature
    use_logn_attn: bool,           // Qwen1 feature
};

pub const QwenAttention = struct {
    config: QwenConfig,
    c_attn: LinearLayer,      // Combined Q,K,V projection
    c_proj: LinearLayer,      // Output projection
    attn_dropout: DropoutLayer,
    resid_dropout: DropoutLayer,
    head_dim: u32,
    kv_head_dim: u32,
    logn_list: ?[]f32,        // LogN scaling factors
};

pub fn forward(self: *Self, hidden_states: Tensor(f32), layer_idx: u32,
               attention_mask: ?Tensor(f32), position_ids: ?Tensor(u32)) !Tensor(f32) {
    // 1. Single QKV projection (GPT-2 style)
    const qkv = try self.c_attn.forward(hidden_states);

    // 2. Split into Q, K, V with GQA support
    var query = try self.extractTensor(qkv, 0, q_size);
    var key = try self.extractTensor(qkv, q_size, kv_size);
    var value = try self.extractTensor(qkv, q_size + kv_size, kv_size);

    // 3. Apply RoPE with Qwen-specific scaling
    if (position_ids) |pos_ids| {
        query = try self.applyQwenRoPE(query, pos_ids,
            self.config.rope_scaling, self.config.use_dynamic_ntk);
        key = try self.applyQwenRoPE(key, pos_ids,
            self.config.rope_scaling, self.config.use_dynamic_ntk);
    }

    // 4. Grouped-query attention
    var attn_output = try self.computeGroupedQueryAttention(
        query, key, value, attention_mask);

    // 5. LogN attention scaling (Qwen 1.x only)
    if (self.config.use_logn_attn and self.logn_list != null) {
        attn_output = try self.scaleAttentionOutput(
            attn_output, self.logn_list.?[layer_idx]);
    }

    // 6. Output projection
    return try self.c_proj.forward(attn_output);
}

pub const QwenMLP = struct {
    w1: LinearLayer,       // Gate projection [d, d_ff]
    w2: LinearLayer,       // Up projection [d, d_ff]
    c_proj: LinearLayer,   // Down projection [d_ff, d]

    pub fn forward(self: *Self, hidden_states: Tensor(f32)) !Tensor(f32) {
        const gate = try self.w1.forward(hidden_states);
        const up = try self.w2.forward(hidden_states);
        const gate_swish = try neural_primitives.swish(gate, self.allocator);
        const gated = try self.elementwiseMultiply(gate_swish, up);
        return try self.c_proj.forward(gated);
    }
};

pub const QwenBlock = struct {
    ln_1: RMSNorm,           // Pre-attention norm
    attn: QwenAttention,     // Attention with RoPE scaling
    ln_2: RMSNorm,           // Pre-MLP norm
    mlp: QwenMLP,            // SwiGLU FFN

    pub fn forward(self: *Self, hidden_states: Tensor(f32), layer_idx: u32,
                   attention_mask: ?Tensor(f32), position_ids: ?Tensor(u32)) !Tensor(f32) {
        var residual = hidden_states;
        const normed = try self.ln_1.forward(hidden_states);
        const attn_out = try self.attn.forward(normed, layer_idx,
            attention_mask, position_ids);
        hidden_states = try self.addTensors(residual, attn_out);

        residual = hidden_states;
        const normed2 = try self.ln_2.forward(hidden_states);
        const mlp_out = try self.mlp.forward(normed2);
        return try self.addTensors(residual, mlp_out);
    }
};

pub const QwenModel = struct {
    config: QwenConfig,
    wte: EmbeddingLayer,      // Token embeddings
    h: []QwenBlock,           // Transformer blocks
    ln_f: RMSNorm,            // Final normalization
    lm_head: LinearLayer,     // Output projection

    pub fn forward(self: *Self, input_ids: Tensor(u32),
                   attention_mask: ?Tensor(f32),
                   position_ids: ?Tensor(u32)) !Tensor(f32) {
        var hidden = try self.wte.forward(input_ids);
        for (self.h, 0..) |*block, i| {
            hidden = try block.forward(hidden, @intCast(i),
                attention_mask, position_ids);
        }
        hidden = try self.ln_f.forward(hidden);
        return try self.lm_head.forward(hidden);
    }
};