从头写 Llama

model = AutoModelForCausalLM.from_pretrained(
  model_dir, dtype=torch.float32, attn_implementation="eager"
).eval()
ids = tok("The capital of France is", return_tensors="pt").input_ids
with torch.inference_mode():
  logits = model(ids).logits[0, -1]
  out = model.generate(ids, max_new_tokens=20, do_sample=False)
json.dump({"logits": logits.tolist(), "greedy": out[0, ids.shape[1]:].tolist()}, sys.stdout)

let mut cache = Cache::new(cfg.layers);
let logits = model.forward(&Tensor::new(ids.as_slice(), &dev)?.unsqueeze(0)?, &mut cache)?;
let last = logits.i((0, ids.len() - 1))?.to_vec1::<f32>()?;
// ...
while out.len() < steps && Some(out[out.len() - 1]) != cfg.eos {
  let x = Tensor::new(&[out[out.len() - 1]], &dev)?.unsqueeze(0)?;
  let l = model.forward(&x, &mut cache)?.i((0, 0))?.to_vec1::<f32>()?;
  out.push(argmax(&l));
}

HF_HUB_OFFLINE=1 python -I ref.py SmolLM2-135M-Instruct > pt.json
handson-scratch cpu SmolLM2-135M-Instruct > cpu.json   # release 构建的二进制
python -I compare.py pt.json cpu.json 1e-4

rope_theta=100000 head_dim=64 heads=9/3
prompt ids [504, 3575, 282, 4649, 314]
text " Paris. Paris is the largest city in France and the capital of the French department of the Espace"
load 99 ms, prefill 5 tokens 27.3 ms, decode 19 tokens 277.9 ms (68.4 tokens/s)
logits                   n=49152  max|diff|=3.67e-05
greedy                   n= 20  max|diff|=0.00e+00
worst=3.67e-05 tol=1e-04 OK

logits                   n=49152  max|diff|=4.10e-05
greedy                   n= 20  max|diff|=0.00e+00
worst=4.10e-05 tol=1e-04 OK

cpu ['4.77e-05', '4.20e-05', '4.48e-05', '4.39e-05', '3.67e-05'] argmax same [True, True, True, True, True]
metal ['1.72e-05', '3.29e-05', '5.44e-05', '4.15e-05', '4.10e-05'] argmax same [True, True, True, True, True]

    // transformers 4 写在顶层；transformers 5 存出来的 config 放在 rope_parameters 里
    let rope = &v["rope_parameters"];
    let rope_theta = v["rope_theta"]
      .as_f64()
      .or(rope["rope_theta"].as_f64())
      .context("config.json: no rope_theta, neither top level nor in rope_parameters")?;
    let rope_type = rope["rope_type"].as_str().or(v["rope_scaling"]["rope_type"].as_str());
    ensure!(matches!(rope_type, None | Some("default")), "rope_type {rope_type:?} not implemented");
    let act = v["hidden_act"].as_str().unwrap_or("silu");
    ensure!(act == "silu", "hidden_act {act} not implemented");
    for key in ["attention_bias", "mlp_bias"] {
      ensure!(!v[key].as_bool().unwrap_or(false), "{key}=true not implemented");
    }

top-level rope_theta: None
rope_parameters: {'rope_theta': 100000, 'rope_type': 'default'}

Error: config.json: no rope_theta, neither top level nor in rope_parameters

{'rope_type': 'default', 'rope_theta': 10000.0}

hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)

  fn forward(&self, x: &Tensor) -> Result<Tensor> {
    let inv_rms = (x.sqr()?.mean_keepdim(D::Minus1)? + self.eps)?.sqrt()?.recip()?;
    Ok(x.broadcast_mul(&inv_rms)?.broadcast_mul(&self.weight)?)
  }

fn rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result<Tensor> {
  let half = x.dim(D::Minus1)? / 2;
  let x1 = x.narrow(D::Minus1, 0, half)?;
  let x2 = x.narrow(D::Minus1, half, half)?;
  let y1 = (x1.broadcast_mul(cos)? - x2.broadcast_mul(sin)?)?;
  let y2 = (x1.broadcast_mul(sin)? + x2.broadcast_mul(cos)?)?;
  Ok(Tensor::cat(&[y1, y2], D::Minus1)?)
}

Cpu hand rope vs rope: 0.000e0  hand interleaved vs rope_i: 0.000e0  hand rope vs rope_i: 5.461e0
Metal { gpu_id: 4294968354 } hand rope vs rope: 2.384e-7  hand interleaved vs rope_i: 2.384e-7  hand rope vs rope_i: 6.039e0

hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

  let (b, kv, t, d) = x.dims4()?;
  Ok(x.unsqueeze(2)?.broadcast_as((b, kv, n, t, d))?.reshape((b, kv * n, t, d))?)

fn causal_mask(t: usize, past: usize, dev: &Device) -> Result<Option<Tensor>> {
  if t == 1 {
    return Ok(None); // 只有一个查询，它能看全部缓存
  }
  let total = past + t;
  let m: Vec<f32> = (0..t)
    .flat_map(|i| (0..total).map(move |j| if j > past + i { f32::NEG_INFINITY } else { 0.0 }))
    .collect();
  Ok(Some(Tensor::from_vec(m, (t, total), dev)?))
}

    let (k, v) = match kv.take() {
      Some((pk, pv)) => (Tensor::cat(&[&pk, &k], 2)?, Tensor::cat(&[&pv, &v], 2)?),
      None => (k.contiguous()?, v.contiguous()?),
    };
    *kv = Some((k.clone(), v.clone()));

  fn forward(&self, ids: &Tensor, cache: &mut Cache) -> Result<Tensor> {
    let (b, t) = ids.dims2()?;
    let past = cache.len()?;
    let rot = self.rope_table(past, t, ids.device())?;
    let mask = causal_mask(t, past, ids.device())?;

fn rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result<Tensor> {
  // 交错配对：第 2i 维和第 2i + 1 维（candle 的 rope_i，GGUF 里的 Llama 用这种）
  let (b, h, t, d) = x.dims4()?;
  let x = x.reshape((b, h, t, d / 2, 2))?;
  let x1 = x.narrow(4, 0, 1)?.squeeze(4)?;
  let x2 = x.narrow(4, 1, 1)?.squeeze(4)?;
  let y1 = (x1.broadcast_mul(cos)? - x2.broadcast_mul(sin)?)?;
  let y2 = (x1.broadcast_mul(sin)? + x2.broadcast_mul(cos)?)?;
  Ok(Tensor::stack(&[y1, y2], 4)?.reshape((b, h, t, d))?)
}

logits                   n=49152  max|diff|=7.89e+00
greedy                   n= 20  max|diff|=1.13e+04
worst=1.13e+04 tol=1e-04 MISMATCH

text " Paris.\n\n\n\nThe capital of France is Paris.\n\n\n\nThe capital of"

Error: cannot find tensor model.layers.30.input_layernorm.weight

[package]
name = "handson-scratch"
version = "0.1.0"
edition = "2024"

[dependencies]
candle-core = { version = "0.11", features = ["metal", "accelerate"] }
candle-nn = { version = "0.11", features = ["metal", "accelerate"] }
anyhow = "1"
serde_json = "1"
tokenizers = { version = "*", default-features = false, features = ["onig"] }

//! SmolLM2-135M（Llama 结构）从头写，只用 candle-core 和 candle-nn。
//! 用法：scratch <cpu|metal> <model_dir> [steps] [prompt]
//! stdout：{"logits": 最后一个位置的 logits, "greedy": 贪心生成的 id}；stderr：文本和计时。
use anyhow::{Context, Result, bail, ensure};
use candle_core::{D, DType, Device, IndexOp, Tensor};
use candle_nn::{Linear, VarBuilder, linear_no_bias};
use serde_json::Value;
use std::time::Instant;

struct Config {
  hidden: usize,
  intermediate: usize,
  layers: usize,
  heads: usize,
  kv_heads: usize,
  head_dim: usize,
  vocab: usize,
  eps: f64,
  rope_theta: f32,
  tie: bool,
  eos: Option<u32>,
}

fn int(v: &Value, key: &str) -> Result<usize> {
  let n = v[key].as_u64().with_context(|| format!("config.json: no integer `{key}`"))?;
  Ok(n as usize)
}

impl Config {
  fn from_json(v: &Value) -> Result<Self> {
    let (hidden, heads) = (int(v, "hidden_size")?, int(v, "num_attention_heads")?);
    // transformers 4 写在顶层；transformers 5 存出来的 config 放在 rope_parameters 里
    let rope = &v["rope_parameters"];
    let rope_theta = v["rope_theta"]
      .as_f64()
      .or(rope["rope_theta"].as_f64())
      .context("config.json: no rope_theta, neither top level nor in rope_parameters")?;
    let rope_type = rope["rope_type"].as_str().or(v["rope_scaling"]["rope_type"].as_str());
    ensure!(matches!(rope_type, None | Some("default")), "rope_type {rope_type:?} not implemented");
    let act = v["hidden_act"].as_str().unwrap_or("silu");
    ensure!(act == "silu", "hidden_act {act} not implemented");
    for key in ["attention_bias", "mlp_bias"] {
      ensure!(!v[key].as_bool().unwrap_or(false), "{key}=true not implemented");
    }
    Ok(Self {
      hidden,
      intermediate: int(v, "intermediate_size")?,
      layers: int(v, "num_hidden_layers")?,
      heads,
      kv_heads: int(v, "num_key_value_heads").unwrap_or(heads),
      head_dim: int(v, "head_dim").unwrap_or(hidden / heads),
      vocab: int(v, "vocab_size")?,
      eps: v["rms_norm_eps"].as_f64().context("config.json: no rms_norm_eps")?,
      rope_theta: rope_theta as f32,
      tie: v["tie_word_embeddings"].as_bool().unwrap_or(false),
      eos: v["eos_token_id"].as_u64().map(|e| e as u32),
    })
  }
}

struct RmsNorm {
  weight: Tensor,
  eps: f64,
}

impl RmsNorm {
  fn load(vb: VarBuilder, size: usize, eps: f64) -> Result<Self> {
    Ok(Self { weight: vb.get(size, "weight")?, eps })
  }

  fn forward(&self, x: &Tensor) -> Result<Tensor> {
    let inv_rms = (x.sqr()?.mean_keepdim(D::Minus1)? + self.eps)?.sqrt()?.recip()?;
    Ok(x.broadcast_mul(&inv_rms)?.broadcast_mul(&self.weight)?)
  }
}

/// 非交错 RoPE，和 transformers 的 rotate_half 同一种配对：第 i 维和第 i + d/2 维。
/// x: [b, heads, t, d]；cos、sin: [t, d/2]
fn rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result<Tensor> {
  let half = x.dim(D::Minus1)? / 2;
  let x1 = x.narrow(D::Minus1, 0, half)?;
  let x2 = x.narrow(D::Minus1, half, half)?;
  let y1 = (x1.broadcast_mul(cos)? - x2.broadcast_mul(sin)?)?;
  let y2 = (x1.broadcast_mul(sin)? + x2.broadcast_mul(cos)?)?;
  Ok(Tensor::cat(&[y1, y2], D::Minus1)?)
}

/// GQA：[b, kv, t, d] -> [b, kv * n, t, d]，KV 头 j 给查询头 j*n .. j*n+n-1 用。
fn repeat_kv(x: &Tensor, n: usize) -> Result<Tensor> {
  if n == 1 {
    return Ok(x.contiguous()?);
  }
  let (b, kv, t, d) = x.dims4()?;
  Ok(x.unsqueeze(2)?.broadcast_as((b, kv, n, t, d))?.reshape((b, kv * n, t, d))?)
}

/// 加法因果 mask：查询 i（绝对位置 past + i）只能看键 0 ..= past + i。
fn causal_mask(t: usize, past: usize, dev: &Device) -> Result<Option<Tensor>> {
  if t == 1 {
    return Ok(None); // 只有一个查询，它能看全部缓存
  }
  let total = past + t;
  let m: Vec<f32> = (0..t)
    .flat_map(|i| (0..total).map(move |j| if j > past + i { f32::NEG_INFINITY } else { 0.0 }))
    .collect();
  Ok(Some(Tensor::from_vec(m, (t, total), dev)?))
}

struct Attention {
  q_proj: Linear,
  k_proj: Linear,
  v_proj: Linear,
  o_proj: Linear,
  heads: usize,
  kv_heads: usize,
  head_dim: usize,
}

impl Attention {
  fn load(vb: VarBuilder, c: &Config) -> Result<Self> {
    let (q, kv) = (c.heads * c.head_dim, c.kv_heads * c.head_dim);
    Ok(Self {
      q_proj: linear_no_bias(c.hidden, q, vb.pp("q_proj"))?,
      k_proj: linear_no_bias(c.hidden, kv, vb.pp("k_proj"))?,
      v_proj: linear_no_bias(c.hidden, kv, vb.pp("v_proj"))?,
      o_proj: linear_no_bias(q, c.hidden, vb.pp("o_proj"))?,
      heads: c.heads,
      kv_heads: c.kv_heads,
      head_dim: c.head_dim,
    })
  }

  fn forward(&self, x: &Tensor, rot: &(Tensor, Tensor), mask: Option<&Tensor>, kv: &mut Option<(Tensor, Tensor)>) -> Result<Tensor> {
    let (b, t, _) = x.dims3()?;
    let split = |proj: &Linear, n: usize| -> Result<Tensor> {
      Ok(x.apply(proj)?.reshape((b, t, n, self.head_dim))?.transpose(1, 2)?)
    };
    let q = rope(&split(&self.q_proj, self.heads)?, &rot.0, &rot.1)?;
    let k = rope(&split(&self.k_proj, self.kv_heads)?, &rot.0, &rot.1)?;
    let v = split(&self.v_proj, self.kv_heads)?;
    // KV 缓存：存 GQA 复制之前的 k、v，沿序列维（第 2 维）接上
    let (k, v) = match kv.take() {
      Some((pk, pv)) => (Tensor::cat(&[&pk, &k], 2)?, Tensor::cat(&[&pv, &v], 2)?),
      None => (k.contiguous()?, v.contiguous()?),
    };
    *kv = Some((k.clone(), v.clone()));
    let n = self.heads / self.kv_heads;
    let (k, v) = (repeat_kv(&k, n)?, repeat_kv(&v, n)?);
    let att = (q.matmul(&k.t()?)? * (self.head_dim as f64).powf(-0.5))?;
    let att = match mask {
      Some(m) => att.broadcast_add(m)?,
      None => att,
    };
    let att = candle_nn::ops::softmax_last_dim(&att)?;
    let y = att.matmul(&v)?.transpose(1, 2)?.reshape((b, t, self.heads * self.head_dim))?;
    Ok(y.apply(&self.o_proj)?)
  }
}

struct Block {
  input_layernorm: RmsNorm,
  self_attn: Attention,
  post_attention_layernorm: RmsNorm,
  gate_proj: Linear,
  up_proj: Linear,
  down_proj: Linear,
}

impl Block {
  fn load(vb: VarBuilder, c: &Config) -> Result<Self> {
    let mlp = vb.pp("mlp");
    Ok(Self {
      input_layernorm: RmsNorm::load(vb.pp("input_layernorm"), c.hidden, c.eps)?,
      self_attn: Attention::load(vb.pp("self_attn"), c)?,
      post_attention_layernorm: RmsNorm::load(vb.pp("post_attention_layernorm"), c.hidden, c.eps)?,
      gate_proj: linear_no_bias(c.hidden, c.intermediate, mlp.pp("gate_proj"))?,
      up_proj: linear_no_bias(c.hidden, c.intermediate, mlp.pp("up_proj"))?,
      down_proj: linear_no_bias(c.intermediate, c.hidden, mlp.pp("down_proj"))?,
    })
  }

  fn forward(&self, x: &Tensor, rot: &(Tensor, Tensor), mask: Option<&Tensor>, kv: &mut Option<(Tensor, Tensor)>) -> Result<Tensor> {
    let h = self.input_layernorm.forward(x)?;
    let x = (x + self.self_attn.forward(&h, rot, mask, kv)?)?;
    let h = self.post_attention_layernorm.forward(&x)?;
    // SwiGLU：down(silu(gate(h)) * up(h))
    let m = (h.apply(&self.gate_proj)?.silu()? * h.apply(&self.up_proj)?)?.apply(&self.down_proj)?;
    Ok((x + m)?)
  }
}

/// 每层一对 (k, v)，形状 [b, kv_heads, 已缓存的长度, head_dim]。
struct Cache(Vec<Option<(Tensor, Tensor)>>);

impl Cache {
  fn new(layers: usize) -> Self {
    Self(vec![None; layers])
  }

  fn len(&self) -> Result<usize> {
    Ok(match &self.0[0] {
      Some((k, _)) => k.dim(2)?,
      None => 0,
    })
  }
}

struct Model {
  embed_tokens: Tensor,
  layers: Vec<Block>,
  norm: RmsNorm,
  lm_head: Linear,
  inv_freq: Vec<f32>,
}

impl Model {
  fn load(vb: VarBuilder, c: &Config) -> Result<Self> {
    let embed_tokens = vb.get((c.vocab, c.hidden), "model.embed_tokens.weight")?;
    let lm_head = if c.tie {
      Linear::new(embed_tokens.clone(), None) // 和词嵌入共用同一个张量
    } else {
      linear_no_bias(c.hidden, c.vocab, vb.pp("lm_head"))?
    };
    let layers = (0..c.layers)
      .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), c))
      .collect::<Result<Vec<_>>>()?;
    let inv_freq = (0..c.head_dim / 2)
      .map(|i| 1.0 / c.rope_theta.powf((2 * i) as f32 / c.head_dim as f32))
      .collect();
    Ok(Self { embed_tokens, layers, norm: RmsNorm::load(vb.pp("model.norm"), c.hidden, c.eps)?, lm_head, inv_freq })
  }

  /// 位置 start .. start + t 的 cos、sin，形状 [t, head_dim / 2]
  fn rope_table(&self, start: usize, t: usize, dev: &Device) -> Result<(Tensor, Tensor)> {
    let half = self.inv_freq.len();
    let angles: Vec<f32> = (start..start + t)
      .flat_map(|p| self.inv_freq.iter().map(move |f| p as f32 * f))
      .collect();
    let angles = Tensor::from_vec(angles, (t, half), dev)?;
    Ok((angles.cos()?, angles.sin()?))
  }

  /// ids: [b, t] -> 每个位置的 logits [b, t, vocab]。位置从缓存长度推出来，不用调用方传。
  fn forward(&self, ids: &Tensor, cache: &mut Cache) -> Result<Tensor> {
    let (b, t) = ids.dims2()?;
    let past = cache.len()?;
    let rot = self.rope_table(past, t, ids.device())?;
    let mask = causal_mask(t, past, ids.device())?;
    let mut x = self.embed_tokens.index_select(&ids.flatten_all()?, 0)?.reshape((b, t, ()))?;
    for (layer, kv) in self.layers.iter().zip(cache.0.iter_mut()) {
      x = layer.forward(&x, &rot, mask.as_ref(), kv)?;
    }
    Ok(self.norm.forward(&x)?.apply(&self.lm_head)?)
  }
}

