xserv/crates/xserv-model/src/gpt2.rs

use std::collections::HashMap;
use xserv_kernels::*;
use xserv_tensor::{DType, Device, Tensor};

use crate::config::ModelConfig;

pub struct GPT2 {
    pub config: ModelConfig,
    wte: Tensor,
    wpe: Tensor,
    layers: Vec<GPT2Block>,
    ln_f_g: Tensor,
    ln_f_b: Tensor,
    lm_head: Tensor, // precomputed wte^T
}

struct GPT2Block {
    ln_1_g: Tensor,
    ln_1_b: Tensor,
    attn_qkv_w: Tensor,
    attn_qkv_b: Tensor,
    attn_out_w: Tensor,
    attn_out_b: Tensor,
    ln_2_g: Tensor,
    ln_2_b: Tensor,
    mlp_fc_w: Tensor,
    mlp_fc_b: Tensor,
    mlp_proj_w: Tensor,
    mlp_proj_b: Tensor,
}

pub struct KVCache {
    // Per layer, per head: raw bytes (works for both f32 and bf16)
    k: Vec<Vec<Vec<u8>>>,  // [num_layers][num_heads][seq_len * head_dim * elem_size]
    v: Vec<Vec<Vec<u8>>>,
    len: usize,
    num_heads: usize,
    head_dim: usize,
    elem_size: usize,
    dtype: DType,
    device: Device,
}

impl KVCache {
    pub fn new(num_layers: usize, num_heads: usize, head_dim: usize, dtype: DType, device: Device) -> Self {
        Self {
            k: (0..num_layers).map(|_| vec![vec![]; num_heads]).collect(),
            v: (0..num_layers).map(|_| vec![vec![]; num_heads]).collect(),
            len: 0,
            num_heads,
            head_dim,
            elem_size: dtype.size_bytes(),
            dtype,
            device,
        }
    }

    pub fn seq_len(&self) -> usize { self.len }

    /// Append from a CPU tensor with shape [1, H, new_tokens, D].
    pub fn append_kv_tensor(&mut self, layer: usize, k_cpu: &Tensor, v_cpu: &Tensor, new_tokens: usize) {
        let hd = self.head_dim;
        let es = self.elem_size;
        let k_bytes = k_cpu.storage().as_cpu_bytes();
        let v_bytes = v_cpu.storage().as_cpu_bytes();
        let chunk = new_tokens * hd * es;
        for h in 0..self.num_heads {
            let off = h * chunk;
            self.k[layer][h].extend_from_slice(&k_bytes[off..off + chunk]);
            self.v[layer][h].extend_from_slice(&v_bytes[off..off + chunk]);
        }
        if layer == 0 {
            self.len += new_tokens;
        }
    }

    /// Reconstruct [1, H, seq_len, D] tensors.
    pub fn get_kv_tensors(&self, layer: usize) -> (Tensor, Tensor) {
        let sl = self.len;
        let hd = self.head_dim;
        let nh = self.num_heads;
        let es = self.elem_size;
        let head_bytes = sl * hd * es;
        let total = nh * head_bytes;
        let mut k_data = vec![0u8; total];
        let mut v_data = vec![0u8; total];
        for h in 0..nh {
            let off = h * head_bytes;
            k_data[off..off + head_bytes].copy_from_slice(&self.k[layer][h]);
            v_data[off..off + head_bytes].copy_from_slice(&self.v[layer][h]);
        }
        let shape = &[1, nh, sl, hd];
        let k = tensor_from_raw_bytes(&k_data, shape, self.dtype).to_device(self.device);
        let v = tensor_from_raw_bytes(&v_data, shape, self.dtype).to_device(self.device);
        (k, v)
    }
}

fn tensor_from_raw_bytes(bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
    match dtype {
        DType::F32 => {
            let data: &[f32] = unsafe {
                std::slice::from_raw_parts(bytes.as_ptr() as *const f32, bytes.len() / 4)
            };
            Tensor::from_slice(data, shape)
        }
        DType::BF16 => {
            let data: &[half::bf16] = unsafe {
                std::slice::from_raw_parts(bytes.as_ptr() as *const half::bf16, bytes.len() / 2)
            };
            Tensor::from_slice(data, shape)
        }
        _ => panic!("unsupported dtype for KV cache"),
    }
}

