feat: Phase 3 — decoder-only transformer forward pass

[aarambh-darshan]

Summary

Implements Phase 3 — Model Architecture from ROADMAP.md (lines 203–243).

Full decoder-only transformer forward pass: token IDs in, [batch, seq, vocab_size] logits out — verified on both nano and small configs with random weights.

---

Architecture (ARCHITECTURE.md §5)

token_ids [B, T]
  → TokenPositionEmbedding   token embed + learned absolute position embed (summed)
  → TransformerBlock × N     pre-norm RMSNorm → CausalSelfAttention → residual
                             pre-norm RMSNorm → SwiGluFfn             → residual
  → RMSNorm (final)
  → LM head                  n_embd → vocab_size (weight-tied to token embedding)
  → logits [B, T, vocab_size]

---

Files Added / Changed

File	Role
src/model/mod.rs	Module declarations + PocModel re-export
src/model/embedding.rs	TokenPositionEmbedding — two Embedding tables summed element-wise
src/model/ffn.rs	SwiGluFfn — silu(x@Wgate)*(x@Wup)@Wdown, no bias
src/model/attention.rs	CausalSelfAttention — standard MHA + additive causal mask, no GQA
src/model/block.rs	TransformerBlock — pre-norm + residual for both sublayers
src/model/transformer.rs	PocModel — full pipeline + PocModel::random_init test helper
src/main.rs	Added mod model; declaration

---

Key Design Decisions

Weight tying

When cfg.tie_embeddings == true (default for both presets), the LM head reuses the token embedding weight matrix via reshape + matmul:

x [B,T,C] → [B*T,C] @ embed_weight^T [C,V] → [B*T,V] → reshape [B,T,V]

No separate lm_head.weight tensor is allocated, halving the output-layer parameter count.

Causal mask

Built inside forward from x.device() so it always lives on the correct device. Uses additive -inf mask (not multiplicative) applied before softmax. Explicitly expand()-ed to [B, n_head, T, T] — candle does not auto-broadcast via +.

No GQA, no RoPE

Standard MHA with full Q/K/V per head. Learned absolute position embeddings. Both choices follow ARCHITECTURE.md §5.2 and §5.4 — they minimise confounds in the forgetting measurement that starts at Phase 7.

---

Tests — 27/27 pass

Test	What it verifies
test_forward_shape_nano	[2,16] → [2,16,2000] ✅
test_forward_shape_small	[1,32] → [1,32,4000] ✅
test_causal_mask_end_to_end	logits at pos 0–1 bit-identical when only pos 2+ tokens differ ✅
test_no_nan_inf_nano	no NaN/Inf on non-trivial token IDs, nano ✅
test_no_nan_inf_small	no NaN/Inf on non-trivial token IDs, small ✅
test_weight_tying_enabled_by_default	nano preset has tie_embeddings=true ✅
test_causal_mask_isolation (attention unit)	attention-level mask isolation ✅
+ 20 prior Phase 0–2 tests	all still passing ✅

cargo test   →  27 passed; 0 failed
cargo clippy -- -D warnings  →  0 warnings

---

CI

The existing CI (.github/workflows/ci.yml) runs on this PR:

Job	What it runs
msrv	cargo check on Rust 1.89.0
quality	cargo fmt --check → cargo check → cargo clippy -D warnings → cargo test → release build → CLI smoke
security-audit	rustsec/audit-check on dependencies

All three jobs are expected to pass — local verification already confirms build, clippy, and test results match CI requirements.

---

Next Phase

Phase 4 — Stage 1 training + checkpoint save

• src/checkpoint.rs — save_checkpoint / load_checkpoint (real .safetensors)
• src/train.rs — AdamW training loop, train_stage1
• main.rs — train-stage1 CLI subcommand

Milestone: cargo run -- train-stage1 --config nano --steps 2000 converges and writes checkpoints/stage1/.