Function forward 

pub fn forward(
    token_id: u32,
    pos: usize,
    graph: &Graph,
    weights: &ModelWeights,
    cache: &mut KVCache,
) -> Vec<f32>

Run a single forward pass for one token through the model.

Returns logits of shape [vocab_size].