diff --git a/crates/spark-model/examples/rmsnorm_vanilla_microtest.rs b/crates/spark-model/examples/rmsnorm_vanilla_microtest.rs new file mode 100644 index 000000000..7dff08083 --- /dev/null +++ b/crates/spark-model/examples/rmsnorm_vanilla_microtest.rs @@ -0,0 +1,385 @@ +// SPDX-License-Identifier: AGPL-3.0-only +//! LEG-3 device verification for the V4 RMSNorm substitution (pre-reg +//! `DSPARK-V4-PHASE2-RMSNORM-PREREG.md`). +//! +//! Runs the PRODUCTION kernels against an F64 host formula: +//! V2 `rms_norm_vanilla` == x * w / sqrt(mean(x^2) + eps) (the new path) +//! I1 `rms_norm` + zero weight == pure normalize (1 + 0 = 1) (the q_b_norm invariant) +//! V3 `rms_norm` + offset weight == x * (1 + w) / rms (the non-V4 regression) +//! V1d the OLD V4 path (`1 + bf16(w-1)`) vs the NEW one (`w`), on real V4 weight vectors: +//! both are run on-device and compared to F64 truth computed from the EXACT weight. +//! +//! Weight regimes covered: near 0, near 1, negative / sign-straddling, large, and +//! the real V4 `q_norm` / `attn_norm` / `compressor.norm` distributions. +//! +//! cargo run -p spark-model --release --example rmsnorm_vanilla_microtest \ +//! --features cuda,gpu-examples +use anyhow::{Result, bail}; +use half::bf16; +use spark_runtime::cuda_backend::AtlasCudaBackend; +use spark_runtime::gpu::{DevicePtr, GpuBackend}; +use spark_runtime::kernel_args::KernelLaunch; + +const EPS: f32 = 1e-6; + +struct Lcg(u64); +impl Lcg { + fn f(&mut self) -> f64 { + self.0 = self + .0 + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((self.0 >> 11) as f64) / ((1u64 << 53) as f64) + } + fn r(&mut self, a: f64, b: f64) -> f64 { + a + (b - a) * self.f() + } +} + +fn ub(g: &dyn GpuBackend, d: &[bf16]) -> Result { + let b: Vec = d.iter().flat_map(|x| x.to_bits().to_le_bytes()).collect(); + let p = g.alloc(b.len())?; + g.copy_h2d(&b, p)?; + Ok(p) +} +fn db(g: &dyn GpuBackend, p: DevicePtr, n: usize) -> Result> { + let mut b = vec![0u8; n * 2]; + g.copy_d2h(p, &mut b)?; + Ok(b.chunks_exact(2) + .map(|c| bf16::from_bits(u16::from_le_bytes([c[0], c[1]])).to_f64()) + .collect()) +} + +/// F64 ground truth. `w_eff` is the weight the formula should apply. +/// Inputs are the exact BF16 values the device sees, so any delta is the kernel's. +fn truth(x: &[bf16], w_eff: &[f64], hidden: usize, tokens: usize) -> Vec { + let mut out = vec![0.0f64; tokens * hidden]; + for t in 0..tokens { + let xs = &x[t * hidden..(t + 1) * hidden]; + let ms: f64 = xs.iter().map(|v| v.to_f64() * v.to_f64()).sum::() / hidden as f64; + let rms = 1.0 / (ms + EPS as f64).sqrt(); + for i in 0..hidden { + // Round through BF16 exactly like the kernel's store does. + out[t * hidden + i] = bf16::from_f32((xs[i].to_f64() * rms * w_eff[i]) as f32).to_f64(); + } + } + out +} + +fn launch( + g: &dyn GpuBackend, + module: &str, + func: &str, + x: DevicePtr, + w: DevicePtr, + o: DevicePtr, + tokens: u32, + hidden: u32, +) -> Result<()> { + let k = g.kernel(module, func)?; + KernelLaunch::new(g, k) + .grid([tokens, 1, 1]) + .block([hidden.min(1024), 1, 1]) + .arg_ptr(x) + .arg_ptr(w) + .arg_ptr(o) + .arg_u32(hidden) + .arg_f32(EPS) + .launch(0)?; + g.synchronize(0)?; + Ok(()) +} + +/// max |device - truth|, and the max relative error against the truth magnitude. +fn cmp(dev: &[f64], truth_v: &[f64]) -> (f64, f64) { + let mut amax = 0.0f64; + let mut rmax = 0.0f64; + for (d, t) in dev.iter().zip(truth_v) { + let a = (d - t).abs(); + amax = amax.max(a); + if t.abs() > 1e-9 { + rmax = rmax.max(a / t.abs()); + } + } + (amax, rmax) +} + +/// bf16(w - 1) then +1 — exactly what the OLD loader stored and the offset kernel recovered. +fn old_effective(w: f64) -> f64 { + 1.0 + bf16::from_f32((w - 1.0) as f32).to_f64() +} + +struct Case { + name: &'static str, + lo: f64, + hi: f64, +} + +fn main() -> Result<()> { + let gpu = AtlasCudaBackend::new(0, &atlas_kernels::ptx_modules())?; + let g: &dyn GpuBackend = &gpu; + + let dims = [512usize, 1024, 4096, 7168]; + let seeds = [1u64, 7, 42, 1337, 90210]; + // Weight regimes. `v4_qnorm`/`v4_attn` mirror the measured V4 distributions; + // `straddle_zero` mirrors compressor.norm (which sign-flips under the old path). + let cases = [ + Case { + name: "v4_qnorm (~0.04)", + lo: 0.0098, + hi: 0.2402, + }, + Case { + name: "v4_attnnorm (~0.03)", + lo: 0.0084, + hi: 0.0400, + }, + Case { + name: "straddle_zero ", + lo: -0.2256, + hi: 0.2256, + }, + Case { + name: "near_one ", + lo: 0.6300, + hi: 0.9900, + }, + Case { + name: "large ", + lo: 1.0000, + hi: 3.8281, + }, + Case { + name: "negative ", + lo: -1.5000, + hi: -0.0100, + }, + ]; + + let mut fail = 0usize; + println!("=== V2 — production `rms_norm_vanilla` vs F64 truth (out = x*w/rms) ==="); + println!( + "{:<22} {:>6} {:>6} {:>12} {:>12}", + "weight regime", "hidden", "seed", "max |abs|", "max rel" + ); + println!("{}", "-".repeat(64)); + for c in &cases { + for &hidden in &dims { + for &seed in &seeds { + let tokens = 3usize; + let mut r = Lcg(seed); + let x: Vec = (0..tokens * hidden) + .map(|_| bf16::from_f32(r.r(-3.0, 3.0) as f32)) + .collect(); + let w: Vec = (0..hidden) + .map(|_| bf16::from_f32(r.r(c.lo, c.hi) as f32)) + .collect(); + let w64: Vec = w.iter().map(|v| v.to_f64()).collect(); + + let xp = ub(g, &x)?; + let wp = ub(g, &w)?; + let op = g.alloc(tokens * hidden * 2)?; + launch( + g, + "rms_norm_vanilla", + "rms_norm_vanilla", + xp, + wp, + op, + tokens as u32, + hidden as u32, + )?; + let dev = db(g, op, tokens * hidden)?; + let truth_v = truth(&x, &w64, hidden, tokens); + let (a, rel) = cmp(&dev, &truth_v); + // BF16 store granularity: one ULP at magnitude ~4 is ~0.03. The kernel + // reduces in F32, so allow one BF16 ULP of slack, no more. + let bad = rel > 0.01; + if bad { + fail += 1; + } + if seed == seeds[0] && hidden == dims[0] { + println!( + "{:<22} {:>6} {:>6} {:>12.3e} {:>12.3e}{}", + c.name, + hidden, + seed, + a, + rel, + if bad { " <<< FAIL" } else { "" } + ); + } + if bad { + println!( + " FAIL {:<20} hidden={} seed={} abs={:.3e} rel={:.3e}", + c.name, hidden, seed, a, rel + ); + } + } + } + } + println!( + "V2: {} regimes x {} dims x {} seeds = {} cells, {} failures\n", + cases.len(), + dims.len(), + seeds.len(), + cases.len() * dims.len() * seeds.len(), + fail + ); + + // ── I1 — the q_b_norm invariant: offset kernel + ZERO weight == pure normalize ── + println!("=== I1 — `rms_norm` + zero weight (norm_unit_w) == pure normalize (unchanged) ==="); + let hidden = 512usize; + let tokens = 4usize; + let mut r = Lcg(2026); + let x: Vec = (0..tokens * hidden) + .map(|_| bf16::from_f32(r.r(-3.0, 3.0) as f32)) + .collect(); + let wz: Vec = vec![bf16::from_f32(0.0); hidden]; + let xp = ub(g, &x)?; + let wp = ub(g, &wz)?; + let op = g.alloc(tokens * hidden * 2)?; + launch( + g, + "norm", + "rms_norm", + xp, + wp, + op, + tokens as u32, + hidden as u32, + )?; + let dev = db(g, op, tokens * hidden)?; + let ones = vec![1.0f64; hidden]; + let truth_v = truth(&x, &ones, hidden, tokens); + let (a, rel) = cmp(&dev, &truth_v); + let i1_ok = a == 0.0; + println!( + " offset kernel, w=0 -> max|abs| {:.3e} max rel {:.3e} {}", + a, + rel, + if i1_ok { + "BIT-EXACT — invariant holds" + } else { + "<<< FAIL" + } + ); + if !i1_ok { + fail += 1; + } + + // ── V3 — non-V4 regression: offset kernel still computes x*(1+w)/rms ── + println!("\n=== V3 — regression: `rms_norm` (offset-from-1) is UNCHANGED for non-V4 ==="); + let mut r = Lcg(5150); + let wo: Vec = (0..hidden) + .map(|_| bf16::from_f32(r.r(-0.5, 0.5) as f32)) + .collect(); + // Offset-from-1 model semantics: effective weight = 1 + stored. + let w_eff: Vec = wo.iter().map(|v| 1.0 + v.to_f64()).collect(); + let wp = ub(g, &wo)?; + let op = g.alloc(tokens * hidden * 2)?; + launch( + g, + "norm", + "rms_norm", + xp, + wp, + op, + tokens as u32, + hidden as u32, + )?; + let dev = db(g, op, tokens * hidden)?; + let truth_v = truth(&x, &w_eff, hidden, tokens); + let (a, rel) = cmp(&dev, &truth_v); + let v3_ok = rel < 0.01; + println!( + " offset kernel, w~U(-0.5,0.5) -> max|abs| {:.3e} max rel {:.3e} {}", + a, + rel, + if v3_ok { + "UNCHANGED — offset path intact" + } else { + "<<< FAIL" + } + ); + if !v3_ok { + fail += 1; + } + + // ── V1d — OLD vs NEW on device, both scored against F64 truth from the EXACT weight ── + println!( + "\n=== V1d — OLD path (1 + bf16(w-1)) vs NEW path (exact w), device, vs F64 truth ===" + ); + println!( + "{:<22} {:>13} {:>13} {:>13} {:>13}", + "weight regime", "OLD max|abs|", "OLD max rel", "NEW max|abs|", "NEW max rel" + ); + println!("{}", "-".repeat(80)); + for c in &cases { + let mut r = Lcg(31337); + let x: Vec = (0..tokens * hidden) + .map(|_| bf16::from_f32(r.r(-3.0, 3.0) as f32)) + .collect(); + let w: Vec = (0..hidden) + .map(|_| bf16::from_f32(r.r(c.lo, c.hi) as f32)) + .collect(); + let w64: Vec = w.iter().map(|v| v.to_f64()).collect(); + let truth_v = truth(&x, &w64, hidden, tokens); // truth uses the EXACT checkpoint weight + let xp = ub(g, &x)?; + + // OLD: loader stored bf16(w-1); offset kernel adds 1 back. + let w_old: Vec = w + .iter() + .map(|v| bf16::from_f32((v.to_f64() - 1.0) as f32)) + .collect(); + let wp_old = ub(g, &w_old)?; + let op_old = g.alloc(tokens * hidden * 2)?; + launch( + g, + "norm", + "rms_norm", + xp, + wp_old, + op_old, + tokens as u32, + hidden as u32, + )?; + let (a_old, r_old) = cmp(&db(g, op_old, tokens * hidden)?, &truth_v); + + // NEW: exact weight, vanilla kernel. + let wp_new = ub(g, &w)?; + let op_new = g.alloc(tokens * hidden * 2)?; + launch( + g, + "rms_norm_vanilla", + "rms_norm_vanilla", + xp, + wp_new, + op_new, + tokens as u32, + hidden as u32, + )?; + let (a_new, r_new) = cmp(&db(g, op_new, tokens * hidden)?, &truth_v); + + // Sanity: the host-side prediction of the old effective weight. + let pred: Vec = w64.iter().map(|v| old_effective(*v)).collect(); + let pred_truth = truth(&x, &pred, hidden, tokens); + let (a_pred, _) = cmp(&db(g, op_old, tokens * hidden)?, &pred_truth); + + println!( + "{:<22} {:>13.3e} {:>13.3e} {:>13.3e} {:>13.3e} [old==1+bf16(w-1) model: {:.1e}]", + c.name, a_old, r_old, a_new, r_new, a_pred + ); + if r_new > r_old && r_old > 1e-6 { + println!(" <<< FAIL: the new path is not better than the old one"); + fail += 1; + } + } + + println!(); + if fail > 0 { + bail!("{fail} failure(s) — READING R1: STOP AND REPAIR. No behavioral eval."); + } + println!("ALL DEVICE CHECKS PASS (V2 + I1 + V3 + V1d)."); + Ok(()) +} diff --git a/crates/spark-model/src/layers/deepseek_v4_mtp.rs b/crates/spark-model/src/layers/deepseek_v4_mtp.rs index ef73b02de..51a5f1903 100644 --- a/crates/spark-model/src/layers/deepseek_v4_mtp.rs +++ b/crates/spark-model/src/layers/deepseek_v4_mtp.rs @@ -150,7 +150,9 @@ impl DeepseekV4MtpHead { lm_head, mtp_vocab_size, kv_cache: Mutex::new(kv_cache), - rms_norm_k: gpu.kernel("norm", "rms_norm")?, + // V4 ships HF-vanilla norm weights (enorm/hnorm/norm are loaded + // exactly) — the offset-from-1 kernel would apply `1 + w`. + rms_norm_k: gpu.kernel("rms_norm_vanilla", "rms_norm_vanilla")?, dense_gemv_k: gpu.kernel("gemv", "dense_gemv_bf16")?, residual_add_k: gpu.kernel("residual_add", "bf16_residual_add")?, hc_expand_k: gpu.kernel("hyper_connection", "hc_expand")?, diff --git a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward.rs b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward.rs index de63afa8e..a86d92501 100644 --- a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward.rs +++ b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward.rs @@ -240,7 +240,7 @@ impl Qwen3AttentionLayer { if let Some(ref q_norm_full) = self.attn.q_norm_full { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_out, q_norm_full, q_out, @@ -252,7 +252,7 @@ impl Qwen3AttentionLayer { } else if !self.attn.q_norm.weight.is_null() { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_out, &self.attn.q_norm, q_out, @@ -265,7 +265,7 @@ impl Qwen3AttentionLayer { if let Some(ref k_norm_full) = self.attn.k_norm_full { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_out, k_norm_full, k_out, @@ -277,7 +277,7 @@ impl Qwen3AttentionLayer { } else if !self.attn.k_norm.weight.is_null() { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_out, &self.attn.k_norm, k_out, @@ -300,7 +300,7 @@ impl Qwen3AttentionLayer { if let Some(v_norm_w) = self.v_norm_weight.as_ref() { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, v_out, v_norm_w, v_out, diff --git a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_mla.rs b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_mla.rs index e4a4d1730..87b501efc 100644 --- a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_mla.rs +++ b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_mla.rs @@ -116,7 +116,7 @@ impl Qwen3AttentionLayer { prof!("q_norm", { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_latent, &mla.q_a_norm, q_latent, @@ -262,7 +262,7 @@ impl Qwen3AttentionLayer { } ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, kv_latent, &mla.kv_a_norm, kv_latent, diff --git a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_v4.rs b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_v4.rs index 9c0b7356a..4b6451fb6 100644 --- a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_v4.rs +++ b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_v4.rs @@ -177,7 +177,7 @@ impl Qwen3AttentionLayer { let block = compressed.offset(tgt as usize * hd_mla as usize * 2); ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, block, &comp.norm, block, @@ -308,7 +308,7 @@ impl Qwen3AttentionLayer { prof!("q_norm", { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_latent, &mla.q_a_norm, q_latent, @@ -426,7 +426,7 @@ impl Qwen3AttentionLayer { // → attention score overflow → NaN. nkv heads × (kv_dim/nkv) each. ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_out, &mla.kv_a_norm, k_out, diff --git a/crates/spark-model/src/layers/qwen3_attention/init.rs b/crates/spark-model/src/layers/qwen3_attention/init.rs index 7bd6738cc..0d9480083 100644 --- a/crates/spark-model/src/layers/qwen3_attention/init.rs +++ b/crates/spark-model/src/layers/qwen3_attention/init.rs @@ -173,6 +173,12 @@ impl Qwen3AttentionLayer { "fp8_gemm_t_blockscaled", ), rms_norm_k: gpu.kernel("norm", "rms_norm")?, + rms_norm_w_k: if crate::ships_vanilla_norm_weights(config) { + gpu.kernel("rms_norm_vanilla", "rms_norm_vanilla")? + } else { + gpu.kernel("norm", "rms_norm")? + }, + norm_vanilla: crate::ships_vanilla_norm_weights(config), rms_norm_residual_k: gpu.kernel("norm", "rms_norm_residual")?, dense_gemv_k: gpu.kernel("gemv", "dense_gemv_bf16")?, w4a16_gemv_k: gpu.kernel("w4a16_gemv", "w4a16_gemv")?, diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip.rs index 69229036a..9b7d390cd 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip.rs @@ -250,7 +250,7 @@ impl Qwen3AttentionLayer { if let Some(ref q_norm_full) = self.attn.q_norm_full { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_contiguous, q_norm_full, q_contiguous, @@ -262,7 +262,7 @@ impl Qwen3AttentionLayer { } else if !self.attn.q_norm.weight.is_null() { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_contiguous, &self.attn.q_norm, q_contiguous, @@ -276,7 +276,7 @@ impl Qwen3AttentionLayer { if let Some(ref k_norm_full) = self.attn.k_norm_full { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_contiguous, k_norm_full, k_contiguous, @@ -288,7 +288,7 @@ impl Qwen3AttentionLayer { } else if !self.attn.k_norm.weight.is_null() { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_contiguous, &self.attn.k_norm, k_contiguous, @@ -326,7 +326,7 @@ impl Qwen3AttentionLayer { if let Some(v_norm_w) = self.v_norm_weight.as_ref() { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, v_contiguous, v_norm_w, v_contiguous, diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_mla.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_mla.rs index eec760296..56b97d7dd 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_mla.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_mla.rs @@ -91,7 +91,7 @@ impl Qwen3AttentionLayer { } ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_latent, &mla.q_a_norm, q_latent, @@ -156,7 +156,7 @@ impl Qwen3AttentionLayer { } ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, kv_latent, &mla.kv_a_norm, kv_latent, diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_v4.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_v4.rs index d1dcca756..68ba3945d 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_v4.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_v4.rs @@ -108,7 +108,7 @@ impl Qwen3AttentionLayer { .map_err(|e| anyhow::anyhow!("V4 attn: q_latent gemm sync failed: {e}"))?; ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_latent, &mla.q_a_norm, q_latent, @@ -243,7 +243,7 @@ impl Qwen3AttentionLayer { // kv_latent so the cached latent, k_out and v_out are all normalized. ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, kv_latent, &mla.kv_a_norm, kv_latent, @@ -489,7 +489,7 @@ impl Qwen3AttentionLayer { .launch(stream)?; ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, compressed, &comp.norm, compressed, diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/paged.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/paged.rs index 40965a768..bd3d6fcde 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/paged.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/paged.rs @@ -204,7 +204,7 @@ impl Qwen3AttentionLayer { // never reach this branch — they early-return above. ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_contiguous, q_norm_full, q_contiguous, @@ -216,7 +216,7 @@ impl Qwen3AttentionLayer { } else if !self.attn.q_norm.weight.is_null() { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_contiguous, &self.attn.q_norm, q_contiguous, @@ -230,7 +230,7 @@ impl Qwen3AttentionLayer { if let Some(ref k_norm_full) = self.attn.k_norm_full { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_contiguous, k_norm_full, k_contiguous, @@ -242,7 +242,7 @@ impl Qwen3AttentionLayer { } else if !self.attn.k_norm.weight.is_null() { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_contiguous, &self.attn.k_norm, k_contiguous, @@ -264,7 +264,7 @@ impl Qwen3AttentionLayer { if let Some(v_norm_w) = self.v_norm_weight.as_ref() { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, v_contiguous, v_norm_w, v_contiguous, diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_mla.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_mla.rs index c7999e66f..7832b7318 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_mla.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_mla.rs @@ -79,7 +79,7 @@ impl Qwen3AttentionLayer { )?; ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_latent, &mla.q_a_norm, q_latent, @@ -116,7 +116,7 @@ impl Qwen3AttentionLayer { )?; ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, kv_latent, &mla.kv_a_norm, kv_latent, diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_v4.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_v4.rs index ac88366c7..6784fab9a 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_v4.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_v4.rs @@ -69,7 +69,7 @@ impl Qwen3AttentionLayer { )?; ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_latent, &mla.q_a_norm, q_latent, @@ -127,7 +127,7 @@ impl Qwen3AttentionLayer { // (DeepSeek-V4: kv = kv_norm(kv_proj(h))). n*nkv rows of kv_lora dims. ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_out, &mla.kv_a_norm, k_out, diff --git a/crates/spark-model/src/layers/qwen3_attention/trait_impl/decode_inner.rs b/crates/spark-model/src/layers/qwen3_attention/trait_impl/decode_inner.rs index 7719aa382..d2d3ee652 100644 --- a/crates/spark-model/src/layers/qwen3_attention/trait_impl/decode_inner.rs +++ b/crates/spark-model/src/layers/qwen3_attention/trait_impl/decode_inner.rs @@ -125,7 +125,7 @@ impl Qwen3AttentionLayer { if let Some(ref post_norm) = self.post_attn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, attn_out, post_norm, attn_out, @@ -184,7 +184,7 @@ impl Qwen3AttentionLayer { if let Some(ref post_norm) = self.post_ffn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, moe_out, post_norm, moe_out, @@ -265,7 +265,7 @@ impl Qwen3AttentionLayer { // post-MoE norm (in-place on moe_output) ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, moe_out, post_norm, moe_out, @@ -284,7 +284,7 @@ impl Qwen3AttentionLayer { // post-dense norm (layernorm_1) ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, dense_out, dense_norm, dense_out, @@ -308,7 +308,7 @@ impl Qwen3AttentionLayer { if let Some(ref combined_norm) = self.post_ffn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, dense_out, combined_norm, dense_out, @@ -352,7 +352,7 @@ impl Qwen3AttentionLayer { if let Some(ref post_norm) = self.post_ffn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, dense_out, post_norm, dense_out, @@ -500,7 +500,7 @@ impl Qwen3AttentionLayer { let normed = ctx.buffers.norm_output(); ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, hidden, &self.input_norm, normed, @@ -531,7 +531,7 @@ impl Qwen3AttentionLayer { if let Some(ref post_norm) = self.post_attn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, attn_out, post_norm, attn_out, @@ -657,7 +657,7 @@ impl Qwen3AttentionLayer { let normed2 = ctx.buffers.norm_output(); ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, hidden, &self.post_attn_norm, normed2, @@ -672,7 +672,7 @@ impl Qwen3AttentionLayer { if let Some(ref post_norm) = self.post_ffn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, ffn_out, post_norm, ffn_out, diff --git a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mla.rs b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mla.rs index 206a74b57..2f4e73eae 100644 --- a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mla.rs +++ b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mla.rs @@ -264,7 +264,7 @@ impl Qwen3AttentionLayer { } ops::rms_norm( gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_latent, &mla.q_a_norm, q_latent, @@ -365,7 +365,7 @@ impl Qwen3AttentionLayer { } ops::rms_norm( gpu, - self.rms_norm_k, + self.rms_norm_w_k, kv_latent, &mla.kv_a_norm, kv_latent, diff --git a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mod.rs b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mod.rs index 44996608a..463f5ba55 100644 --- a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mod.rs +++ b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mod.rs @@ -191,7 +191,7 @@ impl Qwen3AttentionLayer { } ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, c.hidden, &self.input_norm, c.normed, @@ -337,7 +337,7 @@ impl Qwen3AttentionLayer { } ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, c.hidden, &self.post_attn_norm, c.normed, diff --git a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/qkv.rs b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/qkv.rs index e6c525719..97bfff2f4 100644 --- a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/qkv.rs +++ b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/qkv.rs @@ -70,7 +70,7 @@ impl Qwen3AttentionLayer { if !self.attn.q_norm.weight.is_null() { ops::rms_norm( fwd.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_out_i, &self.attn.q_norm, q_out_i, @@ -83,7 +83,7 @@ impl Qwen3AttentionLayer { if !self.attn.k_norm.weight.is_null() { ops::rms_norm( fwd.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_out_i, &self.attn.k_norm, k_out_i, @@ -186,7 +186,7 @@ impl Qwen3AttentionLayer { if !self.attn.q_norm.weight.is_null() { ops::rms_norm( fwd.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_out_i, &self.attn.q_norm, q_out_i, @@ -199,7 +199,7 @@ impl Qwen3AttentionLayer { if !self.attn.k_norm.weight.is_null() { ops::rms_norm( fwd.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_out_i, &self.attn.k_norm, k_out_i, @@ -301,7 +301,7 @@ impl Qwen3AttentionLayer { if !self.attn.q_norm.weight.is_null() { ops::rms_norm( fwd.gpu, - self.rms_norm_k, + self.rms_norm_w_k, q_out_i, &self.attn.q_norm, q_out_i, @@ -314,7 +314,7 @@ impl Qwen3AttentionLayer { if !self.attn.k_norm.weight.is_null() { ops::rms_norm( fwd.gpu, - self.rms_norm_k, + self.rms_norm_w_k, k_out_i, &self.attn.k_norm, k_out_i, diff --git a/crates/spark-model/src/layers/qwen3_attention/trait_impl/prefill_inner.rs b/crates/spark-model/src/layers/qwen3_attention/trait_impl/prefill_inner.rs index 48c5dc095..16903d405 100644 --- a/crates/spark-model/src/layers/qwen3_attention/trait_impl/prefill_inner.rs +++ b/crates/spark-model/src/layers/qwen3_attention/trait_impl/prefill_inner.rs @@ -257,7 +257,7 @@ impl Qwen3AttentionLayer { if let Some(ref post_norm) = self.post_attn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, attn_out, post_norm, attn_out, @@ -350,7 +350,7 @@ impl Qwen3AttentionLayer { // 1. Norm dense MLP output with post_feedforward_layernorm_1 ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, dense_out, dense_norm, dense_out, @@ -372,7 +372,7 @@ impl Qwen3AttentionLayer { let moe_out = ctx.buffers.moe_output(); ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, moe_out, post_norm, moe_out, @@ -396,7 +396,7 @@ impl Qwen3AttentionLayer { if let Some(ref combined_norm) = self.post_ffn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, moe_out, combined_norm, moe_out, @@ -421,7 +421,7 @@ impl Qwen3AttentionLayer { if let Some(ref post_norm) = self.post_ffn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, dense_out, post_norm, dense_out, @@ -553,7 +553,7 @@ impl Qwen3AttentionLayer { let normed = ctx.buffers.norm_output(); ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, hidden, &self.input_norm, normed, @@ -631,7 +631,7 @@ impl Qwen3AttentionLayer { if let Some(ref post_norm) = self.post_attn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, attn_out, post_norm, attn_out, @@ -755,7 +755,7 @@ impl Qwen3AttentionLayer { let normed2 = ctx.buffers.norm_output(); ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, hidden, &self.post_attn_norm, normed2, @@ -774,7 +774,7 @@ impl Qwen3AttentionLayer { if let Some(ref post_norm) = self.post_ffn_out_norm { ops::rms_norm( ctx.gpu, - self.rms_norm_k, + self.rms_norm_w_k, dense_out, post_norm, dense_out, diff --git a/crates/spark-model/src/layers/qwen3_attention/types.rs b/crates/spark-model/src/layers/qwen3_attention/types.rs index 34ae5cd5e..21c35b7db 100644 --- a/crates/spark-model/src/layers/qwen3_attention/types.rs +++ b/crates/spark-model/src/layers/qwen3_attention/types.rs @@ -98,7 +98,7 @@ pub struct CompressorWeights { pub wkv: DenseWeight, /// gate_proj: same shape as wkv. pub wgate: DenseWeight, - /// kv_norm weight `[head_dim]` — STANDARD RMSNorm (loaded via dense_minus_one). + /// kv_norm weight `[head_dim]` — HF-vanilla RMSNorm (loaded exactly). pub norm: DenseWeight, /// position_bias / ape: [ratio, proj_dim] BF16, added to the gate before softmax. pub ape: spark_runtime::gpu::DevicePtr, @@ -281,7 +281,16 @@ pub struct Qwen3AttentionLayer { pub(super) per_token_group_quant_fp8_k: KernelHandle, pub(super) fp8_gemm_t_blockscaled_k: KernelHandle, // Kernels — decode (GEMV M=1) + /// Offset-from-1 `rms_norm` (`out = x * (1 + w) / rms`). Used ONLY for the + /// unweighted normalize (`norm_unit_w()` is zero-filled, so `1 + 0 = 1`). pub(super) rms_norm_k: KernelHandle, + /// The norm kernel for every weight that comes from the CHECKPOINT. + /// Same handle as `rms_norm_k` for offset-from-1 models; `rms_norm_vanilla` + /// (`out = x * w / rms`) for models that ship HF-vanilla norm weights. + pub(super) rms_norm_w_k: KernelHandle, + /// True when `rms_norm_w_k` is the vanilla kernel — i.e. the checkpoint's + /// norm weights are loaded exactly, with no `-1` pre-subtraction. + pub(super) norm_vanilla: bool, pub(super) rms_norm_residual_k: KernelHandle, /// Gemma-4 FP32-input rms_norm (absolute formula). pub(super) rms_norm_f32_in_k: KernelHandle, diff --git a/crates/spark-model/src/lib.rs b/crates/spark-model/src/lib.rs index f32aed3d6..3bce5be66 100644 --- a/crates/spark-model/src/lib.rs +++ b/crates/spark-model/src/lib.rs @@ -33,3 +33,124 @@ pub mod traits; pub mod vision_preprocess; pub mod weight_loader; pub mod weight_map; + +/// True when the checkpoint ships **HF-vanilla** RMSNorm weights — i.e. the norm +/// weight is used as `out = x * w / rms`, not Qwen3-Next's offset-from-1 +/// `out = x * (1 + w) / rms`. +/// +/// Such a model must load its norm weights **exactly** and dispatch +/// `rms_norm_vanilla`. The alternative — pre-subtracting 1.0 and storing +/// `bf16(w - 1)` for the offset kernel — is only lossless when `w ≈ 1`. +/// DeepSeek-V4's norm weights are ≈ 0.03, so `w - 1 ≈ -0.97`, and BF16's +/// rounding error there (~1.9e-3 absolute) becomes a **1.8-3.4 % relative error +/// on the weight itself** once 1 is added back — catastrophic cancellation. +/// Measured over all 249 V4 norm tensors: up to 19 % on `q_norm`, and 100 % +/// with sign flips on the compressor norms. +/// +/// This is an explicit model dispatch, NOT an inference from weight statistics. +pub fn ships_vanilla_norm_weights(config: &atlas_core::config::ModelConfig) -> bool { + model_type_ships_vanilla_norm_weights(&config.model_type) +} + +/// The dispatch predicate itself, on the bare `model_type`, so it is unit-testable +/// without constructing a full `ModelConfig`. +pub fn model_type_ships_vanilla_norm_weights(model_type: &str) -> bool { + model_type == "deepseek_v4" +} + +#[cfg(test)] +mod norm_convention_tests { + use super::model_type_ships_vanilla_norm_weights as vanilla; + use half::bf16; + + /// Only DeepSeek-V4 takes the vanilla path. Every other family keeps the + /// offset-from-1 convention it was loaded and validated under. + #[test] + fn only_deepseek_v4_uses_vanilla_norm_weights() { + assert!(vanilla("deepseek_v4")); + for other in [ + "qwen3_next", + "qwen3_5_moe", + "qwen3_moe", + "deepseek_v3", + "llama", + "mistral", + "nemotron", + "", + ] { + assert!(!vanilla(other), "{other} must keep offset-from-1 semantics"); + } + } + + /// The norm weights ship as BF16 in the checkpoint, so the value the loader + /// actually sees is `bf16(x)`. Model that exactly. + fn ckpt(x: f32) -> f32 { + bf16::from_f32(x).to_f32() + } + /// OLD loader: store `bf16(w - 1)`, so the offset-from-1 kernel applies + /// `1 + bf16(w - 1)`. + fn old_effective(w: f32) -> f32 { + 1.0 + bf16::from_f32(w - 1.0).to_f32() + } + fn rel_err(got: f32, want: f32) -> f32 { + ((got - want) / want).abs() + } + + /// The whole reason for this change: the offset round-trip is EXACTLY lossless + /// when `w` is near 1 — which is why every other model family was unaffected — + /// and badly lossy when `w` is near 0, which is what DeepSeek-V4 ships. + #[test] + fn offset_roundtrip_is_lossless_near_one_and_lossy_near_zero() { + // Near 1 and above (every offset-convention model, and V4's own kv_norm in + // most layers): `w - 1` stays exactly representable, so the round-trip is + // bit-exact. This is the regression guarantee for other model families. + for x in [0.63_f32, 0.75, 0.875, 0.99, 1.0, 1.5, 2.0, 3.828125] { + let w = ckpt(x); + assert_eq!( + old_effective(w), + w, + "w={w} must round-trip EXACTLY under the offset convention" + ); + } + + // Near 0 — real `attn_norm` values from DeepSeek-V4-Flash layers.0. + // Catastrophic cancellation: a large RELATIVE error on the weight itself. + for x in [0.028931_f32, 0.030762, 0.032471, 0.033447] { + let w = ckpt(x); + let old_err = rel_err(old_effective(w), w); + assert!( + old_err > 0.01, + "attn_norm w={w}: the old path should be >1% wrong, got {:.2}%", + old_err * 100.0 + ); + // The exact load is lossless by construction: the checkpoint is BF16. + assert_eq!(ckpt(w), w); + } + } + + /// `q_norm`'s small tail — the old path is ~20 % wrong on the weight. + #[test] + fn q_norm_worst_case_is_badly_wrong_under_the_offset_convention() { + let w = ckpt(0.009766); + let old_err = rel_err(old_effective(w), w); + assert!( + old_err > 0.10, + "expected >10% relative error, got {:.2}%", + old_err * 100.0 + ); + assert_eq!(ckpt(w), w); + } + + /// Compressor / indexer norms straddle zero. The offset round-trip does not + /// merely degrade them — `bf16(w - 1)` rounds to exactly -1.0, so `1 + (-1.0)` + /// collapses the weight to ZERO, destroying sign and magnitude together. + /// No downstream precision can recover that. + #[test] + fn compressor_norm_sign_collapse_under_the_offset_convention() { + let w = ckpt(-0.001); + let old = old_effective(w); + assert_eq!(old, 0.0, "expected collapse to zero, got {old}"); + assert!(w < 0.0, "the true weight is negative"); + assert_eq!(ckpt(w), w, "the exact load preserves it"); + } +} diff --git a/crates/spark-model/src/model/impl_a1.rs b/crates/spark-model/src/model/impl_a1.rs index 55e5cb220..7bcadc6ba 100644 --- a/crates/spark-model/src/model/impl_a1.rs +++ b/crates/spark-model/src/model/impl_a1.rs @@ -61,7 +61,14 @@ impl TransformerModel { ssm_cache_slots: usize, ssm_checkpoint_interval: usize, ) -> Result { - let rms_norm_kernel = gpu.kernel("norm", "rms_norm")?; + // `rms_norm_kernel` normalizes exactly one weight: `final_norm` (a + // checkpoint tensor). Models that ship HF-vanilla norm weights load it + // exactly and must use the vanilla kernel. + let rms_norm_kernel = if crate::ships_vanilla_norm_weights(&config) { + gpu.kernel("rms_norm_vanilla", "rms_norm_vanilla")? + } else { + gpu.kernel("norm", "rms_norm")? + }; let dense_gemv_kernel = gpu.kernel("gemv", "dense_gemv_bf16")?; // FP32-output dense GEMV — the FP32 logits path required an FP32 // residual stream, which no longer exists, so this stays diff --git a/crates/spark-model/src/weight_loader/deepseek_v4.rs b/crates/spark-model/src/weight_loader/deepseek_v4.rs index 9b9d1f942..f572cd1da 100644 --- a/crates/spark-model/src/weight_loader/deepseek_v4.rs +++ b/crates/spark-model/src/weight_loader/deepseek_v4.rs @@ -20,7 +20,7 @@ use spark_runtime::weights::WeightStore; use super::ModelWeightLoader; use crate::layer::TransformerLayer; -use crate::weight_map::{DenseWeight, MtpWeights, dense, dense_minus_one}; +use crate::weight_map::{DenseWeight, MtpWeights, dense, dense_auto}; pub struct DeepSeekV4WeightLoader; @@ -66,15 +66,16 @@ impl ModelWeightLoader for DeepSeekV4WeightLoader { _config: &ModelConfig, _gpu: &dyn GpuBackend, ) -> Result { - // DeepSeek-V4 uses STANDARD RMSNorm (scale = weight); subtract 1.0 so the - // offset-from-1 rms_norm kernel computes `1 + (w-1) = w`. See dense_minus_one. - if let Ok(w) = dense_minus_one(store, "norm.weight", _gpu) { + // DeepSeek-V4 ships HF-vanilla RMSNorm weights (scale = weight). Load them + // EXACTLY; the model dispatches `rms_norm_vanilla` (see + // `crate::ships_vanilla_norm_weights`). + if let Ok(w) = dense_auto(store, "norm.weight", _gpu) { return Ok(w); } - if let Ok(w) = dense_minus_one(store, "model.norm.weight", _gpu) { + if let Ok(w) = dense_auto(store, "model.norm.weight", _gpu) { return Ok(w); } - dense_minus_one(store, "final_norm.weight", _gpu) + dense_auto(store, "final_norm.weight", _gpu) .context("DeepSeek-V4: no final norm tensor found (tried norm.weight, model.norm.weight, final_norm.weight)") } diff --git a/crates/spark-model/src/weight_loader/deepseek_v4/assemble.rs b/crates/spark-model/src/weight_loader/deepseek_v4/assemble.rs index 2eb30c6c0..2f82c6070 100644 --- a/crates/spark-model/src/weight_loader/deepseek_v4/assemble.rs +++ b/crates/spark-model/src/weight_loader/deepseek_v4/assemble.rs @@ -14,8 +14,8 @@ use crate::layers::qwen3_attention::{ CompressorWeights, HcHeadWeights, HcSiteWeights, HcWeights, MlaWeights, Qwen3AttentionLayer, }; use crate::weight_map::{ - AttentionWeights, DenseWeight, ExpertWeight, MoeWeights, QuantizedWeight, dense, - dense_minus_one, quantized, quantized_v2, + AttentionWeights, DenseWeight, ExpertWeight, MoeWeights, QuantizedWeight, dense, dense_auto, + quantized, quantized_v2, }; /// Load one MoE expert projection, dispatching by the on-disk format so the V4 @@ -375,7 +375,7 @@ pub fn assemble_layer( let wkv = dense(store, &format!("{cp}.wkv.weight"))?; let wgate = dense(store, &format!("{cp}.wgate.weight"))?; // compressor.norm is a STANDARD RMSNorm → subtract 1 for the offset kernel. - let norm = dense_minus_one(store, &format!("{cp}.norm.weight"), gpu)?; + let norm = dense_auto(store, &format!("{cp}.norm.weight"), gpu)?; let ape = store.get(&format!("{cp}.ape"))?.ptr; // 4b: allocate the persistent flat compressed-KV pool for this layer. // Sized to the full context (max_position_embeddings // ratio blocks) @@ -549,6 +549,19 @@ pub fn assemble_layer( }); } + // V4 loads its norm weights EXACTLY (no -1 pre-subtraction) and normalizes + // with `rms_norm_vanilla`. The only paths that would bypass that kernel are + // the fused residual+norm ones, taken when the mHC highway is absent — and + // they are offset-from-1 only, with no vanilla twin in-tree. They would + // silently compute `x * (1 + w)` on an exact weight. Fail loudly instead. + anyhow::ensure!( + layer.hc.is_some(), + "DeepSeek-V4 requires the mHC highway (hc_mult > 0, got {}): without it the fused \ + residual+norm kernels would apply the offset-from-1 convention to exactly-loaded \ + norm weights", + config.hc_mult + ); + Ok(Box::new(layer)) } diff --git a/crates/spark-model/src/weight_loader/deepseek_v4/load_layers.rs b/crates/spark-model/src/weight_loader/deepseek_v4/load_layers.rs index d5f1d1510..41d75fc25 100644 --- a/crates/spark-model/src/weight_loader/deepseek_v4/load_layers.rs +++ b/crates/spark-model/src/weight_loader/deepseek_v4/load_layers.rs @@ -8,7 +8,7 @@ use spark_runtime::weights::WeightStore; use crate::layer::TransformerLayer; use crate::layers::qwen3_attention::HcHeadWeights; -use crate::weight_map::{DenseWeight, dense_auto, dense_minus_one}; +use crate::weight_map::{DenseWeight, dense_auto}; pub fn load_all_layers( store: &WeightStore, @@ -104,8 +104,8 @@ pub fn load_all_layers( let lp = format!("layers.{i}"); let ap = format!("{lp}.attn"); - let input_norm = dense_minus_one(store, &format!("{lp}.attn_norm.weight"), gpu)?; - let post_attn_norm = dense_minus_one(store, &format!("{lp}.ffn_norm.weight"), gpu)?; + let input_norm = dense_auto(store, &format!("{lp}.attn_norm.weight"), gpu)?; + let post_attn_norm = dense_auto(store, &format!("{lp}.ffn_norm.weight"), gpu)?; // DeepSeek-V4-Flash attention weights are FP8 block-quantized in the // checkpoint (config quant group_0: float-quantized, block [128,128]); @@ -117,10 +117,10 @@ pub fn load_all_layers( // already fall back to dense_gemv/dense_gemm when the nvfp4 view is None. let wq_a = dense_auto(store, &format!("{ap}.wq_a.weight"), gpu)?; let wq_b = dense_auto(store, &format!("{ap}.wq_b.weight"), gpu)?; - let q_a_norm = dense_minus_one(store, &format!("{ap}.q_norm.weight"), gpu)?; + let q_a_norm = dense_auto(store, &format!("{ap}.q_norm.weight"), gpu)?; let wkv_a = dense_auto(store, &format!("{ap}.wkv.weight"), gpu)?; - let kv_a_norm = dense_minus_one(store, &format!("{ap}.kv_norm.weight"), gpu)?; + let kv_a_norm = dense_auto(store, &format!("{ap}.kv_norm.weight"), gpu)?; let is_v4_flash = config.o_lora_rank > 0; let null = DenseWeight { diff --git a/crates/spark-model/src/weight_loader/deepseek_v4/mtp.rs b/crates/spark-model/src/weight_loader/deepseek_v4/mtp.rs index 7257a212c..39438fa22 100644 --- a/crates/spark-model/src/weight_loader/deepseek_v4/mtp.rs +++ b/crates/spark-model/src/weight_loader/deepseek_v4/mtp.rs @@ -26,7 +26,7 @@ use spark_runtime::weights::WeightStore; use crate::layer::TransformerLayer; use crate::layers::qwen3_attention::HcHeadWeights; -use crate::weight_map::{DenseWeight, dense_auto, dense_minus_one}; +use crate::weight_map::{DenseWeight, dense_auto}; /// A loaded DeepSeek-V4 MTP draft module: the reused V4 transformer body plus the /// MTP-specific input combiner and final norm. Embedding + lm_head are shared with @@ -91,13 +91,13 @@ pub fn load_v4_mtp_module( // Mirrors the V4-Flash branch of `load_all_layers` (o_lora_rank > 0): direct // KV projection + grouped low-rank O projection, so the V3 absorption tensors // (w_uk_t / w_uv / w_qk_absorbed / block-diagonals) are NULL/unused. - let input_norm = dense_minus_one(store, &format!("{prefix}.attn_norm.weight"), gpu)?; - let post_attn_norm = dense_minus_one(store, &format!("{prefix}.ffn_norm.weight"), gpu)?; + let input_norm = dense_auto(store, &format!("{prefix}.attn_norm.weight"), gpu)?; + let post_attn_norm = dense_auto(store, &format!("{prefix}.ffn_norm.weight"), gpu)?; let wq_a = dense_auto(store, &format!("{ap}.wq_a.weight"), gpu)?; let wq_b = dense_auto(store, &format!("{ap}.wq_b.weight"), gpu)?; - let q_a_norm = dense_minus_one(store, &format!("{ap}.q_norm.weight"), gpu)?; + let q_a_norm = dense_auto(store, &format!("{ap}.q_norm.weight"), gpu)?; let wkv_a = dense_auto(store, &format!("{ap}.wkv.weight"), gpu)?; - let kv_a_norm = dense_minus_one(store, &format!("{ap}.kv_norm.weight"), gpu)?; + let kv_a_norm = dense_auto(store, &format!("{ap}.kv_norm.weight"), gpu)?; let wo_a = dense_auto(store, &format!("{ap}.wo_a.weight"), gpu)?; let wo_b = dense_auto(store, &format!("{ap}.wo_b.weight"), gpu)?; @@ -164,11 +164,12 @@ pub fn load_v4_mtp_module( )?; // ── MTP-specific combiner + final norm ── - // enorm/hnorm/norm are standard RMSNorms (offset-from-1 kernel → dense_minus_one). + // enorm/hnorm/norm are HF-vanilla RMSNorms — loaded exactly, normalized by + // `rms_norm_vanilla` (see `DeepseekV4MtpHead::rms_norm_k`). // e_proj/h_proj are FP8 block-scaled linears in the checkpoint (dense_auto dequants). - let enorm = dense_minus_one(store, &format!("{prefix}.enorm.weight"), gpu)?; - let hnorm = dense_minus_one(store, &format!("{prefix}.hnorm.weight"), gpu)?; - let norm = dense_minus_one(store, &format!("{prefix}.norm.weight"), gpu)?; + let enorm = dense_auto(store, &format!("{prefix}.enorm.weight"), gpu)?; + let hnorm = dense_auto(store, &format!("{prefix}.hnorm.weight"), gpu)?; + let norm = dense_auto(store, &format!("{prefix}.norm.weight"), gpu)?; let e_proj = dense_auto(store, &format!("{prefix}.e_proj.weight"), gpu)?; let h_proj = dense_auto(store, &format!("{prefix}.h_proj.weight"), gpu)?; diff --git a/crates/spark-model/src/weight_map/model_a.rs b/crates/spark-model/src/weight_map/model_a.rs index 534a8c749..eb0fd2129 100644 --- a/crates/spark-model/src/weight_map/model_a.rs +++ b/crates/spark-model/src/weight_map/model_a.rs @@ -166,40 +166,20 @@ pub(crate) fn dense(store: &WeightStore, name: &str) -> Result { Ok(DenseWeight { weight: w.ptr }) } -/// Load a BF16 norm weight and subtract 1.0 from every element. -/// -/// Atlas's `rms_norm` kernel uses the Qwen3-Next "offset-from-1" convention -/// (`out = x * (1 + weight)`). Models with STANDARD RMSNorm (`out = x * weight`, -/// e.g. DeepSeek-V4: `DeepseekV4RMSNorm` = T5LayerNorm) must pre-subtract 1.0 so -/// the kernel computes `1 + (w - 1) = w`. Without this, every norm is scaled -/// wrong (e.g. kv_norm 2.6x, attn_norm ~30x too large) → attention overflow. -pub(crate) fn dense_minus_one( - store: &WeightStore, - name: &str, - gpu: &dyn GpuBackend, -) -> Result { - let w = store.get(name)?; - let n = w.num_elements(); - let mut bf16_buf = vec![0u8; n * 2]; - gpu.copy_d2h(w.ptr, &mut bf16_buf)?; - let adjusted: Vec = bf16_buf - .chunks_exact(2) - .flat_map(|c| { - let bits = u16::from_le_bytes([c[0], c[1]]); - let v = f32::from_bits((bits as u32) << 16) - 1.0; - // round-to-nearest-even bf16 truncation - let u = v.to_bits(); - let round_bit = (u >> 15) & 1; - let sticky = (u & 0x7FFF != 0) as u32; - let bf = - ((u >> 16) as u16).wrapping_add((round_bit & (sticky | ((u >> 16) & 1))) as u16); - bf.to_le_bytes() - }) - .collect(); - let ptr = gpu.alloc(adjusted.len())?; - gpu.copy_h2d(&adjusted, ptr)?; - Ok(DenseWeight { weight: ptr }) -} +// REMOVED: `dense_minus_one` — the offset-from-1 norm loader. +// +// It pre-subtracted 1.0 and stored `bf16(w - 1)` so the offset-from-1 `rms_norm` +// kernel would recover `1 + (w - 1) = w`. That round-trip is only lossless when +// `w ≈ 1`. DeepSeek-V4 — its ONLY caller — ships HF-vanilla norm weights of +// ≈ 0.03, so it stored ≈ −0.97 and BF16's rounding error there (~1.9e-3 absolute) +// became a 1.8–3.4 % RELATIVE error on the weight once 1 was added back +// (catastrophic cancellation; up to 19 % on `q_norm`, 100 % with sign flips on the +// compressor norms — measured over all 249 V4 norm tensors, 2026-07-13). +// +// V4 now loads its norm weights exactly (`dense_auto`) and dispatches +// `rms_norm_vanilla`. See `crate::ships_vanilla_norm_weights`. Models whose norm +// weights are genuinely stored as an offset (Qwen3-Next, init 0) use `dense` and +// keep the offset kernel — they never needed this function. /// Load a weight, auto-dequanting FP8 block-scaled to BF16 when needed. ///