docs: Phase T9 — export to xserv

Architecture diff table (xtrain TinyTransformer vs xserv qwen3.rs), the
QK-norm structural decision + BF16 acceptance criterion, the tensor-name +
layout mapping table, and the dash5 closed-loop verification recipe.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Cargo.lock

@@ -109,6 +109,16 @@ version = "0.8.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
 
+[[package]]
+name = "safetensors"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cc0cdb7198d738a111f6df8fef42cb175412c311d0c4ac9126ff4e550ad1a0e8"
+dependencies = [
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "serde"
 version = "1.0.228"
@@ -249,6 +259,8 @@ dependencies = [
 name = "xtrain-train"
 version = "0.1.0"
 dependencies = [
+ "half",
+ "safetensors",
  "xserv-tokenizer",
  "xtrain-autodiff",
  "xtrain-cuda",

crates/xtrain-model/src/config.rs

@@ -42,6 +42,7 @@ impl Config {
     /// Total learnable parameter count (for logging / sanity).
     pub fn num_params(&self) -> usize {
         let per_layer = 2 * self.dim                       // 2 rmsnorm gammas
+            + 2 * self.head_dim                             // q/k per-head norm gammas
             + 3 * self.dim * self.dim                       // q/k/v proj
             + self.dim * self.dim                           // out proj
             + 2 * self.dim * self.ffn_hidden                // gate/up proj

crates/xtrain-model/src/lib.rs

@@ -2,9 +2,10 @@
 //!
 //! A from-scratch decoder built entirely from the [`xtrain_autodiff`] op set:
 //! token embedding → `n_layers` × {pre-RMSNorm → multi-head causal attention
-//! (RoPE) → residual; pre-RMSNorm → SwiGLU MLP → residual} → final RMSNorm →
-//! LM-head matmul. The forward builds an autograd graph; calling `.backward()`
-//! on the cross-entropy loss fills every parameter's `.grad()`.
+//! (per-head QK-norm + RoPE) → residual; pre-RMSNorm → SwiGLU MLP → residual} →
+//! final RMSNorm → LM-head matmul. The forward builds an autograd graph; calling
+//! `.backward()` on the cross-entropy loss fills every parameter's `.grad()`.
+//! Per-head QK-norm (Qwen3-style) makes the architecture xserv-compatible (T9).
 //!
 //! Conventions (matching the engine, not HuggingFace):
 //! - Linear weights are `[in, out]` and applied as `x @ W` (no transpose), since

crates/xtrain-model/src/model.rs

@@ -13,6 +13,8 @@ struct Block {
     wq: Var,        // [dim, dim]
     wk: Var,        // [dim, dim]
     wv: Var,        // [dim, dim]
+    q_norm: Var,    // [head_dim] — per-head QK-norm (Qwen3-style)
+    k_norm: Var,    // [head_dim]
     wo: Var,        // [dim, dim]
     ffn_norm: Var,  // [dim]
     w_gate: Var,    // [dim, ffn_hidden]
@@ -52,6 +54,8 @@ impl TinyTransformer {
                 wq: mk(&[cfg.dim, cfg.dim]),
                 wk: mk(&[cfg.dim, cfg.dim]),
                 wv: mk(&[cfg.dim, cfg.dim]),
+                q_norm: mk(&[cfg.head_dim]),
+                k_norm: mk(&[cfg.head_dim]),
                 wo: mk(&[cfg.dim, cfg.dim]),
                 ffn_norm: mk(&[cfg.dim]),
                 w_gate: mk(&[cfg.dim, cfg.ffn_hidden]),
@@ -87,6 +91,8 @@ impl TinyTransformer {
                 b.wq.clone(),
                 b.wk.clone(),
                 b.wv.clone(),
+                b.q_norm.clone(),
+                b.k_norm.clone(),
                 b.wo.clone(),
                 b.ffn_norm.clone(),
                 b.w_gate.clone(),
@@ -136,22 +142,29 @@ impl TinyTransformer {
         // Project, then lay out as per-head [seq, head_dim] tensors.
         //   [seq,dim] @ [dim,dim] = [seq,dim]
         //   reshape   [seq, nh, hd]
+        //   qk-norm   per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
         //   rope      (kernel expects exactly [tokens, heads, head_dim])
         //   transpose [nh, seq, hd]  → split into nh × [seq, hd]
-        let to_heads = |proj: Var, rope: bool| -> Vec<Var> {
+        let to_heads = |proj: Var, norm: Option<&Var>| -> Vec<Var> {
             let r = ops::reshape(&proj, &[seq, nh, hd]);
-            let r = if rope {
-                ops::rope(&r, self.cfg.rope_theta)
-            } else {
-                r
+            let r = match norm {
+                // Per-head RMSNorm: flatten the (seq,nh) head rows, norm over hd,
+                // restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
+                Some(gamma) => {
+                    let flat = ops::reshape(&r, &[seq * nh, hd]);
+                    let normed = ops::rms_norm(&flat, gamma, self.cfg.eps);
+                    let r = ops::reshape(&normed, &[seq, nh, hd]);
+                    ops::rope(&r, self.cfg.rope_theta)
+                }
+                None => r,
             };
             let t = ops::transpose_3d01(&r); // [nh, seq, hd]
             ops::split_heads(&t)
         };
 
-        let q = to_heads(ops::matmul(x, &b.wq), true);
-        let k = to_heads(ops::matmul(x, &b.wk), true);
-        let v = to_heads(ops::matmul(x, &b.wv), false);
+        let q = to_heads(ops::matmul(x, &b.wq), Some(&b.q_norm));
+        let k = to_heads(ops::matmul(x, &b.wk), Some(&b.k_norm));
+        let v = to_heads(ops::matmul(x, &b.wv), None);
 
         // Per-head scaled-dot-product attention with causal mask.
         let heads_out: Vec<Var> = (0..nh)

crates/xtrain-model/tests/parity.py

@@ -98,7 +98,7 @@ lm_head = load("lm_head")
 layers = []
 for l in range(NL):
     layers.append({p: load(f"l{l}_{p}") for p in
-                   ["attn_norm", "wq", "wk", "wv", "wo",
+                   ["attn_norm", "wq", "wk", "wv", "q_norm", "k_norm", "wo",
                     "ffn_norm", "w_gate", "w_up", "w_down"]})
 
 idx = torch.tensor(ids, dtype=torch.long)
@@ -111,6 +111,9 @@ for L in layers:
     q = (x @ L["wq"]).reshape(SEQ, NH, HD)
     k = (x @ L["wk"]).reshape(SEQ, NH, HD)
     v = (x @ L["wv"]).reshape(SEQ, NH, HD)
+    # Per-head QK-norm (Qwen3-style), before RoPE.
+    q = rms_norm(q, L["q_norm"])
+    k = rms_norm(k, L["k_norm"])
     q = rope(q).transpose(0, 1)  # [nh, seq, hd]
     k = rope(k).transpose(0, 1)
     v = v.transpose(0, 1)

crates/xtrain-model/tests/parity_dump.rs

@@ -145,6 +145,8 @@ fn param_names(cfg: &Config) -> Vec<String> {
             "wq",
             "wk",
             "wv",
+            "q_norm",
+            "k_norm",
             "wo",
             "ffn_norm",
             "w_gate",

crates/xtrain-train/Cargo.toml

@@ -14,7 +14,14 @@ xtrain-cuda = { path = "../xtrain-cuda" }
 # crate inherits xserv's workspace for its own deps (serde/regex) — Cargo reads
 # the target package's workspace, not ours.
 xserv-tokenizer = { path = "../../../xserv/crates/xserv-tokenizer" }
+# T9 export to xserv: HF Qwen3 safetensors + BF16 weight cast.
+half.workspace = true
+safetensors = "0.5"
 
 [[bin]]
 name = "train"
 path = "src/bin/train.rs"
+
+[[bin]]
+name = "export_safetensors"
+path = "src/bin/export_safetensors.rs"

crates/xtrain-train/src/bin/export_safetensors.rs

@@ -0,0 +1,258 @@
+//! Phase T9 — export a trained xtrain checkpoint into the format xserv loads:
+//! an HF Qwen3-style `config.json` + `model.safetensors` (+ a copy of the GPT-2
+//! `tokenizer.json`), so xserv's `Qwen3` loader can serve the same weights.
+//!
+//! xtrain's `TinyTransformer` is (after T9) architecturally a tiny Qwen3:
+//! RoPE (rotate_half, pos=row) + RMSNorm + per-head QK-norm + SwiGLU + separate
+//! lm_head, MHA (n_kv_heads = n_heads). The only deltas to xserv are mechanical:
+//!   - tensor NAMES  → HF Qwen3 names (`model.layers.{i}.self_attn.q_proj.weight` …)
+//!   - 2D proj LAYOUT → xtrain stores `[in,out]` (computes `x@W`); xserv/HF want
+//!     `[out,in]` (computes `x@Wᵀ`) → transpose every 2D projection weight.
+//!     1D norms and the `[vocab,dim]` embedding/lm_head rows are unchanged.
+//!   - DTYPE → xserv's Qwen3 forward is BF16-only, so weights are written as BF16.
+//!
+//! See `docs/08-export-xserv.md` for the full architecture diff + mapping table.
+//!
+//! Run on dash5 (needs a GPU to materialise the checkpoint params):
+//!   export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
+//!   cargo run -p xtrain-train --release --bin export_safetensors -- \
+//!       /tmp/xtrain_tinystories.ckpt \
+//!       /opt/wjh/models/gpt2/tokenizer.json \
+//!       /tmp/xtrain_export
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("export_safetensors: built without CUDA (no_cuda); run on a GPU host (dash5).");
+}
+
+#[cfg(not(no_cuda))]
+use std::path::{Path, PathBuf};
+
+#[cfg(not(no_cuda))]
+use half::bf16;
+#[cfg(not(no_cuda))]
+use xtrain_autodiff::tape::Var;
+#[cfg(not(no_cuda))]
+use xtrain_cuda::device;
+#[cfg(not(no_cuda))]
+use xtrain_model::{Config, TinyTransformer, param_to_host};
+#[cfg(not(no_cuda))]
+use xtrain_tensor::Device;
+
+// Same deterministic init scheme as bin/train.rs, so a freshly-built model has
+// the right shapes before `load_into` overwrites the values from the checkpoint.
+#[cfg(not(no_cuda))]
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+/// A param ready to serialize: HF name + the (possibly transposed) row-major
+/// data + its shape. Stored as BF16 (xserv's Qwen3 forward is BF16-only).
+#[cfg(not(no_cuda))]
+struct Export {
+    name: String,
+    data: Vec<bf16>,
+    shape: Vec<usize>,
+}
+
+/// 1D norm / embedding row-table: keep layout, just cast to BF16.
+#[cfg(not(no_cuda))]
+fn keep(name: &str, v: &Var) -> Export {
+    let host = param_to_host(v);
+    let shape = v.value().shape().to_vec();
+    Export {
+        name: name.to_string(),
+        data: host.iter().map(|&x| bf16::from_f32(x)).collect(),
+        shape,
+    }
+}
+
+/// 2D projection weight: xtrain `[in,out]` (x@W) → HF `[out,in]` (x@Wᵀ). Transpose
+/// the row-major matrix and cast to BF16.
+#[cfg(not(no_cuda))]
+fn transpose(name: &str, v: &Var) -> Export {
+    let host = param_to_host(v);
+    let shape = v.value().shape().to_vec();
+    assert_eq!(shape.len(), 2, "transpose expects a 2D weight: {name}");
+    let (rows, cols) = (shape[0], shape[1]); // [in, out]
+    let mut out = vec![bf16::ZERO; rows * cols];
+    for r in 0..rows {
+        for c in 0..cols {
+            // out[c, r] = in[r, c]
+            out[c * rows + r] = bf16::from_f32(host[r * cols + c]);
+        }
+    }
+    Export {
+        name: name.to_string(),
+        data: out,
+        shape: vec![cols, rows], // [out, in]
+    }
+}
+
+/// Assemble every export tensor in HF Qwen3 naming, reading the xtrain params in
+/// their stable `params()` order:
+///   embed → per block [attn_norm, wq, wk, wv, q_norm, k_norm, wo, ffn_norm,
+///                      w_gate, w_up, w_down] → final_norm → lm_head
+#[cfg(not(no_cuda))]
+fn build_exports(model: &TinyTransformer) -> Vec<Export> {
+    let cfg = model.config();
+    let p = model.params();
+    let mut it = p.iter();
+    let mut next = || it.next().expect("params() ran short");
+
+    let mut ex = Vec::new();
+    ex.push(keep("model.embed_tokens.weight", next())); // [vocab, dim]
+    for l in 0..cfg.n_layers {
+        let b = format!("model.layers.{l}");
+        ex.push(keep(&format!("{b}.input_layernorm.weight"), next()));
+        ex.push(transpose(&format!("{b}.self_attn.q_proj.weight"), next()));
+        ex.push(transpose(&format!("{b}.self_attn.k_proj.weight"), next()));
+        ex.push(transpose(&format!("{b}.self_attn.v_proj.weight"), next()));
+        ex.push(keep(&format!("{b}.self_attn.q_norm.weight"), next()));
+        ex.push(keep(&format!("{b}.self_attn.k_norm.weight"), next()));
+        ex.push(transpose(&format!("{b}.self_attn.o_proj.weight"), next()));
+        ex.push(keep(
+            &format!("{b}.post_attention_layernorm.weight"),
+            next(),
+        ));
+        ex.push(transpose(&format!("{b}.mlp.gate_proj.weight"), next()));
+        ex.push(transpose(&format!("{b}.mlp.up_proj.weight"), next()));
+        ex.push(transpose(&format!("{b}.mlp.down_proj.weight"), next()));
+    }
+    ex.push(keep("model.norm.weight", next())); // [dim]
+    ex.push(transpose("lm_head.weight", next())); // [dim,vocab] → [vocab,dim]
+    assert!(it.next().is_none(), "params() had extra tensors");
+    ex
+}
+
+/// config.json matching xserv's `ModelConfig` for a Qwen3 with xtrain's dims and
+/// reconciled fields (eps, rope theta, head_dim, n_kv_heads = n_heads, untied).
+#[cfg(not(no_cuda))]
+fn config_json(cfg: &Config) -> String {
+    format!(
+        r#"{{
+  "architectures": ["Qwen3ForCausalLM"],
+  "model_type": "qwen3",
+  "vocab_size": {vocab},
+  "hidden_size": {dim},
+  "intermediate_size": {ffn},
+  "num_hidden_layers": {layers},
+  "num_attention_heads": {heads},
+  "num_key_value_heads": {kv_heads},
+  "head_dim": {head_dim},
+  "max_position_embeddings": 2048,
+  "rms_norm_eps": {eps},
+  "rope_theta": {theta},
+  "tie_word_embeddings": false,
+  "attention_bias": false,
+  "hidden_act": "silu"
+}}
+"#,
+        vocab = cfg.vocab,
+        dim = cfg.dim,
+        ffn = cfg.ffn_hidden,
+        layers = cfg.n_layers,
+        heads = cfg.n_heads,
+        kv_heads = cfg.n_heads, // xtrain is MHA → kv heads == query heads
+        head_dim = cfg.head_dim,
+        eps = cfg.eps,
+        theta = cfg.rope_theta,
+    )
+}
+
+#[cfg(not(no_cuda))]
+fn main() {
+    use safetensors::tensor::{Dtype, TensorView};
+    use xserv_tokenizer::Tokenizer;
+
+    let args: Vec<String> = std::env::args().collect();
+    let ckpt = args
+        .get(1)
+        .map(PathBuf::from)
+        .unwrap_or_else(|| PathBuf::from("/tmp/xtrain_tinystories.ckpt"));
+    let tok_path = args
+        .get(2)
+        .map(PathBuf::from)
+        .unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
+    let out_dir = args
+        .get(3)
+        .map(PathBuf::from)
+        .unwrap_or_else(|| PathBuf::from("/tmp/xtrain_export"));
+
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let dev = Device::Cuda(0);
+
+    // Size the model exactly like bin/train.rs: gpt2 vocab + n_layers = 4.
+    let tok = Tokenizer::from_file(&tok_path);
+    let vocab = tok.vocab_size();
+    let mut cfg = Config::tiny();
+    cfg.vocab = vocab;
+    cfg.n_layers = 4;
+    println!(
+        "export: ckpt {} → {} (vocab {}, dim {}, layers {}, heads {}, head_dim {})",
+        ckpt.display(),
+        out_dir.display(),
+        cfg.vocab,
+        cfg.dim,
+        cfg.n_layers,
+        cfg.n_heads,
+        cfg.head_dim,
+    );
+
+    let mut seed = 1u64;
+    let model = TinyTransformer::new(cfg, dev, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.04)
+        }
+    });
+    xtrain_train::checkpoint::load_into(&ckpt, &model.params()).expect("load checkpoint");
+
+    let exports = build_exports(&model);
+    println!("export: {} tensors", exports.len());
+
+    // Serialize to safetensors. Each TensorView borrows the raw BF16 bytes.
+    let views: Vec<(String, TensorView)> = exports
+        .iter()
+        .map(|e| {
+            let bytes = unsafe {
+                std::slice::from_raw_parts(e.data.as_ptr() as *const u8, e.data.len() * 2)
+            };
+            let view = TensorView::new(Dtype::BF16, e.shape.clone(), bytes)
+                .unwrap_or_else(|err| panic!("bad tensor view {}: {err}", e.name));
+            (e.name.clone(), view)
+        })
+        .collect();
+
+    std::fs::create_dir_all(&out_dir).expect("mkdir out_dir");
+    let st = safetensors::tensor::serialize(views.iter().map(|(n, v)| (n.as_str(), v)), &None)
+        .expect("serialize safetensors");
+    std::fs::write(out_dir.join("model.safetensors"), st).expect("write model.safetensors");
+    std::fs::write(out_dir.join("config.json"), config_json(&cfg)).expect("write config.json");
+    copy_tokenizer(&tok_path, &out_dir);
+
+    println!(
+        "export: wrote config.json + model.safetensors + tokenizer.json to {}",
+        out_dir.display()
+    );
+}
+
+/// Place the tokenizer beside the weights so xserv loads it from the model dir.
+#[cfg(not(no_cuda))]
+fn copy_tokenizer(tok_path: &Path, out_dir: &Path) {
+    std::fs::copy(tok_path, out_dir.join("tokenizer.json")).expect("copy tokenizer.json");
+}

crates/xtrain-train/tests/adamw_parity.py

@@ -74,7 +74,7 @@ SEQ = len(ids)
 
 NAMES = ["embed"]
 for l in range(NL):
-    for p in ["attn_norm", "wq", "wk", "wv", "wo",
+    for p in ["attn_norm", "wq", "wk", "wv", "q_norm", "k_norm", "wo",
               "ffn_norm", "w_gate", "w_up", "w_down"]:
         NAMES.append(f"l{l}_{p}")
 NAMES += ["final_norm", "lm_head"]
@@ -115,6 +115,9 @@ def forward():
         q = (x @ P[f"l{l}_wq"]).reshape(SEQ, NH, HD)
         k = (x @ P[f"l{l}_wk"]).reshape(SEQ, NH, HD)
         v = (x @ P[f"l{l}_wv"]).reshape(SEQ, NH, HD)
+        # Per-head QK-norm (Qwen3-style), before RoPE.
+        q = rms_norm(q, P[f"l{l}_q_norm"])
+        k = rms_norm(k, P[f"l{l}_k_norm"])
         q = rope(q).transpose(0, 1)
         k = rope(k).transpose(0, 1)
         v = v.transpose(0, 1)

crates/xtrain-train/tests/adamw_parity_dump.rs

@@ -156,6 +156,8 @@ fn param_names(cfg: &Config) -> Vec<String> {
             "wq",
             "wk",
             "wv",
+            "q_norm",
+            "k_norm",
             "wo",
             "ffn_norm",
             "w_gate",

docs/08-export-xserv.md

@@ -0,0 +1,138 @@
+# Phase T9: Export to xserv — Design Document
+
+## Goal
+
+闭环 xtrain ↔ xserv：把 xtrain 练出的权重导成 **xserv 的 Qwen3 loader 能直接加载并服务**的格式
+（HF 命名的 `config.json` + `model.safetensors` + 复用的 gpt2 `tokenizer.json`），让推理侧 xserv
+跑出**与 xtrain 自身一致**的生成结果。
+
+验收信号：**同一份权重 + 同一 prompt，xserv 的贪心生成 token 序列对住 xtrain 的贪心生成**
+（logits 在浮点容差内吻合）。这是整条 P0→P6 学习链的收口——训练栈练出来的东西，推理栈真的能用。
+
+> 范围与诚实边界：**不改 xserv**。xserv 是独立项目，T9 只调整 xtrain 的导出去适配 xserv 既有 loader。
+> 架构差异中只有一项是「结构性」的（QK-norm，见下），处理方式见 **Key Design Decision 1**。
+
+## 关键第一步：架构 diff（xtrain TinyTransformer vs xserv qwen3.rs）
+
+逐 op 读 `crates/xtrain-model/src/model.rs` 对 `~/projects/xserv/crates/xserv-model/src/qwen3.rs`
+（forward 路径 `forward_with_cache`，即 `dump-logits` 走的 prefill）：
+
+| 维度 | xtrain TinyTransformer | xserv Qwen3 | 是否兼容 / 处理 |
+|------|------------------------|-------------|----------------|
+| RMSNorm 公式 | `x*rsqrt(mean(x²)+eps)*γ`（无均值减） | 同 | ✅ 一致 |
+| RMSNorm eps | `cfg.eps`（tiny=1e-5） | `rms_norm_eps`（config.json 提供） | ✅ 导出 eps 到 config |
+| **QK-norm** | **T9 前没有** | **强制** per-head RMSNorm(Q,K)（`q_norm`/`k_norm`，shape `[head_dim]`） | **结构性差异** → 见 Decision 1 |
+| RoPE 约定 | rotate_half，`freq=θ^(-2i/hd)`，pos=行号 | rotate_half，cos/sin cache `freq=1/θ^(2i/hd)`，pos=token idx | ✅ **逐式一致**（见 rope.cu）|
+| RoPE θ | `cfg.rope_theta`（10000） | `rope_theta`（config 提供，默认 1e6） | ✅ 导出 θ 到 config |
+| Attention scale | `1/√head_dim` | `1/√head_dim`（attention.rs / flash） | ✅ 一致 |
+| Causal mask | 加性 `-1e9` 上三角 | causal flag（online softmax） | ✅ 同义 |
+| GQA | MHA（无 kv 分组） | 支持 GQA（`num_kv_heads`） | ✅ 设 `num_key_value_heads = num_attention_heads`（退化为 MHA）|
+| SwiGLU | `down(silu(gate)∘up)`，gate/up 独立 proj | 同（融合存 `gate_up_proj`，loader 内部 cat）| ✅ 一致，导出仍分开 gate/up（loader 自己 cat）|
+| 偏置 | 无 | `attention_bias=false`，不读 bias | ✅ 一致 |
+| final norm | `final_norm` 后接 lm_head | `model.norm` 后接 lm_head | ✅ 一致 |
+| embedding tying | 独立 `lm_head` | 独立（`tie_word_embeddings=false`）| ✅ 一致 |
+| 2D 权重 layout | `[in,out]`，算 `x@W` | `[out,in]`，算 `x@Wᵀ`（loader `.transpose(0,1)`）| ⚠️ 导出须转置 |
+| **dtype** | **f32**（训练 + 自身推理） | **BF16 only**（kernel 全 assert BF16）| ⚠️ 导出转 BF16；数值上不可能 bit-exact，见 Decision 2 |
+| vocab | gpt2 BPE，50257 | config 提供 | ✅ 导出 vocab |
+
+**结论**：除 QK-norm 一项外，其余差异都是机械的（命名 / 转置 / dtype / 退化 GQA）。QK-norm 是
+xserv 对 Qwen3 的**强制**步骤（`head_rmsnorm(q,q_norm)` / `head_rmsnorm(k,k_norm)` 无条件执行，
+且 γ=1 也不是恒等——它仍按每个 head 向量的 RMS 做归一），xtrain 训练时从未施加 → 若不处理，
+xserv 的前向数学与 xtrain 不同，闭环不可能成立。
+
+## Key Design Decisions
+
+### Decision 1：给 xtrain 加 per-head QK-norm（而不是伪造 match，也不是停在 blocker）
+
+逃生舱里给了三条路：① 给 xtrain 补 QK-norm 重训以对齐 Qwen3；② 改 xserv（被禁止）；
+③ 当作 blocker 报告。选 ①——因为它**真的能闭环**且改动**外科手术级**：
+
+- xserv 的顺序是 `reshape → head_rmsnorm → transpose_for_rope → rope`，即 **QK-norm 在 RoPE 之前**，
+  作用在每个 `[head_dim]` 的 head 向量上。
+- xtrain 复用既有的 2D `rms_norm` op：把 `[seq,nh,hd]` reshape 成 `[seq*nh, hd]`，用 `[hd]` 的 γ
+  做 rms_norm，再 reshape 回去，**插在 reshape 与 rope 之间**——与 xserv 逐步对齐。autograd 全自动
+  （rms_norm/reshape 都已有 backward），优化器/checkpoint/DDP 都按 `params()` 泛化迭代，自动兼容。
+- 每 block 新增 `q_norm`/`k_norm` 两个 `[head_dim]` leaf；`params()` 顺序在 `wv` 之后、`wo` 之前插入
+  `q_norm,k_norm`；`num_params()` 加 `2*head_dim/layer`；PyTorch 对拍参考（parity.py / adamw_parity.py）
+  同步加 QK-norm，名单同步——改完仍全套自洽。
+
+这样导出的权重是**真·Qwen3 兼容**的（训练时就带 QK-norm），不是凑出来的假象。
+
+### Decision 2：BF16 是 xserv 的硬约束 → 闭环判据用「贪心 token 一致」而非 bit-exact
+
+xserv 的 Qwen3 前向 **只支持 BF16**（embedding/rmsnorm/gemm/silu/rope kernel 全部
+`assert_eq!(dtype, BF16)`，`dump-logits` 也按 bf16 读 logits）。xtrain 是 f32。所以：
+
+- 导出时把所有权重 `bf16::from_f32` 转成 BF16 写入 safetensors。
+- **数值上不可能 bit-exact**：BF16 仅 8 位尾数，权重舍入 + 前向全程 BF16 累加，相对 xtrain 的 f32
+  会有 ~1e-2 量级的 logits 漂移。因此**闭环判据定为「同 prompt 下贪心 argmax token 序列一致」**
+  （+ 报告 logits 的 top-1/分布吻合度），这是 BF16 推理对 f32 训练能给出的最强、且诚实的判据。
+
+### Decision 3：导出复用 train.rs 的 config + checkpoint，零猜测
+
+导出器 `bin/export_safetensors.rs` 用与 `bin/train.rs` **完全相同**的 `Config`（gpt2 vocab、
+`n_layers=4`、`Config::tiny()` 其余字段）建空模型 → `checkpoint::load_into` 灌入训练权重 →
+按 `params()` 稳定序映射。tokenizer.json 直接 copy 进导出目录，两侧用同一份 BPE。
+
+## 张量名 + layout 映射表
+
+xtrain `params()` 序（T9 后）：
+`embed[vocab,dim]` → 每 block `[attn_norm, wq, wk, wv, q_norm, k_norm, wo, ffn_norm, w_gate, w_up, w_down]`
+→ `final_norm[dim]` → `lm_head[dim,vocab]`。
+
+| xtrain 参数 | shape | → HF Qwen3 名 | HF shape | 操作 |
+|-------------|-------|---------------|----------|------|
+| `embed` | `[vocab,dim]` | `model.embed_tokens.weight` | `[vocab,dim]` | keep（行索引两侧同）|
+| `attn_norm` | `[dim]` | `model.layers.{i}.input_layernorm.weight` | `[dim]` | keep |
+| `wq` | `[dim,dim]` | `…self_attn.q_proj.weight` | `[dim,dim]` | **transpose** |
+| `wk` | `[dim,dim]` | `…self_attn.k_proj.weight` | `[dim,dim]` | **transpose** |
+| `wv` | `[dim,dim]` | `…self_attn.v_proj.weight` | `[dim,dim]` | **transpose** |
+| `q_norm` | `[head_dim]` | `…self_attn.q_norm.weight` | `[head_dim]` | keep |
+| `k_norm` | `[head_dim]` | `…self_attn.k_norm.weight` | `[head_dim]` | keep |
+| `wo` | `[dim,dim]` | `…self_attn.o_proj.weight` | `[dim,dim]` | **transpose** |
+| `ffn_norm` | `[dim]` | `…post_attention_layernorm.weight` | `[dim]` | keep |
+| `w_gate` | `[dim,ffn]` | `…mlp.gate_proj.weight` | `[ffn,dim]` | **transpose** |
+| `w_up` | `[dim,ffn]` | `…mlp.up_proj.weight` | `[ffn,dim]` | **transpose** |
+| `w_down` | `[ffn,dim]` | `…mlp.down_proj.weight` | `[dim,ffn]` | **transpose** |
+| `final_norm` | `[dim]` | `model.norm.weight` | `[dim]` | keep |
+| `lm_head` | `[dim,vocab]` | `lm_head.weight` | `[vocab,dim]` | **transpose** |
+
+全部以 **BF16** dtype 写入。config.json 字段：`architectures=["Qwen3ForCausalLM"]`、`model_type="qwen3"`、
+`vocab_size`、`hidden_size=dim`、`intermediate_size=ffn`、`num_hidden_layers`、`num_attention_heads`、
+`num_key_value_heads=num_attention_heads`、`head_dim`、`rms_norm_eps=eps`、`rope_theta`、
+`tie_word_embeddings=false`、`attention_bias=false`、`hidden_act="silu"`。
+
+## Module Layout
+
+```
+crates/xtrain-model/src/model.rs            # +q_norm/k_norm leaf；attention 插 per-head QK-norm；params() 序更新
+crates/xtrain-model/src/config.rs           # num_params() 计入 QK-norm γ
+crates/xtrain-train/src/bin/export_safetensors.rs   # 导出器（本 Phase 核心实现）
+crates/xtrain-train/Cargo.toml              # +half, +safetensors="0.5"（对齐 xserv）, +bin
+crates/xtrain-model/tests/parity{.py,_dump.rs}      # PyTorch 对拍同步加 QK-norm
+crates/xtrain-train/tests/adamw_parity{.py,_dump.rs}# 同上
+```
+
+## 验证方法
+
+1. **本地**：`cargo check --workspace` + `cargo fmt --all -- --check` 过（导出器 GPU 体 gated 在
+   `not(no_cuda)`，host 侧只 check）。
+2. **dash5（闭环）**：
+   ```bash
+   export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
+   # ① 训练一个小模型 → checkpoint（或复用已训的）
+   cargo run -p xtrain-train --release --bin train -- \
+       /opt/wjh/models/gpt2/tokenizer.json data/tinystories-valid-3mb.txt \
+       <steps> /tmp/xtrain_tinystories.ckpt
+   # ② 导出
+   cargo run -p xtrain-train --release --bin export_safetensors -- \
+       /tmp/xtrain_tinystories.ckpt /opt/wjh/models/gpt2/tokenizer.json /tmp/xtrain_export
+   # ③ xserv 加载 + dump logits（同 prompt）
+   #   (在 /opt/wjh/projects/xserv) cargo run -p xserv-model --release --bin dump-logits -- /tmp/xtrain_export "<prompt>"
+   # ④ 对拍：xserv 贪心 token 序列 vs xtrain 自身贪心（sample.rs generate, temp=0）
+   ```
+   判据：**贪心 token 序列一致**（BF16 推理 vs f32 训练，logits top-1 吻合；分布在 BF16 容差内）。
+
+## 验证结果
+
+> 待 dash5 实跑回填（push origin main → dash5 pull → 跑 ②③④ → capture）。