impl GPT2 {
    pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
        crate::init_kernels();
        let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
            w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
        };

        let wte = take(&mut w, "wte.weight");
        let wpe = take(&mut w, "wpe.weight");
        let ln_f_g = take(&mut w, "ln_f.weight");
        let ln_f_b = take(&mut w, "ln_f.bias");
        let lm_head = wte.transpose(0, 1).contiguous();

        let num_layers = config.num_layers();
        let mut layers = Vec::with_capacity(num_layers);
        for i in 0..num_layers {
            let p = format!("h.{i}");
            layers.push(GPT2Block {
                ln_1_g: take(&mut w, &format!("{p}.ln_1.weight")),
                ln_1_b: take(&mut w, &format!("{p}.ln_1.bias")),
                attn_qkv_w: take(&mut w, &format!("{p}.attn.c_attn.weight")),
                attn_qkv_b: take(&mut w, &format!("{p}.attn.c_attn.bias")),
                attn_out_w: take(&mut w, &format!("{p}.attn.c_proj.weight")),
                attn_out_b: take(&mut w, &format!("{p}.attn.c_proj.bias")),
                ln_2_g: take(&mut w, &format!("{p}.ln_2.weight")),
                ln_2_b: take(&mut w, &format!("{p}.ln_2.bias")),
                mlp_fc_w: take(&mut w, &format!("{p}.mlp.c_fc.weight")),
                mlp_fc_b: take(&mut w, &format!("{p}.mlp.c_fc.bias")),
                mlp_proj_w: take(&mut w, &format!("{p}.mlp.c_proj.weight")),
                mlp_proj_b: take(&mut w, &format!("{p}.mlp.c_proj.bias")),
            });
        }

        Self { config, wte, wpe, layers, ln_f_g, ln_f_b, lm_head }
    }

    /// Full forward pass without KV cache (for testing / correctness comparison).
    pub fn forward(&self, token_ids: &[u32]) -> Tensor {
        let seq_len = token_ids.len();
        let hidden = self.config.hidden();
        let num_heads = self.config.num_heads();
        let head_dim = self.config.head_dim();

        let tok_emb = embedding(&self.wte, token_ids);
        let pos_ids: Vec<u32> = (0..seq_len as u32).collect();
        let pos_emb = embedding(&self.wpe, &pos_ids);
        let mut x = add_tensors(&tok_emb, &pos_emb);

        for layer in &self.layers {
            x = self.transformer_block(layer, &x, None, 0, seq_len, num_heads, head_dim, hidden);
        }

        let x = layernorm(&x, &self.ln_f_g, &self.ln_f_b, self.config.ln_eps());
        matmul_2d(&x, &self.lm_head)
    }

    /// Forward pass with KV cache. First call = prefill, subsequent = decode.
    pub fn forward_with_cache(&self, token_ids: &[u32], cache: &mut KVCache) -> Tensor {
        let new_tokens = token_ids.len();
        let pos_offset = cache.seq_len();
        let hidden = self.config.hidden();
        let num_heads = self.config.num_heads();
        let head_dim = self.config.head_dim();

        let tok_emb = embedding(&self.wte, token_ids);
        let pos_ids: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
        let pos_emb = embedding(&self.wpe, &pos_ids);
        let mut x = add_tensors(&tok_emb, &pos_emb);

        for (layer_idx, layer) in self.layers.iter().enumerate() {
            x = self.transformer_block(
                layer, &x, Some((cache, layer_idx)),
                pos_offset, new_tokens, num_heads, head_dim, hidden,
            );
        }

        let x = layernorm(&x, &self.ln_f_g, &self.ln_f_b, self.config.ln_eps());
        matmul_2d(&x, &self.lm_head)
    }

    fn transformer_block(
        &self,
        layer: &GPT2Block,
        x: &Tensor,
        cache: Option<(&mut KVCache, usize)>,
        pos_offset: usize,
        new_tokens: usize,
        num_heads: usize,
        head_dim: usize,
        hidden: usize,
    ) -> Tensor {
        let residual = x.clone();
        let normed = layernorm(x, &layer.ln_1_g, &layer.ln_1_b, self.config.ln_eps());

        let qkv = linear(&normed, &layer.attn_qkv_w, Some(&layer.attn_qkv_b));
        let (q, k_new, v_new) = split_qkv(&qkv, num_heads, head_dim, new_tokens);

        let (k_full, v_full) = if let Some((cache, layer_idx)) = cache {
            let k_cpu = k_new.to_device(Device::Cpu);
            let v_cpu = v_new.to_device(Device::Cpu);
            cache.append_kv_tensor(layer_idx, &k_cpu, &v_cpu, new_tokens);
            cache.get_kv_tensors(layer_idx)
        } else {
            (k_new, v_new)
        };

        let attn_out = attention(&q, &k_full, &v_full, true);
        let attn_out = merge_heads(&attn_out, new_tokens, hidden);
        let attn_out = linear(&attn_out, &layer.attn_out_w, Some(&layer.attn_out_b));
        let x = add_tensors(&residual, &attn_out);

        let residual = x.clone();
        let normed = layernorm(&x, &layer.ln_2_g, &layer.ln_2_b, self.config.ln_eps());
        let fc = linear(&normed, &layer.mlp_fc_w, Some(&layer.mlp_fc_b));
        let activated = gelu(&fc);
        let proj = linear(&activated, &layer.mlp_proj_w, Some(&layer.mlp_proj_b));
        add_tensors(&residual, &proj)
    }
}