fn main() -> Result<()> {
  let args: Vec<String> = std::env::args().collect();
  let usage = "usage: scratch <cpu|metal> <model_dir> [steps] [prompt]";
  let dev = match args.get(1).map(String::as_str) {
    Some("cpu") => Device::Cpu,
    Some("metal") => Device::new_metal(0)?,
    _ => bail!(usage),
  };
  let model_dir = args.get(2).context(usage)?;
  let steps: usize = args.get(3).map_or(Ok(20), |s| s.parse())?;
  let prompt = args.get(4).map_or("The capital of France is", String::as_str);

  let t_load = Instant::now();
  let json: Value = serde_json::from_slice(&std::fs::read(format!("{model_dir}/config.json"))?)?;
  let cfg = Config::from_json(&json)?;
  eprintln!("rope_theta={} head_dim={} heads={}/{}", cfg.rope_theta, cfg.head_dim, cfg.heads, cfg.kv_heads);
  let file = format!("{model_dir}/model.safetensors");
  let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[file], DType::F32, &dev)? };
  let model = Model::load(vb, &cfg)?;
  let tok = tokenizers::Tokenizer::from_file(format!("{model_dir}/tokenizer.json")).map_err(anyhow::Error::msg)?;
  let ids = tok.encode(prompt, true).map_err(anyhow::Error::msg)?.get_ids().to_vec();
  let load = t_load.elapsed();

  // prefill：整段 prompt 一次送进去
  let mut cache = Cache::new(cfg.layers);
  let t_prefill = Instant::now();
  let logits = model.forward(&Tensor::new(ids.as_slice(), &dev)?.unsqueeze(0)?, &mut cache)?;
  let last = logits.i((0, ids.len() - 1))?.to_vec1::<f32>()?;
  let prefill = t_prefill.elapsed();

  // decode：每步只送一个 token，前面的 k、v 从缓存里取
  let argmax = |v: &[f32]| (0..v.len()).fold(0, |best, i| if v[i] > v[best] { i } else { best }) as u32;
  let mut out = vec![argmax(&last)];
  let t_decode = Instant::now();
  while out.len() < steps && Some(out[out.len() - 1]) != cfg.eos {
    let x = Tensor::new(&[out[out.len() - 1]], &dev)?.unsqueeze(0)?;
    let l = model.forward(&x, &mut cache)?.i((0, 0))?.to_vec1::<f32>()?;
    out.push(argmax(&l));
  }
  let decode = t_decode.elapsed();

  eprintln!("prompt ids {ids:?}");
  eprintln!("text {:?}", tok.decode(&out, false).map_err(anyhow::Error::msg)?);
  eprintln!(
    "load {:.0} ms, prefill {} tokens {:.1} ms, decode {} tokens {:.1} ms ({:.1} tokens/s)",
    load.as_secs_f64() * 1e3,
    ids.len(),
    prefill.as_secs_f64() * 1e3,
    out.len() - 1,
    decode.as_secs_f64() * 1e3,
    (out.len() - 1) as f64 / decode.as_secs_f64()
  );
  println!("{}", serde_json::json!({ "logits": last, "greedy": out }));
  Ok(())
}