diff --git a/.gitignore b/.gitignore
index 42ec8e0..06b7cbb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,8 +2,10 @@
 
 # Large model weights (download separately)
 codegen_weights/pytorch_model.bin
+codegen_weights/*.safetensors
 codegen_weights/.cache/
 glm_checkpoint.txt
+glm_checkpoint/
 
 # Editor/OS
 *.swp
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 72f48d0..35f4fe8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,31 +7,164 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed — found while validating against the real 350M checkpoint
+- **`codegen download` reported success without downloading anything.**
+  `huggingface-cli` has been renamed to `hf`; the old name now prints a deprecation
+  notice, downloads nothing, and exits 0. The command believed the exit code and
+  printed "✓ Download complete!" over an empty directory. It now tries `hf` first and
+  verifies the file exists rather than trusting the status.
+- **`complete` printed nothing in its default streaming mode.**
+  `CodeGenGenerator::with_tokenizer` existed but was never called, so `decode_token`
+  returned an empty string for every token.
+- **`chat` sliced off the wrong tokens** — the same generated-only slicing bug fixed in
+  `complete`, `repl` and the server last pass, missed in this fourth call site.
+- **`-t` was bound to both `--temperature` and `--template`**, and `--stream` was a flag
+  with `default_value = "true"`, so it could never be false and the non-streaming branch
+  was unreachable. `--template` loses its short flag; `--stream` becomes `--no-stream`.
+- **`info` printed hardcoded model values** that had drifted from the checkpoint
+  (`max_seq_len` 1024 against the real 2048), while parsing the real config a few lines
+  later, and only looked for `pytorch_model.bin` so it missed a converted safetensors file.
+
+### Fixed — training
+- **`glm-train` did not train.** Both cross-entropy paths accumulated into an `f32` and
+  returned `Tensor::new(total, device)` — a fresh leaf with no autograd history — so
+  `optimizer.backward_step(&loss)` received an empty gradient store and every step was a
+  no-op. Rebuilt from tensor operations around `candle_nn::loss::cross_entropy`. A real
+  80-step run now goes from loss 10.80 to 4.99; before the fix, 60 optimizer steps left
+  the loss at 4.1934047 both before and after.
+- **Every checkpoint save failed.** `save_safetensors` called `tensor.to_vec1::<u8>()`,
+  which errors on anything of rank > 1, so training ended in
+  "unexpected rank, expected: 1, got: 2" as soon as it reached a weight matrix.
+- **`configs/train.yaml` did not deserialize.** `eval_steps` sat under a separate
+  `evaluation:` section and `tokenizer_path` was absent. Nothing had ever loaded the file
+  because `glm-train` had no `--config` flag. Config structs now use `#[serde(default)]`,
+  so a partial YAML falls back to defaults field by field.
+- `gradient_accumulation_steps` and `max_grad_norm` were declared in `TrainingConfig` and
+  claimed here as delivered, but read nowhere. Both are now implemented:
+  `gradient_accumulation_steps` micro-batch losses are averaged into one backward, and
+  `clip_grad_norm` scales gradients to the global-norm bound.
+- The learning rate was set *after* `optimizer.step`, so every step ran on the previous
+  step's rate.
+- Masking used a hand-rolled LCG re-seeded from the step counter, so consecutive steps drew
+  near-identical corruption patterns. Replaced with one `StdRng` seeded from
+  `TrainingConfig::seed`.
+
+### Added — training
+- `--config <path>` on `glm-train`, using the existing `TrainConfig::from_file`. It was
+  marked done in the project notes but had never been wired to the CLI.
+- `src/training/loss.rs` with tests that fail if gradients stop reaching parameters, if a
+  tiny model stops memorising a batch, or if an optimizer step leaves weights unchanged
+- Gradient-clipping tests, and a test that `configs/train.yaml` actually parses
+
+### Removed — training
+- `src/training/checkpoint.rs` (192 lines): never declared in `mod.rs`, so it had never
+  been compiled, and it duplicated the checkpoint code in `train.rs`
+- Per-tensor `println!` debug output from `TrainableGLMModel::save_safetensors`
+- Three copies of the sequence-corruption logic and two of the tensor-to-safetensors
+  conversion, collapsed into one each
+
+### Performance
+- **KV cache no longer rewrites itself on every token.** `KVCache::append` used
+  `Tensor::slice_assign`, which is not a targeted write: candle zero-pads the source out
+  to the full destination shape, allocates a full-size mask, pads that too, and runs
+  `where_cond` over every element — the entire `[1, heads, max_seq_len, head_dim]` buffer,
+  for `k` and `v`, in every layer, for every token. Replaced with `Tensor::slice_set`,
+  which copies only the new tokens in place. On the benchmark model: prefill
+  88.5 ms → 36.2 ms, and decode 58.4 ms → 5.6 ms per token, a 10.4× improvement.
+- **`lm_head` no longer runs over the whole prompt.** Generation reads only the final
+  position but paid the `hidden_dim × vocab_size` projection for every prefill token —
+  6.0 ms of a 35.8 ms prefill. `forward_hidden` and `project_logits` are now separate, and
+  generation narrows before projecting.
+- `Embedding::forward` uses one `index_select` instead of a `get` plus `stack` per token.
+  No measurable difference at benchmark sizes; it is simply less code.
+- Removed the per-token `eprintln!` timing lines from `generate_stream` — a stderr flush
+  per token inside the generation loop.
+
+### Fixed
+- **FP16 models built by `CodeGenModel::new_blank` panicked on the first forward.**
+  `MultiHeadAttention::new_blank`, the FFN constructors and `Embedding::zeros` hardcoded
+  F32 regardless of `config.dtype`, giving `dtype mismatch in mul, lhs: F32, rhs: F16`.
+  Loading real weights masked this, because every tensor was replaced on the way in.
+- **`cargo clippy --all-targets --all-features -- -D warnings` failed on every commit.**
+  `main.rs` re-declared the whole module tree instead of using the library crate, so the
+  binary compiled a second copy and every `pub` item the CLI did not call was reported as
+  dead code — 34 warnings from one cause. The CI lint gate now passes.
+- `KVCache::append` reports a clear error when a prompt exceeds `max_seq_len` instead of
+  failing inside candle.
+- **CodeGen produced incorrect output with real weights.** The fused `qkv_proj`
+  tensor was split as `[all q | all v | all k]`, but CodeGen stores it as four
+  interleaved model-parallel groups of `[q | v | k]` (`mp_num = 4` upstream).
+  Every attention head was reading the wrong slice of the projection.
+- **Sampling was always greedy.** The hand-rolled LCG in `sampling::sample`
+  divided a `u64` state by `u32::MAX`, producing values around 6e9, so the
+  selection loop always fell through to argmax — `temperature`, `top_k` and
+  `top_p` had no effect. Replaced with a `StdRng` seeded once per generation.
+- **`complete`, `repl` and `/generate` sliced off the wrong tokens.**
+  `CodeGenGenerator::generate` returns generated tokens only, but all three
+  callers stripped `prompt.len()` tokens from the front, dropping output or
+  panicking when the prompt was longer than the completion.
+- **The HTTP server ignored `config.json` and `--f16`**, loading with
+  `CodeGenConfig::default()` (`rotary_dim` 64 rather than the checkpoint's 32).
+  It now shares `ModelContext` with the CLI.
+- **Missing checkpoint tensors were silently ignored**, leaving zeroed layers
+  that produce nonsense instead of an error. `WeightLoader` now names the first
+  missing tensor and fails.
+- `TrainingConfig::to_glm_config()` now uses `self` instead of hardcoded defaults
+- Hardcoded `0..6` layer loops replaced with `self.config.num_layers` / `model.num_layers`
+- `GLMTrainer::from_config` now accepts `TrainConfig` (outer struct) so `to_glm_config()` resolves correctly
+- Hardcoded tokenizer path now configurable via `TrainingConfig::tokenizer_path`
+- Removed unused `_dtype_str` variable in `model.rs`
+
 ### Added
+- `CodeGenTokenizer::vocab_size`, reported by `info`. CodeGen pads `vocab_size` to 51200
+  while the tokenizer stops at 50294; the 905 untrained rows never win a sample, verified
+  at temperature 2.0 with no top-k or nucleus cut, and an integration test keeps that true
+- Integration tests that run against the real checkpoint: sampling stays inside the
+  tokenizer's vocabulary, and a fixed seed reproduces while a different seed diverges
+- `WeightLoader::load_from_safetensors` and `WeightLoader::load` (format picked by
+  extension). `examples/convert_codegen_to_safetensors.rs` wrote a file nothing in the repo
+  could read; `ModelContext` now prefers `model.safetensors` over `pytorch_model.bin`
+- Round-trip test covering the converter's output
+- `KVCache` test asserting token-by-token appends match a single block append
+- `--seed` global flag for reproducible sampling
+- Numerical parity tests (`tests/codegen_parity.rs`) against the HuggingFace
+  CodeGen reference, using committed tiny-model fixtures generated by
+  `scripts/gen_parity_fixture.py` — no weight download needed
 - MIT License file
 - Crate-level rustdoc for `lib.rs`
 - `rust-version = "1.75.0"` MSRV in `Cargo.toml`
 - `tokenizer_path` field in `TrainingConfig` (configurable, defaults to `codegen_weights/tokenizer.json`)
 - `num_layers` field stored in `TrainableGLMModel` for param enumeration
 - `Default` impl for `TrainConfig` (outer config struct)
-- 69 unit tests covering config, data, layers, GLM, codegen, sampling, training, and generation
+- Unit tests covering config, data, layers, GLM, codegen, sampling, training, and generation
 - 2 integration tests for CodeGen weight loading and tokenizer roundtrip
 - CI workflows for benchmarks, security audit, and code coverage
 - Branch triggers for `codegen` on all CI workflows
 
-### Fixed
-- `TrainingConfig::to_glm_config()` now uses `self` instead of hardcoded defaults
-- Hardcoded `0..6` layer loops replaced with `self.config.num_layers` / `model.num_layers`
-- `GLMTrainer::from_config` now accepts `TrainConfig` (outer struct) so `to_glm_config()` resolves correctly
-- Hardcoded tokenizer path now configurable via `TrainingConfig::tokenizer_path`
-- Removed unused `_dtype_str` variable in `model.rs`
-
 ### Changed
+- **`benches/transformer.rs` now benchmarks this crate.** Every benchmark was a hand-inlined
+  copy of the model, including a QKV split carrying the bug the real model no longer has, so
+  the suite could report neither a speedup nor a regression. Replaced with benchmarks that
+  call the public API: prefill, prefill without the vocabulary projection, forward passes,
+  the full generator, weight loading, and F32 against F16.
+- Blank constructors (`Embedding::zeros`, `MultiHeadAttention::new_blank`, the FFN
+  variants) take an explicit `DType`
+- `sampling::sample` takes an RNG instead of a `u64` seed, so one stream spans a
+  whole generation
+- `server::start_server` takes the weights directory and `--f16` flag rather than
+  a `.bin` path
 - Removed blanket `#![allow(dead_code)]` from `lib.rs` and `main.rs`
 - `cli` and `model` modules now re-exported from `lib.rs`
-- README updated with 69 tests, full CLI reference, deps table, mermaid project structure
+- README updated with full CLI reference, deps table, mermaid project structure
 - Docs updated: `run-codegen-on-raspberry-pi.md`, `train-code-infill.md` (correct CLI commands)
 
+### Removed
+- **`src/codegen/quantized.rs`.** Never constructed outside its own tests, and
+  `QuantizedLinear::forward` dequantized the whole weight matrix on every call, making it
+  strictly slower than the F32 path it would have replaced. A real speedup needs int8
+  matmul kernels, which candle does not expose. The README and docs no longer claim it.
+- `CodeGenBlock::norm2`, unused: CodeGen blocks are parallel and have one norm
+
 ## [0.1.0] - 2024-01-01
 
 ### Added
diff --git a/README.md b/README.md
index fea77ef..771bd6f 100644
--- a/README.md
+++ b/README.md
@@ -9,12 +9,12 @@
 [![Rust](https://img.shields.io/badge/Rust-2021-orange?logo=rust)](https://www.rust-lang.org/)
 [![Candle](https://img.shields.io/badge/Candle-0.8-blue)](https://github.com/huggingface/candle)
 [![License](https://img.shields.io/badge/License-MIT-green)](LICENSE)
-[![Tests](https://img.shields.io/badge/Tests-69%2F69%20✓-brightgreen)]()
+[![Tests](https://img.shields.io/badge/Tests-102%2F102%20✓-brightgreen)]()
 [![CI](https://img.shields.io/badge/CI-GitHub%20Actions-blue?logo=githubactions)](.github/workflows/ci.yml)
 
 <br>
 
-*No GPU required. No Python runtime. Pure Rust tensor ops — 350M params on an i5-6600.*
+*No GPU required. No Python runtime. Pure Rust tensor ops — 350M params, 20 ms/token on a laptop CPU.*
 
 </div>
 
@@ -49,7 +49,6 @@ graph LR
         C1[RoPE Rotary Embedding]
         C2[KV Cache]
         C3[Parallel Attn + FFN]
-        C4[INT8 Quantization]
     end
 
     E --> TB
@@ -155,15 +154,43 @@ graph TD
 
 ## ⚡ Performance
 
-*Benchmarked on i5-6600 (4C/4T, 7.5GB RAM) — release mode*
+`cargo bench` drives the library directly. The benchmark model is a stand-in for
+CodeGen-350M — `hidden_dim` 512, 12 layers, 8 heads, vocab 16384 — but keeps the real
+`max_seq_len` of 2048, because KV-cache cost scales with the cache buffer rather than with
+parameter count.
 
-| Operation | Time |
+*Apple M1 Pro, release mode:*
+
+| Benchmark | Time |
 |:----------|:-----|
-| Weight loading (350M) | ~0.5s |
-| Prefill (7 tokens) | ~0.3s |
-| Autoregressive step | ~0.1s/token |
-| 37-token generation | ~3.2s |
-| GLM training step | ~0.02s |
+| `prefill/32_tokens` | 35.8 ms |
+| `prefill/32_tokens_hidden_only` (no vocabulary projection) | 29.8 ms |
+| `generator/32_prompt_16_new` (prefill + decode + sampling) | 110.9 ms |
+| Decode, per token | ~5.6 ms |
+| `weight_load/tiny_h4` | 1.38 ms |
+| `prefill_dtype/f32` vs `prefill_dtype/f16` | 41.8 ms vs 29.4 ms |
+
+Reproduce with `cargo bench --bench transformer`, or refresh this table with
+`bash scripts/update-readme-benchmarks.sh`.
+
+### CodeGen-350M, real weights
+
+`Salesforce/codegen-350M-multi`, 64 tokens from `def quicksort(arr):`, greedy,
+Apple M1 Pro:
+
+| | F32 | F16 |
+|:--|--:|--:|
+| Generation, 64 tokens | 4.4 s | **1.3 s** |
+| Per token | 68.8 ms | **20.3 ms** |
+| Wall clock including weight load | 5.0 s | 1.7 s |
+| Peak resident memory | 1.63 GB | 1.08 GB |
+
+Reproduce with:
+
+```bash
+cargo run --release -- download
+cargo run --release -- --f16 complete "def quicksort(arr):" --max-tokens 64 --temperature 0.0 --no-stream
+```
 
 > ⚠️ Debug builds are ~20× slower. Always use `--release`.
 
@@ -210,7 +237,6 @@ graph TD
         CM["model.rs"]
         CR["rotary.rs<br/>RoPE"]
         CW["weights.rs<br/>PyTorch Loader"]
-        CQ["quantized.rs<br/>INT8 Quantization"]
         CK["kv_cache.rs"]
     end
 
@@ -299,13 +325,14 @@ cargo run --release -- glm-train --data-path data --steps 500
 |:-----|:------------|
 | `--f16` | Use FP16 precision (faster, less memory) |
 | `--weights-dir DIR` | Path to weights directory (default: `codegen_weights`) |
+| `--seed N` | Fixed sampling seed for reproducible output |
 
 ### Subcommand Reference
 
 | Command | Description | Options |
 |:--------|:------------|:--------|
 | `chat` | Multi-turn conversational code generation | `--system <prompt>` |
-| `complete <prompt>` | Single-shot code generation | `--max-tokens`, `--temperature`, `--template`, `--stream` |
+| `complete <prompt>` | Single-shot code generation | `--max-tokens`, `--temperature`, `--template`, `--no-stream` |
 | `repl` | Interactive REPL (single-turn) | — |
 | `info` | Print model info and weight status | — |
 | `download` | Download CodeGen-350M weights from HuggingFace | — |
@@ -349,19 +376,13 @@ curl -X POST http://localhost:8080/generate \
 
 ---
 
-## ⚖️ Quantization
-
-### INT8 Dynamic Quantization
-
-Per-channel INT8 quantization reduces model size by ~4x with minimal quality loss:
-
-- **Method**: Per-channel symmetric quantization with offset (u8 + 128)
-- **Compression**: ~4x for large linear layers
-- **Ranking**: Preserves relative token rankings
+## ⚖️ Precision
 
 ### FP16 Inference
 
-Full dtype propagation through all layers — 23% speedup on CPU:
+Full dtype propagation through all layers, including blank-initialised models.
+On real CodeGen-350M weights it is **3.4× faster** than F32 (20.3 ms/token against
+68.8 ms) and uses a third less memory:
 
 ```bash
 cargo run --release -- --f16 chat
@@ -377,24 +398,26 @@ The GLM model supports training from scratch with a production-grade pipeline:
 ### Configuration
 
 ```yaml
-# configs/train.yaml
+# configs/train.yaml — every field is optional and falls back to its default
 model:
+  vocab_size: 51200
   hidden_dim: 256
   num_layers: 6
   num_heads: 8
   ffn_dim: 1024
-  max_seq_len: 128
-  vocab_size: 16384
+  max_seq_len: 512
 
 training:
-  batch_size: 8
-  learning_rate: 0.0003
-  num_steps: 1000
+  learning_rate: 1e-4
+  max_grad_norm: 1.0
+  micro_batch_size: 1              # sequences per forward pass
+  gradient_accumulation_steps: 32  # forward passes per optimizer step
+  max_steps: 10000
 ```
 
 ### Features
 
-- **YAML config** — Declarative training configuration
+- **YAML config** — Declarative training configuration via `--config`
 - **DataLoader** — Train/eval split, shuffling, random windowing
 - **LR Scheduler** — Cosine decay with linear warmup
 - **Safetensors Checkpoints** — Save/load model + optimizer state
@@ -405,7 +428,11 @@ training:
 ### Quick Start
 
 ```bash
+# defaults
 cargo run --release -- glm-train --data-path data --steps 500
+
+# or drive it from a YAML config
+cargo run --release -- glm-train --config configs/train.yaml --steps 500
 ```
 
 ---
@@ -446,8 +473,7 @@ cargo run --release -- glm-train --data-path data --steps 500
 - **Blank-Infilling** — Bidirectional context with causal within-blank masking
 - **Sampling Pipeline** — Repetition penalty → temperature → top-k → top-p → random sample
 - **Zero-Init Loading** — Avoids allocating 350M random floats before overwriting with weights
-- **FP16 Inference** — Full dtype propagation for 23% speedup
-- **INT8 Quantization** — Per-channel symmetric quantization, ~4x compression
+- **FP16 Inference** — Full dtype propagation through every layer
 - **Token Streaming** — Stream tokens as they're generated
 - **Prompt Templates** — Completion, instruct, and chat templates
 - **Multi-Turn Chat** — Conversational code generation with history
@@ -499,7 +525,6 @@ cargo test
 | CodeGen Config | Defaults, head_dim, HF config parsing |
 | CodeGen KV Cache | New, append, reset, dtype support |
 | CodeGen Model | Blank-forward, RoPE no-segfault |
-| Quantized | INT8 quantized linear roundtrip, ranking preservation |
 | Sampling | Argmax, temperature-zero |
 | Training Config | Defaults, YAML serialization, GLM config conversion |
 | Training Data | DataLoader, batch, truncation, split, reset |
diff --git a/benches/transformer.rs b/benches/transformer.rs
index 016088d..c75bcc7 100644
--- a/benches/transformer.rs
+++ b/benches/transformer.rs
@@ -1,561 +1,171 @@
-use candle_core::{DType, Device, Tensor};
-use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
-
-fn bench_rope_2d(c: &mut Criterion) {
-    let device = Device::Cpu;
-    let mut group = c.benchmark_group("rope_2d");
-
-    for &hidden_dim in &[256, 512, 1024] {
-        let max_positions = 512;
-        let pos_1 = Tensor::randn(0.0f32, 0.02f32, (max_positions, hidden_dim), &device).unwrap();
-        let pos_2 = Tensor::randn(0.0f32, 0.02f32, (max_positions, hidden_dim), &device).unwrap();
-
-        for &seq_len in &[32, 64, 128] {
-            let pos_1_ids: Vec<u32> = (0..seq_len as u32).map(|i| i % 4).collect();
-            let pos_2_ids: Vec<u32> = (0..seq_len as u32).collect();
-
-            group.throughput(Throughput::Elements((seq_len * hidden_dim) as u64));
-            group.bench_with_input(
-                BenchmarkId::new("f32", format!("h{}_s{}", hidden_dim, seq_len)),
-                &(&pos_1, &pos_2, &pos_1_ids, &pos_2_ids, hidden_dim, seq_len),
-                |b, (pos_1, pos_2, pos_1_ids, pos_2_ids, hidden_dim, seq_len)| {
-                    b.iter(|| {
-                        let p1_tensor = Tensor::from_slice(pos_1_ids, *seq_len, &device).unwrap();
-                        let p2_tensor = Tensor::from_slice(pos_2_ids, *seq_len, &device).unwrap();
-                        let e1 = pos_1.index_select(&p1_tensor, 0).unwrap();
-                        let e2 = pos_2.index_select(&p2_tensor, 0).unwrap();
-                        e1.broadcast_add(&e2).unwrap().unsqueeze(0).unwrap()
-                    });
-                },
-            );
-        }
-    }
-    group.finish();
-}
-
-fn bench_kv_cache(c: &mut Criterion) {
-    let device = Device::Cpu;
-    let mut group = c.benchmark_group("kv_cache");
-
-    for &(n_heads, head_dim) in &[(8, 64), (16, 64), (8, 128)] {
-        let max_seq = 512;
-
-        group.bench_function(
-            BenchmarkId::new("append_1", format!("nh{}_hd{}", n_heads, head_dim)),
-            |b| {
-                let mut cache = candle_core::Tensor::zeros(
-                    (1, n_heads, max_seq, head_dim),
-                    DType::F32,
-                    &device,
-                )
-                .unwrap();
-                let mut pos = 0usize;
-                let k_new =
-                    Tensor::randn(0.0f32, 1.0f32, (1, n_heads, 1, head_dim), &device).unwrap();
-                let v_new =
-                    Tensor::randn(0.0f32, 1.0f32, (1, n_heads, 1, head_dim), &device).unwrap();
-
-                b.iter(|| {
-                    let end = pos + 1;
-                    cache = cache
-                        .slice_assign(&[0..1, 0..n_heads, pos..end, 0..head_dim], &k_new)
-                        .unwrap();
-                    pos = end;
-                    pos
-                });
-            },
-        );
-
-        group.bench_function(
-            BenchmarkId::new("append_8", format!("nh{}_hd{}", n_heads, head_dim)),
-            |b| {
-                let mut cache = candle_core::Tensor::zeros(
-                    (1, n_heads, max_seq, head_dim),
-                    DType::F32,
-                    &device,
-                )
-                .unwrap();
-                let mut pos = 0usize;
-                let k_new =
-                    Tensor::randn(0.0f32, 1.0f32, (1, n_heads, 8, head_dim), &device).unwrap();
-                let v_new =
-                    Tensor::randn(0.0f32, 1.0f32, (1, n_heads, 8, head_dim), &device).unwrap();
-
-                b.iter(|| {
-                    let end = pos + 8;
-                    cache = cache
-                        .slice_assign(&[0..1, 0..n_heads, pos..end, 0..head_dim], &k_new)
-                        .unwrap();
-                    pos = end;
-                    pos
-                });
-            },
-        );
-    }
-    group.finish();
-}
-
-fn bench_quantized_matmul(c: &mut Criterion) {
-    let device = Device::Cpu;
-    let mut group = c.benchmark_group("quantized_matmul");
-
-    for &hidden_dim in &[256, 512, 1024] {
-        let weight = Tensor::randn(0.0f32, 0.02f32, (hidden_dim, hidden_dim), &device).unwrap();
-
-        group.throughput(Throughput::Elements(hidden_dim as u64));
-        group.bench_with_input(
-            BenchmarkId::new("f32", format!("h{}", hidden_dim)),
-            &(&weight, hidden_dim),
-            |b, (weight, hidden_dim)| {
-                b.iter(|| {
-                    let x = Tensor::randn(0.0f32, 1.0f32, (1, 1, *hidden_dim), &device).unwrap();
-                    x.broadcast_matmul(&weight.unsqueeze(0).unwrap()).unwrap()
-                });
-            },
-        );
-
-        // Simulate INT8 quantized matmul: round weights to simulate int8 precision loss
-        let weight_q = weight.round().unwrap();
-        group.throughput(Throughput::Elements(hidden_dim as u64));
-        group.bench_with_input(
-            BenchmarkId::new("i8_simulated", format!("h{}", hidden_dim)),
-            &(&weight_q, hidden_dim),
-            |b, (weight_q, hidden_dim)| {
-                b.iter(|| {
-                    let x = Tensor::randn(0.0f32, 1.0f32, (1, 1, *hidden_dim), &device).unwrap();
-                    let x_q = x.round().unwrap();
-                    x_q.broadcast_matmul(&weight_q.unsqueeze(0).unwrap())
-                        .unwrap()
-                });
-            },
-        );
+//! Benchmarks for the CodeGen inference path.
+//!
+//! These call the library directly. An earlier version of this file re-implemented
+//! the model inline, which meant it measured a copy of the code rather than the
+//! code — including a QKV split that had a bug the real model no longer has.
+
+use std::path::Path;
+
+use candle_core::{DType, Device};
+use criterion::{criterion_group, criterion_main, Criterion, Throughput};
+
+use rust_transformer::codegen::config::CodeGenConfig;
+use rust_transformer::codegen::kv_cache::KVCache;
+use rust_transformer::codegen::model::CodeGenModel;
+use rust_transformer::codegen::weights::WeightLoader;
+use rust_transformer::generation::codegen_generate::CodeGenGenerator;
+
+/// A stand-in for CodeGen-350M: small enough to build quickly, but with the real
+/// `max_seq_len`. KV-cache cost scales with the cache buffer, not with parameter
+/// count, so shrinking `max_seq_len` would hide exactly what these measure.
+fn bench_config() -> CodeGenConfig {
+    CodeGenConfig {
+        vocab_size: 16384,
+        hidden_dim: 512,
+        num_layers: 12,
+        num_heads: 8,
+        ffn_dim: 2048,
+        max_seq_len: 2048,
+        rotary_dim: 32,
+        ..Default::default()
     }
-    group.finish();
 }
 
-pub fn bench_attention(c: &mut Criterion) {
-    let device = Device::Cpu;
-    let mut group = c.benchmark_group("attention");
-
-    for &hidden_dim in &[256, 512, 1024] {
-        for &num_heads in &[4, 8, 16] {
-            if hidden_dim % num_heads != 0 {
-                continue;
-            }
-            let head_dim = hidden_dim / num_heads;
-            let seq_len = 64;
-
-            let qkv_weight =
-                Tensor::randn(0.0f32, 0.02f32, (hidden_dim, hidden_dim * 3), &device).unwrap();
-            let out_weight =
-                Tensor::randn(0.0f32, 0.02f32, (hidden_dim, hidden_dim), &device).unwrap();
-
-            group.throughput(Throughput::Elements(
-                (seq_len * hidden_dim * num_heads) as u64,
-            ));
-            group.bench_with_input(
-                BenchmarkId::new("f32", format!("h{}_nh{}", hidden_dim, num_heads)),
-                &(
-                    &qkv_weight,
-                    &out_weight,
-                    hidden_dim,
-                    num_heads,
-                    head_dim,
-                    seq_len,
-                ),
-                |b, (qkv_weight, out_weight, hidden_dim, num_heads, head_dim, seq_len)| {
-                    b.iter(|| {
-                        let x = Tensor::randn(0.0f32, 1.0f32, (1, *seq_len, *hidden_dim), &device)
-                            .unwrap();
-                        let qkv = x
-                            .broadcast_matmul(&qkv_weight.unsqueeze(0).unwrap())
-                            .unwrap();
-                        let qkv = qkv
-                            .reshape((1, *seq_len, 3, *num_heads, *head_dim))
-                            .unwrap();
-                        let qkv = qkv.permute((0, 3, 2, 1, 4)).unwrap();
-                        let q = qkv.get_on_dim(2, 0).unwrap();
-                        let v = qkv.get_on_dim(2, 1).unwrap();
-                        let k = qkv.get_on_dim(2, 2).unwrap();
-
-                        let scale = 1.0 / (*head_dim as f64).sqrt();
-                        let scores = q.broadcast_matmul(&k.transpose(2, 3).unwrap()).unwrap();
-                        let scores = (scores * scale).unwrap();
+const PROMPT_LEN: usize = 32;
+const NEW_TOKENS: usize = 16;
 
-                        let weights = candle_nn::ops::softmax(&scores, 3).unwrap();
-                        let context = weights.broadcast_matmul(&v).unwrap();
-                        let context = context
-                            .permute((0, 2, 1, 3))
-                            .unwrap()
-                            .reshape((1, *seq_len, *hidden_dim))
-                            .unwrap();
-                        context.broadcast_matmul(&out_weight.unsqueeze(0).unwrap())
-                    });
-                },
-            );
-        }
-    }
-    group.finish();
+fn prompt_tokens(config: &CodeGenConfig) -> Vec<u32> {
+    (0..PROMPT_LEN)
+        .map(|i| (i * 7 % config.vocab_size) as u32)
+        .collect()
 }
 
-pub fn bench_ffn(c: &mut Criterion) {
+/// Prefill: one forward pass over the whole prompt, empty cache.
+fn bench_prefill(c: &mut Criterion) {
     let device = Device::Cpu;
-    let mut group = c.benchmark_group("ffn");
-
-    for &hidden_dim in &[256, 512, 1024] {
-        for &ffn_dim in &[1024, 2048, 4096] {
-            let seq_len = 64;
-
-            let fc_in = Tensor::randn(0.0f32, 0.02f32, (hidden_dim, ffn_dim), &device).unwrap();
-            let fc_out = Tensor::randn(0.0f32, 0.02f32, (ffn_dim, hidden_dim), &device).unwrap();
-
-            group.throughput(Throughput::Elements((seq_len * hidden_dim) as u64));
-            group.bench_with_input(
-                BenchmarkId::new("gelu", format!("h{}_ffn{}", hidden_dim, ffn_dim)),
-                &(&fc_in, &fc_out, hidden_dim, ffn_dim, seq_len),
-                |b, (fc_in, fc_out, hidden_dim, _ffn_dim, seq_len)| {
-                    b.iter(|| {
-                        let x = Tensor::randn(0.0f32, 1.0f32, (1, *seq_len, *hidden_dim), &device)
-                            .unwrap();
-                        let hidden = x.broadcast_matmul(&fc_in.unsqueeze(0).unwrap()).unwrap();
-                        let activated = hidden.gelu().unwrap();
-                        activated.broadcast_matmul(&fc_out.unsqueeze(0).unwrap())
-                    });
-                },
-            );
-        }
-    }
+    let config = bench_config();
+    let model = CodeGenModel::new_blank(config.clone(), &device).unwrap();
+    let tokens = prompt_tokens(&config);
+    let positions: Vec<usize> = (0..tokens.len()).collect();
+
+    let mut group = c.benchmark_group("prefill");
+    group.throughput(Throughput::Elements(PROMPT_LEN as u64));
+    group.bench_function("32_tokens", |b| {
+        b.iter(|| {
+            let mut cache: Option<Vec<KVCache>> = None;
+            model
+                .forward_with_cache(&tokens, &positions, &mut cache)
+                .unwrap()
+        });
+    });
+    // Same work minus the vocabulary projection. The gap is what generation
+    // saves by projecting only the last position.
+    group.bench_function("32_tokens_hidden_only", |b| {
+        b.iter(|| {
+            let mut cache: Option<Vec<KVCache>> = None;
+            model
+                .forward_hidden(&tokens, &positions, &mut cache)
+                .unwrap()
+        });
+    });
     group.finish();
 }
 
-pub fn bench_layernorm(c: &mut Criterion) {
+/// The path `codegen complete` actually takes: prefill, decode, sampling.
+fn bench_generator(c: &mut Criterion) {
     let device = Device::Cpu;
-    let mut group = c.benchmark_group("layernorm");
-
-    for &hidden_dim in &[256, 512, 1024] {
-        let seq_len = 64;
-        let weight = Tensor::ones(hidden_dim, DType::F32, &device).unwrap();
-        let bias = Tensor::zeros(hidden_dim, DType::F32, &device).unwrap();
-        let eps = 1e-5;
-
-        group.throughput(Throughput::Elements((seq_len * hidden_dim) as u64));
-        group.bench_with_input(
-            BenchmarkId::new("f32", format!("h{}", hidden_dim)),
-            &(&weight, &bias, hidden_dim, seq_len, eps),
-            |b, (weight, bias, hidden_dim, seq_len, eps)| {
-                b.iter(|| {
-                    let x =
-                        Tensor::randn(0.0f32, 1.0f32, (1, *seq_len, *hidden_dim), &device).unwrap();
-                    let last_dim = x.dims().len() - 1;
-                    let mean = x.mean(last_dim).unwrap();
-                    let mean = mean.unsqueeze(last_dim).unwrap();
-                    let x_centered = x.broadcast_sub(&mean).unwrap();
-                    let variance = x_centered.sqr().unwrap().mean(last_dim).unwrap();
-                    let std = (variance + *eps).unwrap().sqrt().unwrap();
-                    let std = std.unsqueeze(last_dim).unwrap();
-                    let normalized = x_centered.broadcast_div(&std).unwrap();
-                    let weight = if weight.dtype() != normalized.dtype() {
-                        weight.to_dtype(normalized.dtype()).unwrap()
-                    } else {
-                        (*weight).clone()
-                    };
-                    let bias = if bias.dtype() != normalized.dtype() {
-                        bias.to_dtype(normalized.dtype()).unwrap()
-                    } else {
-                        (*bias).clone()
-                    };
-                    normalized
-                        .broadcast_mul(&weight)
-                        .unwrap()
-                        .broadcast_add(&bias)
-                        .unwrap()
-                });
-            },
-        );
-    }
+    let config = bench_config();
+    let model = CodeGenModel::new_blank(config.clone(), &device).unwrap();
+    let tokens = prompt_tokens(&config);
+    let generator = CodeGenGenerator::new(model, 0.8, 40, 0.9, 1.2, NEW_TOKENS);
+
+    let mut group = c.benchmark_group("generator");
+    group.throughput(Throughput::Elements(NEW_TOKENS as u64));
+    group.bench_function("32_prompt_16_new", |b| {
+        b.iter(|| generator.generate(&tokens).unwrap());
+    });
     group.finish();
 }
 
-pub fn bench_full_block(c: &mut Criterion) {
+/// Prefill plus autoregressive decode — the number that matters for `complete`.
+/// Subtract the `prefill` result to isolate per-token decode cost.
+fn bench_generate(c: &mut Criterion) {
     let device = Device::Cpu;
-    let mut group = c.benchmark_group("full_block");
-
-    for &hidden_dim in &[256, 512, 1024] {
-        for &num_layers in &[1, 6, 12] {
-            let seq_len = 64;
-            let num_heads = 8;
-            let head_dim = hidden_dim / num_heads;
-            let ffn_dim = hidden_dim * 4;
-
-            let mut layers = Vec::new();
-            for _ in 0..num_layers {
-                let qkv_weight =
-                    Tensor::randn(0.0f32, 0.02f32, (hidden_dim, hidden_dim * 3), &device).unwrap();
-                let out_weight =
-                    Tensor::randn(0.0f32, 0.02f32, (hidden_dim, hidden_dim), &device).unwrap();
-                let ln1_w = Tensor::ones(hidden_dim, DType::F32, &device).unwrap();
-                let ln1_b = Tensor::zeros(hidden_dim, DType::F32, &device).unwrap();
-                let ln2_w = Tensor::ones(hidden_dim, DType::F32, &device).unwrap();
-                let ln2_b = Tensor::zeros(hidden_dim, DType::F32, &device).unwrap();
-                let fc_in = Tensor::randn(0.0f32, 0.02f32, (hidden_dim, ffn_dim), &device).unwrap();
-                let fc_out =
-                    Tensor::randn(0.0f32, 0.02f32, (ffn_dim, hidden_dim), &device).unwrap();
-                layers.push((
-                    qkv_weight, out_weight, ln1_w, ln1_b, ln2_w, ln2_b, fc_in, fc_out,
-                ));
+    let config = bench_config();
+    let model = CodeGenModel::new_blank(config.clone(), &device).unwrap();
+    let tokens = prompt_tokens(&config);
+    let positions: Vec<usize> = (0..tokens.len()).collect();
+
+    let mut group = c.benchmark_group("generate");
+    group.throughput(Throughput::Elements(NEW_TOKENS as u64));
+    group.bench_function("32_prompt_16_new", |b| {
+        b.iter(|| {
+            let mut cache: Option<Vec<KVCache>> = None;
+            model
+                .forward_with_cache(&tokens, &positions, &mut cache)
+                .unwrap();
+            for step in 0..NEW_TOKENS {
+                let pos = tokens.len() + step;
+                model
+                    .forward_with_cache(&[1u32], &[pos], &mut cache)
+                    .unwrap();
             }
-
-            group.throughput(Throughput::Elements((seq_len * hidden_dim) as u64));
-            group.bench_with_input(
-                BenchmarkId::new("f32", format!("h{}_l{}", hidden_dim, num_layers)),
-                &(&layers, hidden_dim, num_layers, num_heads, head_dim, ffn_dim, seq_len),
-                |b, (layers, hidden_dim, num_layers, num_heads, head_dim, _ffn_dim, seq_len)| {
-                    b.iter(|| {
-                        let mut x = Tensor::randn(0.0f32, 1.0f32, (1, *seq_len, *hidden_dim), &device).unwrap();
-                        for layer in layers.iter().take(*num_layers) {
-                            let (qkv_weight, out_weight, ln1_w, ln1_b, ln2_w, ln2_b, fc_in, fc_out) = layer;
-
-                            // LN1
-                            let last_dim = x.dims().len() - 1;
-                            let mean = x.mean(last_dim).unwrap();
-                            let mean = mean.unsqueeze(last_dim).unwrap();
-                            let x_centered = x.broadcast_sub(&mean).unwrap();
-                            let variance = x_centered.sqr().unwrap().mean(last_dim).unwrap();
-                            let std = (variance + 1e-5).unwrap().sqrt().unwrap();
-                            let std = std.unsqueeze(last_dim).unwrap();
-                            let mut normed = x_centered.broadcast_div(&std).unwrap();
-                            let ln1_w = if ln1_w.dtype() != normed.dtype() {
-                                ln1_w.to_dtype(normed.dtype()).unwrap()
-                            } else {
-                                ln1_w.clone()
-                            };
-                            let ln1_b = if ln1_b.dtype() != normed.dtype() {
-                                ln1_b.to_dtype(normed.dtype()).unwrap()
-                            } else {
-                                ln1_b.clone()
-                            };
-                            normed = normed.broadcast_mul(&ln1_w).unwrap().broadcast_add(&ln1_b).unwrap();
-
-                            // Attention
-                            let qkv = normed.broadcast_matmul(&qkv_weight.unsqueeze(0).unwrap()).unwrap();
-                            let qkv = qkv.reshape((1, *seq_len, 3, *num_heads, *head_dim)).unwrap();
-                            let qkv = qkv.permute((0, 3, 2, 1, 4)).unwrap();
-                            let q = qkv.get_on_dim(2, 0).unwrap();
-                            let v = qkv.get_on_dim(2, 1).unwrap();
-                            let k = qkv.get_on_dim(2, 2).unwrap();
-
-                            let scale = 1.0 / (*head_dim as f64).sqrt();
-                            let scores = q.broadcast_matmul(&k.transpose(2, 3).unwrap()).unwrap();
-                            let scores = (scores * scale).unwrap();
-
-                            let weights = candle_nn::ops::softmax(&scores, 3).unwrap();
-                            let context = weights.broadcast_matmul(&v).unwrap();
-                            let context = context.permute((0, 2, 1, 3)).unwrap()
-                                .reshape((1, *seq_len, *hidden_dim)).unwrap();
-                            let attn_out = context.broadcast_matmul(&out_weight.unsqueeze(0).unwrap()).unwrap();
-
-                            // Residual 1
-                            x = x.broadcast_add(&attn_out).unwrap();
-
-                            // LN2
-                            let last_dim = x.dims().len() - 1;
-                            let mean = x.mean(last_dim).unwrap();
-                            let mean = mean.unsqueeze(last_dim).unwrap();
-                            let x_centered = x.broadcast_sub(&mean).unwrap();
-                            let variance = x_centered.sqr().unwrap().mean(last_dim).unwrap();
-                            let std = (variance + 1e-5).unwrap().sqrt().unwrap();
-                            let std = std.unsqueeze(last_dim).unwrap();
-                            let mut normed = x_centered.broadcast_div(&std).unwrap();
-                            let ln2_w = if ln2_w.dtype() != normed.dtype() {
-                                ln2_w.to_dtype(normed.dtype()).unwrap()
-                            } else {
-                                ln2_w.clone()
-                            };
-                            let ln2_b = if ln2_b.dtype() != normed.dtype() {
-                                ln2_b.to_dtype(normed.dtype()).unwrap()
-                            } else {
-                                ln2_b.clone()
-                            };
-                            normed = normed.broadcast_mul(&ln2_w).unwrap().broadcast_add(&ln2_b).unwrap();
-
-                            // FFN
-                            let hidden = normed.broadcast_matmul(&fc_in.unsqueeze(0).unwrap()).unwrap();
-                            let activated = hidden.gelu().unwrap();
-                            let ffn_out = activated.broadcast_matmul(&fc_out.unsqueeze(0).unwrap()).unwrap();
-
-                            // Residual 2
-                            x = x.broadcast_add(&ffn_out).unwrap();
-                        }
-                    });
-                },
-            );
-        }
-    }
+        });
+    });
     group.finish();
 }
 
-pub fn bench_e2e_inference(c: &mut Criterion) {
+/// Weight loading, on the committed parity fixture.
+fn bench_weight_load(c: &mut Criterion) {
     let device = Device::Cpu;
-    let mut group = c.benchmark_group("e2e_inference");
-
-    for &hidden_dim in &[256, 512, 1024] {
-        let seq_len = 7; // prefill
-        let num_heads = 8;
-        let head_dim = hidden_dim / num_heads;
-        let ffn_dim = hidden_dim * 4;
-        let num_layers = 6;
-
-        let mut layers = Vec::new();
-        for _ in 0..num_layers {
-            let qkv_weight =
-                Tensor::randn(0.0f32, 0.02f32, (hidden_dim, hidden_dim * 3), &device).unwrap();
-            let out_weight =
-                Tensor::randn(0.0f32, 0.02f32, (hidden_dim, hidden_dim), &device).unwrap();
-            let ln1_w = Tensor::ones(hidden_dim, DType::F32, &device).unwrap();
-            let ln1_b = Tensor::zeros(hidden_dim, DType::F32, &device).unwrap();
-            let ln2_w = Tensor::ones(hidden_dim, DType::F32, &device).unwrap();
-            let ln2_b = Tensor::zeros(hidden_dim, DType::F32, &device).unwrap();
-            let fc_in = Tensor::randn(0.0f32, 0.02f32, (hidden_dim, ffn_dim), &device).unwrap();
-            let fc_out = Tensor::randn(0.0f32, 0.02f32, (ffn_dim, hidden_dim), &device).unwrap();
-            layers.push((
-                qkv_weight, out_weight, ln1_w, ln1_b, ln2_w, ln2_b, fc_in, fc_out,
-            ));
-        }
-
-        group.throughput(Throughput::Elements(1));
-        group.bench_with_input(
-            BenchmarkId::new("prefill", format!("h{}", hidden_dim)),
-            &(&layers, hidden_dim, num_layers, num_heads, head_dim, ffn_dim, seq_len),
-            |b, (layers, hidden_dim, num_layers, num_heads, head_dim, _ffn_dim, seq_len)| {
-                b.iter(|| {
-                    let mut x = Tensor::randn(0.0f32, 1.0f32, (1, *seq_len, *hidden_dim), &device).unwrap();
-                    for layer in layers.iter().take(*num_layers) {
-                        let (qkv_weight, out_weight, ln1_w, ln1_b, _ln2_w, _ln2_b, _fc_in, _fc_out) = layer;
-
-                        let last_dim = x.dims().len() - 1;
-                        let mean = x.mean(last_dim).unwrap();
-                        let mean = mean.unsqueeze(last_dim).unwrap();
-                        let x_centered = x.broadcast_sub(&mean).unwrap();
-                        let variance = x_centered.sqr().unwrap().mean(last_dim).unwrap();
-                        let std = (variance + 1e-5).unwrap().sqrt().unwrap();
-                        let std = std.unsqueeze(last_dim).unwrap();
-                        let mut normed = x_centered.broadcast_div(&std).unwrap();
-                        let ln1_w = if ln1_w.dtype() != normed.dtype() {
-                            ln1_w.to_dtype(normed.dtype()).unwrap()
-                        } else {
-                            ln1_w.clone()
-                        };
-                        let ln1_b = if ln1_b.dtype() != normed.dtype() {
-                            ln1_b.to_dtype(normed.dtype()).unwrap()
-                        } else {
-                            ln1_b.clone()
-                        };
-                        normed = normed.broadcast_mul(&ln1_w).unwrap().broadcast_add(&ln1_b).unwrap();
-
-                        let qkv = normed.broadcast_matmul(&qkv_weight.unsqueeze(0).unwrap()).unwrap();
-                        let qkv = qkv.reshape((1, *seq_len, 3, *num_heads, *head_dim)).unwrap();
-                        let qkv = qkv.permute((0, 3, 2, 1, 4)).unwrap();
-                        let q = qkv.get_on_dim(2, 0).unwrap();
-                        let v = qkv.get_on_dim(2, 1).unwrap();
-                        let k = qkv.get_on_dim(2, 2).unwrap();
-
-                        let scale = 1.0 / (*head_dim as f64).sqrt();
-                        let scores = q.broadcast_matmul(&k.transpose(2, 3).unwrap()).unwrap();
-                        let scores = (scores * scale).unwrap();
-
-                        let weights = candle_nn::ops::softmax(&scores, 3).unwrap();
-                        let context = weights.broadcast_matmul(&v).unwrap();
-                        let context = context.permute((0, 2, 1, 3)).unwrap()
-                            .reshape((1, *seq_len, *hidden_dim)).unwrap();
-                        let attn_out = context.broadcast_matmul(&out_weight.unsqueeze(0).unwrap()).unwrap();
-
-                        x = x.broadcast_add(&attn_out).unwrap();
-                    }
-                });
-            },
-        );
+    let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
+    let path = dir.join("tiny_h4.pth");
+    if !path.exists() {
+        eprintln!("skipping weight_load bench: {} not found", path.display());
+        return;
     }
-    group.finish();
-}
-
-pub fn bench_f16_vs_f32(c: &mut Criterion) {
-    let device = Device::Cpu;
-    let mut group = c.benchmark_group("f16_vs_f32");
-
-    for &dtype in &[DType::F32, DType::F16] {
-        let hidden_dim = 1024;
-        let num_heads = 16;
-        let head_dim = hidden_dim / num_heads;
-        let seq_len = 64;
-
-        let qkv_weight = Tensor::randn(0.0f32, 0.02f32, (hidden_dim, hidden_dim * 3), &device)
-            .unwrap()
-            .to_dtype(dtype)
+    let config_json: serde_json::Value =
+        serde_json::from_str(&std::fs::read_to_string(dir.join("tiny_h4_config.json")).unwrap())
             .unwrap();
-        let out_weight = Tensor::randn(0.0f32, 0.02f32, (hidden_dim, hidden_dim), &device)
-            .unwrap()
-            .to_dtype(dtype)
-            .unwrap();
-
-        group.throughput(Throughput::Elements(
-            (seq_len * hidden_dim * num_heads) as u64,
-        ));
-        group.bench_with_input(
-            BenchmarkId::new("attention", format!("{:?}", dtype)),
-            &(
-                &qkv_weight,
-                &out_weight,
-                hidden_dim,
-                num_heads,
-                head_dim,
-                seq_len,
-            ),
-            |b, (qkv_weight, out_weight, hidden_dim, num_heads, head_dim, seq_len)| {
-                b.iter(|| {
-                    let x = Tensor::randn(0.0f32, 1.0f32, (1, *seq_len, *hidden_dim), &device)
-                        .unwrap()
-                        .to_dtype(dtype)
-                        .unwrap();
-                    let qkv = x
-                        .broadcast_matmul(&qkv_weight.unsqueeze(0).unwrap())
-                        .unwrap();
-                    let qkv = qkv
-                        .reshape((1, *seq_len, 3, *num_heads, *head_dim))
-                        .unwrap();
-                    let qkv = qkv.permute((0, 3, 2, 1, 4)).unwrap();
-                    let q = qkv.get_on_dim(2, 0).unwrap();
-                    let v = qkv.get_on_dim(2, 1).unwrap();
-                    let k = qkv.get_on_dim(2, 2).unwrap();
+    let config = CodeGenConfig::from_hf_config(&config_json);
 
-                    let scale = 1.0 / (*head_dim as f64).sqrt();
-                    let scores = q.broadcast_matmul(&k.transpose(2, 3).unwrap()).unwrap();
-                    let scores = (scores * scale).unwrap();
+    c.bench_function("weight_load/tiny_h4", |b| {
+        b.iter(|| WeightLoader::load_from_pytorch(&path, &config, &device).unwrap());
+    });
+}
 
-                    let weights = candle_nn::ops::softmax(&scores, 3).unwrap();
-                    let context = weights.broadcast_matmul(&v).unwrap();
-                    let context = context
-                        .permute((0, 2, 1, 3))
-                        .unwrap()
-                        .reshape((1, *seq_len, *hidden_dim))
-                        .unwrap();
-                    context.broadcast_matmul(&out_weight.unsqueeze(0).unwrap())
-                });
-            },
-        );
+/// F32 against F16 on the same shapes, to keep the README's precision claim honest.
+fn bench_dtype(c: &mut Criterion) {
+    let device = Device::Cpu;
+    let mut group = c.benchmark_group("prefill_dtype");
+
+    for (name, dtype) in [("f32", DType::F32), ("f16", DType::F16)] {
+        let config = CodeGenConfig {
+            dtype,
+            ..bench_config()
+        };
+        let model = CodeGenModel::new_blank(config.clone(), &device).unwrap();
+        let tokens = prompt_tokens(&config);
+        let positions: Vec<usize> = (0..tokens.len()).collect();
+
+        group.bench_function(name, |b| {
+            b.iter(|| {
+                let mut cache: Option<Vec<KVCache>> = None;
+                model
+                    .forward_with_cache(&tokens, &positions, &mut cache)
+                    .unwrap()
+            });
+        });
     }
     group.finish();
 }
 
 criterion_group!(
     benches,
-    bench_attention,
-    bench_ffn,
-    bench_layernorm,
-    bench_full_block,
-    bench_e2e_inference,
-    bench_f16_vs_f32,
-    bench_rope_2d,
-    bench_kv_cache,
-    bench_quantized_matmul
+    bench_prefill,
+    bench_generate,
+    bench_generator,
+    bench_weight_load,
+    bench_dtype
 );
 criterion_main!(benches);
diff --git a/configs/train.yaml b/configs/train.yaml
index a31be32..26a33fb 100644
--- a/configs/train.yaml
+++ b/configs/train.yaml
@@ -1,5 +1,5 @@
 # GLM Training Configuration
-# Run with: cargo run --release -- --glm-train --config configs/train.yaml
+# Run with: cargo run --release -- glm-train --config configs/train.yaml
 
 # Model configuration (matches GLMConfig)
 model:
@@ -38,8 +38,6 @@ training:
     max_steps: 10000
     min_lr_ratio: 0.1
 
-  # BATCH  # minimum LR as ratio of base LR
-
   # Batch/accumulation
   micro_batch_size: 1    # sequences per forward pass (keep 1 for CPU)
   gradient_accumulation_steps: 32  # simulate batch=32
@@ -50,6 +48,7 @@ training:
   eval_every: 500
   save_every: 1000
   log_every: 10
+  eval_steps: 100
 
   # Checkpointing
   checkpoint_dir: "glm_checkpoint"
@@ -58,11 +57,7 @@ training:
 
   # Precision
   dtype: "f32"  # "f32" or "f16"
-
-# Evaluation
-evaluation:
-  eval_steps: 100
-  blank_infill_samples: 10
+  tokenizer_path: "codegen_weights/tokenizer.json"
 
 # Logging
 logging:
diff --git a/docs/architecture-deep-dive.md b/docs/architecture-deep-dive.md
index 443af47..8a07a72 100644
--- a/docs/architecture-deep-dive.md
+++ b/docs/architecture-deep-dive.md
@@ -55,16 +55,18 @@ Key details:
 - **Causal mask**: Upper-triangular mask with -inf in the upper triangle
 - **No bias**: Following GPT-2/CodeGen convention (LayerNorm provides bias)
 
+**QKV layout**: CodeGen does *not* store the fused projection as
+`[all q | all v | all k]`. It keeps the sharded layout of the original TPU
+implementation — `mp_num = 4` consecutive groups, each holding `[q | v | k]` for
+its slice of the heads. Splitting it any other way gives every head the wrong
+slice of the projection, which is silent: shapes still line up, output is
+garbage. See `split_qkv` in `src/codegen/model.rs`.
+
 ```rust
-// Simplified forward
-fn forward(&self, x: &Tensor) -> Result<Tensor> {
-    let qkv = x.matmul(&self.qkv_weight)?;
-    let q = qkv.narrow(1, 0, hidden_dim);
-    let k = qkv.narrow(1, hidden_dim, hidden_dim);
-    let v = qkv.narrow(1, 2 * hidden_dim, hidden_dim);
-    // ... split heads, compute attention, concat heads
-    output.matmul(&self.out_weight)
-}
+// [bs, sl, MP_NUM, 3, local_dim] — dim 3 selects q, v, k in that order
+let grouped = qkv.reshape((bs, sl, MP_NUM, 3, local_dim))?;
+let q = grouped.get_on_dim(3, 0)?.reshape((bs, sl, num_heads, head_dim))?;
+// flattening group and per-group-head axes yields head `group * (heads / MP_NUM) + i`
 ```
 
 **Causal mask optimization**: Instead of creating a full `[seq, seq]` mask each
@@ -245,15 +247,14 @@ This reduces the per-token computation from O(seq²) to O(seq).
 Weights are loaded from HuggingFace PyTorch checkpoints:
 
 ```rust
-// Loads from .bin (pickle) or .safetensors
-let pth = PthTensors::new(path)?;
-// Map PyTorch names → model parameters
-model.embedding = load_tensor("transformer.wte.weight");
-model.blocks[i].attn = load_tensor("transformer.h.i.attn.qkv_proj.weight");
+// WeightLoader::load picks the format from the extension
+WeightLoader::load(path, &config, &device)?;      // .safetensors or .bin
 ```
 
-The `WeightLoader` handles name mapping and dtype conversion, with zero-init
-to avoid allocating 350M random floats.
+`WeightLoader` handles name mapping and dtype conversion, with zero-init to avoid
+allocating 350M random floats. A tensor named in the model but absent from the
+checkpoint is an error, not a silently zeroed layer. `ModelContext` prefers
+`model.safetensors` over `pytorch_model.bin` when both are present.
 
 ### 3.5 Sampling Pipeline
 
@@ -334,7 +335,7 @@ Auto-cleanup keeps only the N most recent checkpoints.
 
 All weights and activations can use F16 (half precision):
 - **Memory**: ~50% reduction (700MB → 350MB for CodeGen)
-- **Speed**: ~23% faster on i5-6600 with `gemm` F16 support
+- **Speed**: 3.4x faster than F32 on CodeGen-350M (20.3 ms/token against 68.8 ms, Apple M1 Pro)
 - **Quality**: Negligible degradation for inference
 
 ```rust
diff --git a/docs/comparison-frameworks.md b/docs/comparison-frameworks.md
index 32a945b..e40d88c 100644
--- a/docs/comparison-frameworks.md
+++ b/docs/comparison-frameworks.md
@@ -15,7 +15,7 @@ Comparing this project against other Rust ML frameworks for transformer inferenc
 | **Build Time** | ~2 min | ~3 min | ~5 min | ~10 min |
 | **Model Support** | GLM + CodeGen | Llama, GPT2, etc. | Custom | Any PyTorch model |
 | **Training** | ✅ GLM only | ❌ No | ✅ Yes | ✅ Yes |
-| **Quantization** | ✅ INT8 (manual) | ✅ GGML/GGUF | ✅ INT8/FP4 | ❌ No |
+| **Quantization** | ❌ FP16 only | ✅ GGML/GGUF | ✅ INT8/FP4 | ❌ No |
 
 ---
 
@@ -30,12 +30,12 @@ Comparing this project against other Rust ML frameworks for transformer inferenc
 - Educational value — everything is hand-written for clarity
 - Real CodeGen-350M inference on CPU
 - Working training pipeline with safetensors checkpoints
-- FP16 support (23% speedup on i5-6600)
+- FP16 support (3.4x speedup on CodeGen-350M, Apple M1 Pro)
 
 **Weaknesses**:
 - CPU-only (no CUDA/Metal backend)
 - Limited model zoo (GLM + CodeGen only)
-- No quantization via GGML (manual INT8 only)
+- No quantization (FP32 and FP16 only)
 - No distributed training
 
 **Best for**: Learning, CPU-only deployment, edge devices, custom architectures
@@ -89,16 +89,17 @@ Comparing this project against other Rust ML frameworks for transformer inferenc
 
 ## 3. Performance Benchmarks
 
-*Measured on i5-6600 (4C/4T, 7.5GB RAM), CPU-only, release mode*
+*This project measured on Apple M1 Pro, CPU-only, release mode, greedy decoding.
+The comparison rows are not re-measured here and are indicative only.*
 
 ### CodeGen-350M Inference (FP32)
 
-| Framework | Load Time | Prefill (7 tok) | Per Token | 50 Tokens |
-|:----------|:---------:|:----------------:|:---------:|:---------:|
-| **This project** | 0.5s | 0.30s | 0.15s | 7.8s |
-| candle (llama example) | 0.8s* | 0.25s | 0.12s | 6.5s |
-| burn | N/A | N/A | N/A | N/A (no CodeGen port) |
-| tch-rs | 0.4s | 0.20s | 0.10s | 5.2s |
+| Framework | Load Time | Per Token | 64 Tokens |
+|:----------|:---------:|:---------:|:---------:|
+| **This project** | ~0.6s | 68.8 ms | 4.4s |
+| candle (llama example) | 0.8s* | ~120 ms | ~6.5s |
+| burn | N/A | N/A | N/A (no CodeGen port) |
+| tch-rs | 0.4s | ~100 ms | ~5.2s |
 
 *\* candle-examples uses GGML quantization by default, not FP32*
 
@@ -106,14 +107,14 @@ Comparing this project against other Rust ML frameworks for transformer inferenc
 
 | Framework | Per Token | Speedup vs FP32 |
 |:----------|:---------:|:----------------:|
-| **This project** | 0.12s | 1.23× |
-| tch-rs | 0.08s | 1.25× |
+| **This project** | 20.3 ms | 3.4x |
+| tch-rs | ~80 ms | ~1.25x |
 
 ### Memory Usage (CodeGen-350M)
 
 | Framework | FP32 Loading | FP32 Runtime | FP16 Runtime |
 |:----------|:-----------:|:-----------:|:-----------:|
-| **This project** | 700 MB | ~1.0 GB | ~600 MB |
+| **This project** | 760 MB | 1.63 GB | 1.08 GB |
 | tch-rs | 700 MB | ~1.2 GB | ~700 MB |
 
 ---
diff --git a/docs/run-codegen-on-raspberry-pi.md b/docs/run-codegen-on-raspberry-pi.md
index 136862e..5c1e2a4 100644
--- a/docs/run-codegen-on-raspberry-pi.md
+++ b/docs/run-codegen-on-raspberry-pi.md
@@ -170,4 +170,4 @@ sudo systemctl start codegen
 - Generation speed is ~0.3-0.8 tokens/second (usable for short completions)
 - Not suitable for real-time interactive use
 - Large context windows (>512 tokens) may cause OOM
-- INT8 quantization can reduce memory by ~4x with minimal quality loss
+- `--f16` halves weight memory and is wired through every layer
diff --git a/docs/train-code-infill.md b/docs/train-code-infill.md
index ae6f0dc..00e20fc 100644
--- a/docs/train-code-infill.md
+++ b/docs/train-code-infill.md
@@ -47,21 +47,28 @@ model:
   vocab_size: 16384
 
 training:
-  batch_size: 8
   learning_rate: 0.001
   max_steps: 5000
-  warmup_steps: 200
-  lr_schedule: cosine
-  eval_interval: 100
-  save_interval: 500
-  keep_last_n: 3
   max_seq_len: 128
-  gradient_accumulation_steps: 4
+  micro_batch_size: 1              # sequences per forward pass
+  gradient_accumulation_steps: 4   # forward passes per optimizer step
   max_grad_norm: 1.0
+  eval_every: 100
+  save_every: 500
+  log_every: 5
+  keep_last_n_checkpoints: 3
+  checkpoint_dir: glm_checkpoint
   data_dir: training_data
   download_if_empty: true
+  lr_schedule:
+    type: cosine
+    warmup_steps: 200
+    max_steps: 5000
+    min_lr_ratio: 0.1
 ```
 
+Every field is optional — anything you leave out falls back to its default.
+
 ### Key Parameters
 
 | Parameter | Value | What it controls |
@@ -69,7 +76,7 @@ training:
 | `hidden_dim` | 256 | Width of the transformer (small = fast) |
 | `num_layers` | 6 | Depth of the transformer |
 | `max_steps` | 5000 | Total training steps (~14M param model) |
-| `gradient_accumulation_steps` | 4 | Simulates larger batch size on limited memory |
+| `gradient_accumulation_steps` | 4 | Forward passes averaged into one optimizer step |
 | `warmup_steps` | 200 | Gradually increases LR to avoid instability |
 
 ---
@@ -79,7 +86,7 @@ training:
 ### Start training with YAML config:
 
 ```bash
-cargo run --release -- glm-train --data-path training_data --steps 5000
+cargo run --release -- glm-train --config configs/train.yaml --data-path training_data --steps 5000
 ```
 
 ### Or with defaults (auto-downloads training data):
@@ -91,23 +98,26 @@ cargo run --release -- glm-train
 ### What you'll see:
 
 ```
-Step    100 | loss=4.21 | lr=0.00050 | 0.12s/step
-Step    200 | loss=3.85 | lr=0.00100 | 0.11s/step
-Step    300 | loss=3.52 | lr=0.00087 | 0.11s/step
-...
-Step   1000 | loss=2.14 | lr=0.00012 | 0.11s/step
-Step   2000 | loss=1.43 | lr=0.00004 | 0.11s/step
-...
-Eval  | train_loss=1.02 | val_loss=1.31
+Step 5: loss = 10.8024 (avg = 10.8313), lr = 5.00e-4
+Step 10: loss = 10.2722 (avg = 10.7133), lr = 1.00e-3
+Step 20: loss = 9.2113 (avg = 10.2314), lr = 9.73e-4
+Step 40: loss = 5.5989 (avg = 8.6961), lr = 7.75e-4
+Step 60: loss = 4.3074 (avg = 7.5089), lr = 4.72e-4
+Step 80: loss = 4.9949 (avg = 6.8133), lr = 2.05e-4
+  Checkpoint saved to "glm_checkpoint/model_step_000080.safetensors" (step 80)
 ```
 
-Loss should drop steadily from ~4.5 to ~1.0 over 5000 steps.
+That is a real 80-step run on a 2-layer, `hidden_dim` 128 model over a single source
+file — small enough to overfit quickly, which is what you want when checking that the
+setup works at all. The per-step loss is noisy because each step sees a fresh random
+masking of one sequence; watch the running average instead. On a real corpus with the
+config above, expect a slower but steadier decline.
 
 ### Checkpoints
 
-Checkpoints are saved to `checkpoints/`:
+Checkpoints are saved to `checkpoint_dir` (default `glm_checkpoint/`):
 ```
-checkpoints/
+glm_checkpoint/
 ├── model_step_0000500.safetensors
 ├── optimizer_step_0000500.json
 ├── model_step_0001000.safetensors
@@ -170,6 +180,6 @@ This loads the latest checkpoint and generates code completions.
 |---------|-------|-----|
 | `No data files found` | Empty data directory | Use `--data-path` with a valid directory or populate `training_data/` |
 | Loss not decreasing | Learning rate too high/low | Adjust `learning_rate` (try 3e-4) |
-| Out of memory | Batch/sequence too large | Reduce `batch_size` or `max_seq_len` |
-| NaN loss | Gradient explosion | Reduce `learning_rate` or increase `warmup_steps` |
+| Out of memory | Batch/sequence too large | Reduce `micro_batch_size` or `max_seq_len` |
+| NaN loss | Gradient explosion | Reduce `learning_rate`, lower `max_grad_norm`, or increase `warmup_steps` |
 | Slow training | Debug build | Always use `--release` for training |
diff --git a/examples/convert_codegen_to_safetensors.rs b/examples/convert_codegen_to_safetensors.rs
index fde8a1f..89cc98f 100644
--- a/examples/convert_codegen_to_safetensors.rs
+++ b/examples/convert_codegen_to_safetensors.rs
@@ -9,8 +9,10 @@
 //!       /path/to/codegen_weights/pytorch_model.bin \
 //!       /path/to/codegen_weights/model.safetensors
 //!
-//! Then load in inference:
-//!   cargo run --release -- --codegen --weights /path/to/codegen_weights/model.safetensors
+//! The result is picked up automatically: `ModelContext` prefers
+//! `model.safetensors` over `pytorch_model.bin` in the weights directory.
+//!
+//!   cargo run --release -- --weights-dir /path/to/codegen_weights complete "def f():"
 
 use std::fs;
 use std::path::PathBuf;
diff --git a/scripts/gen_parity_fixture.py b/scripts/gen_parity_fixture.py
new file mode 100644
index 0000000..ee99ea1
--- /dev/null
+++ b/scripts/gen_parity_fixture.py
@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Generate tiny CodeGen fixtures for the Rust numerical parity test.
+
+Builds small randomly-initialised CodeGen models with HuggingFace transformers,
+saves each state dict in torch pickle format (what `candle_core::pickle::PthTensors`
+reads), and records the reference logits for a fixed input sequence.
+
+The fixtures are committed so CI needs no Python. Regenerate with:
+
+    python3 scripts/gen_parity_fixture.py
+"""
+
+import json
+import pathlib
+
+import torch
+from transformers import CodeGenConfig, CodeGenForCausalLM
+
+OUT = pathlib.Path(__file__).resolve().parent.parent / "tests" / "fixtures"
+
+# Two shapes on purpose:
+#   n_head == mp_num  -> exactly one head per model-parallel group
+#   n_head == 2*mp_num -> several heads per group, so a "one head per group"
+#                         shortcut cannot pass the test
+VARIANTS = [
+    # name,        n_embd, n_layer, n_head, rotary_dim
+    ("tiny_h4", 128, 2, 4, 16),   # head_dim 32, partial rotary
+    ("tiny_h8", 128, 2, 8, 16),   # head_dim 16, full rotary
+]
+
+TOKENS = [3, 17, 42, 8, 255, 1, 99, 128]
+
+
+def build(name, n_embd, n_layer, n_head, rotary_dim):
+    torch.manual_seed(0)
+    config = CodeGenConfig(
+        vocab_size=256,
+        n_positions=64,
+        n_ctx=64,
+        n_embd=n_embd,
+        n_layer=n_layer,
+        n_head=n_head,
+        rotary_dim=rotary_dim,
+        n_inner=None,
+        activation_function="gelu_new",
+        resid_pdrop=0.0,
+        embd_pdrop=0.0,
+        attn_pdrop=0.0,
+        layer_norm_epsilon=1e-5,
+        tie_word_embeddings=False,
+    )
+    model = CodeGenForCausalLM(config)
+    model.eval()
+
+    # LayerNorm weights default to exactly 1.0 / bias 0.0, which a buggy loader
+    # that silently drops tensors could accidentally match. Perturb them so a
+    # missed norm tensor shows up as a logit mismatch.
+    torch.manual_seed(1)
+    with torch.no_grad():
+        for module in model.modules():
+            if isinstance(module, torch.nn.LayerNorm):
+                module.weight.add_(torch.randn_like(module.weight) * 0.05)
+                module.bias.add_(torch.randn_like(module.bias) * 0.05)
+
+    input_ids = torch.tensor([TOKENS], dtype=torch.long)
+    with torch.no_grad():
+        logits = model(input_ids).logits
+
+    torch.save(model.state_dict(), OUT / f"{name}.pth")
+    (OUT / f"{name}_config.json").write_text(
+        json.dumps(config.to_dict(), indent=2, sort_keys=True, default=str) + "\n"
+    )
+    (OUT / f"{name}_logits.json").write_text(
+        json.dumps(
+            {
+                "tokens": TOKENS,
+                "shape": list(logits.shape),
+                "logits": logits.reshape(-1).tolist(),
+            }
+        )
+        + "\n"
+    )
+    print(f"{name}: logits {tuple(logits.shape)}  head_dim {n_embd // n_head}")
+
+
+if __name__ == "__main__":
+    OUT.mkdir(parents=True, exist_ok=True)
+    for variant in VARIANTS:
+        build(*variant)
diff --git a/scripts/update-readme-benchmarks.sh b/scripts/update-readme-benchmarks.sh
index 8125e6e..9882fa9 100755
--- a/scripts/update-readme-benchmarks.sh
+++ b/scripts/update-readme-benchmarks.sh
@@ -37,10 +37,6 @@ done < /tmp/bench_output.txt
 
 echo "Found ${#results[@]} benchmark results"
 
-# Generate benchmark table rows
-TABLE_ROWS=""
-TABLE_ROWS+="| `$name` | $value $unit |\n"
-
 # Now update README.md
 README="$PROJECT_DIR/README.md"
 
@@ -52,13 +48,13 @@ for key in "${!results[@]}"; do
     # Try to map benchmark names to user-friendly labels
     label=""
     case "$key" in
-        *"attention"*) label="Attention forward" ;;
-        *"ffn"*) label="FFN forward" ;;
-        *"layernorm"*|*"layer_norm"*) label="LayerNorm forward" ;;
-        *"block"*) label="Transformer block" ;;
-        *"e2e"*) label="E2E inference (prefill + N tokens)" ;;
-        *"f16"*|*"fp16"*) label="FP16 inference" ;;
-        *"f32"*|*"fp32"*) label="FP32 inference" ;;
+        *"hidden_only"*) label="Prefill, no vocabulary projection" ;;
+        *"prefill_dtype/f16"*) label="Prefill (FP16)" ;;
+        *"prefill_dtype/f32"*) label="Prefill (FP32)" ;;
+        *"prefill"*) label="Prefill (32 tokens)" ;;
+        *"generator"*) label="Generate (32 prompt + 16 new)" ;;
+        *"generate"*) label="Forward passes (32 prompt + 16 new)" ;;
+        *"weight_load"*) label="Weight loading" ;;
         *) label="$key" ;;
     esac
     BENCH_TABLE+="| $label | ${results[$key]} |\n"
diff --git a/src/cli.rs b/src/cli.rs
index c61a770..4fc3cdd 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -10,6 +10,7 @@ use std::path::PathBuf;
 ///
 /// - `--f16` — Use FP16 precision (faster, less memory)
 /// - `--weights-dir` — Path to weights directory (default: `codegen_weights`)
+/// - `--seed` — Fixed sampling seed for reproducible output
 #[derive(Parser)]
 #[command(
     name = "codegen",
@@ -25,6 +26,10 @@ pub struct Cli {
     #[arg(long, global = true, default_value = "codegen_weights")]
     pub weights_dir: PathBuf,
 
+    /// Fixed sampling seed for reproducible output (default: random each run)
+    #[arg(long, global = true)]
+    pub seed: Option<u64>,
+
     #[command(subcommand)]
     pub command: Commands,
 }
@@ -53,12 +58,14 @@ pub enum Commands {
         temperature: f64,
 
         /// Prompt template: completion, instruct, chat
-        #[arg(short, long, default_value = "completion")]
+        ///
+        /// No short flag: `-t` already belongs to `--temperature`.
+        #[arg(long, default_value = "completion")]
         template: String,
 
-        /// Stream tokens as they're generated
-        #[arg(long, default_value = "true")]
-        stream: bool,
+        /// Print the completion only when it is finished, instead of streaming it
+        #[arg(long)]
+        no_stream: bool,
     },
 
     /// Interactive REPL (single-turn)
@@ -86,5 +93,9 @@ pub enum Commands {
         /// Number of training steps
         #[arg(short, long, default_value = "1000")]
         steps: usize,
+
+        /// YAML training config (see configs/train.yaml). Defaults are used if omitted.
+        #[arg(short, long)]
+        config: Option<PathBuf>,
     },
 }
diff --git a/src/codegen/kv_cache.rs b/src/codegen/kv_cache.rs
index 4ae7b97..1abc594 100644
--- a/src/codegen/kv_cache.rs
+++ b/src/codegen/kv_cache.rs
@@ -4,8 +4,7 @@ pub struct KVCache {
     k: Tensor,
     v: Tensor,
     pos: usize,
-    n_heads: usize,
-    head_dim: usize,
+    max_seq_len: usize,
 }
 
 impl KVCache {
@@ -22,26 +21,31 @@ impl KVCache {
             k,
             v,
             pos: 0,
-            n_heads,
-            head_dim,
+            max_seq_len,
         })
     }
 
+    /// Write `k_new`/`v_new` at the current position and return views over the
+    /// cache up to and including them.
+    ///
+    /// Uses [`Tensor::slice_set`], which copies only the new tokens into the
+    /// existing buffer. The previous `slice_assign` was not a targeted write: it
+    /// zero-padded the source out to the full buffer shape, built a full-size
+    /// mask, and ran `where_cond` over every element — the whole
+    /// `[1, heads, max_seq_len, head_dim]` cache, twice per layer per token.
     pub fn append(&mut self, k_new: &Tensor, v_new: &Tensor) -> Result<(Tensor, Tensor)> {
         let seq_len = k_new.dim(2)?;
         let end = self.pos + seq_len;
-
-        let k = self.k.slice_assign(
-            &[0..1, 0..self.n_heads, self.pos..end, 0..self.head_dim],
-            k_new,
-        )?;
-        let v = self.v.slice_assign(
-            &[0..1, 0..self.n_heads, self.pos..end, 0..self.head_dim],
-            v_new,
-        )?;
-
-        self.k = k;
-        self.v = v;
+        if end > self.max_seq_len {
+            return Err(candle_core::Error::Msg(format!(
+                "KV cache overflow: {seq_len} more token(s) at position {} exceeds max_seq_len {}",
+                self.pos, self.max_seq_len
+            )));
+        }
+
+        // slice_set requires both sides contiguous; the rotary output may not be.
+        self.k.slice_set(&k_new.contiguous()?, 2, self.pos)?;
+        self.v.slice_set(&v_new.contiguous()?, 2, self.pos)?;
         self.pos = end;
 
         let k_out = self.k.narrow(2, 0, self.pos)?;
@@ -125,6 +129,50 @@ mod tests {
         assert_eq!(cache.position(), 0);
     }
 
+    #[test]
+    fn test_kv_cache_incremental_matches_block() {
+        let device = Device::Cpu;
+        let k = Tensor::randn(0.0f32, 1.0f32, (1, 2, 5, 4), &device).unwrap();
+        let v = Tensor::randn(0.0f32, 1.0f32, (1, 2, 5, 4), &device).unwrap();
+
+        let mut block = KVCache::new(16, 2, 4, DType::F32, &device).unwrap();
+        let (k_block, v_block) = block.append(&k, &v).unwrap();
+
+        let mut incremental = KVCache::new(16, 2, 4, DType::F32, &device).unwrap();
+        let mut last = None;
+        for i in 0..5 {
+            let k_i = k.narrow(2, i, 1).unwrap();
+            let v_i = v.narrow(2, i, 1).unwrap();
+            last = Some(incremental.append(&k_i, &v_i).unwrap());
+        }
+        let (k_inc, v_inc) = last.unwrap();
+
+        assert_eq!(k_inc.dims(), k_block.dims());
+        for (a, b) in [(&k_inc, &k_block), (&v_inc, &v_block)] {
+            let diff = (a - b)
+                .unwrap()
+                .abs()
+                .unwrap()
+                .flatten_all()
+                .unwrap()
+                .max(0)
+                .unwrap()
+                .to_scalar::<f32>()
+                .unwrap();
+            assert!(diff < 1e-6, "incremental append diverged by {diff}");
+        }
+    }
+
+    #[test]
+    fn test_kv_cache_rejects_overflow() {
+        let device = Device::Cpu;
+        let mut cache = KVCache::new(4, 2, 4, DType::F32, &device).unwrap();
+        let k = Tensor::zeros((1, 2, 5, 4), DType::F32, &device).unwrap();
+        let v = Tensor::zeros((1, 2, 5, 4), DType::F32, &device).unwrap();
+        let err = cache.append(&k, &v).unwrap_err().to_string();
+        assert!(err.contains("KV cache overflow"), "unexpected error: {err}");
+    }
+
     #[test]
     fn test_kv_cache_dtype_f16() {
         let device = Device::Cpu;
diff --git a/src/codegen/mod.rs b/src/codegen/mod.rs
index b876022..ed82d06 100644
--- a/src/codegen/mod.rs
+++ b/src/codegen/mod.rs
@@ -9,11 +9,10 @@
 //! - Rotary Position Embedding (RoPE)
 //! - KV cache for efficient autoregressive generation
 //! - Parallel attention + FFN (GPT-J style)
-//! - INT8 quantization support
+//! - FP16 inference
 
 pub mod config;
 pub mod kv_cache;
 pub mod model;
-pub mod quantized;
 pub mod rotary;
 pub mod weights;
diff --git a/src/codegen/model.rs b/src/codegen/model.rs
index 99ec3a2..2a6e7c7 100644
--- a/src/codegen/model.rs
+++ b/src/codegen/model.rs
@@ -9,6 +9,44 @@ use super::config::CodeGenConfig;
 use super::kv_cache::KVCache;
 use super::rotary::RotaryEmbedding;
 
+/// Number of model-parallel shards baked into CodeGen's fused `qkv_proj` weight.
+///
+/// The projection is not stored as `[all q | all v | all k]`. It keeps the
+/// sharded layout of the original TPU implementation: `MP_NUM` consecutive
+/// groups, each holding `[q | v | k]` for its slice of the heads. Mirrors
+/// `mp_num = 4` in HuggingFace `models/codegen/modeling_codegen.py`.
+const MP_NUM: usize = 4;
+
+/// Split CodeGen's fused QKV projection into `[batch, heads, seq, head_dim]`
+/// tensors, undoing the model-parallel group interleaving described on
+/// [`MP_NUM`].
+fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize) -> Result<(Tensor, Tensor, Tensor)> {
+    if num_heads % MP_NUM != 0 {
+        return Err(candle_core::Error::Msg(format!(
+            "num_heads ({num_heads}) must be divisible by the CodeGen model-parallel group count ({MP_NUM})"
+        )));
+    }
+    let (bs, sl, _) = qkv.dims3()?;
+    let local_dim = num_heads * head_dim / MP_NUM;
+    // [bs, sl, MP_NUM, 3, local_dim] — dim 3 selects q, v, k in that order.
+    let grouped = qkv.reshape((bs, sl, MP_NUM, 3, local_dim))?;
+
+    let mut split = Vec::with_capacity(3);
+    for slot in 0..3 {
+        // Flattening the group and per-group-head axes yields head index
+        // `group * (num_heads / MP_NUM) + i`, matching `_split_heads` upstream.
+        let part = grouped
+            .get_on_dim(3, slot)?
+            .contiguous()?
+            .reshape((bs, sl, num_heads, head_dim))?;
+        split.push(part.permute((0, 2, 1, 3))?.contiguous()?);
+    }
+    let k = split.pop().unwrap();
+    let v = split.pop().unwrap();
+    let q = split.pop().unwrap();
+    Ok((q, k, v))
+}
+
 pub struct CodeGenModel {
     pub embedding: Embedding,
     pub blocks: Vec<CodeGenBlock>,
@@ -19,10 +57,14 @@ pub struct CodeGenModel {
     config: CodeGenConfig,
 }
 
+/// A CodeGen transformer block.
+///
+/// Attention and FFN are *parallel* (GPT-J style): both read the same
+/// `norm1` output and their results are summed into one residual, so unlike a
+/// serial pre-norm block there is only one LayerNorm per block.
 pub struct CodeGenBlock {
     pub norm1: LayerNorm,
     pub attn: MultiHeadAttention,
-    pub norm2: LayerNorm,
     pub ffn: FeedForward,
 }
 
@@ -30,23 +72,23 @@ impl CodeGenBlock {
     pub fn new_blank(config: &CodeGenConfig, device: &Device) -> Result<Self> {
         let dtype = config.dtype;
         let norm1 = LayerNorm::zeros_with_dtype(config.hidden_dim, config.eps, dtype, device)?;
-        let norm2 = LayerNorm::zeros_with_dtype(config.hidden_dim, config.eps, dtype, device)?;
-        let attn = MultiHeadAttention::new_blank(config.hidden_dim, config.num_heads, device)?;
-        let ffn =
-            FeedForward::new_blank(Activation::GELU, config.hidden_dim, config.ffn_dim, device)?;
-        Ok(Self {
-            norm1,
-            norm2,
-            attn,
-            ffn,
-        })
+        let attn =
+            MultiHeadAttention::new_blank(config.hidden_dim, config.num_heads, dtype, device)?;
+        let ffn = FeedForward::new_blank(
+            Activation::GELU,
+            config.hidden_dim,
+            config.ffn_dim,
+            dtype,
+            device,
+        )?;
+        Ok(Self { norm1, attn, ffn })
     }
 }
 
 impl CodeGenModel {
     pub fn new_blank(config: CodeGenConfig, device: &Device) -> Result<Self> {
         let dtype = config.dtype;
-        let embedding = Embedding::zeros(config.vocab_size, config.hidden_dim, device)?;
+        let embedding = Embedding::zeros(config.vocab_size, config.hidden_dim, dtype, device)?;
         let rotary = RotaryEmbedding::new(config.rotary_dim, config.max_seq_len, dtype, device)?;
 
         let mut blocks = Vec::new();
@@ -69,11 +111,28 @@ impl CodeGenModel {
         })
     }
 
+    /// Full forward pass, logits for every input position.
     pub fn forward_with_cache(
         &self,
         token_ids: &[u32],
         positions: &[usize],
         cache: &mut Option<Vec<KVCache>>,
+    ) -> Result<Tensor> {
+        let hidden = self.forward_hidden(token_ids, positions, cache)?;
+        self.project_logits(&hidden)
+    }
+
+    /// Everything up to and including the final norm, without the vocabulary
+    /// projection.
+    ///
+    /// Generation only ever reads the last position, so it calls this, narrows,
+    /// and projects one row — the `lm_head` matmul is `hidden_dim × vocab_size`
+    /// per position and would otherwise be paid for the whole prompt at prefill.
+    pub fn forward_hidden(
+        &self,
+        token_ids: &[u32],
+        positions: &[usize],
+        cache: &mut Option<Vec<KVCache>>,
     ) -> Result<Tensor> {
         let mut x = self.embedding.forward(token_ids)?;
         let max_seq = self.config.max_seq_len;
@@ -96,11 +155,7 @@ impl CodeGenModel {
             let attn_out = {
                 let qkv = normed.broadcast_matmul(&block.attn.qkv_weight().unsqueeze(0)?)?;
                 let (bs, sl, _) = qkv.dims3()?;
-                let qkv = qkv.reshape((bs, sl, 3, self.config.num_heads, head_dim))?;
-                let qkv = qkv.permute((0, 3, 2, 1, 4))?;
-                let q = qkv.get_on_dim(2, 0)?;
-                let v = qkv.get_on_dim(2, 1)?;
-                let k = qkv.get_on_dim(2, 2)?;
+                let (q, k, v) = split_qkv(&qkv, self.config.num_heads, head_dim)?;
 
                 let q_rot = self.rotary.apply_rotary(&q, positions)?;
                 let k_rot = self.rotary.apply_rotary(&k, positions)?;
@@ -136,8 +191,12 @@ impl CodeGenModel {
             x = (x + attn_out + ffn_out)?;
         }
 
-        let x = self.final_norm.forward(&x)?;
-        let logits = x.broadcast_matmul(&self.lm_head.unsqueeze(0)?)?;
+        self.final_norm.forward(&x)
+    }
+
+    /// Project hidden states onto the vocabulary.
+    pub fn project_logits(&self, hidden: &Tensor) -> Result<Tensor> {
+        let logits = hidden.broadcast_matmul(&self.lm_head.unsqueeze(0)?)?;
         logits.broadcast_add(&self.lm_head_bias)
     }
 }
@@ -157,4 +216,27 @@ mod tests {
         assert_eq!(logits.dims(), &[1, 5, 50400]);
         Ok(())
     }
+
+    /// Regression: `new_blank` hardcoded F32 weights regardless of `config.dtype`,
+    /// so an F16 model panicked with "dtype mismatch in mul" on the first forward.
+    #[test]
+    fn test_codegen_blank_forward_f16() -> Result<()> {
+        let device = Device::Cpu;
+        let config = CodeGenConfig {
+            dtype: candle_core::DType::F16,
+            vocab_size: 64,
+            hidden_dim: 32,
+            num_layers: 1,
+            num_heads: 4,
+            ffn_dim: 64,
+            max_seq_len: 16,
+            rotary_dim: 8,
+            ..Default::default()
+        };
+        let model = CodeGenModel::new_blank(config, &device)?;
+        let logits = model.forward_with_cache(&[1u32, 2, 3], &[0usize, 1, 2], &mut None)?;
+        assert_eq!(logits.dims(), &[1, 3, 64]);
+        assert_eq!(logits.dtype(), candle_core::DType::F16);
+        Ok(())
+    }
 }
diff --git a/src/codegen/quantized.rs b/src/codegen/quantized.rs
deleted file mode 100644
index 8a26593..0000000
--- a/src/codegen/quantized.rs
+++ /dev/null
@@ -1,236 +0,0 @@
-//! INT8 Dynamic Quantization for CodeGen-350M
-//!
-//! Quantizes linear layer weights from F32 to INT8 with per-channel scaling,
-//! reducing memory footprint by ~4× and potentially improving cache performance.
-//!
-//! Quantization scheme:
-//!   weight_int8 = round(weight_f32 / scale)
-//!   scale = max(abs(weight_f32), epsilon) / 127.0
-//!
-//! Forward: out = (weight_int8 * scale) @ input
-
-use candle_core::{DType, Device, Result, Tensor};
-
-/// Quantized linear layer with U8 (offset) weights.
-pub struct QuantizedLinear {
-    /// Quantized weights (u8 with 128 offset), shape [out_dim, in_dim]
-    weight_int8: Tensor,
-    /// Per-output-channel scale factors, shape [out_dim, 1]
-    scale: Tensor,
-    /// Original shape
-    out_dim: usize,
-    in_dim: usize,
-}
-
-impl QuantizedLinear {
-    /// Create a quantized linear layer from F32 weights.
-    ///
-    /// `weight` should have shape [out_dim, in_dim].
-    pub fn from_f32(weight: &Tensor, device: &Device) -> Result<Self> {
-        let dims = weight.dims();
-        if dims.len() != 2 {
-            return Err(candle_core::Error::Msg(format!(
-                "Expected 2D weight tensor, got {}D",
-                dims.len()
-            )));
-        }
-        let out_dim = dims[0];
-        let in_dim = dims[1];
-
-        // Move to F32 on CPU for quantization
-        let w = weight.to_dtype(DType::F32)?.to_device(&Device::Cpu)?;
-        let w_vec: Vec<f32> = w.to_vec2::<f32>()?.into_iter().flatten().collect();
-
-        // Per-output-channel quantization
-        // Store quantized weights as u8 (offset by 128 to handle negative values)
-        let mut quantized: Vec<u8> = Vec::with_capacity(out_dim * in_dim);
-        let mut scales: Vec<f32> = Vec::with_capacity(out_dim);
-
-        for row in 0..out_dim {
-            let start = row * in_dim;
-            let end = start + in_dim;
-            let row_slice = &w_vec[start..end];
-
-            // Find scale for this row (per-channel)
-            let max_abs = row_slice
-                .iter()
-                .map(|v| v.abs())
-                .fold(f32::MIN_POSITIVE, f32::max);
-            let scale = (max_abs / 127.0).max(f32::MIN_POSITIVE);
-
-            // Quantize: map from [-128, 127] to [0, 255] range
-            for &val in row_slice {
-                let q = (val / scale).round().clamp(-128.0, 127.0) as i8;
-                quantized.push((q as i16 + 128) as u8); // offset by 128 to make non-negative
-            }
-            scales.push(scale);
-        }
-
-        let weight_q = Tensor::from_vec(quantized, (out_dim, in_dim), &Device::Cpu)?;
-        let scale_t = Tensor::from_vec(scales, (out_dim, 1), &Device::Cpu)?.to_dtype(DType::F32)?;
-
-        Ok(Self {
-            weight_int8: weight_q.to_device(device)?,
-            scale: scale_t.to_device(device)?,
-            out_dim,
-            in_dim,
-        })
-    }
-
-    /// Forward pass: dequantize on-the-fly and compute matmul.
-    ///
-    /// `input` shape: \[batch, seq\_len, `in_dim`\] or \[`in_dim`\]
-    /// Returns shape: \[batch, seq\_len, `out_dim`\] or \[`out_dim`\]
-    pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
-        // Dequantize: weight_f32 = (weight_u8 - 128) * scale
-        // weight_int8 shape: [out_dim, in_dim]
-        // scale shape: [out_dim, 1]
-        let w_f32 = self.weight_int8.to_dtype(DType::F32)?;
-        let offset = Tensor::new(128.0f32, w_f32.device())?;
-        let w_centered = w_f32.broadcast_sub(&offset)?;
-        let w_deq = w_centered.broadcast_mul(&self.scale)?;
-
-        // w_deq shape: [out_dim, in_dim], we need [in_dim, out_dim] for matmul
-        let w_t = w_deq.t()?;
-
-        if input.dims().len() == 2 {
-            // [batch, in_dim] @ [in_dim, out_dim] = [batch, out_dim]
-            input.matmul(&w_t)
-        } else {
-            // [in_dim] @ [in_dim, out_dim] = [out_dim]
-            input.matmul(&w_t)
-        }
-    }
-
-    /// Estimated size in bytes (INT8 weights + F32 scales)
-    pub fn estimated_size(&self) -> usize {
-        self.out_dim * self.in_dim         // INT8 weights: 1 byte each
-            + self.out_dim * 4 // F32 scales: 4 bytes each
-    }
-
-    /// Original F32 size for comparison
-    pub fn original_size(&self) -> usize {
-        self.out_dim * self.in_dim * 4 // F32: 4 bytes each
-    }
-}
-
-/// A quantized version of the CodeGen attention weights (QKV projection).
-/// In the real model, QKV weight is [hidden_dim, 3 * hidden_dim] (transposed).
-pub struct QuantizedAttention {
-    pub qkv: Option<QuantizedLinear>,
-    pub out: Option<QuantizedLinear>,
-}
-
-impl QuantizedAttention {
-    pub fn new(
-        qkv_weight: Option<&Tensor>,
-        out_weight: Option<&Tensor>,
-        device: &Device,
-    ) -> Result<Self> {
-        let qkv = if let Some(w) = qkv_weight {
-            Some(QuantizedLinear::from_f32(w, device)?)
-        } else {
-            None
-        };
-        let out = if let Some(w) = out_weight {
-            Some(QuantizedLinear::from_f32(w, device)?)
-        } else {
-            None
-        };
-        Ok(Self { qkv, out })
-    }
-}
-
-/// A quantized version of the CodeGen FFN weights.
-pub struct QuantizedFFN {
-    pub fc_in: Option<QuantizedLinear>,
-    pub fc_out: Option<QuantizedLinear>,
-}
-
-impl QuantizedFFN {
-    pub fn new(
-        fc_in_weight: Option<&Tensor>,
-        fc_out_weight: Option<&Tensor>,
-        device: &Device,
-    ) -> Result<Self> {
-        let fc_in = if let Some(w) = fc_in_weight {
-            Some(QuantizedLinear::from_f32(w, device)?)
-        } else {
-            None
-        };
-        let fc_out = if let Some(w) = fc_out_weight {
-            Some(QuantizedLinear::from_f32(w, device)?)
-        } else {
-            None
-        };
-        Ok(Self { fc_in, fc_out })
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_quantized_linear_roundtrip() -> Result<()> {
-        let device = Device::Cpu;
-
-        // Create a simple weight matrix
-        let w = Tensor::from_vec(vec![1.0f32, -2.0, 3.0, -4.0, 5.0, -6.0], (2, 3), &device)?;
-
-        let qlinear = QuantizedLinear::from_f32(&w, &device)?;
-
-        // Create input
-        let input = Tensor::from_vec(vec![1.0f32, 2.0, 3.0], (1, 3), &device)?;
-
-        let output = qlinear.forward(&input)?;
-        let out_vec: Vec<f32> = output.to_vec2::<f32>()?.into_iter().flatten().collect();
-
-        // Should be approximate (INT8 quantization introduces some error)
-        assert_eq!(out_vec.len(), 2, "Output should have 2 values (out_dim=2)");
-
-        // Verify quantization reduces memory
-        let estimated = qlinear.estimated_size();
-        let original = qlinear.original_size();
-        assert!(estimated < original, "Quantized should be smaller");
-        // For large matrices this is ~4x, for small ones scales overhead reduces it
-        assert!(
-            (original as f64 / estimated as f64) > 1.0,
-            "Should have some compression, got {}/{} = {:.1}x",
-            original,
-            estimated,
-            original as f64 / estimated as f64
-        );
-
-        println!("  Original size: {original} bytes");
-        println!("  Quantized size: {estimated} bytes");
-        println!("  Compression: {:.1}x", original as f64 / estimated as f64);
-        println!("  Output: {out_vec:?}");
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_quantized_preserves_ranking() -> Result<()> {
-        let device = Device::Cpu;
-
-        // Create weights with clear pattern
-        let w = Tensor::from_vec(vec![10.0f32, 0.0, 0.0, 0.0, 20.0, 0.0], (2, 3), &device)?;
-
-        let qlinear = QuantizedLinear::from_f32(&w, &device)?;
-
-        let input = Tensor::from_vec(vec![1.0f32, 2.0, 3.0], (1, 3), &device)?;
-        let output = qlinear.forward(&input)?;
-        let out_vec: Vec<f32> = output.to_vec2::<f32>()?.into_iter().flatten().collect();
-
-        // output[0] should be ~10 (first row * input), output[1] should be ~40 (second row * input)
-        assert!(
-            out_vec[1] > out_vec[0],
-            "Second output should be larger (20*2=40 vs 10*1=10), got {:?}",
-            out_vec
-        );
-
-        println!("  Ranking preserved: {out_vec:?}");
-        Ok(())
-    }
-}
diff --git a/src/codegen/weights.rs b/src/codegen/weights.rs
index d403a1b..5a703ea 100644
--- a/src/codegen/weights.rs
+++ b/src/codegen/weights.rs
@@ -1,5 +1,5 @@
 use candle_core::pickle::PthTensors;
-use candle_core::{DType, Device, Result, Tensor};
+use candle_core::{DType, Device, Error, Result, Tensor};
 use std::path::Path;
 
 use super::config::CodeGenConfig;
@@ -9,124 +9,112 @@ use crate::layers::embedding::Embedding;
 use crate::layers::ffn::{FeedForward, GELUNewFFN};
 use crate::layers::norm::LayerNorm;
 
-fn load_tensor(pth: &PthTensors, name: &str, dtype: DType, device: &Device) -> Option<Tensor> {
-    let t = pth.get(name).ok()?;
-    let t = t?;
+fn to(t: Tensor, dtype: DType, device: &Device) -> Result<Tensor> {
     let t = if t.dtype() != dtype {
-        t.to_dtype(dtype).ok()?
+        t.to_dtype(dtype)?
     } else {
         t
     };
-    t.to_device(device).ok()
+    t.to_device(device)
+}
+
+fn missing(name: &str) -> Error {
+    Error::Msg(format!("checkpoint is missing tensor `{name}`"))
 }
 
 pub struct WeightLoader;
 
 impl WeightLoader {
+    /// Load a HuggingFace PyTorch checkpoint (`pytorch_model.bin`).
     pub fn load_from_pytorch(
         path: &Path,
         config: &CodeGenConfig,
         device: &Device,
     ) -> Result<CodeGenModel> {
         let pth = PthTensors::new(path, None)?;
-        let mut model = CodeGenModel::new_blank(config.clone(), device)?;
         let dtype = config.dtype;
+        Self::build(
+            |name| {
+                let t = pth.get(name)?.ok_or_else(|| missing(name))?;
+                to(t, dtype, device)
+            },
+            config,
+            device,
+        )
+    }
 
-        if let Some(w) = load_tensor(&pth, "transformer.wte.weight", dtype, device) {
-            model.embedding = Embedding::from_tensor(w);
+    /// Load a safetensors checkpoint (`model.safetensors`), as produced by
+    /// `examples/convert_codegen_to_safetensors.rs`.
+    pub fn load_from_safetensors(
+        path: &Path,
+        config: &CodeGenConfig,
+        device: &Device,
+    ) -> Result<CodeGenModel> {
+        let tensors = candle_core::safetensors::load(path, device)?;
+        let dtype = config.dtype;
+        Self::build(
+            |name| {
+                let t = tensors.get(name).ok_or_else(|| missing(name))?.clone();
+                to(t, dtype, device)
+            },
+            config,
+            device,
+        )
+    }
+
+    /// Load whichever checkpoint format `path` is, decided by extension.
+    pub fn load(path: &Path, config: &CodeGenConfig, device: &Device) -> Result<CodeGenModel> {
+        match path.extension().and_then(|e| e.to_str()) {
+            Some("safetensors") => Self::load_from_safetensors(path, config, device),
+            _ => Self::load_from_pytorch(path, config, device),
         }
+    }
+
+    /// Assemble a model from a checkpoint accessor.
+    ///
+    /// `get` must fail loudly for an absent tensor: a missing one used to be
+    /// swallowed, leaving that layer at its all-zero initial value, which
+    /// silently zeroes its output instead of reporting a bad checkpoint.
+    fn build<F>(get: F, config: &CodeGenConfig, device: &Device) -> Result<CodeGenModel>
+    where
+        F: Fn(&str) -> Result<Tensor>,
+    {
+        let mut model = CodeGenModel::new_blank(config.clone(), device)?;
+
+        model.embedding = Embedding::from_tensor(get("transformer.wte.weight")?);
 
         for i in 0..config.num_layers {
             let block = &mut model.blocks[i];
 
-            if let Some(w) = load_tensor(
-                &pth,
-                &format!("transformer.h.{i}.ln_1.weight"),
-                dtype,
-                device,
-            ) {
-                let b = load_tensor(&pth, &format!("transformer.h.{i}.ln_1.bias"), dtype, device)
-                    .unwrap_or(Tensor::zeros(config.hidden_dim, dtype, device).unwrap());
-                block.norm1 = LayerNorm::from_tensor(w, b, config.eps);
-            }
-
-            if let Some(w) = load_tensor(
-                &pth,
-                &format!("transformer.h.{i}.attn.qkv_proj.weight"),
-                dtype,
-                device,
-            ) {
-                let w_t = w.transpose(0, 1)?;
-                let out_w = load_tensor(
-                    &pth,
-                    &format!("transformer.h.{i}.attn.out_proj.weight"),
-                    dtype,
-                    device,
-                )
-                .unwrap()
-                .transpose(0, 1)?;
-                block.attn = MultiHeadAttention::from_tensors(w_t, out_w, config.num_heads);
-            }
-
-            if let Some(w) = load_tensor(
-                &pth,
-                &format!("transformer.h.{i}.ln_2.weight"),
-                dtype,
-                device,
-            ) {
-                let b = load_tensor(&pth, &format!("transformer.h.{i}.ln_2.bias"), dtype, device)
-                    .unwrap_or(Tensor::zeros(config.hidden_dim, dtype, device).unwrap());
-                block.norm2 = LayerNorm::from_tensor(w, b, config.eps);
-            }
-
-            if let Some(fc_in) = load_tensor(
-                &pth,
-                &format!("transformer.h.{i}.mlp.fc_in.weight"),
-                dtype,
-                device,
-            ) {
-                let fc_in_t = fc_in.transpose(0, 1)?;
-                let fc_in_b = load_tensor(
-                    &pth,
-                    &format!("transformer.h.{i}.mlp.fc_in.bias"),
-                    dtype,
-                    device,
-                )
-                .unwrap_or(Tensor::zeros(config.ffn_dim, dtype, device).unwrap());
-                let fc_out = load_tensor(
-                    &pth,
-                    &format!("transformer.h.{i}.mlp.fc_out.weight"),
-                    dtype,
-                    device,
-                )
-                .unwrap();
-                let fc_out_t = fc_out.transpose(0, 1)?;
-                let fc_out_b = load_tensor(
-                    &pth,
-                    &format!("transformer.h.{i}.mlp.fc_out.bias"),
-                    dtype,
-                    device,
-                )
-                .unwrap_or(Tensor::zeros(config.hidden_dim, dtype, device).unwrap());
-                block.ffn = FeedForward::GELUNew(GELUNewFFN::from_tensors_with_bias(
-                    fc_in_t, fc_in_b, fc_out_t, fc_out_b,
-                ));
-            }
-        }
+            block.norm1 = LayerNorm::from_tensor(
+                get(&format!("transformer.h.{i}.ln_1.weight"))?,
+                get(&format!("transformer.h.{i}.ln_1.bias"))?,
+                config.eps,
+            );
 
-        if let Some(w) = load_tensor(&pth, "transformer.ln_f.weight", dtype, device) {
-            let b = load_tensor(&pth, "transformer.ln_f.bias", dtype, device)
-                .unwrap_or(Tensor::zeros(config.hidden_dim, dtype, device).unwrap());
-            model.final_norm = LayerNorm::from_tensor(w, b, config.eps);
-        }
+            block.attn = MultiHeadAttention::from_tensors(
+                get(&format!("transformer.h.{i}.attn.qkv_proj.weight"))?.transpose(0, 1)?,
+                get(&format!("transformer.h.{i}.attn.out_proj.weight"))?.transpose(0, 1)?,
+                config.num_heads,
+            );
 
-        if let Some(w) = load_tensor(&pth, "lm_head.weight", dtype, device) {
-            model.lm_head = w.transpose(0, 1)?;
-            if let Some(b) = load_tensor(&pth, "lm_head.bias", dtype, device) {
-                model.lm_head_bias = b;
-            }
+            block.ffn = FeedForward::GELUNew(GELUNewFFN::from_tensors_with_bias(
+                get(&format!("transformer.h.{i}.mlp.fc_in.weight"))?.transpose(0, 1)?,
+                get(&format!("transformer.h.{i}.mlp.fc_in.bias"))?,
+                get(&format!("transformer.h.{i}.mlp.fc_out.weight"))?.transpose(0, 1)?,
+                get(&format!("transformer.h.{i}.mlp.fc_out.bias"))?,
+            ));
         }
 
+        model.final_norm = LayerNorm::from_tensor(
+            get("transformer.ln_f.weight")?,
+            get("transformer.ln_f.bias")?,
+            config.eps,
+        );
+
+        model.lm_head = get("lm_head.weight")?.transpose(0, 1)?;
+        model.lm_head_bias = get("lm_head.bias")?;
+
         Ok(model)
     }
 }
diff --git a/src/commands/chat.rs b/src/commands/chat.rs
index 18cbef8..34b1f3d 100644
--- a/src/commands/chat.rs
+++ b/src/commands/chat.rs
@@ -9,6 +9,7 @@ use crate::model::ModelContext;
 
 pub fn run(cli: &Cli, system: Option<String>) -> Result<()> {
     let mut ctx = ModelContext::load_default(&cli.weights_dir, cli.f16)?;
+    ctx.generator.set_seed(cli.seed);
 
     let mut session = ChatSession::new(1024);
     if let Some(sys) = system {
@@ -133,9 +134,9 @@ pub fn run(cli: &Cli, system: Option<String>) -> Result<()> {
         ctx.generator.generate_stream(&prompt_ids, &mut collector)?;
         let elapsed = start.elapsed();
 
-        let all_tokens = collector.tokens;
-        let generated_ids = &all_tokens[prompt_ids.len()..];
-        let response = ctx.tokenizer.decode(generated_ids)?;
+        // `generate_stream` reports generated tokens only; the prompt is not replayed.
+        let generated_ids = collector.tokens;
+        let response = ctx.tokenizer.decode(&generated_ids)?;
 
         if !response.is_empty() {
             println!("{response}");
diff --git a/src/commands/complete.rs b/src/commands/complete.rs
index 6150820..ea4b96a 100644
--- a/src/commands/complete.rs
+++ b/src/commands/complete.rs
@@ -16,6 +16,7 @@ pub fn run(
 ) -> Result<()> {
     let mut ctx = ModelContext::load(&cli.weights_dir, cli.f16, temperature)?;
     ctx.generator.set_max_new_tokens(max_tokens);
+    ctx.generator.set_seed(cli.seed);
 
     // Apply prompt template
     let formatted = match template {
@@ -49,7 +50,7 @@ pub fn run(
         ctx.generator.generate_stream(&token_ids, &mut collector)?;
         let elapsed = start.elapsed();
 
-        let generated_ids = &collector.tokens[token_ids.len()..];
+        let generated_ids = &collector.tokens[..];
         let output = ctx.tokenizer.decode(generated_ids)?;
         println!("{output}");
         println!(
diff --git a/src/commands/download.rs b/src/commands/download.rs
index 7ababf5..9c11bf8 100644
--- a/src/commands/download.rs
+++ b/src/commands/download.rs
@@ -1,18 +1,35 @@
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
-use anyhow::Result;
+use anyhow::{bail, Result};
 
 use crate::cli::Cli;
 
-pub fn run(cli: &Cli) -> Result<()> {
-    let weights_path = cli.weights_dir.join("pytorch_model.bin");
+/// The CLI names to try, in order.
+///
+/// `huggingface-cli` was renamed to `hf`. The old name still exists but now
+/// prints a deprecation notice, downloads nothing, and **exits 0** — so trusting
+/// the exit status reported a successful download of no files at all.
+const DOWNLOAD_TOOLS: &[&str] = &["hf", "huggingface-cli"];
+
+const REPO: &str = "Salesforce/codegen-350M-multi";
 
-    if weights_path.exists() {
-        let meta = std::fs::metadata(&weights_path)?;
-        let size_mb = meta.len() as f64 / (1024.0 * 1024.0);
+/// The checkpoint file, if one is already there. Mirrors the order
+/// `ModelContext::load` prefers.
+fn existing_weights(dir: &Path) -> Option<(PathBuf, f64)> {
+    for name in ["model.safetensors", "pytorch_model.bin"] {
+        let path = dir.join(name);
+        if let Ok(meta) = std::fs::metadata(&path) {
+            return Some((path, meta.len() as f64 / (1024.0 * 1024.0)));
+        }
+    }
+    None
+}
+
+pub fn run(cli: &Cli) -> Result<()> {
+    if let Some((path, size_mb)) = existing_weights(&cli.weights_dir) {
         println!(
             "Weights already present at: {} ({:.1} MB)",
-            weights_path.display(),
+            path.display(),
             size_mb
         );
         return Ok(());
@@ -20,46 +37,55 @@ pub fn run(cli: &Cli) -> Result<()> {
 
     println!("\x1b[1mDownloading CodeGen-350M-multi from HuggingFace...\x1b[0m\n");
 
-    // Try huggingface-cli first
-    let status = std::process::Command::new("huggingface-cli")
-        .args([
-            "download",
-            "Salesforce/codegen-350M-multi",
-            "--local-dir",
-            cli.weights_dir.to_str().unwrap(),
-        ])
-        .status();
+    let dir = cli
+        .weights_dir
+        .to_str()
+        .ok_or_else(|| anyhow::anyhow!("weights directory path is not valid UTF-8"))?;
 
-    match status {
-        Ok(s) if s.success() => {
+    for tool in DOWNLOAD_TOOLS {
+        let status = std::process::Command::new(tool)
+            .args(["download", REPO, "--local-dir", dir])
+            .status();
+
+        let Ok(status) = status else {
+            continue; // not installed, try the next name
+        };
+
+        // Check for the file rather than believing the exit code.
+        if let Some((path, size_mb)) = existing_weights(&cli.weights_dir) {
             println!("\n\x1b[32m✓ Download complete!\x1b[0m");
+            println!("  {} ({:.1} MB)", path.display(), size_mb);
             println!("  Run `codegen info` to verify.");
+            return Ok(());
         }
-        Ok(_) => {
-            println!("\n\x1b[31mhuggingface-cli failed.\x1b[0m");
-            print_manual_instructions(&cli.weights_dir);
-        }
-        Err(_) => {
-            println!("huggingface-cli not found. Install it:");
-            println!("  pip install huggingface_hub\n");
-            print_manual_instructions(&cli.weights_dir);
+
+        if status.success() {
+            println!("\n\x1b[33m`{tool}` reported success but no weights appeared.\x1b[0m");
+        } else {
+            println!("\n\x1b[31m`{tool}` failed.\x1b[0m");
         }
     }
 
-    Ok(())
+    print_manual_instructions(&cli.weights_dir);
+    bail!("could not download weights automatically");
 }
 
 fn print_manual_instructions(weights_dir: &Path) {
-    println!("\nManual download:");
+    println!("\nInstall the HuggingFace CLI:");
+    println!("  pip install -U huggingface_hub\n");
+    println!("Or download manually:");
     println!("  1. Install git-lfs:");
     println!("     git lfs install\n");
     println!("  2. Clone the repo:");
     println!(
-        "     git clone https://huggingface.co/Salesforce/codegen-350M-multi {}\n",
+        "     git clone https://huggingface.co/{REPO} {}\n",
         weights_dir.display()
     );
-    println!("  3. Or use wget for individual files:");
-    println!("     wget -P {}/ https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/pytorch_model.bin", weights_dir.display());
-    println!("     wget -P {}/ https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/tokenizer.json", weights_dir.display());
-    println!("     wget -P {}/ https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json", weights_dir.display());
+    println!("  3. Or fetch the individual files:");
+    for file in ["pytorch_model.bin", "tokenizer.json", "config.json"] {
+        println!(
+            "     wget -P {}/ https://huggingface.co/{REPO}/resolve/main/{file}",
+            weights_dir.display()
+        );
+    }
 }
diff --git a/src/commands/glm_train.rs b/src/commands/glm_train.rs
index 70569f0..1b56ce0 100644
--- a/src/commands/glm_train.rs
+++ b/src/commands/glm_train.rs
@@ -6,32 +6,36 @@ use candle_core::Device;
 use crate::cli::Cli;
 use crate::training::{GLMTrainer, TrainConfig};
 
-pub fn run(cli: &Cli, data_path: &str, steps: usize) -> Result<()> {
+pub fn run(cli: &Cli, data_path: &str, steps: usize, config_path: Option<&Path>) -> Result<()> {
     let device = Device::Cpu;
     let dtype_str = if cli.f16 { "FP16" } else { "FP32" };
 
     println!("\x1b[1mGLM Training Demo ({dtype_str})\x1b[0m\n");
 
-    let config = TrainConfig {
-        training: crate::training::TrainingConfig {
-            data_dir: PathBuf::from(data_path),
-            max_steps: steps,
-            dtype: if cli.f16 {
-                "f16".to_string()
-            } else {
-                "f32".to_string()
-            },
-            ..crate::training::TrainingConfig::default()
-        },
-        ..TrainConfig::default()
+    let mut config = match config_path {
+        Some(path) => {
+            println!("Config: {}", path.display());
+            TrainConfig::from_file(path).map_err(|e| anyhow::anyhow!("{e}"))?
+        }
+        None => TrainConfig::default(),
     };
 
+    // CLI flags win over the file.
+    config.training.data_dir = PathBuf::from(data_path);
+    config.training.max_steps = steps;
+    config.training.dtype = if cli.f16 { "f16" } else { "f32" }.to_string();
+
+    let model = &config.model;
     println!("GLM Config:");
-    println!("  vocab_size:   50257");
-    println!("  hidden_dim:   1024");
-    println!("  num_layers:   4");
-    println!("  num_heads:    8");
-    println!("  ffn_dim:      4096");
+    println!("  vocab_size:   {}", model.vocab_size);
+    println!("  hidden_dim:   {}", model.hidden_dim);
+    println!("  num_layers:   {}", model.num_layers);
+    println!("  num_heads:    {}", model.num_heads);
+    println!("  ffn_dim:      {}", model.ffn_dim);
+    println!(
+        "  batch:        {} x {} (micro x accumulation)",
+        config.training.micro_batch_size, config.training.gradient_accumulation_steps
+    );
 
     // Check if data exists
     let data_dir = Path::new(data_path);
diff --git a/src/commands/info.rs b/src/commands/info.rs
index 0676b9d..7400eab 100644
--- a/src/commands/info.rs
+++ b/src/commands/info.rs
@@ -3,6 +3,7 @@ use anyhow::Result;
 use crate::cli::Cli;
 use crate::codegen::config::CodeGenConfig;
 use crate::glm::GLMConfig;
+use crate::tokenizer::CodeGenTokenizer;
 
 pub fn run(cli: &Cli) -> Result<()> {
     // GLM info
@@ -14,53 +15,64 @@ pub fn run(cli: &Cli) -> Result<()> {
     println!("  num_heads:    {}", glm_config.num_heads);
     println!("  ffn_dim:      {}", glm_config.ffn_dim);
 
-    // CodeGen info
-    println!("\n\x1b[1mCodeGen-350M:\x1b[0m");
-    println!("  vocab_size:   51200");
-    println!("  hidden_dim:   1024");
-    println!("  num_layers:   20");
-    println!("  num_heads:    16");
-    println!("  ffn_dim:      4096");
-    println!("  max_seq_len:  1024");
-    println!("  params:       ~350M");
-
-    // Weight status
-    let weights_path = cli.weights_dir.join("pytorch_model.bin");
     let config_path = cli.weights_dir.join("config.json");
     let tokenizer_path = cli.weights_dir.join("tokenizer.json");
 
-    println!("\n\x1b[1mWeight Status:\x1b[0m");
-    if weights_path.exists() {
-        let meta = std::fs::metadata(&weights_path)?;
-        let size_mb = meta.len() as f64 / (1024.0 * 1024.0);
-        println!(
-            "  weights:    \x1b[32m✓\x1b[0m {} ({:.1} MB)",
-            weights_path.display(),
-            size_mb
-        );
+    // Report the checkpoint's own config rather than hardcoded numbers, which had
+    // drifted from it (max_seq_len was printed as 1024; the checkpoint says 2048).
+    let config = if config_path.exists() {
+        let config_str = std::fs::read_to_string(&config_path)?;
+        let config_json: serde_json::Value = serde_json::from_str(&config_str)?;
+        CodeGenConfig::from_hf_config(&config_json)
     } else {
-        println!(
-            "  weights:    \x1b[31m✗\x1b[0m Not found at {}",
-            weights_path.display()
-        );
+        CodeGenConfig::default()
+    };
+
+    println!("\n\x1b[1mCodeGen-350M:\x1b[0m");
+    println!("  vocab_size:   {}", config.vocab_size);
+    println!("  hidden_dim:   {}", config.hidden_dim);
+    println!("  num_layers:   {}", config.num_layers);
+    println!("  num_heads:    {}", config.num_heads);
+    println!("  ffn_dim:      {}", config.ffn_dim);
+    println!("  max_seq_len:  {}", config.max_seq_len);
+    println!("  rotary_dim:   {}", config.rotary_dim);
+    println!("  dtype:        {:?}", if cli.f16 { "f16" } else { "f32" });
+
+    println!("\n\x1b[1mWeight Status:\x1b[0m");
+
+    // Same order of preference as `ModelContext::load`.
+    let weights = ["model.safetensors", "pytorch_model.bin"]
+        .into_iter()
+        .map(|name| cli.weights_dir.join(name))
+        .find(|path| path.exists());
+
+    match &weights {
+        Some(path) => {
+            let size_mb = std::fs::metadata(path)?.len() as f64 / (1024.0 * 1024.0);
+            println!(
+                "  weights:    \x1b[32m✓\x1b[0m {} ({:.1} MB)",
+                path.display(),
+                size_mb
+            );
+        }
+        None => println!(
+            "  weights:    \x1b[31m✗\x1b[0m No model.safetensors or pytorch_model.bin in {}",
+            cli.weights_dir.display()
+        ),
     }
 
     if config_path.exists() {
-        let config_str = std::fs::read_to_string(&config_path)?;
-        let config_json: serde_json::Value = serde_json::from_str(&config_str)?;
-        let config = CodeGenConfig::from_hf_config(&config_json);
-        println!(
-            "  config:     \x1b[32m✓\x1b[0m {} (dtype: {:?})",
-            config_path.display(),
-            config.dtype
-        );
+        println!("  config:     \x1b[32m✓\x1b[0m {}", config_path.display());
     } else {
         println!("  config:     \x1b[33m-\x1b[0m Not found (using defaults)");
     }
 
     if tokenizer_path.exists() {
+        let vocab = CodeGenTokenizer::from_file(tokenizer_path.to_str().unwrap())
+            .map(|t| t.vocab_size().to_string())
+            .unwrap_or_else(|_| "unreadable".to_string());
         println!(
-            "  tokenizer:  \x1b[32m✓\x1b[0m {}",
+            "  tokenizer:  \x1b[32m✓\x1b[0m {} ({vocab} tokens)",
             tokenizer_path.display()
         );
     } else {
@@ -70,8 +82,7 @@ pub fn run(cli: &Cli) -> Result<()> {
         );
     }
 
-    let all_present = weights_path.exists() && tokenizer_path.exists();
-    if all_present {
+    if weights.is_some() && tokenizer_path.exists() {
         println!(
             "\n  \x1b[32mReady to run: `codegen chat` or `codegen complete \"prompt\"`\x1b[0m"
         );
diff --git a/src/commands/repl.rs b/src/commands/repl.rs
index a96edce..7114ffd 100644
--- a/src/commands/repl.rs
+++ b/src/commands/repl.rs
@@ -7,7 +7,8 @@ use crate::generation::codegen_generate::CollectStream;
 use crate::model::ModelContext;
 
 pub fn run(cli: &Cli) -> Result<()> {
-    let ctx = ModelContext::load_default(&cli.weights_dir, cli.f16)?;
+    let mut ctx = ModelContext::load_default(&cli.weights_dir, cli.f16)?;
+    ctx.generator.set_seed(cli.seed);
 
     println!("CodeGen-350M REPL — type a prompt and get generated code.");
     println!("Type 'exit' or 'quit' to stop.\n");
@@ -35,7 +36,7 @@ pub fn run(cli: &Cli) -> Result<()> {
         ctx.generator.generate_stream(&token_ids, &mut collector)?;
         let elapsed = start.elapsed();
 
-        let generated_ids = &collector.tokens[token_ids.len()..];
+        let generated_ids = &collector.tokens[..];
         let output = ctx.tokenizer.decode(generated_ids)?;
         println!("\n{output}");
         println!(
diff --git a/src/commands/serve.rs b/src/commands/serve.rs
index da24ca8..a5d9251 100644
--- a/src/commands/serve.rs
+++ b/src/commands/serve.rs
@@ -10,7 +10,9 @@ pub fn run(cli: &Cli, port: u16) -> Result<()> {
         let addr: SocketAddr = ([0, 0, 0, 0], port).into();
         let rt = tokio::runtime::Runtime::new()?;
         rt.block_on(server::start_server(
-            cli.weights_dir.join("pytorch_model.bin").to_str().unwrap(),
+            &cli.weights_dir,
+            cli.f16,
+            cli.seed,
             addr,
         ))
     }
diff --git a/src/generation/codegen_generate.rs b/src/generation/codegen_generate.rs
index 8f9dbd5..71b6f40 100644
--- a/src/generation/codegen_generate.rs
+++ b/src/generation/codegen_generate.rs
@@ -1,4 +1,6 @@
 use candle_core::Result;
+use rand::rngs::StdRng;
+use rand::SeedableRng;
 
 use crate::codegen::kv_cache::KVCache;
 use crate::codegen::model::CodeGenModel;
@@ -93,6 +95,8 @@ pub struct CodeGenGenerator {
     tokenizer: Option<crate::tokenizer::CodeGenTokenizer>,
     /// Prompt template to apply
     template: PromptTemplate,
+    /// Fixed RNG seed for reproducible sampling. `None` seeds from entropy.
+    seed: Option<u64>,
 }
 
 impl CodeGenGenerator {
@@ -113,6 +117,7 @@ impl CodeGenGenerator {
             max_new_tokens,
             tokenizer: None,
             template: PromptTemplate::Completion,
+            seed: None,
         }
     }
 
@@ -143,6 +148,10 @@ impl CodeGenGenerator {
     pub fn set_repetition_penalty(&mut self, p: f64) {
         self.repetition_penalty = p;
     }
+    /// Fix the sampling seed so a run is reproducible. `None` seeds from entropy.
+    pub fn set_seed(&mut self, seed: Option<u64>) {
+        self.seed = seed;
+    }
 
     pub fn temperature(&self) -> f64 {
         self.temperature
@@ -153,6 +162,7 @@ impl CodeGenGenerator {
 
     // ── Standard generate (collects all tokens) ──
 
+    /// Returns the **generated** tokens only — the prompt is not included.
     pub fn generate(&self, prompt_token_ids: &[u32]) -> Result<Vec<u32>> {
         let mut collector = CollectStream::new();
         self.generate_stream(prompt_token_ids, &mut collector)?;
@@ -163,20 +173,29 @@ impl CodeGenGenerator {
 
     /// Generate tokens, calling `handler.on_token()` for each new token.
     /// The handler can return `false` to stop generation early.
+    ///
+    /// The handler only sees generated tokens; the prompt is never replayed.
     pub fn generate_stream(
         &self,
         prompt_token_ids: &[u32],
         handler: &mut dyn StreamHandler,
     ) -> Result<()> {
-        let gen_start = std::time::Instant::now();
+        // One RNG for the whole generation. Re-seeding per token would draw the
+        // same value at every step.
+        let mut rng = match self.seed {
+            Some(seed) => StdRng::seed_from_u64(seed),
+            None => StdRng::from_entropy(),
+        };
 
         let mut cache: Option<Vec<KVCache>> = None;
         let positions: Vec<usize> = (0..prompt_token_ids.len()).collect();
-        let logits = self
+        // Only the final position drives the next token, so project just that row
+        // instead of running lm_head over the whole prompt.
+        let hidden = self
             .model
-            .forward_with_cache(prompt_token_ids, &positions, &mut cache)?;
-
-        let last_logits = logits.get(0)?.get(prompt_token_ids.len() - 1)?;
+            .forward_hidden(prompt_token_ids, &positions, &mut cache)?;
+        let last_hidden = hidden.narrow(1, prompt_token_ids.len() - 1, 1)?;
+        let last_logits = self.model.project_logits(&last_hidden)?.get(0)?.get(0)?;
         let first_token = sample(
             &last_logits,
             self.temperature,
@@ -184,7 +203,7 @@ impl CodeGenGenerator {
             self.top_p,
             self.repetition_penalty,
             prompt_token_ids,
-            42,
+            &mut rng,
         )?;
 
         let mut generated = prompt_token_ids.to_vec();
@@ -196,18 +215,12 @@ impl CodeGenGenerator {
             return Ok(());
         }
 
-        eprintln!(
-            "  [prefill done in {:.1}s]",
-            gen_start.elapsed().as_secs_f64()
-        );
-
-        for step in 1..self.max_new_tokens {
-            let step_start = std::time::Instant::now();
+        for _step in 1..self.max_new_tokens {
             let input_id = [generated.last().copied().unwrap()];
             let pos = vec![generated.len() - 1];
 
-            let logits = self.model.forward_with_cache(&input_id, &pos, &mut cache)?;
-            let token_logits = logits.get(0)?.get(0)?;
+            let hidden = self.model.forward_hidden(&input_id, &pos, &mut cache)?;
+            let token_logits = self.model.project_logits(&hidden)?.get(0)?.get(0)?;
 
             let token_id = sample(
                 &token_logits,
@@ -216,7 +229,7 @@ impl CodeGenGenerator {
                 self.top_p,
                 self.repetition_penalty,
                 &generated,
-                42,
+                &mut rng,
             )?;
 
             if token_id == EOS_TOKEN_ID {
@@ -230,12 +243,6 @@ impl CodeGenGenerator {
             if !handler.on_token(token_id, &text) {
                 break; // handler requested early stop
             }
-
-            eprintln!(
-                "  [token {step}/{len} in {:.1}s]",
-                step_start.elapsed().as_secs_f64(),
-                len = self.max_new_tokens
-            );
         }
 
         Ok(())
diff --git a/src/generation/glm_generate.rs b/src/generation/glm_generate.rs
index 8eeee87..89e87d3 100644
--- a/src/generation/glm_generate.rs
+++ b/src/generation/glm_generate.rs
@@ -3,6 +3,8 @@ use candle_core::Result;
 use crate::generation::sampling::sample;
 use crate::glm::attention_mask::build_glm_mask;
 use crate::glm::model::GLMModel;
+use rand::rngs::StdRng;
+use rand::SeedableRng;
 
 pub struct GLMGenerator {
     model: GLMModel,
@@ -46,6 +48,7 @@ impl GLMGenerator {
     /// No blanks — generates new tokens after the prompt.
     pub fn generate(&self, prompt_token_ids: &[u32]) -> Result<Vec<u32>> {
         let device = self.model.embedding.weight.device();
+        let mut rng = StdRng::from_entropy();
         let mut generated = prompt_token_ids.to_vec();
 
         for _step in 0..self.max_new_tokens {
@@ -63,7 +66,7 @@ impl GLMGenerator {
                 self.top_p,
                 self.repetition_penalty,
                 &generated,
-                42,
+                &mut rng,
             )?;
 
             if Some(token_id) == self.eos_token_id {
@@ -86,6 +89,7 @@ impl GLMGenerator {
     /// - earlier positions within the same blank (causal within-span)
     pub fn fill_blanks(&self, context: &[u32], blank_lens: &[usize]) -> Result<Vec<u32>> {
         let device = self.model.embedding.weight.device();
+        let mut rng = StdRng::from_entropy();
         let total_blank_tokens: usize = blank_lens.iter().sum();
 
         let mut all_tokens = context.to_vec();
@@ -110,7 +114,7 @@ impl GLMGenerator {
                     self.top_p,
                     self.repetition_penalty,
                     &all_tokens,
-                    42,
+                    &mut rng,
                 )?;
 
                 all_tokens[position] = token_id;
diff --git a/src/generation/sampling.rs b/src/generation/sampling.rs
index 52640de..7236878 100644
--- a/src/generation/sampling.rs
+++ b/src/generation/sampling.rs
@@ -1,5 +1,10 @@
 use candle_core::{DType, Result, Tensor};
+use rand::Rng;
 
+/// Sample one token from `logits`.
+///
+/// `rng` is supplied by the caller so a whole generation can share one stream —
+/// re-seeding per token would draw the same value every step.
 pub fn sample(
     logits: &Tensor,
     temperature: f64,
@@ -7,7 +12,7 @@ pub fn sample(
     top_p: f64,
     repetition_penalty: f64,
     seen_tokens: &[u32],
-    seed: u64,
+    rng: &mut impl Rng,
 ) -> Result<u32> {
     let logits = if logits.dtype() != DType::F32 {
         logits.to_dtype(DType::F32)?
@@ -89,11 +94,7 @@ pub fn sample(
         *p /= sum;
     }
 
-    let mut state = seed;
-    let r = {
-        state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
-        (state as f32) / (u32::MAX as f32)
-    };
+    let r: f32 = rng.gen();
 
     cumulative = 0.0;
     for &(idx, p) in &indexed {
@@ -110,6 +111,8 @@ pub fn sample(
 mod tests {
     use super::*;
     use candle_core::Device;
+    use rand::rngs::StdRng;
+    use rand::SeedableRng;
 
     #[test]
     fn test_argmax_sampling() -> Result<()> {
@@ -117,7 +120,8 @@ mod tests {
         let mut data = vec![0.0f32; 100];
         data[42] = 10.0;
         let logits = Tensor::from_vec(data, 100, &device)?;
-        let token = sample(&logits, 0.0, 1, 1.0, 1.0, &[], 0)?;
+        let mut rng = StdRng::seed_from_u64(7);
+        let token = sample(&logits, 0.0, 1, 1.0, 1.0, &[], &mut rng)?;
         assert_eq!(token, 42);
         Ok(())
     }
@@ -129,7 +133,8 @@ mod tests {
         data[99] = 5.0;
         data[50] = 3.0;
         let logits = Tensor::from_vec(data, 100, &device)?;
-        let token = sample(&logits, 0.0, 50, 1.0, 1.0, &[], 0)?;
+        let mut rng = StdRng::seed_from_u64(7);
+        let token = sample(&logits, 0.0, 50, 1.0, 1.0, &[], &mut rng)?;
         assert_eq!(token, 99);
         Ok(())
     }
@@ -139,7 +144,8 @@ mod tests {
         let device = Device::Cpu;
         let data = vec![1.0f32; 10];
         let logits = Tensor::from_vec(data, 10, &device)?;
-        let token = sample(&logits, 0.0, 10, 1.0, 2.0, &[5], 0)?;
+        let mut rng = StdRng::seed_from_u64(7);
+        let token = sample(&logits, 0.0, 10, 1.0, 2.0, &[5], &mut rng)?;
         assert_ne!(token, 5);
         Ok(())
     }
@@ -152,7 +158,8 @@ mod tests {
         data[1] = 9.0;
         data[2] = 8.0;
         let logits = Tensor::from_vec(data, 100, &device)?;
-        let token = sample(&logits, 0.0, 3, 1.0, 1.0, &[], 0)?;
+        let mut rng = StdRng::seed_from_u64(7);
+        let token = sample(&logits, 0.0, 3, 1.0, 1.0, &[], &mut rng)?;
         assert!(token <= 2);
         Ok(())
     }
@@ -164,7 +171,8 @@ mod tests {
         data[0] = 10.0;
         data[1] = 9.0;
         let logits = Tensor::from_vec(data, 100, &device)?;
-        let token = sample(&logits, 1.0, 100, 0.5, 1.0, &[], 42)?;
+        let mut rng = StdRng::seed_from_u64(7);
+        let token = sample(&logits, 1.0, 100, 0.5, 1.0, &[], &mut rng)?;
         assert!(token <= 1);
         Ok(())
     }
@@ -174,7 +182,8 @@ mod tests {
         let device = Device::Cpu;
         let data = vec![1.0f32; 10];
         let logits = Tensor::from_vec(data, 10, &device)?;
-        let token = sample(&logits, 1.0, 10, 1.0, 1.0, &[], 0)?;
+        let mut rng = StdRng::seed_from_u64(7);
+        let token = sample(&logits, 1.0, 10, 1.0, 1.0, &[], &mut rng)?;
         assert!(token < 10);
         Ok(())
     }
@@ -184,7 +193,8 @@ mod tests {
         let device = Device::Cpu;
         let data = vec![5.0f32];
         let logits = Tensor::from_vec(data, 1, &device)?;
-        let token = sample(&logits, 1.0, 1, 1.0, 1.0, &[], 0)?;
+        let mut rng = StdRng::seed_from_u64(7);
+        let token = sample(&logits, 1.0, 1, 1.0, 1.0, &[], &mut rng)?;
         assert_eq!(token, 0);
         Ok(())
     }
@@ -195,8 +205,46 @@ mod tests {
         let mut data = vec![-10.0f32; 50];
         data[25] = -1.0;
         let logits = Tensor::from_vec(data, 50, &device)?;
-        let token = sample(&logits, 0.0, 50, 1.0, 1.0, &[], 0)?;
+        let mut rng = StdRng::seed_from_u64(7);
+        let token = sample(&logits, 0.0, 50, 1.0, 1.0, &[], &mut rng)?;
         assert_eq!(token, 25);
         Ok(())
     }
+
+    /// Regression: sampling used to fall through to argmax for every token
+    /// because the RNG produced values far outside [0, 1).
+    #[test]
+    fn test_sampling_is_not_always_argmax() -> Result<()> {
+        let device = Device::Cpu;
+        let logits = Tensor::from_vec(vec![0.0f32; 32], 32, &device)?;
+        let mut rng = StdRng::seed_from_u64(7);
+
+        let mut seen = std::collections::HashSet::new();
+        for _ in 0..50 {
+            seen.insert(sample(&logits, 1.0, 32, 1.0, 1.0, &[], &mut rng)?);
+        }
+        assert!(
+            seen.len() > 1,
+            "uniform logits should not always yield the same token, got {seen:?}"
+        );
+        Ok(())
+    }
+
+    /// Same seed, same sequence — reproducibility for `--seed`.
+    #[test]
+    fn test_sampling_is_reproducible_for_a_given_seed() -> Result<()> {
+        let device = Device::Cpu;
+        let logits = Tensor::from_vec(vec![0.0f32; 32], 32, &device)?;
+
+        let draw = |seed: u64| -> Result<Vec<u32>> {
+            let mut rng = StdRng::seed_from_u64(seed);
+            (0..10)
+                .map(|_| sample(&logits, 1.0, 32, 1.0, 1.0, &[], &mut rng))
+                .collect()
+        };
+
+        assert_eq!(draw(11)?, draw(11)?);
+        assert_ne!(draw(11)?, draw(12)?);
+        Ok(())
+    }
 }
diff --git a/src/glm/model.rs b/src/glm/model.rs
index fbdb8ea..4d467f2 100644
--- a/src/glm/model.rs
+++ b/src/glm/model.rs
@@ -105,28 +105,7 @@ impl GLMModel {
         labels: &[i64],
     ) -> Result<Tensor> {
         let logits = self.forward(token_ids, context_len, blank_lens, mask)?;
-
-        // Compute cross-entropy on non -1 labels
-        let seq_len = token_ids.len();
-        let mut total_loss = 0.0f32;
-        let mut count = 0;
-
-        for (i, &label) in labels.iter().enumerate().take(seq_len) {
-            if label >= 0 {
-                let logits_i = logits.get(0)?.get(i)?;
-                let ce = candle_nn::ops::log_softmax(&logits_i, 0)?
-                    .get(label as usize)?
-                    .neg()?;
-                total_loss += ce.to_scalar::<f32>()?;
-                count += 1;
-            }
-        }
-
-        if count > 0 {
-            total_loss /= count as f32;
-        }
-
-        Tensor::new(total_loss, logits.device())
+        crate::training::loss::cross_entropy_loss(&logits, labels)
     }
 }
 
diff --git a/src/glm/trainable.rs b/src/glm/trainable.rs
index 1d51757..4f4985c 100644
--- a/src/glm/trainable.rs
+++ b/src/glm/trainable.rs
@@ -1,4 +1,4 @@
-use candle_core::{DType, Device, Result, Tensor, Var};
+use candle_core::{Device, Result, Tensor, Var};
 
 use safetensors::{serialize, SafeTensors};
 use std::collections::HashMap;
@@ -333,67 +333,19 @@ impl TrainableGLMModel {
         let vars = self.param_vars();
         let names = self.param_names();
 
-        // Collect all tensor data first to keep it alive
         let mut tensor_data = Vec::new();
         for (name, var) in names.iter().zip(vars.iter()) {
             let tensor = var.as_tensor();
-            let dtype = tensor.dtype();
             let shape = tensor.shape().dims().to_vec();
-            println!(
-                "Processing tensor: {} shape={:?} dtype={:?}",
-                name, shape, dtype
-            );
-
-            // Convert tensor to bytes based on its dtype
-            let data: Vec<u8> = match dtype {
-                DType::F32 => {
-                    let flat = tensor.flatten_all()?;
-                    let vec: Vec<f32> = flat.to_vec1()?;
-                    bytemuck::cast_slice(&vec).to_vec()
-                }
-                DType::F16 => {
-                    let flat = tensor.flatten_all()?;
-                    let vec: Vec<half::f16> = flat.to_vec1()?;
-                    bytemuck::cast_slice(&vec).to_vec()
-                }
-                DType::BF16 => {
-                    let flat = tensor.flatten_all()?;
-                    let vec: Vec<half::bf16> = flat.to_vec1()?;
-                    bytemuck::cast_slice(&vec).to_vec()
-                }
-                _ => {
-                    return Err(candle_core::Error::Msg(format!(
-                        "Unsupported dtype: {:?}",
-                        dtype
-                    )))
-                }
-            };
-            println!("  -> data len: {}", data.len());
-            tensor_data.push((name.clone(), data, shape, dtype));
+            let (data, st_dtype) = crate::training::train::tensor_to_bytes(tensor)?;
+            tensor_data.push((name.clone(), data, shape, st_dtype));
         }
 
         let mut tensors = HashMap::new();
-        for (name, data, shape, dtype) in &tensor_data {
-            let st_dtype = match dtype {
-                DType::F32 => safetensors::Dtype::F32,
-                DType::F16 => safetensors::Dtype::F16,
-                DType::BF16 => safetensors::Dtype::BF16,
-                _ => {
-                    return Err(candle_core::Error::Msg(format!(
-                        "Unsupported dtype: {:?}",
-                        dtype
-                    )))
-                }
-            };
-            println!(
-                "Creating TensorView for {}: shape={:?}, data_len={}",
-                name,
-                shape,
-                data.len()
-            );
+        for (name, data, shape, st_dtype) in &tensor_data {
             tensors.insert(
                 name.clone(),
-                safetensors::tensor::TensorView::new(st_dtype, shape.to_vec(), data)?,
+                safetensors::tensor::TensorView::new(*st_dtype, shape.to_vec(), data)?,
             );
         }
 
diff --git a/src/layers/attention.rs b/src/layers/attention.rs
index aaf7e96..84c603e 100644
--- a/src/layers/attention.rs
+++ b/src/layers/attention.rs
@@ -9,15 +9,16 @@ pub struct MultiHeadAttention {
 }
 
 impl MultiHeadAttention {
-    pub fn new_blank(hidden_dim: usize, num_heads: usize, device: &Device) -> Result<Self> {
+    pub fn new_blank(
+        hidden_dim: usize,
+        num_heads: usize,
+        dtype: DType,
+        device: &Device,
+    ) -> Result<Self> {
         assert_eq!(hidden_dim % num_heads, 0);
         let head_dim = hidden_dim / num_heads;
-        let qkv_weight = Tensor::zeros(
-            (hidden_dim, hidden_dim * 3),
-            candle_core::DType::F32,
-            device,
-        )?;
-        let out_weight = Tensor::zeros((hidden_dim, hidden_dim), candle_core::DType::F32, device)?;
+        let qkv_weight = Tensor::zeros((hidden_dim, hidden_dim * 3), dtype, device)?;
+        let out_weight = Tensor::zeros((hidden_dim, hidden_dim), dtype, device)?;
         let scale = 1.0 / (head_dim as f64).sqrt();
         Ok(Self {
             qkv_weight,
diff --git a/src/layers/embedding.rs b/src/layers/embedding.rs
index e58f6e5..1fb4c1c 100644
--- a/src/layers/embedding.rs
+++ b/src/layers/embedding.rs
@@ -1,12 +1,17 @@
-use candle_core::{Device, Result, Tensor};
+use candle_core::{DType, Device, Result, Tensor};
 
 pub struct Embedding {
     pub weight: Tensor,
 }
 
 impl Embedding {
-    pub fn zeros(vocab_size: usize, hidden_dim: usize, device: &Device) -> Result<Self> {
-        let weight = Tensor::zeros((vocab_size, hidden_dim), candle_core::DType::F32, device)?;
+    pub fn zeros(
+        vocab_size: usize,
+        hidden_dim: usize,
+        dtype: DType,
+        device: &Device,
+    ) -> Result<Self> {
+        let weight = Tensor::zeros((vocab_size, hidden_dim), dtype, device)?;
         Ok(Self { weight })
     }
 
@@ -16,11 +21,8 @@ impl Embedding {
     }
 
     pub fn forward(&self, ids: &[u32]) -> Result<Tensor> {
-        let mut rows = Vec::new();
-        for &id in ids {
-            rows.push(self.weight.get(id as usize)?);
-        }
-        Tensor::stack(&rows, 0)?.unsqueeze(0)
+        let ids = Tensor::from_slice(ids, ids.len(), self.weight.device())?;
+        self.weight.index_select(&ids, 0)?.unsqueeze(0)
     }
 
     pub fn from_tensor(weight: Tensor) -> Self {
diff --git a/src/layers/ffn.rs b/src/layers/ffn.rs
index a981e2d..971543a 100644
--- a/src/layers/ffn.rs
+++ b/src/layers/ffn.rs
@@ -1,4 +1,4 @@
-use candle_core::{Device, Result, Tensor};
+use candle_core::{DType, Device, Result, Tensor};
 
 #[allow(dead_code)]
 #[allow(clippy::upper_case_acronyms)]
@@ -37,10 +37,15 @@ pub struct SwiGLU {
 }
 
 impl SwiGLU {
-    pub fn new_blank(hidden_dim: usize, ffn_dim: usize, device: &Device) -> Result<Self> {
-        let gate = Tensor::zeros((hidden_dim, ffn_dim), candle_core::DType::F32, device)?;
-        let up = Tensor::zeros((hidden_dim, ffn_dim), candle_core::DType::F32, device)?;
-        let down = Tensor::zeros((ffn_dim, hidden_dim), candle_core::DType::F32, device)?;
+    pub fn new_blank(
+        hidden_dim: usize,
+        ffn_dim: usize,
+        dtype: DType,
+        device: &Device,
+    ) -> Result<Self> {
+        let gate = Tensor::zeros((hidden_dim, ffn_dim), dtype, device)?;
+        let up = Tensor::zeros((hidden_dim, ffn_dim), dtype, device)?;
+        let down = Tensor::zeros((ffn_dim, hidden_dim), dtype, device)?;
         Ok(Self { gate, up, down })
     }
 
@@ -68,9 +73,14 @@ pub struct GELUFFN {
 }
 
 impl GELUFFN {
-    pub fn new_blank(hidden_dim: usize, ffn_dim: usize, device: &Device) -> Result<Self> {
-        let fc_in = Tensor::zeros((hidden_dim, ffn_dim), candle_core::DType::F32, device)?;
-        let fc_out = Tensor::zeros((ffn_dim, hidden_dim), candle_core::DType::F32, device)?;
+    pub fn new_blank(
+        hidden_dim: usize,
+        ffn_dim: usize,
+        dtype: DType,
+        device: &Device,
+    ) -> Result<Self> {
+        let fc_in = Tensor::zeros((hidden_dim, ffn_dim), dtype, device)?;
+        let fc_out = Tensor::zeros((ffn_dim, hidden_dim), dtype, device)?;
         Ok(Self { fc_in, fc_out })
     }
 
@@ -101,11 +111,16 @@ pub struct GELUNewFFN {
 }
 
 impl GELUNewFFN {
-    pub fn new_blank(hidden_dim: usize, ffn_dim: usize, device: &Device) -> Result<Self> {
-        let fc_in = Tensor::zeros((hidden_dim, ffn_dim), candle_core::DType::F32, device)?;
-        let fc_in_bias = Tensor::zeros(ffn_dim, candle_core::DType::F32, device)?;
-        let fc_out = Tensor::zeros((ffn_dim, hidden_dim), candle_core::DType::F32, device)?;
-        let fc_out_bias = Tensor::zeros(hidden_dim, candle_core::DType::F32, device)?;
+    pub fn new_blank(
+        hidden_dim: usize,
+        ffn_dim: usize,
+        dtype: DType,
+        device: &Device,
+    ) -> Result<Self> {
+        let fc_in = Tensor::zeros((hidden_dim, ffn_dim), dtype, device)?;
+        let fc_in_bias = Tensor::zeros(ffn_dim, dtype, device)?;
+        let fc_out = Tensor::zeros((ffn_dim, hidden_dim), dtype, device)?;
+        let fc_out_bias = Tensor::zeros(hidden_dim, dtype, device)?;
         Ok(Self {
             fc_in,
             fc_in_bias,
@@ -166,15 +181,18 @@ impl FeedForward {
         activation: Activation,
         hidden_dim: usize,
         ffn_dim: usize,
+        dtype: DType,
         device: &Device,
     ) -> Result<Self> {
         match activation {
             Activation::SwiGLU => Ok(Self::SwiGLU(SwiGLU::new_blank(
-                hidden_dim, ffn_dim, device,
+                hidden_dim, ffn_dim, dtype, device,
+            )?)),
+            Activation::GELU => Ok(Self::GELU(GELUFFN::new_blank(
+                hidden_dim, ffn_dim, dtype, device,
             )?)),
-            Activation::GELU => Ok(Self::GELU(GELUFFN::new_blank(hidden_dim, ffn_dim, device)?)),
             Activation::GELUNew => Ok(Self::GELUNew(GELUNewFFN::new_blank(
-                hidden_dim, ffn_dim, device,
+                hidden_dim, ffn_dim, dtype, device,
             )?)),
         }
     }
@@ -189,7 +207,10 @@ impl FeedForward {
             Activation::SwiGLU => Ok(Self::SwiGLU(SwiGLU::new(hidden_dim, ffn_dim, device)?)),
             Activation::GELU => Ok(Self::GELU(GELUFFN::new(hidden_dim, ffn_dim, device)?)),
             Activation::GELUNew => Ok(Self::GELUNew(GELUNewFFN::new_blank(
-                hidden_dim, ffn_dim, device,
+                hidden_dim,
+                ffn_dim,
+                DType::F32,
+                device,
             )?)),
         }
     }
diff --git a/src/lib.rs b/src/lib.rs
index a94207b..2d3135f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,7 +3,8 @@
 //! This crate provides:
 //! - GLM (General Language Model) architecture for infilling
 //! - CodeGen-350M architecture for code generation
-//! - INT8 dynamic quantization support
+//! - FP16 inference
+//! - CLI subcommand implementations
 //! - HTTP server for model serving (feature-gated)
 //! - Training pipeline with YAML config
 //!
@@ -13,6 +14,7 @@
 
 pub mod cli;
 pub mod codegen;
+pub mod commands;
 pub mod generation;
 pub mod glm;
 pub mod layers;
diff --git a/src/main.rs b/src/main.rs
index f798180..369ace3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,19 +1,10 @@
-mod cli;
-mod commands;
-mod model;
-
-mod codegen;
-mod generation;
-mod glm;
-mod layers;
-#[cfg(feature = "server")]
-mod server;
-mod tokenizer;
-mod training;
+//! CLI entry point. Everything it drives lives in the library crate, so the
+//! binary does not compile a second copy of the module tree.
 
 use clap::Parser;
 
-use cli::{Cli, Commands};
+use rust_transformer::cli::{Cli, Commands};
+use rust_transformer::commands;
 
 fn main() -> anyhow::Result<()> {
     let cli = Cli::parse();
@@ -25,14 +16,23 @@ fn main() -> anyhow::Result<()> {
             max_tokens,
             temperature,
             template,
-            stream,
-        } => commands::complete::run(&cli, prompt, *max_tokens, *temperature, template, *stream),
+            no_stream,
+        } => commands::complete::run(
+            &cli,
+            prompt,
+            *max_tokens,
+            *temperature,
+            template,
+            !*no_stream,
+        ),
         Commands::Repl => commands::repl::run(&cli),
         Commands::Info => commands::info::run(&cli),
         Commands::Serve { port } => commands::serve::run(&cli, *port),
         Commands::Download => commands::download::run(&cli),
-        Commands::GlmTrain { data_path, steps } => {
-            commands::glm_train::run(&cli, data_path, *steps)
-        }
+        Commands::GlmTrain {
+            data_path,
+            steps,
+            config,
+        } => commands::glm_train::run(&cli, data_path, *steps, config.as_deref()),
     }
 }
diff --git a/src/model.rs b/src/model.rs
index f1368ec..eea333b 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -23,19 +23,24 @@ pub struct ModelContext {
 impl ModelContext {
     /// Load CodeGen-350M weights, tokenizer, and build generator.
     pub fn load(weights_dir: &Path, use_f16: bool, temperature: f64) -> Result<Self> {
-        let weights_path = weights_dir.join("pytorch_model.bin");
         let config_path = weights_dir.join("config.json");
         let tokenizer_path = weights_dir.join("tokenizer.json");
 
-        if !weights_path.exists() {
+        // safetensors loads faster and needs no pickle, so prefer it when present.
+        let safetensors_path = weights_dir.join("model.safetensors");
+        let pytorch_path = weights_dir.join("pytorch_model.bin");
+        let weights_path = if safetensors_path.exists() {
+            safetensors_path
+        } else if pytorch_path.exists() {
+            pytorch_path
+        } else {
             bail!(
-                "Weights not found at: {}\n\
+                "Weights not found in {dir}: expected model.safetensors or pytorch_model.bin\n\
                  Run `codegen download` or:\n  \
-                 huggingface-cli download Salesforce/codegen-350M-multi --local-dir {}",
-                weights_path.display(),
-                weights_dir.display()
+                 huggingface-cli download Salesforce/codegen-350M-multi --local-dir {dir}",
+                dir = weights_dir.display()
             );
-        }
+        };
 
         let mut config = if config_path.exists() {
             let config_str = std::fs::read_to_string(&config_path)?;
@@ -52,9 +57,12 @@ impl ModelContext {
         let device = Device::Cpu;
 
         let tokenizer = CodeGenTokenizer::from_file(tokenizer_path.to_str().unwrap())?;
-        let model = WeightLoader::load_from_pytorch(&weights_path, &config, &device)?;
+        let model = WeightLoader::load(&weights_path, &config, &device)?;
 
-        let generator = CodeGenGenerator::new(model, temperature, 40, 0.9, 1.2, 256);
+        // Without this the generator has no tokenizer, so the streaming callback
+        // receives an empty string for every token and `complete` prints nothing.
+        let generator = CodeGenGenerator::new(model, temperature, 40, 0.9, 1.2, 256)
+            .with_tokenizer(tokenizer.clone());
 
         Ok(Self {
             generator,
diff --git a/src/server.rs b/src/server.rs
index d9f628b..11b3199 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -9,6 +9,7 @@
 //! - `POST /generate` — Generate code from prompt
 
 use std::net::SocketAddr;
+use std::path::Path;
 use std::sync::Arc;
 
 use axum::extract::State;
@@ -18,9 +19,8 @@ use axum::{Json, Router};
 use serde::{Deserialize, Serialize};
 use tokio::sync::Mutex;
 
-use crate::codegen::config::CodeGenConfig;
-use crate::codegen::weights::WeightLoader;
 use crate::generation::codegen_generate::CodeGenGenerator;
+use crate::model::ModelContext;
 use crate::tokenizer::CodeGenTokenizer;
 
 pub struct AppState {
@@ -67,28 +67,20 @@ pub struct HealthResponse {
     pub model: String,
 }
 
-pub async fn start_server(weights_path: &str, addr: SocketAddr) -> anyhow::Result<()> {
-    let device = candle_core::Device::Cpu;
-    let config = CodeGenConfig::default();
-
-    println!("Loading CodeGen model from {weights_path}...");
-    let model =
-        WeightLoader::load_from_pytorch(std::path::Path::new(weights_path), &config, &device)?;
-
-    let tokenizer = CodeGenTokenizer::from_file("codegen_weights/tokenizer.json")
-        .map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {e}"))?;
-
-    let generator = CodeGenGenerator::new(
-        model, 0.8, // temperature
-        40,  // top_k
-        0.9, // top_p
-        1.1, // repetition_penalty
-        256, // max_new_tokens
-    );
+pub async fn start_server(
+    weights_dir: &Path,
+    use_f16: bool,
+    seed: Option<u64>,
+    addr: SocketAddr,
+) -> anyhow::Result<()> {
+    println!("Loading CodeGen model from {}...", weights_dir.display());
+    // Reuse the CLI loader so the server honours config.json and --f16 too.
+    let mut ctx = ModelContext::load(weights_dir, use_f16, 0.8)?;
+    ctx.generator.set_seed(seed);
 
     let state = Arc::new(AppState {
-        generator: Arc::new(Mutex::new(generator)),
-        tokenizer,
+        generator: Arc::new(Mutex::new(ctx.generator)),
+        tokenizer: ctx.tokenizer,
     });
 
     let app = Router::new()
@@ -132,17 +124,15 @@ async fn generate(
         )
     })?;
 
-    let generated = state
-        .tokenizer
-        .decode(&tokens[prompt_ids.len()..])
-        .map_err(|e| {
-            (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                format!("Decode error: {e}"),
-            )
-        })?;
-
-    let token_count = tokens.len() - prompt_ids.len();
+    // `generate` returns the generated tokens only — the prompt is not included.
+    let generated = state.tokenizer.decode(&tokens).map_err(|e| {
+        (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            format!("Decode error: {e}"),
+        )
+    })?;
+
+    let token_count = tokens.len();
 
     Ok(Json(GenerateResponse {
         generated,
diff --git a/src/tokenizer.rs b/src/tokenizer.rs
index 4a6e758..0505d7a 100644
--- a/src/tokenizer.rs
+++ b/src/tokenizer.rs
@@ -1,6 +1,7 @@
 use anyhow::Result;
 use tokenizers::Tokenizer as HFTokenizer;
 
+#[derive(Clone)]
 pub struct CodeGenTokenizer {
     tokenizer: HFTokenizer,
 }
@@ -19,6 +20,16 @@ impl CodeGenTokenizer {
         Ok(encoded.get_ids().to_vec())
     }
 
+    /// Number of ids the tokenizer can actually decode, including added tokens.
+    ///
+    /// CodeGen's `config.json` declares `vocab_size: 51200`, but the tokenizer tops
+    /// out at id 50294 — the embedding and `lm_head` are padded to a round number.
+    /// Those padding rows are untrained, so their logits are arbitrary and sampling
+    /// can land on an id that decodes to nothing.
+    pub fn vocab_size(&self) -> usize {
+        self.tokenizer.get_vocab_size(true)
+    }
+
     pub fn decode(&self, ids: &[u32]) -> Result<String> {
         self.tokenizer
             .decode(ids, false)
@@ -59,6 +70,27 @@ mod tests {
         assert!(text.is_empty());
     }
 
+    /// The model's declared `vocab_size` is padded well beyond what the tokenizer
+    /// can decode, which is why generation trims logits before sampling.
+    #[test]
+    fn test_vocab_size_is_below_the_models_padded_vocab() {
+        if !Path::new("codegen_weights/tokenizer.json").exists() {
+            eprintln!("Skipping: tokenizer not found");
+            return;
+        }
+        let tok = CodeGenTokenizer::from_file("codegen_weights/tokenizer.json").unwrap();
+        let config: serde_json::Value =
+            serde_json::from_str(&std::fs::read_to_string("codegen_weights/config.json").unwrap())
+                .unwrap();
+        let declared = config["vocab_size"].as_u64().unwrap() as usize;
+
+        assert!(
+            tok.vocab_size() < declared,
+            "expected padding: tokenizer {} vs config {declared}",
+            tok.vocab_size()
+        );
+    }
+
     #[test]
     fn test_encode_special_characters() {
         if !Path::new("codegen_weights/tokenizer.json").exists() {
diff --git a/src/training/checkpoint.rs b/src/training/checkpoint.rs
deleted file mode 100644
index 497e337..0000000
--- a/src/training/checkpoint.rs
+++ /dev/null
@@ -1,193 +0,0 @@
-use std::path::{Path, PathBuf};
-use candle_core::{Device, Result, Tensor, DType};
-use candle_nn::Var;
-use safetensors::tensor::TensorView;
-use safetensors::serialize;
-
-pub struct CheckpointManager {
-    checkpoint_dir: PathBuf,
-    save_optimizer_state: bool,
-    keep_last_n: usize,
-}
-
-impl CheckpointManager {
-    pub fn new<P: AsRef<Path>>(checkpoint_dir: P, save_optimizer_state: bool, keep_last_n: usize) -> Self {
-        Self {
-            checkpoint_dir: checkpoint_dir.as_ref().to_path_buf(),
-            save_optimizer_state,
-            keep_last_n,
-        }
-    }
-
-    pub fn save_model(
-        &self,
-        step: usize,
-        model_vars: &[(&str, &Var)],
-    ) -> Result<()> {
-        std::fs::create_dir_all(&self.checkpoint_dir)
-            .map_err(|e| candle_core::Error::Msg(format!("Failed to create checkpoint dir: {e}")))?;
-
-        let model_path = self.checkpoint_dir.join(format!("model_step_{:07}.safetensors", step));
-        let mut tensors = Vec::new();
-
-        for (name, var) in model_vars {
-            let tensor = var.as_tensor();
-            let tensor_view = TensorView::new(tensor.dtype(), tensor.shape(), tensor.as_slice()?)?;
-            tensors.push((name.to_string(), tensor_view));
-        }
-
-        let bytes = serialize(&tensors, &None)
-            .map_err(|e| candle_core::Error::Msg(format!("Failed to serialize: {e}")))?;
-
-        std::fs::write(&model_path, bytes)
-            .map_err(|e| candle_core::Error::Msg(format!("Failed to write model: {e}")))?;
-
-        // Save step metadata
-        let step_path = self.checkpoint_dir.join(format!("step_{:07}.txt", step));
-        std::fs::write(&step_path, step.to_string())
-            .map_err(|e| candle_core::Error::Msg(format!("Failed to write step: {e}")))?;
-
-        // Clean old checkpoints
-        self.cleanup_old_checkpoints(step)?;
-
-        println!("  Checkpoint saved: {:?} (step {})", model_path, step);
-        Ok(())
-    }
-
-    pub fn save_optimizer_state(
-        &self,
-        step: usize,
-        optimizer_state: &[(&str, Vec<Tensor>)],
-    ) -> Result<()> {
-        if !self.save_optimizer_state {
-            return Ok(());
-        }
-
-        let opt_path = self.checkpoint_dir.join(format!("optimizer_step_{:07}.safetensors", step));
-        let mut tensors = Vec::new();
-
-        for (name, state_tensors) in optimizer_state {
-            for (i, tensor) in state_tensors.iter().enumerate() {
-                let tensor_name = format!("{}_{}", name, i);
-                let tensor_view = TensorView::new(tensor.dtype(), tensor.shape(), tensor.as_slice()?)?;
-                tensors.push((tensor_name, tensor_view));
-            }
-        }
-
-        let bytes = serialize(&tensors, &None)
-            .map_err(|e| candle_core::Error::Msg(format!("Failed to serialize optimizer: {e}")))?;
-
-        std::fs::write(&opt_path, bytes)
-            .map_err(|e| candle_core::Error::Msg(format!("Failed to write optimizer: {e}")))?;
-
-        Ok(())
-    }
-
-    pub fn load_model<P: AsRef<Path>>(
-        &self,
-        path: P,
-        model_vars: &mut [(&str, &mut Var)],
-        device: &Device,
-    ) -> Result<()> {
-        let data = std::fs::read(path)
-            .map_err(|e| candle_core::Error::Msg(format!("Failed to read checkpoint: {e}")))?;
-
-        let tensors = safetensors::SafeTensors::deserialize(&data)
-            .map_err(|e| candle_core::Error::Msg(format!("Failed to deserialize: {e}")))?;
-
-        for (name, var) in model_vars {
-            if let Ok(tensor_view) = tensors.tensor(name) {
-                let tensor = Tensor::from_slice(
-                    tensor_view.data(),
-                    tensor_view.shape(),
-                    device,
-                )?;
-                var.set(&tensor)?;
-            }
-        }
-
-        Ok(())
-    }
-
-    pub fn load_optimizer_state<P: AsRef<Path>>(
-        &self,
-        path: P,
-        optimizer_vars: &mut [(&str, &mut [Tensor])],
-    ) -> Result<()> {
-        let data = std::fs::read(path)
-            .map_err(|e| candle_core::Error::Msg(format!("Failed to read optimizer: {e}")))?;
-
-        let tensors = safetensors::SafeTensors::deserialize(&data)
-            .map_err(|e| candle_core::Error::Msg(format!("Failed to deserialize optimizer: {e}")))?;
-
-        for (prefix, var_tensors) in optimizer_vars {
-            for (i, var) in var_tensors.iter_mut().enumerate() {
-                let name = format!("{}_{}", prefix, i);
-                if let Ok(tensor_view) = tensors.tensor(&name) {
-                    let tensor = Tensor::from_slice(
-                        tensor_view.data(),
-                        tensor_view.shape(),
-                        var.device(),
-                    )?;
-                    *var = tensor;
-                }
-            }
-        }
-
-        Ok(())
-    }
-
-    pub fn find_latest_checkpoint(&self) -> Option<(usize, PathBuf)> {
-        let mut checkpoints = Vec::new();
-
-        if let Ok(entries) = std::fs::read_dir(&self.checkpoint_dir) {
-            for entry in entries.flatten() {
-                let path = entry.path();
-                if let Some(name) = path.file_name().and_then(|s| s.to_str()) {
-                    if name.starts_with("model_step_") && name.ends_with(".safetensors") {
-                        if let Some(step_str) = name.strip_prefix("model_step_").and_then(|s| s.strip_suffix(".safetensors")) {
-                            if let Ok(step) = step_str.parse::<usize>() {
-                                checkpoints.push((step, path));
-                            }
-                        }
-                    }
-                }
-            }
-        }
-
-        checkpoints.sort_by_key(|&(step, _)| step);
-        checkpoints.into_iter().next_back()
-    }
-
-    fn cleanup_old_checkpoints(&self, current_step: usize) -> Result<()> {
-        let mut checkpoints = Vec::new();
-
-        if let Ok(entries) = std::fs::read_dir(&self.checkpoint_dir) {
-            for entry in entries.flatten() {
-                let path = entry.path();
-                if let Some(name) = path.file_name().and_then(|s| s.to_str()) {
-                    if name.starts_with("model_step_") && name.ends_with(".safetensors") {
-                        if let Some(step_str) = name.strip_prefix("model_step_").and_then(|s| s.strip_suffix(".safetensors")) {
-                            if let Ok(step) = step_str.parse::<usize>() {
-                                checkpoints.push((step, path));
-                            }
-                        }
-                    }
-                }
-            }
-        }
-
-        checkpoints.sort_by_key(|&(step, _)| step);
-
-        if checkpoints.len() > self.keep_last_n {
-            for (step, path) in checkpoints.iter().take(checkpoints.len() - self.keep_last_n) {
-                let _ = std::fs::remove_file(path);
-                // Also remove corresponding step file and optimizer file
-                let _ = std::fs::remove_file(self.checkpoint_dir.join(format!("step_{:07}.txt", step)));
-                let _ = std::fs::remove_file(self.checkpoint_dir.join(format!("optimizer_step_{:07}.safetensors", step)));
-            }
-        }
-
-        Ok(())
-    }
-}
\ No newline at end of file
diff --git a/src/training/config.rs b/src/training/config.rs
index d1680dc..c804cba 100644
--- a/src/training/config.rs
+++ b/src/training/config.rs
@@ -2,12 +2,14 @@ use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
 
 #[derive(Debug, Clone, Default, Serialize, Deserialize)]
+#[serde(default)]
 pub struct TrainConfig {
     pub model: ModelConfig,
     pub training: TrainingConfig,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(default)]
 pub struct ModelConfig {
     pub vocab_size: usize,
     pub hidden_dim: usize,
@@ -41,6 +43,7 @@ impl Default for ModelConfig {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(default)]
 pub struct TrainingConfig {
     pub data_dir: PathBuf,
     pub download_if_empty: bool,
@@ -110,6 +113,7 @@ impl TrainingConfig {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(default)]
 pub struct LrScheduleConfig {
     #[serde(rename = "type")]
     pub schedule_type: String,
@@ -250,4 +254,28 @@ mod tests {
         assert_eq!(config.model.hidden_dim, cloned.model.hidden_dim);
         assert_eq!(config.training.max_steps, cloned.training.max_steps);
     }
+
+    /// The config the repo ships must actually deserialize. It did not: `eval_steps`
+    /// lived under a separate `evaluation:` section and `tokenizer_path` was absent,
+    /// and nothing loaded the file because `glm-train` had no `--config` flag.
+    #[test]
+    fn shipped_config_loads() {
+        let path = concat!(env!("CARGO_MANIFEST_DIR"), "/configs/train.yaml");
+        let config = TrainConfig::from_file(path).expect("configs/train.yaml must parse");
+        assert_eq!(config.model.hidden_dim, 256);
+        assert_eq!(config.training.gradient_accumulation_steps, 32);
+        assert_eq!(config.training.max_grad_norm, 1.0);
+    }
+
+    /// A config naming only a couple of fields should fall back to defaults.
+    #[test]
+    fn partial_config_falls_back_to_defaults() {
+        let yaml = "training:\n  learning_rate: 0.5\n";
+        let config: TrainConfig = serde_yaml::from_str(yaml).expect("partial config must parse");
+        assert_eq!(config.training.learning_rate, 0.5);
+        assert_eq!(
+            config.training.micro_batch_size,
+            TrainingConfig::default().micro_batch_size
+        );
+    }
 }
diff --git a/src/training/loss.rs b/src/training/loss.rs
new file mode 100644
index 0000000..4e39d08
--- /dev/null
+++ b/src/training/loss.rs
@@ -0,0 +1,194 @@
+//! Cross-entropy loss for GLM training.
+
+use candle_core::{Result, Tensor};
+
+/// Mean cross-entropy over the positions carrying a label.
+///
+/// `logits` is `[1, seq_len, vocab_size]`; `labels[i]` is the target token for
+/// position `i`, or `-1` for positions that should not contribute. Returns a
+/// scalar with no labelled positions, in which case there is nothing to learn
+/// from and the caller should skip the optimizer step.
+///
+/// This is built from tensor operations on purpose. The previous version summed
+/// `ce.to_scalar::<f32>()` into an `f32` and returned `Tensor::new(total, device)`
+/// — a fresh leaf with no autograd history — so `loss.backward()` produced an
+/// empty gradient store and every training step was a no-op.
+pub fn cross_entropy_loss(logits: &Tensor, labels: &[i64]) -> Result<Tensor> {
+    let device = logits.device();
+    let seq_len = logits.dim(1)?;
+
+    let positions: Vec<u32> = labels
+        .iter()
+        .take(seq_len)
+        .enumerate()
+        .filter(|(_, &label)| label >= 0)
+        .map(|(i, _)| i as u32)
+        .collect();
+    let targets: Vec<u32> = labels
+        .iter()
+        .take(seq_len)
+        .filter(|&&label| label >= 0)
+        .map(|&label| label as u32)
+        .collect();
+
+    if positions.is_empty() {
+        return Tensor::new(0.0f32, device);
+    }
+
+    let n = positions.len();
+    let index = Tensor::from_vec(positions, n, device)?;
+    let selected = logits.squeeze(0)?.index_select(&index, 0)?;
+    let targets = Tensor::from_vec(targets, n, device)?;
+
+    candle_nn::loss::cross_entropy(&selected.to_dtype(candle_core::DType::F32)?, &targets)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::{DType, Device};
+    use candle_nn::{AdamW, Optimizer, ParamsAdamW};
+
+    use crate::glm::config::GLMConfig;
+    use crate::glm::trainable::TrainableGLMModel;
+
+    fn tiny_config() -> GLMConfig {
+        GLMConfig {
+            vocab_size: 64,
+            hidden_dim: 32,
+            num_layers: 2,
+            num_heads: 4,
+            ffn_dim: 64,
+            max_seq_len: 16,
+            ..Default::default()
+        }
+    }
+
+    /// A masked-denoising batch in the shape the trainer produces: some positions
+    /// replaced by the mask token, labels holding the originals.
+    fn corrupted_batch(config: &GLMConfig) -> (Vec<u32>, Vec<i64>) {
+        let tokens: Vec<u32> = vec![3, 9, 14, 2, 41, 7, 33, 5];
+        let mask_id = config.vocab_size as u32 - 1;
+        let mut inputs = tokens.clone();
+        let mut labels = vec![-1i64; tokens.len()];
+        for i in [2usize, 5, 7] {
+            labels[i] = tokens[i] as i64;
+            inputs[i] = mask_id;
+        }
+        (inputs, labels)
+    }
+
+    #[test]
+    fn no_labelled_positions_is_zero() -> Result<()> {
+        let device = Device::Cpu;
+        let logits = Tensor::randn(0.0f32, 1.0f32, (1, 4, 10), &device)?;
+        let loss = cross_entropy_loss(&logits, &[-1, -1, -1, -1])?;
+        assert_eq!(loss.to_scalar::<f32>()?, 0.0);
+        Ok(())
+    }
+
+    #[test]
+    fn matches_hand_computed_log_softmax() -> Result<()> {
+        let device = Device::Cpu;
+        // One position, one clear winner: loss is -log_softmax(logits)[label].
+        let logits = Tensor::from_vec(vec![1.0f32, 2.0, 3.0], (1, 1, 3), &device)?;
+        let loss = cross_entropy_loss(&logits, &[2])?.to_scalar::<f32>()?;
+
+        let denom: f32 = [1.0f32, 2.0, 3.0].iter().map(|v| (v - 3.0f32).exp()).sum();
+        let expected = -((3.0f32 - 3.0) - denom.ln());
+        assert!(
+            (loss - expected).abs() < 1e-5,
+            "loss {loss} vs expected {expected}"
+        );
+        Ok(())
+    }
+
+    /// The direct assertion for the detached-loss bug: backprop must reach the
+    /// model's parameters. On the previous implementation the gradient store came
+    /// back empty.
+    #[test]
+    fn gradients_reach_the_parameters() -> Result<()> {
+        let device = Device::Cpu;
+        let config = tiny_config();
+        let model = TrainableGLMModel::new(config.clone(), &device)?;
+        let (inputs, labels) = corrupted_batch(&config);
+
+        let logits = model.forward_causal(&inputs)?;
+        let loss = cross_entropy_loss(&logits, &labels)?;
+        let grads = loss.backward()?;
+
+        let params = model.param_vars();
+        let with_grad = params.iter().filter(|v| grads.get(v).is_some()).count();
+        assert!(
+            with_grad >= params.len() - 1,
+            "only {with_grad} of {} parameters received gradients",
+            params.len()
+        );
+        Ok(())
+    }
+
+    /// End-to-end proof that optimisation works: a two-layer model memorises one
+    /// batch. Fails flat if gradients stop flowing for any reason.
+    #[test]
+    fn training_reduces_loss() -> Result<()> {
+        let device = Device::Cpu;
+        let config = tiny_config();
+        let model = TrainableGLMModel::new(config.clone(), &device)?;
+        let (inputs, labels) = corrupted_batch(&config);
+
+        let mut optimizer = AdamW::new(
+            model.param_vars(),
+            ParamsAdamW {
+                lr: 0.05,
+                ..Default::default()
+            },
+        )?;
+
+        let first =
+            cross_entropy_loss(&model.forward_causal(&inputs)?, &labels)?.to_scalar::<f32>()?;
+
+        let mut last = first;
+        for _ in 0..60 {
+            let loss = cross_entropy_loss(&model.forward_causal(&inputs)?, &labels)?;
+            last = loss.to_scalar::<f32>()?;
+            optimizer.backward_step(&loss)?;
+        }
+
+        assert!(
+            last < first * 0.3,
+            "loss went {first} -> {last}; the model is not learning"
+        );
+        Ok(())
+    }
+
+    /// Parameters must actually change value, not merely receive gradients.
+    #[test]
+    fn optimizer_updates_parameters() -> Result<()> {
+        let device = Device::Cpu;
+        let config = tiny_config();
+        let model = TrainableGLMModel::new(config.clone(), &device)?;
+        let (inputs, labels) = corrupted_batch(&config);
+
+        let before = model.lm_head.as_tensor().copy()?;
+        let mut optimizer = AdamW::new(
+            model.param_vars(),
+            ParamsAdamW {
+                lr: 0.05,
+                ..Default::default()
+            },
+        )?;
+        let loss = cross_entropy_loss(&model.forward_causal(&inputs)?, &labels)?;
+        optimizer.backward_step(&loss)?;
+
+        let delta = (model.lm_head.as_tensor() - &before)?
+            .abs()?
+            .max_all()?
+            .to_dtype(DType::F32)?
+            .to_scalar::<f32>()?;
+        assert!(
+            delta > 0.0,
+            "lm_head did not change after an optimizer step"
+        );
+        Ok(())
+    }
+}
diff --git a/src/training/mod.rs b/src/training/mod.rs
index 1271940..692428c 100644
--- a/src/training/mod.rs
+++ b/src/training/mod.rs
@@ -4,11 +4,12 @@
 //! - YAML-based configuration ([`config::TrainConfig`])
 //! - Data loading with train/eval split ([`data::DataLoader`])
 //! - Learning rate scheduling with warmup ([`lr_scheduler::LrScheduler`])
-//! - Gradient accumulation and clipping
+//! - Gradient accumulation and clipping ([`train::clip_grad_norm`])
 //! - Safetensors checkpoint save/load
 
 pub mod config;
 pub mod data;
+pub mod loss;
 pub mod lr_scheduler;
 pub mod train;
 
diff --git a/src/training/train.rs b/src/training/train.rs
index 70a6ced..716a511 100644
--- a/src/training/train.rs
+++ b/src/training/train.rs
@@ -2,9 +2,12 @@ use std::collections::VecDeque;
 use std::path::{Path, PathBuf};
 
 use bytemuck;
-use candle_core::{DType, Device, Result, Tensor};
+use candle_core::backprop::GradStore;
+use candle_core::{DType, Device, Result, Tensor, Var};
 use candle_nn::{AdamW, Optimizer, ParamsAdamW};
 use half;
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
 use safetensors::tensor::TensorView;
 use safetensors::SafeTensors;
 
@@ -15,6 +18,7 @@ use crate::training::config::{TrainConfig, TrainingConfig};
 use crate::training::data::{
     download_default_data, split_train_eval, DataLoader, TrainingExample as DataTrainingExample,
 };
+use crate::training::loss::cross_entropy_loss;
 use crate::training::lr_scheduler::LrScheduler;
 
 #[allow(dead_code)]
@@ -30,6 +34,28 @@ pub struct GLMTrainer {
     tokenizer: CodeGenTokenizer,
     dtype: DType,
     loss_history: VecDeque<f64>,
+    rng: StdRng,
+}
+
+/// Corrupt a sequence for masked denoising: some positions are replaced by the
+/// mask token or a random token, and their originals become the labels.
+/// Positions left alone get `-1`, which [`cross_entropy_loss`] ignores.
+fn corrupt(config: &GLMConfig, tokens: &[u32], rng: &mut impl Rng) -> (Vec<u32>, Vec<i64>) {
+    let mask_token_id = config.vocab_size as u32 - 1;
+    let mut inputs = tokens.to_vec();
+    let mut labels = vec![-1i64; tokens.len()];
+
+    for i in 0..tokens.len() {
+        if rng.gen::<f64>() < config.blank_ratio {
+            labels[i] = tokens[i] as i64;
+            inputs[i] = if rng.gen::<f64>() < config.mask_ratio {
+                mask_token_id
+            } else {
+                rng.gen_range(0..config.vocab_size as u32)
+            };
+        }
+    }
+    (inputs, labels)
 }
 
 impl GLMTrainer {
@@ -74,48 +100,66 @@ impl GLMTrainer {
             tokenizer,
             dtype,
             loss_history: VecDeque::with_capacity(100),
+            rng: StdRng::seed_from_u64(tc.seed),
         })
     }
 
-    pub fn train_step(&mut self, token_ids: &[u32], _device: &Device) -> Result<f64> {
-        let n = token_ids.len();
-        let mask_token_id = self.glm_config.vocab_size as u32 - 1;
+    /// Corrupt, forward, loss. The returned tensor stays connected to the
+    /// autograd graph — the caller owns the optimizer step.
+    pub fn compute_loss(&mut self, token_ids: &[u32]) -> Result<Tensor> {
+        let (input_ids, labels) = corrupt(&self.glm_config, token_ids, &mut self.rng);
+        let logits = self.model.forward_causal(&input_ids)?;
+        cross_entropy_loss(&logits, &labels)
+    }
 
-        let mut rng_state = self.step as u64;
-        let mut rand = || -> f64 {
-            rng_state = rng_state.wrapping_mul(6364136223846793005).wrapping_add(1);
-            (rng_state as f64) / (u64::MAX as f64)
-        };
+    /// Backward, clip by global norm, step. Advances the step counter and the
+    /// learning-rate schedule.
+    fn apply_gradients(&mut self, loss: &Tensor) -> Result<()> {
+        let mut grads = loss.backward()?;
+        let params = self.model.param_vars();
 
-        let mut input_ids = token_ids.to_vec();
-        let mut labels = vec![-1i64; n];
-
-        for i in 0..n {
-            let r = rand();
-            if r < self.glm_config.blank_ratio {
-                labels[i] = token_ids[i] as i64;
-                if rand() < self.glm_config.mask_ratio {
-                    input_ids[i] = mask_token_id;
-                } else {
-                    input_ids[i] = (rand() * self.glm_config.vocab_size as f64) as u32;
-                }
-            }
+        if self.config.max_grad_norm > 0.0 {
+            clip_grad_norm(&mut grads, &params, self.config.max_grad_norm)?;
         }
 
-        let logits = self.model.forward_causal(&input_ids)?;
-        let loss = cross_entropy_loss(&logits, &labels)?;
+        // Set the rate this step will use, then step — the previous code set it
+        // afterwards, so every step ran on the preceding step's rate.
+        self.optimizer.set_learning_rate(self.lr_scheduler.get_lr());
+        self.optimizer.step(&grads)?;
 
-        let loss_scalar = loss.to_scalar::<f32>()? as f64;
+        self.step += 1;
+        self.lr_scheduler.step();
+        Ok(())
+    }
 
-        // backward_step does backward + step in one call
-        self.optimizer.backward_step(&loss)?;
+    /// One optimizer step over `gradient_accumulation_steps` micro-batches of
+    /// `micro_batch_size` sequences each.
+    ///
+    /// Averaging the micro-batch losses and running a single backward is
+    /// equivalent to accumulating their gradients, and needs no `GradStore`
+    /// bookkeeping.
+    pub fn train_step(&mut self, loader: &mut DataLoader, _device: &Device) -> Result<f64> {
+        let accum_steps = self.config.gradient_accumulation_steps.max(1);
+        let mut losses: Vec<Tensor> = Vec::new();
+
+        for _ in 0..accum_steps {
+            for tokens in loader.next_batch(self.config.micro_batch_size) {
+                losses.push(self.compute_loss(&tokens)?);
+            }
+        }
 
-        // Update learning rate for next step
-        let current_lr = self.lr_scheduler.get_lr();
-        self.optimizer.set_learning_rate(current_lr);
+        if losses.is_empty() {
+            return Ok(0.0);
+        }
 
-        self.step += 1;
-        self.lr_scheduler.step();
+        let mut total = losses[0].clone();
+        for loss in &losses[1..] {
+            total = (total + loss)?;
+        }
+        let loss = (total / losses.len() as f64)?;
+        let loss_scalar = loss.to_scalar::<f32>()? as f64;
+
+        self.apply_gradients(&loss)?;
 
         self.loss_history.push_back(loss_scalar);
         if self.loss_history.len() > 100 {
@@ -175,10 +219,8 @@ impl GLMTrainer {
 
         // Training loop
         while self.step < self.config.max_steps {
-            let batch = train_loader.next_batch(self.config.micro_batch_size);
-
-            for tokens in batch {
-                let loss = self.train_step(&tokens, device)?;
+            {
+                let loss = self.train_step(&mut train_loader, device)?;
 
                 // Logging
                 if self.step % self.config.log_every == 0 {
@@ -224,29 +266,9 @@ impl GLMTrainer {
         for _ in 0..self.config.eval_steps.min(eval_loader.len()) {
             let batch = eval_loader.next_batch(self.config.micro_batch_size);
             for tokens in batch {
-                let n = tokens.len();
-                let mask_token_id = self.glm_config.vocab_size as u32 - 1;
-
-                let mut input_ids = tokens.clone();
-                let mut labels = vec![-1i64; n];
-
-                // Use fixed seed for reproducible eval
-                let mut rng_state = 12345 + self.step as u64;
-                let mut rand = || -> f64 {
-                    rng_state = rng_state.wrapping_mul(6364136223846793005).wrapping_add(1);
-                    (rng_state as f64) / (u64::MAX as f64)
-                };
-
-                for i in 0..n {
-                    if rand() < self.glm_config.blank_ratio {
-                        labels[i] = tokens[i] as i64;
-                        if rand() < self.glm_config.mask_ratio {
-                            input_ids[i] = mask_token_id;
-                        } else {
-                            input_ids[i] = (rand() * self.glm_config.vocab_size as f64) as u32;
-                        }
-                    }
-                }
+                // Fixed seed so eval loss is comparable across runs.
+                let mut eval_rng = StdRng::seed_from_u64(12345);
+                let (input_ids, labels) = corrupt(&self.glm_config, &tokens, &mut eval_rng);
 
                 let logits = self.model.forward_causal(&input_ids)?;
                 let loss = cross_entropy_loss(&logits, &labels)?;
@@ -410,27 +432,35 @@ impl GLMTrainer {
     }
 }
 
-fn cross_entropy_loss(logits: &Tensor, labels: &[i64]) -> Result<Tensor> {
-    let seq_len = logits.dims()[1];
-    let mut total_loss = 0.0f32;
-    let mut count = 0usize;
-
-    for (i, &label) in labels.iter().enumerate().take(seq_len) {
-        if label >= 0 {
-            let logits_i = logits.get(0)?.get(i)?;
-            let ce = candle_nn::ops::log_softmax(&logits_i, 0)?
-                .get(label as usize)?
-                .neg()?;
-            total_loss += ce.to_scalar::<f32>()?;
-            count += 1;
+/// Scale every gradient so the global L2 norm across `params` is at most
+/// `max_norm`. Returns the norm before clipping.
+pub fn grad_global_norm(grads: &GradStore, params: &[Var]) -> Result<f64> {
+    let mut sum_sq = 0f64;
+    for var in params {
+        if let Some(grad) = grads.get(var) {
+            sum_sq += grad
+                .sqr()?
+                .sum_all()?
+                .to_dtype(DType::F32)?
+                .to_scalar::<f32>()? as f64;
         }
     }
+    Ok(sum_sq.sqrt())
+}
 
-    if count > 0 {
-        total_loss /= count as f32;
+/// Clip gradients in place by global norm. Returns the norm before clipping.
+pub fn clip_grad_norm(grads: &mut GradStore, params: &[Var], max_norm: f64) -> Result<f64> {
+    let norm = grad_global_norm(grads, params)?;
+    if norm.is_finite() && norm > max_norm {
+        let scale = max_norm / norm;
+        for var in params {
+            let scaled = grads.get(var).map(|grad| grad * scale).transpose()?;
+            if let Some(scaled) = scaled {
+                grads.insert(var, scaled);
+            }
+        }
     }
-
-    Tensor::new(total_loss, logits.device())
+    Ok(norm)
 }
 
 fn param_names(num_layers: usize) -> Vec<String> {
@@ -452,6 +482,41 @@ fn param_names(num_layers: usize) -> Vec<String> {
     names
 }
 
+/// Flatten a tensor into safetensors bytes.
+///
+/// The previous version called `tensor.to_vec1::<u8>()`, which fails on anything
+/// of rank > 1 — so every checkpoint save errored with
+/// "unexpected rank, expected: 1, got: 2" the moment it reached a weight matrix.
+pub fn tensor_to_bytes(tensor: &Tensor) -> Result<(Vec<u8>, safetensors::Dtype)> {
+    let flat = tensor.flatten_all()?;
+    match tensor.dtype() {
+        DType::F32 => {
+            let values: Vec<f32> = flat.to_vec1()?;
+            Ok((
+                bytemuck::cast_slice(&values).to_vec(),
+                safetensors::Dtype::F32,
+            ))
+        }
+        DType::F16 => {
+            let values: Vec<half::f16> = flat.to_vec1()?;
+            Ok((
+                bytemuck::cast_slice(&values).to_vec(),
+                safetensors::Dtype::F16,
+            ))
+        }
+        DType::BF16 => {
+            let values: Vec<half::bf16> = flat.to_vec1()?;
+            Ok((
+                bytemuck::cast_slice(&values).to_vec(),
+                safetensors::Dtype::BF16,
+            ))
+        }
+        dtype => Err(candle_core::Error::Msg(format!(
+            "unsupported checkpoint dtype: {dtype:?}"
+        ))),
+    }
+}
+
 fn save_safetensors(path: &Path, tensors: &[(String, &Tensor)]) -> Result<()> {
     use safetensors::serialize;
     use std::collections::HashMap;
@@ -460,20 +525,8 @@ fn save_safetensors(path: &Path, tensors: &[(String, &Tensor)]) -> Result<()> {
     let mut tensor_data = Vec::new();
 
     for (name, tensor) in tensors {
-        let data = tensor.to_vec1::<u8>()?;
+        let (data, st_dtype) = tensor_to_bytes(tensor)?;
         let shape = tensor.shape().dims().to_vec();
-        let dtype = tensor.dtype();
-        let st_dtype = match dtype {
-            DType::F32 => safetensors::Dtype::F32,
-            DType::F16 => safetensors::Dtype::F16,
-            DType::BF16 => safetensors::Dtype::BF16,
-            _ => {
-                return Err(candle_core::Error::Msg(format!(
-                    "Unsupported dtype: {:?}",
-                    dtype
-                )))
-            }
-        };
         tensor_data.push((name.clone(), data, shape, st_dtype));
     }
 
@@ -563,3 +616,70 @@ fn load_data(data_dir: &Path, tokenizer: &CodeGenTokenizer) -> Result<Vec<DataTr
 
     Ok(examples)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::training::loss::cross_entropy_loss;
+
+    #[test]
+    fn clipping_bounds_the_global_norm() -> Result<()> {
+        let device = Device::Cpu;
+        let config = GLMConfig {
+            vocab_size: 64,
+            hidden_dim: 32,
+            num_layers: 2,
+            num_heads: 4,
+            ffn_dim: 64,
+            max_seq_len: 16,
+            ..Default::default()
+        };
+        let model = TrainableGLMModel::new(config, &device)?;
+        let params = model.param_vars();
+
+        let logits = model.forward_causal(&[3u32, 9, 14, 2, 41])?;
+        let loss = cross_entropy_loss(&logits, &[1i64, 2, 3, 4, 5])?;
+        let mut grads = loss.backward()?;
+
+        let before = grad_global_norm(&grads, &params)?;
+        assert!(before > 0.0, "no gradients to clip");
+
+        let max_norm = before / 10.0;
+        let reported = clip_grad_norm(&mut grads, &params, max_norm)?;
+        assert!((reported - before).abs() < 1e-6);
+
+        let after = grad_global_norm(&grads, &params)?;
+        assert!(
+            (after - max_norm).abs() < 1e-4,
+            "clipped norm {after} should sit at the bound {max_norm}"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn clipping_leaves_small_gradients_alone() -> Result<()> {
+        let device = Device::Cpu;
+        let config = GLMConfig {
+            vocab_size: 64,
+            hidden_dim: 32,
+            num_layers: 1,
+            num_heads: 4,
+            ffn_dim: 64,
+            max_seq_len: 16,
+            ..Default::default()
+        };
+        let model = TrainableGLMModel::new(config, &device)?;
+        let params = model.param_vars();
+
+        let logits = model.forward_causal(&[3u32, 9, 14])?;
+        let loss = cross_entropy_loss(&logits, &[1i64, 2, 3])?;
+        let mut grads = loss.backward()?;
+
+        let before = grad_global_norm(&grads, &params)?;
+        clip_grad_norm(&mut grads, &params, before * 10.0)?;
+        let after = grad_global_norm(&grads, &params)?;
+
+        assert!((after - before).abs() < 1e-6, "{before} changed to {after}");
+        Ok(())
+    }
+}
diff --git a/tests/codegen_integration.rs b/tests/codegen_integration.rs
index e604948..969931c 100644
--- a/tests/codegen_integration.rs
+++ b/tests/codegen_integration.rs
@@ -5,8 +5,34 @@
 
 use std::path::Path;
 
+use rust_transformer::codegen::config::CodeGenConfig;
+use rust_transformer::codegen::weights::WeightLoader;
+use rust_transformer::generation::codegen_generate::CodeGenGenerator;
+use rust_transformer::tokenizer::CodeGenTokenizer;
+
 fn weights_available() -> bool {
     Path::new("codegen_weights/pytorch_model.bin").exists()
+        || Path::new("codegen_weights/model.safetensors").exists()
+}
+
+fn weights_path() -> &'static Path {
+    if Path::new("codegen_weights/model.safetensors").exists() {
+        Path::new("codegen_weights/model.safetensors")
+    } else {
+        Path::new("codegen_weights/pytorch_model.bin")
+    }
+}
+
+/// The checkpoint's own config — the library defaults differ from it.
+fn real_config() -> CodeGenConfig {
+    let json: serde_json::Value =
+        serde_json::from_str(&std::fs::read_to_string("codegen_weights/config.json").unwrap())
+            .expect("config.json is not valid JSON");
+    CodeGenConfig::from_hf_config(&json)
+}
+
+fn real_tokenizer() -> CodeGenTokenizer {
+    CodeGenTokenizer::from_file("codegen_weights/tokenizer.json").expect("Failed to load tokenizer")
 }
 
 #[test]
@@ -17,7 +43,12 @@ fn codegen_forward_pass_with_real_weights() {
     }
 
     let device = candle_core::Device::Cpu;
-    let config = rust_transformer::codegen::config::CodeGenConfig::default();
+    // The checkpoint's own config.json — the defaults differ from it
+    // (rotary_dim 64 vs 32, vocab_size 50400 vs 51200).
+    let config_json: serde_json::Value =
+        serde_json::from_str(&std::fs::read_to_string("codegen_weights/config.json").unwrap())
+            .expect("config.json is not valid JSON");
+    let config = rust_transformer::codegen::config::CodeGenConfig::from_hf_config(&config_json);
     let model = rust_transformer::codegen::weights::WeightLoader::load_from_pytorch(
         Path::new("codegen_weights/pytorch_model.bin"),
         &config,
@@ -36,15 +67,13 @@ fn codegen_forward_pass_with_real_weights() {
         model, 0.0, 1, 1.0, 1.0, 64,
     );
 
+    // `generate` returns the generated tokens only, not prompt + generated.
     let generated = gen.generate(&token_ids).expect("Generation failed");
     let output = tokenizer.decode(&generated).expect("Decode failed");
 
     println!("Prompt: {prompt}");
     println!("Generated: {output}");
-    assert!(
-        generated.len() > token_ids.len(),
-        "Should generate at least one token"
-    );
+    assert!(!generated.is_empty(), "Should generate at least one token");
 }
 
 #[test]
@@ -68,3 +97,76 @@ fn codegen_tokenizer_encode_decode_roundtrip() {
     assert!(!ids.is_empty(), "Should produce tokens");
     assert!(!decoded.is_empty(), "Should decode back to text");
 }
+
+/// Sampling must never emit an id the tokenizer cannot decode.
+///
+/// `config.json` declares `vocab_size: 51200` while the tokenizer stops at 50294, so
+/// 905 untrained rows sit at the top of `lm_head`. Their logits turn out to be low
+/// enough that they never win — verified at temperature 2.0 with no top-k or nucleus
+/// cut — so generation does not mask them. This test is what keeps that true.
+#[test]
+fn sampling_stays_within_the_tokenizer_vocabulary() {
+    if !weights_available() {
+        eprintln!("Skipping: weights not found");
+        return;
+    }
+    let device = candle_core::Device::Cpu;
+    let config = real_config();
+    let tokenizer = real_tokenizer();
+    let model = WeightLoader::load(weights_path(), &config, &device).expect("Failed to load model");
+
+    assert!(
+        tokenizer.vocab_size() < config.vocab_size,
+        "this test is pointless unless the model vocabulary is padded"
+    );
+
+    let prompt = tokenizer
+        .encode("def fibonacci(n):")
+        .expect("encode failed");
+    let mut generator =
+        CodeGenGenerator::new(model, 0.9, 40, 0.95, 1.1, 32).with_tokenizer(tokenizer.clone());
+    generator.set_seed(Some(7));
+
+    let generated = generator.generate(&prompt).expect("generation failed");
+    let out_of_range: Vec<u32> = generated
+        .iter()
+        .copied()
+        .filter(|&t| (t as usize) >= tokenizer.vocab_size())
+        .collect();
+    assert!(
+        out_of_range.is_empty(),
+        "sampled ids outside the tokenizer vocabulary: {out_of_range:?}"
+    );
+
+    let text = tokenizer.decode(&generated).expect("decode failed");
+    println!("Sampled: {text}");
+    assert!(!text.trim().is_empty(), "sampled output decoded to nothing");
+}
+
+/// The sampler used to be greedy no matter what, so every run matched. Now a fixed
+/// seed must reproduce and a different seed must diverge.
+#[test]
+fn sampling_is_reproducible_and_seed_dependent() {
+    if !weights_available() {
+        eprintln!("Skipping: weights not found");
+        return;
+    }
+    let device = candle_core::Device::Cpu;
+    let config = real_config();
+    let tokenizer = real_tokenizer();
+    let prompt = tokenizer
+        .encode("def quicksort(arr):")
+        .expect("encode failed");
+
+    let run = |seed: u64| -> Vec<u32> {
+        let model =
+            WeightLoader::load(weights_path(), &config, &device).expect("Failed to load model");
+        let mut generator =
+            CodeGenGenerator::new(model, 0.9, 40, 0.95, 1.1, 16).with_tokenizer(tokenizer.clone());
+        generator.set_seed(Some(seed));
+        generator.generate(&prompt).expect("generation failed")
+    };
+
+    assert_eq!(run(1), run(1), "same seed must reproduce");
+    assert_ne!(run(1), run(2), "different seeds must diverge");
+}
diff --git a/tests/codegen_parity.rs b/tests/codegen_parity.rs
new file mode 100644
index 0000000..fd588ea
--- /dev/null
+++ b/tests/codegen_parity.rs
@@ -0,0 +1,277 @@
+//! Numerical parity against the HuggingFace CodeGen reference implementation.
+//!
+//! Fixtures are tiny randomly-initialised models generated by
+//! `scripts/gen_parity_fixture.py` and committed under `tests/fixtures/`, so this
+//! test needs no model download and no Python.
+
+use std::path::Path;
+
+use candle_core::Device;
+use rust_transformer::codegen::config::CodeGenConfig;
+use rust_transformer::codegen::kv_cache::KVCache;
+use rust_transformer::codegen::model::CodeGenModel;
+use rust_transformer::codegen::weights::WeightLoader;
+use rust_transformer::generation::codegen_generate::CodeGenGenerator;
+
+struct Fixture {
+    config: CodeGenConfig,
+    model: CodeGenModel,
+    tokens: Vec<u32>,
+    reference: Vec<f32>,
+    shape: Vec<usize>,
+}
+
+fn load(name: &str) -> Fixture {
+    let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
+    let device = Device::Cpu;
+
+    let config_json: serde_json::Value = serde_json::from_str(
+        &std::fs::read_to_string(dir.join(format!("{name}_config.json")))
+            .expect("fixture config missing — run scripts/gen_parity_fixture.py"),
+    )
+    .expect("fixture config is not valid JSON");
+    let config = CodeGenConfig::from_hf_config(&config_json);
+
+    let model = WeightLoader::load_from_pytorch(&dir.join(format!("{name}.pth")), &config, &device)
+        .expect("failed to load fixture weights");
+
+    let logits_json: serde_json::Value = serde_json::from_str(
+        &std::fs::read_to_string(dir.join(format!("{name}_logits.json")))
+            .expect("fixture logits missing"),
+    )
+    .expect("fixture logits are not valid JSON");
+
+    let tokens: Vec<u32> = logits_json["tokens"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|v| v.as_u64().unwrap() as u32)
+        .collect();
+    let reference: Vec<f32> = logits_json["logits"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|v| v.as_f64().unwrap() as f32)
+        .collect();
+    let shape: Vec<usize> = logits_json["shape"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|v| v.as_u64().unwrap() as usize)
+        .collect();
+
+    Fixture {
+        config,
+        model,
+        tokens,
+        reference,
+        shape,
+    }
+}
+
+fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
+    assert_eq!(a.len(), b.len(), "logit count mismatch");
+    a.iter()
+        .zip(b)
+        .map(|(x, y)| (x - y).abs())
+        .fold(0.0f32, f32::max)
+}
+
+fn assert_prefill_matches_reference(name: &str) {
+    let f = load(name);
+    let positions: Vec<usize> = (0..f.tokens.len()).collect();
+
+    let logits = f
+        .model
+        .forward_with_cache(&f.tokens, &positions, &mut None)
+        .expect("forward pass failed");
+
+    assert_eq!(logits.dims(), f.shape.as_slice(), "logit shape mismatch");
+
+    let ours: Vec<f32> = logits
+        .flatten_all()
+        .unwrap()
+        .to_vec1::<f32>()
+        .expect("logits are not f32");
+
+    let diff = max_abs_diff(&ours, &f.reference);
+    assert!(
+        diff < 1e-5,
+        "{name}: max abs logit difference vs HuggingFace reference is {diff}, expected < 1e-5"
+    );
+}
+
+#[test]
+fn parity_prefill_one_head_per_group() {
+    assert_prefill_matches_reference("tiny_h4");
+}
+
+#[test]
+fn parity_prefill_multiple_heads_per_group() {
+    assert_prefill_matches_reference("tiny_h8");
+}
+
+/// Feeding tokens one at a time through the KV cache must reproduce the logits
+/// of a single-shot prefill over the same tokens.
+#[test]
+fn parity_incremental_decode_matches_prefill() {
+    let f = load("tiny_h4");
+    let positions: Vec<usize> = (0..f.tokens.len()).collect();
+
+    let prefill = f
+        .model
+        .forward_with_cache(&f.tokens, &positions, &mut None)
+        .expect("prefill failed");
+    let last = f.tokens.len() - 1;
+    let prefill_last: Vec<f32> = prefill
+        .get(0)
+        .unwrap()
+        .get(last)
+        .unwrap()
+        .to_vec1::<f32>()
+        .unwrap();
+
+    let mut cache: Option<Vec<KVCache>> = None;
+    let mut step_last = Vec::new();
+    for (i, &token) in f.tokens.iter().enumerate() {
+        let logits = f
+            .model
+            .forward_with_cache(&[token], &[i], &mut cache)
+            .expect("incremental step failed");
+        step_last = logits
+            .get(0)
+            .unwrap()
+            .get(0)
+            .unwrap()
+            .to_vec1::<f32>()
+            .unwrap();
+    }
+
+    let diff = max_abs_diff(&step_last, &prefill_last);
+    assert!(
+        diff < 1e-5,
+        "incremental decode diverges from prefill by {diff}, expected < 1e-5"
+    );
+
+    // Guard against the config being silently wrong in a way both paths share.
+    assert_eq!(f.config.vocab_size, prefill_last.len());
+}
+
+/// `generate` returns generated tokens only. The `complete`, `repl` and server
+/// paths slice on that assumption, so pin it here: they used to strip
+/// `prompt.len()` tokens off the front, dropping output or panicking outright.
+#[test]
+fn generate_returns_generated_tokens_only() {
+    let f = load("tiny_h4");
+    let prompt = f.tokens.clone();
+    let max_new_tokens = 5;
+
+    // Greedy so the result does not depend on the RNG. The fixture vocabulary is
+    // 256 tokens, so the real EOS id (50256) can never be produced and exactly
+    // `max_new_tokens` tokens come back.
+    let gen = CodeGenGenerator::new(f.model, 0.0, 1, 1.0, 1.0, max_new_tokens);
+    let generated = gen.generate(&prompt).expect("generation failed");
+
+    assert_eq!(
+        generated.len(),
+        max_new_tokens,
+        "expected generated tokens only, got {} for a {}-token prompt",
+        generated.len(),
+        prompt.len()
+    );
+    assert!(
+        generated
+            .iter()
+            .all(|t| (*t as usize) < f.config.vocab_size),
+        "generated an out-of-vocabulary token: {generated:?}"
+    );
+}
+
+/// The safetensors converter's output must load back to the same model.
+/// Before `WeightLoader::load_from_safetensors` existed, nothing in the repo
+/// could read what `examples/convert_codegen_to_safetensors.rs` wrote.
+#[test]
+fn safetensors_round_trip_matches_pytorch() {
+    let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
+    let device = Device::Cpu;
+
+    let f = load("tiny_h4");
+    let positions: Vec<usize> = (0..f.tokens.len()).collect();
+    let from_pth = f
+        .model
+        .forward_with_cache(&f.tokens, &positions, &mut None)
+        .unwrap()
+        .flatten_all()
+        .unwrap()
+        .to_vec1::<f32>()
+        .unwrap();
+
+    // Same conversion the example performs: read every tensor, write safetensors.
+    let pth = candle_core::pickle::PthTensors::new(dir.join("tiny_h4.pth"), None).unwrap();
+    let names: Vec<String> = pth.tensor_infos().keys().cloned().collect();
+    let tensors: std::collections::HashMap<String, candle_core::Tensor> = names
+        .iter()
+        .filter_map(|name| pth.get(name).ok().flatten().map(|t| (name.clone(), t)))
+        .collect();
+    assert!(!tensors.is_empty(), "no tensors read from the fixture");
+
+    let tmp = tempfile::tempdir().unwrap();
+    let out = tmp.path().join("model.safetensors");
+    candle_core::safetensors::save(&tensors, &out).unwrap();
+
+    let model = WeightLoader::load_from_safetensors(&out, &f.config, &device)
+        .expect("failed to load converted safetensors");
+    let from_st = model
+        .forward_with_cache(&f.tokens, &positions, &mut None)
+        .unwrap()
+        .flatten_all()
+        .unwrap()
+        .to_vec1::<f32>()
+        .unwrap();
+
+    let diff = max_abs_diff(&from_st, &from_pth);
+    assert!(
+        diff < 1e-6,
+        "safetensors round trip changed the logits by {diff}"
+    );
+}
+
+/// The streaming callback must receive decoded text, not empty strings.
+///
+/// `CodeGenGenerator::with_tokenizer` existed but was never called, so
+/// `decode_token` always returned `""` and `complete` — which streams by default —
+/// printed nothing at all.
+#[test]
+fn streaming_handler_receives_decoded_text() {
+    let tokenizer_path = "codegen_weights/tokenizer.json";
+    if !Path::new(tokenizer_path).exists() {
+        eprintln!("Skipping: tokenizer not found");
+        return;
+    }
+    let tokenizer =
+        rust_transformer::tokenizer::CodeGenTokenizer::from_file(tokenizer_path).unwrap();
+
+    let f = load("tiny_h4");
+    let gen = CodeGenGenerator::new(f.model, 0.0, 1, 1.0, 1.0, 4).with_tokenizer(tokenizer);
+
+    #[derive(Default)]
+    struct Recorder {
+        texts: Vec<String>,
+    }
+    impl rust_transformer::generation::codegen_generate::StreamHandler for Recorder {
+        fn on_token(&mut self, _token: u32, text: &str) -> bool {
+            self.texts.push(text.to_string());
+            true
+        }
+    }
+
+    let mut recorder = Recorder::default();
+    gen.generate_stream(&f.tokens, &mut recorder)
+        .expect("generation failed");
+
+    assert_eq!(recorder.texts.len(), 4, "expected one callback per token");
+    assert!(
+        recorder.texts.iter().any(|t| !t.is_empty()),
+        "every streamed token decoded to an empty string"
+    );
+}
diff --git a/tests/fixtures/tiny_h4.pth b/tests/fixtures/tiny_h4.pth
new file mode 100644
index 0000000..1e79624
Binary files /dev/null and b/tests/fixtures/tiny_h4.pth differ
diff --git a/tests/fixtures/tiny_h4_config.json b/tests/fixtures/tiny_h4_config.json
new file mode 100644
index 0000000..490b60c
--- /dev/null
+++ b/tests/fixtures/tiny_h4_config.json
@@ -0,0 +1,39 @@
+{
+  "_name_or_path": "",
+  "activation_function": "gelu_new",
+  "architectures": null,
+  "attn_pdrop": 0.0,
+  "bos_token_id": 50256,
+  "chunk_size_feed_forward": 0,
+  "dtype": null,
+  "embd_pdrop": 0.0,
+  "eos_token_id": 50256,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "initializer_range": 0.02,
+  "is_encoder_decoder": false,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "codegen",
+  "n_ctx": 64,
+  "n_embd": 128,
+  "n_head": 4,
+  "n_inner": null,
+  "n_layer": 2,
+  "n_positions": 64,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "problem_type": null,
+  "resid_pdrop": 0.0,
+  "return_dict": true,
+  "rotary_dim": 16,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.12.1",
+  "use_cache": true,
+  "vocab_size": 256
+}
diff --git a/tests/fixtures/tiny_h4_logits.json b/tests/fixtures/tiny_h4_logits.json
new file mode 100644
index 0000000..235c22b
--- /dev/null
+++ b/tests/fixtures/tiny_h4_logits.json
@@ -0,0 +1 @@
+{"tokens": [3, 17, 42, 8, 255, 1, 99, 128], "shape": [1, 8, 256], "logits": [0.13797208666801453, -0.12160780280828476, 0.015172661282122135, 0.12064576894044876, 0.05379147827625275, -0.29649820923805237, 0.2050786018371582, -0.1565551459789276, -0.06716244667768478, 0.03846554830670357, 0.3188575804233551, -0.17997929453849792, -0.1908494383096695, -0.04417635127902031, -0.08735319972038269, 0.40719372034072876, -0.24718672037124634, 0.14152072370052338, -0.12805403769016266, -0.06504087150096893, -0.09366442263126373, -0.18055735528469086, -0.2347811460494995, 0.29062095284461975, 0.04116421565413475, 0.06698659062385559, 0.13029389083385468, 0.03848961740732193, 0.2765219509601593, -0.2865062355995178, 0.20257458090782166, 0.1723998636007309, 0.2137833684682846, 0.4326269030570984, -0.0009396881214343011, 0.17215093970298767, 0.2375749796628952, -0.11660584807395935, 0.04120151698589325, -0.3389374613761902, -0.16921766102313995, -0.09879470616579056, -0.18560680747032166, 0.33959048986434937, 0.10422768443822861, -0.08966687321662903, -0.6132300496101379, 0.255328893661499, 0.16417236626148224, -0.24560371041297913, 0.07819544523954391, -0.10255435854196548, -0.4769432842731476, 0.033296238631010056, 0.3392610549926758, 0.2763991057872772, -0.1838977336883545, 0.20095357298851013, 0.1466151475906372, -0.19600743055343628, 0.06724176555871964, -0.041446734219789505, -0.2139338105916977, -0.13107256591320038, 0.1255294233560562, -0.15353916585445404, -0.1614842414855957, 0.05491883307695389, -0.08805611729621887, 0.30660784244537354, -0.2726972699165344, -0.22136205434799194, 0.017901619896292686, -0.2540002763271332, 0.12028636783361435, 0.055389653891325, -0.051570210605859756, 0.0692351832985878, -0.1784837543964386, 0.023299047723412514, 0.20437733829021454, 0.2109411507844925, -0.06377265602350235, 0.04263145849108696, 0.39125245809555054, 0.13542288541793823, 0.1764322817325592, 0.28708139061927795, 0.3067520260810852, -0.024627884849905968, 0.08944395929574966, 0.3394125998020172, 0.022080564871430397, 0.2146754115819931, 0.1351061314344406, 0.14937809109687805, -0.4144514799118042, 0.25638189911842346, -0.16860802471637726, 0.019572222605347633, 0.508650541305542, -0.1423056721687317, -0.010113397613167763, -0.20181059837341309, 0.08867400884628296, 0.06273669749498367, 0.23533694446086884, 0.14392374455928802, -0.04992162808775902, 0.08846477419137955, -0.05817421153187752, 0.26681283116340637, -0.27640360593795776, -0.29602283239364624, 0.02165238745510578, -0.008114050142467022, -0.28540971875190735, 0.019992755725979805, -0.3857240378856659, -0.21671807765960693, -0.007429229561239481, 0.2322276085615158, 0.12162502110004425, 0.09285692870616913, -0.3684569299221039, -0.06472666561603546, -0.003905214834958315, -0.26022034883499146, -0.011866025626659393, 0.02485169656574726, -0.15326546132564545, -0.2529875934123993, 0.11364205181598663, -0.08141594380140305, 0.07734422385692596, 0.052367035299539566, 0.10160461068153381, 0.11050158739089966, 0.0920712798833847, -0.28216618299484253, 0.2772896885871887, 0.271312415599823, 0.19137859344482422, -0.5577556490898132, -0.6188551187515259, -0.5914639234542847, 0.1950666904449463, -0.010510235093533993, 0.004505499731749296, 0.1058182567358017, -0.15907728672027588, -0.11431925743818283, -0.47599345445632935, 0.07211202383041382, 0.27881208062171936, -0.4005069434642792, -0.03024817816913128, -0.14754502475261688, -0.24671459197998047, -0.16086988151073456, -0.709753692150116, -0.35539573431015015, -0.13390493392944336, 0.0928003191947937, 0.04863446578383446, -0.0495501346886158, 0.18421092629432678, -0.018080364912748337, 0.5899195075035095, 0.365138977766037, 0.1963970959186554, 0.1972842514514923, -0.08086569607257843, 0.2867452800273895, 0.36433976888656616, -0.0034080499317497015, 0.3998028039932251, 0.1984650045633316, -0.002680029021576047, -0.06290017068386078, -0.026031747460365295, -0.10381772369146347, 0.13147704303264618, 0.12688888609409332, 0.027697419747710228, -0.25628265738487244, -0.09270302951335907, -0.30371055006980896, -0.17885839939117432, 0.2448159009218216, 0.019905803725123405, 0.31289583444595337, 0.6091428399085999, -0.29281556606292725, -0.09114749729633331, -0.20409734547138214, -0.10191241651773453, 0.07138940691947937, 0.34490522742271423, -0.25056400895118713, 0.12000268697738647, 0.0363888256251812, 0.15915948152542114, 0.3538452088832855, 0.21529997885227203, -0.053232014179229736, 0.5123010873794556, -0.3488830626010895, -0.16908413171768188, -0.22125312685966492, 0.14842158555984497, 0.11073613166809082, 0.021522222086787224, -0.19074460864067078, 0.2579523026943207, 0.3705143332481384, 0.15849003195762634, 0.16991882026195526, 0.06226116791367531, -0.06926238536834717, -0.4051671028137207, 0.31475120782852173, -0.2134520709514618, 0.21438013017177582, -0.02561035566031933, 0.11234936118125916, -0.016122084110975266, -0.4536351263523102, -0.1443435549736023, -0.05539524927735329, -0.17995263636112213, -0.11866766214370728, 0.09042748808860779, 0.05701519921422005, 0.2888790965080261, -0.023825109004974365, -0.06706663221120834, 0.34762290120124817, -0.23093216121196747, 0.0027742611709982157, -0.13100560009479523, -0.1632840782403946, -0.2922150492668152, 0.1936868578195572, -0.30890393257141113, 0.17172551155090332, 0.36000964045524597, -0.1943063735961914, 0.0905841737985611, -0.19335758686065674, -0.17201192677021027, 0.03831101208925247, -0.026263123378157616, 0.23113606870174408, -0.3485087752342224, 0.09443095326423645, 0.11230451613664627, -0.1538846641778946, 0.6527838706970215, -0.07330699265003204, 0.6585939526557922, 0.18166093528270721, 0.22973082959651947, 0.4143810272216797, 0.08408821374177933, -0.18594226241111755, -0.19325894117355347, -0.12203656136989594, 0.31579774618148804, -0.21571511030197144, -0.18997636437416077, 0.46451535820961, 0.09896256774663925, -0.3819522261619568, -0.057064007967710495, 0.3411034345626831, 0.10054958611726761, 0.05066530033946037, -0.35839149355888367, -0.053074367344379425, -0.3858056366443634, 0.09979639202356339, 0.08741068840026855, -0.06904573738574982, 0.10780099034309387, 0.3676793575286865, -0.06574448198080063, 0.133367657661438, 0.011892307549715042, -0.009637637995183468, 0.16948889195919037, 0.18666408956050873, -0.07560979574918747, -0.20428653061389923, -0.13375520706176758, 0.2661072611808777, 0.16493961215019226, -0.03417146950960159, -0.16584378480911255, -0.09299695491790771, 0.1512826681137085, 0.28490209579467773, -0.2359611541032791, -0.5328914523124695, -0.275350958108902, 0.11128020286560059, -0.11670579761266708, 0.13909827172756195, -0.4446939527988434, -0.143888458609581, 0.35516414046287537, -0.08010520786046982, -0.013807199895381927, -0.29741424322128296, 0.5260536670684814, -0.044274020940065384, 0.24142663180828094, -0.08907772600650787, -0.29276710748672485, 0.32396334409713745, -0.021903546527028084, -0.26519834995269775, 0.0856281965970993, -0.1632746011018753, -0.044272441416978836, -0.10231183469295502, 0.028264617547392845, 0.11770310997962952, 0.17787589132785797, 0.06981819868087769, -0.13995875418186188, -0.0372747927904129, -0.1708967238664627, -0.09157034754753113, -0.20730772614479065, -0.17059366405010223, -0.34517452120780945, 0.08556120097637177, 0.4133048951625824, 0.10640161484479904, -0.21350108087062836, -0.14047075808048248, 0.28017252683639526, -0.1867925077676773, -0.2767457962036133, 0.42247089743614197, -0.011200746521353722, -0.07419400662183762, 0.03608120605349541, -0.17860639095306396, 0.04689581319689751, -0.20695346593856812, -0.013390627689659595, -0.18870525062084198, -0.08851992338895798, -0.2981280982494354, -0.0807793140411377, -0.04121050983667374, 0.06908270716667175, -0.0409683920443058, 0.12120222300291061, 0.034107089042663574, 0.18572326004505157, 0.590471088886261, -0.45534199476242065, -0.0801481232047081, -0.39978504180908203, -0.06894753128290176, 0.1731927990913391, -0.0093894824385643, 0.09983590245246887, 0.13170002400875092, -0.18126973509788513, -0.17244108021259308, -0.43173032999038696, 0.029740437865257263, 0.08199063688516617, -0.18150891363620758, -0.14963634312152863, -0.3662071228027344, 0.08197434991598129, 0.051793552935123444, 0.20771971344947815, -0.0920766294002533, 0.1959647834300995, 0.32539546489715576, 0.17856252193450928, -0.6734200716018677, -0.641075074672699, 0.4333253502845764, 0.2678319215774536, -0.05155530199408531, -0.02104790136218071, 0.3848993480205536, 0.28120583295822144, -0.2654270529747009, 0.2599720060825348, 0.2615697383880615, 0.04135964438319206, 0.14102204144001007, -0.2933970093727112, 0.3136025369167328, 0.16355553269386292, -0.002891572192311287, -0.4069441556930542, 0.11593437939882278, 0.11810056120157242, -0.02709425799548626, -0.19092397391796112, 0.05262255668640137, -0.047022510319948196, 0.07308406382799149, 0.25999051332473755, -0.20469015836715698, -0.1616048961877823, 0.32732999324798584, -0.08644941449165344, -0.07819141447544098, -0.06374135613441467, 0.1674523949623108, -0.15234646201133728, -0.14367817342281342, -0.005968580022454262, -0.0038029998540878296, 0.21332783997058868, 0.06426537781953812, -0.0521220937371254, -0.0013084083329886198, 0.11900461465120316, -0.05915044620633125, -0.09773025661706924, -0.054381269961595535, 0.11837377399206161, 0.13561590015888214, 0.24038557708263397, 0.023996254429221153, -0.39097219705581665, -0.0678715929389, -0.23275701701641083, 0.28837403655052185, 0.4594753086566925, 0.34222114086151123, -0.225957453250885, 0.04399102181196213, -0.09129935503005981, 0.035645272582769394, -0.03285100311040878, -0.22219084203243256, -0.23309728503227234, -0.10083010792732239, -0.06457238644361496, -0.023981982842087746, -0.28208041191101074, -0.08091793954372406, -0.060667917132377625, -0.1735491007566452, -0.11396098881959915, -0.47893446683883667, -0.17736142873764038, 0.2003777027130127, 0.20018436014652252, 0.07551862299442291, 0.30486685037612915, 0.08622340857982635, 0.0011746714590117335, -0.3898710012435913, 0.10835777223110199, -0.28995177149772644, -0.20937153697013855, 0.23967792093753815, 0.3102099299430847, 0.399117648601532, 0.263426274061203, -0.024005098268389702, 0.29754963517189026, 0.08551881462335587, -0.10599511116743088, 0.015001863241195679, 0.060382869094610214, -0.150295227766037, 0.034161049872636795, -0.06295378506183624, 0.10308922082185745, 0.204490527510643, 0.030635442584753036, 0.021293042227625847, -0.12706011533737183, -0.06197863444685936, 0.2626688778400421, -0.19129115343093872, 0.22021807730197906, 0.05858926475048065, -0.058575984090566635, -0.19565598666667938, 0.29368817806243896, -0.12312636524438858, 0.35392090678215027, 0.29508015513420105, -0.02815718576312065, -0.12028196454048157, 0.027273911982774734, -0.2145763337612152, -0.05555100739002228, 0.1663387417793274, -0.08715866506099701, -0.27873265743255615, -0.08636936545372009, 0.05133362114429474, 0.07966428250074387, -0.14460258185863495, 0.22641943395137787, 0.22503001987934113, 0.4570649266242981, -0.36357489228248596, 0.08248382061719894, -0.09642182290554047, -0.3318933844566345, -0.012439023703336716, 0.3023649752140045, 0.0012982247862964869, -0.025836031883955002, -0.311245858669281, -0.18457110226154327, 0.25812825560569763, 0.3368228077888489, 0.29223763942718506, 0.10281623899936676, -0.09649816155433655, 0.3434634208679199, -0.07284535467624664, -0.0828649252653122, -0.25983360409736633, -0.19592790305614471, 0.17540572583675385, 0.2842106819152832, -0.007276078220456839, -0.23288841545581818, 0.014300515875220299, 0.16241855919361115, 0.37890568375587463, 0.19551359117031097, 0.2988224923610687, 0.062271714210510254, 0.27850085496902466, 0.18463081121444702, 0.06591840833425522, 0.2824695408344269, 0.13272659480571747, 0.11020039021968842, 0.07221277058124542, -0.15831783413887024, -0.0643276497721672, 0.15242783725261688, -0.05916016921401024, -0.06371211260557175, 0.18343834578990936, -0.1718055158853531, -0.15983855724334717, 0.1015930101275444, 0.16928242146968842, -0.09237965941429138, 0.06419386714696884, -0.39357471466064453, -0.15530015528202057, -0.3959938585758209, -0.18180160224437714, 0.1489635407924652, -0.17110322415828705, -0.17660385370254517, -0.03427605330944061, 0.18506668508052826, 0.4471210539340973, 0.42980167269706726, 0.111565500497818, 0.09676375985145569, -0.06098480895161629, 0.04592487961053848, -0.2007003128528595, -0.023303115740418434, 0.03029496781527996, -0.29984235763549805, 0.2654820382595062, -0.5322842001914978, 0.23558714985847473, 0.30861586332321167, 0.14994053542613983, -0.21627141535282135, -0.1315939873456955, -0.16238388419151306, -0.22082629799842834, -0.0625302717089653, 0.18891479074954987, -0.1920023262500763, 0.039497941732406616, 0.09821950644254684, 0.20352384448051453, -0.162011981010437, -0.0001428985269740224, 0.2803606390953064, 0.24644726514816284, -0.1704164296388626, 0.4739779829978943, 0.06632119417190552, -0.2490990310907364, 0.02002951130270958, -0.33889737725257874, -0.13887730240821838, 0.20084166526794434, -0.13267259299755096, -0.20740853250026703, -0.16962316632270813, 0.2212701141834259, -0.13810628652572632, -0.025836598128080368, 0.007588289212435484, -0.3357228636741638, -0.030586445704102516, -0.32043442130088806, 0.23881827294826508, 0.08565063774585724, -0.14609947800636292, -0.045959800481796265, 0.18279628455638885, -0.1612376719713211, 0.1668296605348587, 0.18477563560009003, -0.021192137151956558, 0.1937129944562912, 0.00896142516285181, 0.19050319492816925, -0.12331343442201614, -0.3890042006969452, 0.0328926220536232, -0.36068108677864075, -0.23757030069828033, -0.09863688796758652, -0.22431187331676483, 0.17604295909404755, 0.2755296528339386, 0.15277263522148132, 0.2757103741168976, 0.47979989647865295, 0.13906553387641907, 0.07783946394920349, -0.06763409823179245, -0.12024787068367004, 0.3559828996658325, -0.2248033881187439, -0.2707900404930115, 0.23217880725860596, -0.17645034193992615, -0.12484882771968842, 0.09285150468349457, -0.07128874212503433, 0.16424980759620667, 0.29791510105133057, -0.20045660436153412, 0.37126782536506653, 0.07500802725553513, -0.12394223362207413, -0.3391548693180084, 0.0033112461678683758, 0.15995003283023834, -0.06794382631778717, -0.2513236999511719, 0.01060274988412857, 0.04944679141044617, -0.10767608880996704, -0.06480609625577927, -0.0005666770739480853, -0.24972206354141235, -0.037655096501111984, 0.010137236677110195, 0.05900650843977928, -0.10445363819599152, -0.1309271901845932, 0.2252497673034668, -0.06689074635505676, 0.015401826240122318, 0.176425501704216, 0.016921399161219597, -0.2796393036842346, -0.01271600741893053, 0.4550013244152069, -0.027980132028460503, 0.01344845350831747, 0.35556790232658386, 0.5186659097671509, 0.19090719521045685, -0.4696762263774872, -0.35049521923065186, -0.3602648079395294, 0.23641003668308258, 0.02258492261171341, -0.02355692908167839, -0.24813854694366455, 0.2369951754808426, 0.057410456240177155, 0.15081848204135895, 0.15781693160533905, 0.10607361793518066, -0.04771288484334946, -0.24064388871192932, -0.010049204342067242, 0.02772049605846405, 0.04106789082288742, 0.2115369588136673, -0.42391496896743774, 0.37044671177864075, 0.11429581791162491, -0.052308302372694016, -0.07245279848575592, -0.07232723385095596, 0.22448879480361938, -0.02341577038168907, -0.08484914898872375, -0.21119697391986847, -0.3060718774795532, -0.18266622722148895, 0.18282121419906616, 0.2294582724571228, 0.043393444269895554, -0.15766780078411102, -0.4007110595703125, 0.22883687913417816, 0.1570768654346466, 0.2030973732471466, 0.128989577293396, 0.423714280128479, -0.1583908200263977, 0.07832053303718567, -0.14644865691661835, 0.2580159306526184, 0.17284752428531647, -0.06958986818790436, 0.011864102445542812, -0.18590402603149414, -0.043440159410238266, 0.013218197971582413, 0.12371540069580078, 0.20074798166751862, -0.25197482109069824, -0.320940762758255, 0.5651822686195374, 0.18148896098136902, -0.33774706721305847, 0.08395498991012573, -0.3266913890838623, -0.0666591227054596, -0.019509922713041306, -0.010901929810643196, 0.24310599267482758, 0.1502653956413269, 0.1756867915391922, -0.05507582426071167, -0.14482833445072174, 0.07060935348272324, 0.13741806149482727, -0.3451942503452301, -0.20501278340816498, -0.3016970455646515, -0.19355548918247223, 0.34951403737068176, -0.11094190925359726, 0.30309051275253296, 0.03160431981086731, -0.5028238296508789, 0.06456613540649414, 0.5471810102462769, -0.05658369138836861, 0.13723602890968323, 0.09142395853996277, -0.015407303348183632, -0.022394968196749687, -0.08434309810400009, 0.05049850791692734, 0.04473499581217766, -0.4622526466846466, -0.04984290152788162, 0.11986816674470901, -0.3493078351020813, 0.24754080176353455, 0.17947930097579956, 0.2521968185901642, 0.348464697599411, -0.16694925725460052, -0.005244084168225527, 0.08055635541677475, -0.2675238847732544, 0.2486725151538849, 0.27256080508232117, 0.2496766746044159, -0.033478014171123505, -0.2605714201927185, 0.0921252965927124, 0.07038817554712296, -0.3656310439109802, -0.1537981629371643, 0.04908651486039162, 0.10100474208593369, 0.2997700870037079, 0.004640668164938688, 0.5057780742645264, 0.2785108983516693, 0.1374833583831787, 0.25271135568618774, -0.5460724234580994, 0.09079844504594803, 0.14453503489494324, -0.1346917748451233, -0.16755682229995728, 0.23564288020133972, -0.269066721200943, -0.2892857491970062, -0.15001341700553894, 0.3866274654865265, 0.17778606712818146, -0.08344589918851852, -0.07467872649431229, -0.3509157598018646, -0.5346744060516357, -0.10124004632234573, 0.2883478105068207, -0.08453106135129929, -0.12160143256187439, -0.07825160771608353, 0.54541015625, -0.07863668352365494, 0.3171530067920685, -0.21781328320503235, -0.27427342534065247, -0.032456833869218826, 0.02829727716743946, -0.48655855655670166, 0.1830202043056488, 0.09098141640424728, -0.03544985502958298, 0.22672806680202484, -0.324095219373703, 0.15576350688934326, 0.16658572852611542, -0.14009393751621246, 0.33569011092185974, 0.2068876475095749, 0.3805713951587677, -0.02142568677663803, 0.34842872619628906, 0.20569869875907898, 0.29127606749534607, 0.4712377190589905, -0.08093085885047913, 0.31948262453079224, 0.1657782942056656, -0.23990441858768463, -0.26159000396728516, -0.06318029761314392, -0.03956446424126625, -0.06536612659692764, -0.15550698339939117, 0.09400029480457306, 0.09627417474985123, -0.32104721665382385, 0.04134408012032509, 0.29689058661460876, -0.11313915252685547, -0.27745845913887024, 0.18948720395565033, 0.0710681825876236, 0.09767882525920868, -0.10305707901716232, 0.05818010866641998, -0.10769274085760117, -0.08762200176715851, -0.396599143743515, 0.3974556028842926, 0.4637082815170288, -0.17687484622001648, 0.192091703414917, 0.20246753096580505, -0.05056514963507652, 0.03675537556409836, 0.09243063628673553, 0.11120465397834778, 0.029207753017544746, -0.21643824875354767, 0.2651476263999939, 0.28266721963882446, -0.07300429046154022, -0.006493838038295507, -0.25386276841163635, 0.08341214060783386, 0.21187053620815277, -0.07184035331010818, 0.25222495198249817, 0.4037500023841858, 0.02791490964591503, 0.40293607115745544, 0.15702995657920837, -0.04420260712504387, -0.06018448993563652, 0.19080418348312378, -0.112275630235672, 0.13983166217803955, 0.2256273478269577, -0.19744323194026947, 0.1630644053220749, 0.09769744426012039, 0.21278266608715057, 0.13482245802879333, 0.16992011666297913, 0.04735538363456726, 0.07954015582799911, 0.13635724782943726, 0.15392634272575378, -0.03579705208539963, -0.2289511263370514, -0.34496477246284485, 0.037119310349226, 0.06418383121490479, 0.11970329284667969, -0.15436235070228577, 0.27817976474761963, 0.43157604336738586, 0.09672033786773682, -0.08673971146345139, -0.35478973388671875, 0.2358928769826889, 0.025091297924518585, -0.09803817421197891, 0.34233754873275757, 0.15862126648426056, -0.399463415145874, 0.0714215412735939, -0.06673096865415573, -0.15527252852916718, 0.09956493228673935, 0.09821203351020813, -0.12821440398693085, 0.04464436322450638, 0.5211008191108704, 0.21397460997104645, -0.03022950142621994, 0.4559059739112854, 0.1600365936756134, 0.03629790619015694, -0.1696312576532364, -0.3626648485660553, -0.21402746438980103, -0.027754394337534904, -0.21266897022724152, 0.10829539597034454, 0.17135462164878845, -0.005158039275556803, -0.08488848805427551, -0.2735466957092285, 0.039001837372779846, -0.21710485219955444, 0.25524550676345825, -0.4746195077896118, 0.34093549847602844, 0.1843951791524887, -0.001988593488931656, 0.12418776750564575, 0.03457724675536156, -0.1272900551557541, 0.18058881163597107, -0.3911055624485016, 0.004619804676622152, 0.07135342061519623, 0.1682513803243637, 0.23508808016777039, -0.13547806441783905, -0.1881614327430725, -0.3511948883533478, -0.05144886299967766, -0.024926675483584404, 0.21134696900844574, -0.21954695880413055, -0.5541378855705261, -0.1498831808567047, -0.03932051733136177, -0.018357833847403526, 0.015230671502649784, -0.24479350447654724, 0.19421648979187012, -0.07452015578746796, 0.030247030779719353, 0.04997779801487923, -0.08032412081956863, 0.17609860002994537, 0.04773733764886856, 0.21117627620697021, -0.0434027723968029, 0.28671565651893616, 0.395160973072052, -0.01068263128399849, -0.12878894805908203, -0.12175583839416504, -0.5627642869949341, 0.4152143895626068, 0.05398310720920563, -0.5084017515182495, 0.14412517845630646, 0.0013986852718517184, -0.17103908956050873, 0.11972671747207642, -0.14429634809494019, 0.13582506775856018, -0.3808746933937073, 0.007659114431589842, 0.21368615329265594, -0.005105776246637106, 0.18463093042373657, -0.04490243270993233, -0.23119790852069855, -0.030697792768478394, -0.11668961495161057, 0.269471138715744, 0.3876360356807709, 0.0061350055038928986, 0.06573937833309174, 0.15153621137142181, -0.16785326600074768, -0.02458394691348076, 0.01575670763850212, 0.15413536131381989, 0.39806023240089417, 0.27672600746154785, 0.01747356355190277, -0.2313486635684967, -0.05504807084798813, 0.284742534160614, -0.26622048020362854, 0.10963558405637741, 0.14563138782978058, 0.3589685261249542, -0.14507801830768585, -0.0015110732056200504, 0.045892950147390366, 0.13884061574935913, 0.4234732687473297, -0.025821831077337265, 0.07326806336641312, -0.10178418457508087, -0.21023111045360565, -0.2793581783771515, -0.08339869976043701, 0.2754431664943695, 0.05904866009950638, -0.30401891469955444, -0.1392781138420105, -0.3694321811199188, -0.025997500866651535, -0.15052998065948486, 0.17050588130950928, -0.03404858335852623, 0.3864196538925171, 0.44843563437461853, 0.2949526906013489, 0.07512057572603226, -0.09680856019258499, 0.0869416669011116, -0.06311952322721481, -0.0965830534696579, 0.18054568767547607, -0.29095232486724854, -0.2308036834001541, -0.019912004470825195, -0.2800006866455078, 0.2037179321050644, -0.2158188670873642, -0.08019885420799255, -0.06963054835796356, 0.1048179566860199, -0.10654757171869278, 0.2033763825893402, -0.11859297752380371, -0.01585397869348526, -0.0908111035823822, 0.10354913026094437, -0.6378620862960815, -0.005265220999717712, 0.05308053642511368, -0.10869542509317398, -0.19602878391742706, 0.015404962003231049, -0.3730745315551758, 0.10167336463928223, 0.1819157898426056, -0.348254919052124, 0.29131731390953064, 0.12113802134990692, -0.4847511053085327, -0.036215294152498245, 0.13096748292446136, -0.37685030698776245, 0.3699994385242462, -0.0027151890099048615, 0.6169676780700684, 0.36679959297180176, 0.02598322369158268, -0.1990974247455597, -0.12571147084236145, -0.020678071305155754, 0.1452460139989853, 0.2273503541946411, 0.36885473132133484, 0.07086660712957382, -0.0008883221889846027, 0.08506348729133606, 0.3406333029270172, 0.09483325481414795, -0.4519844055175781, -0.002182470867410302, 0.07158346474170685, -0.05259435996413231, 0.25525277853012085, 0.4263688623905182, -0.14534735679626465, 0.09979727119207382, 0.03198561444878578, 0.003319048322737217, 0.25923100113868713, 0.009479174390435219, -0.3665257692337036, 0.40459927916526794, -0.12994760274887085, -0.10230016708374023, -0.017253493890166283, -0.45655134320259094, 0.5030069351196289, 0.2625327706336975, -0.11628741025924683, -0.2132185697555542, -0.020758701488375664, -0.21009540557861328, 0.18591292202472687, -0.2741595208644867, 0.12829627096652985, 0.21784020960330963, 0.22891642153263092, -0.025495627894997597, 0.22447335720062256, -0.2524319589138031, 0.17787779867649078, 0.1596207469701767, -0.07467305660247803, 0.10759252309799194, 0.032125379890203476, 0.13156037032604218, 0.20401333272457123, -0.03809856250882149, 0.07744371145963669, 0.2568513751029968, -0.14913006126880646, -0.10831383615732193, 0.04972991347312927, 0.10479779541492462, -0.051978085190057755, 0.29928791522979736, -0.11767341941595078, 0.32789888978004456, 0.2038828432559967, 0.0743916854262352, 0.6412511467933655, 0.13583099842071533, 0.10759548097848892, 0.2308519333600998, -0.24664299190044403, 0.3721102178096771, -0.02351945824921131, -0.10112743079662323, -0.3047333359718323, 0.11707509309053421, 0.5186108350753784, 0.0031921269837766886, -0.09495160728693008, -0.15443941950798035, 0.6329295635223389, -0.3190315067768097, -0.03332437574863434, -0.05968708544969559, 0.1580001711845398, 0.046429023146629333, -0.030800672248005867, -0.14754512906074524, -0.2264290750026703, -0.2552430331707001, -0.034399956464767456, -0.06951862573623657, -0.17723046243190765, 0.24144957959651947, -0.0984838530421257, -0.3775452673435211, -0.4145618677139282, 0.3133876323699951, 0.006170050241053104, 0.05176849290728569, 0.12698917090892792, -0.09283670783042908, -0.2570161521434784, -0.20972304046154022, -0.31641465425491333, 0.11202407628297806, 0.049365703016519547, -0.16444464027881622, 0.39471811056137085, 0.19668307900428772, 0.05227060616016388, -0.41797223687171936, -0.06235254183411598, -0.3339781165122986, 0.3990107476711273, -0.24495214223861694, -0.038360241800546646, -0.4009758234024048, 0.3424265384674072, -0.07786083221435547, -0.24237488210201263, 0.1574542373418808, 0.08686242252588272, 0.24037156999111176, 0.1450793445110321, 0.25008824467658997, 0.17372776567935944, -0.018059546127915382, 0.15159101784229279, 0.1114591509103775, 0.1478556990623474, -0.10112257301807404, -0.013197009451687336, -0.06391160935163498, -0.20079490542411804, -0.3639146089553833, -0.06534230709075928, 0.2550210952758789, -0.08081948012113571, -0.16230757534503937, 0.10611914843320847, 0.04778743535280228, 0.11844739317893982, 0.4641798734664917, -0.0007247965550050139, 0.16326065361499786, -0.21439619362354279, 0.23889705538749695, -0.02909567579627037, -0.00022772418742533773, -0.42967739701271057, -0.1250351220369339, 0.11708596348762512, -0.33639827370643616, 0.03511737287044525, -0.1958734393119812, -0.22355115413665771, 0.2393595278263092, 0.4238606095314026, 0.06315309554338455, 0.29709017276763916, 0.6720890402793884, -0.40645506978034973, -0.2724353075027466, -0.020272059366106987, -0.15409187972545624, 0.07740417867898941, -0.11998358368873596, 0.22535136342048645, -0.1987643986940384, 0.2711876630783081, 0.04510762169957161, -0.5711377263069153, -0.16831141710281372, 0.39593705534935, -0.18316037952899933, 0.30402058362960815, -0.3299526870250702, 0.3411722779273987, -0.13877275586128235, -0.07594024389982224, -0.049145132303237915, 0.3575536906719208, -0.07139462232589722, 0.07482044398784637, -0.006590346340090036, 0.1982530653476715, 0.02596045844256878, 0.27477139234542847, 0.14888174831867218, 0.11608318239450455, -0.6615954041481018, 0.03868667408823967, -0.19130051136016846, 0.09754057973623276, 0.6332378387451172, -0.04146476462483406, 0.5417518615722656, -0.23385271430015564, -0.051694106310606, -0.09825295954942703, -0.03700282797217369, -0.5766159296035767, -0.09951458126306534, -0.17383863031864166, -0.33959388732910156, 0.11280610412359238, 0.30803900957107544, 0.2508973181247711, 0.07892360538244247, 0.3767714202404022, -0.08807791024446487, -0.13401342928409576, 0.05688486620783806, -0.19457398355007172, 0.2069525122642517, 0.35432106256484985, 0.10146773606538773, 0.12748022377490997, 0.05974818393588066, -0.08832797408103943, -0.15850767493247986, -0.21839496493339539, -0.04893484711647034, -0.2725726068019867, -0.0679113045334816, 0.09000633656978607, 0.10210023075342178, -0.4546208083629608, 0.04069598764181137, 0.1829143464565277, -0.13967332243919373, 0.009866458363831043, 0.01788308098912239, -0.27501749992370605, -0.04595281928777695, 0.15212202072143555, 0.019543282687664032, -0.45904541015625, -0.27408120036125183, 0.19381023943424225, 0.21112160384655, -0.08168908953666687, -0.12243496626615524, -0.2764289677143097, 0.024611055850982666, 0.010280454531311989, -0.1351429969072342, -0.08700072765350342, 0.4585425853729248, -0.09888572990894318, -0.004357435740530491, -0.21492382884025574, 0.11265389621257782, -0.09238718450069427, -0.384973406791687, 0.22795791923999786, -0.029424795880913734, -0.11859522759914398, -0.07303991168737411, 0.09781509637832642, 0.033325307071208954, 0.029625672847032547, -0.14974023401737213, 0.29529649019241333, 0.0723930224776268, 0.15067040920257568, 0.1979246884584427, -0.21585416793823242, 0.23143981397151947, -0.1148027703166008, 0.2561602294445038, -0.15317288041114807, -0.2542095482349396, 0.0792994275689125, -0.29255327582359314, -0.28851667046546936, 0.0034326177556067705, -0.036514732986688614, -0.014896251261234283, 0.28638139367103577, 0.009342114441096783, 0.13587680459022522, 0.4643132984638214, -0.1630719006061554, 0.03222288563847542, -0.02854287251830101, -0.6401428580284119, 0.30053040385246277, 0.24753817915916443, -0.20973384380340576, -0.34270063042640686, -0.04455482214689255, 0.2251274287700653, 0.16872258484363556, -0.037617579102516174, -0.2054091989994049, 0.15615913271903992, -0.1754268854856491, -0.18638300895690918, -0.01820538192987442, -0.1877502202987671, -0.14585933089256287, 0.07476029545068741, 0.07397663593292236, 0.22332610189914703, -0.014326357282698154, 0.0747845470905304, 0.5896537899971008, -0.17506828904151917, 0.2693387567996979, 0.2449544370174408, 0.05220824480056763, -0.10486685484647751, -0.27120673656463623, 0.0662393569946289, -0.22422710061073303, -0.081029012799263, 0.07503752410411835, -0.05053665116429329, -0.3124209940433502, 0.03394806757569313, 0.06145913526415825, 0.37737834453582764, 0.02951560914516449, 0.2129962146282196, -0.7124789357185364, -0.18582683801651, 0.08653475344181061, 0.2126546949148178, -0.20666740834712982, 0.27595236897468567, -0.02102145366370678, -0.19767366349697113, -0.2996633052825928, 0.4048396348953247, 0.039066459983587265, 0.17336903512477875, -0.04301892966032028, -0.2856467068195343, 0.33582425117492676, -0.01441101636737585, -0.030352110043168068, 0.038449130952358246, 0.022197052836418152, 0.334929496049881, 0.010943532921373844, -0.05871725082397461, -0.018744690343737602, -0.20752432942390442, 0.05608788877725601, 0.24174907803535461, -0.10894539952278137, 0.2034904658794403, 0.20740267634391785, 0.13239935040473938, -0.04345891624689102, -0.33313310146331787, 0.09389914572238922, -0.07092995941638947, -0.12753203511238098, 0.22795604169368744, -0.5624366998672485, -0.16503210365772247, -0.09571953117847443, -0.2373831868171692, 0.18216875195503235, -0.5074917078018188, -0.0026591948699206114, -0.14069029688835144, -0.25143155455589294, -0.12354934215545654, -0.09241624176502228, -0.13307498395442963, 0.20033827424049377, -0.08337816596031189, -0.2295752316713333, -0.09313782304525375, 0.3933560848236084, -0.1179254874587059, 0.09733954817056656, -0.3089309632778168, 0.16302154958248138, 0.21664763987064362, 0.25733160972595215, -0.07830451428890228, 0.03143015503883362, -0.2881809175014496, 0.2966748774051666, 0.09826602041721344, 0.08921930938959122, -0.4100605845451355, -0.11924587935209274, -0.20243017375469208, 0.37739425897598267, -0.03274588659405708, 0.12339536845684052, 0.22139427065849304, 0.10051485151052475, 0.28454136848449707, 0.30116891860961914, 0.06216886267066002, -0.07886812835931778, -0.2862052023410797, -0.09715639799833298, 0.22043989598751068, 0.3474111557006836, 0.31355175375938416, 0.20504070818424225, -0.15421131253242493, -0.07879829406738281, 0.17593391239643097, -0.30611759424209595, 0.3327474892139435, 0.15850050747394562, 0.33628252148628235, -0.11261149495840073, -0.4145216941833496, 0.442872017621994, 0.04342631623148918, -0.35810598731040955, 0.12295138835906982, 0.038954317569732666, -0.04912494868040085, -0.015561711974442005, 0.09662165492773056, -0.22968244552612305, -0.08043761551380157, -0.05205916613340378, -0.2341342270374298, -0.11822197586297989, 0.14870253205299377, 0.010974317789077759, 0.0846526175737381, -0.0745944008231163, -0.0029328023083508015, -0.28185516595840454, 0.42502403259277344, 0.08692428469657898, -0.1518162339925766, 0.46795526146888733, -0.02010761946439743, -0.07644113153219223, 0.06880263984203339, 0.09215308725833893, -0.09094592928886414, -0.052914153784513474, 0.12693381309509277, 0.3135060966014862, 0.21571144461631775, -0.41715240478515625, -0.3710979223251343, 0.15795017778873444, 0.07629584521055222, 0.19387176632881165, 0.08862168341875076, 0.18029342591762543, 0.40940436720848083, 0.03349858149886131, -0.014827312901616096, -0.04850732535123825, -0.19595947861671448, 0.0844983384013176, 0.18754175305366516, -0.12028379738330841, -0.1343783438205719, 0.19737032055854797, -0.11402186751365662, 0.07672974467277527, -0.2184404730796814, -0.21982403099536896, 0.07026565074920654, -0.09336826205253601, 0.2474168837070465, -0.04494660347700119, -0.004553758539259434, -0.49873095750808716, 0.037144772708415985, -0.31616345047950745, -0.19044211506843567, 0.1785377860069275, 0.3150533139705658, -0.05522512272000313, -0.13861830532550812, -0.07817687839269638, 0.0751267597079277, 0.38302168250083923, -0.12581339478492737, -0.01829366944730282, -0.15940892696380615, -0.05902543663978577, -0.025067077949643135, -0.18378806114196777, 0.005217286292463541, -0.14806538820266724, -0.36358511447906494, 0.058632705360651016, -0.2222817838191986, 0.1266707330942154, -0.06486767530441284, -0.3353508710861206, 0.34174486994743347, -0.16403935849666595, 0.003993190824985504, -0.28437578678131104, -0.053566824644804, 0.2577468454837799, 0.07448665052652359, -0.18107354640960693, 0.1156330332159996, -0.4084545373916626, -0.049463171511888504, 0.2806776762008667, 0.06804341077804565, 0.19486239552497864, -0.2670835852622986, -0.23804014921188354, -0.09265869110822678, 0.004533577244728804, 0.6394497752189636, -0.1078929752111435, 0.29055848717689514, 0.33537760376930237, -0.2548307180404663, -0.39849722385406494, 0.43656235933303833, 0.057926394045352936, -0.06263807415962219, 0.06250933557748795, 0.4665600657463074, 0.22083671391010284, 0.09799487143754959, 0.6195300221443176, -0.026363417506217957, -0.18050339818000793, -0.1619247943162918, -0.017730316147208214, -0.25793853402137756, 0.1244640052318573, -0.09367793053388596, -0.17001669108867645, 0.08115778863430023, 0.0965321883559227, 0.3433254063129425, -0.21176251769065857, 0.034540992230176926, -0.3491550087928772, 0.4622664451599121, 0.1671515703201294, -0.15585850179195404, -0.07531314343214035, -0.01268692035228014, -0.07268738001585007, 0.1695394068956375, -0.3712298572063446, -0.019290341064333916, 0.4501163959503174, -0.07804549485445023, 0.10564486682415009, -0.36733484268188477, -0.5134638547897339, -0.15085315704345703, -0.05766366794705391, 0.2839277684688568, -0.039917442947626114, 0.002464322606101632, 0.15434332191944122, 0.06693709641695023, -0.044904861599206924, -0.4727996587753296, 0.2287757396697998, 0.07398524135351181, 0.14052830636501312, 0.371290385723114, -0.04773787036538124, -0.0533120222389698, 0.1975872963666916, -0.07798238098621368, 0.2817493677139282, -0.17458154261112213, -0.46510568261146545, -0.1428234577178955, -0.09347642958164215, 0.2093307226896286, -0.1396270990371704, -0.19714398682117462, -0.38171011209487915, 0.012170447036623955, 0.4910549819469452, 0.17575320601463318, 0.241937518119812, -0.1937774419784546, -0.029226092621684074, 0.07547789812088013, -0.17394103109836578, 0.061186596751213074, -0.32634925842285156, -0.14357119798660278, -0.02200428396463394, -0.20684002339839935, 0.2145439088344574, -0.022394854575395584, -0.02275504358112812, -0.04784916341304779, 0.09367402642965317, 0.038120683282613754, 0.08591125905513763, -0.1366461217403412, -0.09722235798835754, -0.0986277163028717, -0.20601528882980347, -0.07371743768453598, 0.28274470567703247, 0.023698413744568825, -0.09299236536026001, 0.10031933337450027, -0.15373210608959198, 0.08874983340501785, -0.12071976810693741, -0.016329750418663025, 0.030793851241469383, 0.14248153567314148, 0.24902917444705963, -0.08586367219686508, -0.032679587602615356, -0.4468836486339569, 0.0649845227599144, 0.4187562167644501, -0.2575630247592926, 0.1612449586391449, 0.2806607484817505, 0.25186607241630554, -0.18802230060100555, -0.1641024351119995, -0.15155890583992004, -0.11297833174467087, 0.18683739006519318, 0.18661582469940186, 0.1841040998697281, 0.034915827214717865, 0.14198429882526398, 0.1809442639350891, -0.00284718070179224, -0.21217162907123566, -0.1330992579460144, 0.005542280152440071, -0.11966531723737717, 0.18252889811992645, 0.026674170047044754, -0.3690573275089264, -0.45218801498413086, -0.21554875373840332, -0.19318002462387085, -0.31044796109199524, 0.22007852792739868, -0.0991736352443695, -0.05229809880256653, -0.12433364242315292, -0.06435205042362213, 0.45587560534477234, 0.017262635752558708, 0.42667078971862793, 0.5317025780677795, 0.18526968359947205, -0.08383453637361526, 0.11155389994382858, 0.0188443586230278, -0.06600974500179291, 0.04185177758336067, -0.3549390733242035, -0.053183991461992264, -0.10848186910152435, -0.08795752376317978, 0.21250340342521667, -0.14107352495193481, -0.0011400262592360377, 0.39878883957862854, 0.23509334027767181, -0.4224924147129059, 0.15322303771972656, 0.28449997305870056, -0.1846538782119751, -0.08085231482982635, -0.09617929905653, 0.3012271225452423, -0.24097523093223572, -0.2511283755302429, 0.07765209674835205, 0.22644370794296265, -0.24334602057933807, 0.2552917003631592, 0.061097778379917145, -0.08640249818563461, -0.1836627870798111, 0.36664879322052, -0.25842395424842834, 0.28168970346450806, 0.11775891482830048, -0.13175146281719208, 0.3049212694168091, 0.3759821653366089, -0.2934805750846863, 0.24080343544483185, 0.05031620338559151, -0.1546012908220291, 0.16539068520069122, -0.16234932839870453, 0.04960669204592705, 0.04545273259282112, 0.02244069240987301, 0.3382852375507355, -0.09488473832607269, -0.09629730135202408, -0.18673236668109894, -0.11188986152410507, -0.08622637391090393, -0.1441119760274887, 0.2856958508491516, 0.08405803889036179, -0.29654744267463684, -0.2261635661125183, -0.20700007677078247, -0.1864265352487564, -0.33735549449920654, 0.05152876675128937, -0.016044115647673607, 0.44890621304512024, 0.6223233342170715, -0.31260934472084045, 0.19532647728919983, -0.22434481978416443, 0.0677201896905899, -0.38060063123703003, -0.1491813063621521, 0.09035596996545792, 0.05618315190076828, 0.21031999588012695, -0.01565629057586193, 0.3211047649383545, -0.24846264719963074, 0.15252405405044556, 0.32450512051582336, -0.24208669364452362, 0.14003659784793854, -0.08408191055059433, 0.05215233191847801, 0.004425026476383209, -0.43147897720336914, -0.06158427894115448, -0.15501955151557922, -0.307052880525589, 0.6805815696716309, 0.15920576453208923, 0.153838649392128, 0.19582322239875793, -0.07174821943044662, 0.009960738942027092, 0.2325565367937088, -0.06306822597980499, -0.17565688490867615, 0.15634645521640778, -0.03780302032828331, 0.3915519416332245, 0.00023100913676898926, -0.30778390169143677, -0.08396747708320618, 0.19523707032203674, -0.16481070220470428, 0.10846290737390518, 0.321951299905777, -0.0054274266585707664, -0.03594231233000755, 0.3486814498901367, -0.10092350840568542, 0.44017669558525085, 0.2724825143814087, 0.13226713240146637, 0.2087632417678833, 0.541908323764801, 0.1010870561003685, 0.26421087980270386, 0.13724088668823242, -0.15543261170387268, 0.11957918852567673, -0.3819953501224518, 0.2681388258934021, 0.038335252553224564, -0.14927010238170624, -0.12001821398735046, 0.10033683478832245, -0.3448009788990021, 0.4414682388305664, 0.002026696689426899, 0.09886753559112549, 0.19212636351585388, 0.051119204610586166, 0.13111227750778198, 0.39653751254081726, 0.35668647289276123, -0.21606826782226562, -0.2661657929420471, -0.07273141294717789, 0.1785300374031067, -0.12034109979867935, 0.6376805305480957, -0.3669545650482178, 0.35525721311569214, 0.23681290447711945, 0.11149488389492035, 0.0009487506467849016, 0.23230040073394775, 0.12854397296905518, -0.4854184091091156, -0.007641731761395931, -0.039015572518110275, -0.550255298614502, 0.034043088555336, -0.015992727130651474, -0.22703790664672852, -0.1733274608850479, -0.2143063247203827, 0.12683576345443726, -0.12746652960777283, 0.3714889883995056, -0.14113090932369232, 0.1461343914270401, -0.1047661304473877, 0.006553714629262686, -0.015459248796105385, -0.1993369460105896, 0.14841949939727783, -0.12969298660755157, 0.16662177443504333, 0.07223477214574814, 0.22970931231975555, -0.1699347347021103, 0.22280219197273254, 0.052692580968141556, 0.30829668045043945, -0.11267919838428497, 0.005308744963258505, -0.11129885166883469, -0.2809552848339081, -0.08258936554193497, 0.07430632412433624, 0.09139449149370193, -0.36024680733680725, 0.2726401686668396, 0.348268985748291, 0.24737775325775146, -0.308316171169281, -0.06702796369791031, -0.1664225161075592, 0.08103512972593307, -0.25600019097328186, 0.2549786865711212, 0.19929906725883484, -0.45650729537010193, 0.17165100574493408, 0.03431098908185959, -0.004508157726377249, 0.23794148862361908, 0.22548924386501312, -0.4006038010120392, 0.004861735738813877, -0.3583067059516907, -0.15687201917171478, 0.2862221300601959, 0.19862304627895355, -0.017098726704716682, 0.1498011350631714, -0.16238683462142944, 0.06850370019674301, -0.11610788106918335, 0.20272761583328247, -0.10752753913402557, 0.26559925079345703, -0.20066295564174652, -0.21649320423603058, 0.3223603665828705, -0.22452743351459503, -0.023795953020453453, 0.11615648120641708, 0.19631323218345642, -0.2107708603143692, -0.29961177706718445, 0.058438464999198914, 0.24339361488819122, 0.15352578461170197, -9.867265907814726e-05, -0.2029186189174652, 0.2910192906856537, -0.06201080605387688, 0.6055508852005005, 0.36252710223197937, -0.2541446089744568, 0.34554266929626465, -0.004870417062193155, 0.21045732498168945, -0.24053840339183807, -0.28676241636276245, 0.016104767099022865, 0.13112136721611023, -0.2205459624528885, 0.10129981487989426, 0.23806189000606537, 0.0773029774427414, -0.15682947635650635, -0.2897454500198364, 0.21097812056541443, -0.07592977583408356, 0.07636670768260956, 0.2903869152069092, -0.24407972395420074, 0.26638802886009216, -0.03807307779788971, 0.04416986554861069, 0.4984849691390991, 0.15725287795066833, 0.38291093707084656, 0.3352759778499603, 0.1742335557937622, 0.10592364519834518, 0.3040756583213806, 0.0019543843809515238, 0.6564239859580994, 0.02589373104274273, -0.10833204537630081, 0.1987474262714386, -0.16524940729141235, 0.00021965797350276262, 0.16463090479373932, -0.24028006196022034, 0.3058965504169464, -0.23893782496452332, 0.11180949956178665, 0.24736779928207397, 0.03774884343147278, 0.13389380276203156, -0.19293901324272156, -0.0525573194026947, -0.010163001716136932, -0.2981116473674774, 0.009471491910517216, -0.08888828754425049, -0.4862995147705078, 0.003414287930354476, -0.3598620295524597, -0.013856817036867142, -0.014671380631625652]}
diff --git a/tests/fixtures/tiny_h8.pth b/tests/fixtures/tiny_h8.pth
new file mode 100644
index 0000000..5f9814a
Binary files /dev/null and b/tests/fixtures/tiny_h8.pth differ
diff --git a/tests/fixtures/tiny_h8_config.json b/tests/fixtures/tiny_h8_config.json
new file mode 100644
index 0000000..5b8122f
--- /dev/null
+++ b/tests/fixtures/tiny_h8_config.json
@@ -0,0 +1,39 @@
+{
+  "_name_or_path": "",
+  "activation_function": "gelu_new",
+  "architectures": null,
+  "attn_pdrop": 0.0,
+  "bos_token_id": 50256,
+  "chunk_size_feed_forward": 0,
+  "dtype": null,
+  "embd_pdrop": 0.0,
+  "eos_token_id": 50256,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "initializer_range": 0.02,
+  "is_encoder_decoder": false,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "codegen",
+  "n_ctx": 64,
+  "n_embd": 128,
+  "n_head": 8,
+  "n_inner": null,
+  "n_layer": 2,
+  "n_positions": 64,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "problem_type": null,
+  "resid_pdrop": 0.0,
+  "return_dict": true,
+  "rotary_dim": 16,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.12.1",
+  "use_cache": true,
+  "vocab_size": 256
+}
diff --git a/tests/fixtures/tiny_h8_logits.json b/tests/fixtures/tiny_h8_logits.json
new file mode 100644
index 0000000..de0ba37
--- /dev/null
+++ b/tests/fixtures/tiny_h8_logits.json
@@ -0,0 +1 @@
+{"tokens": [3, 17, 42, 8, 255, 1, 99, 128], "shape": [1, 8, 256], "logits": [0.13797208666801453, -0.12160780280828476, 0.015172661282122135, 0.12064576894044876, 0.05379147827625275, -0.29649820923805237, 0.2050786018371582, -0.1565551459789276, -0.06716244667768478, 0.03846554830670357, 0.3188575804233551, -0.17997929453849792, -0.1908494383096695, -0.04417635127902031, -0.08735319972038269, 0.40719372034072876, -0.24718672037124634, 0.14152072370052338, -0.12805403769016266, -0.06504087150096893, -0.09366442263126373, -0.18055735528469086, -0.2347811460494995, 0.29062095284461975, 0.04116421565413475, 0.06698659062385559, 0.13029389083385468, 0.03848961740732193, 0.2765219509601593, -0.2865062355995178, 0.20257458090782166, 0.1723998636007309, 0.2137833684682846, 0.4326269030570984, -0.0009396881214343011, 0.17215093970298767, 0.2375749796628952, -0.11660584807395935, 0.04120151698589325, -0.3389374613761902, -0.16921766102313995, -0.09879470616579056, -0.18560680747032166, 0.33959048986434937, 0.10422768443822861, -0.08966687321662903, -0.6132300496101379, 0.255328893661499, 0.16417236626148224, -0.24560371041297913, 0.07819544523954391, -0.10255435854196548, -0.4769432842731476, 0.033296238631010056, 0.3392610549926758, 0.2763991057872772, -0.1838977336883545, 0.20095357298851013, 0.1466151475906372, -0.19600743055343628, 0.06724176555871964, -0.041446734219789505, -0.2139338105916977, -0.13107256591320038, 0.1255294233560562, -0.15353916585445404, -0.1614842414855957, 0.05491883307695389, -0.08805611729621887, 0.30660784244537354, -0.2726972699165344, -0.22136205434799194, 0.017901619896292686, -0.2540002763271332, 0.12028636783361435, 0.055389653891325, -0.051570210605859756, 0.0692351832985878, -0.1784837543964386, 0.023299047723412514, 0.20437733829021454, 0.2109411507844925, -0.06377265602350235, 0.04263145849108696, 0.39125245809555054, 0.13542288541793823, 0.1764322817325592, 0.28708139061927795, 0.3067520260810852, -0.024627884849905968, 0.08944395929574966, 0.3394125998020172, 0.022080564871430397, 0.2146754115819931, 0.1351061314344406, 0.14937809109687805, -0.4144514799118042, 0.25638189911842346, -0.16860802471637726, 0.019572222605347633, 0.508650541305542, -0.1423056721687317, -0.010113397613167763, -0.20181059837341309, 0.08867400884628296, 0.06273669749498367, 0.23533694446086884, 0.14392374455928802, -0.04992162808775902, 0.08846477419137955, -0.05817421153187752, 0.26681283116340637, -0.27640360593795776, -0.29602283239364624, 0.02165238745510578, -0.008114050142467022, -0.28540971875190735, 0.019992755725979805, -0.3857240378856659, -0.21671807765960693, -0.007429229561239481, 0.2322276085615158, 0.12162502110004425, 0.09285692870616913, -0.3684569299221039, -0.06472666561603546, -0.003905214834958315, -0.26022034883499146, -0.011866025626659393, 0.02485169656574726, -0.15326546132564545, -0.2529875934123993, 0.11364205181598663, -0.08141594380140305, 0.07734422385692596, 0.052367035299539566, 0.10160461068153381, 0.11050158739089966, 0.0920712798833847, -0.28216618299484253, 0.2772896885871887, 0.271312415599823, 0.19137859344482422, -0.5577556490898132, -0.6188551187515259, -0.5914639234542847, 0.1950666904449463, -0.010510235093533993, 0.004505499731749296, 0.1058182567358017, -0.15907728672027588, -0.11431925743818283, -0.47599345445632935, 0.07211202383041382, 0.27881208062171936, -0.4005069434642792, -0.03024817816913128, -0.14754502475261688, -0.24671459197998047, -0.16086988151073456, -0.709753692150116, -0.35539573431015015, -0.13390493392944336, 0.0928003191947937, 0.04863446578383446, -0.0495501346886158, 0.18421092629432678, -0.018080364912748337, 0.5899195075035095, 0.365138977766037, 0.1963970959186554, 0.1972842514514923, -0.08086569607257843, 0.2867452800273895, 0.36433976888656616, -0.0034080499317497015, 0.3998028039932251, 0.1984650045633316, -0.002680029021576047, -0.06290017068386078, -0.026031747460365295, -0.10381772369146347, 0.13147704303264618, 0.12688888609409332, 0.027697419747710228, -0.25628265738487244, -0.09270302951335907, -0.30371055006980896, -0.17885839939117432, 0.2448159009218216, 0.019905803725123405, 0.31289583444595337, 0.6091428399085999, -0.29281556606292725, -0.09114749729633331, -0.20409734547138214, -0.10191241651773453, 0.07138940691947937, 0.34490522742271423, -0.25056400895118713, 0.12000268697738647, 0.0363888256251812, 0.15915948152542114, 0.3538452088832855, 0.21529997885227203, -0.053232014179229736, 0.5123010873794556, -0.3488830626010895, -0.16908413171768188, -0.22125312685966492, 0.14842158555984497, 0.11073613166809082, 0.021522222086787224, -0.19074460864067078, 0.2579523026943207, 0.3705143332481384, 0.15849003195762634, 0.16991882026195526, 0.06226116791367531, -0.06926238536834717, -0.4051671028137207, 0.31475120782852173, -0.2134520709514618, 0.21438013017177582, -0.02561035566031933, 0.11234936118125916, -0.016122084110975266, -0.4536351263523102, -0.1443435549736023, -0.05539524927735329, -0.17995263636112213, -0.11866766214370728, 0.09042748808860779, 0.05701519921422005, 0.2888790965080261, -0.023825109004974365, -0.06706663221120834, 0.34762290120124817, -0.23093216121196747, 0.0027742611709982157, -0.13100560009479523, -0.1632840782403946, -0.2922150492668152, 0.1936868578195572, -0.30890393257141113, 0.17172551155090332, 0.36000964045524597, -0.1943063735961914, 0.0905841737985611, -0.19335758686065674, -0.17201192677021027, 0.03831101208925247, -0.026263123378157616, 0.23113606870174408, -0.3485087752342224, 0.09443095326423645, 0.11309722810983658, -0.15473923087120056, 0.662071168422699, -0.06606530398130417, 0.6609460711479187, 0.18824854493141174, 0.2294606864452362, 0.4111865758895874, 0.09342775493860245, -0.18579526245594025, -0.19412420690059662, -0.1244269534945488, 0.31320658326148987, -0.21811898052692413, -0.19284780323505402, 0.4705081284046173, 0.0988529622554779, -0.3770406246185303, -0.058348484337329865, 0.3337516486644745, 0.10379016399383545, 0.05062384530901909, -0.3621722459793091, -0.04999285936355591, -0.3783969283103943, 0.10267385095357895, 0.08293101191520691, -0.06442992389202118, 0.10337621718645096, 0.3607746362686157, -0.0646272599697113, 0.12985685467720032, 0.016312792897224426, -0.004679706413298845, 0.17402033507823944, 0.195995032787323, -0.08374378830194473, -0.20531730353832245, -0.1385122388601303, 0.2672266662120819, 0.16584455966949463, -0.027894511818885803, -0.1507774144411087, -0.09081410616636276, 0.14935290813446045, 0.28635624051094055, -0.2348891943693161, -0.5304549336433411, -0.27226758003234863, 0.10707123577594757, -0.11309769004583359, 0.14964033663272858, -0.4480280578136444, -0.14622420072555542, 0.35938701033592224, -0.08414198458194733, -0.013268403708934784, -0.3053439259529114, 0.5279374122619629, -0.04263276234269142, 0.24758297204971313, -0.09561405330896378, -0.2904089093208313, 0.3222983479499817, -0.029253017157316208, -0.2712303698062897, 0.08924799412488937, -0.16471461951732635, -0.049886103719472885, -0.10957154631614685, 0.030337616801261902, 0.11476067453622818, 0.17604531347751617, 0.07316932082176208, -0.13722996413707733, -0.035167768597602844, -0.1711834818124771, -0.08807266503572464, -0.208733931183815, -0.16633684933185577, -0.33437037467956543, 0.08371637761592865, 0.4117291569709778, 0.11732731014490128, -0.22609634697437286, -0.14739570021629333, 0.2800227105617523, -0.19284120202064514, -0.2760176658630371, 0.4322760999202728, -0.019549401476979256, -0.07597988843917847, 0.04100493714213371, -0.18073026835918427, 0.04292677715420723, -0.21047134697437286, -0.02622130699455738, -0.18790338933467865, -0.0920541062951088, -0.29893338680267334, -0.07757652550935745, -0.045110639184713364, 0.07920093834400177, -0.03261323273181915, 0.12029965966939926, 0.03322594612836838, 0.18971964716911316, 0.5931516885757446, -0.4545058310031891, -0.07274221628904343, -0.3987809419631958, -0.06465096771717072, 0.17316977679729462, -0.005919998977333307, 0.09789375960826874, 0.131780743598938, -0.1911570429801941, -0.17051228880882263, -0.43521904945373535, 0.031356602907180786, 0.0726364329457283, -0.17543958127498627, -0.15176057815551758, -0.3644367754459381, 0.09173732250928879, 0.058867815881967545, 0.20303648710250854, -0.08989720046520233, 0.19463984668254852, 0.31619155406951904, 0.17896072566509247, -0.6792815327644348, -0.6332765221595764, 0.4353419244289398, 0.2696858048439026, -0.04295768588781357, -0.022547010332345963, 0.3821994960308075, 0.2782876491546631, -0.2642652094364166, 0.25823697447776794, 0.26448917388916016, 0.04724867269396782, 0.14197054505348206, -0.2868606150150299, 0.30502668023109436, 0.16510632634162903, -0.003995835315436125, -0.4035305678844452, 0.12874513864517212, 0.10935413837432861, -0.020559852942824364, -0.18537940084934235, 0.0471443273127079, -0.04914327710866928, 0.07496724277734756, 0.26874852180480957, -0.19698554277420044, -0.16418491303920746, 0.3232228457927704, -0.08573740720748901, -0.07388627529144287, -0.06117302551865578, 0.1626524180173874, -0.15429188311100006, -0.13958142697811127, -0.005142178852111101, 0.0015309692826122046, 0.20949555933475494, 0.06730081140995026, -0.04640423506498337, 1.1059491953346878e-05, 0.11994519084692001, -0.0586535707116127, -0.0989355742931366, -0.05414619669318199, 0.12388554960489273, 0.13790486752986908, 0.23918454349040985, 0.02772393822669983, -0.38159364461898804, -0.06234785541892052, -0.22729270160198212, 0.28907889127731323, 0.4611451327800751, 0.3389657139778137, -0.23100700974464417, 0.039184149354696274, -0.09227710962295532, 0.032157715409994125, -0.03919070214033127, -0.22379657626152039, -0.2248854786157608, -0.10089784115552902, -0.056316912174224854, -0.023091543465852737, -0.2870732545852661, -0.07996683567762375, -0.062263570725917816, -0.17002655565738678, -0.12187619507312775, -0.47767794132232666, -0.18775568902492523, 0.2058379054069519, 0.2081000804901123, 0.06751962751150131, 0.3068946301937103, 0.07464830577373505, 0.007350680883973837, -0.39154812693595886, 0.09970088303089142, -0.28536584973335266, -0.2037181258201599, 0.23812280595302582, 0.3087700307369232, 0.40525051951408386, 0.26627418398857117, -0.02140399068593979, 0.3023695647716522, 0.08339347690343857, -0.11029431968927383, 0.0126644903793931, 0.06003294885158539, -0.14977054297924042, 0.038256414234638214, -0.06137093901634216, 0.11202659457921982, 0.206180602312088, 0.03345143049955368, 0.031536318361759186, -0.1287703961133957, -0.07086833566427231, 0.25338318943977356, -0.2010396122932434, 0.2196558713912964, 0.056152548640966415, -0.05940153822302818, -0.20116490125656128, 0.28619396686553955, -0.11873779445886612, 0.35017040371894836, 0.2875915467739105, -0.023660089820623398, -0.1158359944820404, 0.02419712021946907, -0.21636950969696045, -0.053557202219963074, 0.16536392271518707, -0.08526304364204407, -0.2825116217136383, -0.10012324154376984, 0.05390790477395058, 0.07165170460939407, -0.14770987629890442, 0.22471977770328522, 0.231315016746521, 0.46060076355934143, -0.36644190549850464, 0.08478888869285583, -0.10112893581390381, -0.33486559987068176, -0.01644526608288288, 0.3037440776824951, 0.0011272924020886421, -0.030057281255722046, -0.30786994099617004, -0.18525277078151703, 0.2585781216621399, 0.33867430686950684, 0.2894304096698761, 0.10698139667510986, -0.10338841378688812, 0.3430311381816864, -0.07448221743106842, -0.08568823337554932, -0.256818026304245, -0.19260022044181824, 0.1753675639629364, 0.2812010943889618, -0.012797174043953419, -0.2319236397743225, 0.015557916834950447, 0.16760432720184326, 0.3821672201156616, 0.19823461771011353, 0.30726850032806396, 0.060834068804979324, 0.276072233915329, 0.1893700808286667, 0.06375132501125336, 0.2811519205570221, 0.12656602263450623, 0.11339154094457626, 0.07140684872865677, -0.1601482480764389, -0.06109786033630371, 0.15411868691444397, -0.06164822354912758, -0.06502798199653625, 0.18413615226745605, -0.16802147030830383, -0.15177111327648163, 0.10412481427192688, 0.1678478866815567, -0.09589812904596329, 0.06816836446523666, -0.3964516818523407, -0.15999983251094818, -0.3907450735569, -0.1852341741323471, 0.15068507194519043, -0.1667487621307373, -0.18001464009284973, -0.030477475374937057, 0.18382778763771057, 0.45157942175865173, 0.4262046813964844, 0.11685487627983093, 0.09861814230680466, -0.060508690774440765, 0.05304969847202301, -0.19721242785453796, -0.02288387343287468, 0.028281984850764275, -0.295274019241333, 0.2659447193145752, -0.5266457200050354, 0.23883000016212463, 0.306633323431015, 0.1477511078119278, -0.22015877068042755, -0.1351996213197708, -0.16631540656089783, -0.2213287502527237, -0.06901919841766357, 0.1861397624015808, -0.19645890593528748, 0.04045708104968071, 0.09949818253517151, 0.202034130692482, -0.15276026725769043, -0.0051469625905156136, 0.2856571674346924, 0.2468234896659851, -0.16640424728393555, 0.47486090660095215, 0.07597853988409042, -0.24814702570438385, 0.021254675462841988, -0.33491966128349304, -0.13893108069896698, 0.1987626552581787, -0.12858758866786957, -0.20276989042758942, -0.1764591485261917, 0.21836018562316895, -0.14063186943531036, -0.02324821613729, 0.01200202014297247, -0.33661481738090515, -0.028010103851556778, -0.30980831384658813, 0.2365112453699112, 0.0864056870341301, -0.1482081115245819, -0.05156503617763519, 0.1830853670835495, -0.1705993413925171, 0.1715056300163269, 0.18507468700408936, -0.02392854541540146, 0.19670455157756805, 0.017262298613786697, 0.18932297825813293, -0.12605059146881104, -0.38753893971443176, 0.03485049307346344, -0.3609975278377533, -0.24316932260990143, -0.10204458236694336, -0.22364841401576996, 0.16607287526130676, 0.2675122320652008, 0.14984501898288727, 0.2709795832633972, 0.4844171106815338, 0.14102177321910858, 0.0771946907043457, -0.07807864248752594, -0.12009087204933167, 0.35533395409584045, -0.23073628544807434, -0.2693736255168915, 0.22895869612693787, -0.17428730428218842, -0.11867199093103409, 0.09513570368289948, -0.07562820613384247, 0.16438356041908264, 0.30030179023742676, -0.20357981324195862, 0.37953174114227295, 0.07266756892204285, -0.1230880469083786, -0.3427314758300781, -0.003198185469955206, 0.16185413300991058, -0.06756097823381424, -0.2541719377040863, 0.01146284956485033, 0.0437474250793457, -0.10268678516149521, -0.0687747672200203, -0.008676156401634216, -0.25069838762283325, -0.03210292384028435, 0.0060976482927799225, 0.06129223108291626, -0.10888843983411789, -0.12824542820453644, 0.22091364860534668, -0.07332968711853027, 0.015706755220890045, 0.17723815143108368, 0.01829698495566845, -0.28642839193344116, -0.01664990931749344, 0.4557572901248932, -0.02582811564207077, 0.009577569551765919, 0.3557652235031128, 0.5194957256317139, 0.18505920469760895, -0.4734773337841034, -0.35063424706459045, -0.36339250206947327, 0.23438765108585358, 0.021805303171277046, -0.029108203947544098, -0.24753381311893463, 0.23572483658790588, 0.053390584886074066, 0.150822713971138, 0.158674418926239, 0.10102605819702148, -0.043830692768096924, -0.2402631640434265, -0.01185689028352499, 0.022305386140942574, 0.052540045231580734, 0.205717071890831, -0.425773948431015, 0.37134307622909546, 0.11194020509719849, -0.05192159116268158, -0.07011052966117859, -0.07208608090877533, 0.23014336824417114, -0.016968315467238426, -0.08524180948734283, -0.21515081822872162, -0.30541732907295227, -0.18486833572387695, 0.19032318890094757, 0.22810938954353333, 0.047896116971969604, -0.1554894745349884, -0.40350690484046936, 0.2281004637479782, 0.1544295996427536, 0.2056574523448944, 0.12524454295635223, 0.4287235736846924, -0.1531759351491928, 0.07809318602085114, -0.1477736085653305, 0.25666406750679016, 0.17150050401687622, -0.06970887631177902, 0.010999691672623158, -0.18571047484874725, -0.043932825326919556, 0.0013662331039085984, 0.12450432032346725, 0.2007376104593277, -0.25598248839378357, -0.3218079209327698, 0.566963791847229, 0.1880696415901184, -0.33341485261917114, 0.08863195031881332, -0.32682037353515625, -0.06527630239725113, -0.02008838951587677, -0.006654055323451757, 0.2442128211259842, 0.1500924974679947, 0.1811467409133911, -0.06290586292743683, -0.14387567341327667, 0.07598918676376343, 0.1354294717311859, -0.3395877182483673, -0.2108222246170044, -0.3037504255771637, -0.19261936843395233, 0.3533783257007599, -0.11155732721090317, 0.3080183267593384, 0.03122074156999588, -0.5036177635192871, 0.0586593896150589, 0.5452279448509216, -0.05504472553730011, 0.1353476196527481, 0.09949391335248947, -0.01808970980346203, -0.01982753910124302, -0.08946214616298676, 0.04525788128376007, 0.04272906109690666, -0.4611493945121765, -0.05322571471333504, 0.11778531223535538, -0.3513640761375427, 0.2451709359884262, 0.18506741523742676, 0.25514838099479675, 0.35022515058517456, -0.167572021484375, -0.009060642682015896, 0.08104793727397919, -0.270858496427536, 0.2499336302280426, 0.2720787227153778, 0.24909919500350952, -0.03022707626223564, -0.26501137018203735, 0.09311448037624359, 0.0681208074092865, -0.36627134680747986, -0.15530598163604736, 0.05097350850701332, 0.10604248940944672, 0.3050232529640198, 0.00686998013406992, 0.5068744421005249, 0.27083125710487366, 0.13335971534252167, 0.25759321451187134, -0.5460232496261597, 0.08899059891700745, 0.1437467634677887, -0.13040278851985931, -0.16160088777542114, 0.23711815476417542, -0.2721710205078125, -0.28550705313682556, -0.15045639872550964, 0.3883058428764343, 0.18191641569137573, -0.08159399777650833, -0.07457898557186127, -0.3497689664363861, -0.5310201048851013, -0.10056376457214355, 0.28874626755714417, -0.08447863161563873, -0.11901246011257172, -0.0832962691783905, 0.5455342531204224, -0.08114401251077652, 0.31726354360580444, -0.21916310489177704, -0.2765762209892273, -0.03151125833392143, 0.02979958802461624, -0.48685118556022644, 0.1824142038822174, 0.0929233580827713, -0.037530794739723206, 0.22651875019073486, -0.3250076472759247, 0.15297754108905792, 0.16819463670253754, -0.14484694600105286, 0.34008508920669556, 0.2067747265100479, 0.382371187210083, -0.01901732012629509, 0.34590986371040344, 0.20872347056865692, 0.2852517366409302, 0.4745241403579712, -0.0772872194647789, 0.31533634662628174, 0.1668190360069275, -0.23989474773406982, -0.25901803374290466, -0.06183804199099541, -0.0407821349799633, -0.06747337430715561, -0.15776848793029785, 0.0962514653801918, 0.09535742551088333, -0.32183441519737244, 0.033351290971040726, 0.3008825182914734, -0.11305758357048035, -0.2753670811653137, 0.19662533700466156, 0.072392039000988, 0.0981489047408104, -0.10488250851631165, 0.05623627081513405, -0.11212635785341263, -0.08649247884750366, -0.40001776814460754, 0.39853495359420776, 0.46003052592277527, -0.17701134085655212, 0.19510307908058167, 0.20786789059638977, -0.04760366305708885, 0.042184147983789444, 0.08961888402700424, 0.12033609300851822, 0.03785823658108711, -0.21610543131828308, 0.2640339136123657, 0.2817733585834503, -0.0723705142736435, 0.0007182909757830203, -0.25434553623199463, 0.08602303266525269, 0.21401605010032654, -0.07006165385246277, 0.25724565982818604, 0.4058113992214203, 0.03005906008183956, 0.3989778161048889, 0.1535385400056839, -0.04701404273509979, -0.058875832706689835, 0.19550861418247223, -0.11605629324913025, 0.136036217212677, 0.2252558022737503, -0.19479906558990479, 0.16444607079029083, 0.09362050890922546, 0.21717557311058044, 0.1315596103668213, 0.1703791469335556, 0.05236491933465004, 0.08122515678405762, 0.14203952252864838, 0.1571657508611679, -0.03278452157974243, -0.22945015132427216, -0.352108895778656, 0.03833387419581413, 0.06589394062757492, 0.12287381291389465, -0.15305951237678528, 0.2817755937576294, 0.43052974343299866, 0.09265834093093872, -0.08391295373439789, -0.3578912615776062, 0.23847363889217377, 0.023048264905810356, -0.09503881633281708, 0.34070906043052673, 0.15809760987758636, -0.3954713046550751, 0.07279622554779053, -0.06684569269418716, -0.1559291034936905, 0.09816277027130127, 0.09336099028587341, -0.13170230388641357, 0.048114459961652756, 0.521411657333374, 0.21204710006713867, -0.03029509447515011, 0.4578786790370941, 0.16362275183200836, 0.030445443466305733, -0.16770905256271362, -0.3583342432975769, -0.2146870642900467, -0.02773110382258892, -0.210508793592453, 0.1098756268620491, 0.17095054686069489, -0.012734541669487953, -0.07709736377000809, -0.2749696969985962, 0.042392343282699585, -0.21428550779819489, 0.2537795603275299, -0.47964781522750854, 0.33967822790145874, 0.1847957819700241, -0.0025801747106015682, 0.12302713096141815, 0.03012050688266754, -0.12847298383712769, 0.1779477596282959, -0.3915199637413025, 0.005979795940220356, 0.06635972857475281, 0.17381329834461212, 0.23373648524284363, -0.1352713704109192, -0.18469716608524323, -0.35072049498558044, -0.057265836745500565, -0.027782555669546127, 0.21474167704582214, -0.220360666513443, -0.5559919476509094, -0.1500745266675949, -0.04058489948511124, -0.01975410245358944, 0.015063724480569363, -0.24488505721092224, 0.19088394939899445, -0.08108088374137878, 0.024782365188002586, 0.0504009984433651, -0.08328104764223099, 0.17574334144592285, 0.04205522686243057, 0.20632325112819672, -0.04534284770488739, 0.2838665544986725, 0.39337414503097534, -0.008617287501692772, -0.12490372359752655, -0.12448129802942276, -0.5632677674293518, 0.4132124185562134, 0.05565560609102249, -0.5071448087692261, 0.1447521150112152, 0.0020343931391835213, -0.1740882694721222, 0.1277698278427124, -0.1443924903869629, 0.12890315055847168, -0.389346718788147, 0.012497788295149803, 0.21453270316123962, -0.0049479347653687, 0.18714331090450287, -0.0400487445294857, -0.23097266256809235, -0.02431545779109001, -0.12176081538200378, 0.2685040533542633, 0.3934905230998993, 0.007754985708743334, 0.062460340559482574, 0.1473664939403534, -0.17245668172836304, -0.024080464616417885, 0.012284360826015472, 0.1529606282711029, 0.3920636773109436, 0.28430432081222534, 0.01800810918211937, -0.23816226422786713, -0.060417290776968, 0.27956631779670715, -0.2670403718948364, 0.1096525490283966, 0.1450692117214203, 0.35958167910575867, -0.13933488726615906, -0.011017024517059326, 0.04487098008394241, 0.1340564489364624, 0.4233812093734741, -0.02106393128633499, 0.08135899156332016, -0.0994509607553482, -0.22051574289798737, -0.2785786986351013, -0.08071716129779816, 0.2755527198314667, 0.06539693474769592, -0.3035334646701813, -0.13880281150341034, -0.37248533964157104, -0.019502902403473854, -0.1477678120136261, 0.1661153882741928, -0.04302966222167015, 0.3836665749549866, 0.4525268077850342, 0.29004210233688354, 0.06465601176023483, -0.09940249472856522, 0.08649495989084244, -0.06486691534519196, -0.09454496204853058, 0.17909695208072662, -0.2838318347930908, -0.22912831604480743, -0.012353264726698399, -0.28793570399284363, 0.2116449475288391, -0.2204311490058899, -0.07684657722711563, -0.06366005539894104, 0.10615354776382446, -0.11050670593976974, 0.20243307948112488, -0.11300783604383469, -0.011061429977416992, -0.08530132472515106, 0.10198746621608734, -0.6434175968170166, -0.012322584167122841, 0.05006060376763344, -0.11037357896566391, -0.19836683571338654, 0.01617954671382904, -0.37077611684799194, 0.10221119970083237, 0.18555739521980286, -0.3536912798881531, 0.28991395235061646, 0.12170452624559402, -0.4839293658733368, -0.03759793937206268, 0.13502652943134308, -0.38045117259025574, 0.3738781213760376, -0.009310918860137463, 0.620690107345581, 0.36739587783813477, 0.027000106871128082, -0.19940006732940674, -0.130784273147583, -0.016824116930365562, 0.14018140733242035, 0.2314874678850174, 0.3732298016548157, 0.06782659888267517, -0.0017938533565029502, 0.09110531210899353, 0.3433997631072998, 0.09555798023939133, -0.4527759552001953, 0.003211547154933214, 0.06937957555055618, -0.05693189427256584, 0.25175783038139343, 0.43132543563842773, -0.14016816020011902, 0.1018085852265358, 0.04392167925834656, 0.0013034043367952108, 0.2585051953792572, 0.008525178767740726, -0.3700050413608551, 0.3994216024875641, -0.12866757810115814, -0.10906583070755005, -0.023937424644827843, -0.45951345562934875, 0.501953661441803, 0.257312148809433, -0.12237625569105148, -0.21970732510089874, -0.02212701365351677, -0.20265279710292816, 0.18879015743732452, -0.27936312556266785, 0.1349678635597229, 0.22481398284435272, 0.22980988025665283, -0.028872894123196602, 0.23108641803264618, -0.24799348413944244, 0.18305692076683044, 0.16843485832214355, -0.07318955659866333, 0.10870401561260223, 0.03641321510076523, 0.13466492295265198, 0.2143716961145401, -0.042408961802721024, 0.08154222369194031, 0.2586716413497925, -0.149084210395813, -0.09810827672481537, 0.06033540144562721, 0.10254546254873276, -0.057776108384132385, 0.2976278066635132, -0.11873958259820938, 0.3306417465209961, 0.19747717678546906, 0.08186476677656174, 0.6335793733596802, 0.12999595701694489, 0.11188028752803802, 0.23062676191329956, -0.24835273623466492, 0.3714772164821625, -0.015770243480801582, -0.09920602291822433, -0.3051190972328186, 0.11255602538585663, 0.5123395323753357, 0.005374712869524956, -0.08825579285621643, -0.15170222520828247, 0.6344287395477295, -0.3277178406715393, -0.028070678934454918, -0.054426465183496475, 0.16650432348251343, 0.04648751765489578, -0.026701275259256363, -0.15159642696380615, -0.23069976270198822, -0.2512108087539673, -0.03176489472389221, -0.06145557761192322, -0.17437680065631866, 0.23778267204761505, -0.09511382132768631, -0.37860485911369324, -0.41486260294914246, 0.3110959231853485, 0.013640864752233028, 0.05072733387351036, 0.12941990792751312, -0.08883795142173767, -0.2585447132587433, -0.20693956315517426, -0.31599634885787964, 0.11958229541778564, 0.04853833466768265, -0.16896682977676392, 0.4002004861831665, 0.1946057230234146, 0.05183328315615654, -0.42120710015296936, -0.0655217245221138, -0.33028444647789, 0.41093724966049194, -0.24923910200595856, -0.03265449032187462, -0.4077798128128052, 0.3413347005844116, -0.08796700090169907, -0.2373380959033966, 0.15562397241592407, 0.09584078937768936, 0.24425263702869415, 0.1475476771593094, 0.25085684657096863, 0.17082229256629944, -0.015563084743916988, 0.15146291255950928, 0.11514436453580856, 0.15179680287837982, -0.10195362567901611, -0.017925098538398743, -0.06917615979909897, -0.19774216413497925, -0.36271923780441284, -0.06394775211811066, 0.25435981154441833, -0.07529687881469727, -0.1656830757856369, 0.10637582838535309, 0.04637614265084267, 0.1216297596693039, 0.4668866693973541, -0.0021339592058211565, 0.16076074540615082, -0.21306845545768738, 0.2391972839832306, -0.02989991381764412, -0.008168994449079037, -0.427844375371933, -0.1313236951828003, 0.11788056790828705, -0.33968299627304077, 0.03331360965967178, -0.19193311035633087, -0.22007182240486145, 0.24162185192108154, 0.42644333839416504, 0.06763817369937897, 0.28797388076782227, 0.668995201587677, -0.4047048091888428, -0.2716086506843567, -0.026163622736930847, -0.1488603949546814, 0.07346831262111664, -0.12028740346431732, 0.22896842658519745, -0.1964862197637558, 0.2728445827960968, 0.05373270809650421, -0.5707187652587891, -0.16195259988307953, 0.3896471858024597, -0.18537145853042603, 0.3044551908969879, -0.3331352472305298, 0.34249308705329895, -0.14319993555545807, -0.07530680298805237, -0.04317435622215271, 0.35457703471183777, -0.07736784219741821, 0.0737881287932396, -0.003418406005948782, 0.1910245567560196, 0.027461756020784378, 0.2754681706428528, 0.14740993082523346, 0.1160288155078888, -0.6654649376869202, 0.04085110127925873, -0.18821382522583008, 0.10560272634029388, 0.6328638195991516, -0.04359981790184975, 0.5376655459403992, -0.23090378940105438, -0.053518958389759064, -0.09368325769901276, -0.035687368363142014, -0.5783125758171082, -0.0974419042468071, -0.17652729153633118, -0.34248167276382446, 0.1198473870754242, 0.3132513463497162, 0.251362144947052, 0.07922787964344025, 0.38098201155662537, -0.08811616897583008, -0.13682541251182556, 0.05347563326358795, -0.19889743626117706, 0.2049666792154312, 0.35079893469810486, 0.09824243187904358, 0.1280089169740677, 0.062172479927539825, -0.08750792592763901, -0.15906976163387299, -0.2192315012216568, -0.044837381690740585, -0.2752540707588196, -0.06824228912591934, 0.08822368085384369, 0.10791503638029099, -0.4533323347568512, 0.04257060959935188, 0.18984092772006989, -0.1333230882883072, 0.005950800143182278, 0.021368365734815598, -0.2756708860397339, -0.050660666078329086, 0.15291298925876617, 0.024264957755804062, -0.46145209670066833, -0.2749907374382019, 0.18971268832683563, 0.2158459722995758, -0.07679734379053116, -0.12300625443458557, -0.27498659491539, 0.03125789761543274, 0.017463894560933113, -0.13913686573505402, -0.08544687926769257, 0.4624566435813904, -0.09982538968324661, -0.007967834360897541, -0.21223482489585876, 0.11483140289783478, -0.08859793096780777, -0.3865339756011963, 0.22931402921676636, -0.03084983117878437, -0.1217142790555954, -0.0714973658323288, 0.0913734883069992, 0.031692419201135635, 0.02740675024688244, -0.14807617664337158, 0.2933056056499481, 0.07259352505207062, 0.150038942694664, 0.20389658212661743, -0.2095879167318344, 0.23059092462062836, -0.11141899228096008, 0.2624780237674713, -0.15074850618839264, -0.2476958930492401, 0.07190576940774918, -0.28707557916641235, -0.2841626703739166, 0.010203353129327297, -0.03580302745103836, -0.022564366459846497, 0.28507155179977417, 0.006695019081234932, 0.13581916689872742, 0.46150973439216614, -0.16186277568340302, 0.03777439519762993, -0.030225878581404686, -0.6401489973068237, 0.30239227414131165, 0.24481503665447235, -0.2169131487607956, -0.34458643198013306, -0.045895546674728394, 0.22673271596431732, 0.1682630479335785, -0.0390043668448925, -0.20760641992092133, 0.16144004464149475, -0.18018151819705963, -0.18361453711986542, -0.017317038029432297, -0.191421240568161, -0.14025135338306427, 0.07902661710977554, 0.0754220113158226, 0.21752332150936127, -0.012803039513528347, 0.0774211511015892, 0.5890892744064331, -0.17607739567756653, 0.27382269501686096, 0.24932193756103516, 0.056992318481206894, -0.10868872702121735, -0.2739160358905792, 0.0649508610367775, -0.22561705112457275, -0.08291907608509064, 0.07544154673814774, -0.049379654228687286, -0.31352803111076355, 0.04076319932937622, 0.0549563430249691, 0.37105947732925415, 0.030625058338046074, 0.21595090627670288, -0.7079184055328369, -0.1802363097667694, 0.09123457968235016, 0.21302811801433563, -0.20670568943023682, 0.2741404175758362, -0.024637417867779732, -0.19757501780986786, -0.3012378513813019, 0.4057568609714508, 0.03876390680670738, 0.16843172907829285, -0.0370790921151638, -0.28847476840019226, 0.3370668590068817, -0.019235806539654732, -0.026341354474425316, 0.039102137088775635, 0.021336672827601433, 0.3399619162082672, 0.008782588876783848, -0.0555642731487751, -0.015803012996912003, -0.20690904557704926, 0.06078285723924637, 0.23993942141532898, -0.10998759418725967, 0.20169198513031006, 0.20720556378364563, 0.12724320590496063, -0.04475519806146622, -0.3302326798439026, 0.09620943665504456, -0.07248255610466003, -0.1277698576450348, 0.22776609659194946, -0.5702326893806458, -0.16125471889972687, -0.0926797091960907, -0.23591458797454834, 0.18719488382339478, -0.5090630054473877, 8.624705515103415e-05, -0.13772189617156982, -0.2471744567155838, -0.12411326169967651, -0.09268904477357864, -0.13948358595371246, 0.20179636776447296, -0.08524119108915329, -0.2287042737007141, -0.09486792981624603, 0.39887794852256775, -0.12004146724939346, 0.09652216732501984, -0.30567285418510437, 0.1634252965450287, 0.21806053817272186, 0.25786036252975464, -0.08417965471744537, 0.03463392332196236, -0.2978096008300781, 0.2922680079936981, 0.09859910607337952, 0.0875714123249054, -0.4012094736099243, -0.1180829331278801, -0.20086008310317993, 0.3827769160270691, -0.02966497652232647, 0.12432833760976791, 0.2163151651620865, 0.1049681082367897, 0.2804584205150604, 0.2993391454219818, 0.060952529311180115, -0.07689746469259262, -0.28560805320739746, -0.0924665629863739, 0.2123049944639206, 0.34827154874801636, 0.3108852207660675, 0.20688988268375397, -0.15706351399421692, -0.0801946371793747, 0.17795152962207794, -0.3054390847682953, 0.33463284373283386, 0.15933990478515625, 0.3370025157928467, -0.11468599736690521, -0.4142603874206543, 0.43951913714408875, 0.040413618087768555, -0.36411213874816895, 0.12548542022705078, 0.039100080728530884, -0.05175444111227989, -0.013089758343994617, 0.09843230992555618, -0.2279045581817627, -0.08150480687618256, -0.05419597774744034, -0.23237672448158264, -0.11533608287572861, 0.1469910442829132, 0.016510535031557083, 0.087582066655159, -0.07245096564292908, -0.007324876729398966, -0.27845925092697144, 0.4279777407646179, 0.08992302417755127, -0.1539613902568817, 0.46780118346214294, -0.020996809005737305, -0.07692563533782959, 0.06766723096370697, 0.09129513800144196, -0.09024716913700104, -0.05205707624554634, 0.12949904799461365, 0.3138450086116791, 0.2147463858127594, -0.4237498641014099, -0.36574897170066833, 0.1590256243944168, 0.0716824159026146, 0.1964207887649536, 0.08330993354320526, 0.18056462705135345, 0.4125373661518097, 0.03199251741170883, -0.014201803132891655, -0.04786520451307297, -0.1973641812801361, 0.08327662199735641, 0.18733075261116028, -0.1196286603808403, -0.13152454793453217, 0.1972564458847046, -0.11437661200761795, 0.07769683748483658, -0.21692024171352386, -0.21537750959396362, 0.07370525598526001, -0.09096450358629227, 0.24400578439235687, -0.04425204172730446, -0.001715398859232664, -0.5019264817237854, 0.03904663026332855, -0.31542330980300903, -0.1955813467502594, 0.17944589257240295, 0.3158412575721741, -0.051372647285461426, -0.1406637579202652, -0.07807744294404984, 0.07771122455596924, 0.38187727332115173, -0.12212510406970978, -0.02049827016890049, -0.15709830820560455, -0.05735661834478378, -0.03108682855963707, -0.18736524879932404, 0.007145440671592951, -0.14526218175888062, -0.36052238941192627, 0.06165940687060356, -0.22556337714195251, 0.12641359865665436, -0.0632370337843895, -0.3335878551006317, 0.3422716557979584, -0.1660255491733551, 0.003225923515856266, -0.2865915298461914, -0.05257180705666542, 0.2598038613796234, 0.07454393804073334, -0.18490248918533325, 0.11369923502206802, -0.4055955410003662, -0.05017711967229843, 0.2790580689907074, 0.0707625225186348, 0.19450512528419495, -0.2669900953769684, -0.23600581288337708, -0.09219469130039215, -0.001194056705571711, 0.6422613859176636, -0.10753321647644043, 0.2948673367500305, 0.3371433615684509, -0.25945842266082764, -0.39526134729385376, 0.43433552980422974, 0.057605091482400894, -0.06400303542613983, 0.06513926386833191, 0.466957151889801, 0.22010594606399536, 0.09701237827539444, 0.6181700825691223, -0.025219671428203583, -0.18128883838653564, -0.16145087778568268, -0.017693834379315376, -0.25732651352882385, 0.12695005536079407, -0.09514722228050232, -0.1678393930196762, 0.07769102603197098, 0.10048671811819077, 0.3448786437511444, -0.21441207826137543, 0.032956186681985855, -0.34543633460998535, 0.4631159007549286, 0.16866780817508698, -0.16043764352798462, -0.07552498579025269, -0.010367482900619507, -0.06890632957220078, 0.16693125665187836, -0.3656156361103058, -0.018262360244989395, 0.4490826427936554, -0.07620551437139511, 0.10491639375686646, -0.37003570795059204, -0.5146818161010742, -0.1526218205690384, -0.05699790269136429, 0.27691319584846497, -0.04207991436123848, 0.0035703186877071857, 0.15147139132022858, 0.06838071346282959, -0.045601729303598404, -0.4753073751926422, 0.22643736004829407, 0.07241035252809525, 0.14090555906295776, 0.37025174498558044, -0.04914121329784393, -0.05388611555099487, 0.19829817116260529, -0.07414809614419937, 0.2852654457092285, -0.17624127864837646, -0.46711695194244385, -0.13952450454235077, -0.09475569427013397, 0.21062198281288147, -0.13872599601745605, -0.20019495487213135, -0.38410457968711853, 0.011939293704926968, 0.49206116795539856, 0.17529694736003876, 0.24048827588558197, -0.19081276655197144, -0.030594680458307266, 0.0771908089518547, -0.17978617548942566, 0.06187446042895317, -0.3235628008842468, -0.1415555328130722, -0.021755022928118706, -0.19942909479141235, 0.21415609121322632, -0.025940844789147377, -0.0224665068089962, -0.051242727786302567, 0.09434128552675247, 0.03880259394645691, 0.08484309166669846, -0.13915029168128967, -0.09709453582763672, -0.09798453748226166, -0.20233628153800964, -0.07568106055259705, 0.2830343544483185, 0.026006391271948814, -0.09774927794933319, 0.09761378914117813, -0.16060298681259155, 0.0883295089006424, -0.1172109916806221, -0.018764888867735863, 0.028517430648207664, 0.14510777592658997, 0.24448107182979584, -0.08580934256315231, -0.028442105278372765, -0.44492220878601074, 0.06536547839641571, 0.42115411162376404, -0.2581290304660797, 0.16185811161994934, 0.2784891426563263, 0.25662800669670105, -0.18918468058109283, -0.16250449419021606, -0.14768941700458527, -0.11208032816648483, 0.1832655817270279, 0.18813066184520721, 0.18338118493556976, 0.03502741456031799, 0.14037439227104187, 0.1759331375360489, -0.0060289883986115456, -0.20951054990291595, -0.13219277560710907, 0.008557484485208988, -0.11949580162763596, 0.18252050876617432, 0.024084098637104034, -0.36383217573165894, -0.45233598351478577, -0.21205833554267883, -0.19642876088619232, -0.307856023311615, 0.21651718020439148, -0.0996658056974411, -0.05421178787946701, -0.12564200162887573, -0.062369607388973236, 0.4586743712425232, 0.017445553094148636, 0.4289401173591614, 0.5302051305770874, 0.18625003099441528, -0.08875242620706558, 0.11472655832767487, 0.01903630420565605, -0.06652329862117767, 0.04081876948475838, -0.3547857105731964, -0.05746148154139519, -0.10648404806852341, -0.08411262929439545, 0.21160826086997986, -0.14399567246437073, -0.001533661619760096, 0.402569055557251, 0.23118293285369873, -0.42335113883018494, 0.15734195709228516, 0.2807837724685669, -0.18674631416797638, -0.07638344913721085, -0.09469461441040039, 0.2996332347393036, -0.2397388219833374, -0.2524983882904053, 0.07910165935754776, 0.22440893948078156, -0.24777723848819733, 0.2518589198589325, 0.062079012393951416, -0.08729743212461472, -0.18292109668254852, 0.3681671619415283, -0.25812309980392456, 0.2777676582336426, 0.1170569434762001, -0.12515251338481903, 0.30695345997810364, 0.37907880544662476, -0.2951100468635559, 0.24224039912223816, 0.0562015101313591, -0.15310490131378174, 0.16725729405879974, -0.16514886915683746, 0.05014282837510109, 0.047883979976177216, 0.030458010733127594, 0.337740033864975, -0.09268610179424286, -0.10392095893621445, -0.1853087842464447, -0.10743589699268341, -0.08417680114507675, -0.14629536867141724, 0.28254029154777527, 0.08360442519187927, -0.29671624302864075, -0.22271020710468292, -0.2068423628807068, -0.1879844069480896, -0.33452680706977844, 0.05302631855010986, -0.015878498554229736, 0.4459949731826782, 0.6217730045318604, -0.31203997135162354, 0.1947515606880188, -0.22410471737384796, 0.06667597591876984, -0.38273078203201294, -0.14765357971191406, 0.09039947390556335, 0.05403902009129524, 0.21308015286922455, -0.015996692702174187, 0.3191845118999481, -0.2492191195487976, 0.15279227495193481, 0.3262900710105896, -0.24421937763690948, 0.14222607016563416, -0.08758469671010971, 0.05400446057319641, 0.0044992780312895775, -0.43212568759918213, -0.06014556437730789, -0.15350529551506042, -0.30881109833717346, 0.684353232383728, 0.15833599865436554, 0.150883287191391, 0.19537995755672455, -0.0714106336236, 0.0024355920031666756, 0.23413798213005066, -0.06315486878156662, -0.17479084432125092, 0.15844647586345673, -0.036659203469753265, 0.3946676552295685, 0.0006852307124063373, -0.30835700035095215, -0.08664325624704361, 0.19467194378376007, -0.17225390672683716, 0.10923860222101212, 0.3232629597187042, -0.004813152365386486, -0.03592580184340477, 0.3495360314846039, -0.09980570524930954, 0.4423748850822449, 0.27199676632881165, 0.13414743542671204, 0.20537199079990387, 0.5442022681236267, 0.10399258881807327, 0.2637161314487457, 0.1376948356628418, -0.154172882437706, 0.12310938537120819, -0.38073065876960754, 0.26749661564826965, 0.04061388596892357, -0.14770062267780304, -0.11814047396183014, 0.09793497622013092, -0.3438791334629059, 0.44314736127853394, 0.003363360883668065, 0.09836671501398087, 0.19237782061100006, 0.04710410162806511, 0.12668199837207794, 0.39392587542533875, 0.3594183325767517, -0.2144288271665573, -0.2674466669559479, -0.07416475564241409, 0.17704449594020844, -0.11840175092220306, 0.6387115120887756, -0.3668268322944641, 0.35043808817863464, 0.2349214404821396, 0.10484906286001205, 0.0018725874833762646, 0.23220402002334595, 0.12798872590065002, -0.4888915419578552, -0.006456371862441301, -0.04029998183250427, -0.5537481904029846, 0.035348761826753616, -0.010908723808825016, -0.23117420077323914, -0.17765213549137115, -0.21690583229064941, 0.12727268040180206, -0.12442623823881149, 0.3653775751590729, -0.13992388546466827, 0.14835986495018005, -0.10545650124549866, 0.00677597476169467, -0.014606463722884655, -0.20343326032161713, 0.14445710182189941, -0.1322217434644699, 0.1730964481830597, 0.07335352152585983, 0.23473715782165527, -0.1697155088186264, 0.21863339841365814, 0.04846953973174095, 0.30712586641311646, -0.10963138192892075, 0.0010880609042942524, -0.11430948227643967, -0.2796919047832489, -0.08629263937473297, 0.0771900936961174, 0.091727115213871, -0.3595750033855438, 0.26998239755630493, 0.34682390093803406, 0.24977745115756989, -0.3160820007324219, -0.06467965990304947, -0.1672583669424057, 0.08543020486831665, -0.2589549124240875, 0.25261378288269043, 0.20616373419761658, -0.4576788544654846, 0.1715049147605896, 0.037419918924570084, -0.004635381046682596, 0.2325759381055832, 0.222879096865654, -0.40177658200263977, 0.005732292775064707, -0.3574826717376709, -0.15282116830348969, 0.2869740128517151, 0.20015020668506622, -0.012383422814309597, 0.14824539422988892, -0.16018062829971313, 0.07428795099258423, -0.11738135665655136, 0.2022092491388321, -0.10460886359214783, 0.2642272710800171, -0.2035832703113556, -0.21699115633964539, 0.3207097053527832, -0.22375233471393585, -0.02517789602279663, 0.11197695881128311, 0.1935439109802246, -0.2135145217180252, -0.3005884289741516, 0.0554325096309185, 0.24299226701259613, 0.15143999457359314, -0.00042552861850708723, -0.21006333827972412, 0.29073581099510193, -0.060575149953365326, 0.6028258800506592, 0.36874571442604065, -0.2577216625213623, 0.3448943495750427, -0.004511699546128511, 0.20920781791210175, -0.2354143261909485, -0.28416579961776733, 0.02212844230234623, 0.12905481457710266, -0.22686147689819336, 0.10779892653226852, 0.2333153933286667, 0.07520298659801483, -0.1544494777917862, -0.29509440064430237, 0.2116313874721527, -0.07793357223272324, 0.07072737067937851, 0.29300782084465027, -0.24422629177570343, 0.26565271615982056, -0.041896749287843704, 0.044637031853199005, 0.5018559694290161, 0.15858100354671478, 0.3866262137889862, 0.33520352840423584, 0.17222119867801666, 0.10377871245145798, 0.3036658763885498, 0.0032014460302889347, 0.6580135226249695, 0.02165387198328972, -0.10519301891326904, 0.20042367279529572, -0.15982472896575928, 0.0032899542711675167, 0.17089803516864777, -0.2411164939403534, 0.30369752645492554, -0.24009795486927032, 0.11853630840778351, 0.24648228287696838, 0.035253364592790604, 0.13861504197120667, -0.19350986182689667, -0.05480891466140747, -0.006002613343298435, -0.2968408465385437, 0.01168200932443142, -0.08150807023048401, -0.48666301369667053, 0.002623066073283553, -0.35812386870384216, -0.011602930724620819, -0.012528179213404655]}