// --- Helper ops (unchanged) ---

fn linear(x: &Tensor, weight: &Tensor, bias: Option<&Tensor>) -> Tensor {
    let out = matmul_2d(x, weight);
    if let Some(b) = bias { add_bias(&out, b) } else { out }
}

fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
    assert_eq!(a.ndim(), 2);
    assert_eq!(b.ndim(), 2);
    matmul(a, b, GemmBackend::CuBlas)
}

fn add_tensors(a: &Tensor, b: &Tensor) -> Tensor {
    xserv_kernels::add(a, b)
}

fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
    // bias: [N], x: [S, N] — broadcast add via reshape
    assert_eq!(x.ndim(), 2);
    assert_eq!(bias.ndim(), 1);
    let n = bias.shape()[0];
    assert_eq!(x.shape()[1], n);
    let rows = x.shape()[0];
    // Broadcast: tile bias to [S, N] on CPU, then GPU add
    let b_cpu = bias.to_device(Device::Cpu);
    match x.dtype() {
        DType::F32 => {
            let bd = b_cpu.as_slice::<f32>();
            let tiled: Vec<f32> = (0..rows).flat_map(|_| bd.iter().copied()).collect();
            let b_full = Tensor::from_slice(&tiled, x.shape()).to_device(x.device());
            xserv_kernels::add(x, &b_full)
        }
        DType::BF16 => {
            let bd = b_cpu.as_slice::<half::bf16>();
            let tiled: Vec<half::bf16> = (0..rows).flat_map(|_| bd.iter().copied()).collect();
            let b_full = Tensor::from_slice(&tiled, x.shape()).to_device(x.device());
            xserv_kernels::add(x, &b_full)
        }
        _ => panic!("unsupported dtype"),
    }
}

fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) -> (Tensor, Tensor, Tensor) {
    let hidden = num_heads * head_dim;
    let qkv_cpu = qkv.to_device(Device::Cpu);
    let data = qkv_cpu.as_slice::<f32>();

    let mut q_data = vec![0.0f32; num_heads * seq_len * head_dim];
    let mut k_data = vec![0.0f32; num_heads * seq_len * head_dim];
    let mut v_data = vec![0.0f32; num_heads * seq_len * head_dim];

    for s in 0..seq_len {
        let row = &data[s * 3 * hidden..(s + 1) * 3 * hidden];
        for h in 0..num_heads {
            let src_off = h * head_dim;
            let dst_off = (h * seq_len + s) * head_dim;
            q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
            k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
            v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
        }
    }

    let device = qkv.device();
    let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
    let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
    let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
    (q, k, v)
}

fn merge_heads(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
    let num_heads = x.shape()[1];
    let head_dim = x.shape()[3];
    let x_cpu = x.to_device(Device::Cpu);
    let src = x_cpu.as_slice::<f32>();

    let mut out = vec![0.0f32; seq_len * hidden];
    for s in 0..seq_len {
        for h in 0..num_heads {
            let src_off = (h * seq_len + s) * head_dim;
            let dst_off = s * hidden + h * head_dim;
            out[dst_off..dst_off + head_dim].copy_from_slice(&src[src_off..src_off + head_dim]);
        }
    }
    Tensor::from_slice(&out, &[seq_len, hidden]).to_device(x.device())
}

/// Greedy sampling: return the argmax token ID from the last position's logits.
pub fn sample_greedy(logits: &Tensor) -> u32 {
    assert_eq!(logits.ndim(), 2);
    let logits_cpu = logits.to_device(Device::Cpu);
    let data = logits_cpu.as_slice::<f32>();
    let vocab_size = logits.shape()[1];
    let seq_len = logits.shape()[0];
    let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
    last_row.iter()
        .enumerate()
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
        .map(|(idx, _)| idx as u32)
        .unwrap()
}