From 4d7cd50c524eef66d4a03858b2317f5f7e8fc2e2 Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:56:47 -0400 Subject: [PATCH 01/26] atlas-core: host decoders for ModelOpt FP4-Local storage formats CSF-compressed NVFP4 block scales (lossless, per the layout vLLM's csf_scales.py specifies), exact MXFP8 -> BF16 expansion, and NVFP4 rows -> BF16 for row-gathered tables. Pure arithmetic, no env lookups in the hot loops. Checked on nvidia/Qwen3.8-Flash-Next-FP4-Local with examples/fp4local_decode: CSF bytes identical to a numpy port of the vLLM reference decoder on four expert tensors (81-762 exceptions each); dequantized experts cos 0.994, MXFP8 projections cos 0.9996 and PLE rows cos 0.995 against the NVFP4 pack. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01L4aUGgtiZX94pLr63HXQzn --- crates/atlas-core/examples/fp4local_decode.rs | 116 +++++++++++ crates/atlas-core/src/lib.rs | 1 + crates/atlas-core/src/modelopt_codec.rs | 194 ++++++++++++++++++ crates/atlas-core/src/modelopt_codec_tests.rs | 153 ++++++++++++++ 4 files changed, 464 insertions(+) create mode 100644 crates/atlas-core/examples/fp4local_decode.rs create mode 100644 crates/atlas-core/src/modelopt_codec.rs create mode 100644 crates/atlas-core/src/modelopt_codec_tests.rs diff --git a/crates/atlas-core/examples/fp4local_decode.rs b/crates/atlas-core/examples/fp4local_decode.rs new file mode 100644 index 000000000..3c9388fca --- /dev/null +++ b/crates/atlas-core/examples/fp4local_decode.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Decode tensors of a ModelOpt FP4-Local checkpoint with +//! `atlas_core::modelopt_codec` and write the raw result, so an independent +//! reference can check it byte for byte. +//! +//! ```text +//! fp4local_decode csf # E4M3 bytes +//! fp4local_decode mxfp8 # BF16 bytes +//! fp4local_decode ple # BF16 row +//! ``` + +use anyhow::{Context, Result, bail}; +use atlas_core::modelopt_codec::{decode_csf_scales, mxfp8_to_bf16, nvfp4_rows_to_bf16}; +use std::io::{Read, Seek, SeekFrom}; +use std::path::Path; + +struct Tensor { + file: std::path::PathBuf, + offset: u64, + len: usize, + shape: Vec, +} + +fn locate(snap: &Path, name: &str) -> Result { + let index: serde_json::Value = + serde_json::from_slice(&std::fs::read(snap.join("model.safetensors.index.json"))?)?; + let shard = index["weight_map"][name] + .as_str() + .with_context(|| format!("{name} is not in the index"))?; + let file = snap.join(shard); + let mut f = std::fs::File::open(&file)?; + let mut n = [0u8; 8]; + f.read_exact(&mut n)?; + let hlen = u64::from_le_bytes(n); + let mut header = vec![0u8; hlen as usize]; + f.read_exact(&mut header)?; + let header: serde_json::Value = serde_json::from_slice(&header)?; + let t = &header[name]; + let offs = t["data_offsets"].as_array().context("data_offsets")?; + let (a, b) = (offs[0].as_u64().unwrap(), offs[1].as_u64().unwrap()); + let shape = t["shape"] + .as_array() + .context("shape")? + .iter() + .map(|v| v.as_u64().unwrap() as usize) + .collect(); + Ok(Tensor { + file, + offset: 8 + hlen + a, + len: (b - a) as usize, + shape, + }) +} + +fn read(t: &Tensor, start: usize, len: usize) -> Result> { + let mut f = std::fs::File::open(&t.file)?; + f.seek(SeekFrom::Start(t.offset + start as u64))?; + let mut buf = vec![0u8; len]; + f.read_exact(&mut buf)?; + Ok(buf) +} + +fn main() -> Result<()> { + let args: Vec = std::env::args().collect(); + let snap = Path::new(args.get(1).context("snapshot dir")?); + let out = match args.get(2).map(String::as_str) { + Some("csf") => { + let name = &args[3]; + let fixed = locate(snap, &format!("{name}.nvfp4_csf_fixed"))?; + let exc = locate(snap, &format!("{name}.nvfp4_csf_exceptions"))?; + decode_csf_scales( + &read(&fixed, 0, fixed.len)?, + fixed.shape[0], + fixed.shape[1], + &read(&exc, 0, exc.len)?, + )? + } + Some("mxfp8") => { + let m = &args[3]; + let w = locate(snap, &format!("{m}.weight"))?; + let s = locate(snap, &format!("{m}.weight_scale"))?; + mxfp8_to_bf16( + &read(&w, 0, w.len)?, + &read(&s, 0, s.len)?, + w.shape[0], + w.shape[1], + )? + } + Some("ple") => { + let (base, row): (&str, usize) = (&args[3], args[4].parse()?); + let w = locate(snap, &format!("{base}.weight"))?; + let s = locate(snap, &format!("{base}.weight_scale"))?; + let table = base.rsplit_once(".shard_").context("shard base")?.0; + let g = locate(snap, &format!("{table}.weight_scale_2"))?; + let g = read(&g, 0, 4)?; + let global = f32::from_le_bytes([g[0], g[1], g[2], g[3]]); + let (pw, sw) = (w.shape[1], s.shape[1]); + let dim = pw * 2; + let mut dst = vec![0u8; dim * 2]; + nvfp4_rows_to_bf16( + &read(&w, row * pw, pw)?, + &read(&s, row * sw, sw)?, + global, + dim, + dim / sw, + &mut dst, + )?; + dst + } + _ => bail!("mode must be csf | mxfp8 | ple"), + }; + std::fs::write(args.last().unwrap(), &out)?; + println!("wrote {} bytes", out.len()); + Ok(()) +} diff --git a/crates/atlas-core/src/lib.rs b/crates/atlas-core/src/lib.rs index d83e7e49d..935d93a04 100644 --- a/crates/atlas-core/src/lib.rs +++ b/crates/atlas-core/src/lib.rs @@ -9,6 +9,7 @@ pub mod config; pub mod dtype; pub mod error; pub mod fault; +pub mod modelopt_codec; pub mod ngram_table; pub mod numeric; pub mod qwen4exp_reference; diff --git a/crates/atlas-core/src/modelopt_codec.rs b/crates/atlas-core/src/modelopt_codec.rs new file mode 100644 index 000000000..dc950d359 --- /dev/null +++ b/crates/atlas-core/src/modelopt_codec.rs @@ -0,0 +1,194 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Host-side decoders for the storage formats NVIDIA's ModelOpt "FP4-Local" +//! packs use and Atlas kernels do not read directly: +//! +//! * **CSF NVFP4 block scales** (`weight_scale_encoding: "csf"`, codec +//! `byte-window4-fixed-stream-u24-exceptions/1`). Each E4M3 block-scale +//! matrix `[rows, cols]` ships as two tensors: +//! - `.nvfp4_csf_fixed`, U8 `[rows / 16, 16 * (1 + cols / 2)]`: per +//! 16-row slab, one base byte per row, then those rows' 4-bit codes, two +//! per byte, low nibble first. A scale byte is `base + code`. +//! - `.nvfp4_csf_exceptions`, U32: words whose low 24 bits are a flat +//! `row * cols + col` position and whose high 8 bits replace the byte +//! there. +//! +//! The decode is lossless: the bytes equal the uncompressed E4M3 scales, so +//! every NVFP4 kernel consumes them unchanged. Format as specified by +//! vLLM's `csf_scales.py` (Apache-2.0, vllm-project/vllm#60646); this is an +//! independent Rust implementation of that published layout. +//! +//! * **MXFP8** (OCP microscaling): an E4M3 weight `[out, in]` with one E8M0 +//! exponent per 32 inputs, `[out, in / 32]`. Expanded to BF16 EXACTLY: an +//! E4M3 value has 4 significant bits and the scale is a power of two, so the +//! product always fits BF16's 8. +//! +//! * **NVFP4 rows → BF16**, for a table gathered by row (the PLE n-gram +//! embedding), where the block scales live in a separate tensor. +//! +//! Pure arithmetic, no CUDA, no environment lookups in the hot loops. + +use crate::numeric::{FP4_E2M1, FP8_E4M3_LUT}; +use anyhow::{Result, bail, ensure}; + +/// Rows per CSF slab; every slab stores 16 base bytes ahead of its codes. +pub const CSF_SLAB_ROWS: usize = 16; +/// Largest legal row base: a base above 240 cannot reach 255 with a nibble. +pub const CSF_MAX_BASE: u8 = 240; +/// Inputs per MXFP8 scale. +pub const MXFP8_BLOCK: usize = 32; + +/// Shape `[rows, cols]` of the E4M3 scale matrix a CSF fixed stream encodes. +pub fn csf_scale_shape(fixed_rows: usize, fixed_cols: usize) -> Result<(usize, usize)> { + ensure!( + fixed_cols.is_multiple_of(CSF_SLAB_ROWS) && fixed_cols >= CSF_SLAB_ROWS, + "CSF fixed stream width {fixed_cols} is not 16 * (1 + cols / 2)" + ); + let cols = (fixed_cols / CSF_SLAB_ROWS - 1) * 2; + ensure!(cols > 0, "CSF fixed stream encodes zero columns"); + Ok((fixed_rows * CSF_SLAB_ROWS, cols)) +} + +/// Rebuild a `[rows, cols]` E4M3 block-scale matrix (as raw bytes) from its +/// CSF fixed stream `[fixed_rows, fixed_cols]` and little-endian U32 +/// exception words. +pub fn decode_csf_scales( + fixed: &[u8], + fixed_rows: usize, + fixed_cols: usize, + exceptions: &[u8], +) -> Result> { + let (rows, cols) = csf_scale_shape(fixed_rows, fixed_cols)?; + ensure!( + fixed.len() == fixed_rows * fixed_cols, + "CSF fixed stream is {} bytes, shape [{fixed_rows}, {fixed_cols}] needs {}", + fixed.len(), + fixed_rows * fixed_cols + ); + ensure!( + exceptions.len().is_multiple_of(4), + "CSF exception stream is {} bytes, not whole U32 words", + exceptions.len() + ); + let half = cols / 2; + let mut out = vec![0u8; rows * cols]; + for (slab, src) in fixed.chunks_exact(fixed_cols).enumerate() { + let (bases, codes) = src.split_at(CSF_SLAB_ROWS); + for (i, &base) in bases.iter().enumerate() { + if base > CSF_MAX_BASE { + bail!("CSF row base {base} exceeds {CSF_MAX_BASE}"); + } + let row = slab * CSF_SLAB_ROWS + i; + let packed = &codes[i * half..(i + 1) * half]; + let dst = &mut out[row * cols..(row + 1) * cols]; + for (pair, &byte) in dst.chunks_exact_mut(2).zip(packed) { + pair[0] = base + (byte & 0x0F); + pair[1] = base + (byte >> 4); + } + } + } + for word in exceptions.chunks_exact(4) { + let w = u32::from_le_bytes([word[0], word[1], word[2], word[3]]); + let pos = (w & 0x00FF_FFFF) as usize; + ensure!( + pos < out.len(), + "CSF exception position {pos} is outside the {rows}x{cols} scale matrix" + ); + out[pos] = (w >> 24) as u8; + } + Ok(out) +} + +/// The top 16 bits of an f32 whose low 16 bits are already zero (or whose +/// rounding the caller has accepted). Round-to-nearest-even; NaN stays NaN. +#[inline(always)] +fn bf16_rne(val: f32) -> u16 { + let bits = val.to_bits(); + if val.is_nan() { + return ((bits >> 16) as u16 & 0x8000) | 0x7FC0; + } + let lsb = (bits >> 16) & 1; + (bits.wrapping_add(0x7FFF + lsb) >> 16) as u16 +} + +/// `2^(e - 127)` for an E8M0 exponent byte. `0xFF` is NaN in OCP MX; it maps +/// to zero for the same reason the E4M3 table maps NaN to zero. +#[inline(always)] +fn e8m0_to_f32(e: u8) -> f32 { + match e { + 0xFF => 0.0, + // 2^-127 is an f32 subnormal: build it from its bits. + 0 => f32::from_bits(0x0040_0000), + _ => f32::from_bits((e as u32) << 23), + } +} + +/// Expand an MXFP8 weight `[out, inp]` (E4M3 bytes) with its E8M0 scales +/// `[out, inp / 32]` into little-endian BF16 bytes `[out, inp]`. Exact. +pub fn mxfp8_to_bf16(weight: &[u8], scale: &[u8], out: usize, inp: usize) -> Result> { + ensure!( + inp.is_multiple_of(MXFP8_BLOCK), + "MXFP8 input width {inp} is not a multiple of {MXFP8_BLOCK}" + ); + let blocks = inp / MXFP8_BLOCK; + ensure!( + weight.len() == out * inp && scale.len() == out * blocks, + "MXFP8 weight {} / scale {} bytes do not match [{out}, {inp}]", + weight.len(), + scale.len() + ); + let mut dst = vec![0u8; out * inp * 2]; + for ((w, s), d) in weight + .chunks_exact(MXFP8_BLOCK) + .zip(scale) + .zip(dst.chunks_exact_mut(MXFP8_BLOCK * 2)) + { + let mult = e8m0_to_f32(*s); + for (&b, o) in w.iter().zip(d.chunks_exact_mut(2)) { + o.copy_from_slice(&bf16_rne(FP8_E4M3_LUT[b as usize] * mult).to_le_bytes()); + } + } + Ok(dst) +} + +/// Dequantize NVFP4 rows to little-endian BF16. `packed` is `[rows, dim / 2]` +/// (low nibble first), `scale` is `[rows, dim / group]` E4M3, `global` is the +/// tensor's `weight_scale_2`. Writes `rows * dim` BF16 values into `dst`. +pub fn nvfp4_rows_to_bf16( + packed: &[u8], + scale: &[u8], + global: f32, + dim: usize, + group: usize, + dst: &mut [u8], +) -> Result<()> { + ensure!( + dim.is_multiple_of(2) && group > 0 && dim.is_multiple_of(group), + "NVFP4 row width {dim} must be even and a multiple of group {group}" + ); + let rows = packed.len() / (dim / 2); + ensure!( + packed.len() == rows * dim / 2 + && scale.len() == rows * dim / group + && dst.len() == rows * dim * 2, + "NVFP4 rows: packed {} / scale {} / dst {} bytes disagree on {rows} rows of {dim}", + packed.len(), + scale.len(), + dst.len() + ); + for (i, (&byte, o)) in packed.iter().zip(dst.chunks_exact_mut(4)).enumerate() { + let elem = i * 2; + // `(fp4 * block) * global`, the order `numeric::nvfp4_dequant` uses: + // the first product is exact, so there is exactly one f32 rounding. + let block = FP8_E4M3_LUT[scale[elem / group] as usize]; + let lo = bf16_rne(FP4_E2M1[(byte & 0x0F) as usize] * block * global); + let hi = bf16_rne(FP4_E2M1[(byte >> 4) as usize] * block * global); + o[..2].copy_from_slice(&lo.to_le_bytes()); + o[2..].copy_from_slice(&hi.to_le_bytes()); + } + Ok(()) +} + +#[cfg(test)] +#[path = "modelopt_codec_tests.rs"] +mod tests; diff --git a/crates/atlas-core/src/modelopt_codec_tests.rs b/crates/atlas-core/src/modelopt_codec_tests.rs new file mode 100644 index 000000000..4a9d152cd --- /dev/null +++ b/crates/atlas-core/src/modelopt_codec_tests.rs @@ -0,0 +1,153 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +use super::*; +use crate::numeric::{f32_to_bf16, nvfp4_dequant}; + +/// Deterministic xorshift so the tests need no RNG dependency. +struct Rng(u64); +impl Rng { + fn next(&mut self) -> u64 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + self.0 + } + fn byte_in(&mut self, lo: u8, hi: u8) -> u8 { + lo + (self.next() % u64::from(hi - lo)) as u8 + } +} + +/// Reference CSF encoder: row base = row minimum capped at 240, any byte the +/// 4-bit window cannot reach becomes an exception. Mirrors the encoder in +/// vLLM's `test_csf_scales.py`, which is how the published layout is pinned. +fn encode(raw: &[u8], rows: usize, cols: usize) -> (Vec, usize, usize, Vec) { + assert!(rows.is_multiple_of(CSF_SLAB_ROWS) && cols.is_multiple_of(2)); + let fixed_cols = CSF_SLAB_ROWS * (1 + cols / 2); + let mut fixed = vec![0u8; rows / CSF_SLAB_ROWS * fixed_cols]; + let mut exceptions = Vec::new(); + for r in 0..rows { + let row = &raw[r * cols..(r + 1) * cols]; + let base = (*row.iter().min().unwrap()).min(CSF_MAX_BASE); + let (slab, i) = (r / CSF_SLAB_ROWS, r % CSF_SLAB_ROWS); + let slab_bytes = &mut fixed[slab * fixed_cols..(slab + 1) * fixed_cols]; + slab_bytes[i] = base; + for (c, &byte) in row.iter().enumerate() { + let code = i32::from(byte) - i32::from(base); + let code = if (0..=15).contains(&code) { + code as u8 + } else { + let pos = (r * cols + c) as u32; + exceptions.extend_from_slice(&((u32::from(byte) << 24) | pos).to_le_bytes()); + 0 + }; + let at = CSF_SLAB_ROWS + i * (cols / 2) + c / 2; + slab_bytes[at] |= if c % 2 == 0 { code } else { code << 4 }; + } + } + (fixed, rows / CSF_SLAB_ROWS, fixed_cols, exceptions) +} + +/// The two routed-expert scale shapes Flash-Next ships: gate/up `[640, 160]` +/// and down `[2560, 40]`. +#[test] +fn csf_round_trips_the_flash_next_expert_shapes() { + for (rows, cols) in [(640, 160), (2560, 40)] { + let mut rng = Rng(0x9E37_79B9_7F4A_7C15 ^ (rows * cols) as u64); + let mut raw: Vec = (0..rows * cols).map(|_| rng.byte_in(0x30, 0x40)).collect(); + for _ in 0..64 { + let at = (rng.next() as usize) % raw.len(); + raw[at] = 0x7E; // far outside every row's window + } + let (fixed, fr, fc, exc) = encode(&raw, rows, cols); + assert_eq!(csf_scale_shape(fr, fc).unwrap(), (rows, cols)); + assert!(exc.len() / 4 >= 64, "outliers must travel as exceptions"); + assert_eq!(decode_csf_scales(&fixed, fr, fc, &exc).unwrap(), raw); + } +} + +/// The real checkpoint's down_proj stream: `[160, 336]` fixed encodes +/// `[2560, 40]` scales. +#[test] +fn csf_shape_matches_the_checkpoint_header() { + assert_eq!(csf_scale_shape(160, 336).unwrap(), (2560, 40)); + assert_eq!(csf_scale_shape(40, 1296).unwrap(), (640, 160)); + assert!(csf_scale_shape(40, 1295).is_err()); +} + +#[test] +fn csf_rejects_bad_streams() { + let raw = vec![0x38u8; 16 * 4]; + let (mut fixed, fr, fc, _) = encode(&raw, 16, 4); + // Exception past the end of the matrix. + let bad = ((0x38u32 << 24) | 64).to_le_bytes(); + assert!(decode_csf_scales(&fixed, fr, fc, &bad).is_err()); + // A torn U32 word. + assert!(decode_csf_scales(&fixed, fr, fc, &[1, 2, 3]).is_err()); + // Wrong byte count for the declared shape. + assert!(decode_csf_scales(&fixed[1..], fr, fc, &[]).is_err()); + // A base no nibble can lift to 255 is not a legal stream. + fixed[0] = 241; + assert!(decode_csf_scales(&fixed, fr, fc, &[]).is_err()); +} + +/// MXFP8 -> BF16 must be EXACT: compare against an f64 reconstruction that +/// shares no code with the decoder. +#[test] +fn mxfp8_expansion_is_exact() { + let (out, inp) = (8, 96); + let mut rng = Rng(42); + let weight: Vec = (0..out * inp) + .map(|_| { + let b = (rng.next() & 0xFF) as u8; + if b & 0x7F == 0x7F { 0x38 } else { b } // skip NaN codes + }) + .collect(); + let scale: Vec = (0..out * inp / 32).map(|_| rng.byte_in(110, 135)).collect(); + let got = mxfp8_to_bf16(&weight, &scale, out, inp).unwrap(); + for (i, &w) in weight.iter().enumerate() { + let (sign, exp, man) = (w >> 7, (w >> 3) & 0xF, w & 7); + let mag = if exp == 0 { + f64::from(man) / 8.0 * 2f64.powi(-6) + } else { + (1.0 + f64::from(man) / 8.0) * 2f64.powi(i32::from(exp) - 7) + }; + let signed = if sign == 1 { -mag } else { mag }; + let v = signed * 2f64.powi(i32::from(scale[i / 32]) - 127); + let bits = u16::from_le_bytes([got[2 * i], got[2 * i + 1]]); + let back = f64::from(f32::from_bits(u32::from(bits) << 16)); + assert_eq!( + back, + v, + "element {i}: byte {w:#04x} scale {}", + scale[i / 32] + ); + } +} + +#[test] +fn mxfp8_rejects_misaligned_widths() { + assert!(mxfp8_to_bf16(&[0; 48], &[127; 1], 1, 48).is_err()); + assert!(mxfp8_to_bf16(&[0; 64], &[127; 1], 1, 64).is_err()); +} + +/// NVFP4 rows decode to the same values as the existing whole-matrix +/// dequant, cast through the shared BF16 rounding. +#[test] +fn nvfp4_rows_match_the_reference_dequant() { + let (rows, dim, group) = (5, 160, 16); + let mut rng = Rng(7); + let packed: Vec = (0..rows * dim / 2) + .map(|_| (rng.next() & 0xFF) as u8) + .collect(); + let scale: Vec = (0..rows * dim / group) + .map(|_| rng.byte_in(0x20, 0x60)) + .collect(); + let global = 3.1e-4f32; + let mut got = vec![0u8; rows * dim * 2]; + nvfp4_rows_to_bf16(&packed, &scale, global, dim, group, &mut got).unwrap(); + let want = nvfp4_dequant(&packed, &scale, global, rows, dim, group).unwrap(); + for (i, w) in want.iter().enumerate() { + let bits = u16::from_le_bytes([got[2 * i], got[2 * i + 1]]); + assert_eq!(bits, f32_to_bf16(*w), "element {i}"); + } +} From 4f731d6761e022f0fefee41e1b1c22ef74899132 Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:56:48 -0400 Subject: [PATCH 02/26] spark-runtime: load FP4-Local checkpoints without making their compressed tensors resident - U32 maps to a raw 4-byte container (CSF exception streams); the fast loader's header now defers to the shared dtype table instead of a copy. - The fast loader defers CSF streams and the weight/scale pair of every module in host_decode_modules, so the build expands them from disk. - is_ngram_table also takes the per-shard NVFP4 block scales of the PLE table, which are row data like the weight. - An index that names shards under a directory (tensors/model-*.safetensors) routed zero tensors to every shard; match shard paths, not file names. - DeferredTensor::read_bytes / WeightStore::undefer for build-time consumers. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01L4aUGgtiZX94pLr63HXQzn --- .../spark-runtime/src/fast_weights/header.rs | 40 ++++++++++------- .../src/fast_weights/header_tests.rs | 45 +++++++++++++++++++ crates/spark-runtime/src/fast_weights/mod.rs | 45 ++++++++++++++++--- crates/spark-runtime/src/fast_weights/skip.rs | 20 ++++++++- crates/spark-runtime/src/weights.rs | 8 +++- .../spark-runtime/src/weights/name_utils.rs | 21 ++++++++- .../spark-runtime/src/weights/store_edit.rs | 38 ++++++++++++++++ .../src/weights/store_edit_tests.rs | 28 ++++++++++++ 8 files changed, 220 insertions(+), 25 deletions(-) diff --git a/crates/spark-runtime/src/fast_weights/header.rs b/crates/spark-runtime/src/fast_weights/header.rs index 41a788c78..5af91ce8a 100644 --- a/crates/spark-runtime/src/fast_weights/header.rs +++ b/crates/spark-runtime/src/fast_weights/header.rs @@ -32,6 +32,24 @@ pub(super) struct TensorMeta { pub(super) len: usize, } +/// The tensors an index routes to `shard_path`. +/// +/// Compared as PATHS under `model_dir`, not by file name: an index may name +/// its shards with a directory (nvidia's FP4-Local pack routes every tensor to +/// `tensors/model-000NN-of-00021.safetensors`), and a bare-name comparison then +/// matches nothing — every shard "loads" zero tensors and the build fails far +/// from the cause, on a missing embedding. +pub(super) fn tensors_routed_to( + map: &HashMap, + model_dir: &Path, + shard_path: &Path, +) -> Vec { + map.iter() + .filter(|(_, s)| model_dir.join(s.as_str()) == shard_path) + .map(|(t, _)| t.clone()) + .collect() +} + /// Discover which safetensor files to load. Mirrors the resolution order in /// `crate::weights::SafetensorsLoader::load`. /// @@ -124,10 +142,6 @@ pub(super) fn parse_header(file: &mut File) -> Result> { } let dtype_str = info["dtype"].as_str().unwrap_or("BF16"); let (dtype, from_f16) = match dtype_str { - "F32" => (WeightDtype::FP32, false), - "BF16" => (WeightDtype::BF16, false), - "I16" => (WeightDtype::Int16, false), - "I32" => (WeightDtype::Int32, false), // F16 is not store-legal (WeightDtype is closed to store dtypes): // stage as BF16 and mark for byte conversion in the copy loop. // centml modelopt W4A4 exports ship all unquantized tensors as F16. @@ -139,17 +153,13 @@ pub(super) fn parse_header(file: &mut File) -> Result> { (WeightDtype::BF16, true) } } - "U8" => (WeightDtype::UInt8, false), - // I8 is a 1-byte raw container; DeepSeek-V4-Flash-NVFP4 ships its MTP - // experts' 4-bit-packed weights as I8 (vs U8 for the main layers). - // Signedness is irrelevant for packed FP4 — the dequant kernel extracts - // nibbles by bit ops — so treat I8 as raw bytes (UInt8), matching the - // NVFP4 expert path. - "I8" => (WeightDtype::UInt8, false), - "F8_E4M3" => (WeightDtype::FP8E4M3, false), - "F8_E8M0" => (WeightDtype::FP8E8M0, false), - "I64" => (WeightDtype::Int64, false), - other => bail!("Unsupported safetensors dtype '{other}' for tensor {name}"), + // Every other dtype maps exactly as on the buffered and RDMA paths: + // one table (`from_safetensors_str`), so the loaders cannot drift. + other => ( + WeightDtype::from_safetensors_str(other) + .with_context(|| format!("tensor {name}"))?, + false, + ), }; let shape: Vec = info["shape"] .as_array() diff --git a/crates/spark-runtime/src/fast_weights/header_tests.rs b/crates/spark-runtime/src/fast_weights/header_tests.rs index f2715250a..222ff22d1 100644 --- a/crates/spark-runtime/src/fast_weights/header_tests.rs +++ b/crates/spark-runtime/src/fast_weights/header_tests.rs @@ -73,3 +73,48 @@ fn unknown_dtype_is_still_refused() { }; assert!(format!("{err:#}").contains("C64"), "{err:#}"); } + +/// FP4-Local's CSF exception streams are U32; they stage as a raw 4-byte +/// container exactly as the buffered loader maps them. +#[test] +fn u32_stages_as_a_raw_four_byte_container() { + let mut file = write_safetensors( + "u32", + &[("e.weight_scale.nvfp4_csf_exceptions", "U32", &[81], 4)], + ); + let metas = parse_header(&mut file).unwrap(); + assert_eq!(metas[0].dtype, WeightDtype::Int32); + assert_eq!(metas[0].len, 324); + assert!(!metas[0].from_f16); +} + +/// FP4-Local's index names its shards under `tensors/`; each shard must still +/// receive exactly the tensors routed to it. +#[test] +fn shards_in_a_subdirectory_receive_their_tensors() { + let dir = Path::new("/models/fp4-local"); + let map: HashMap = [ + ("a.weight", "tensors/model-00001-of-00002.safetensors"), + ("b.weight", "tensors/model-00002-of-00002.safetensors"), + ("c.weight", "tensors/model-00001-of-00002.safetensors"), + ] + .into_iter() + .map(|(t, s)| (t.to_string(), s.to_string())) + .collect(); + let mut got = tensors_routed_to( + &map, + dir, + &dir.join("tensors/model-00001-of-00002.safetensors"), + ); + got.sort(); + assert_eq!(got, ["a.weight", "c.weight"]); + // A flat index still works, and a same-named file elsewhere matches nothing. + let flat: HashMap = [("x.weight".to_string(), "model.safetensors".to_string())] + .into_iter() + .collect(); + assert_eq!( + tensors_routed_to(&flat, dir, &dir.join("model.safetensors")), + ["x.weight"] + ); + assert!(tensors_routed_to(&map, dir, &dir.join("model-00001-of-00002.safetensors")).is_empty()); +} diff --git a/crates/spark-runtime/src/fast_weights/mod.rs b/crates/spark-runtime/src/fast_weights/mod.rs index 44b5fe7f3..731fa0969 100644 --- a/crates/spark-runtime/src/fast_weights/mod.rs +++ b/crates/spark-runtime/src/fast_weights/mod.rs @@ -28,7 +28,7 @@ mod direct_io; mod header; mod shard; -use header::resolve_shards; +use header::{resolve_shards, tensors_routed_to}; /// Run `f` over `len` bytes at `start` within a deferred tensor, read with /// `O_DIRECT` when the filesystem allows it (so the bytes never sit in the @@ -141,6 +141,12 @@ pub struct FastSafetensorsLoader { /// packed weights and block scales, which the model loader reads from /// disk as this rank's slice (see `should_defer_tensor`). pub expert_tp: bool, + /// Modules whose `.weight` and `.weight_scale` are stored in a format the + /// kernels do not read (MXFP8, or NVFP4 in a tower that runs BF16). They + /// are DEFERRED — located, not uploaded — and the model build expands + /// them on the host, so the originals never occupy device memory. CSF + /// scale streams are deferred regardless of this list. + pub host_decode_modules: std::collections::HashSet, } /// Default tensor-count cap for per-shard `O_DIRECT`. Above this, the fast @@ -176,6 +182,7 @@ impl FastSafetensorsLoader { prefetch_shards: false, demand_paged_patterns: Vec::new(), expert_tp: false, + host_decode_modules: std::collections::HashSet::new(), } } @@ -197,6 +204,7 @@ impl FastSafetensorsLoader { prefetch_shards: false, demand_paged_patterns: Vec::new(), expert_tp: false, + host_decode_modules: std::collections::HashSet::new(), } } } @@ -269,12 +277,9 @@ impl WeightLoader for FastSafetensorsLoader { .file_name() .and_then(|n| n.to_str()) .unwrap_or_default(); - let tensor_filter: Option> = tensor_to_shard.as_ref().map(|map| { - map.iter() - .filter(|(_, s)| *s == shard_name) - .map(|(t, _)| t.clone()) - .collect() - }); + let tensor_filter: Option> = tensor_to_shard + .as_ref() + .map(|map| tensors_routed_to(map, model_dir, shard_path)); tracing::info!( "Fast-loading shard {}/{}: {}{}", @@ -377,6 +382,32 @@ mod expert_filter_tests { assert!(loader.should_skip_tensor(target)); } + /// FP4-Local: CSF streams and listed MXFP8 modules are located, not + /// uploaded; their scalar siblings and every other tensor load as usual. + #[test] + fn host_decoded_tensors_are_deferred() { + let mut loader = FastSafetensorsLoader::new(); + let attn = "model.language_model.layers.3.self_attn.q_proj"; + loader.host_decode_modules.insert(attn.to_string()); + let expert = "model.language_model.layers.0.mlp.experts.3.down_proj"; + for name in [ + format!("{attn}.weight"), + format!("{attn}.weight_scale"), + format!("{expert}.weight_scale.nvfp4_csf_fixed"), + format!("{expert}.weight_scale.nvfp4_csf_exceptions"), + ] { + assert!(loader.should_defer_tensor(&name), "{name}"); + } + for name in [ + format!("{expert}.weight"), + format!("{expert}.weight_scale_2"), + format!("{attn}.weight_scale_2"), + "model.language_model.layers.3.self_attn.k_proj.weight".to_string(), + ] { + assert!(!loader.should_defer_tensor(&name), "{name}"); + } + } + #[test] fn expert_tp_keeps_every_expert_and_defers_only_its_bulk() { let mut loader = FastSafetensorsLoader::with_ep(1, 2, 288); diff --git a/crates/spark-runtime/src/fast_weights/skip.rs b/crates/spark-runtime/src/fast_weights/skip.rs index a73161197..a2a37a108 100644 --- a/crates/spark-runtime/src/fast_weights/skip.rs +++ b/crates/spark-runtime/src/fast_weights/skip.rs @@ -24,10 +24,28 @@ use super::FastSafetensorsLoader; use crate::weights::parse_expert_index; impl FastSafetensorsLoader { + /// Tensors located on disk but not uploaded: host-decoded formats and + /// expert-TP slices. Deferred tensors are also left out of the pre-flight. + pub(super) fn should_defer_tensor(&self, name: &str) -> bool { + self.is_host_decoded(name) || self.is_expert_tp_slice(name) + } + + /// CSF scale streams, and the weight/scale pair of every module the model + /// build expands on the host (`host_decode_modules`). + fn is_host_decoded(&self, name: &str) -> bool { + if crate::weights::is_csf_stream(name) { + return true; + } + let module = name + .strip_suffix(".weight") + .or_else(|| name.strip_suffix(".weight_scale")); + module.is_some_and(|m| self.host_decode_modules.contains(m)) + } + /// Expert TP: a routed expert's packed weight or block scale, which the /// model loader slices straight from disk. The appended predictor layer /// keeps its own replication rules, and the scalars load as usual. - pub(super) fn should_defer_tensor(&self, name: &str) -> bool { + fn is_expert_tp_slice(&self, name: &str) -> bool { self.expert_tp && !name.starts_with("mtp.") && (name.ends_with(".weight") || name.ends_with(".weight_scale")) diff --git a/crates/spark-runtime/src/weights.rs b/crates/spark-runtime/src/weights.rs index 1b6db9a94..07ef150c5 100644 --- a/crates/spark-runtime/src/weights.rs +++ b/crates/spark-runtime/src/weights.rs @@ -87,6 +87,10 @@ impl WeightDtype { // (`keeps_raw_f16` → FP16), everything else converts to BF16. safetensors::Dtype::I16 => Ok(Self::Int16), safetensors::Dtype::I32 => Ok(Self::Int32), + // U32: raw 4-byte container. Only the CSF NVFP4 scale exception + // streams (`*.nvfp4_csf_exceptions`) use it, and the loader + // decodes those on the host; signedness is never interpreted. + safetensors::Dtype::U32 => Ok(Self::Int32), safetensors::Dtype::F32 => Ok(Self::FP32), safetensors::Dtype::U8 => Ok(Self::UInt8), // I8: raw 1-byte container for 4-bit-packed NVFP4 (DeepSeek-V4 MTP @@ -110,6 +114,8 @@ impl WeightDtype { "BF16" => Self::BF16, "I16" => Self::Int16, "I32" => Self::Int32, + // U32: raw 4-byte container, as on the disk path. + "U32" => Self::Int32, "U8" => Self::UInt8, // I8 is a 1-byte raw container (packed NVFP4); signedness is // irrelevant, treat as raw bytes exactly like the disk path. @@ -511,7 +517,7 @@ mod side_files; mod store_edit; pub(crate) use name_utils::keeps_raw_f16; pub(crate) use name_utils::split_trailing_index; -pub use name_utils::{is_ngram_table, parse_expert_index}; +pub use name_utils::{is_csf_stream, is_ngram_table, parse_expert_index}; pub(crate) use side_files::merge_exl3_side_files; #[cfg(test)] diff --git a/crates/spark-runtime/src/weights/name_utils.rs b/crates/spark-runtime/src/weights/name_utils.rs index e39f664a5..7a051f5fa 100644 --- a/crates/spark-runtime/src/weights/name_utils.rs +++ b/crates/spark-runtime/src/weights/name_utils.rs @@ -45,17 +45,31 @@ pub(crate) fn keeps_raw_f16(name: &str) -> bool { /// filename rule would work for one release and silently fail the other, /// where "silently" means a 221 GB OOM pre-flight on a 121 GB box. /// +/// NVIDIA's FP4-Local pack stores the table as NVFP4, so each shard also has +/// a per-row block-scale tensor `shard_{i}.weight_scale` (E4M3 `[R, 10]`). +/// Those are table data read by row alongside the weight, so they match too. +/// /// What must NOT match, because the loader needs these resident: /// - LongCat's small `post_projs` /// - Qwen's `ngram_embedding.weight_scale` (one BF16 scalar) /// - Qwen's `ngram_heads_offsets` / `ngram_heads_vocab_sizes` (I64, 16 each) pub fn is_ngram_table(name: &str) -> bool { if name.contains("ngram_embeddings.embedders.") || name.contains("ngram_embedding.shard_") { - return name.ends_with(".weight") || exl3_trellis_tail(name); + return name.ends_with(".weight") + || (name.contains("ngram_embedding.shard_") && name.ends_with(".weight_scale")) + || exl3_trellis_tail(name); } name.ends_with(".ngram_embedding.trellis") } +/// A CSF-compressed NVFP4 block-scale stream (`weight_scale_encoding: "csf"`). +/// +/// Never uploaded: the loader defers both halves and the model build decodes +/// them into a plain E4M3 `weight_scale` (see `atlas_core::modelopt_codec`). +pub fn is_csf_stream(name: &str) -> bool { + name.ends_with(".nvfp4_csf_fixed") || name.ends_with(".nvfp4_csf_exceptions") +} + /// True when the tail after the last `ngram_embedding.shard_` is a /// bare `.trellis` (EXL3's sharded n-gram layout) — i.e. the segment(s) after /// `shard_` are exactly `trellis`. @@ -97,6 +111,11 @@ mod ngram_defer_tests { let base = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"; assert!(is_ngram_table(&format!("{base}.shard_0.weight"))); assert!(is_ngram_table(&format!("{base}.shard_127.weight"))); + // FP4-Local: each NVFP4 shard carries its own per-row block scales. + assert!(is_ngram_table(&format!("{base}.shard_0.weight_scale"))); + assert!(is_ngram_table(&format!("{base}.shard_127.weight_scale"))); + // ...but its per-table global scale is one resident F32. + assert!(!is_ngram_table(&format!("{base}.weight_scale_2"))); // The per-table scalar scale (RadixArk's FP8 release) must stay // resident — the row cache needs it to dequantize. diff --git a/crates/spark-runtime/src/weights/store_edit.rs b/crates/spark-runtime/src/weights/store_edit.rs index 30d7de611..35ce4cbf6 100644 --- a/crates/spark-runtime/src/weights/store_edit.rs +++ b/crates/spark-runtime/src/weights/store_edit.rs @@ -53,6 +53,12 @@ impl WeightStore { Some(tensor) } + /// Take a deferred tensor's locator out of the store, once a build step + /// has consumed it (so later passes over `deferred_sorted` never see it). + pub fn undefer(&mut self, name: &str) -> Option { + self.deferred.remove(name) + } + /// True if this name's tensor memory was already freed by /// [`reclaim`](Self::reclaim) (its map entry may still exist). Call this /// before [`remove`](Self::remove) to decide whether the returned buffer must be freed. @@ -66,3 +72,35 @@ impl WeightStore { .unwrap_or(false) } } + +impl DeferredTensor { + /// On-disk byte length: element count times the element size. Packed + /// formats with no per-element size (`PackedQ2_0`) are never deferred. + pub fn byte_len(&self) -> Result { + let elem = self.dtype.byte_size(); + if elem == 0 { + bail!( + "deferred tensor dtype {:?} has no fixed element size", + self.dtype + ); + } + Ok(self.shape.iter().product::() * elem) + } + + /// Read the tensor's raw bytes from its shard with one positioned read. + pub fn read_bytes(&self) -> Result> { + use std::io::{Read, Seek, SeekFrom}; + let len = self.byte_len()?; + let mut file = std::fs::File::open(&self.path)?; + file.seek(SeekFrom::Start(self.offset))?; + let mut buf = vec![0u8; len]; + file.read_exact(&mut buf).map_err(|e| { + anyhow::anyhow!( + "reading {len} bytes at {} from {}: {e}", + self.offset, + self.path.display() + ) + })?; + Ok(buf) + } +} diff --git a/crates/spark-runtime/src/weights/store_edit_tests.rs b/crates/spark-runtime/src/weights/store_edit_tests.rs index f5a215de7..d51479285 100644 --- a/crates/spark-runtime/src/weights/store_edit_tests.rs +++ b/crates/spark-runtime/src/weights/store_edit_tests.rs @@ -134,3 +134,31 @@ fn remove_of_reclaimed_tensor_reports_it_and_release_does_not_double_free() { ); assert!(gpu.read_alloc(b_ptr).is_none(), "b was freed at release"); } + +#[test] +fn deferred_tensor_reads_its_bytes_and_undefers() { + let dir = std::env::temp_dir().join(format!("atlas-defer-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("shard.bin"); + // 5 junk header bytes, then a [2, 3] U32 tensor. + let mut bytes = vec![0xAAu8; 5]; + for v in 0u32..6 { + bytes.extend_from_slice(&v.to_le_bytes()); + } + std::fs::write(&path, &bytes).unwrap(); + let d = DeferredTensor { + path: path.clone(), + offset: 5, + shape: vec![2, 3], + dtype: WeightDtype::Int32, + }; + assert_eq!(d.byte_len().unwrap(), 24); + assert_eq!(d.read_bytes().unwrap(), bytes[5..]); + + let mut store = WeightStore::empty(); + store.defer("x.nvfp4_csf_exceptions".into(), d); + assert!(store.undefer("x.nvfp4_csf_exceptions").is_some()); + assert!(store.deferred("x.nvfp4_csf_exceptions").is_none()); + assert!(store.undefer("x.nvfp4_csf_exceptions").is_none()); + std::fs::remove_dir_all(&dir).ok(); +} From 560642d893f8c31f1dea77eca0b0a2946ee77152 Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:56:48 -0400 Subject: [PATCH 03/26] qwen4_exp: serve nvidia/Qwen3.8-Flash-Next-FP4-Local The pack differs from the NVFP4 one only in how it stores tensors: - weight_map::modelopt_local decodes the 73,728 CSF expert scales into plain E4M3, expands 438 MXFP8 projections and 56 vision NVFP4 projections to BF16, and rounds the 100 F32 norms (hc_norm, PLE norms) to the BF16 the kernels read. Left F32, those norms produced NaN logits: every request answered with one empty token. - The PLE table is NVFP4: a second row cache serves the block-scale rows and the gather decodes both on the host, the trellis path's contract. Verified on a GB10 (sparky): boots in 65 s, coherent chat, correct tool call, needle passes at 4.4k prompt tokens, 16.9 tok/s decode without MTP (the NVFP4 pack measures 16.14 without it). MTP experts are per-expert NVFP4 here and are not wired yet, so --speculative is not supported. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01L4aUGgtiZX94pLr63HXQzn --- crates/spark-model/src/factory/build.rs | 3 + crates/spark-model/src/layers/ple.rs | 2 + crates/spark-model/src/layers/ple/batch.rs | 6 +- .../spark-model/src/layers/ple/batch_tests.rs | 1 + crates/spark-model/src/layers/ple/layer.rs | 16 +- .../spark-model/src/layers/ple/ngram_nvfp4.rs | 117 +++++++ .../src/layers/ple/ngram_nvfp4_tests.rs | 42 +++ crates/spark-model/src/layers/ple/verify.rs | 24 +- .../spark-model/src/layers/qwen3_ssm/tests.rs | 1 + .../src/weight_loader/qwen4_exp/ple.rs | 91 ++--- .../src/weight_loader/qwen4_exp/ple_host.rs | 100 ++++++ .../src/weight_loader/qwen4_exp/ple_nvfp4.rs | 98 ++++++ crates/spark-model/src/weight_map.rs | 2 + .../src/weight_map/modelopt_local.rs | 317 ++++++++++++++++++ .../src/weight_map/modelopt_local_tests.rs | 312 +++++++++++++++++ .../src/main_modules/serve_phases/weights.rs | 8 + 16 files changed, 1061 insertions(+), 79 deletions(-) create mode 100644 crates/spark-model/src/layers/ple/ngram_nvfp4.rs create mode 100644 crates/spark-model/src/layers/ple/ngram_nvfp4_tests.rs create mode 100644 crates/spark-model/src/weight_loader/qwen4_exp/ple_host.rs create mode 100644 crates/spark-model/src/weight_loader/qwen4_exp/ple_nvfp4.rs create mode 100644 crates/spark-model/src/weight_map/modelopt_local.rs create mode 100644 crates/spark-model/src/weight_map/modelopt_local_tests.rs diff --git a/crates/spark-model/src/factory/build.rs b/crates/spark-model/src/factory/build.rs index 0f59aace8..d16dd3117 100644 --- a/crates/spark-model/src/factory/build.rs +++ b/crates/spark-model/src/factory/build.rs @@ -156,6 +156,9 @@ pub fn build_model( crate::weight_map::exl3::exl3_materialize_dense(&mut store, gpu.as_ref())?; config.attn_proj_bf16 = true; } + // FP4-Local: CSF expert scales and MXFP8 / vision-NVFP4 projections become + // the plain NVFP4 / BF16 layout the loaders below read. No-op otherwise. + crate::weight_map::modelopt_local::materialize(&mut store, &config, gpu.as_ref())?; // ── LoRA adapter load (pre-arena, pre-KV-sizing) ── // MUST run before `BufferArena::new` and the `gpu.free_memory()` diff --git a/crates/spark-model/src/layers/ple.rs b/crates/spark-model/src/layers/ple.rs index f4f1230ea..6478c68ce 100644 --- a/crates/spark-model/src/layers/ple.rs +++ b/crates/spark-model/src/layers/ple.rs @@ -72,4 +72,6 @@ pub(crate) use layer::verify_snapshots_enabled; // The weight loader builds an EXL3 trellis spec. Non-CUDA builds do not // construct PLE layers, so the import is unused there. #[cfg_attr(not(feature = "cuda"), allow(unused_imports))] +pub(crate) use layer::ngram_nvfp4; +#[cfg_attr(not(feature = "cuda"), allow(unused_imports))] pub(crate) use layer::ngram_trellis; diff --git a/crates/spark-model/src/layers/ple/batch.rs b/crates/spark-model/src/layers/ple/batch.rs index 3d4abed4f..caa89e01a 100644 --- a/crates/spark-model/src/layers/ple/batch.rs +++ b/crates/spark-model/src/layers/ple/batch.rs @@ -190,8 +190,8 @@ impl PleLayer { /// /// Falls back to exactly that loop where the batched form does not /// apply: a consumable decode prestage on any member (the per-sequence - /// forward decides whether to use it), a trellis table (its gather - /// writes `emb` from the host per call), more rows than one scratch span, + /// forward decides whether to use it), a host-decoded table (trellis or + /// NVFP4: its gather writes `emb` from the host per call), more rows than one scratch span, /// a member wide enough to skip the per-row conv snapshots, or the kill /// switch. pub fn forward_seqs( @@ -212,7 +212,7 @@ impl PleLayer { let batched = seq_batch_enabled() && seqs.len() > 1 && contiguous - && self.trellis.is_none() + && !self.host_decoded() && total <= self.scratch_tokens && total <= self.max_tokens && verify_snapshots_enabled() diff --git a/crates/spark-model/src/layers/ple/batch_tests.rs b/crates/spark-model/src/layers/ple/batch_tests.rs index 52bb1a433..07835995a 100644 --- a/crates/spark-model/src/layers/ple/batch_tests.rs +++ b/crates/spark-model/src/layers/ple/batch_tests.rs @@ -42,6 +42,7 @@ fn layer(gpu: &MockGpuBackend) -> PleLayer { 32, 32, None, + None, gpu, ) .unwrap() diff --git a/crates/spark-model/src/layers/ple/layer.rs b/crates/spark-model/src/layers/ple/layer.rs index 129c172c6..b2d4a1002 100644 --- a/crates/spark-model/src/layers/ple/layer.rs +++ b/crates/spark-model/src/layers/ple/layer.rs @@ -63,6 +63,8 @@ pub struct PleLayer { scale_va: Option, /// Packed EXL3 n-gram rows. Host-dequanted into `emb` before the GEMM. trellis: Option, + /// NVFP4 n-gram rows (FP4-Local). Host-dequanted into `emb` likewise. + nvfp4: Option, gemm_k: KernelHandle, gate_k: KernelHandle, conv_k: KernelHandle, @@ -96,6 +98,12 @@ pub struct PleLayer { } impl PleLayer { + /// Rows decoded on the host into `emb` (EXL3 trellis, FP4-Local NVFP4) + /// rather than gathered by the embedding kernel. + pub(super) fn host_decoded(&self) -> bool { + self.trellis.is_some() || self.nvfp4.is_some() + } + #[allow(clippy::too_many_arguments)] pub fn new( dims: PleIdDims, @@ -110,6 +118,7 @@ impl PleLayer { max_tokens: usize, scratch_tokens: usize, trellis: Option, + nvfp4: Option, gpu: &dyn GpuBackend, ) -> Result { dims.validate()?; @@ -132,7 +141,9 @@ impl PleLayer { if let NgramTable::Cached(cache) = &table { let stride = cache.row_stride(); let scaled = cache.scale_dev_va()?.is_some(); - if let Some(t) = trellis.as_ref() { + if let Some(n) = nvfp4.as_ref() { + ngram_nvfp4::check_packed_stride(stride, &n.geom, scaled)?; + } else if let Some(t) = trellis.as_ref() { ngram_trellis::check_packed_stride(stride, t, scaled)?; } else { gather_matches_element_size(stride, head_dim, scaled)?; @@ -162,6 +173,7 @@ impl PleLayer { _ => None, }, trellis, + nvfp4, embed_fp8_k: gpu.kernel("embed_from_argmax", "batched_embed_fp8")?, table: std::sync::Arc::new(std::sync::Mutex::new(table)), embed_k: gpu.kernel("embed_from_argmax", "batched_embed")?, @@ -396,6 +408,8 @@ pub(crate) mod warm; #[path = "gather_guard.rs"] mod gather_guard; +#[path = "ngram_nvfp4.rs"] +pub(crate) mod ngram_nvfp4; #[path = "ngram_trellis.rs"] pub(crate) mod ngram_trellis; #[cfg(feature = "cuda")] diff --git a/crates/spark-model/src/layers/ple/ngram_nvfp4.rs b/crates/spark-model/src/layers/ple/ngram_nvfp4.rs new file mode 100644 index 000000000..a47ca19be --- /dev/null +++ b/crates/spark-model/src/layers/ple/ngram_nvfp4.rs @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Host dequant of NVFP4 n-gram rows (NVIDIA's Qwen3.8-Flash-Next-FP4-Local). +//! +//! That pack stores the PLE table as NVFP4: each shard is a U8 `[R, dim/2]` +//! weight (two E2M1 values per byte, low nibble first) with an E4M3 +//! `[R, dim/16]` block-scale tensor beside it, and the table carries one F32 +//! `weight_scale_2`. The main row cache holds the packed weight rows; a +//! second cache, segmented identically, holds the block-scale rows. A gather +//! resolves both, decodes on the host to the BF16 `[n, dim]` embedding the +//! PLE GEMM expects, and uploads it into `emb` — the same contract as the +//! EXL3 trellis path, so `gather_embed` skips its kernel for both. +//! +//! 26.8 GiB on disk against the NVFP4 pack's 47.7 GiB FP8 table, with the +//! same row count: fewer bytes per fault, more rows per page-cache GiB. + +use anyhow::{Context, Result, ensure}; +use atlas_core::modelopt_codec::nvfp4_rows_to_bf16; + +/// Row geometry of an NVFP4 n-gram table. +#[derive(Debug, Clone, Copy)] +pub struct Nvfp4RowGeom { + /// Decoded row width (160 for Flash-Next). + pub dim: usize, + /// Elements per E4M3 block scale (16). + pub group: usize, + /// The table's `weight_scale_2`. + pub global: f32, +} + +/// An NVFP4 n-gram table's geometry plus its block-scale row cache. +pub struct NgramNvfp4 { + pub geom: Nvfp4RowGeom, + /// Block-scale rows, `dim / group` bytes each, at the same row ids as the + /// weight cache. + #[cfg(feature = "cuda")] + pub scales: std::sync::Mutex, +} + +impl Nvfp4RowGeom { + pub fn packed_stride(&self) -> usize { + self.dim / 2 + } + + pub fn scale_stride(&self) -> usize { + self.dim / self.group + } +} + +/// The weight cache must hold packed rows and nothing else: a scale array +/// would select the FP8 gather, and a wrong stride reads the next row's bytes. +pub fn check_packed_stride(stride: usize, spec: &Nvfp4RowGeom, scaled: bool) -> Result<()> { + ensure!( + !scaled, + "PLE: NVFP4 n-gram rows take their scales from a block-scale table, \ + not a cache scale array" + ); + ensure!( + stride == spec.packed_stride(), + "PLE: NVFP4 row stride {stride} != {} packed bytes for {} dims", + spec.packed_stride(), + spec.dim + ); + Ok(()) +} + +/// Decode gathered rows: `packed[i]` and `scales[i]` are row `i`'s bytes from +/// the weight and block-scale caches. Returns little-endian BF16 `[n, dim]`. +pub fn decode_rows_bf16( + spec: &Nvfp4RowGeom, + packed: &[Vec], + scales: &[Vec], +) -> Result> { + ensure!( + packed.len() == scales.len(), + "PLE nvfp4: {} weight rows but {} scale rows", + packed.len(), + scales.len() + ); + let (pw, sw) = (spec.packed_stride(), spec.scale_stride()); + let row_bytes = spec.dim * 2; + let mut out = vec![0u8; packed.len() * row_bytes]; + for (i, (p, s)) in packed.iter().zip(scales).enumerate() { + ensure!( + p.len() == pw && s.len() == sw, + "PLE nvfp4: row {i} is {}+{} bytes, expected {pw}+{sw}", + p.len(), + s.len() + ); + let dst = &mut out[i * row_bytes..(i + 1) * row_bytes]; + nvfp4_rows_to_bf16(p, s, spec.global, spec.dim, spec.group, dst) + .with_context(|| format!("PLE nvfp4: row {i}"))?; + } + Ok(out) +} + +/// Resolve `ids` in the block-scale cache, copy each row out, and release the +/// pins at once: the decode that consumes them runs synchronously on the host. +#[cfg(feature = "cuda")] +pub fn gather_scales(spec: &NgramNvfp4, ids: &[u64]) -> Result>> { + let mut cache = spec + .scales + .lock() + .map_err(|_| anyhow::anyhow!("PLE nvfp4 scale cache mutex poisoned"))?; + let mut slots = Vec::with_capacity(ids.len()); + cache.resolve(ids, &mut slots)?; + let rows = slots + .iter() + .map(|&s| cache.copy_slot(s)) + .collect::>>(); + cache.end_batch(); + rows +} + +#[cfg(test)] +#[path = "ngram_nvfp4_tests.rs"] +mod tests; diff --git a/crates/spark-model/src/layers/ple/ngram_nvfp4_tests.rs b/crates/spark-model/src/layers/ple/ngram_nvfp4_tests.rs new file mode 100644 index 000000000..5e34acb4c --- /dev/null +++ b/crates/spark-model/src/layers/ple/ngram_nvfp4_tests.rs @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +use super::*; + +fn spec() -> Nvfp4RowGeom { + Nvfp4RowGeom { + dim: 32, + group: 16, + global: 0.25, + } +} + +#[test] +fn decodes_rows_with_block_and_global_scales() { + let s = spec(); + // Row 0: nibble 2 (=1.0) in block 0, nibble 7 (=6.0) in block 1. + // Row 1: nibble 9 (=-0.5) everywhere. + let packed = vec![ + [vec![0x22u8; 8], vec![0x77u8; 8]].concat(), + vec![0x99u8; 16], + ]; + // Block scales: E4M3 0x38 = 1.0, 0x40 = 2.0. + let scales = vec![vec![0x38u8, 0x40], vec![0x40u8, 0x40]]; + let out = decode_rows_bf16(&s, &packed, &scales).unwrap(); + let at = |i: usize| { + f32::from_bits(u32::from(u16::from_le_bytes([out[2 * i], out[2 * i + 1]])) << 16) + }; + assert_eq!(out.len(), 2 * 32 * 2); + assert!((0..16).all(|i| at(i) == 0.25)); // 1.0 * 1.0 * 0.25 + assert!((16..32).all(|i| at(i) == 3.0)); // 6.0 * 2.0 * 0.25 + assert!((32..64).all(|i| at(i) == -0.25)); // -0.5 * 2.0 * 0.25 +} + +#[test] +fn rejects_mismatched_rows_and_strides() { + let s = spec(); + assert!(decode_rows_bf16(&s, &[vec![0; 16]], &[]).is_err()); + assert!(decode_rows_bf16(&s, &[vec![0; 15]], &[vec![0; 2]]).is_err()); + assert!(check_packed_stride(16, &s, false).is_ok()); + assert!(check_packed_stride(32, &s, false).is_err()); + assert!(check_packed_stride(16, &s, true).is_err()); +} diff --git a/crates/spark-model/src/layers/ple/verify.rs b/crates/spark-model/src/layers/ple/verify.rs index 1707a2f91..7a3c45ac9 100644 --- a/crates/spark-model/src/layers/ple/verify.rs +++ b/crates/spark-model/src/layers/ple/verify.rs @@ -139,7 +139,21 @@ impl PleLayer { let (h0, m0, _) = cache.stats(); let t0 = std::time::Instant::now(); cache.resolve(ids, &mut slots)?; - if let Some(spec) = &self.trellis { + if let Some(spec) = &self.nvfp4 { + // FP4-Local: weight rows from this cache, block-scale rows + // from the companion cache at the same ids. + let mut raws = Vec::with_capacity(slots.len()); + for &slot in &slots { + raws.push(cache.copy_slot(slot)?); + } + let scales = super::ngram_nvfp4::gather_scales(spec, ids)?; + let decoded = super::ngram_nvfp4::decode_rows_bf16(&spec.geom, &raws, &scales) + .map_err(|e| anyhow::anyhow!("PLE: ngram nvfp4 dequant: {e:#}"))?; + gpu.copy_h2d_async(&decoded, self.emb, stream)?; + // Same contract as the trellis branch below: `decoded` is a + // temporary, and a graph replay must see the finished upload. + gpu.synchronize(stream)?; + } else if let Some(spec) = &self.trellis { let mut raws = Vec::with_capacity(slots.len()); for &slot in &slots { raws.push(cache.copy_slot(slot)?); @@ -228,10 +242,10 @@ impl PleLayer { gpu: &dyn GpuBackend, stream: u64, ) -> Result<()> { - // Trellis rows are not embeddings. `gather_host` already wrote the - // decoded BF16 matrix into `emb`; launching `batched_embed` would - // overwrite it with packed I16 interpreted as BF16. - if self.trellis.is_none() { + // Trellis and NVFP4 rows are not embeddings. `gather_host` already + // wrote the decoded BF16 matrix into `emb`; launching `batched_embed` + // would overwrite it with packed bytes interpreted as BF16. + if !self.host_decoded() { self.gather_embed_dispatch(table_va, num_tokens, heads, gpu, stream)?; } // The pins taken by this gather must outlive THIS kernel: record an event diff --git a/crates/spark-model/src/layers/qwen3_ssm/tests.rs b/crates/spark-model/src/layers/qwen3_ssm/tests.rs index eae90fb54..bc9d009d1 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/tests.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/tests.rs @@ -347,6 +347,7 @@ fn active_row_without_ple_carry_is_refused() { 1, 1, None, + None, &gpu, ) .unwrap(); diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs b/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs index 6a189d893..7ea179c21 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs @@ -43,6 +43,8 @@ use crate::layers::ple::PleLayer; use crate::layers::ple::{PleIdDims, PleWeights}; #[cfg(feature = "cuda")] use crate::weight_map::dense; +#[cfg(feature = "cuda")] +use host::{bf16_scalar, fp16_host, i64_host}; /// Resident rows in the pinned arena. A forward's gather pins at most one /// scratch span of rows at once — `scratch_tokens * ngram_heads` — because @@ -110,75 +112,6 @@ fn ple_i64_name(store: &WeightStore, lp: &str, exl3: &str, nvfp4: &str) -> Strin .to_string() } -/// Read a small I64 device tensor back to the host. -/// -/// `layer_multipliers` and the two per-head tables are 3 and 16 elements — -/// they are uploaded like any other weight, and the id hash needs them on the -/// host. Reading them back beats adding a host-side path to `WeightStore` for -/// 280 bytes. -#[cfg(feature = "cuda")] -fn i64_host(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result> { - let t = store.get(name).with_context(|| format!("PLE: {name}"))?; - let n = t.num_elements(); - let mut raw = vec![0u8; n * 8]; - gpu.copy_d2h(t.ptr, &mut raw) - .with_context(|| format!("PLE: reading {name} back to host"))?; - Ok(raw - .chunks_exact(8) - .map(|b| u64::from_le_bytes(b.try_into().unwrap())) - .collect()) -} - -/// A one-element BF16 tensor, read back to the host as f32. -/// -/// The FP8 n-gram table's dequant scale is stored this way — `shape: [1]`, -/// dtype BF16 — rather than as a per-row scale file. -#[cfg(feature = "cuda")] -fn bf16_scalar(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result { - let t = store.get(name).with_context(|| format!("PLE: {name}"))?; - anyhow::ensure!( - t.num_elements() == 1, - "PLE: {name} has {} elements, expected 1 (a per-tensor scale)", - t.num_elements() - ); - let mut raw = [0u8; 2]; - gpu.copy_d2h(t.ptr, &mut raw) - .with_context(|| format!("PLE: reading {name} back to host"))?; - let bits = u16::from_le_bytes(raw); - let v = f32::from_bits((bits as u32) << 16); - anyhow::ensure!( - v.is_finite() && v > 0.0, - "PLE: {name} is {v}, which cannot be a dequant scale" - ); - Ok(v) -} - -/// A resident FP16 vector, read back as raw fp16 bits. -/// -/// EXL3 `head_bias` stays FP16 (`keeps_raw_f16`) because the n-gram dequant -/// adds it in the same units as `decode_mul1` * scale. -#[cfg(feature = "cuda")] -fn fp16_host(store: &WeightStore, name: &str, n: usize, gpu: &dyn GpuBackend) -> Result> { - let t = store.get(name).with_context(|| format!("PLE: {name}"))?; - anyhow::ensure!( - t.num_elements() == n, - "PLE: {name} has {} elements, expected {n}", - t.num_elements() - ); - anyhow::ensure!( - matches!(t.dtype, spark_runtime::weights::WeightDtype::FP16), - "PLE: {name} is {:?}, expected FP16 (raw head_bias)", - t.dtype - ); - let mut raw = vec![0u8; n * 2]; - gpu.copy_d2h(t.ptr, &mut raw) - .with_context(|| format!("PLE: reading {name} back to host"))?; - Ok(raw - .chunks_exact(2) - .map(|b| u16::from_le_bytes([b[0], b[1]])) - .collect()) -} - /// Build the PLE layer for `layer_idx`, or `None` if this model has none. #[cfg(feature = "cuda")] pub(super) fn load( @@ -323,6 +256,9 @@ pub(super) fn load( // cache copies those bytes; host dequant expands each row to 160 // BF16 before the GEMM. See `layers/ple/ngram_trellis.rs`. spark_runtime::weights::WeightDtype::Int16 => 2, + // FP4-Local NVFP4: U8 rows of packed E2M1, host-decoded with their + // block-scale siblings (`ple_nvfp4`). + spark_runtime::weights::WeightDtype::UInt8 => 1, other => anyhow::bail!( "PLE: n-gram table is {other:?}; the row cache gathers BF16, \ F8_E4M3, or I16 rows (`batched_embed` / `batched_embed_fp8`)" @@ -364,10 +300,13 @@ pub(super) fn load( ); } + let n_shards = shards.len(); + let nvfp4 = nvfp4::open_if_packed(store, &lp, dtype, n_shards, rows_per, head_dim, slots, gpu)?; + // FP8 rows need their dequant scale, or the gather returns raw E4M3 // magnitudes and the whole n-gram contribution is off by a constant // factor — fluent output, wrong logits. - if elem == 1 { + if matches!(dtype, spark_runtime::weights::WeightDtype::FP8E4M3) { let name = format!("{lp}.ple_embedding.ngram_embedding.weight_scale"); let scale = bf16_scalar(store, &name, gpu)?; cache @@ -428,6 +367,8 @@ pub(super) fn load( head_vocab_sizes: dims.head_vocab_sizes.clone(), }), ) + } else if let Some(n) = &nvfp4 { + (n.geom.dim, None) } else { (head_dim, None) }; @@ -445,6 +386,7 @@ pub(super) fn load( max_tokens, span, trellis, + nvfp4, gpu, ) .map(Some) @@ -490,3 +432,12 @@ mod slots_tests { assert!(derived_slots(20481, 16) >= 20481 * 16); } } + +/// Host readback of the small resident PLE tensors. +#[cfg(feature = "cuda")] +#[path = "ple_host.rs"] +mod host; +/// FP4-Local's NVFP4 n-gram table: block-scale cache and global scale. +#[cfg(feature = "cuda")] +#[path = "ple_nvfp4.rs"] +mod nvfp4; diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/ple_host.rs b/crates/spark-model/src/weight_loader/qwen4_exp/ple_host.rs new file mode 100644 index 000000000..d9820ee27 --- /dev/null +++ b/crates/spark-model/src/weight_loader/qwen4_exp/ple_host.rs @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Small resident PLE tensors read back to the host: the id-hash tables, the +//! per-tensor dequant scalars, and EXL3's FP16 head bias. + +use anyhow::{Context, Result, ensure}; +use spark_runtime::gpu::GpuBackend; +use spark_runtime::weights::{WeightDtype, WeightStore}; + +/// Read a small I64 device tensor back to the host. +/// +/// `layer_multipliers` and the two per-head tables are 3 and 16 elements — +/// they are uploaded like any other weight, and the id hash needs them on the +/// host. Reading them back beats adding a host-side path to `WeightStore` for +/// 280 bytes. +pub(super) fn i64_host(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result> { + let t = store.get(name).with_context(|| format!("PLE: {name}"))?; + let n = t.num_elements(); + let mut raw = vec![0u8; n * 8]; + gpu.copy_d2h(t.ptr, &mut raw) + .with_context(|| format!("PLE: reading {name} back to host"))?; + Ok(raw + .chunks_exact(8) + .map(|b| u64::from_le_bytes(b.try_into().unwrap())) + .collect()) +} + +/// A one-element BF16 tensor, read back to the host as f32. +/// +/// The FP8 n-gram table's dequant scale is stored this way — `shape: [1]`, +/// dtype BF16 — rather than as a per-row scale file. +pub(super) fn bf16_scalar(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result { + let t = store.get(name).with_context(|| format!("PLE: {name}"))?; + anyhow::ensure!( + t.num_elements() == 1, + "PLE: {name} has {} elements, expected 1 (a per-tensor scale)", + t.num_elements() + ); + let mut raw = [0u8; 2]; + gpu.copy_d2h(t.ptr, &mut raw) + .with_context(|| format!("PLE: reading {name} back to host"))?; + let bits = u16::from_le_bytes(raw); + let v = f32::from_bits((bits as u32) << 16); + anyhow::ensure!( + v.is_finite() && v > 0.0, + "PLE: {name} is {v}, which cannot be a dequant scale" + ); + Ok(v) +} + +/// A resident FP16 vector, read back as raw fp16 bits. +/// +/// EXL3 `head_bias` stays FP16 (`keeps_raw_f16`) because the n-gram dequant +/// adds it in the same units as `decode_mul1` * scale. +pub(super) fn fp16_host( + store: &WeightStore, + name: &str, + n: usize, + gpu: &dyn GpuBackend, +) -> Result> { + let t = store.get(name).with_context(|| format!("PLE: {name}"))?; + anyhow::ensure!( + t.num_elements() == n, + "PLE: {name} has {} elements, expected {n}", + t.num_elements() + ); + anyhow::ensure!( + matches!(t.dtype, spark_runtime::weights::WeightDtype::FP16), + "PLE: {name} is {:?}, expected FP16 (raw head_bias)", + t.dtype + ); + let mut raw = vec![0u8; n * 2]; + gpu.copy_d2h(t.ptr, &mut raw) + .with_context(|| format!("PLE: reading {name} back to host"))?; + Ok(raw + .chunks_exact(2) + .map(|b| u16::from_le_bytes([b[0], b[1]])) + .collect()) +} + +/// A resident one-element F32 tensor (`weight_scale_2`), read to the host. +pub(super) fn f32_scalar(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result { + let t = store + .get(name) + .with_context(|| format!("PLE nvfp4: {name}"))?; + ensure!( + t.dtype == WeightDtype::FP32 && t.num_elements() == 1, + "PLE nvfp4: {name} is {:?} {:?}, expected one F32", + t.dtype, + t.shape + ); + let mut raw = [0u8; 4]; + gpu.copy_d2h(t.ptr, &mut raw)?; + let v = f32::from_le_bytes(raw); + ensure!( + v.is_finite() && v > 0.0, + "PLE nvfp4: {name} is {v}, not a dequant scale" + ); + Ok(v) +} diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/ple_nvfp4.rs b/crates/spark-model/src/weight_loader/qwen4_exp/ple_nvfp4.rs new file mode 100644 index 000000000..0a55e6f00 --- /dev/null +++ b/crates/spark-model/src/weight_loader/qwen4_exp/ple_nvfp4.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! The NVFP4 n-gram table of nvidia/Qwen3.8-Flash-Next-FP4-Local. +//! +//! Its 128 shards are U8 `[R, 80]` (160 E2M1 values per row) with a sibling +//! E4M3 `shard_{i}.weight_scale` `[R, 10]`, and the table has one F32 +//! `weight_scale_2`. The packed rows go through the regular row cache; this +//! module opens the identically segmented block-scale cache and reads the +//! global scale, which together are what `layers::ple::ngram_nvfp4` decodes. + +use anyhow::{Context, Result, ensure}; +use spark_runtime::gpu::GpuBackend; +use spark_runtime::weights::{WeightDtype, WeightStore}; + +use super::host::f32_scalar; +use crate::layers::ple::ngram_nvfp4::{NgramNvfp4, Nvfp4RowGeom}; + +/// The NVFP4 decoder when the table's shards are packed U8, `None` for every +/// other element type. `rows_per` and `packed_cols` come from shard 0's +/// weight; `shards` is how many weight shards were found. +#[allow(clippy::too_many_arguments)] +pub(super) fn open_if_packed( + store: &WeightStore, + lp: &str, + dtype: WeightDtype, + shards: usize, + rows_per: usize, + packed_cols: usize, + slots: usize, + gpu: &dyn GpuBackend, +) -> Result> { + if dtype != WeightDtype::UInt8 { + return Ok(None); + } + open(store, lp, shards, rows_per, packed_cols, slots, gpu).map(Some) +} + +fn open( + store: &WeightStore, + lp: &str, + shards: usize, + rows_per: usize, + packed_cols: usize, + slots: usize, + gpu: &dyn GpuBackend, +) -> Result { + let table = format!("{lp}.ple_embedding.ngram_embedding"); + let mut scale_shards = Vec::with_capacity(shards); + let mut scale_cols = 0; + for i in 0..shards { + let name = format!("{table}.shard_{i}.weight_scale"); + let d = store.deferred(&name).with_context(|| { + format!("PLE nvfp4: shard {i} has packed rows but `{name}` was not deferred") + })?; + ensure!( + d.dtype == WeightDtype::FP8E4M3 && d.shape.len() == 2 && d.shape[0] == rows_per, + "PLE nvfp4: `{name}` is {:?} {:?}, expected E4M3 [{rows_per}, _]", + d.dtype, + d.shape + ); + if i == 0 { + scale_cols = d.shape[1]; + } + ensure!( + d.shape[1] == scale_cols, + "PLE nvfp4: `{name}` width differs from shard 0" + ); + scale_shards.push((d.path.clone(), d.offset)); + } + let dim = packed_cols * 2; + ensure!( + scale_cols > 0 && dim.is_multiple_of(scale_cols), + "PLE nvfp4: {dim} dims do not split into {scale_cols} block scales" + ); + let geom = Nvfp4RowGeom { + dim, + group: dim / scale_cols, + global: f32_scalar(store, &format!("{table}.weight_scale_2"), gpu)?, + }; + let scales = spark_storage::NgramRowCache::open_segmented( + &scale_shards, + rows_per as u64, + None, + scale_cols, + slots, + ) + .context("PLE nvfp4: block-scale row cache")?; + tracing::info!( + "PLE n-gram table: NVFP4 ({dim} dims, group {}, global {:.6e}), block scales \ + served off NVMe from {shards} sibling shards", + geom.group, + geom.global + ); + Ok(NgramNvfp4 { + geom, + scales: std::sync::Mutex::new(scales), + }) +} diff --git a/crates/spark-model/src/weight_map.rs b/crates/spark-model/src/weight_map.rs index 31fdfa54c..53293ad2b 100644 --- a/crates/spark-model/src/weight_map.rs +++ b/crates/spark-model/src/weight_map.rs @@ -25,6 +25,8 @@ mod loaders_mtp; mod model_a; #[path = "weight_map/model_b.rs"] mod model_b; +#[path = "weight_map/modelopt_local.rs"] +pub mod modelopt_local; #[path = "weight_map/moe.rs"] mod moe; #[path = "weight_map/mxfp8.rs"] diff --git a/crates/spark-model/src/weight_map/modelopt_local.rs b/crates/spark-model/src/weight_map/modelopt_local.rs new file mode 100644 index 000000000..6595aa54e --- /dev/null +++ b/crates/spark-model/src/weight_map/modelopt_local.rs @@ -0,0 +1,317 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Load-time expansion of NVIDIA's ModelOpt "FP4-Local" storage formats into +//! the layouts Atlas already serves (nvidia/Qwen3.8-Flash-Next-FP4-Local). +//! +//! That pack differs from the NVFP4 pack Atlas runs only in HOW it stores +//! three things, never in what they compute: +//! +//! * routed-expert NVFP4 block scales are CSF-compressed +//! (`weight_scale.nvfp4_csf_{fixed,exceptions}`) — decoded losslessly into +//! the plain E4M3 `weight_scale` every NVFP4 kernel reads; +//! * GDN / attention / QSA-indexer / shared-expert projections are MXFP8 — +//! expanded exactly to the BF16 `weight` the NVFP4 pack ships; +//! * vision MLP/merger projections are W4A16 NVFP4 — dequantized to the BF16 +//! the vision tower reads. +//! +//! The loader DEFERS every source tensor (`FastSafetensorsLoader:: +//! host_decode_modules`, `is_csf_stream`), so the compressed originals are +//! read from disk here and never occupy device memory; only the expanded +//! tensor is uploaded. A loader that uploaded them anyway (the non-fast path) +//! still works: resident sources are read back and freed. +//! +//! The PLE n-gram table's NVFP4 rows are a different consumer (gathered by +//! row off NVMe) and are handled by the PLE loader, not here. + +use std::collections::HashSet; + +use anyhow::{Context, Result, ensure}; +use atlas_core::config::ModelConfig; +use atlas_core::modelopt_codec::{decode_csf_scales, mxfp8_to_bf16, nvfp4_rows_to_bf16}; +use spark_runtime::gpu::GpuBackend; +use spark_runtime::weights::{WeightDtype, WeightStore, WeightTensor}; + +const CSF_FIXED: &str = ".nvfp4_csf_fixed"; +const CSF_EXCEPTIONS: &str = ".nvfp4_csf_exceptions"; + +/// Modules whose weight/scale pair the loader must defer for host expansion: +/// every `MXFP8` entry of the ModelOpt `quantized_layers` map, and NVFP4 +/// entries of the vision tower (which runs BF16). Empty for any other pack. +pub fn host_decode_modules(config: &ModelConfig) -> HashSet { + let Some(qc) = config.quantization_config.as_ref() else { + return HashSet::new(); + }; + qc.quantized_layers + .iter() + .filter(|(path, spec)| { + let algo = spec.quant_algo.to_ascii_uppercase(); + algo == "MXFP8" + || (path.starts_with("model.visual.") + && atlas_core::config::is_nvfp4_quant_algo(&algo)) + }) + .map(|(path, _)| path.clone()) + .collect() +} + +/// What [`materialize`] did, for the load log and for tests. +#[derive(Debug, Default, PartialEq, Eq)] +pub struct Materialized { + pub csf_scales: usize, + pub mxfp8_to_bf16: usize, + pub nvfp4_to_bf16: usize, + pub f32_norms_to_bf16: usize, + pub uploaded_bytes: u64, +} + +/// A norm weight the FP4-Local pack keeps in F32 where every other release of +/// the model ships BF16: the 97 hyper-connection `hc_norm`s and the three PLE +/// norms (model card: "100 FP32 tensors"; 103 with the MTP block's). The +/// loaders hand these pointers straight to kernels that read BF16, so left +/// as F32 every normalization reads garbage — the server answers with one +/// empty token. Rounded to BF16 here, which is also what the pack's own vLLM +/// recipe (`--dtype bfloat16`) does; they sit within 0.07% of the NVFP4 +/// pack's BF16 values. +fn is_f32_norm(name: &str, t: &WeightTensor) -> bool { + // `hc_norm`, and the PLE's `norm_conv` / `norm_key` / `norm_query`: the + // last module segment names a norm. + let names_norm = name + .strip_suffix(".weight") + .and_then(|m| m.rsplit('.').next()) + .is_some_and(|seg| seg.contains("norm")); + t.dtype == WeightDtype::FP32 && t.shape.len() == 1 && names_norm +} + +fn norms_to_bf16(store: &mut WeightStore, gpu: &dyn GpuBackend) -> Result<(usize, u64)> { + let mut names: Vec = store + .names() + .filter(|n| store.get(n).is_ok_and(|t| is_f32_norm(n, t))) + .map(str::to_string) + .collect(); + names.sort(); + let mut bytes = 0; + for name in &names { + let src = take(store, gpu, name)?.context("norm vanished")?; + let bf16: Vec = src + .bytes + .chunks_exact(4) + .flat_map(|b| { + let v = f32::from_le_bytes([b[0], b[1], b[2], b[3]]); + atlas_core::numeric::f32_to_bf16(v).to_le_bytes() + }) + .collect(); + bytes += upload( + store, + gpu, + name.clone(), + &bf16, + src.shape, + WeightDtype::BF16, + )?; + } + Ok((names.len(), bytes)) +} + +/// Raw bytes and geometry of a source tensor, from disk or from the device. +struct Source { + bytes: Vec, + shape: Vec, + dtype: WeightDtype, +} + +/// Take `name` out of the store, preferring its deferred on-disk copy. A +/// resident copy is read back and its buffer freed (unless already reclaimed). +fn take(store: &mut WeightStore, gpu: &dyn GpuBackend, name: &str) -> Result> { + if let Some(d) = store.undefer(name) { + let bytes = d + .read_bytes() + .with_context(|| format!("reading deferred {name}"))?; + return Ok(Some(Source { + bytes, + shape: d.shape, + dtype: d.dtype, + })); + } + let reclaimed = store.was_reclaimed(name); + let Some(t) = store.remove(name) else { + return Ok(None); + }; + ensure!( + !reclaimed, + "{name} was reclaimed before it could be expanded" + ); + let mut bytes = vec![0u8; t.byte_size()]; + gpu.copy_d2h(t.ptr, &mut bytes)?; + gpu.free(t.ptr)?; + Ok(Some(Source { + bytes, + shape: t.shape, + dtype: t.dtype, + })) +} + +/// Read a resident F32 scalar (`weight_scale_2`), removing it from the store. +fn take_scalar(store: &mut WeightStore, gpu: &dyn GpuBackend, name: &str) -> Result { + let s = take(store, gpu, name)?.with_context(|| format!("{name} is missing"))?; + ensure!( + s.dtype == WeightDtype::FP32 && s.bytes.len() == 4, + "{name}: expected one F32, got {:?} {:?}", + s.dtype, + s.shape + ); + Ok(f32::from_le_bytes([ + s.bytes[0], s.bytes[1], s.bytes[2], s.bytes[3], + ])) +} + +fn upload( + store: &mut WeightStore, + gpu: &dyn GpuBackend, + name: String, + bytes: &[u8], + shape: Vec, + dtype: WeightDtype, +) -> Result { + let ptr = gpu.alloc(bytes.len())?; + gpu.copy_h2d(bytes, ptr)?; + store.insert(name, WeightTensor { ptr, shape, dtype })?; + Ok(bytes.len() as u64) +} + +/// Names of every CSF fixed stream the store knows, deferred or resident. +fn csf_bases(store: &WeightStore) -> Vec { + let mut bases: Vec = store + .deferred_sorted() + .into_iter() + .map(|(n, _)| n.as_str()) + .chain(store.names()) + .filter_map(|n| n.strip_suffix(CSF_FIXED).map(str::to_string)) + .collect(); + bases.sort(); + bases.dedup(); + bases +} + +fn decode_one_csf(store: &mut WeightStore, gpu: &dyn GpuBackend, base: &str) -> Result { + let fixed = take(store, gpu, &format!("{base}{CSF_FIXED}"))?.context("fixed stream")?; + let exc = take(store, gpu, &format!("{base}{CSF_EXCEPTIONS}"))? + .with_context(|| format!("{base}: CSF fixed stream without its exceptions partner"))?; + ensure!( + fixed.dtype == WeightDtype::UInt8 && fixed.shape.len() == 2, + "{base}: CSF fixed stream is {:?} {:?}, expected 2-D U8", + fixed.dtype, + fixed.shape + ); + let (fr, fc) = (fixed.shape[0], fixed.shape[1]); + let (rows, cols) = atlas_core::modelopt_codec::csf_scale_shape(fr, fc)?; + let scales = decode_csf_scales(&fixed.bytes, fr, fc, &exc.bytes) + .with_context(|| format!("decoding CSF scales of {base}"))?; + upload( + store, + gpu, + base.to_string(), + &scales, + vec![rows, cols], + WeightDtype::FP8E4M3, + ) +} + +/// Expand one host-decoded module into a BF16 `.weight`. Returns +/// `(was_mxfp8, bytes)`, or `None` when the checkpoint does not carry it. +fn expand_module( + store: &mut WeightStore, + gpu: &dyn GpuBackend, + module: &str, +) -> Result> { + let Some(w) = take(store, gpu, &format!("{module}.weight"))? else { + return Ok(None); + }; + let s = take(store, gpu, &format!("{module}.weight_scale"))? + .with_context(|| format!("{module}.weight has no weight_scale"))?; + ensure!( + w.shape.len() == 2, + "{module}.weight is {:?}, expected 2-D", + w.shape + ); + let out = w.shape[0]; + let (bytes, inp, mx) = match w.dtype { + WeightDtype::FP8E4M3 => { + let inp = w.shape[1]; + (mxfp8_to_bf16(&w.bytes, &s.bytes, out, inp)?, inp, true) + } + WeightDtype::UInt8 => { + let inp = w.shape[1] * 2; + ensure!( + s.shape.len() == 2 && s.shape[1] > 0, + "{module}: bad NVFP4 scale shape" + ); + let group = inp / s.shape[1]; + let global = take_scalar(store, gpu, &format!("{module}.weight_scale_2"))?; + // The activation scale is a W4A4 input; a BF16 projection has none. + if let Some(t) = store.remove(&format!("{module}.input_scale")) { + gpu.free(t.ptr)?; + } + let mut dst = vec![0u8; out * inp * 2]; + nvfp4_rows_to_bf16(&w.bytes, &s.bytes, global, inp, group, &mut dst)?; + (dst, inp, false) + } + other => anyhow::bail!("{module}.weight is {other:?}; expected MXFP8 E4M3 or NVFP4 U8"), + }; + let n = upload( + store, + gpu, + format!("{module}.weight"), + &bytes, + vec![out, inp], + WeightDtype::BF16, + )?; + Ok(Some((mx, n))) +} + +/// Expand every FP4-Local storage format present in `store`. A no-op (and +/// cheap) for any checkpoint without CSF streams or host-decoded modules. +pub fn materialize( + store: &mut WeightStore, + config: &ModelConfig, + gpu: &dyn GpuBackend, +) -> Result { + let mut done = Materialized::default(); + for base in csf_bases(store) { + done.uploaded_bytes += decode_one_csf(store, gpu, &base)?; + done.csf_scales += 1; + } + let mut modules: Vec = host_decode_modules(config).into_iter().collect(); + modules.sort(); + for module in modules { + let Some((mx, n)) = expand_module(store, gpu, &module)? else { + continue; + }; + if mx { + done.mxfp8_to_bf16 += 1; + } else { + done.nvfp4_to_bf16 += 1; + } + done.uploaded_bytes += n; + } + if done != Materialized::default() { + // Only an FP4-Local checkpoint reaches here: another release's F32 + // tensors are never touched. + let (n, b) = norms_to_bf16(store, gpu)?; + done.f32_norms_to_bf16 = n; + done.uploaded_bytes += b; + tracing::info!( + "FP4-Local: decoded {} CSF scale tensors, expanded {} MXFP8 and {} NVFP4 \ + projections to BF16, rounded {} F32 norms to BF16 ({:.2} GiB uploaded; \ + compressed originals never resident)", + done.csf_scales, + done.mxfp8_to_bf16, + done.nvfp4_to_bf16, + done.f32_norms_to_bf16, + done.uploaded_bytes as f64 / (1u64 << 30) as f64 + ); + } + Ok(done) +} + +#[cfg(test)] +#[path = "modelopt_local_tests.rs"] +mod tests; diff --git a/crates/spark-model/src/weight_map/modelopt_local_tests.rs b/crates/spark-model/src/weight_map/modelopt_local_tests.rs new file mode 100644 index 000000000..511b0a7d9 --- /dev/null +++ b/crates/spark-model/src/weight_map/modelopt_local_tests.rs @@ -0,0 +1,312 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +use super::*; +use atlas_core::config::{QuantLayerSpec, QuantizationConfig}; +use spark_runtime::gpu::mock::MockGpuBackend; +use spark_runtime::weights::DeferredTensor; + +const ATTN: &str = "model.language_model.layers.3.self_attn.q_proj"; +const VIS: &str = "model.visual.blocks.0.mlp.linear_fc1"; +const HC_NORM: &str = "model.language_model.layers.0.attn_hyper_connection.hc_norm.weight"; +const PLE_NORM: &str = "model.language_model.layers.1.ple.norm_conv.weight"; +const EXPERT: &str = "model.language_model.layers.0.mlp.experts.3.down_proj.weight_scale"; + +fn spec(algo: &str, group: usize) -> QuantLayerSpec { + QuantLayerSpec { + quant_algo: algo.into(), + group_size: group, + } +} + +/// A config carrying the FP4-Local `quantized_layers` shape: CSF NVFP4 +/// experts, an MXFP8 projection, a W4A16 NVFP4 vision MLP, and the PLE table. +fn fp4_local_config() -> ModelConfig { + let mut cfg = ModelConfig::qwen3_next_80b_nvfp4(); + let mut map = std::collections::BTreeMap::new(); + map.insert( + "model.language_model.layers.0.mlp.experts".into(), + spec("NVFP4", 16), + ); + map.insert(ATTN.into(), spec("MXFP8", 32)); + map.insert(VIS.into(), spec("W4A16_NVFP4", 16)); + map.insert( + "model.language_model.layers.1.ple.ple_embedding.ngram_embedding".into(), + spec("NVFP4", 16), + ); + map.insert("mtp.layers.0.mlp.experts".into(), spec("W4A16_NVFP4", 16)); + cfg.quantization_config = Some(QuantizationConfig { + quant_method: "modelopt".into(), + quant_algo: "MIXED_PRECISION".into(), + format: String::new(), + ignore_modules: vec![], + weight_block_size: vec![], + group_size: 0, + quantized_layers: map, + exl3: None, + }); + cfg +} + +/// Only MXFP8 modules and the vision tower's NVFP4 are host-expanded: the +/// trunk experts, MTP experts and PLE table keep their own NVFP4 paths. +#[test] +fn host_decode_modules_pick_mxfp8_and_vision_nvfp4_only() { + let got = host_decode_modules(&fp4_local_config()); + let want: HashSet = [ATTN, VIS].iter().map(|s| s.to_string()).collect(); + assert_eq!(got, want); + assert!(host_decode_modules(&ModelConfig::qwen3_next_80b_nvfp4()).is_empty()); +} + +/// Writes tensors into one scratch "shard" and defers each at its offset, +/// exactly as the fast loader records them. +struct Shard { + path: std::path::PathBuf, + bytes: Vec, + entries: Vec<(String, u64, Vec, WeightDtype)>, +} +impl Shard { + fn new(tag: &str) -> Self { + let dir = std::env::temp_dir().join(format!("atlas-fp4l-{tag}-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + Self { + path: dir.join("s.bin"), + bytes: vec![0; 8], + entries: Vec::new(), + } + } + fn add(&mut self, name: &str, data: &[u8], shape: &[usize], dtype: WeightDtype) { + self.entries + .push((name.into(), self.bytes.len() as u64, shape.to_vec(), dtype)); + self.bytes.extend_from_slice(data); + } + fn defer_into(self, store: &mut WeightStore) { + std::fs::write(&self.path, &self.bytes).unwrap(); + for (name, offset, shape, dtype) in self.entries { + store.defer( + name, + DeferredTensor { + path: self.path.clone(), + offset, + shape, + dtype, + }, + ); + } + } +} + +fn resident( + store: &mut WeightStore, + gpu: &MockGpuBackend, + name: &str, + data: &[u8], + dtype: WeightDtype, +) { + let ptr = gpu.alloc(data.len()).unwrap(); + gpu.copy_h2d(data, ptr).unwrap(); + store + .insert( + name.into(), + WeightTensor { + ptr, + shape: vec![], + dtype, + }, + ) + .unwrap(); +} + +fn bf16_at(bytes: &[u8], i: usize) -> f32 { + f32::from_bits(u32::from(u16::from_le_bytes([bytes[2 * i], bytes[2 * i + 1]])) << 16) +} + +#[test] +fn materialize_expands_every_fp4_local_format() { + let gpu = MockGpuBackend::new(); + let mut store = WeightStore::empty(); + let mut shard = Shard::new("all"); + + // CSF: a [16, 4] scale matrix whose bytes are all 0x38 except one outlier. + let mut fixed = vec![0u8; 16 * 3]; + fixed[..16].fill(0x38); // bases; every code nibble 0 + let outlier = ((0x7Eu32 << 24) | 5).to_le_bytes(); + shard.add( + &format!("{EXPERT}.nvfp4_csf_fixed"), + &fixed, + &[1, 48], + WeightDtype::UInt8, + ); + shard.add( + &format!("{EXPERT}.nvfp4_csf_exceptions"), + &outlier, + &[1], + WeightDtype::Int32, + ); + + // MXFP8 [1, 32]: E4M3 1.0 (0x38) everywhere, E8M0 scale 2^-3. + shard.add( + &format!("{ATTN}.weight"), + &[0x38; 32], + &[1, 32], + WeightDtype::FP8E4M3, + ); + shard.add( + &format!("{ATTN}.weight_scale"), + &[124], + &[1, 1], + WeightDtype::UInt8, + ); + + // Vision NVFP4 [1, 16]: nibble 2 (=1.0) everywhere, block 2.0, global 0.5. + shard.add( + &format!("{VIS}.weight"), + &[0x22; 8], + &[1, 8], + WeightDtype::UInt8, + ); + shard.add( + &format!("{VIS}.weight_scale"), + &[0x40], + &[1, 1], + WeightDtype::FP8E4M3, + ); + shard.defer_into(&mut store); + resident( + &mut store, + &gpu, + &format!("{VIS}.weight_scale_2"), + &0.5f32.to_le_bytes(), + WeightDtype::FP32, + ); + resident( + &mut store, + &gpu, + &format!("{VIS}.input_scale"), + &1f32.to_le_bytes(), + WeightDtype::FP32, + ); + + // An F32 hc_norm [2] = [1.0, 2.0], as the pack ships it. + let norm: Vec = [1.0f32, 2.0].iter().flat_map(|v| v.to_le_bytes()).collect(); + let ptr = gpu.alloc(norm.len()).unwrap(); + gpu.copy_h2d(&norm, ptr).unwrap(); + let t = WeightTensor { + ptr, + shape: vec![2], + dtype: WeightDtype::FP32, + }; + store.insert(HC_NORM.into(), t).unwrap(); + let ptr = gpu.alloc(norm.len()).unwrap(); + gpu.copy_h2d(&norm, ptr).unwrap(); + let t = WeightTensor { + ptr, + shape: vec![2], + dtype: WeightDtype::FP32, + }; + store.insert(PLE_NORM.into(), t).unwrap(); + // A 1-D F32 that is not a norm (a bias, say) is left alone. + let ptr = gpu.alloc(norm.len()).unwrap(); + let t = WeightTensor { + ptr, + shape: vec![2], + dtype: WeightDtype::FP32, + }; + store + .insert("model.language_model.layers.1.ple.conv1d.bias".into(), t) + .unwrap(); + let done = materialize(&mut store, &fp4_local_config(), &gpu).unwrap(); + assert_eq!( + ( + done.csf_scales, + done.mxfp8_to_bf16, + done.nvfp4_to_bf16, + done.f32_norms_to_bf16 + ), + (1, 1, 1, 2) + ); + let norm = store.get(HC_NORM).unwrap(); + assert_eq!(norm.dtype, WeightDtype::BF16); + assert_eq!(gpu.read_alloc(norm.ptr).unwrap(), [0x80, 0x3F, 0x00, 0x40]); // 1.0, 2.0 + assert_eq!(store.get(PLE_NORM).unwrap().dtype, WeightDtype::BF16); + let bias = store + .get("model.language_model.layers.1.ple.conv1d.bias") + .unwrap(); + assert_eq!(bias.dtype, WeightDtype::FP32); + assert!( + store.deferred_sorted().is_empty(), + "every source must be consumed" + ); + + let scale = store.get(EXPERT).unwrap(); + assert_eq!( + (scale.dtype, scale.shape.clone()), + (WeightDtype::FP8E4M3, vec![16, 4]) + ); + let mut want = vec![0x38u8; 64]; + want[5] = 0x7E; + assert_eq!(gpu.read_alloc(scale.ptr).unwrap(), want); + + let attn = store.get(&format!("{ATTN}.weight")).unwrap(); + assert_eq!( + (attn.dtype, attn.shape.clone()), + (WeightDtype::BF16, vec![1, 32]) + ); + let a = gpu.read_alloc(attn.ptr).unwrap(); + assert!((0..32).all(|i| bf16_at(&a, i) == 0.125)); + assert!(!store.contains(&format!("{ATTN}.weight_scale"))); + + let vis = store.get(&format!("{VIS}.weight")).unwrap(); + assert_eq!( + (vis.dtype, vis.shape.clone()), + (WeightDtype::BF16, vec![1, 16]) + ); + let v = gpu.read_alloc(vis.ptr).unwrap(); + assert!((0..16).all(|i| bf16_at(&v, i) == 1.0)); + for gone in ["weight_scale", "weight_scale_2", "input_scale"] { + assert!( + !store.contains(&format!("{VIS}.{gone}")), + "{gone} must be consumed" + ); + } +} + +/// A CSF fixed stream with no exceptions partner is a broken checkpoint, not +/// something to serve with a half-decoded scale. +#[test] +fn unpaired_csf_stream_is_refused() { + let gpu = MockGpuBackend::new(); + let mut store = WeightStore::empty(); + let mut shard = Shard::new("unpaired"); + shard.add( + &format!("{EXPERT}.nvfp4_csf_fixed"), + &[0x38; 48], + &[1, 48], + WeightDtype::UInt8, + ); + shard.defer_into(&mut store); + let err = materialize(&mut store, &fp4_local_config(), &gpu).unwrap_err(); + assert!( + format!("{err:#}").contains("without its exceptions partner"), + "{err:#}" + ); +} + +/// Every other checkpoint passes through untouched. +#[test] +fn materialize_is_a_no_op_without_fp4_local_tensors() { + let gpu = MockGpuBackend::new(); + let mut store = WeightStore::empty(); + resident(&mut store, &gpu, "a.weight", &[0; 4], WeightDtype::BF16); + resident( + &mut store, + &gpu, + "b.norm.weight", + &[0; 4], + WeightDtype::FP32, + ); + let done = materialize(&mut store, &ModelConfig::qwen3_next_80b_nvfp4(), &gpu).unwrap(); + assert_eq!(done, Materialized::default()); + assert!(store.contains("a.weight")); + // Another release's F32 tensors are never touched. + assert_eq!(store.get("b.norm.weight").unwrap().dtype, WeightDtype::FP32); +} diff --git a/crates/spark-server/src/main_modules/serve_phases/weights.rs b/crates/spark-server/src/main_modules/serve_phases/weights.rs index de19280a3..12cae302b 100644 --- a/crates/spark-server/src/main_modules/serve_phases/weights.rs +++ b/crates/spark-server/src/main_modules/serve_phases/weights.rs @@ -153,6 +153,14 @@ pub(crate) fn load_weight_store( loader.demand_paged_patterns ); } + loader.host_decode_modules = + spark_model::weight_map::modelopt_local::host_decode_modules(config); + if !loader.host_decode_modules.is_empty() { + tracing::info!( + "Host-decoded at build, never resident: {} MXFP8 / vision-NVFP4 modules", + loader.host_decode_modules.len() + ); + } loader.skip_activation_scales = skip_activation_scales(config); loader.skip_mtp = skip_mtp(config, args); loader.skip_layer_prefix = unused_mtp_prefix.clone(); From 2c6b24d80095d112744f4010ae270fdfa90d1ca4 Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sat, 17 Oct 2026 21:03:43 -0400 Subject: [PATCH 04/26] qwen4_exp MTP: per-expert NVFP4 experts (FP4-Local) nvidia/Qwen3.8-Flash-Next-FP4-Local quantizes mtp.layers.0.mlp.experts as W4A16 NVFP4 with plain block scales, the trunk experts' own format, so a PerExpertNvfp4 layout loads them through quantized_any as shipped. The load log now names the detected layout instead of always claiming a stacked BF16 requantization. mtp.rs's tests move to mtp_tests.rs. Verified on a GB10: --speculative --num-drafts 1 boots, K2 acceptance 73-75%, decode 18.9-19.8 tok/s (16.9 without), outputs coherent and the tool call unchanged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01L4aUGgtiZX94pLr63HXQzn --- .../src/weight_loader/qwen4_exp/mtp.rs | 77 ++++--------------- .../src/weight_loader/qwen4_exp/mtp_nvfp4.rs | 47 +++++++++++ .../src/weight_loader/qwen4_exp/mtp_tests.rs | 74 ++++++++++++++++++ 3 files changed, 138 insertions(+), 60 deletions(-) create mode 100644 crates/spark-model/src/weight_loader/qwen4_exp/mtp_nvfp4.rs create mode 100644 crates/spark-model/src/weight_loader/qwen4_exp/mtp_tests.rs diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/mtp.rs b/crates/spark-model/src/weight_loader/qwen4_exp/mtp.rs index 2c2c6e778..d3e4011c8 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/mtp.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/mtp.rs @@ -80,6 +80,9 @@ enum MtpExpertLayout { /// `weight_scale_inv` (FP8_PB_WO, g=128) — declared by the pack's /// `quantized_layers` map entry `mtp.layers.0.mlp.experts`. PerExpertFp8BlockScaled, + /// nvidia FP4-Local: per-expert NVFP4 (`weight_scale` + `weight_scale_2`), + /// the trunk experts' own format — see `nvfp4`. + PerExpertNvfp4, } /// Detect the MTP expert storage layout from the store. `weight_scale_inv` @@ -92,6 +95,9 @@ fn mtp_expert_layout(store: &WeightStore, mlp: &str) -> MtpExpertLayout { if store.contains(&format!("{mlp}.experts.0.gate_proj.weight_scale_inv")) { return MtpExpertLayout::PerExpertFp8BlockScaled; } + if store.contains(&format!("{mlp}.experts.0.gate_proj.weight_scale_2")) { + return MtpExpertLayout::PerExpertNvfp4; + } // Default to the stacked path so its own error names what it expected. MtpExpertLayout::StackedBf16 } @@ -177,6 +183,9 @@ fn build_mtp_moe( }); } } + MtpExpertLayout::PerExpertNvfp4 => { + experts = nvfp4::experts(store, &mlp, config, gpu, absmax_k, quantize_k, stream)?; + } // nvidia pack: per-expert FP8_PB_WO (weight + weight_scale_inv, g=128 // blocks) instead of RadixArk's fused BF16 pair. Dequant to BF16 then // requant to NVFP4 so the body gets the same MoeLayer either way. @@ -420,71 +429,19 @@ pub fn load_qwen4exp_mtp_module( let spent = free_before.saturating_sub(gpu.free_memory().unwrap_or(0) as u64); tracing::info!( "qwen4_exp MTP draft module loaded: 1 reused full-attention layer \ - (gated attn, dense — no QSA indexer + mHC + {}-expert MoE requantized from stacked BF16), \ + (gated attn, dense — no QSA indexer + mHC + {}-expert MoE, experts {:?}), \ shared embed/lm_head, own head mixer. Construction cost {:.2} GB.", config.num_experts, + mtp_expert_layout(store, &format!("{MTP_LAYER_PREFIX}.mlp")), spent as f64 / 1e9, ); Ok(Some(module)) } -#[cfg(test)] -mod tests { - use super::*; - use std::collections::HashMap; - - fn store_with(names: &[&str]) -> WeightStore { - let map: HashMap = names - .iter() - .map(|n| { - ( - n.to_string(), - spark_runtime::weights::WeightTensor { - ptr: spark_runtime::gpu::DevicePtr::NULL, - shape: vec![1], - dtype: spark_runtime::weights::WeightDtype::BF16, - }, - ) - }) - .collect(); - WeightStore::from_map(map) - } - - /// RadixArk pack: the fused pair selects the stacked path. - #[test] - fn stacked_markers_select_stacked_layout() { - let store = store_with(&[ - "mtp.layers.0.mlp.experts.gate_up_proj", - "mtp.layers.0.mlp.experts.down_proj", - ]); - assert_eq!( - mtp_expert_layout(&store, "mtp.layers.0.mlp"), - MtpExpertLayout::StackedBf16 - ); - } +/// Per-expert NVFP4 MTP experts (FP4-Local). +#[path = "mtp_nvfp4.rs"] +mod nvfp4; - /// nvidia pack: per-expert weight_scale_inv selects the FP8 block path — - /// and must win even when a stray fused name is absent. - #[test] - fn per_expert_scale_inv_selects_fp8_layout() { - let store = store_with(&[ - "mtp.layers.0.mlp.experts.0.gate_proj.weight", - "mtp.layers.0.mlp.experts.0.gate_proj.weight_scale_inv", - ]); - assert_eq!( - mtp_expert_layout(&store, "mtp.layers.0.mlp"), - MtpExpertLayout::PerExpertFp8BlockScaled - ); - } - - /// Neither marker → stacked default, whose loader errors with the names - /// it probed (better than guessing the other format). - #[test] - fn no_markers_defaults_to_stacked() { - let store = store_with(&[]); - assert_eq!( - mtp_expert_layout(&store, "mtp.layers.0.mlp"), - MtpExpertLayout::StackedBf16 - ); - } -} +#[cfg(test)] +#[path = "mtp_tests.rs"] +mod tests; diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/mtp_nvfp4.rs b/crates/spark-model/src/weight_loader/qwen4_exp/mtp_nvfp4.rs new file mode 100644 index 000000000..7f475e679 --- /dev/null +++ b/crates/spark-model/src/weight_loader/qwen4_exp/mtp_nvfp4.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! MTP routed experts of nvidia/Qwen3.8-Flash-Next-FP4-Local. +//! +//! That pack quantizes `mtp.layers.0.mlp.experts` as W4A16 NVFP4 with plain +//! E4M3 block scales — the exact per-expert format of the trunk experts — so +//! they load through the same `quantized_any` path, uploaded as shipped with +//! no requantization (the NVFP4 pack's FP8 MTP experts need a BF16 round +//! trip; these do not). + +use anyhow::{Context, Result}; +use atlas_core::config::ModelConfig; +use spark_runtime::gpu::{GpuBackend, KernelHandle}; +use spark_runtime::weights::WeightStore; + +use crate::weight_map::{ExpertWeight, Nvfp4Variant, QuantizeCtx, quantized_any}; + +pub(super) fn experts( + store: &WeightStore, + mlp: &str, + config: &ModelConfig, + gpu: &dyn GpuBackend, + absmax_k: KernelHandle, + quantize_k: KernelHandle, + stream: u64, +) -> Result> { + let (h, inter) = (config.hidden_size, config.moe_intermediate_size); + let qctx = QuantizeCtx { + absmax_k, + quantize_k, + stream, + }; + let load = |prefix: String, n: usize, k: usize| { + quantized_any(store, &prefix, n, k, gpu, Nvfp4Variant::Standard, qctx) + .with_context(|| format!("qwen4_exp MTP: NVFP4 {prefix}")) + }; + (0..config.num_experts) + .map(|e| { + let p = |proj: &str| format!("{mlp}.experts.{e}.{proj}_proj"); + Ok(ExpertWeight { + gate_proj: load(p("gate"), inter, h)?, + up_proj: load(p("up"), inter, h)?, + down_proj: load(p("down"), h, inter)?, + }) + }) + .collect() +} diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/mtp_tests.rs b/crates/spark-model/src/weight_loader/qwen4_exp/mtp_tests.rs new file mode 100644 index 000000000..3d5c28cb3 --- /dev/null +++ b/crates/spark-model/src/weight_loader/qwen4_exp/mtp_tests.rs @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +use super::*; +use std::collections::HashMap; + +fn store_with(names: &[&str]) -> WeightStore { + let map: HashMap = names + .iter() + .map(|n| { + ( + n.to_string(), + spark_runtime::weights::WeightTensor { + ptr: spark_runtime::gpu::DevicePtr::NULL, + shape: vec![1], + dtype: spark_runtime::weights::WeightDtype::BF16, + }, + ) + }) + .collect(); + WeightStore::from_map(map) +} + +/// RadixArk pack: the fused pair selects the stacked path. +#[test] +fn stacked_markers_select_stacked_layout() { + let store = store_with(&[ + "mtp.layers.0.mlp.experts.gate_up_proj", + "mtp.layers.0.mlp.experts.down_proj", + ]); + assert_eq!( + mtp_expert_layout(&store, "mtp.layers.0.mlp"), + MtpExpertLayout::StackedBf16 + ); +} + +/// nvidia pack: per-expert weight_scale_inv selects the FP8 block path — +/// and must win even when a stray fused name is absent. +#[test] +fn per_expert_scale_inv_selects_fp8_layout() { + let store = store_with(&[ + "mtp.layers.0.mlp.experts.0.gate_proj.weight", + "mtp.layers.0.mlp.experts.0.gate_proj.weight_scale_inv", + ]); + assert_eq!( + mtp_expert_layout(&store, "mtp.layers.0.mlp"), + MtpExpertLayout::PerExpertFp8BlockScaled + ); +} + +/// Neither marker → stacked default, whose loader errors with the names +/// it probed (better than guessing the other format). +#[test] +fn no_markers_defaults_to_stacked() { + let store = store_with(&[]); + assert_eq!( + mtp_expert_layout(&store, "mtp.layers.0.mlp"), + MtpExpertLayout::StackedBf16 + ); +} + +/// nvidia FP4-Local: per-expert NVFP4 (a `weight_scale_2` and no +/// `weight_scale_inv`) selects the trunk-format NVFP4 path. +#[test] +fn per_expert_global_scale_selects_nvfp4_layout() { + let store = store_with(&[ + "mtp.layers.0.mlp.experts.0.gate_proj.weight", + "mtp.layers.0.mlp.experts.0.gate_proj.weight_scale", + "mtp.layers.0.mlp.experts.0.gate_proj.weight_scale_2", + ]); + assert_eq!( + mtp_expert_layout(&store, "mtp.layers.0.mlp"), + MtpExpertLayout::PerExpertNvfp4 + ); +} From 9786ebe2912e5325972b14d6446846a3b0822946 Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:37:46 -0400 Subject: [PATCH 05/26] feat(qwen4_exp): serve FP4-Local's GDN projections on their native MXFP8 bytes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit nvidia/Qwen3.8-Flash-Next-FP4-Local ships the GDN in_proj_qkv, in_proj_z and out_proj (36 layers, 2.08 GB) as MXFP8: E4M3 plus one E8M0 scale per 32 inputs, QAD-trained that way. They were expanded to 4.15 GB of BF16 at load and read by dense_gemv_bf16 — 44% of a decode token's GPU time. - kernels: mxfp8_proj.cu — mxfp8_proj_gemv (1 row), mxfp8_proj_batchm (2..8 rows, one weight pass, every row bit-identical to the 1-row kernel) and mxfp8_proj_dequant (exact BF16 expansion for prefill). - ops::mxfp8_proj picks by row count; prefill expands one projection at a time into a single 84 MB scratch shared by all GDN layers, then bf16_gemm. - materialize uploads native_mxfp8() modules byte for byte; ATLAS_FP4LOCAL_EXPAND_MXFP8=1 restores the BF16 expansion for A/B. - qwen4_exp::gdn_mx8 builds the layer on those bytes; Qwen3SsmLayer::mx8 is the only copy of both projections and is checked first at every dispatch site (serial decode, batched verify, multi-seq decode, prefill). Measured on GB10 (sparky): - microtest: dequant byte-identical to modelopt_codec::mxfp8_to_bf16 over 42M values; GEMV vs f64 rel_l2 1.7e-3; batchm M=1..8 0 mismatches; 230 GB/s, 1.89x dense_gemv_bf16 on [16384, 2560]; M=8 costs 1.09x M=1. - serve: pre-KV 97.6 -> 95.0 GB; decode with --speculative --num-drafts 1 18.9-19.8 -> 24.2-28.6 tok/s; coherence prompts and tool call correct. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01L4aUGgtiZX94pLr63HXQzn --- crates/spark-model/Cargo.toml | 4 + .../examples/mxfp8_proj_microtest.rs | 258 ++++++++++++++++++ crates/spark-model/src/layers/ops.rs | 3 + .../spark-model/src/layers/ops/mxfp8_proj.rs | 134 +++++++++ .../src/layers/qwen3_ssm/exact_rows.rs | 25 ++ .../spark-model/src/layers/qwen3_ssm/init.rs | 1 + .../spark-model/src/layers/qwen3_ssm/mod.rs | 6 + .../spark-model/src/layers/qwen3_ssm/mx8.rs | 104 +++++++ .../src/layers/qwen3_ssm/mx8_tests.rs | 137 ++++++++++ .../src/layers/qwen3_ssm/ssm_decode_tail.rs | 4 +- .../src/layers/qwen3_ssm/ssm_forward.rs | 5 +- .../spark-model/src/layers/qwen3_ssm/tests.rs | 4 +- .../ssm_batched/impl.rs | 5 +- .../layers/qwen3_ssm/trait_prefill_helper.rs | 4 + .../layers/qwen3_ssm/trait_prefill_proj.rs | 4 + .../src/weight_loader/qwen4_exp.rs | 24 +- .../src/weight_loader/qwen4_exp/gdn_fp8.rs | 12 +- .../src/weight_loader/qwen4_exp/gdn_mx8.rs | 194 +++++++++++++ .../src/weight_map/modelopt_local.rs | 78 +++++- .../src/weight_map/modelopt_local_tests.rs | 76 +++++- .../qwen3.8-flash-next/nvfp4/mxfp8_proj.cu | 214 +++++++++++++++ 21 files changed, 1269 insertions(+), 27 deletions(-) create mode 100644 crates/spark-model/examples/mxfp8_proj_microtest.rs create mode 100644 crates/spark-model/src/layers/ops/mxfp8_proj.rs create mode 100644 crates/spark-model/src/layers/qwen3_ssm/mx8.rs create mode 100644 crates/spark-model/src/layers/qwen3_ssm/mx8_tests.rs create mode 100644 crates/spark-model/src/weight_loader/qwen4_exp/gdn_mx8.rs create mode 100644 kernels/gb10/qwen3.8-flash-next/nvfp4/mxfp8_proj.cu diff --git a/crates/spark-model/Cargo.toml b/crates/spark-model/Cargo.toml index 2e83c23fd..86e72f238 100644 --- a/crates/spark-model/Cargo.toml +++ b/crates/spark-model/Cargo.toml @@ -207,6 +207,10 @@ required-features = ["cuda", "gpu-examples"] name = "w8a16_gemv_batch4_microtest" required-features = ["cuda", "gpu-examples"] +[[example]] +name = "mxfp8_proj_microtest" +required-features = ["cuda", "gpu-examples"] + [[example]] name = "q2_0_gemv_microtest" required-features = ["cuda", "gpu-examples"] diff --git a/crates/spark-model/examples/mxfp8_proj_microtest.rs b/crates/spark-model/examples/mxfp8_proj_microtest.rs new file mode 100644 index 000000000..49b132f6e --- /dev/null +++ b/crates/spark-model/examples/mxfp8_proj_microtest.rs @@ -0,0 +1,258 @@ +// SPDX-License-Identifier: AGPL-3.0-only +//! Correctness and speed of the native MXFP8 kernels (`mxfp8_proj.cu`). +//! +//! Checks, on random E4M3 bytes and E8M0 scales of a real GDN shape: +//! 1. `mxfp8_proj_dequant` is byte-identical to `atlas_core::modelopt_codec::mxfp8_to_bf16`; +//! 2. `mxfp8_proj_gemv` matches an f64 reference over that exact BF16 weight; +//! 3. every row of `mxfp8_proj_batchm` (M = 1..=8) is bit-identical to `mxfp8_proj_gemv`; +//! +//! then times `mxfp8_proj_gemv` against `dense_gemv_bf16` over the dequantized +//! weight (the path FP4-Local ran before this kernel). Run: +//! ATLAS_TARGET_HW=gb10 ATLAS_TARGET_MODEL=qwen3.8-flash-next ATLAS_TARGET_QUANT=nvfp4 \ +//! cargo run -p spark-model --release --features cuda,gpu-examples \ +//! --example mxfp8_proj_microtest [-- N K] + +use anyhow::{Result, ensure}; +use atlas_core::modelopt_codec::mxfp8_to_bf16; +use spark_runtime::cuda_backend::AtlasCudaBackend; +use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle}; +use spark_runtime::kernel_args::{KernelLaunch, div_ceil}; + +const MAXM: usize = 8; + +struct Lcg(u64); +impl Lcg { + fn next(&mut self) -> u64 { + self.0 = self + .0 + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + self.0 >> 11 + } + fn unit(&mut self) -> f32 { + ((self.next() as f64) / ((1u64 << 53) as f64)) as f32 + } +} + +fn up(g: &dyn GpuBackend, d: &[u8]) -> Result { + let p = g.alloc(d.len().max(1))?; + g.copy_h2d(d, p)?; + Ok(p) +} + +fn down(g: &dyn GpuBackend, p: DevicePtr, bytes: usize) -> Result> { + let mut b = vec![0u8; bytes]; + g.copy_d2h(p, &mut b)?; + Ok(b) +} + +fn bf16_at(b: &[u8], i: usize) -> f32 { + f32::from_bits(u32::from(u16::from_le_bytes([b[2 * i], b[2 * i + 1]])) << 16) +} + +fn gemv( + g: &dyn GpuBackend, + k_gemv: KernelHandle, + (a, w, s, c): (DevicePtr, DevicePtr, DevicePtr, DevicePtr), + n: usize, + k: usize, +) -> Result<()> { + KernelLaunch::new(g, k_gemv) + .grid([div_ceil(n as u32, 4), 1, 1]) + .block([256, 1, 1]) + .arg_ptr(a) + .arg_ptr(w) + .arg_ptr(s) + .arg_ptr(c) + .arg_u32(n as u32) + .arg_u32(k as u32) + .launch(0) +} + +fn time_ms(g: &dyn GpuBackend, reps: usize, mut f: impl FnMut() -> Result<()>) -> Result { + f()?; + g.synchronize(0)?; + let t = std::time::Instant::now(); + for _ in 0..reps { + f()?; + } + g.synchronize(0)?; + Ok(t.elapsed().as_secs_f64() * 1e3 / reps as f64) +} + +fn main() -> Result<()> { + let args: Vec = std::env::args().collect(); + // GDN in_proj_qkvz of Qwen3.8-Flash-Next: [16384, 2560]. + let n: usize = args.get(1).map_or(Ok(16384), |v| v.parse())?; + let k: usize = args.get(2).map_or(Ok(2560), |v| v.parse())?; + ensure!(k.is_multiple_of(32), "K must be a multiple of 32"); + + let backend = AtlasCudaBackend::new(0, &atlas_kernels::ptx_modules())?; + let g: &dyn GpuBackend = &backend; + let k_gemv = g.kernel("mxfp8_proj", "mxfp8_proj_gemv")?; + let k_batch = g.kernel("mxfp8_proj", "mxfp8_proj_batchm")?; + let k_deq = g.kernel("mxfp8_proj", "mxfp8_proj_dequant")?; + let k_bf16 = g.kernel("gemv", "dense_gemv_bf16")?; + + let mut rng = Lcg(0x4d58_4650_3800_0001); + // Avoid the NaN bytes 0x7f/0xff: the reference maps them to 0 as well, + // but real checkpoints never carry them. + let weight: Vec = (0..n * k) + .map(|_| match (rng.next() & 0xFF) as u8 { + 0x7F | 0xFF => 0, + b => b, + }) + .collect(); + // Scales around the checkpoint's range: 2^-12 .. 2^-4. + let scale: Vec = (0..n * k / 32) + .map(|_| 115 + (rng.next() % 9) as u8) + .collect(); + let act: Vec = (0..MAXM * k) + .flat_map(|_| { + let v = half::bf16::from_f32(rng.unit() * 2.0 - 1.0); + v.to_bits().to_le_bytes() + }) + .collect(); + + let (w_d, s_d, a_d) = (up(g, &weight)?, up(g, &scale)?, up(g, &act)?); + let deq_d = g.alloc(n * k * 2)?; + let c1 = g.alloc(MAXM * n * 2)?; + let cb = g.alloc(MAXM * n * 2)?; + let mut pass = true; + + // 1. Dequant is the codec, byte for byte. + KernelLaunch::new(g, k_deq) + .grid([div_ceil((n * k / 16) as u32, 256), 1, 1]) + .block([256, 1, 1]) + .arg_ptr(w_d) + .arg_ptr(s_d) + .arg_ptr(deq_d) + .arg_u32(n as u32) + .arg_u32(k as u32) + .launch(0)?; + g.synchronize(0)?; + let deq = down(g, deq_d, n * k * 2)?; + let want = mxfp8_to_bf16(&weight, &scale, n, k)?; + let bad = deq + .chunks_exact(2) + .zip(want.chunks_exact(2)) + .filter(|(a, b)| a != b) + .count(); + println!( + "dequant vs codec: {bad} of {} values differ {}", + n * k, + verdict(bad == 0) + ); + pass &= bad == 0; + + // 2. GEMV against an f64 reference over the same exact weight. + gemv(g, k_gemv, (a_d, w_d, s_d, c1), n, k)?; + g.synchronize(0)?; + let got = down(g, c1, n * 2)?; + let (mut worst, mut num, mut den) = (0f64, 0f64, 0f64); + for row in 0..n { + let r: f64 = (0..k) + .map(|i| f64::from(bf16_at(&act, i)) * f64::from(bf16_at(&want, row * k + i))) + .sum(); + let e = (f64::from(bf16_at(&got, row)) - r).abs(); + worst = worst.max(e / r.abs().max(1.0)); + num += e * e; + den += r * r; + } + let rel_l2 = (num / den.max(1e-30)).sqrt(); + // BF16 output rounding alone is 2^-9 relative. + let ok = rel_l2 < 4e-3 && worst < 1.6e-2; + println!( + "gemv vs f64: rel_l2 {rel_l2:.3e}, worst {worst:.3e} {}", + verdict(ok) + ); + pass &= ok; + + // 3. Batched rows are the single-row kernel's bits. + for m in 1..=MAXM { + KernelLaunch::new(g, k_batch) + .grid([div_ceil(n as u32, 4), 1, 1]) + .block([256, 1, 1]) + .arg_ptr(a_d) + .arg_ptr(w_d) + .arg_ptr(s_d) + .arg_ptr(cb) + .arg_u32(m as u32) + .arg_u32(n as u32) + .arg_u32(k as u32) + .launch(0)?; + for t in 0..m { + gemv( + g, + k_gemv, + (a_d.offset(t * k * 2), w_d, s_d, c1.offset(t * n * 2)), + n, + k, + )?; + } + g.synchronize(0)?; + let (b, r) = (down(g, cb, m * n * 2)?, down(g, c1, m * n * 2)?); + let bad = b + .chunks_exact(2) + .zip(r.chunks_exact(2)) + .filter(|(x, y)| x != y) + .count(); + println!( + "batchm M={m}: {bad} mismatches vs mxfp8_proj_gemv {}", + verdict(bad == 0) + ); + pass &= bad == 0; + } + + // Speed: native MXFP8 vs the BF16 expansion it replaces. + let reps = 200; + let t_mx = time_ms(g, reps, || gemv(g, k_gemv, (a_d, w_d, s_d, c1), n, k))?; + let t_bf = time_ms(g, reps, || { + KernelLaunch::new(g, k_bf16) + .grid([div_ceil(n as u32, 4), 1, 1]) + .block([256, 1, 1]) + .arg_ptr(a_d) + .arg_ptr(deq_d) + .arg_ptr(c1) + .arg_u32(n as u32) + .arg_u32(k as u32) + .launch(0) + })?; + let gbs = |bytes: usize, ms: f64| bytes as f64 / ms / 1e6; + println!( + "mxfp8_proj_gemv {t_mx:.4} ms ({:.0} GB/s) dense_gemv_bf16 {t_bf:.4} ms ({:.0} GB/s) speedup {:.2}x", + gbs(n * k + n * k / 32, t_mx), + gbs(n * k * 2, t_bf), + t_bf / t_mx + ); + for m in [2usize, 4, 8] { + let t = time_ms(g, reps, || { + KernelLaunch::new(g, k_batch) + .grid([div_ceil(n as u32, 4), 1, 1]) + .block([256, 1, 1]) + .arg_ptr(a_d) + .arg_ptr(w_d) + .arg_ptr(s_d) + .arg_ptr(cb) + .arg_u32(m as u32) + .arg_u32(n as u32) + .arg_u32(k as u32) + .launch(0) + })?; + println!( + "mxfp8_proj_batchm M={m} {t:.4} ms ({:.2}x of M=1)", + t / t_mx + ); + } + + if pass { + println!("ALL PASS"); + Ok(()) + } else { + std::process::exit(1) + } +} + +fn verdict(ok: bool) -> &'static str { + if ok { "PASS" } else { "FAIL" } +} diff --git a/crates/spark-model/src/layers/ops.rs b/crates/spark-model/src/layers/ops.rs index 3abdfabf0..9114738ec 100644 --- a/crates/spark-model/src/layers/ops.rs +++ b/crates/spark-model/src/layers/ops.rs @@ -29,6 +29,8 @@ mod dispatch_helpers; mod dispatch_proj; #[path = "ops/mxfp8.rs"] mod mxfp8; +#[path = "ops/mxfp8_proj.rs"] +mod mxfp8_proj; // Row-wise FP8 routing, split out when it took dispatch_proj.rs over the cap. #[path = "ops/dispatch_proj_rowwise.rs"] mod dispatch_proj_rowwise; @@ -334,6 +336,7 @@ pub use moe_lora_grouped::*; pub use moe_nvfp4_mmq::*; pub use moe_prefill::*; pub use mxfp8::*; +pub use mxfp8_proj::*; pub use norm::*; pub use nvfp4_mmq::*; pub use ple::*; diff --git a/crates/spark-model/src/layers/ops/mxfp8_proj.rs b/crates/spark-model/src/layers/ops/mxfp8_proj.rs new file mode 100644 index 000000000..652dc62d2 --- /dev/null +++ b/crates/spark-model/src/layers/ops/mxfp8_proj.rs @@ -0,0 +1,134 @@ +// SPDX-License-Identifier: AGPL-3.0-only +//! Native MXFP8 projections over checkpoint bytes (`mxfp8_proj.cu`). +//! +//! NVIDIA's Qwen3.8-Flash-Next-FP4-Local ships the GDN, attention, indexer +//! and shared-expert projections as MXFP8: E4M3 `[n, k]` plus one E8M0 scale +//! per 32 inputs `[n, k / 32]`. [`mxfp8_proj`] multiplies by those bytes as +//! they are: +//! +//! * 1 row: `mxfp8_proj_gemv`, one pass over the weight; +//! * 2..=8 rows (verify windows, small batches): `mxfp8_proj_batchm`, still +//! one pass, every row bit-identical to the 1-row kernel; +//! * more rows (prefill): `mxfp8_proj_dequant` expands the weight into a BF16 +//! scratch, exactly (E4M3 times a power of two is a BF16 value), and +//! [`super::bf16_gemm`] multiplies by it. The scratch is transient and +//! shared, so no BF16 copy of the weight stays resident. + +use anyhow::{Result, ensure}; +use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle}; +use spark_runtime::kernel_args::{KernelLaunch, div_ceil}; + +use crate::weight_map::Mxfp8Weight; + +/// Widest row count the weight-streaming GEMVs serve. +pub const MXFP8_PROJ_MAX_GEMV_ROWS: usize = 8; + +/// Kernel handles of `mxfp8_proj.cu`. +#[derive(Debug, Clone, Copy)] +pub struct Mxfp8ProjKernels { + pub gemv: KernelHandle, + pub batchm: KernelHandle, + pub dequant: KernelHandle, +} + +impl Mxfp8ProjKernels { + /// Resolve all three or fail: an MXFP8 weight with no kernel to read it + /// must not load. + pub fn resolve(gpu: &dyn GpuBackend) -> Result { + Ok(Self { + gemv: gpu.kernel("mxfp8_proj", "mxfp8_proj_gemv")?, + batchm: gpu.kernel("mxfp8_proj", "mxfp8_proj_batchm")?, + dequant: gpu.kernel("mxfp8_proj", "mxfp8_proj_dequant")?, + }) + } +} + +/// One MXFP8 projection and its geometry. +#[derive(Debug, Clone, Copy)] +pub struct Mxfp8Proj { + pub w: Mxfp8Weight, + /// Output features (weight rows). + pub n: u32, + /// Input features; a multiple of 32. + pub k: u32, +} + +impl Mxfp8Proj { + /// Bytes of the BF16 expansion the prefill path writes into scratch. + pub fn bf16_bytes(&self) -> usize { + self.n as usize * self.k as usize * 2 + } +} + +/// A BF16 scratch for the prefill expansion: device pointer and capacity. +#[derive(Debug, Clone, Copy)] +pub struct Mxfp8Scratch { + pub ptr: DevicePtr, + pub bytes: usize, +} + +/// `out[m, n] = act[m, k] @ W[n, k]^T` for a contiguous `act` (`k` BF16 per +/// row) and `out` (`n` BF16 per row). +#[allow(clippy::too_many_arguments)] +pub fn mxfp8_proj( + gpu: &dyn GpuBackend, + kernels: &Mxfp8ProjKernels, + proj: &Mxfp8Proj, + act: DevicePtr, + out: DevicePtr, + m: usize, + scratch: Mxfp8Scratch, + stream: u64, +) -> Result<()> { + let Mxfp8Proj { w, n, k } = *proj; + ensure!( + m > 0 && n > 0 && k.is_multiple_of(32), + "mxfp8_proj: m={m} n={n} k={k} (k must be a multiple of 32)" + ); + if m == 1 { + return KernelLaunch::new(gpu, kernels.gemv) + .grid([div_ceil(n, 4), 1, 1]) + .block([256, 1, 1]) + .arg_ptr(act) + .arg_ptr(w.data) + .arg_ptr(w.scales) + .arg_ptr(out) + .arg_u32(n) + .arg_u32(k) + .launch(stream); + } + if m <= MXFP8_PROJ_MAX_GEMV_ROWS { + return KernelLaunch::new(gpu, kernels.batchm) + .grid([div_ceil(n, 4), 1, 1]) + .block([256, 1, 1]) + .arg_ptr(act) + .arg_ptr(w.data) + .arg_ptr(w.scales) + .arg_ptr(out) + .arg_u32(m as u32) + .arg_u32(n) + .arg_u32(k) + .launch(stream); + } + ensure!( + !scratch.ptr.is_null() && scratch.bytes >= proj.bf16_bytes(), + "mxfp8_proj: {m} rows need a {} B BF16 scratch, have {} B", + proj.bf16_bytes(), + scratch.bytes + ); + let chunks = u64::from(n) * u64::from(k / 16); + ensure!( + chunks.div_ceil(256) <= u64::from(u32::MAX), + "mxfp8_proj: {n}x{k} is too large for one dequant grid" + ); + KernelLaunch::new(gpu, kernels.dequant) + .grid([chunks.div_ceil(256) as u32, 1, 1]) + .block([256, 1, 1]) + .arg_ptr(w.data) + .arg_ptr(w.scales) + .arg_ptr(scratch.ptr) + .arg_u32(n) + .arg_u32(k) + .launch(stream)?; + super::bf16_gemm(act, scratch.ptr.0, out, m as u32, n, k, stream) +} diff --git a/crates/spark-model/src/layers/qwen3_ssm/exact_rows.rs b/crates/spark-model/src/layers/qwen3_ssm/exact_rows.rs index 205098474..ffb03721c 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/exact_rows.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/exact_rows.rs @@ -36,6 +36,10 @@ impl Qwen3SsmLayer { /// the switch is off or serial decode would take an arm this has no /// exact batched twin for (Q2, interleaved QKVZ, `ATLAS_GDN_FP8_DECODE`): /// the caller keeps its own dispatch. + /// + /// Native MXFP8 projections (`mx8`) are handled here whatever the switch + /// says: they are the layer's only copy, and `mxfp8_proj_batchm` rows are + /// the single-row kernel's bytes, so the exact path is the only path. pub(super) fn exact_rows_proj( &self, proj: GdnProj, @@ -45,6 +49,27 @@ impl Qwen3SsmLayer { ctx: &ForwardContext, stream: u64, ) -> Result { + if let Some(mx8) = self.mx8.as_ref() { + let p = match proj { + GdnProj::Qkvz => &mx8.qkvz, + GdnProj::Out => &mx8.out, + }; + let (in_row, out_row) = (p.k as usize * 2, p.n as usize * 2); + for first in (0..rows).step_by(ops::MXFP8_PROJ_MAX_GEMV_ROWS) { + let m = (rows - first).min(ops::MXFP8_PROJ_MAX_GEMV_ROWS); + ops::mxfp8_proj( + ctx.gpu, + &mx8.kernels, + p, + input.offset(first * in_row), + output.offset(first * out_row), + m, + mx8.scratch, + stream, + )?; + } + return Ok(true); + } if !ctx.levers.qwen4exp_batch_fast || rows == 0 || !self.sequential_qkvz diff --git a/crates/spark-model/src/layers/qwen3_ssm/init.rs b/crates/spark-model/src/layers/qwen3_ssm/init.rs index 8fa434739..bd6652d68 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/init.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/init.rs @@ -48,6 +48,7 @@ impl Qwen3SsmLayer { out_proj_fp8w_t: None, qkvz_fp8w_rowwise: None, out_proj_fp8w_rowwise: None, + mx8: None, qkvz_q2: None, q2_0_gemv_k: super::super::try_kernel(gpu, "q2_0_gemv_vec", "q2_0_gemv_vec"), dequant_q2_0_gn_k: super::super::try_kernel( diff --git a/crates/spark-model/src/layers/qwen3_ssm/mod.rs b/crates/spark-model/src/layers/qwen3_ssm/mod.rs index 226ec0705..bb2bb19f1 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/mod.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/mod.rs @@ -88,6 +88,9 @@ pub struct Qwen3SsmLayer { /// block-scaled `w8a16_*` arms). qkvz_fp8w_rowwise: Option, out_proj_fp8w_rowwise: Option, + /// Native MXFP8 `[Q|K|V|Z]` and out_proj (FP4-Local). When set, it is + /// the only copy of both projections; see `mx8`. + mx8: Option, /// Tier-1c keep-packed ternary Q2_0 fused in_proj_qkvz (`ATLAS_GGUF_NATIVE_Q2`). /// [Q|K|V|Z] rows byte-concatenated from packed `in_proj_qkv` (V-region /// row-permuted) + `in_proj_z` (row-permuted) at load, so the 2-bit weight is @@ -429,6 +432,7 @@ mod init_fp8; mod init_q2; mod kernel_select; mod lora; +pub mod mx8; mod ssm_decode_tail; mod ssm_forward; pub(crate) mod ssm_h_fp16; @@ -462,6 +466,8 @@ pub use gdn_flags::{ // ── TransformerLayer impl (delegates to per-file inherent _inner methods) ── +#[cfg(test)] +mod mx8_tests; #[cfg(test)] mod tests; diff --git a/crates/spark-model/src/layers/qwen3_ssm/mx8.rs b/crates/spark-model/src/layers/qwen3_ssm/mx8.rs new file mode 100644 index 000000000..4a648e50d --- /dev/null +++ b/crates/spark-model/src/layers/qwen3_ssm/mx8.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Native MXFP8 GDN projections (nvidia/Qwen3.8-Flash-Next-FP4-Local). +//! +//! That checkpoint ships `in_proj_qkv`, `in_proj_z` and `out_proj` as MXFP8. +//! A layer built from them holds [`Mx8Gdn`] and nothing else for those two +//! projections: `ssm.in_proj_qkvz` is NULL and every BF16 / FP8 / NVFP4 copy +//! is `None`. Each dispatch site checks `self.mx8` first, so the bytes the +//! checkpoint ships are the bytes every forward reads; a site that forgot +//! the check would launch on a NULL weight and fault, not compute quietly +//! with another copy. + +use anyhow::{Result, ensure}; +use spark_runtime::gpu::DevicePtr; + +use super::Qwen3SsmLayer; +use crate::layer::ForwardContext; +use crate::layers::ops::{self, Mxfp8Proj, Mxfp8ProjKernels, Mxfp8Scratch}; + +/// The MXFP8 `[Q|K|V|Z]` input projection (rows of `in_proj_qkv` then +/// `in_proj_z`, sequential) and `out_proj`, with the kernels that read them +/// and the shared BF16 scratch prefill expands into. +#[derive(Debug, Clone, Copy)] +pub struct Mx8Gdn { + pub qkvz: Mxfp8Proj, + pub out: Mxfp8Proj, + pub kernels: Mxfp8ProjKernels, + pub scratch: Mxfp8Scratch, +} + +impl Qwen3SsmLayer { + /// Install the MXFP8 projections. Refused unless the layer carries no + /// other copy of either projection and its QKVZ output is sequential, + /// which is the layout `in_proj_qkv || in_proj_z` rows produce. + pub fn set_mx8_projections(&mut self, mx8: Mx8Gdn) -> Result<()> { + ensure!( + self.sequential_qkvz, + "MXFP8 GDN: the QKVZ projection must be sequential [Q|K|V|Z]" + ); + ensure!( + self.ssm.in_proj_qkvz.weight.is_null() + && self.out_proj_dense.is_none() + && self.qkvz_nvfp4.is_none() + && self.qkvz_fp8w.is_none() + && self.out_proj_fp8w.is_none() + && self.qkvz_fp8w_rowwise.is_none() + && self.out_proj_fp8w_rowwise.is_none() + && self.qkvz_q2.is_none(), + "MXFP8 GDN: the layer already holds another copy of its projections" + ); + let largest = mx8.qkvz.bf16_bytes().max(mx8.out.bf16_bytes()); + ensure!( + !mx8.scratch.ptr.is_null() && mx8.scratch.bytes >= largest, + "MXFP8 GDN: prefill scratch {} B is smaller than the {largest} B expansion", + mx8.scratch.bytes + ); + self.mx8 = Some(mx8); + Ok(()) + } + + /// `dst[m, qkvz] = act[m, hidden] @ W_qkvz^T`, sequential `[Q|K|V|Z]`. + pub(super) fn mx8_qkvz( + &self, + mx8: &Mx8Gdn, + ctx: &ForwardContext, + act: DevicePtr, + dst: DevicePtr, + m: usize, + stream: u64, + ) -> Result<()> { + ops::mxfp8_proj( + ctx.gpu, + &mx8.kernels, + &mx8.qkvz, + act, + dst, + m, + mx8.scratch, + stream, + ) + } + + /// `dst[m, hidden] = act[m, value_dim] @ W_out^T`. + pub(super) fn mx8_out( + &self, + mx8: &Mx8Gdn, + ctx: &ForwardContext, + act: DevicePtr, + dst: DevicePtr, + m: usize, + stream: u64, + ) -> Result<()> { + ops::mxfp8_proj( + ctx.gpu, + &mx8.kernels, + &mx8.out, + act, + dst, + m, + mx8.scratch, + stream, + ) + } +} diff --git a/crates/spark-model/src/layers/qwen3_ssm/mx8_tests.rs b/crates/spark-model/src/layers/qwen3_ssm/mx8_tests.rs new file mode 100644 index 000000000..87cf4df61 --- /dev/null +++ b/crates/spark-model/src/layers/qwen3_ssm/mx8_tests.rs @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Dispatch of native MXFP8 GDN projections (`mx8.rs`) on the mock backend. + +use super::mx8::Mx8Gdn; +use super::tests::{gpu_launch_grids, run_batched_verify}; +use super::*; +use crate::layers::ops::{Mxfp8Proj, Mxfp8ProjKernels, Mxfp8Scratch}; +use crate::weight_map::{Mxfp8Weight, WeightQuantFormat}; +use atlas_core::config::ModelConfig; +use spark_runtime::gpu::mock::MockGpuBackend; + +/// A GDN layer wired like `gdn_mx8::build`: every other projection slot +/// empty, the MXFP8 pair installed. +fn mx8_layer(gpu: &MockGpuBackend, config: &ModelConfig) -> Qwen3SsmLayer { + let mut layer = bare_layer(gpu, config); + layer.set_mx8_projections(mx8(gpu, config)).unwrap(); + layer +} + +fn bare_layer(gpu: &MockGpuBackend, config: &ModelConfig) -> Qwen3SsmLayer { + let h = config.hidden_size; + let value_dim = config.linear_num_value_heads * config.linear_value_head_dim; + let dw = |bytes: usize| DenseWeight { + weight: gpu.alloc(bytes).unwrap(), + }; + let ssm = SsmWeights { + in_proj_qkvz: DenseWeight { + weight: DevicePtr::NULL, + }, + in_proj_ba: dw(config.ssm_ba_size() * h * 2), + conv1d: dw( + (2 * config.linear_num_key_heads * config.linear_key_head_dim + value_dim) + * config.linear_conv_kernel_dim + * 2, + ), + a_log: dw(config.linear_num_value_heads * 4), + dt_bias: dw(config.linear_num_value_heads * 4), + norm: dw(config.linear_value_head_dim * 2), + out_proj: QuantizedWeight::null(), + }; + Qwen3SsmLayer::new_sequential( + dw(h * 2), + ssm, + dw(h * 2), + FfnComponent::None, + None, + None, + None, + config, + gpu, + ) + .unwrap() +} + +fn mx8(gpu: &MockGpuBackend, config: &ModelConfig) -> Mx8Gdn { + let h = config.hidden_size; + let value_dim = config.linear_num_value_heads * config.linear_value_head_dim; + let proj = |n: usize, k: usize| Mxfp8Proj { + w: Mxfp8Weight { + data: gpu.alloc(n * k).unwrap(), + scales: gpu.alloc(n * k / 32).unwrap(), + }, + n: n as u32, + k: k as u32, + }; + let qkvz = proj(config.ssm_qkvz_size(), h); + let out = proj(h, value_dim); + let bytes = qkvz.bf16_bytes().max(out.bf16_bytes()); + Mx8Gdn { + qkvz, + out, + kernels: Mxfp8ProjKernels::resolve(gpu).unwrap(), + scratch: Mxfp8Scratch { + ptr: gpu.alloc(bytes).unwrap(), + bytes, + }, + } +} + +/// A 7-row verify (two sequences, K=4 and K=3) reads each MXFP8 projection +/// with the weight-streaming GEMV — one launch of grid ceil(N/4) per +/// projection — and never expands a weight. +#[test] +fn verify_rows_stream_the_mxfp8_weights() { + let config = ModelConfig::qwen3_next_80b_nvfp4(); + let gpu = MockGpuBackend::new(); + let layer = mx8_layer(&gpu, &config); + run_batched_verify(&gpu, &config, &layer, &[4, 3]).unwrap(); + let seen = gpu_launch_grids(&gpu); + let qkvz = config.ssm_qkvz_size() as u32; + let h = config.hidden_size as u32; + for (what, n) in [("QKVZ", qkvz), ("out_proj", h)] { + assert!( + seen.contains(&([n / 4, 1, 1], [256, 1, 1])), + "{what} MXFP8 GEMV launch missing; grids seen: {seen:?}" + ); + } + let dequant_grid = (qkvz * h / 16).div_ceil(256); + assert!( + !seen.iter().any(|(g, _)| g[0] == dequant_grid), + "a verify window must not expand the weight: {seen:?}" + ); +} + +/// The MXFP8 pair must be the layer's only copy of its projections. +#[test] +fn a_second_projection_copy_is_refused() { + let config = ModelConfig::qwen3_next_80b_nvfp4(); + let gpu = MockGpuBackend::new(); + let h = config.hidden_size; + let mut layer = bare_layer(&gpu, &config); + layer.set_fp8_decode_weights( + Some(Fp8Weight { + weight: gpu.alloc(config.ssm_qkvz_size() * h).unwrap(), + row_scale: gpu.alloc(4).unwrap(), + n: config.ssm_qkvz_size() as u32, + k: h as u32, + scale_format: WeightQuantFormat::Fp8BlockScaled, + }), + None, + ); + let err = layer.set_mx8_projections(mx8(&gpu, &config)).unwrap_err(); + assert!(err.to_string().contains("another copy"), "{err:#}"); +} + +/// A scratch too small for the prefill expansion is refused at build. +#[test] +fn an_undersized_scratch_is_refused() { + let config = ModelConfig::qwen3_next_80b_nvfp4(); + let gpu = MockGpuBackend::new(); + let mut layer = bare_layer(&gpu, &config); + let mut m = mx8(&gpu, &config); + m.scratch.bytes = m.qkvz.bf16_bytes() - 2; + let err = layer.set_mx8_projections(m).unwrap_err(); + assert!(err.to_string().contains("prefill scratch"), "{err:#}"); +} diff --git a/crates/spark-model/src/layers/qwen3_ssm/ssm_decode_tail.rs b/crates/spark-model/src/layers/qwen3_ssm/ssm_decode_tail.rs index a9345aa79..696716442 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/ssm_decode_tail.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/ssm_decode_tail.rs @@ -90,7 +90,9 @@ impl Qwen3SsmLayer { let value_dim = ctx.config.linear_num_value_heads * ctx.config.linear_value_head_dim; // ── 8. Output projection: [value_dim → hidden_size] ── let out = ctx.buffers.moe_output(); - if let Some(ref fp8) = self.out_proj_fp8w { + if let Some(mx8) = self.mx8.as_ref() { + self.mx8_out(mx8, ctx, normed_out, out, 1, stream)?; + } else if let Some(ref fp8) = self.out_proj_fp8w { // `w8a16_gemv` consumes `[N/BS,K/BS] BF16` block scales — // the canonical Qwen FP8 release format. fp8.scale_format.expect( diff --git a/crates/spark-model/src/layers/qwen3_ssm/ssm_forward.rs b/crates/spark-model/src/layers/qwen3_ssm/ssm_forward.rs index 5d26e63bb..e948a2f15 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/ssm_forward.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/ssm_forward.rs @@ -43,7 +43,10 @@ impl Qwen3SsmLayer { let deinterleaved = ctx.buffers.ssm_deinterleaved(); let qkvz_size = ctx.config.ssm_qkvz_size() as u32; prof!("qkvz", { - if let Some(ref fp8) = self.qkvz_fp8w { + if let Some(mx8) = self.mx8.as_ref() { + // Native MXFP8 (FP4-Local): sequential, straight into place. + self.mx8_qkvz(mx8, ctx, normed, deinterleaved, 1, stream) + } else if let Some(ref fp8) = self.qkvz_fp8w { // FP8 native: w8a16_gemv + deinterleave (no fused QKVZ variant yet). // `w8a16_gemv` consumes `[N/BS,K/BS] BF16` block scales // (see `kernels/gb10/common/w8a16_gemv.cu` line 5). diff --git a/crates/spark-model/src/layers/qwen3_ssm/tests.rs b/crates/spark-model/src/layers/qwen3_ssm/tests.rs index bc9d009d1..13f53ec73 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/tests.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/tests.rs @@ -132,7 +132,7 @@ fn mk_state(gpu: &MockGpuBackend, layer: &Qwen3SsmLayer, n_inter: usize) -> SsmL /// Drive the batched verify body at ragged ks through `decode_verify_multi` /// (the verify_e.rs entry point) and return the result. -fn run_batched_verify( +pub(super) fn run_batched_verify( gpu: &MockGpuBackend, config: &ModelConfig, layer: &Qwen3SsmLayer, @@ -396,7 +396,7 @@ fn active_row_without_ple_carry_is_refused() { } /// Grid/block pairs of every launch recorded by the mock. -fn gpu_launch_grids(gpu: &MockGpuBackend) -> Vec<([u32; 3], [u32; 3])> { +pub(super) fn gpu_launch_grids(gpu: &MockGpuBackend) -> Vec<([u32; 3], [u32; 3])> { gpu.launches_snapshot() .into_iter() .map(|l| (l.grid, l.block)) diff --git a/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched/impl.rs b/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched/impl.rs index 706d5bdf0..3d3cf7e45 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched/impl.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched/impl.rs @@ -72,9 +72,12 @@ impl Qwen3SsmLayer { || (self.qkvz_fp8w.is_none() && self.dense_gemm_k.0 != 0))) || (self.qkvz_nvfp4.is_some() && ((self.w4a16_batchm.has_base() && n <= 16) || tc_wide_ok)); + // Native MXFP8 serves every row count through `exact_rows_proj`. + let qkvz_ok = qkvz_ok || self.mx8.is_some(); let out_ok = self.out_proj_fp8w.is_some() || self.out_proj_dense.is_some() - || self.qkvz_nvfp4.is_some(); + || self.qkvz_nvfp4.is_some() + || self.mx8.is_some(); // Tier-1c keep-packed Q2_0 has no batched packed GEMM; decline so the // per-seq fallback (`ssm_forward` → `q2_0_gemv_vec`) handles it. if n < 2 diff --git a/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_helper.rs b/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_helper.rs index cd39f1868..5e340fc8c 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_helper.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_helper.rs @@ -108,6 +108,10 @@ impl Qwen3SsmLayer { value_dim: usize, stream: u64, ) -> Result<()> { + // Native MXFP8 (FP4-Local): the layer's only out_proj. + if let Some(mx8) = self.mx8.as_ref() { + return self.mx8_out(mx8, ctx, normed_out_buf, out_proj_buf, k as usize, stream); + } let force_w8a8 = matches!(std::env::var("ATLAS_FP8_W8A8").ok().as_deref(), Some("1")); // PER-ROW FP8 from the checkpoint (`ATLAS_FP8_ROWWISE=1`), dequantised // once to BF16 — see the matching arm in `trait_prefill_proj.rs` for diff --git a/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_proj.rs b/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_proj.rs index 92f1fa292..e44c54eda 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_proj.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_proj.rs @@ -67,6 +67,10 @@ impl Qwen3SsmLayer { // GEMM. Bonsai is `sequential_qkvz`, so `proj_dst == deinterleaved` and // no post-deinterleave is needed. Highest priority (all other weight // slots are NULL on this path). + // Native MXFP8 (FP4-Local): the layer's only copy, sequential layout. + if let Some(mx8) = self.mx8.as_ref() { + return self.mx8_qkvz(mx8, ctx, normed, proj_dst, k as usize, stream); + } if self.qkvz_q2.is_some() { let scratch = ctx.buffers.q2_dequant_scratch(); let act_q8 = ctx.buffers.q2_act_q8(); diff --git a/crates/spark-model/src/weight_loader/qwen4_exp.rs b/crates/spark-model/src/weight_loader/qwen4_exp.rs index 0c6c585bd..0bfde92ee 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp.rs @@ -79,6 +79,7 @@ use crate::weight_map::{DenseWeight, MtpWeights, dense}; mod attach; mod ffn; mod gdn_fp8; +mod gdn_mx8; mod hc; mod mtp; mod ple; @@ -270,6 +271,7 @@ impl ModelWeightLoader for Qwen4ExpWeightLoader { // shards, and nothing said which arm spent it. Summed here and logged // once, so the answer is read rather than guessed. let (mut moe_bytes, mut arm_bytes, mut hc_bytes) = (0u64, 0u64, 0u64); + let mut mx8_scratch = None; let free_now = |g: &dyn GpuBackend| g.free_memory().unwrap_or(0) as u64; for i in 0..config.num_hidden_layers { @@ -289,20 +291,16 @@ impl ModelWeightLoader for Qwen4ExpWeightLoader { let post_attn_norm = ones_norm(h, gpu)?; let layer = match config.layer_types[i] { + // FP4-Local ships these MXFP8: served on those bytes. + LayerType::LinearAttention if gdn_mx8::present(store, &lp) => { + let scratch = gdn_mx8::scratch_for(&mut mx8_scratch, config, gpu)?; + gdn_mx8::build( + i, store, &lp, gpu, config, input_norm, post_attn_norm, ffn, scratch, + ) + .with_context(|| format!("qwen4_exp: GDN layer {i} (MXFP8)"))? + } + // BF16 as shipped, not requantized: see `GdnProjections::Bf16`. LayerType::LinearAttention if gdn == GdnProjections::Bf16 => { - // Keep the GDN projections BF16 instead of requantizing - // them to NVFP4 at load. - // - // Two reasons, and the second is the interesting one. - // (1) MEMORY: the requantization is where this model's - // build spends its 7.34 GB (152.8 MB/layer, measured — - // the MoE costs zero because its experts ship NVFP4 and - // upload straight through). - // (2) PRECISION: these tensors ship as BF16 in this - // checkpoint. Only the routed experts are quantized. So - // BF16 -> NVFP4 here is a lossy round trip we chose, not - // one the checkpoint forced, and it lands on the GDN - // projections of 36 of 48 layers. crate::weight_loader::qwen35::load_layers::linear_attn_arms::build_linear_attention_dense_bf16( i, store, &lp, gpu, variant, config, h, input_norm, post_attn_norm, ffn, diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/gdn_fp8.rs b/crates/spark-model/src/weight_loader/qwen4_exp/gdn_fp8.rs index c1b3c0687..ae821fefa 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/gdn_fp8.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/gdn_fp8.rs @@ -43,7 +43,12 @@ use crate::weight_map::{DenseWeight, Nvfp4Variant, free_loader_source}; /// The format of the GDN `in_proj_qkvz` / `out_proj` weights. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum GdnProjections { - /// BF16 as shipped. The default. + /// BF16 as shipped. The default, for two reasons. MEMORY: requantizing + /// is where this model's build spent 7.34 GB (152.8 MB a layer, + /// measured; the MoE costs nothing, its experts upload as shipped). + /// PRECISION: only the routed experts ship quantized, so BF16 -> NVFP4 + /// would be a lossy round trip chosen, not forced, on 36 of 48 layers. + /// (FP4-Local's MXFP8 GDN projections bypass this enum: `gdn_mx8`.) Bf16, /// Requantized to NVFP4 (`ATLAS_QWEN4EXP_BF16_GDN=0`). Nvfp4, @@ -94,7 +99,10 @@ impl GdnProjections { pub(super) fn describe(self) -> &'static str { match self { - Self::Bf16 => "BF16 as shipped (no runtime NVFP4 requantization)", + Self::Bf16 => { + "BF16 as shipped (no runtime NVFP4 requantization), or native MXFP8 where \ + the checkpoint ships MXFP8" + } Self::Nvfp4 => "requantized to NVFP4 (ATLAS_QWEN4EXP_BF16_GDN=0)", Self::Fp8Decode => { "BF16 + a 128x128 block-scaled FP8 copy for decode (ATLAS_QWEN4EXP_FP8_GDN=1)" diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/gdn_mx8.rs b/crates/spark-model/src/weight_loader/qwen4_exp/gdn_mx8.rs new file mode 100644 index 000000000..46f0704ca --- /dev/null +++ b/crates/spark-model/src/weight_loader/qwen4_exp/gdn_mx8.rs @@ -0,0 +1,194 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! GDN layers of nvidia/Qwen3.8-Flash-Next-FP4-Local, on their MXFP8 bytes. +//! +//! That checkpoint ships `in_proj_qkv`, `in_proj_z` and `out_proj` as MXFP8 +//! (E4M3 + one E8M0 scale per 32 inputs), trained that way by NVIDIA's QAD, +//! and `weight_map::modelopt_local` uploads them unchanged. This builds the +//! layer around those bytes: `in_proj_qkv || in_proj_z` rows (and their +//! scale rows) become the sequential `[Q|K|V|Z]` projection, `out_proj` is +//! used in place, and the layer reads both through `mxfp8_proj.cu`. No BF16 +//! copy exists; prefill expands one projection at a time into a scratch +//! shared by every layer. +//! +//! `in_proj_a` / `in_proj_b` (4 MB a layer in all) stay on the BF16 +//! interleaved β/α path the fused gate kernel reads. + +use anyhow::{Context, Result, ensure}; +use atlas_core::config::ModelConfig; +use spark_runtime::gpu::{DevicePtr, GpuBackend}; +use spark_runtime::weights::{WeightDtype, WeightStore}; + +use crate::layer::TransformerLayer; +use crate::layers::FfnComponent; +use crate::layers::Qwen3SsmLayer; +use crate::layers::ops::{Mxfp8Proj, Mxfp8ProjKernels, Mxfp8Scratch}; +use crate::layers::qwen3_ssm::mx8::Mx8Gdn; +use crate::tp_shard::TpGdnDims; +use crate::weight_map::{ + DenseWeight, Mxfp8Weight, QuantizedWeight, SsmWeights, dense_auto, dense_f32_safe, + dense_keep_f32, free_loader_source, interleave_ba, +}; + +const MX8_GROUP: usize = 32; + +/// True when layer `lp`'s GDN projections are MXFP8 in the store. +pub(super) fn present(store: &WeightStore, lp: &str) -> bool { + let w = format!("{lp}.linear_attn.in_proj_qkv"); + store + .get(&format!("{w}.weight")) + .is_ok_and(|t| t.dtype == WeightDtype::FP8E4M3) + && store.contains(&format!("{w}.weight_scale")) +} + +/// The weight and scale buffers of one MXFP8 module, shapes checked. +fn module(store: &WeightStore, name: &str, rows: usize, k: usize) -> Result<[DevicePtr; 2]> { + let w = store.get(&format!("{name}.weight"))?; + let s = store.get(&format!("{name}.weight_scale"))?; + ensure!( + w.dtype == WeightDtype::FP8E4M3 + && w.shape == [rows, k] + && s.dtype == WeightDtype::UInt8 + && s.shape == [rows, k / MX8_GROUP], + "{name}: expected MXFP8 [{rows}, {k}] + [{rows}, {}], got {:?} {:?} + {:?} {:?}", + k / MX8_GROUP, + w.dtype, + w.shape, + s.dtype, + s.shape + ); + Ok([w.ptr, s.ptr]) +} + +/// Copy `a` then `b` (each `rows * row_bytes`) into one new buffer. +fn concat( + gpu: &dyn GpuBackend, + (a, a_rows): (DevicePtr, usize), + (b, b_rows): (DevicePtr, usize), + row_bytes: usize, +) -> Result { + let buf = gpu.alloc((a_rows + b_rows) * row_bytes)?; + gpu.copy_d2d(a, buf, a_rows * row_bytes)?; + gpu.copy_d2d(b, buf.offset(a_rows * row_bytes), b_rows * row_bytes)?; + Ok(buf) +} + +/// The prefill scratch every MXFP8 GDN layer shares: sized for the larger +/// projection, allocated by the first layer that needs it. +pub(super) fn scratch_for( + shared: &mut Option, + config: &ModelConfig, + gpu: &dyn GpuBackend, +) -> Result { + let h = config.hidden_size; + let value_dim = config.linear_num_value_heads * config.linear_value_head_dim; + let bytes = (config.ssm_qkvz_size() * h).max(h * value_dim) * 2; + if let Some(s) = shared.filter(|s| s.bytes >= bytes) { + return Ok(s); + } + ensure!(shared.is_none(), "MXFP8 GDN scratch sized twice"); + let s = Mxfp8Scratch { + ptr: gpu.alloc(bytes)?, + bytes, + }; + tracing::info!( + "MXFP8 GDN: one {:.1} MB BF16 prefill scratch shared by every GDN layer", + bytes as f64 / 1e6 + ); + *shared = Some(s); + Ok(s) +} + +/// Build GDN layer `layer_idx` on its MXFP8 projections. +#[allow(clippy::too_many_arguments)] +pub(super) fn build( + layer_idx: usize, + store: &WeightStore, + lp: &str, + gpu: &dyn GpuBackend, + config: &ModelConfig, + input_norm: DenseWeight, + post_attn_norm: DenseWeight, + ffn: FfnComponent, + scratch: Mxfp8Scratch, +) -> Result> { + ensure!( + config.tp_world_size.max(1) == 1, + "MXFP8 GDN layer {layer_idx}: tensor parallelism is not supported yet" + ); + let h = config.hidden_size; + let dims = TpGdnDims::from_config(config); + let (qkv_rows, z_rows) = (dims.full_conv_dim(), dims.full_value_dim()); + let value_dim = config.linear_num_value_heads * config.linear_value_head_dim; + ensure!( + qkv_rows + z_rows == config.ssm_qkvz_size() && value_dim == z_rows, + "MXFP8 GDN layer {layer_idx}: projection rows {qkv_rows}+{z_rows} do not match \ + the config's QKVZ width {}", + config.ssm_qkvz_size() + ); + let la = format!("{lp}.linear_attn"); + let [qkv_w, qkv_s] = module(store, &format!("{la}.in_proj_qkv"), qkv_rows, h)?; + let [z_w, z_s] = module(store, &format!("{la}.in_proj_z"), z_rows, h)?; + let [out_w, out_s] = module(store, &format!("{la}.out_proj"), h, value_dim)?; + + let qkvz = Mxfp8Weight { + data: concat(gpu, (qkv_w, qkv_rows), (z_w, z_rows), h)?, + scales: concat(gpu, (qkv_s, qkv_rows), (z_s, z_rows), h / MX8_GROUP)?, + }; + for part in ["in_proj_qkv", "in_proj_z"] { + for t in ["weight", "weight_scale"] { + store.reclaim(gpu, &format!("{la}.{part}.{t}"))?; + } + } + + // β/α gates: BF16 (expanded at load), interleaved per key-head group. + let a = dense_auto(store, &format!("{la}.in_proj_a.weight"), gpu)?; + let b = dense_auto(store, &format!("{la}.in_proj_b.weight"), gpu)?; + let ba = interleave_ba(&a, &b, dims.full_nv, dims.full_nk, h, gpu)?; + free_loader_source(store, gpu, &format!("{la}.in_proj_a.weight"), a)?; + free_loader_source(store, gpu, &format!("{la}.in_proj_b.weight"), b)?; + + let ssm = SsmWeights { + in_proj_qkvz: DenseWeight { + weight: DevicePtr::NULL, + }, + in_proj_ba: ba, + conv1d: dense_auto(store, &format!("{la}.conv1d.weight"), gpu)?, + // A_log and dt_bias stay FP32: the decay gate amplifies BF16 error. + a_log: dense_keep_f32(store, &format!("{la}.A_log"), gpu)?, + dt_bias: dense_keep_f32(store, &format!("{la}.dt_bias"), gpu)?, + norm: dense_f32_safe(store, &format!("{la}.norm.weight"), gpu)?, + out_proj: QuantizedWeight::null(), + }; + let mut layer = Qwen3SsmLayer::new_sequential( + input_norm, + ssm, + post_attn_norm, + ffn, + None, + None, + None, + config, + gpu, + )?; + layer + .set_mx8_projections(Mx8Gdn { + qkvz: Mxfp8Proj { + w: qkvz, + n: config.ssm_qkvz_size() as u32, + k: h as u32, + }, + out: Mxfp8Proj { + w: Mxfp8Weight { + data: out_w, + scales: out_s, + }, + n: h as u32, + k: value_dim as u32, + }, + kernels: Mxfp8ProjKernels::resolve(gpu)?, + scratch, + }) + .with_context(|| format!("MXFP8 GDN layer {layer_idx}"))?; + Ok(Box::new(layer)) +} diff --git a/crates/spark-model/src/weight_map/modelopt_local.rs b/crates/spark-model/src/weight_map/modelopt_local.rs index 6595aa54e..014cf4bc3 100644 --- a/crates/spark-model/src/weight_map/modelopt_local.rs +++ b/crates/spark-model/src/weight_map/modelopt_local.rs @@ -9,8 +9,11 @@ //! * routed-expert NVFP4 block scales are CSF-compressed //! (`weight_scale.nvfp4_csf_{fixed,exceptions}`) — decoded losslessly into //! the plain E4M3 `weight_scale` every NVFP4 kernel reads; -//! * GDN / attention / QSA-indexer / shared-expert projections are MXFP8 — -//! expanded exactly to the BF16 `weight` the NVFP4 pack ships; +//! * GDN / attention / QSA-indexer / shared-expert projections are MXFP8. +//! Those with a native MXFP8 consumer ([`native_mxfp8`]: the GDN +//! projections) are uploaded byte for byte, `weight` E4M3 + `weight_scale` +//! E8M0, and read by `mxfp8_proj.cu`; the rest are expanded exactly to the +//! BF16 `weight` the NVFP4 pack ships; //! * vision MLP/merger projections are W4A16 NVFP4 — dequantized to the BF16 //! the vision tower reads. //! @@ -53,10 +56,24 @@ pub fn host_decode_modules(config: &ModelConfig) -> HashSet { .collect() } +/// MXFP8 modules a layer reads natively, so they are uploaded as shipped. +/// `ATLAS_FP4LOCAL_EXPAND_MXFP8=1` expands them to BF16 instead (A/B only: +/// both are exact, the native bytes are half the size). +pub fn native_mxfp8(module: &str) -> bool { + if std::env::var("ATLAS_FP4LOCAL_EXPAND_MXFP8").as_deref() == Ok("1") { + return false; + } + module.starts_with("model.language_model.layers.") + && ["in_proj_qkv", "in_proj_z", "out_proj"] + .iter() + .any(|p| module.ends_with(&format!(".linear_attn.{p}"))) +} + /// What [`materialize`] did, for the load log and for tests. #[derive(Debug, Default, PartialEq, Eq)] pub struct Materialized { pub csf_scales: usize, + pub mxfp8_native: usize, pub mxfp8_to_bf16: usize, pub nvfp4_to_bf16: usize, pub f32_norms_to_bf16: usize, @@ -215,6 +232,49 @@ fn decode_one_csf(store: &mut WeightStore, gpu: &dyn GpuBackend, base: &str) -> ) } +/// Upload one MXFP8 module unchanged: `weight` E4M3 `[n, k]` and +/// `weight_scale` E8M0 `[n, k / 32]`. `None` when the checkpoint lacks it. +fn upload_native_mxfp8( + store: &mut WeightStore, + gpu: &dyn GpuBackend, + module: &str, +) -> Result> { + let Some(w) = take(store, gpu, &format!("{module}.weight"))? else { + return Ok(None); + }; + let s = take(store, gpu, &format!("{module}.weight_scale"))? + .with_context(|| format!("{module}.weight has no weight_scale"))?; + ensure!( + w.dtype == WeightDtype::FP8E4M3 + && w.shape.len() == 2 + && w.shape[1] % 32 == 0 + && s.dtype == WeightDtype::UInt8 + && s.shape == [w.shape[0], w.shape[1] / 32], + "{module}: expected MXFP8 (E4M3 [n, k] + U8 [n, k/32]), got {:?} {:?} + {:?} {:?}", + w.dtype, + w.shape, + s.dtype, + s.shape + ); + let mut n = upload( + store, + gpu, + format!("{module}.weight"), + &w.bytes, + w.shape, + w.dtype, + )?; + n += upload( + store, + gpu, + format!("{module}.weight_scale"), + &s.bytes, + s.shape, + s.dtype, + )?; + Ok(Some(n)) +} + /// Expand one host-decoded module into a BF16 `.weight`. Returns /// `(was_mxfp8, bytes)`, or `None` when the checkpoint does not carry it. fn expand_module( @@ -282,6 +342,13 @@ pub fn materialize( let mut modules: Vec = host_decode_modules(config).into_iter().collect(); modules.sort(); for module in modules { + if native_mxfp8(&module) { + if let Some(n) = upload_native_mxfp8(store, gpu, &module)? { + done.mxfp8_native += 1; + done.uploaded_bytes += n; + } + continue; + } let Some((mx, n)) = expand_module(store, gpu, &module)? else { continue; }; @@ -299,10 +366,11 @@ pub fn materialize( done.f32_norms_to_bf16 = n; done.uploaded_bytes += b; tracing::info!( - "FP4-Local: decoded {} CSF scale tensors, expanded {} MXFP8 and {} NVFP4 \ - projections to BF16, rounded {} F32 norms to BF16 ({:.2} GiB uploaded; \ - compressed originals never resident)", + "FP4-Local: decoded {} CSF scale tensors, kept {} MXFP8 projections native, \ + expanded {} MXFP8 and {} NVFP4 projections to BF16, rounded {} F32 norms to \ + BF16 ({:.2} GiB uploaded; compressed originals never resident)", done.csf_scales, + done.mxfp8_native, done.mxfp8_to_bf16, done.nvfp4_to_bf16, done.f32_norms_to_bf16, diff --git a/crates/spark-model/src/weight_map/modelopt_local_tests.rs b/crates/spark-model/src/weight_map/modelopt_local_tests.rs index 511b0a7d9..00750f455 100644 --- a/crates/spark-model/src/weight_map/modelopt_local_tests.rs +++ b/crates/spark-model/src/weight_map/modelopt_local_tests.rs @@ -10,6 +10,7 @@ const VIS: &str = "model.visual.blocks.0.mlp.linear_fc1"; const HC_NORM: &str = "model.language_model.layers.0.attn_hyper_connection.hc_norm.weight"; const PLE_NORM: &str = "model.language_model.layers.1.ple.norm_conv.weight"; const EXPERT: &str = "model.language_model.layers.0.mlp.experts.3.down_proj.weight_scale"; +const GDN_OUT: &str = "model.language_model.layers.0.linear_attn.out_proj"; fn spec(algo: &str, group: usize) -> QuantLayerSpec { QuantLayerSpec { @@ -28,6 +29,7 @@ fn fp4_local_config() -> ModelConfig { spec("NVFP4", 16), ); map.insert(ATTN.into(), spec("MXFP8", 32)); + map.insert(GDN_OUT.into(), spec("MXFP8", 32)); map.insert(VIS.into(), spec("W4A16_NVFP4", 16)); map.insert( "model.language_model.layers.1.ple.ple_embedding.ngram_embedding".into(), @@ -52,7 +54,7 @@ fn fp4_local_config() -> ModelConfig { #[test] fn host_decode_modules_pick_mxfp8_and_vision_nvfp4_only() { let got = host_decode_modules(&fp4_local_config()); - let want: HashSet = [ATTN, VIS].iter().map(|s| s.to_string()).collect(); + let want: HashSet = [ATTN, GDN_OUT, VIS].iter().map(|s| s.to_string()).collect(); assert_eq!(got, want); assert!(host_decode_modules(&ModelConfig::qwen3_next_80b_nvfp4()).is_empty()); } @@ -218,11 +220,12 @@ fn materialize_expands_every_fp4_local_format() { assert_eq!( ( done.csf_scales, + done.mxfp8_native, done.mxfp8_to_bf16, done.nvfp4_to_bf16, done.f32_norms_to_bf16 ), - (1, 1, 1, 2) + (1, 0, 1, 1, 2) ); let norm = store.get(HC_NORM).unwrap(); assert_eq!(norm.dtype, WeightDtype::BF16); @@ -270,6 +273,75 @@ fn materialize_expands_every_fp4_local_format() { } } +/// The GDN projections have a native MXFP8 kernel: their bytes are uploaded +/// exactly as shipped, weight and scale, and nothing is widened. +#[test] +fn gdn_mxfp8_stays_native() { + assert!(native_mxfp8(GDN_OUT)); + assert!(native_mxfp8( + "model.language_model.layers.7.linear_attn.in_proj_qkv" + )); + assert!(!native_mxfp8(ATTN)); + assert!(!native_mxfp8( + "model.language_model.layers.7.linear_attn.in_proj_a" + )); + + let gpu = MockGpuBackend::new(); + let mut store = WeightStore::empty(); + let mut shard = Shard::new("gdn"); + let weight: Vec = (0..64u8).collect(); + shard.add( + &format!("{GDN_OUT}.weight"), + &weight, + &[2, 32], + WeightDtype::FP8E4M3, + ); + shard.add( + &format!("{GDN_OUT}.weight_scale"), + &[120, 125], + &[2, 1], + WeightDtype::UInt8, + ); + shard.defer_into(&mut store); + let done = materialize(&mut store, &fp4_local_config(), &gpu).unwrap(); + assert_eq!((done.mxfp8_native, done.mxfp8_to_bf16), (1, 0)); + let w = store.get(&format!("{GDN_OUT}.weight")).unwrap(); + assert_eq!( + (w.dtype, w.shape.clone()), + (WeightDtype::FP8E4M3, vec![2, 32]) + ); + assert_eq!(gpu.read_alloc(w.ptr).unwrap(), weight); + let sc = store.get(&format!("{GDN_OUT}.weight_scale")).unwrap(); + assert_eq!( + (sc.dtype, sc.shape.clone()), + (WeightDtype::UInt8, vec![2, 1]) + ); + assert_eq!(gpu.read_alloc(sc.ptr).unwrap(), [120, 125]); +} + +/// A malformed MXFP8 scale is refused rather than uploaded. +#[test] +fn gdn_mxfp8_with_wrong_scale_shape_is_refused() { + let gpu = MockGpuBackend::new(); + let mut store = WeightStore::empty(); + let mut shard = Shard::new("gdnbad"); + shard.add( + &format!("{GDN_OUT}.weight"), + &[0; 64], + &[2, 32], + WeightDtype::FP8E4M3, + ); + shard.add( + &format!("{GDN_OUT}.weight_scale"), + &[0; 4], + &[2, 2], + WeightDtype::UInt8, + ); + shard.defer_into(&mut store); + let err = materialize(&mut store, &fp4_local_config(), &gpu).unwrap_err(); + assert!(format!("{err:#}").contains("expected MXFP8"), "{err:#}"); +} + /// A CSF fixed stream with no exceptions partner is a broken checkpoint, not /// something to serve with a half-decoded scale. #[test] diff --git a/kernels/gb10/qwen3.8-flash-next/nvfp4/mxfp8_proj.cu b/kernels/gb10/qwen3.8-flash-next/nvfp4/mxfp8_proj.cu new file mode 100644 index 000000000..655f378ce --- /dev/null +++ b/kernels/gb10/qwen3.8-flash-next/nvfp4/mxfp8_proj.cu @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +// Native MXFP8 projections for nvidia/Qwen3.8-Flash-Next-FP4-Local. +// +// MXFP8 (OCP microscaling): B is [N, K] E4M3 bytes and S is [N, K/32] E8M0 +// bytes, one power-of-two scale 2^(S-127) per 32 consecutive inputs of a row. +// The checkpoint ships the GDN, attention, indexer and shared-expert +// projections this way; these kernels read those bytes as they are on disk. +// +// mxfp8_proj_gemv C[n] = sum_k A[k] * W[n,k] (M = 1) +// mxfp8_proj_batchm C[t, n] = sum_k A[t, k] * W[n,k] (M <= 8) +// mxfp8_proj_dequant D[n, k] = W[n,k] as BF16 (exact) +// +// with W[n,k] = e4m3(B[n,k]) * 2^(S[n, k/32] - 127). +// +// Both GEMVs run the same templated body, so every row of `mxfp8_proj_batchm` is +// accumulated in exactly the order `mxfp8_proj_gemv` uses: a verify row and the +// sequential decode of the same token produce identical bits. +// +// Accumulation: each thread owns 16-input chunks (half an MXFP8 block). It +// sums the 16 products a * e4m3 in a fixed sequential order, scales that +// partial by the block's power of two (exact barring under/overflow), and +// adds it into an FP32 accumulator. Lanes then reduce by warp shuffle and the +// two warps of an output through shared memory. +// +// The dequant kernel feeds prefill: it expands one weight into a BF16 scratch +// that cuBLASLt multiplies. Every E4M3 value times a power of two in range is +// exactly representable in BF16, so the scratch is the checkpoint's weight, +// not an approximation of it. +// +// Grid (GEMVs): (ceil(N / 4)), Block: 256 — four outputs, 64 threads each. +// Grid (dequant): (ceil(N * K / 16 / 256)), Block: 256. +// Requires K % 32 == 0 and 16-byte aligned A, B, D rows. + +#include + +#define MX8_BLOCK_THREADS 256 +#define MX8_N_PER_BLOCK 4 +#define MX8_THREADS_PER_OUT (MX8_BLOCK_THREADS / MX8_N_PER_BLOCK) +#define MX8_GROUP 32 +#define MX8_MAX_M 8 + +// E4M3 byte -> f32. NaN (0x7F / 0xFF) maps to 0, as in the W8A16 tables. +__device__ __forceinline__ float mx8_e4m3_to_f32(unsigned int v) { + const unsigned int e = (v >> 3) & 0xF; + const unsigned int m = v & 0x7; + float r; + if (e == 0) { + r = (float)m * 0.001953125f; // subnormal: m * 2^-9 + } else if (e == 15 && m == 7) { + r = 0.0f; + } else { + r = __int_as_float(((e + 120) << 23) | (m << 20)); + } + return (v & 0x80) ? -r : r; +} + +// E8M0 byte -> 2^(e - 127). 0xFF (NaN) maps to 0; 0 is the subnormal 2^-127. +__device__ __forceinline__ float mx8_e8m0_to_f32(unsigned int e) { + if (e == 0xFF) return 0.0f; + if (e == 0) return __int_as_float(0x00400000); + return __int_as_float(e << 23); +} + +__device__ __forceinline__ float mx8_bf16_lo(unsigned int x) { + return __uint_as_float(x << 16); +} + +__device__ __forceinline__ float mx8_bf16_hi(unsigned int x) { + return __uint_as_float(x & 0xFFFF0000u); +} + +// Sum of the 16 products of one weight chunk and one activation chunk, in +// input order. +__device__ __forceinline__ float mx8_chunk_dot( + const float* __restrict__ lut, uint4 w, uint4 a0, uint4 a1 +) { + const unsigned int wv[4] = {w.x, w.y, w.z, w.w}; + const unsigned int av[8] = {a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w}; + float p = 0.0f; + #pragma unroll + for (int i = 0; i < 4; i++) { + const unsigned int w32 = wv[i]; + const unsigned int lo = av[2 * i]; + const unsigned int hi = av[2 * i + 1]; + p += mx8_bf16_lo(lo) * lut[w32 & 0xFF]; + p += mx8_bf16_hi(lo) * lut[(w32 >> 8) & 0xFF]; + p += mx8_bf16_lo(hi) * lut[(w32 >> 16) & 0xFF]; + p += mx8_bf16_hi(hi) * lut[w32 >> 24]; + } + return p; +} + +template +__device__ __forceinline__ void mx8_gemv_rows( + const __nv_bfloat16* __restrict__ A, + const unsigned char* __restrict__ B, + const unsigned char* __restrict__ S, + __nv_bfloat16* __restrict__ C, + unsigned int N, + unsigned int K +) { + __shared__ float s_lut[256]; + __shared__ float s_red[MX8_N_PER_BLOCK * 2 * M]; + s_lut[threadIdx.x] = mx8_e4m3_to_f32(threadIdx.x); + __syncthreads(); + + const unsigned int local_out = threadIdx.x / MX8_THREADS_PER_OUT; + const unsigned int lane = threadIdx.x % MX8_THREADS_PER_OUT; + const unsigned int n = blockIdx.x * MX8_N_PER_BLOCK + local_out; + const bool live = n < N; + + float acc[M]; + #pragma unroll + for (int t = 0; t < M; t++) acc[t] = 0.0f; + + if (live) { + const unsigned int k16s = K / 16; + const uint4* wrow = (const uint4*)(B + (unsigned long long)n * K); + const unsigned char* srow = S + (unsigned long long)n * (K / MX8_GROUP); + for (unsigned int c = lane; c < k16s; c += MX8_THREADS_PER_OUT) { + const uint4 w = wrow[c]; + const float scale = mx8_e8m0_to_f32(srow[c >> 1]); + #pragma unroll + for (int t = 0; t < M; t++) { + const uint4* arow = (const uint4*)(A + (unsigned long long)t * K); + const float p = mx8_chunk_dot(s_lut, w, arow[2 * c], arow[2 * c + 1]); + acc[t] += p * scale; + } + } + } + + const unsigned int warp_in_out = lane / 32; + #pragma unroll + for (int t = 0; t < M; t++) { + float v = acc[t]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xFFFFFFFF, v, off); + if ((lane & 31) == 0) s_red[(local_out * 2 + warp_in_out) * M + t] = v; + } + __syncthreads(); + if (live && lane < M) { + const unsigned int t = lane; + const float r = s_red[(local_out * 2) * M + t] + s_red[(local_out * 2 + 1) * M + t]; + C[(unsigned long long)t * N + n] = __float2bfloat16_rn(r); + } +} + +extern "C" __global__ void __launch_bounds__(MX8_BLOCK_THREADS) mxfp8_proj_gemv( + const __nv_bfloat16* __restrict__ A, // [1, K] + const unsigned char* __restrict__ B, // [N, K] E4M3 + const unsigned char* __restrict__ S, // [N, K/32] E8M0 + __nv_bfloat16* __restrict__ C, // [1, N] + unsigned int N, + unsigned int K +) { + mx8_gemv_rows<1>(A, B, S, C, N, K); +} + +extern "C" __global__ void __launch_bounds__(MX8_BLOCK_THREADS) mxfp8_proj_batchm( + const __nv_bfloat16* __restrict__ A, // [M, K] + const unsigned char* __restrict__ B, // [N, K] E4M3 + const unsigned char* __restrict__ S, // [N, K/32] E8M0 + __nv_bfloat16* __restrict__ C, // [M, N] + unsigned int M, + unsigned int N, + unsigned int K +) { + switch (M) { + case 1: mx8_gemv_rows<1>(A, B, S, C, N, K); break; + case 2: mx8_gemv_rows<2>(A, B, S, C, N, K); break; + case 3: mx8_gemv_rows<3>(A, B, S, C, N, K); break; + case 4: mx8_gemv_rows<4>(A, B, S, C, N, K); break; + case 5: mx8_gemv_rows<5>(A, B, S, C, N, K); break; + case 6: mx8_gemv_rows<6>(A, B, S, C, N, K); break; + case 7: mx8_gemv_rows<7>(A, B, S, C, N, K); break; + case 8: mx8_gemv_rows<8>(A, B, S, C, N, K); break; + default: break; // the host refuses M outside 1..=8 + } +} + +extern "C" __global__ void __launch_bounds__(MX8_BLOCK_THREADS) mxfp8_proj_dequant( + const unsigned char* __restrict__ B, // [N, K] E4M3 + const unsigned char* __restrict__ S, // [N, K/32] E8M0 + __nv_bfloat16* __restrict__ D, // [N, K] BF16 + unsigned int N, + unsigned int K +) { + const unsigned long long chunk = (unsigned long long)blockIdx.x * blockDim.x + threadIdx.x; + const unsigned long long chunks = (unsigned long long)N * (K / 16); + if (chunk >= chunks) return; + const unsigned int k16s = K / 16; + const unsigned long long n = chunk / k16s; + const unsigned int c = (unsigned int)(chunk % k16s); + const float scale = mx8_e8m0_to_f32(S[n * (K / MX8_GROUP) + (c >> 1)]); + const uint4 w = ((const uint4*)B)[chunk]; + const unsigned int wv[4] = {w.x, w.y, w.z, w.w}; + unsigned int out[8]; + #pragma unroll + for (int i = 0; i < 4; i++) { + #pragma unroll + for (int j = 0; j < 2; j++) { + const unsigned int b0 = (wv[i] >> (16 * j)) & 0xFF; + const unsigned int b1 = (wv[i] >> (16 * j + 8)) & 0xFF; + const __nv_bfloat16 v0 = __float2bfloat16_rn(mx8_e4m3_to_f32(b0) * scale); + const __nv_bfloat16 v1 = __float2bfloat16_rn(mx8_e4m3_to_f32(b1) * scale); + out[2 * i + j] = (unsigned int)__bfloat16_as_ushort(v0) + | ((unsigned int)__bfloat16_as_ushort(v1) << 16); + } + } + uint4* dst = (uint4*)(D + chunk * 16); + dst[0] = make_uint4(out[0], out[1], out[2], out[3]); + dst[1] = make_uint4(out[4], out[5], out[6], out[7]); +} From db6ffc008fc6fa056727867e9cd198be121cc22b Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:57:20 -0400 Subject: [PATCH 06/26] feat(qwen4_exp): serve FP4-Local's attention q/k/v/o on native MXFP8 FP4-Local ships the 12 full-attention layers' q (with its interleaved gate rows), k, v and o as MXFP8. The generic builder expanded them to BF16 and then requantized BF16 -> NVFP4 for decode: a second, lossy quantization of weights NVIDIA's QAD trained in MXFP8, plus a BF16 copy for prefill. - QuantWeight::Mxfp8(Mxfp8Proj); Mxfp8Proj moves to weight_map so the variant carries its geometry. quant_gemv/quant_gemm refuse it by name. - Qwen3AttentionLayer::set_mx8_projections installs q/k/v/o as the only copies (BF16 slots NULL, no NVFP4/FP8 twin) with the kernels and the shared prefill scratch; try_mx8 is checked first at every projection site: serial decode Q (gated + ungated) / K,V / O, the multi-sequence and verify path (qkv_mx8.rs: one weight pass per 8 rows, then the per-row scatter, as qkv_fp8.rs does) and its O, and the three prefill projections through try_unpacked_prefill (MXFP8, then packed Q2_0). - ops::mxfp8_proj: up to 32 rows run as 8-row GEMV chunks (bit-identical to single-token decode per row), more rows expand + bf16_gemm. - qwen4_exp::attn_mx8 builds the layer; mx8_common holds the checked weight/scale lookup and one scratch sized for the largest projection of either layer kind. native_mxfp8() now covers self_attn.{q,k,v,o}_proj. Measured on GB10 (sparky), --speculative --num-drafts 1: pre-KV 95.0 -> 93.4 GB; decode 24.5-26.3 tok/s (the MXFP8 q/k/v/o read 0.6 GB a token vs 0.3 GB for the lossy NVFP4 requant); coherence prompts, tool call, and the ~4.4k-token passcode recall all correct. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01L4aUGgtiZX94pLr63HXQzn --- .../spark-model/src/layers/ops/mxfp8_proj.rs | 97 ++++++++------ .../src/layers/ops/quant_dispatch.rs | 6 + .../decode/attention_forward.rs | 20 ++- .../decode/attention_forward_kv.rs | 7 ++ .../decode/attention_forward_oproj.rs | 4 + .../src/layers/qwen3_attention/init.rs | 1 + .../src/layers/qwen3_attention/mod.rs | 1 + .../src/layers/qwen3_attention/mx8.rs | 113 +++++++++++++++++ .../qwen3_attention/prefill/cache_skip_qkv.rs | 4 +- .../qwen3_attention/prefill/paged_oproj.rs | 4 +- .../qwen3_attention/prefill/paged_qkv.rs | 4 +- .../trait_impl/multi_seq/attn/o_proj.rs | 5 +- .../trait_impl/multi_seq/mod.rs | 3 + .../trait_impl/multi_seq/mx8_tests.rs | 118 ++++++++++++++++++ .../trait_impl/multi_seq/qkv.rs | 2 + .../trait_impl/multi_seq/qkv_mx8.rs | 74 +++++++++++ .../trait_impl/multi_seq/tests.rs | 2 +- .../src/layers/qwen3_attention/types_layer.rs | 2 + .../src/layers/qwen3_ssm/exact_rows.rs | 24 ++-- .../src/weight_loader/qwen4_exp.rs | 37 +++--- .../src/weight_loader/qwen4_exp/attn_mx8.rs | 116 +++++++++++++++++ .../src/weight_loader/qwen4_exp/gdn_fp8.rs | 23 ++++ .../src/weight_loader/qwen4_exp/gdn_mx8.rs | 52 +------- .../src/weight_loader/qwen4_exp/mx8_common.rs | 76 +++++++++++ crates/spark-model/src/weight_map.rs | 2 +- crates/spark-model/src/weight_map/expert.rs | 14 +++ .../src/weight_map/modelopt_local.rs | 17 ++- .../src/weight_map/modelopt_local_tests.rs | 22 ++-- crates/spark-model/src/weight_map/mxfp8.rs | 17 +++ 29 files changed, 725 insertions(+), 142 deletions(-) create mode 100644 crates/spark-model/src/layers/qwen3_attention/mx8.rs create mode 100644 crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mx8_tests.rs create mode 100644 crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/qkv_mx8.rs create mode 100644 crates/spark-model/src/weight_loader/qwen4_exp/attn_mx8.rs create mode 100644 crates/spark-model/src/weight_loader/qwen4_exp/mx8_common.rs diff --git a/crates/spark-model/src/layers/ops/mxfp8_proj.rs b/crates/spark-model/src/layers/ops/mxfp8_proj.rs index 652dc62d2..f2693e116 100644 --- a/crates/spark-model/src/layers/ops/mxfp8_proj.rs +++ b/crates/spark-model/src/layers/ops/mxfp8_proj.rs @@ -7,8 +7,9 @@ //! they are: //! //! * 1 row: `mxfp8_proj_gemv`, one pass over the weight; -//! * 2..=8 rows (verify windows, small batches): `mxfp8_proj_batchm`, still -//! one pass, every row bit-identical to the 1-row kernel; +//! * 2..=32 rows (verify windows, small batches, short prefill chunks): +//! `mxfp8_proj_batchm` over 8-row chunks, one weight pass a chunk, every +//! row bit-identical to the 1-row kernel; //! * more rows (prefill): `mxfp8_proj_dequant` expands the weight into a BF16 //! scratch, exactly (E4M3 times a power of two is a BF16 value), and //! [`super::bf16_gemm`] multiplies by it. The scratch is transient and @@ -18,11 +19,15 @@ use anyhow::{Result, ensure}; use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle}; use spark_runtime::kernel_args::{KernelLaunch, div_ceil}; -use crate::weight_map::Mxfp8Weight; +pub use crate::weight_map::Mxfp8Proj; -/// Widest row count the weight-streaming GEMVs serve. +/// Widest row count one weight-streaming GEMV launch serves. pub const MXFP8_PROJ_MAX_GEMV_ROWS: usize = 8; +/// Up to this many rows run as chunked GEMVs; beyond it the expansion plus a +/// tensor-core GEMM reads fewer bytes than re-streaming the weight per chunk. +pub const MXFP8_PROJ_MAX_CHUNKED_ROWS: usize = 32; + /// Kernel handles of `mxfp8_proj.cu`. #[derive(Debug, Clone, Copy)] pub struct Mxfp8ProjKernels { @@ -43,23 +48,6 @@ impl Mxfp8ProjKernels { } } -/// One MXFP8 projection and its geometry. -#[derive(Debug, Clone, Copy)] -pub struct Mxfp8Proj { - pub w: Mxfp8Weight, - /// Output features (weight rows). - pub n: u32, - /// Input features; a multiple of 32. - pub k: u32, -} - -impl Mxfp8Proj { - /// Bytes of the BF16 expansion the prefill path writes into scratch. - pub fn bf16_bytes(&self) -> usize { - self.n as usize * self.k as usize * 2 - } -} - /// A BF16 scratch for the prefill expansion: device pointer and capacity. #[derive(Debug, Clone, Copy)] pub struct Mxfp8Scratch { @@ -80,36 +68,75 @@ pub fn mxfp8_proj( scratch: Mxfp8Scratch, stream: u64, ) -> Result<()> { - let Mxfp8Proj { w, n, k } = *proj; + let Mxfp8Proj { n, k, .. } = *proj; ensure!( m > 0 && n > 0 && k.is_multiple_of(32), "mxfp8_proj: m={m} n={n} k={k} (k must be a multiple of 32)" ); if m == 1 { - return KernelLaunch::new(gpu, kernels.gemv) - .grid([div_ceil(n, 4), 1, 1]) - .block([256, 1, 1]) - .arg_ptr(act) - .arg_ptr(w.data) - .arg_ptr(w.scales) - .arg_ptr(out) - .arg_u32(n) - .arg_u32(k) - .launch(stream); + return gemv_rows(gpu, kernels, proj, act, out, 1, stream); + } + if m <= MXFP8_PROJ_MAX_CHUNKED_ROWS { + let (in_row, out_row) = (k as usize * 2, n as usize * 2); + for first in (0..m).step_by(MXFP8_PROJ_MAX_GEMV_ROWS) { + let rows = (m - first).min(MXFP8_PROJ_MAX_GEMV_ROWS); + let (a, o) = (act.offset(first * in_row), out.offset(first * out_row)); + gemv_rows(gpu, kernels, proj, a, o, rows, stream)?; + } + return Ok(()); } - if m <= MXFP8_PROJ_MAX_GEMV_ROWS { - return KernelLaunch::new(gpu, kernels.batchm) + expand_and_gemm(gpu, kernels, proj, act, out, m, scratch, stream) +} + +/// One weight pass for `m <= 8` contiguous rows. +fn gemv_rows( + gpu: &dyn GpuBackend, + kernels: &Mxfp8ProjKernels, + proj: &Mxfp8Proj, + act: DevicePtr, + out: DevicePtr, + m: usize, + stream: u64, +) -> Result<()> { + let Mxfp8Proj { w, n, k } = *proj; + if m == 1 { + return KernelLaunch::new(gpu, kernels.gemv) .grid([div_ceil(n, 4), 1, 1]) .block([256, 1, 1]) .arg_ptr(act) .arg_ptr(w.data) .arg_ptr(w.scales) .arg_ptr(out) - .arg_u32(m as u32) .arg_u32(n) .arg_u32(k) .launch(stream); } + KernelLaunch::new(gpu, kernels.batchm) + .grid([div_ceil(n, 4), 1, 1]) + .block([256, 1, 1]) + .arg_ptr(act) + .arg_ptr(w.data) + .arg_ptr(w.scales) + .arg_ptr(out) + .arg_u32(m as u32) + .arg_u32(n) + .arg_u32(k) + .launch(stream) +} + +/// Exact BF16 expansion into `scratch`, then a tensor-core GEMM. +#[allow(clippy::too_many_arguments)] +fn expand_and_gemm( + gpu: &dyn GpuBackend, + kernels: &Mxfp8ProjKernels, + proj: &Mxfp8Proj, + act: DevicePtr, + out: DevicePtr, + m: usize, + scratch: Mxfp8Scratch, + stream: u64, +) -> Result<()> { + let Mxfp8Proj { w, n, k } = *proj; ensure!( !scratch.ptr.is_null() && scratch.bytes >= proj.bf16_bytes(), "mxfp8_proj: {m} rows need a {} B BF16 scratch, have {} B", diff --git a/crates/spark-model/src/layers/ops/quant_dispatch.rs b/crates/spark-model/src/layers/ops/quant_dispatch.rs index d080272bf..7883622ff 100644 --- a/crates/spark-model/src/layers/ops/quant_dispatch.rs +++ b/crates/spark-model/src/layers/ops/quant_dispatch.rs @@ -58,6 +58,9 @@ pub fn quant_gemv( QuantWeight::Exl3(_) => anyhow::bail!( "quant_gemv: EXL3 is not routed through the generic dispatcher yet (reconstruct + GEMM lands in M3)" ), + QuantWeight::Mxfp8(_) => { + anyhow::bail!("quant_gemv: MXFP8 projections dispatch through ops::mxfp8_proj") + } } } @@ -101,6 +104,9 @@ pub fn quant_gemm( QuantWeight::Exl3(_) => anyhow::bail!( "quant_gemm: EXL3 is not routed through the generic dispatcher yet (reconstruct + GEMM lands in M3)" ), + QuantWeight::Mxfp8(_) => { + anyhow::bail!("quant_gemm: MXFP8 projections dispatch through ops::mxfp8_proj") + } } } 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 3a5e711c9..c584cf9e2 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 @@ -96,7 +96,21 @@ impl Qwen3AttentionLayer { if self.gated { // Q+Gate projection with inline deinterleave (output is [Q_all | Gate_all]) - if let Some(q2) = self.q_weight.as_ref().and_then(|w| w.as_packed_q2()) { + if let Some(r) = self.try_mx8(ctx, self.q_weight.as_ref(), normed, q_out, 1, stream) { + // Native MXFP8 (FP4-Local): raw interleaved [Q|gate], then split. + r?; + self.apply_q_lora(ctx, normed, q_out, stream)?; + ops::deinterleave_qg( + ctx.gpu, + self.deinterleave_qg_k, + q_out, + 1, + nq, + hd, + nq * hd * 2, + stream, + )?; + } else if let Some(q2) = self.q_weight.as_ref().and_then(|w| w.as_packed_q2()) { // Keep-packed Q2_0 (Tier-1c): 2-bit GEMV → gated [Q|Gate], then // the same deinterleave the dense fallback uses. ops::q2_0_gemv_vec(ctx.gpu, self.q2_0_gemv_k, normed, q2, q_out, stream)?; @@ -201,7 +215,9 @@ impl Qwen3AttentionLayer { } } else { // Ungated: Q projection only (no gate) - if let Some(q2) = self.q_weight.as_ref().and_then(|w| w.as_packed_q2()) { + if let Some(r) = self.try_mx8(ctx, self.q_weight.as_ref(), normed, q_out, 1, stream) { + r?; + } else if let Some(q2) = self.q_weight.as_ref().and_then(|w| w.as_packed_q2()) { ops::q2_0_gemv_vec(ctx.gpu, self.q2_0_gemv_k, normed, q2, q_out, stream)?; } else if let Some(fp8) = self.q_weight.as_ref().and_then(|w| w.as_fp8()) { ops::w8a16_gemv( diff --git a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_kv.rs b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_kv.rs index 5bec59b1e..d40b8e6a8 100644 --- a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_kv.rs +++ b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_kv.rs @@ -30,6 +30,13 @@ impl Qwen3AttentionLayer { return Ok(()); } + if let Some(r) = self.try_mx8(ctx, self.k_weight.as_ref(), normed, k_out, 1, stream) { + r?; + return self + .try_mx8(ctx, self.v_weight.as_ref(), normed, v_out, 1, stream) + .unwrap_or_else(|| anyhow::bail!("MXFP8 k_proj without an MXFP8 v_proj")); + } + if let (Some(k_q2), Some(v_q2)) = ( self.k_weight.as_ref().and_then(|w| w.as_packed_q2()), self.v_weight.as_ref().and_then(|w| w.as_packed_q2()), diff --git a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_oproj.rs b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_oproj.rs index 41e1ad6fb..3260c76ff 100644 --- a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_oproj.rs +++ b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_oproj.rs @@ -50,6 +50,10 @@ impl Qwen3AttentionLayer { stream, )?; } + } else if let Some(r) = + self.try_mx8(ctx, self.o_weight.as_ref(), attn_out, o_out, 1, stream) + { + r?; } else if let Some(o_bf16) = self.o_dense_bf16.as_ref() { ops::dense_gemv( ctx.gpu, diff --git a/crates/spark-model/src/layers/qwen3_attention/init.rs b/crates/spark-model/src/layers/qwen3_attention/init.rs index 1f8f4e371..60101d913 100644 --- a/crates/spark-model/src/layers/qwen3_attention/init.rs +++ b/crates/spark-model/src/layers/qwen3_attention/init.rs @@ -177,6 +177,7 @@ impl Qwen3AttentionLayer { v_weight: v_nvfp4.map(QuantWeight::Nvfp4), o_weight: None, o_dense_bf16: None, + mx8_rt: None, mla: None, mla_mx: Vec::new(), mla_q4: Vec::new(), diff --git a/crates/spark-model/src/layers/qwen3_attention/mod.rs b/crates/spark-model/src/layers/qwen3_attention/mod.rs index fab513320..62d5e340a 100644 --- a/crates/spark-model/src/layers/qwen3_attention/mod.rs +++ b/crates/spark-model/src/layers/qwen3_attention/mod.rs @@ -30,6 +30,7 @@ mod init_glm; mod init_independent; mod init_kernel_dispatch; mod kernel_requirements; +pub(crate) mod mx8; mod op_dump; // `innerq_driver` calls the CUDA Driver API directly via `atlas_core::registry`, // which is itself gated on the `cuda` feature. Mirror that gate here so the diff --git a/crates/spark-model/src/layers/qwen3_attention/mx8.rs b/crates/spark-model/src/layers/qwen3_attention/mx8.rs new file mode 100644 index 000000000..4db9352b6 --- /dev/null +++ b/crates/spark-model/src/layers/qwen3_attention/mx8.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Native MXFP8 attention projections (nvidia/Qwen3.8-Flash-Next-FP4-Local). +//! +//! That checkpoint ships q/k/v/o as MXFP8. A layer built from them holds +//! each as `QuantWeight::Mxfp8` and nothing else: the BF16 `attn.*_proj` +//! slots are NULL and no NVFP4 / FP8 twin exists. Every projection site asks +//! [`Qwen3AttentionLayer::try_mx8`] first; a site that did not would fall +//! through to a NULL BF16 slot and fault rather than read another copy. + +use anyhow::{Context, Result, ensure}; +use spark_runtime::gpu::DevicePtr; + +use super::Qwen3AttentionLayer; +use crate::layer::ForwardContext; +use crate::layers::ops::{self, Mxfp8ProjKernels, Mxfp8Scratch}; +use crate::weight_map::{Mxfp8Proj, QuantWeight}; + +/// Kernels and the prefill expansion scratch for MXFP8 projections. +#[derive(Debug, Clone, Copy)] +pub struct Mx8AttnRt { + pub kernels: Mxfp8ProjKernels, + pub scratch: Mxfp8Scratch, +} + +impl Qwen3AttentionLayer { + /// Install MXFP8 q/k/v/o. Refused when the BF16 slots still hold + /// weights, when another decode copy exists, or when the scratch cannot + /// take the largest expansion. + pub fn set_mx8_projections( + &mut self, + [q, k, v, o]: [Mxfp8Proj; 4], + rt: Mx8AttnRt, + ) -> Result<()> { + ensure!( + self.attn.q_proj.weight.is_null() + && self.attn.k_proj.weight.is_null() + && self.attn.v_proj.weight.is_null() + && self.attn.o_proj.is_null() + && self.o_dense_bf16.is_none() + && [ + &self.q_weight, + &self.k_weight, + &self.v_weight, + &self.o_weight + ] + .iter() + .all(|w| w.is_none()), + "MXFP8 attention: the layer already holds another copy of its projections" + ); + let largest = [q, k, v, o].iter().map(Mxfp8Proj::bf16_bytes).max(); + ensure!( + !rt.scratch.ptr.is_null() && rt.scratch.bytes >= largest.unwrap_or(0), + "MXFP8 attention: prefill scratch {} B is smaller than the largest expansion", + rt.scratch.bytes + ); + self.q_weight = Some(QuantWeight::Mxfp8(q)); + self.k_weight = Some(QuantWeight::Mxfp8(k)); + self.v_weight = Some(QuantWeight::Mxfp8(v)); + self.o_weight = Some(QuantWeight::Mxfp8(o)); + self.mx8_rt = Some(rt); + Ok(()) + } + + /// `out[m, n] = input[m, k] @ W^T` when `weight` is MXFP8 (`Some(result)`); + /// `None` for any other format, which the caller dispatches itself. + /// `input` and `out` rows are contiguous at the projection's widths. + pub(crate) fn try_mx8( + &self, + ctx: &ForwardContext, + weight: Option<&QuantWeight>, + input: DevicePtr, + out: DevicePtr, + m: usize, + stream: u64, + ) -> Option> { + let proj = weight.and_then(QuantWeight::as_mxfp8)?; + Some( + self.mx8_rt + .as_ref() + .context("MXFP8 attention weight without its kernels") + .and_then(|rt| { + ops::mxfp8_proj( + ctx.gpu, + &rt.kernels, + proj, + input, + out, + m, + rt.scratch, + stream, + ) + }), + ) + } + + /// The prefill guard for weights that are not GEMM-ready as stored: + /// native MXFP8 ([`Self::try_mx8`]), then keep-packed Q2_0 + /// (`try_q2_prefill`). `None` for every other format. Checked first at + /// each prefill projection: the arms after it read NULL slots for both. + pub(crate) fn try_unpacked_prefill( + &self, + ctx: &ForwardContext, + weight: Option<&QuantWeight>, + input: DevicePtr, + out: DevicePtr, + m: u32, + stream: u64, + ) -> Option> { + self.try_mx8(ctx, weight, input, out, m as usize, stream) + .or_else(|| self.try_q2_prefill(ctx, weight, input, out, m, stream)) + } +} diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_qkv.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_qkv.rs index 7e363e283..92c640e8e 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_qkv.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip_qkv.rs @@ -167,8 +167,8 @@ impl Qwen3AttentionLayer { ), }; - // Keep-packed Q2_0 (Tier-1c): transient-dequant to BF16 then dense GEMM. - if let Some(r) = self.try_q2_prefill(ctx, weight_opt, normed, out, n, stream) { + // Native MXFP8 / keep-packed Q2_0: not GEMM-ready as stored. + if let Some(r) = self.try_unpacked_prefill(ctx, weight_opt, normed, out, n, stream) { return r; } diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_oproj.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_oproj.rs index 20c358e98..b6375699e 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_oproj.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_oproj.rs @@ -36,9 +36,9 @@ impl Qwen3AttentionLayer { ops::bf16_gemm(attn_out, w.0, o_out, n, h, nq * hd, stream)?; return Ok(o_out); } - // Keep-packed Q2_0 (Tier-1c): transient-dequant o_proj then dense GEMM. + // Native MXFP8 / keep-packed Q2_0: not GEMM-ready as stored. if let Some(r) = - self.try_q2_prefill(ctx, self.o_weight.as_ref(), attn_out, o_out, n, stream) + self.try_unpacked_prefill(ctx, self.o_weight.as_ref(), attn_out, o_out, n, stream) { r?; return Ok(o_out); diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_qkv.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_qkv.rs index 2990d6752..842a5d50a 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_qkv.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_qkv.rs @@ -211,8 +211,8 @@ impl Qwen3AttentionLayer { ); return Ok(()); } - // Keep-packed Q2_0 (Tier-1c): transient-dequant to BF16 then dense GEMM. - if let Some(r) = self.try_q2_prefill(ctx, weight_opt, normed, out, n, stream) { + // Native MXFP8 / keep-packed Q2_0: not GEMM-ready as stored. + if let Some(r) = self.try_unpacked_prefill(ctx, weight_opt, normed, out, n, stream) { return r; } diff --git a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/attn/o_proj.rs b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/attn/o_proj.rs index 7632a499b..73ffd00b4 100644 --- a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/attn/o_proj.rs +++ b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/attn/o_proj.rs @@ -155,7 +155,10 @@ impl Qwen3AttentionLayer { } let o_out = fwd.buffers.moe_output(); - if let Some(q2) = self.o_weight.as_ref().and_then(|w| w.as_packed_q2()) { + if let Some(r) = self.try_mx8(fwd, self.o_weight.as_ref(), attn_out, o_out, n, stream) { + // Native MXFP8 (FP4-Local): contiguous [n, q_dim] -> [n, h]. + r?; + } else if let Some(q2) = self.o_weight.as_ref().and_then(|w| w.as_packed_q2()) { // Keep-packed Q2_0 (Tier-1c): per-token 2-bit o_proj GEMV. for i in 0..n { let attn_out_i = attn_out.offset(i * q_dim as usize * bf16); 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 64a74ec81..64ccdae3e 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 @@ -28,11 +28,14 @@ mod mla_gemv; mod mla_glm; mod mla_glm_sparse; mod mla_independent; +#[cfg(test)] +mod mx8_tests; mod nemotron_serial; mod qkv; mod qkv_dp4a; mod qkv_exact4; mod qkv_fp8; +mod qkv_mx8; mod qsa_rows; #[cfg(test)] mod tests; diff --git a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mx8_tests.rs b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mx8_tests.rs new file mode 100644 index 000000000..d5848624c --- /dev/null +++ b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/mx8_tests.rs @@ -0,0 +1,118 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Multi-sequence / verify Q, K, V of a native MXFP8 attention layer +//! (`qkv_mx8.rs`) on the mock backend. + +use super::super::super::Qwen3AttentionLayer; +use super::super::super::mx8::Mx8AttnRt; +use super::tests::qkv_phase_d2d; +use crate::layers::FfnComponent; +use crate::layers::ops::{Mxfp8ProjKernels, Mxfp8Scratch}; +use crate::weight_map::{AttentionWeights, DenseWeight, Mxfp8Proj, Mxfp8Weight, QuantizedWeight}; +use atlas_core::config::ModelConfig; +use spark_runtime::gpu::mock::MockGpuBackend; +use spark_runtime::gpu::{DevicePtr, GpuBackend}; +use spark_runtime::kv_cache::KvCacheDtype; + +/// An attention layer as `attn_mx8::build` wires it: NULL BF16 slots, no +/// NVFP4 twin; `bf16_slots` instead fills the BF16 slots. +fn layer(gpu: &MockGpuBackend, config: &ModelConfig, bf16_slots: bool) -> Qwen3AttentionLayer { + let h = config.hidden_size; + let dw = |bytes: usize| DenseWeight { + weight: if bytes == 0 { + DevicePtr::NULL + } else { + gpu.alloc(bytes).unwrap() + }, + }; + let slot = |bytes: usize| dw(if bf16_slots { bytes } else { 0 }); + let attn = AttentionWeights { + q_proj: slot(h * 2), + k_proj: slot(h * 2), + v_proj: slot(h * 2), + o_proj: QuantizedWeight::null(), + q_norm: dw(config.head_dim * 2), + k_norm: dw(config.head_dim * 2), + q_norm_full: None, + k_norm_full: None, + k_scale: 1.0, + v_scale: 1.0, + }; + Qwen3AttentionLayer::new( + dw(h * 2), + attn, + dw(h * 2), + FfnComponent::None, + 0, + None, + None, + None, + gpu, + KvCacheDtype::Bf16, + 0, + config, + ) + .unwrap() +} + +fn projections(gpu: &MockGpuBackend, config: &ModelConfig) -> ([Mxfp8Proj; 4], Mx8AttnRt) { + let h = config.hidden_size; + let q_dim = config.num_attention_heads * config.head_dim; + let kv_dim = config.num_key_value_heads * config.head_dim; + let proj = |n: usize, k: usize| Mxfp8Proj { + w: Mxfp8Weight { + data: gpu.alloc(n * k).unwrap(), + scales: gpu.alloc(n * k / 32).unwrap(), + }, + n: n as u32, + k: k as u32, + }; + let projs = [ + proj(2 * q_dim, h), + proj(kv_dim, h), + proj(kv_dim, h), + proj(h, q_dim), + ]; + let bytes = 2 * q_dim * h * 2; + let rt = Mx8AttnRt { + kernels: Mxfp8ProjKernels::resolve(gpu).unwrap(), + scratch: Mxfp8Scratch { + ptr: gpu.alloc(bytes).unwrap(), + bytes, + }, + }; + (projs, rt) +} + +/// Four verify rows: one weight-streaming launch per projection (grid +/// ceil(N/4)), then three scatter copies a row into `qkv_buf`. +#[test] +fn verify_rows_read_each_mxfp8_projection_once() { + let config = ModelConfig::qwen3_next_80b_nvfp4(); + let gpu = MockGpuBackend::new(); + let mut l = layer(&gpu, &config, false); + let (projs, rt) = projections(&gpu, &config); + l.set_mx8_projections(projs, rt).unwrap(); + let (copies, launches) = qkv_phase_d2d(&gpu, &config, &l, 4); + assert_eq!(copies, 3 * 4, "three scatter copies a row"); + let q_rows = 2 * config.num_attention_heads * config.head_dim; + let kv_rows = config.num_key_value_heads * config.head_dim; + let grids: Vec<[u32; 3]> = launches.iter().map(|l| l.grid).collect(); + for (what, n) in [("q", q_rows), ("k/v", kv_rows)] { + assert!( + grids.contains(&[(n / 4) as u32, 1, 1]), + "{what} MXFP8 GEMV missing; grids: {grids:?}" + ); + } +} + +/// BF16 slots still holding weights means a second copy: refused. +#[test] +fn a_layer_with_bf16_slots_is_refused() { + let config = ModelConfig::qwen3_next_80b_nvfp4(); + let gpu = MockGpuBackend::new(); + let mut l = layer(&gpu, &config, true); + let (projs, rt) = projections(&gpu, &config); + let err = l.set_mx8_projections(projs, rt).unwrap_err(); + assert!(err.to_string().contains("another copy"), "{err:#}"); +} 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 5bd7e29f3..8cfe1b626 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 @@ -46,6 +46,8 @@ impl Qwen3AttentionLayer { // ATLAS_QWEN4EXP_BATCH_FAST_BISECT: the sequential per-row arm. if fwd.levers.batch_bisect(ops::BISECT_ATTN_PROJ) { self.ms_qkv_seq_rows(c)?; + } else if self.ms_qkv_batch_mx8(c)? { + // Native MXFP8 (FP4-Local): `qkv_mx8.rs`. } else if self.ms_qkv_exact4(c)? { // qwen4_exp exact K=4 verify: serial decode's arithmetic per row // (`qkv_exact4.rs`). diff --git a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/qkv_mx8.rs b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/qkv_mx8.rs new file mode 100644 index 000000000..120a4f14b --- /dev/null +++ b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/qkv_mx8.rs @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Multi-sequence / verify Q, K, V on native MXFP8 weights (FP4-Local). +//! +//! The same shape as `qkv_fp8.rs`: each projection runs once over all `n` +//! rows into a contiguous scratch (one weight pass per eight rows, every row +//! bit-identical to single-token decode), the gated `[Q|gate]` rows are +//! split, then rows are copied into the per-sequence `qkv_buf` layout. + +use anyhow::Result; + +use super::ctx::MultiSeqCtx; +use crate::layers::ops; +use crate::layers::qwen3_attention::Qwen3AttentionLayer; + +impl Qwen3AttentionLayer { + /// `Ok(false)`, nothing launched, unless the layer's projections are MXFP8. + pub(super) fn ms_qkv_batch_mx8(&self, c: &MultiSeqCtx<'_>) -> Result { + if self.q_weight.as_ref().and_then(|w| w.as_mxfp8()).is_none() { + return Ok(false); + } + let MultiSeqCtx { + fwd, + n, + stream, + nq, + nkv, + hd, + bf16, + q_proj_dim, + q_proj_bytes, + per_seq_qkv, + normed, + qkv_buf, + .. + } = *c; + let q_scratch = fwd.buffers.ssm_qkvz(); + let kv_bytes = (nkv * hd) as usize * bf16; + let k_scratch = fwd.buffers.attn_output(); + let v_scratch = k_scratch.offset(n * kv_bytes); + for (w, out) in [ + (self.q_weight.as_ref(), q_scratch), + (self.k_weight.as_ref(), k_scratch), + (self.v_weight.as_ref(), v_scratch), + ] { + self.try_mx8(fwd, w, normed, out, n, stream) + .unwrap_or_else(|| anyhow::bail!("MXFP8 q_proj with a non-MXFP8 k/v_proj"))?; + } + if self.gated && !self.q_lora_active() { + ops::deinterleave_qg( + fwd.gpu, + self.deinterleave_qg_k, + q_scratch, + n as u32, + nq, + hd, + q_proj_dim, + stream, + )?; + } + for row in 0..n { + let q_out = qkv_buf.offset(row * per_seq_qkv); + let k_out = q_out.offset(q_proj_bytes); + let v_out = k_out.offset(kv_bytes); + let copy = |src: spark_runtime::gpu::DevicePtr, dst, bytes| { + fwd.gpu.copy_d2d_async(src, dst, bytes, stream) + }; + copy(q_scratch.offset(row * q_proj_bytes), q_out, q_proj_bytes)?; + copy(k_scratch.offset(row * kv_bytes), k_out, kv_bytes)?; + copy(v_scratch.offset(row * kv_bytes), v_out, kv_bytes)?; + } + Ok(true) + } +} diff --git a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/tests.rs b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/tests.rs index 130b52198..85b0b50bd 100644 --- a/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/tests.rs +++ b/crates/spark-model/src/layers/qwen3_attention/trait_impl/multi_seq/tests.rs @@ -69,7 +69,7 @@ fn attn_layer(gpu: &MockGpuBackend, config: &ModelConfig) -> Qwen3AttentionLayer /// Run `ms_phase_qkv` at `n` verify rows and return how many d2d copies /// the phase issued plus every kernel launch it recorded. -fn qkv_phase_d2d( +pub(super) fn qkv_phase_d2d( gpu: &MockGpuBackend, config: &ModelConfig, layer: &Qwen3AttentionLayer, diff --git a/crates/spark-model/src/layers/qwen3_attention/types_layer.rs b/crates/spark-model/src/layers/qwen3_attention/types_layer.rs index 8571a041f..5f2f2d885 100644 --- a/crates/spark-model/src/layers/qwen3_attention/types_layer.rs +++ b/crates/spark-model/src/layers/qwen3_attention/types_layer.rs @@ -104,6 +104,8 @@ pub struct Qwen3AttentionLayer { /// NVFP4 path (`attn.o_proj`). Used by Gemma-4 dense which honors /// Nvidia ModelOpt's official ignore list. pub(super) o_dense_bf16: Option, + /// MXFP8 kernels + prefill scratch for `QuantWeight::Mxfp8` (`mx8.rs`). + pub(super) mx8_rt: Option, // ── MLA (Multi-head Latent Attention) — 2-step decode ── pub(crate) mla: Option, /// MXFP8 twins of MLA weights keyed by the BF16 weight pointer, with the diff --git a/crates/spark-model/src/layers/qwen3_ssm/exact_rows.rs b/crates/spark-model/src/layers/qwen3_ssm/exact_rows.rs index ffb03721c..1cd409024 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/exact_rows.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/exact_rows.rs @@ -54,20 +54,16 @@ impl Qwen3SsmLayer { GdnProj::Qkvz => &mx8.qkvz, GdnProj::Out => &mx8.out, }; - let (in_row, out_row) = (p.k as usize * 2, p.n as usize * 2); - for first in (0..rows).step_by(ops::MXFP8_PROJ_MAX_GEMV_ROWS) { - let m = (rows - first).min(ops::MXFP8_PROJ_MAX_GEMV_ROWS); - ops::mxfp8_proj( - ctx.gpu, - &mx8.kernels, - p, - input.offset(first * in_row), - output.offset(first * out_row), - m, - mx8.scratch, - stream, - )?; - } + ops::mxfp8_proj( + ctx.gpu, + &mx8.kernels, + p, + input, + output, + rows, + mx8.scratch, + stream, + )?; return Ok(true); } if !ctx.levers.qwen4exp_batch_fast diff --git a/crates/spark-model/src/weight_loader/qwen4_exp.rs b/crates/spark-model/src/weight_loader/qwen4_exp.rs index 0bfde92ee..727ebe410 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp.rs @@ -77,11 +77,13 @@ use crate::weight_map::{DenseWeight, MtpWeights, dense}; // compiler ran. `attach` is also the better name: the module is the three // per-layer attach helpers. mod attach; +mod attn_mx8; mod ffn; mod gdn_fp8; mod gdn_mx8; mod hc; mod mtp; +mod mx8_common; mod ple; mod probe; @@ -231,26 +233,8 @@ impl ModelWeightLoader for Qwen4ExpWeightLoader { // With PLE disabled for bisection, skip the 21 MB arena and the // 128-shard open entirely rather than building what we will not run. let ple_off = std::env::var("ATLAS_QWEN4EXP_NO_PLE").as_deref() == Ok("1"); - // GDN projections stay BF16 by DEFAULT on this checkpoint. Measured, - // both arms, same prompt, util 0.85 / 16K / bf16 KV: - // - // requantized NVFP4 BF16 (default) - // layer construction 7.43 GB / 154.8 MB/L 1.39 GB / 28.9 MB/L - // attn/GDN arms 7.34 GB 1.32 GB - // pre-KV 95.8 GB 90.0 GB - // KV budget 3.9 GB / 172144 tok 9.7 GB / 424464 tok - // decode 2.207 tok/s 2.188 tok/s - // - // 6.04 GB back for ~1% of decode. GDN weight bandwidth is simply not - // what bounds decode here at C=1 — something else dominates — so the - // usual w4a16-is-faster argument does not apply yet. Revisit if that - // changes. - // - // And it is not only a memory lever: ONLY the routed experts are - // quantized in this checkpoint. The GDN projections ship BF16, so - // requantizing them was a lossy round trip we chose, on 36 of 48 - // layers. `=0` opts back into it for A/B. `ATLAS_QWEN4EXP_FP8_GDN` is - // the FP8 middle option (`gdn_fp8`). + // GDN projections stay BF16 by DEFAULT (measured trade-off in + // `gdn_fp8`'s docs); FP4-Local's MXFP8 ones are served natively. let gdn = GdnProjections::from_env()?; tracing::info!( "GDN projections: {} on the {} linear-attention layers", @@ -293,7 +277,7 @@ impl ModelWeightLoader for Qwen4ExpWeightLoader { let layer = match config.layer_types[i] { // FP4-Local ships these MXFP8: served on those bytes. LayerType::LinearAttention if gdn_mx8::present(store, &lp) => { - let scratch = gdn_mx8::scratch_for(&mut mx8_scratch, config, gpu)?; + let scratch = mx8_common::scratch_for(&mut mx8_scratch, config, gpu)?; gdn_mx8::build( i, store, &lp, gpu, config, input_norm, post_attn_norm, ffn, scratch, ) @@ -323,6 +307,17 @@ impl ModelWeightLoader for Qwen4ExpWeightLoader { ) .with_context(|| format!("qwen4_exp: GDN layer {i}"))? } + LayerType::FullAttention if attn_mx8::present(store, &lp) => { + let kv_dtype = layer_kv_dtypes.get(attn_idx).copied(); + let scratch = mx8_common::scratch_for(&mut mx8_scratch, config, gpu)?; + let l = attn_mx8::build( + i, store, &lp, gpu, config, kv_dtype.unwrap_or(KvCacheDtype::Bf16), + attn_idx, (input_norm, post_attn_norm, ffn), scratch, + ) + .with_context(|| format!("qwen4_exp: full-attention layer {i} (MXFP8)"))?; + attn_idx += 1; + l + } LayerType::FullAttention => { let kv_dtype = layer_kv_dtypes .get(attn_idx) diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/attn_mx8.rs b/crates/spark-model/src/weight_loader/qwen4_exp/attn_mx8.rs new file mode 100644 index 000000000..a502f6a9a --- /dev/null +++ b/crates/spark-model/src/weight_loader/qwen4_exp/attn_mx8.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Full-attention layers of nvidia/Qwen3.8-Flash-Next-FP4-Local, on their +//! MXFP8 q/k/v/o bytes. +//! +//! The NVFP4 pack ships these projections BF16, and the generic builder +//! (`build_full_attention_nvfp4`) requantizes BF16 to NVFP4 for decode. On +//! FP4-Local that would be a second, lossy quantization of weights NVIDIA's +//! QAD already trained in MXFP8, so this builder skips it: q (with its +//! interleaved gate rows), k, v and o stay exactly as shipped and every +//! forward reads them through `mxfp8_proj.cu`. + +use anyhow::{Context, Result, ensure}; +use atlas_core::config::ModelConfig; +use spark_runtime::gpu::{DevicePtr, GpuBackend}; +use spark_runtime::kv_cache::KvCacheDtype; +use spark_runtime::weights::{WeightDtype, WeightStore}; + +use super::mx8_common::module; +use crate::layer::TransformerLayer; +use crate::layers::ops::{Mxfp8ProjKernels, Mxfp8Scratch}; +use crate::layers::qwen3_attention::mx8::Mx8AttnRt; +use crate::layers::{FfnComponent, Qwen3AttentionLayer}; +use crate::weight_map::{ + AttentionWeights, DenseWeight, Mxfp8Proj, Mxfp8Weight, QuantizedWeight, dense_auto, + load_kv_scales, +}; + +/// True when layer `lp`'s attention projections are MXFP8 in the store. +pub(super) fn present(store: &WeightStore, lp: &str) -> bool { + let q = format!("{lp}.self_attn.q_proj"); + store + .get(&format!("{q}.weight")) + .is_ok_and(|t| t.dtype == WeightDtype::FP8E4M3) + && store.contains(&format!("{q}.weight_scale")) +} + +/// Build full-attention layer `layer_idx` on its MXFP8 projections. +#[allow(clippy::too_many_arguments)] +pub(super) fn build( + layer_idx: usize, + store: &WeightStore, + lp: &str, + gpu: &dyn GpuBackend, + config: &ModelConfig, + kv_dtype: KvCacheDtype, + attn_idx: usize, + (input_norm, post_attn_norm, ffn): (DenseWeight, DenseWeight, FfnComponent), + scratch: Mxfp8Scratch, +) -> Result> { + ensure!( + config.tp_world_size.max(1) == 1, + "MXFP8 attention layer {layer_idx}: tensor parallelism is not supported yet" + ); + let h = config.hidden_size; + let q_dim = config.num_attention_heads * config.head_dim; + let q_rows = if config.attn_gated { 2 * q_dim } else { q_dim }; + let kv_rows = config.num_key_value_heads * config.head_dim; + let p = format!("{lp}.self_attn"); + let proj = |name: &str, n: usize, k: usize| -> Result { + let [data, scales] = module(store, &format!("{p}.{name}"), n, k)?; + Ok(Mxfp8Proj { + w: Mxfp8Weight { data, scales }, + n: n as u32, + k: k as u32, + }) + }; + let projs = [ + proj("q_proj", q_rows, h)?, + proj("k_proj", kv_rows, h)?, + proj("v_proj", kv_rows, h)?, + proj("o_proj", h, q_dim)?, + ]; + + let null = DenseWeight { + weight: DevicePtr::NULL, + }; + let (k_scale, v_scale) = load_kv_scales(store, &p, gpu); + let attn = AttentionWeights { + q_proj: null, + k_proj: null, + v_proj: null, + o_proj: QuantizedWeight::null(), + q_norm: dense_auto(store, &format!("{p}.q_norm.weight"), gpu)?, + k_norm: dense_auto(store, &format!("{p}.k_norm.weight"), gpu)?, + q_norm_full: None, + k_norm_full: None, + k_scale, + v_scale, + }; + let mut layer = Qwen3AttentionLayer::new( + input_norm, + attn, + post_attn_norm, + ffn, + attn_idx, + None, + None, + None, + gpu, + kv_dtype, + config.fp8_kv_calibration_tokens, + config, + )?; + layer + .set_mx8_projections( + projs, + Mx8AttnRt { + kernels: Mxfp8ProjKernels::resolve(gpu)?, + scratch, + }, + ) + .with_context(|| format!("MXFP8 attention layer {layer_idx}"))?; + tracing::info!("Layer {layer_idx}: attention q/k/v/o native MXFP8 (FP4-Local)"); + Ok(Box::new(layer)) +} diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/gdn_fp8.rs b/crates/spark-model/src/weight_loader/qwen4_exp/gdn_fp8.rs index ae821fefa..374efe459 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/gdn_fp8.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/gdn_fp8.rs @@ -28,6 +28,29 @@ //! top: +28.9 MB a layer per TP2 rank. //! `=full` frees the BF16 copies (-57.7 MB a layer at TP1, where there is no //! headroom for `=1`) and prefill reads the FP8 copy too. +//! +//! Why BF16 is the default — measured on the NVFP4 pack, both arms, same +//! prompt, util 0.85 / 16K / bf16 KV: +//! +//! ```text +//! requantized NVFP4 BF16 (default) +//! layer construction 7.43 GB / 154.8 MB/L 1.39 GB / 28.9 MB/L +//! attn/GDN arms 7.34 GB 1.32 GB +//! pre-KV 95.8 GB 90.0 GB +//! KV budget 3.9 GB / 172144 tok 9.7 GB / 424464 tok +//! decode 2.207 tok/s 2.188 tok/s +//! ``` +//! +//! 6.04 GB back for ~1% of decode. GDN weight bandwidth is simply not +//! what bounds decode here at C=1 — something else dominates — so the +//! usual w4a16-is-faster argument does not apply yet. Revisit if that +//! changes. +//! +//! And it is not only a memory lever: ONLY the routed experts are +//! quantized in this checkpoint. The GDN projections ship BF16, so +//! requantizing them was a lossy round trip we chose, on 36 of 48 +//! layers. `=0` opts back into it for A/B. `ATLAS_QWEN4EXP_FP8_GDN` is +//! the FP8 middle option (`gdn_fp8`). use anyhow::{Result, bail, ensure}; use atlas_core::config::ModelConfig; diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/gdn_mx8.rs b/crates/spark-model/src/weight_loader/qwen4_exp/gdn_mx8.rs index 46f0704ca..247273661 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/gdn_mx8.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/gdn_mx8.rs @@ -8,8 +8,8 @@ //! layer around those bytes: `in_proj_qkv || in_proj_z` rows (and their //! scale rows) become the sequential `[Q|K|V|Z]` projection, `out_proj` is //! used in place, and the layer reads both through `mxfp8_proj.cu`. No BF16 -//! copy exists; prefill expands one projection at a time into a scratch -//! shared by every layer. +//! copy exists; prefill expands one projection at a time into the scratch +//! `mx8_common::scratch_for` shares between every MXFP8 layer. //! //! `in_proj_a` / `in_proj_b` (4 MB a layer in all) stay on the BF16 //! interleaved β/α path the fused gate kernel reads. @@ -19,6 +19,7 @@ use atlas_core::config::ModelConfig; use spark_runtime::gpu::{DevicePtr, GpuBackend}; use spark_runtime::weights::{WeightDtype, WeightStore}; +use super::mx8_common::{MX8_GROUP, module}; use crate::layer::TransformerLayer; use crate::layers::FfnComponent; use crate::layers::Qwen3SsmLayer; @@ -30,8 +31,6 @@ use crate::weight_map::{ dense_keep_f32, free_loader_source, interleave_ba, }; -const MX8_GROUP: usize = 32; - /// True when layer `lp`'s GDN projections are MXFP8 in the store. pub(super) fn present(store: &WeightStore, lp: &str) -> bool { let w = format!("{lp}.linear_attn.in_proj_qkv"); @@ -41,25 +40,6 @@ pub(super) fn present(store: &WeightStore, lp: &str) -> bool { && store.contains(&format!("{w}.weight_scale")) } -/// The weight and scale buffers of one MXFP8 module, shapes checked. -fn module(store: &WeightStore, name: &str, rows: usize, k: usize) -> Result<[DevicePtr; 2]> { - let w = store.get(&format!("{name}.weight"))?; - let s = store.get(&format!("{name}.weight_scale"))?; - ensure!( - w.dtype == WeightDtype::FP8E4M3 - && w.shape == [rows, k] - && s.dtype == WeightDtype::UInt8 - && s.shape == [rows, k / MX8_GROUP], - "{name}: expected MXFP8 [{rows}, {k}] + [{rows}, {}], got {:?} {:?} + {:?} {:?}", - k / MX8_GROUP, - w.dtype, - w.shape, - s.dtype, - s.shape - ); - Ok([w.ptr, s.ptr]) -} - /// Copy `a` then `b` (each `rows * row_bytes`) into one new buffer. fn concat( gpu: &dyn GpuBackend, @@ -73,32 +53,6 @@ fn concat( Ok(buf) } -/// The prefill scratch every MXFP8 GDN layer shares: sized for the larger -/// projection, allocated by the first layer that needs it. -pub(super) fn scratch_for( - shared: &mut Option, - config: &ModelConfig, - gpu: &dyn GpuBackend, -) -> Result { - let h = config.hidden_size; - let value_dim = config.linear_num_value_heads * config.linear_value_head_dim; - let bytes = (config.ssm_qkvz_size() * h).max(h * value_dim) * 2; - if let Some(s) = shared.filter(|s| s.bytes >= bytes) { - return Ok(s); - } - ensure!(shared.is_none(), "MXFP8 GDN scratch sized twice"); - let s = Mxfp8Scratch { - ptr: gpu.alloc(bytes)?, - bytes, - }; - tracing::info!( - "MXFP8 GDN: one {:.1} MB BF16 prefill scratch shared by every GDN layer", - bytes as f64 / 1e6 - ); - *shared = Some(s); - Ok(s) -} - /// Build GDN layer `layer_idx` on its MXFP8 projections. #[allow(clippy::too_many_arguments)] pub(super) fn build( diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/mx8_common.rs b/crates/spark-model/src/weight_loader/qwen4_exp/mx8_common.rs new file mode 100644 index 000000000..559b53d6c --- /dev/null +++ b/crates/spark-model/src/weight_loader/qwen4_exp/mx8_common.rs @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Shared pieces of the FP4-Local MXFP8 layer builders (`gdn_mx8`, +//! `attn_mx8`): the checked weight/scale lookup and the one BF16 prefill +//! scratch every MXFP8 layer expands into. + +use anyhow::{Result, ensure}; +use atlas_core::config::ModelConfig; +use spark_runtime::gpu::{DevicePtr, GpuBackend}; +use spark_runtime::weights::{WeightDtype, WeightStore}; + +use crate::layers::ops::Mxfp8Scratch; + +/// Inputs per E8M0 scale. +pub(super) const MX8_GROUP: usize = 32; + +/// The weight and scale buffers of one MXFP8 module, shapes checked. +pub(super) fn module( + store: &WeightStore, + name: &str, + rows: usize, + k: usize, +) -> Result<[DevicePtr; 2]> { + let w = store.get(&format!("{name}.weight"))?; + let s = store.get(&format!("{name}.weight_scale"))?; + ensure!( + w.dtype == WeightDtype::FP8E4M3 + && w.shape == [rows, k] + && s.dtype == WeightDtype::UInt8 + && s.shape == [rows, k / MX8_GROUP], + "{name}: expected MXFP8 [{rows}, {k}] + [{rows}, {}], got {:?} {:?} + {:?} {:?}", + k / MX8_GROUP, + w.dtype, + w.shape, + s.dtype, + s.shape + ); + Ok([w.ptr, s.ptr]) +} + +/// The prefill scratch every MXFP8 layer shares, sized for the largest +/// projection either layer kind holds (GDN `[Q|K|V|Z]` and `out_proj`, +/// attention `q` (with its gate) and `o`) and allocated by the first layer +/// that asks. Prefill runs layers one after another on one stream, so one +/// buffer serves them all. +pub(super) fn scratch_for( + shared: &mut Option, + config: &ModelConfig, + gpu: &dyn GpuBackend, +) -> Result { + if let Some(s) = *shared { + return Ok(s); + } + let h = config.hidden_size; + let value_dim = config.linear_num_value_heads * config.linear_value_head_dim; + let q_dim = config.num_attention_heads * config.head_dim; + let q_rows = if config.attn_gated { 2 * q_dim } else { q_dim }; + let elems = [ + config.ssm_qkvz_size() * h, + h * value_dim, + q_rows * h, + h * q_dim, + ]; + let bytes = elems.into_iter().max().unwrap_or(0) * 2; + ensure!(bytes > 0, "MXFP8 scratch: the config has no projection"); + let s = Mxfp8Scratch { + ptr: gpu.alloc(bytes)?, + bytes, + }; + tracing::info!( + "FP4-Local MXFP8: one {:.1} MB BF16 prefill scratch shared by every MXFP8 layer", + bytes as f64 / 1e6 + ); + *shared = Some(s); + Ok(s) +} diff --git a/crates/spark-model/src/weight_map.rs b/crates/spark-model/src/weight_map.rs index 53293ad2b..2282c6518 100644 --- a/crates/spark-model/src/weight_map.rs +++ b/crates/spark-model/src/weight_map.rs @@ -58,7 +58,7 @@ pub use loaders_fp8::*; pub use loaders_mtp::*; pub use model_a::*; pub use moe::*; -pub use mxfp8::Mxfp8Weight; +pub use mxfp8::{Mxfp8Proj, Mxfp8Weight}; pub use nemotron::*; pub use nvfp4_detect::*; pub use quantize_fns::*; diff --git a/crates/spark-model/src/weight_map/expert.rs b/crates/spark-model/src/weight_map/expert.rs index b5630cb1c..79c558927 100644 --- a/crates/spark-model/src/weight_map/expert.rs +++ b/crates/spark-model/src/weight_map/expert.rs @@ -130,6 +130,11 @@ pub enum QuantWeight { /// EXL3 trellis (mul1 codebook). Reconstruct + GEMM lands in M3. Exl3(super::exl3::Exl3Weight), + + /// MXFP8 as the checkpoint ships it (E4M3 + E8M0 per 32 inputs). + /// Kernels: `mxfp8_proj.cu` through `ops::mxfp8_proj`, which needs a + /// scratch for prefill, so the generic `quant_gemv`/`quant_gemm` refuse it. + Mxfp8(super::Mxfp8Proj), } impl QuantWeight { @@ -146,6 +151,7 @@ impl QuantWeight { Self::Dense(w) => w.weight.is_null(), Self::PackedQ2(w) => w.is_null(), Self::Exl3(w) => w.is_null(), + Self::Mxfp8(p) => p.w.data.is_null(), } } @@ -165,6 +171,14 @@ impl QuantWeight { } } + /// Extract as MXFP8, if this weight is that variant. + pub fn as_mxfp8(&self) -> Option<&super::Mxfp8Proj> { + match self { + Self::Mxfp8(p) => Some(p), + _ => None, + } + } + /// Extract as NVFP4, if this weight is that variant. pub fn as_nvfp4(&self) -> Option<&QuantizedWeight> { match self { diff --git a/crates/spark-model/src/weight_map/modelopt_local.rs b/crates/spark-model/src/weight_map/modelopt_local.rs index 014cf4bc3..5cdd1d2db 100644 --- a/crates/spark-model/src/weight_map/modelopt_local.rs +++ b/crates/spark-model/src/weight_map/modelopt_local.rs @@ -10,8 +10,8 @@ //! (`weight_scale.nvfp4_csf_{fixed,exceptions}`) — decoded losslessly into //! the plain E4M3 `weight_scale` every NVFP4 kernel reads; //! * GDN / attention / QSA-indexer / shared-expert projections are MXFP8. -//! Those with a native MXFP8 consumer ([`native_mxfp8`]: the GDN -//! projections) are uploaded byte for byte, `weight` E4M3 + `weight_scale` +//! Those with a native MXFP8 consumer ([`native_mxfp8`]: the GDN and +//! attention projections) are uploaded byte for byte, `weight` E4M3 + `weight_scale` //! E8M0, and read by `mxfp8_proj.cu`; the rest are expanded exactly to the //! BF16 `weight` the NVFP4 pack ships; //! * vision MLP/merger projections are W4A16 NVFP4 — dequantized to the BF16 @@ -63,10 +63,17 @@ pub fn native_mxfp8(module: &str) -> bool { if std::env::var("ATLAS_FP4LOCAL_EXPAND_MXFP8").as_deref() == Ok("1") { return false; } + const NATIVE: [&str; 7] = [ + ".linear_attn.in_proj_qkv", + ".linear_attn.in_proj_z", + ".linear_attn.out_proj", + ".self_attn.q_proj", + ".self_attn.k_proj", + ".self_attn.v_proj", + ".self_attn.o_proj", + ]; module.starts_with("model.language_model.layers.") - && ["in_proj_qkv", "in_proj_z", "out_proj"] - .iter() - .any(|p| module.ends_with(&format!(".linear_attn.{p}"))) + && NATIVE.iter().any(|suffix| module.ends_with(suffix)) } /// What [`materialize`] did, for the load log and for tests. diff --git a/crates/spark-model/src/weight_map/modelopt_local_tests.rs b/crates/spark-model/src/weight_map/modelopt_local_tests.rs index 00750f455..33787e059 100644 --- a/crates/spark-model/src/weight_map/modelopt_local_tests.rs +++ b/crates/spark-model/src/weight_map/modelopt_local_tests.rs @@ -6,6 +6,7 @@ use spark_runtime::gpu::mock::MockGpuBackend; use spark_runtime::weights::DeferredTensor; const ATTN: &str = "model.language_model.layers.3.self_attn.q_proj"; +const SHARED: &str = "model.language_model.layers.3.mlp.shared_expert.up_proj"; const VIS: &str = "model.visual.blocks.0.mlp.linear_fc1"; const HC_NORM: &str = "model.language_model.layers.0.attn_hyper_connection.hc_norm.weight"; const PLE_NORM: &str = "model.language_model.layers.1.ple.norm_conv.weight"; @@ -28,7 +29,7 @@ fn fp4_local_config() -> ModelConfig { "model.language_model.layers.0.mlp.experts".into(), spec("NVFP4", 16), ); - map.insert(ATTN.into(), spec("MXFP8", 32)); + map.insert(SHARED.into(), spec("MXFP8", 32)); map.insert(GDN_OUT.into(), spec("MXFP8", 32)); map.insert(VIS.into(), spec("W4A16_NVFP4", 16)); map.insert( @@ -54,7 +55,10 @@ fn fp4_local_config() -> ModelConfig { #[test] fn host_decode_modules_pick_mxfp8_and_vision_nvfp4_only() { let got = host_decode_modules(&fp4_local_config()); - let want: HashSet = [ATTN, GDN_OUT, VIS].iter().map(|s| s.to_string()).collect(); + let want: HashSet = [SHARED, GDN_OUT, VIS] + .iter() + .map(|s| s.to_string()) + .collect(); assert_eq!(got, want); assert!(host_decode_modules(&ModelConfig::qwen3_next_80b_nvfp4()).is_empty()); } @@ -147,13 +151,13 @@ fn materialize_expands_every_fp4_local_format() { // MXFP8 [1, 32]: E4M3 1.0 (0x38) everywhere, E8M0 scale 2^-3. shard.add( - &format!("{ATTN}.weight"), + &format!("{SHARED}.weight"), &[0x38; 32], &[1, 32], WeightDtype::FP8E4M3, ); shard.add( - &format!("{ATTN}.weight_scale"), + &format!("{SHARED}.weight_scale"), &[124], &[1, 1], WeightDtype::UInt8, @@ -249,14 +253,14 @@ fn materialize_expands_every_fp4_local_format() { want[5] = 0x7E; assert_eq!(gpu.read_alloc(scale.ptr).unwrap(), want); - let attn = store.get(&format!("{ATTN}.weight")).unwrap(); + let attn = store.get(&format!("{SHARED}.weight")).unwrap(); assert_eq!( (attn.dtype, attn.shape.clone()), (WeightDtype::BF16, vec![1, 32]) ); let a = gpu.read_alloc(attn.ptr).unwrap(); assert!((0..32).all(|i| bf16_at(&a, i) == 0.125)); - assert!(!store.contains(&format!("{ATTN}.weight_scale"))); + assert!(!store.contains(&format!("{SHARED}.weight_scale"))); let vis = store.get(&format!("{VIS}.weight")).unwrap(); assert_eq!( @@ -281,7 +285,11 @@ fn gdn_mxfp8_stays_native() { assert!(native_mxfp8( "model.language_model.layers.7.linear_attn.in_proj_qkv" )); - assert!(!native_mxfp8(ATTN)); + assert!(native_mxfp8(ATTN)); + assert!(!native_mxfp8( + "model.language_model.layers.3.self_attn.indexer.index_qk_proj" + )); + assert!(!native_mxfp8("model.visual.blocks.0.attn.qkv")); assert!(!native_mxfp8( "model.language_model.layers.7.linear_attn.in_proj_a" )); diff --git a/crates/spark-model/src/weight_map/mxfp8.rs b/crates/spark-model/src/weight_map/mxfp8.rs index 843fae203..189289764 100644 --- a/crates/spark-model/src/weight_map/mxfp8.rs +++ b/crates/spark-model/src/weight_map/mxfp8.rs @@ -11,3 +11,20 @@ pub struct Mxfp8Weight { pub data: DevicePtr, pub scales: DevicePtr, } + +/// An MXFP8 projection with its geometry: `out = act @ W^T`, `W` `[n, k]`. +#[derive(Debug, Clone, Copy)] +pub struct Mxfp8Proj { + pub w: Mxfp8Weight, + /// Output features (weight rows). + pub n: u32, + /// Input features; a multiple of 32. + pub k: u32, +} + +impl Mxfp8Proj { + /// Bytes of the exact BF16 expansion a prefill GEMM multiplies by. + pub fn bf16_bytes(&self) -> usize { + self.n as usize * self.k as usize * 2 + } +} From 788404f0e8967e87b8b4d71ad3cd0e41bd3448e6 Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 02:50:03 -0400 Subject: [PATCH 07/26] qwen4_exp: keep hc/PLE norm weights F32 end to end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The FP4-Local pack ships the 100 hyper-connection hc_norm and PLE norm_conv/norm_key/norm_query weights as F32 (the only release that does). modelopt_local rounded them to BF16 because every consumer kernel was declared __nv_bfloat16* — measured lossy on all 100 tensors. The kernels already compute (1 + w) in FP32, so change the norm-weight arguments to const float* (hyper_connection.cu: hc_pre, hc_head, hc_pre_stage, hc_pre_stage_bf16, hc_post_stage_bf16, hc_pre_stage_vec, hc_post_stage_vec; ple.cu: ple_gate; strix-hip mirror), load every site's norm through dense_keep_f32 — F32 native for FP4-Local, exact (bits << 16) widening for the NVFP4 pack — and delete the BF16 rounding. NVFP4 stays bit-identical; FP4-Local keeps its full precision. --- .../examples/qwen4exp_hc_prefill_check.rs | 2 +- .../ops/hyper_connection_lowrank_tests.rs | 14 ++++- .../spark-model/src/layers/ops/ple_tests.rs | 17 +++++- .../layers/ops/qwen4exp_oracle_ple_tests.rs | 7 ++- .../src/layers/ops/qwen4exp_oracle_tests.rs | 5 +- crates/spark-model/src/layers/ple/layer.rs | 4 +- .../layers/qwen3_attention/types_weights.rs | 6 +- .../src/weight_loader/qwen4_exp/hc.rs | 10 +++- .../src/weight_loader/qwen4_exp/ple.rs | 10 ++-- .../src/weight_map/modelopt_local.rs | 60 +++++-------------- .../src/weight_map/modelopt_local_tests.rs | 23 +++++-- crates/spark-model/src/weight_map/tests.rs | 53 +++++++++++++++- .../nvfp4/hyper_connection.cu | 44 +++++++------- kernels/gb10/qwen3.8-flash-next/nvfp4/ple.cu | 12 ++-- .../nvfp4/hyper_connection.cu | 16 ++--- 15 files changed, 180 insertions(+), 103 deletions(-) diff --git a/crates/spark-model/examples/qwen4exp_hc_prefill_check.rs b/crates/spark-model/examples/qwen4exp_hc_prefill_check.rs index c30b5f4b0..e3f26c7f6 100644 --- a/crates/spark-model/examples/qwen4exp_hc_prefill_check.rs +++ b/crates/spark-model/examples/qwen4exp_hc_prefill_check.rs @@ -75,7 +75,7 @@ fn site(g: &dyn GpuBackend, rng: &mut Lcg, inject: bool) -> Result Vec { out } +/// Widen a BF16 blob to F32 by zero extension — what `dense_keep_f32` does +/// at load now that the norm readers are `const float*`. +fn bf16_to_f32_bytes(bytes: &[u8]) -> Vec { + bytes + .chunks_exact(2) + .flat_map(|c| ((u16::from_le_bytes([c[0], c[1]]) as u32) << 16).to_le_bytes()) + .collect() +} + /// One site's low-rank weights, uploaded. fn site_weights(g: &dyn GpuBackend, f: &Fixture, site: &str, inject: bool) -> HcLowRank { let up = f.bytes(&format!("{site}_w_up")); HcLowRank { - norm_w: upload(g, &f.bytes(&format!("{site}_w_hc_norm"))), + norm_w: upload( + g, + &bf16_to_f32_bytes(&f.bytes(&format!("{site}_w_hc_norm"))), + ), down_w: upload(g, &f.bytes(&format!("{site}_w_down"))), up_w: upload(g, &transpose_bf16(&up, f.hc * f.h, f.rank)), inject_w: if inject { diff --git a/crates/spark-model/src/layers/ops/ple_tests.rs b/crates/spark-model/src/layers/ops/ple_tests.rs index 52ce728f9..d4d8e5da3 100644 --- a/crates/spark-model/src/layers/ops/ple_tests.rs +++ b/crates/spark-model/src/layers/ops/ple_tests.rs @@ -71,6 +71,15 @@ impl Fx { } } +/// Widen a BF16 blob to F32 by zero extension — what `dense_keep_f32` does +/// at load now that the norm readers are `const float*`. +fn bf16_to_f32_bytes(bytes: &[u8]) -> Vec { + bytes + .chunks_exact(2) + .flat_map(|c| ((u16::from_le_bytes([c[0], c[1]]) as u32) << 16).to_le_bytes()) + .collect() +} + fn upload(g: &dyn GpuBackend, b: &[u8]) -> DevicePtr { let p = g.alloc(b.len().max(256)).unwrap(); g.copy_h2d_async(b, p, g.default_stream()).unwrap(); @@ -159,9 +168,11 @@ fn ple_kernels_match_reference() { let hidden = upload(g, &f.bytes("hidden")); let key = upload(g, &f.bytes("key_proj_out")); let value = upload(g, &f.bytes("value_proj_out")); - let nq = upload(g, &f.bytes("w_norm_query")); - let nk = upload(g, &f.bytes("w_norm_key")); - let nc = upload(g, &f.bytes("w_norm_conv")); + // The PLE norms feed `const float*` kernels; the fixture ships the + // checkpoint's BF16, so widen it the way `dense_keep_f32` does at load. + let nq = upload(g, &bf16_to_f32_bytes(&f.bytes("w_norm_query"))); + let nk = upload(g, &bf16_to_f32_bytes(&f.bytes("w_norm_key"))); + let nc = upload(g, &bf16_to_f32_bytes(&f.bytes("w_norm_conv"))); let gated = g.alloc(t * c * 4).unwrap(); // FP32: the conv reads this, and BF16 here costs the output 2.7%. let gated_normed = g.alloc(t * c * 4).unwrap(); diff --git a/crates/spark-model/src/layers/ops/qwen4exp_oracle_ple_tests.rs b/crates/spark-model/src/layers/ops/qwen4exp_oracle_ple_tests.rs index 188d8d92e..89394135d 100644 --- a/crates/spark-model/src/layers/ops/qwen4exp_oracle_ple_tests.rs +++ b/crates/spark-model/src/layers/ops/qwen4exp_oracle_ple_tests.rs @@ -116,9 +116,10 @@ fn qwen4exp_oracle_ple_matches_the_cpu_reference() { let d_hidden = upload(g, &f32_bytes(&hidden_host)); let d_key = upload(g, &bf16_bytes(&key_bf16)); let d_value = upload(g, &bf16_bytes(&value_bf16)); - let d_nq = upload(g, &bf16_bytes(&norm_query)); - let d_nk = upload(g, &bf16_bytes(&norm_key)); - let d_nc = upload(g, &bf16_bytes(&norm_conv)); + // The PLE norms feed `const float*` kernels; upload F32 like the loader. + let d_nq = upload(g, &f32_bytes(&norm_query)); + let d_nk = upload(g, &f32_bytes(&norm_key)); + let d_nc = upload(g, &f32_bytes(&norm_conv)); let d_conv_w = upload(g, &bf16_bytes(&conv1d)); let gated = g.alloc(tokens * wide * 4).unwrap(); let gated_normed = g.alloc(tokens * wide * 4).unwrap(); diff --git a/crates/spark-model/src/layers/ops/qwen4exp_oracle_tests.rs b/crates/spark-model/src/layers/ops/qwen4exp_oracle_tests.rs index 9f6ed3c2d..c093feff7 100644 --- a/crates/spark-model/src/layers/ops/qwen4exp_oracle_tests.rs +++ b/crates/spark-model/src/layers/ops/qwen4exp_oracle_tests.rs @@ -208,7 +208,10 @@ fn make_site(g: &dyn GpuBackend, rng: &mut Rng, inject: bool) -> Site { let inject_w = inject.then(|| bf16_round(&rng.vec(HC * wide, (wide as f32).sqrt().recip()))); Site { dev: HcLowRank { - norm_w: upload(g, &bf16_bytes(&norm)), + // `hc_norm` feeds `const float*` kernels; a bf16-rounded host + // value is exactly representable in F32, so this is the same + // weight the reference computes with. + norm_w: upload(g, &f32_bytes(&norm)), down_w: upload(g, &bf16_bytes(&down)), // `up` stays [wide, RANK] for the reference below; the kernels // read `input_mix_weight_up` transposed as [RANK, wide], which is diff --git a/crates/spark-model/src/layers/ple/layer.rs b/crates/spark-model/src/layers/ple/layer.rs index b2d4a1002..5077489a3 100644 --- a/crates/spark-model/src/layers/ple/layer.rs +++ b/crates/spark-model/src/layers/ple/layer.rs @@ -372,7 +372,9 @@ pub(crate) fn bounded_scratch(scratch: usize, max_tokens: usize) -> usize { scratch.clamp(VERIFY_SNAP_SLOTS.min(max_tokens), max_tokens) } -/// The dense weights of one PLE site. +/// The dense weights of one PLE site. The three norms are **F32** — `ple_gate` +/// is declared `const float*` and the loader (`dense_keep_f32`) guarantees the +/// dtype; everything else is BF16. pub struct PleWeights { pub key_proj: DenseWeight, pub value_proj: DenseWeight, diff --git a/crates/spark-model/src/layers/qwen3_attention/types_weights.rs b/crates/spark-model/src/layers/qwen3_attention/types_weights.rs index b42442226..4f126fd4e 100644 --- a/crates/spark-model/src/layers/qwen3_attention/types_weights.rs +++ b/crates/spark-model/src/layers/qwen3_attention/types_weights.rs @@ -170,10 +170,12 @@ pub struct CompressorWeights { /// Both share the `[T, hc_mult, H]` highway and `hc_mult`, so the presence of /// this struct — not the model name — is what selects the kernel. /// -/// All BF16, matching the checkpoint. See `bench/qwen4_exp/ARCHITECTURE.md` §1. +/// All BF16 except `norm_w` (F32). See `bench/qwen4_exp/ARCHITECTURE.md` §1. #[derive(Clone, Copy)] pub struct HcLowRank { - /// `hc_norm` `[hc_mult*hidden]`. A GROUPED RMSNorm scale: the streams + /// `hc_norm` `[hc_mult*hidden]` **F32** — every consumer kernel is + /// declared `const float*`; the loader (`dense_keep_f32`) guarantees the + /// dtype. A GROUPED RMSNorm scale: the streams /// normalize independently inside the vector, `group_size = hidden`. pub norm_w: DevicePtr, /// `input_mix_weight_down` `[rank, hc_mult*hidden]`. diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/hc.rs b/crates/spark-model/src/weight_loader/qwen4_exp/hc.rs index ec327b483..41733593c 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/hc.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/hc.rs @@ -29,7 +29,7 @@ use spark_runtime::kernel_args::KernelLaunch; use spark_runtime::weights::WeightStore; use crate::layers::qwen3_attention::{HcHeadWeights, HcLowRank, HcSiteWeights}; -use crate::weight_map::dense; +use crate::weight_map::{dense, dense_keep_f32}; /// Flips every `input_mix_weight_up` from the checkpoint's `[hc*H, rank]` to /// the `[rank, hc*H]` the collapse kernels read. @@ -127,8 +127,14 @@ fn load_site( let up_w = g("input_mix_weight_up")?; tr.apply(up_w) .with_context(|| format!("qwen4_exp mHC: transposing {prefix}.input_mix_weight_up"))?; + // `hc_norm` feeds kernels declared `const float*`: keep an F32 checkpoint + // native (FP4-Local ships it F32) and widen a BF16 one (NVFP4) by zero + // extension — exact in both directions. + let norm_w = dense_keep_f32(store, &format!("{prefix}.hc_norm.weight"), tr.gpu) + .map(|d| d.weight) + .with_context(|| format!("qwen4_exp mHC: {prefix}.hc_norm.weight"))?; Ok(HcLowRank { - norm_w: g("hc_norm")?, + norm_w, down_w: g("input_mix_weight_down")?, up_w, inject_w: if with_inject { diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs b/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs index 7ea179c21..8dacdbd9e 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs @@ -42,7 +42,7 @@ use crate::layers::ple::PleLayer; #[cfg(feature = "cuda")] use crate::layers::ple::{PleIdDims, PleWeights}; #[cfg(feature = "cuda")] -use crate::weight_map::dense; +use crate::weight_map::{dense, dense_keep_f32}; #[cfg(feature = "cuda")] use host::{bf16_scalar, fp16_host, i64_host}; @@ -318,9 +318,11 @@ pub(super) fn load( let weights = PleWeights { key_proj: dense(store, &format!("{lp}.key_proj.weight"))?, value_proj: dense(store, &format!("{lp}.value_proj.weight"))?, - norm_key: dense(store, &format!("{lp}.norm_key.weight"))?, - norm_query: dense(store, &format!("{lp}.norm_query.weight"))?, - norm_conv: dense(store, &format!("{lp}.norm_conv.weight"))?, + // The PLE norms feed `ple_gate`, declared `const float*`: keep F32 + // native (FP4-Local ships them F32), widen BF16 exactly (NVFP4). + norm_key: dense_keep_f32(store, &format!("{lp}.norm_key.weight"), gpu)?, + norm_query: dense_keep_f32(store, &format!("{lp}.norm_query.weight"), gpu)?, + norm_conv: dense_keep_f32(store, &format!("{lp}.norm_conv.weight"), gpu)?, conv1d: dense(store, &format!("{lp}.conv1d.weight"))?, }; diff --git a/crates/spark-model/src/weight_map/modelopt_local.rs b/crates/spark-model/src/weight_map/modelopt_local.rs index 5cdd1d2db..ca2b2fe19 100644 --- a/crates/spark-model/src/weight_map/modelopt_local.rs +++ b/crates/spark-model/src/weight_map/modelopt_local.rs @@ -15,7 +15,11 @@ //! E8M0, and read by `mxfp8_proj.cu`; the rest are expanded exactly to the //! BF16 `weight` the NVFP4 pack ships; //! * vision MLP/merger projections are W4A16 NVFP4 — dequantized to the BF16 -//! the vision tower reads. +//! the vision tower reads; +//! * the hyper-connection `hc_norm`s and the three PLE norms are F32 where +//! other releases ship BF16 — left exactly as shipped; the kernels that +//! read them are `const float*`, and the other packs' BF16 norms are +//! widened to F32 at layer-load (`dense_keep_f32`). //! //! The loader DEFERS every source tensor (`FastSafetensorsLoader:: //! host_decode_modules`, `is_csf_stream`), so the compressed originals are @@ -83,18 +87,15 @@ pub struct Materialized { pub mxfp8_native: usize, pub mxfp8_to_bf16: usize, pub nvfp4_to_bf16: usize, - pub f32_norms_to_bf16: usize, + pub f32_norms_native: usize, pub uploaded_bytes: u64, } /// A norm weight the FP4-Local pack keeps in F32 where every other release of /// the model ships BF16: the 97 hyper-connection `hc_norm`s and the three PLE /// norms (model card: "100 FP32 tensors"; 103 with the MTP block's). The -/// loaders hand these pointers straight to kernels that read BF16, so left -/// as F32 every normalization reads garbage — the server answers with one -/// empty token. Rounded to BF16 here, which is also what the pack's own vLLM -/// recipe (`--dtype bfloat16`) does; they sit within 0.07% of the NVFP4 -/// pack's BF16 values. +/// kernels reading them are declared `const float*`, so the resident tensor +/// is served as shipped — nothing to do here but count it for the log. fn is_f32_norm(name: &str, t: &WeightTensor) -> bool { // `hc_norm`, and the PLE's `norm_conv` / `norm_key` / `norm_query`: the // last module segment names a norm. @@ -105,36 +106,6 @@ fn is_f32_norm(name: &str, t: &WeightTensor) -> bool { t.dtype == WeightDtype::FP32 && t.shape.len() == 1 && names_norm } -fn norms_to_bf16(store: &mut WeightStore, gpu: &dyn GpuBackend) -> Result<(usize, u64)> { - let mut names: Vec = store - .names() - .filter(|n| store.get(n).is_ok_and(|t| is_f32_norm(n, t))) - .map(str::to_string) - .collect(); - names.sort(); - let mut bytes = 0; - for name in &names { - let src = take(store, gpu, name)?.context("norm vanished")?; - let bf16: Vec = src - .bytes - .chunks_exact(4) - .flat_map(|b| { - let v = f32::from_le_bytes([b[0], b[1], b[2], b[3]]); - atlas_core::numeric::f32_to_bf16(v).to_le_bytes() - }) - .collect(); - bytes += upload( - store, - gpu, - name.clone(), - &bf16, - src.shape, - WeightDtype::BF16, - )?; - } - Ok((names.len(), bytes)) -} - /// Raw bytes and geometry of a source tensor, from disk or from the device. struct Source { bytes: Vec, @@ -368,19 +339,20 @@ pub fn materialize( } if done != Materialized::default() { // Only an FP4-Local checkpoint reaches here: another release's F32 - // tensors are never touched. - let (n, b) = norms_to_bf16(store, gpu)?; - done.f32_norms_to_bf16 = n; - done.uploaded_bytes += b; + // tensors are never counted. + done.f32_norms_native = store + .names() + .filter(|n| store.get(n).is_ok_and(|t| is_f32_norm(n, t))) + .count(); tracing::info!( "FP4-Local: decoded {} CSF scale tensors, kept {} MXFP8 projections native, \ - expanded {} MXFP8 and {} NVFP4 projections to BF16, rounded {} F32 norms to \ - BF16 ({:.2} GiB uploaded; compressed originals never resident)", + expanded {} MXFP8 and {} NVFP4 projections to BF16, kept {} F32 norms \ + native ({:.2} GiB uploaded; compressed originals never resident)", done.csf_scales, done.mxfp8_native, done.mxfp8_to_bf16, done.nvfp4_to_bf16, - done.f32_norms_to_bf16, + done.f32_norms_native, done.uploaded_bytes as f64 / (1u64 << 30) as f64 ); } diff --git a/crates/spark-model/src/weight_map/modelopt_local_tests.rs b/crates/spark-model/src/weight_map/modelopt_local_tests.rs index 33787e059..1241cf748 100644 --- a/crates/spark-model/src/weight_map/modelopt_local_tests.rs +++ b/crates/spark-model/src/weight_map/modelopt_local_tests.rs @@ -227,14 +227,29 @@ fn materialize_expands_every_fp4_local_format() { done.mxfp8_native, done.mxfp8_to_bf16, done.nvfp4_to_bf16, - done.f32_norms_to_bf16 + done.f32_norms_native ), (1, 0, 1, 1, 2) ); + // The F32 norms are served as shipped: same dtype, byte-identical data. let norm = store.get(HC_NORM).unwrap(); - assert_eq!(norm.dtype, WeightDtype::BF16); - assert_eq!(gpu.read_alloc(norm.ptr).unwrap(), [0x80, 0x3F, 0x00, 0x40]); // 1.0, 2.0 - assert_eq!(store.get(PLE_NORM).unwrap().dtype, WeightDtype::BF16); + assert_eq!(norm.dtype, WeightDtype::FP32); + assert_eq!( + gpu.read_alloc(norm.ptr).unwrap(), + [1.0f32, 2.0] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>() + ); + let norm = store.get(PLE_NORM).unwrap(); + assert_eq!(norm.dtype, WeightDtype::FP32); + assert_eq!( + gpu.read_alloc(norm.ptr).unwrap(), + [1.0f32, 2.0] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>() + ); let bias = store .get("model.language_model.layers.1.ple.conv1d.bias") .unwrap(); diff --git a/crates/spark-model/src/weight_map/tests.rs b/crates/spark-model/src/weight_map/tests.rs index ddab38048..080b8bbc5 100644 --- a/crates/spark-model/src/weight_map/tests.rs +++ b/crates/spark-model/src/weight_map/tests.rs @@ -6,7 +6,7 @@ use anyhow::{Context, Result, bail, ensure}; use spark_runtime::gpu::{DevicePtr, GpuBackend}; -use spark_runtime::weights::{WeightDtype, WeightStore}; +use spark_runtime::weights::{WeightDtype, WeightStore, WeightTensor}; use super::*; @@ -15,6 +15,57 @@ use super::*; // tests live beside them there. The identical copy that used to sit here // was deleted rather than kept in sync by hand. +/// `dense_keep_f32` is what the qwen4_exp loaders hand to the `const float*` +/// norm consumers (`hc_norm`, the PLE norms): an F32 source passes through +/// untouched and a BF16 one is widened by zero extension — `(bf16 as u32) +/// << 16`, exact, no rounding. +#[test] +fn dense_keep_f32_passes_f32_and_widens_bf16() { + use spark_runtime::gpu::mock::MockGpuBackend; + + fn put(store: &mut WeightStore, gpu: &MockGpuBackend, name: &str, bytes: &[u8], dtype: WeightDtype) { + let ptr = gpu.alloc(bytes.len()).unwrap(); + gpu.copy_h2d(bytes, ptr).unwrap(); + store + .insert( + name.into(), + WeightTensor { + ptr, + shape: vec![bytes.len() / dtype.byte_size()], + dtype, + }, + ) + .unwrap(); + } + + let gpu = MockGpuBackend::new(); + let mut store = WeightStore::empty(); + + // BF16 [1.0, -2.5, 0x1.fp-4]: each lands as (bits << 16). + let bf: Vec = [0x3F80u16, 0xC020, 0x3D00] + .iter() + .flat_map(|b| b.to_le_bytes()) + .collect(); + put(&mut store, &gpu, "l.attn_hyper_connection.hc_norm.weight", &bf, WeightDtype::BF16); + let w = dense_keep_f32(&store, "l.attn_hyper_connection.hc_norm.weight", &gpu).unwrap(); + let got = gpu.read_alloc(w.weight).unwrap(); + let want: Vec = [0x3F80u32 << 16, 0xC020 << 16, 0x3D00 << 16] + .iter() + .flat_map(|b| b.to_le_bytes()) + .collect(); + assert_eq!(got, want); + + // F32 source: the SAME pointer comes back, no copy. + let f: Vec = [1.5f32, -0.25] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + put(&mut store, &gpu, "l.ple.norm_key.weight", &f, WeightDtype::FP32); + let src = store.get("l.ple.norm_key.weight").unwrap().ptr; + let w = dense_keep_f32(&store, "l.ple.norm_key.weight", &gpu).unwrap(); + assert_eq!(w.weight, src); +} + #[test] fn test_weight_name_patterns() { // Verify our name generation matches actual HF patterns. diff --git a/kernels/gb10/qwen3.8-flash-next/nvfp4/hyper_connection.cu b/kernels/gb10/qwen3.8-flash-next/nvfp4/hyper_connection.cu index dce826d2d..6e53f69f9 100644 --- a/kernels/gb10/qwen3.8-flash-next/nvfp4/hyper_connection.cu +++ b/kernels/gb10/qwen3.8-flash-next/nvfp4/hyper_connection.cu @@ -144,7 +144,7 @@ extern "C" __global__ void hc_expand( __device__ __forceinline__ void qhc_collapse( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, + const float* __restrict__ hc_norm_w, const __nv_bfloat16* __restrict__ down_w, const __nv_bfloat16* __restrict__ up_w, const __nv_bfloat16* __restrict__ inject_w, @@ -193,7 +193,7 @@ __device__ __forceinline__ void qhc_collapse( // ── stage normed = x * rms * (1 + w) once ── for (unsigned int i = tid; i < hc_dim; i += blockDim.x) { - smem_normed[i] = x[i] * smem_rms[i / H] * (1.0f + (float)hc_norm_w[i]); + smem_normed[i] = x[i] * smem_rms[i / H] * (1.0f + hc_norm_w[i]); } __syncthreads(); @@ -252,7 +252,7 @@ __device__ __forceinline__ void qhc_collapse( // streams [T, hc, H] -> y_out [T, H] collapsed, inj_out [T, hc]. extern "C" __global__ void hc_pre( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, // [hc*H] + const float* __restrict__ hc_norm_w, // [hc*H] const __nv_bfloat16* __restrict__ down_w, // [rank, hc*H] const __nv_bfloat16* __restrict__ up_w, // [rank, hc*H] const __nv_bfloat16* __restrict__ inject_w, // [hc, hc*H] @@ -274,7 +274,7 @@ extern "C" __global__ void hc_pre( // role. extern "C" __global__ void hc_head( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, + const float* __restrict__ hc_norm_w, const __nv_bfloat16* __restrict__ down_w, const __nv_bfloat16* __restrict__ up_w, __nv_bfloat16* __restrict__ y_out, @@ -345,7 +345,7 @@ extern "C" __global__ void hc_post( // Stage 1: normed = x * rms * (1 + w) -> global scratch [T, hc*H]. extern "C" __global__ void hc_pre_stage( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, + const float* __restrict__ hc_norm_w, float* __restrict__ normed_out, // [T, hc*H] const unsigned int hidden_size, const unsigned int hc, @@ -385,7 +385,7 @@ extern "C" __global__ void hc_pre_stage( __syncthreads(); } for (unsigned int i = tid; i < hc_dim; i += blockDim.x) { - out[i] = x[i] * smem_rms[i / H] * (1.0f + (float)hc_norm_w[i]); + out[i] = x[i] * smem_rms[i / H] * (1.0f + hc_norm_w[i]); } } @@ -848,7 +848,7 @@ extern "C" __global__ void hc_pre_finish_x4( extern "C" __global__ void hc_pre_stage_bf16( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, + const float* __restrict__ hc_norm_w, __nv_bfloat16* __restrict__ normed_out, // [T, hc*H] BF16 const unsigned int hidden_size, const unsigned int hc, @@ -889,7 +889,7 @@ extern "C" __global__ void hc_pre_stage_bf16( } for (unsigned int i = tid; i < hc_dim; i += blockDim.x) { out[i] = __float2bfloat16( - x[i] * smem_rms[i / H] * (1.0f + (float)hc_norm_w[i])); + x[i] * smem_rms[i / H] * (1.0f + hc_norm_w[i])); } } @@ -918,7 +918,7 @@ extern "C" __global__ void __launch_bounds__(QHC_WBLOCK, 1) hc_post_stage_bf16( const __nv_bfloat16* __restrict__ block_out, // [T, H] float* streams, // [T, hc, H]: residual in, post out const float* __restrict__ inj, // [T, hc] - const __nv_bfloat16* __restrict__ hc_norm_w, // [hc*H] + const float* __restrict__ hc_norm_w, // [hc*H] __nv_bfloat16* __restrict__ normed_out, // [T, hc*H] BF16 const unsigned int hidden_size, const unsigned int hc, @@ -988,7 +988,7 @@ extern "C" __global__ void __launch_bounds__(QHC_WBLOCK, 1) hc_post_stage_bf16( if (s2 < hc) { const unsigned int i = s2 * H + d; out[i] = __float2bfloat16( - v[k][s2] * smem_rms[s2] * (1.0f + (float)hc_norm_w[i])); + v[k][s2] * smem_rms[s2] * (1.0f + hc_norm_w[i])); } } } @@ -2102,7 +2102,7 @@ extern "C" __global__ void hc_pre_down_vec_wide( // columns [y * hc_dim / S, (y + 1) * hc_dim / S) of `normed`. extern "C" __global__ void hc_pre_stage_vec( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, + const float* __restrict__ hc_norm_w, float* __restrict__ normed_out, // [T, hc*H] const unsigned int hidden_size, const unsigned int hc, @@ -2161,13 +2161,13 @@ extern "C" __global__ void hc_pre_stage_vec( const unsigned int c1 = blockIdx.y + 1 == gridDim.y ? hc_dim : c0 + span; for (unsigned int i = c0 + 4u * tid; i < c1; i += 4u * QHC_WBLOCK) { const float4 xv = *reinterpret_cast(x + i); - const uint2 wv = *reinterpret_cast(hc_norm_w + i); + const float4 wv = *reinterpret_cast(hc_norm_w + i); const float rms = smem_rms[i / H]; float4 o; - o.x = xv.x * rms * (1.0f + qhc_bf_lo(wv.x)); - o.y = xv.y * rms * (1.0f + qhc_bf_hi(wv.x)); - o.z = xv.z * rms * (1.0f + qhc_bf_lo(wv.y)); - o.w = xv.w * rms * (1.0f + qhc_bf_hi(wv.y)); + o.x = xv.x * rms * (1.0f + wv.x); + o.y = xv.y * rms * (1.0f + wv.y); + o.z = xv.z * rms * (1.0f + wv.z); + o.w = xv.w * rms * (1.0f + wv.w); *reinterpret_cast(out + i) = o; } } @@ -2234,7 +2234,7 @@ hc_post_stage_vec( const __nv_bfloat16* __restrict__ block_out, // [T, H] float* streams, // [T, hc, H]: residual in, post out const float* __restrict__ inj, // [T, hc] - const __nv_bfloat16* __restrict__ hc_norm_w, // [hc*H] + const float* __restrict__ hc_norm_w, // [hc*H] float* __restrict__ normed_out, // [T, hc*H] const unsigned int hidden_size, const unsigned int hc, @@ -2307,13 +2307,13 @@ hc_post_stage_vec( v.z = r.z + qhc_bf_lo(xv.y) * ws; v.w = r.w + qhc_bf_hi(xv.y) * ws; *reinterpret_cast(x + i) = v; - const uint2 nw = *reinterpret_cast(hc_norm_w + i); + const float4 nw = *reinterpret_cast(hc_norm_w + i); const float rms = smem_rms[s]; float4 o; - o.x = v.x * rms * (1.0f + qhc_bf_lo(nw.x)); - o.y = v.y * rms * (1.0f + qhc_bf_hi(nw.x)); - o.z = v.z * rms * (1.0f + qhc_bf_lo(nw.y)); - o.w = v.w * rms * (1.0f + qhc_bf_hi(nw.y)); + o.x = v.x * rms * (1.0f + nw.x); + o.y = v.y * rms * (1.0f + nw.y); + o.z = v.z * rms * (1.0f + nw.z); + o.w = v.w * rms * (1.0f + nw.w); *reinterpret_cast(out + i) = o; } } diff --git a/kernels/gb10/qwen3.8-flash-next/nvfp4/ple.cu b/kernels/gb10/qwen3.8-flash-next/nvfp4/ple.cu index 199cc4b76..bb2d194ec 100644 --- a/kernels/gb10/qwen3.8-flash-next/nvfp4/ple.cu +++ b/kernels/gb10/qwen3.8-flash-next/nvfp4/ple.cu @@ -91,9 +91,9 @@ extern "C" __global__ void ple_gate( const float* __restrict__ hidden, // [T, hc*H] FP32 highway const __nv_bfloat16* __restrict__ key, // [T, hc*H] key_proj out const __nv_bfloat16* __restrict__ value, // [T, H] value_proj out - const __nv_bfloat16* __restrict__ norm_query_w, // [hc*H] - const __nv_bfloat16* __restrict__ norm_key_w, // [hc*H] - const __nv_bfloat16* __restrict__ norm_conv_w, // [hc*H] + const float* __restrict__ norm_query_w, // [hc*H] + const float* __restrict__ norm_key_w, // [hc*H] + const float* __restrict__ norm_conv_w, // [hc*H] float* __restrict__ gated_out, // [T, hc*H] FP32 float* __restrict__ gated_normed, // [T, hc*H] FP32 — see below const unsigned int hidden_size, // H @@ -137,8 +137,8 @@ extern "C" __global__ void ple_gate( float dot = 0.0f; for (unsigned int d = tid; d < H; d += PLE_BLOCK) { const unsigned int i = s * H + d; - float qn = qs[d] * rq * (1.0f + (float)norm_query_w[i]); - float kn = (float)ks[d] * rk * (1.0f + (float)norm_key_w[i]); + float qn = qs[d] * rq * (1.0f + norm_query_w[i]); + float kn = (float)ks[d] * rk * (1.0f + norm_key_w[i]); dot += qn * kn; } dot = ple_block_sum(dot, smem_red); @@ -172,7 +172,7 @@ extern "C" __global__ void ple_gate( for (unsigned int d = tid; d < H; d += PLE_BLOCK) { const unsigned int i = s * H + d; float x = g_out[i]; - gn_out[i] = x * r * (1.0f + (float)norm_conv_w[i]); + gn_out[i] = x * r * (1.0f + norm_conv_w[i]); } __syncthreads(); } diff --git a/kernels/strix-hip/qwen3.8-flash-next/nvfp4/hyper_connection.cu b/kernels/strix-hip/qwen3.8-flash-next/nvfp4/hyper_connection.cu index 43752f15e..bb92ec65e 100644 --- a/kernels/strix-hip/qwen3.8-flash-next/nvfp4/hyper_connection.cu +++ b/kernels/strix-hip/qwen3.8-flash-next/nvfp4/hyper_connection.cu @@ -142,7 +142,7 @@ extern "C" __global__ void hc_expand( __device__ __forceinline__ void qhc_collapse( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, + const float* __restrict__ hc_norm_w, const __nv_bfloat16* __restrict__ down_w, const __nv_bfloat16* __restrict__ up_w, const __nv_bfloat16* __restrict__ inject_w, @@ -191,7 +191,7 @@ __device__ __forceinline__ void qhc_collapse( // ── stage normed = x * rms * (1 + w) once ── for (unsigned int i = tid; i < hc_dim; i += blockDim.x) { - smem_normed[i] = x[i] * smem_rms[i / H] * (1.0f + (float)hc_norm_w[i]); + smem_normed[i] = x[i] * smem_rms[i / H] * (1.0f + hc_norm_w[i]); } __syncthreads(); @@ -250,7 +250,7 @@ __device__ __forceinline__ void qhc_collapse( // streams [T, hc, H] -> y_out [T, H] collapsed, inj_out [T, hc]. extern "C" __global__ void hc_pre( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, // [hc*H] + const float* __restrict__ hc_norm_w, // [hc*H] const __nv_bfloat16* __restrict__ down_w, // [rank, hc*H] const __nv_bfloat16* __restrict__ up_w, // [rank, hc*H] const __nv_bfloat16* __restrict__ inject_w, // [hc, hc*H] @@ -272,7 +272,7 @@ extern "C" __global__ void hc_pre( // role. extern "C" __global__ void hc_head( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, + const float* __restrict__ hc_norm_w, const __nv_bfloat16* __restrict__ down_w, const __nv_bfloat16* __restrict__ up_w, __nv_bfloat16* __restrict__ y_out, @@ -328,7 +328,7 @@ extern "C" __global__ void hc_post( // Stage 1: normed = x * rms * (1 + w) -> global scratch [T, hc*H]. extern "C" __global__ void hc_pre_stage( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, + const float* __restrict__ hc_norm_w, float* __restrict__ normed_out, // [T, hc*H] const unsigned int hidden_size, const unsigned int hc, @@ -368,7 +368,7 @@ extern "C" __global__ void hc_pre_stage( __syncthreads(); } for (unsigned int i = tid; i < hc_dim; i += blockDim.x) { - out[i] = x[i] * smem_rms[i / H] * (1.0f + (float)hc_norm_w[i]); + out[i] = x[i] * smem_rms[i / H] * (1.0f + hc_norm_w[i]); } } @@ -729,7 +729,7 @@ extern "C" __global__ void hc_pre_finish_x4( extern "C" __global__ void hc_pre_stage_bf16( const float* __restrict__ streams, - const __nv_bfloat16* __restrict__ hc_norm_w, + const float* __restrict__ hc_norm_w, __nv_bfloat16* __restrict__ normed_out, // [T, hc*H] BF16 const unsigned int hidden_size, const unsigned int hc, @@ -770,7 +770,7 @@ extern "C" __global__ void hc_pre_stage_bf16( } for (unsigned int i = tid; i < hc_dim; i += blockDim.x) { out[i] = __float2bfloat16( - x[i] * smem_rms[i / H] * (1.0f + (float)hc_norm_w[i])); + x[i] * smem_rms[i / H] * (1.0f + hc_norm_w[i])); } } From 62ebdf25125cdaeea0b684c86d45dd45c889f112 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 04:15:47 -0400 Subject: [PATCH 08/26] qwen4_exp: keep the FP4-Local shared expert on its native MXFP8 bytes The FP4-Local checkpoint ships every main-layer mlp.shared_expert {gate,up,down}_proj as MXFP8 (E4M3 + E8M0 scales). The loader used to expand them to BF16 and requantize them to NVFP4 for the fused routed kernels -- an 8-bit -> 4-bit loss on the one expert every token goes through. Generalize the Laguna BF16 seam into a separate shared-expert enum (Bf16 | Mxfp8): fused routed kernels run with a NULL shared expert and the native chain (mxfp8_proj + silu_mul) fills shared_out afterward. Gate and up are row-concatenated into one [2*inter, h] projection so the <=32-row bands take one pass plus fused_silu_mul; larger prefills use the unfused trio. Covers decode, MTP K=2 verify, C>1 batched decode (forward_rows, qwen4exp_fast, atomic_c4) and prefill. The MTP block's BF16 shared expert is installed through the BF16 arm. Tracked weight memory: +0.15 GB vs the NVFP4 shared copy. --- crates/spark-model/src/layers/moe/forward.rs | 8 +- .../src/layers/moe/forward_atomic_c4.rs | 25 +- .../src/layers/moe/forward_batched.rs | 8 +- .../spark-model/src/layers/moe/forward_k2.rs | 8 +- .../src/layers/moe/forward_k2/originals.rs | 2 +- .../src/layers/moe/forward_k2/unified_t.rs | 2 +- .../spark-model/src/layers/moe/forward_k3.rs | 6 +- .../src/layers/moe/forward_k3_shared.rs | 2 +- .../spark-model/src/layers/moe/forward_k4.rs | 2 +- .../spark-model/src/layers/moe/forward_k5.rs | 2 +- .../src/layers/moe/forward_prefill_bf16.rs | 4 +- .../src/layers/moe/forward_prefill_fp8.rs | 2 +- .../src/layers/moe/forward_prefill_phase.rs | 2 +- .../src/layers/moe/forward_prefill_q38.rs | 2 +- .../src/layers/moe/forward_rows.rs | 18 +- .../src/layers/moe/forward_token_major.rs | 2 +- .../src/layers/moe/gate_up_btile_preflight.rs | 2 +- .../moe/gate_up_btile_resident_arena.rs | 2 +- .../spark-model/src/layers/moe/helpers_c.rs | 127 +--------- .../src/layers/moe/helpers_shared.rs | 228 ++++++++++++++++++ crates/spark-model/src/layers/moe/init.rs | 2 +- crates/spark-model/src/layers/moe/mod.rs | 35 ++- .../spark-model/src/layers/moe/mod_tests.rs | 122 +++++++++- .../src/layers/moe/prequant_fp4.rs | 2 +- .../src/layers/moe/qwen4exp_fast.rs | 33 ++- .../src/layers/moe/shared_fp8_cache_load.rs | 2 +- .../spark-model/src/layers/moe/shared_m16.rs | 2 +- crates/spark-model/src/layers/moe/types.rs | 7 +- .../spark-model/src/layers/ops/activations.rs | 23 ++ .../src/weight_loader/qwen35/load_layers.rs | 1 + .../src/weight_loader/qwen4_exp.rs | 2 +- .../src/weight_loader/qwen4_exp/ffn.rs | 90 ++++++- .../src/weight_loader/qwen4_exp/mtp.rs | 30 +-- .../src/weight_map/modelopt_local.rs | 5 +- .../src/weight_map/modelopt_local_tests.rs | 23 +- .../spark-model/src/weight_map/ssm_qwen35.rs | 10 +- crates/spark-model/src/weight_map/tests.rs | 24 +- 37 files changed, 659 insertions(+), 208 deletions(-) create mode 100644 crates/spark-model/src/layers/moe/helpers_shared.rs diff --git a/crates/spark-model/src/layers/moe/forward.rs b/crates/spark-model/src/layers/moe/forward.rs index dc66a5c5f..1dc453b7d 100644 --- a/crates/spark-model/src/layers/moe/forward.rs +++ b/crates/spark-model/src/layers/moe/forward.rs @@ -430,11 +430,11 @@ impl MoeLayer { let shared_gate_scratch = ctx.buffers.logits(); let shared_up_scratch = ctx.buffers.ssm_qkvz(); - if let (Some(gp), Some(up), Some(dp), Some(shared)) = ( + if let (Some(gp), Some(up), Some(dp), Some(SeparateSharedExpert::Bf16(shared))) = ( self.bf16_gate_weight_ptrs, self.bf16_up_weight_ptrs, self.bf16_down_weight_ptrs, - self.bf16_shared_expert, + self.separate_shared_expert, ) { // BF16 path: FP8-dequant-on-load. Eliminates the per-layer 0.989 // FP8 cosine ceiling by serving experts as BF16 end-to-end. @@ -730,8 +730,8 @@ impl MoeLayer { )?; } - if self.has_mixed_bf16_shared_expert() { - self.run_bf16_shared_expert( + if self.has_separate_shared_expert() { + self.run_separate_shared_expert( input, 1, h, diff --git a/crates/spark-model/src/layers/moe/forward_atomic_c4.rs b/crates/spark-model/src/layers/moe/forward_atomic_c4.rs index 88916fda5..2a2e8dd9a 100644 --- a/crates/spark-model/src/layers/moe/forward_atomic_c4.rs +++ b/crates/spark-model/src/layers/moe/forward_atomic_c4.rs @@ -28,14 +28,16 @@ impl MoeLayer { ); // Feature-1 phase-1: decode does not yet fold the expert delta. self.reject_decode_lora(ctx, "forward_atomic_c4_decode")?; - let has_shared = self.weights.shared_expert.gate_proj.weight.0 != 0 + // A checkpoint-native separate shared expert still counts: the fused + // kernels NULL-guard it and `shared_down_out` is filled separately. + let has_shared = (self.weights.shared_expert.gate_proj.weight.0 != 0 && self.weights.shared_expert.up_proj.weight.0 != 0 - && self.weights.shared_expert.down_proj.weight.0 != 0; + && self.weights.shared_expert.down_proj.weight.0 != 0) + || self.has_separate_shared_expert(); let nvfp4_supported = num_tokens == 4 && self.moe_decode_atomic_c4_silu_down_accum_k.0 != 0 && self.moe_decode_atomic_c4_finalize_k.0 != 0 && self.bf16_gate_weight_ptrs.is_none() - && !self.has_mixed_bf16_shared_expert() && self.fp8_gate_weight_ptrs.is_none() && !self.use_t_layout_for_decode() && self.pre_expert_norm.is_none() @@ -176,6 +178,23 @@ impl MoeLayer { stream, )?; + // Checkpoint-native separate shared expert: the fused kernels saw + // NULL weights (they skip it); fill `shared_down_out` for all 4 rows + // before the finalize/blend reads it. + if self.has_separate_shared_expert() { + self.run_separate_shared_expert( + input, + n, + h, + ctx.config.shared_expert_intermediate_size as u32, + shared_gate_scratch, + shared_up_scratch, + shared_down_out, + ctx, + stream, + )?; + } + let is_ep = ctx.comm.is_some() && ctx.config.ep_world_size > 1; ops::moe_decode_atomic_c4_finalize( ctx.gpu, diff --git a/crates/spark-model/src/layers/moe/forward_batched.rs b/crates/spark-model/src/layers/moe/forward_batched.rs index 632af4e70..d093479be 100644 --- a/crates/spark-model/src/layers/moe/forward_batched.rs +++ b/crates/spark-model/src/layers/moe/forward_batched.rs @@ -170,11 +170,11 @@ impl MoeLayer { } let shared_out = ctx.buffers.attn_output(); - if let (Some(gp), Some(up), Some(dp), Some(shared)) = ( + if let (Some(gp), Some(up), Some(dp), Some(SeparateSharedExpert::Bf16(shared))) = ( self.bf16_gate_weight_ptrs, self.bf16_up_weight_ptrs, self.bf16_down_weight_ptrs, - self.bf16_shared_expert, + self.separate_shared_expert, ) { // BF16 path (FP8-dequant-on-load): same fused kernels as decode. ops::moe_expert_gate_up_shared_bf16( @@ -454,8 +454,8 @@ impl MoeLayer { stream, )?; - if self.has_mixed_bf16_shared_expert() { - self.run_bf16_shared_expert( + if self.has_separate_shared_expert() { + self.run_separate_shared_expert( input_t, 1, h, diff --git a/crates/spark-model/src/layers/moe/forward_k2.rs b/crates/spark-model/src/layers/moe/forward_k2.rs index 3a8728553..87e70d98e 100644 --- a/crates/spark-model/src/layers/moe/forward_k2.rs +++ b/crates/spark-model/src/layers/moe/forward_k2.rs @@ -70,7 +70,7 @@ impl MoeLayer { self.lora.is_none() && self.bf16_gate_weight_ptrs.is_none() && self.fp8_gate_weight_ptrs.is_none() - && !self.has_mixed_bf16_shared_expert() + && !self.has_separate_shared_expert() && matches!( self.experts_scale_kind, crate::weight_map::WeightQuantFormat::Nvfp4 @@ -129,7 +129,7 @@ impl MoeLayer { // through the _t kernels and run the shared expert as one batched BF16 // GEMM pass afterwards (`mixed_bf16_shared` below). Every other layout // falls back to the per-token loop. - let mixed_bf16_shared = self.has_mixed_bf16_shared_expert(); + let mixed_bf16_shared = self.has_separate_shared_expert(); // Either fused layout serves the mixed config. The originals-layout // kernels are usable only since 37e818ad NULL-guarded their shared // expert (their `_t` siblings always had that guard); before it, this @@ -182,11 +182,11 @@ impl MoeLayer { let output = ctx.buffers.moe_output(); if use_bf16_batch2 - && let (Some(gp), Some(up), Some(dp), Some(shared)) = ( + && let (Some(gp), Some(up), Some(dp), Some(SeparateSharedExpert::Bf16(shared))) = ( self.bf16_gate_weight_ptrs, self.bf16_up_weight_ptrs, self.bf16_down_weight_ptrs, - self.bf16_shared_expert, + self.separate_shared_expert, ) { // BF16 batch2 path (FP8-dequant-on-load experts, MTP K=2 verify). diff --git a/crates/spark-model/src/layers/moe/forward_k2/originals.rs b/crates/spark-model/src/layers/moe/forward_k2/originals.rs index b707b6fac..8fe1bf227 100644 --- a/crates/spark-model/src/layers/moe/forward_k2/originals.rs +++ b/crates/spark-model/src/layers/moe/forward_k2/originals.rs @@ -92,7 +92,7 @@ impl MoeLayer { // replacing the placeholder the kernel was told (via NULL) to skip. if mixed_bf16_shared { let shared_inter = ctx.config.shared_expert_intermediate_size as u32; - self.run_bf16_shared_expert( + self.run_separate_shared_expert( input, 2, h, diff --git a/crates/spark-model/src/layers/moe/forward_k2/unified_t.rs b/crates/spark-model/src/layers/moe/forward_k2/unified_t.rs index 7f9d3a6a6..85a32f3a3 100644 --- a/crates/spark-model/src/layers/moe/forward_k2/unified_t.rs +++ b/crates/spark-model/src/layers/moe/forward_k2/unified_t.rs @@ -156,7 +156,7 @@ impl MoeLayer { // shared scratch buffers for the non-mixed case. if mixed_bf16_shared { let shared_inter = ctx.config.shared_expert_intermediate_size as u32; - self.run_bf16_shared_expert( + self.run_separate_shared_expert( input, 2, h, diff --git a/crates/spark-model/src/layers/moe/forward_k3.rs b/crates/spark-model/src/layers/moe/forward_k3.rs index 48647d431..cc35be8ad 100644 --- a/crates/spark-model/src/layers/moe/forward_k3.rs +++ b/crates/spark-model/src/layers/moe/forward_k3.rs @@ -41,7 +41,7 @@ impl MoeLayer { self.lora.is_none() && self.bf16_gate_weight_ptrs.is_none() && self.fp8_gate_weight_ptrs.is_none() - && !self.has_mixed_bf16_shared_expert() + && !self.has_separate_shared_expert() && matches!( self.experts_scale_kind, crate::weight_map::WeightQuantFormat::Nvfp4 @@ -69,7 +69,7 @@ impl MoeLayer { // Mixed NVFP4-routed / BF16-shared (Laguna): batch the routed half // through the _t kernels and run the shared expert as one batched BF16 // pass afterwards. See forward_k2 for the rationale. - let mixed_bf16_shared = self.has_mixed_bf16_shared_expert(); + let mixed_bf16_shared = self.has_separate_shared_expert(); if mixed_bf16_shared && !(self.use_btile_or_t_decode() && self.moe_expert_gate_up_shared_batch3_t_k.0 != 0 @@ -302,7 +302,7 @@ impl MoeLayer { } if mixed_bf16_shared { let shared_inter = ctx.config.shared_expert_intermediate_size as u32; - self.run_bf16_shared_expert( + self.run_separate_shared_expert( input, 3, h, diff --git a/crates/spark-model/src/layers/moe/forward_k3_shared.rs b/crates/spark-model/src/layers/moe/forward_k3_shared.rs index c23fe8527..7501220af 100644 --- a/crates/spark-model/src/layers/moe/forward_k3_shared.rs +++ b/crates/spark-model/src/layers/moe/forward_k3_shared.rs @@ -20,7 +20,7 @@ impl MoeLayer { let optimized = self.lora.is_none() && self.bf16_gate_weight_ptrs.is_none() && self.fp8_gate_weight_ptrs.is_none() - && !self.has_mixed_bf16_shared_expert() + && !self.has_separate_shared_expert() && matches!( self.experts_scale_kind, crate::weight_map::WeightQuantFormat::Nvfp4 diff --git a/crates/spark-model/src/layers/moe/forward_k4.rs b/crates/spark-model/src/layers/moe/forward_k4.rs index 70e5ae364..bd8592097 100644 --- a/crates/spark-model/src/layers/moe/forward_k4.rs +++ b/crates/spark-model/src/layers/moe/forward_k4.rs @@ -20,7 +20,7 @@ impl MoeLayer { let optimized = self.lora.is_none() && self.bf16_gate_weight_ptrs.is_none() && self.fp8_gate_weight_ptrs.is_none() - && !self.has_mixed_bf16_shared_expert() + && !self.has_separate_shared_expert() && matches!( self.experts_scale_kind, crate::weight_map::WeightQuantFormat::Nvfp4 diff --git a/crates/spark-model/src/layers/moe/forward_k5.rs b/crates/spark-model/src/layers/moe/forward_k5.rs index 5ef6488ad..b7e02bce0 100644 --- a/crates/spark-model/src/layers/moe/forward_k5.rs +++ b/crates/spark-model/src/layers/moe/forward_k5.rs @@ -105,7 +105,7 @@ impl MoeLayer { let optimized = self.lora.is_none() && self.bf16_gate_weight_ptrs.is_none() && self.fp8_gate_weight_ptrs.is_none() - && !self.has_mixed_bf16_shared_expert() + && !self.has_separate_shared_expert() && matches!( self.experts_scale_kind, crate::weight_map::WeightQuantFormat::Nvfp4 diff --git a/crates/spark-model/src/layers/moe/forward_prefill_bf16.rs b/crates/spark-model/src/layers/moe/forward_prefill_bf16.rs index 2ed560659..5ec4a9ce8 100644 --- a/crates/spark-model/src/layers/moe/forward_prefill_bf16.rs +++ b/crates/spark-model/src/layers/moe/forward_prefill_bf16.rs @@ -34,9 +34,9 @@ impl MoeLayer { }; // ── Shared expert (BF16 dense GEMM) ── - let has_shared = shared_inter > 0 && self.bf16_shared_expert.is_some(); + let has_shared = shared_inter > 0 && self.separate_shared_expert.is_some(); if has_shared { - self.run_bf16_shared_expert( + self.run_separate_shared_expert( input, n, h, diff --git a/crates/spark-model/src/layers/moe/forward_prefill_fp8.rs b/crates/spark-model/src/layers/moe/forward_prefill_fp8.rs index b67b801e0..5daded223 100644 --- a/crates/spark-model/src/layers/moe/forward_prefill_fp8.rs +++ b/crates/spark-model/src/layers/moe/forward_prefill_fp8.rs @@ -96,7 +96,7 @@ impl MoeLayer { && self.per_token_group_quant_fp8_k.0 != 0; let has_shared = shared_inter > 0; let bf16_shared = has_shared - && self.run_bf16_shared_expert( + && self.run_separate_shared_expert( input, n, h, diff --git a/crates/spark-model/src/layers/moe/forward_prefill_phase.rs b/crates/spark-model/src/layers/moe/forward_prefill_phase.rs index 47a5dd488..a5ed40537 100644 --- a/crates/spark-model/src/layers/moe/forward_prefill_phase.rs +++ b/crates/spark-model/src/layers/moe/forward_prefill_phase.rs @@ -94,7 +94,7 @@ impl MoeLayer { } return Ok(()); } - if self.run_bf16_shared_expert( + if self.run_separate_shared_expert( input, n, h, diff --git a/crates/spark-model/src/layers/moe/forward_prefill_q38.rs b/crates/spark-model/src/layers/moe/forward_prefill_q38.rs index 6b2e6719a..c9821ccfb 100644 --- a/crates/spark-model/src/layers/moe/forward_prefill_q38.rs +++ b/crates/spark-model/src/layers/moe/forward_prefill_q38.rs @@ -189,7 +189,7 @@ impl MoeLayer { && self.shared_gate_fp8.is_none() && self.shared_up_fp8.is_none() && self.shared_down_fp8.is_none() - && self.bf16_shared_expert.is_none() + && self.separate_shared_expert.is_none() && !self.m5_projections.shared.enabled() && !self.gelu_activation && self.lora.is_none() diff --git a/crates/spark-model/src/layers/moe/forward_rows.rs b/crates/spark-model/src/layers/moe/forward_rows.rs index 7d2d3ff82..2d772233c 100644 --- a/crates/spark-model/src/layers/moe/forward_rows.rs +++ b/crates/spark-model/src/layers/moe/forward_rows.rs @@ -74,7 +74,6 @@ impl MoeLayer { && self.lora.is_none() && self.pre_expert_norm.is_none() && self.weights.router_pre_norm.is_none() - && !self.has_mixed_bf16_shared_expert() && self.tid2eid_dev.is_none() && self.correction_bias_dev.is_none() && self.router_logits_n as usize == ctx.config.num_experts @@ -239,6 +238,9 @@ impl MoeLayer { let (h32, inter32, top_k32, rows32) = (h as u32, inter as u32, top_k as u32, rows as u32); let table = |t: &ExpertPtrTable| (t.packed_ptrs, t.scale_ptrs, t.scale2_vals); let shared = &self.weights.shared_expert; + // `shared` is NULL under a checkpoint-native separate shared expert — + // the fused kernels skip it and the separate chain fills `shared_out` + // for the same rows below. if k.units_ready() { // A launch takes up to QWEN4EXP_MOE_UNITS_MAX_ROWS rows; each row's // bytes are the same in any chunk (and under ATLAS_QWEN4EXP_MOE_TC @@ -285,6 +287,20 @@ impl MoeLayer { )?; } + if self.has_separate_shared_expert() { + self.run_separate_shared_expert( + input, + rows32, + h32, + ctx.config.shared_expert_intermediate_size as u32, + ctx.buffers.logits(), + ctx.buffers.ssm_qkvz(), + shared_out, + ctx, + stream, + )?; + } + // Each row's blend, as `forward_row_local` runs it: under EP against a // zeroed shared row (the EP reduce adds the shared expert once). let is_ep = ctx.comm.is_some() && ctx.config.ep_world_size > 1; diff --git a/crates/spark-model/src/layers/moe/forward_token_major.rs b/crates/spark-model/src/layers/moe/forward_token_major.rs index ddb20df28..d0ad4aaf4 100644 --- a/crates/spark-model/src/layers/moe/forward_token_major.rs +++ b/crates/spark-model/src/layers/moe/forward_token_major.rs @@ -49,7 +49,7 @@ impl MoeLayer { && self.weights.shared_expert.up_proj.weight.0 != 0 && self.weights.shared_expert.down_proj.weight.0 != 0; let nvfp4_supported = self.bf16_gate_weight_ptrs.is_none() - && !self.has_mixed_bf16_shared_expert() + && !self.has_separate_shared_expert() && self.fp8_gate_weight_ptrs.is_none() && !self.use_t_layout_for_decode() && has_shared; diff --git a/crates/spark-model/src/layers/moe/gate_up_btile_preflight.rs b/crates/spark-model/src/layers/moe/gate_up_btile_preflight.rs index cb7d6bdec..cf5f64ffc 100644 --- a/crates/spark-model/src/layers/moe/gate_up_btile_preflight.rs +++ b/crates/spark-model/src/layers/moe/gate_up_btile_preflight.rs @@ -30,7 +30,7 @@ pub(super) fn layer(layer: &MoeLayer, config: &ModelConfig) -> Result<()> { && layer.shared_down_t.is_none() && layer.shared_gate_up_receipt.is_none() && layer.cutlass_grouped_host.is_none() - && layer.bf16_shared_expert.is_none() + && layer.separate_shared_expert.is_none() && layer.fp8_shared_expert.is_none() && layer.bf16_gate_weight_ptrs.is_none() && layer.fp8_gate_weight_ptrs.is_none() diff --git a/crates/spark-model/src/layers/moe/gate_up_btile_resident_arena.rs b/crates/spark-model/src/layers/moe/gate_up_btile_resident_arena.rs index 971d083e3..abf0c48ee 100644 --- a/crates/spark-model/src/layers/moe/gate_up_btile_resident_arena.rs +++ b/crates/spark-model/src/layers/moe/gate_up_btile_resident_arena.rs @@ -139,7 +139,7 @@ impl MoeLayer { anyhow::ensure!( self.bf16_gate_weight_ptrs.is_none() && self.fp8_gate_weight_ptrs.is_none() - && self.bf16_shared_expert.is_none() + && self.separate_shared_expert.is_none() && self.fp8_shared_expert.is_none() && self.gate_fp8.is_none() && self.shared_gate_fp8.is_none() diff --git a/crates/spark-model/src/layers/moe/helpers_c.rs b/crates/spark-model/src/layers/moe/helpers_c.rs index 7fab906c6..9a44a55e6 100644 --- a/crates/spark-model/src/layers/moe/helpers_c.rs +++ b/crates/spark-model/src/layers/moe/helpers_c.rs @@ -35,7 +35,7 @@ impl MoeLayer { // A checkpoint-native BF16 shared expert is the authoritative copy. // Do not manufacture an FP8 prefill variant with different numerics. - if self.bf16_shared_expert.is_none() + if self.separate_shared_expert.is_none() && !self.weights.shared_expert.gate_proj.is_null() && shared_inter > 0 { @@ -109,7 +109,7 @@ impl MoeLayer { self.bf16_up_weight_ptrs = Some(build_bf16_ptr_table(up_experts, gpu)?); self.bf16_down_weight_ptrs = Some(build_bf16_ptr_table(down_experts, gpu)?); if shared_gate.is_null() && shared_up.is_null() && shared_down.is_null() { - self.bf16_shared_expert = None; + self.separate_shared_expert = None; } else { self.set_bf16_shared_expert( DenseWeight { @@ -124,129 +124,6 @@ impl MoeLayer { Ok(()) } - /// Install checkpoint-native BF16 shared-expert weights independently of - /// routed-expert precision. - pub fn set_bf16_shared_expert( - &mut self, - gate_proj: DenseWeight, - up_proj: DenseWeight, - down_proj: DenseWeight, - ) -> Result<()> { - self.btile_storage.require_legacy()?; - self.bf16_shared_expert = Some(Bf16SharedExpert::new(gate_proj, up_proj, down_proj)?); - Ok(()) - } - - /// Whether a BF16 shared expert must overwrite the contribution produced - /// by a quantized fused routed-expert kernel. - pub(super) fn has_mixed_bf16_shared_expert(&self) -> bool { - self.bf16_shared_expert.is_some() && self.bf16_gate_weight_ptrs.is_none() - } - - /// Evaluate a checkpoint-native BF16 shared expert into `down_out`. - /// - /// Callers supply scratch buffers because the safe aliases differ between - /// decode and prefill. Returns `true` when BF16 weights were installed. - #[allow(clippy::too_many_arguments)] - pub(super) fn run_bf16_shared_expert( - &self, - input: DevicePtr, - num_tokens: u32, - hidden_size: u32, - shared_intermediate: u32, - gate_out: DevicePtr, - up_out: DevicePtr, - down_out: DevicePtr, - ctx: &ForwardContext, - stream: u64, - ) -> Result { - let Some(shared) = self.bf16_shared_expert else { - return Ok(false); - }; - anyhow::ensure!( - num_tokens > 0 && shared_intermediate > 0, - "BF16 shared expert requires non-zero token and intermediate dimensions" - ); - - let project = |activation: DevicePtr, - weight: &DenseWeight, - output: DevicePtr, - n: u32, - k: u32| - -> Result<()> { - if num_tokens == 1 { - ops::dense_gemv( - ctx.gpu, - self.dense_gemv, - activation, - weight, - output, - n, - k, - stream, - ) - } else if ctx.dispatch.cublas_gemm { - // Multi-token shared expert: cuBLASLt BF16 beats the hand-written - // mma.sync GEMM. The single-token arm above stays on the GEMV — - // decode-sized shapes do not repay cuBLAS heuristic overhead. - ops::cublas_bf16_proj_dense( - activation, - weight.weight, - output, - num_tokens, - n, - k, - stream, - ) - } else { - ops::dense_gemm_prefill( - ctx.gpu, - self.dense_gemm, - self.dense_gemm_pipelined, - activation, - weight, - output, - num_tokens, - n, - k, - stream, - ) - } - }; - - project( - input, - &shared.gate_proj, - gate_out, - shared_intermediate, - hidden_size, - )?; - project( - input, - &shared.up_proj, - up_out, - shared_intermediate, - hidden_size, - )?; - ops::silu_mul( - ctx.gpu, - self.moe_act_mul, - gate_out, - up_out, - gate_out, - num_tokens * shared_intermediate, - stream, - )?; - project( - gate_out, - &shared.down_proj, - down_out, - hidden_size, - shared_intermediate, - )?; - Ok(true) - } - /// Router gate GEMM for a BF16 (dense) gate weight at prefill: /// `gate_logits[n, num_experts] = router_in @ gate^T`. /// diff --git a/crates/spark-model/src/layers/moe/helpers_shared.rs b/crates/spark-model/src/layers/moe/helpers_shared.rs new file mode 100644 index 000000000..6b01e2a2c --- /dev/null +++ b/crates/spark-model/src/layers/moe/helpers_shared.rs @@ -0,0 +1,228 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Checkpoint-native separate shared expert (BF16 or MXFP8) installed beside +//! the fused routed-expert kernels. + +use super::*; + +impl MoeLayer { + /// Install checkpoint-native BF16 shared-expert weights independently of + /// routed-expert precision. + pub fn set_bf16_shared_expert( + &mut self, + gate_proj: DenseWeight, + up_proj: DenseWeight, + down_proj: DenseWeight, + ) -> Result<()> { + self.btile_storage.require_legacy()?; + self.separate_shared_expert = Some(SeparateSharedExpert::Bf16(Bf16SharedExpert::new( + gate_proj, up_proj, down_proj, + )?)); + Ok(()) + } + + /// Install FP4-Local's checkpoint-native MXFP8 shared expert: the fused + /// routed kernels get a NULL shared expert and this chain projects the + /// E4M3 bytes through `mxfp8_proj.cu` — no BF16 expansion, no 4-bit + /// requant of an 8-bit weight. + pub(crate) fn set_mxfp8_shared_expert( + &mut self, + gate_proj: Mxfp8Proj, + up_proj: Mxfp8Proj, + down_proj: Mxfp8Proj, + gate_up: Option, + kernels: ops::Mxfp8ProjKernels, + scratch: ops::Mxfp8Scratch, + ) -> Result<()> { + self.btile_storage.require_legacy()?; + anyhow::ensure!( + !gate_proj.w.data.is_null() && !up_proj.w.data.is_null() && !down_proj.w.data.is_null(), + "MXFP8 shared expert requires non-null gate/up/down weights" + ); + self.separate_shared_expert = Some(SeparateSharedExpert::Mxfp8 { + gate_proj, + up_proj, + down_proj, + gate_up, + kernels, + scratch, + }); + Ok(()) + } + + /// Whether a checkpoint-native shared expert must overwrite the + /// contribution produced by a quantized fused routed-expert kernel. + pub(super) fn has_separate_shared_expert(&self) -> bool { + self.separate_shared_expert.is_some() && self.bf16_gate_weight_ptrs.is_none() + } + + /// Evaluate a checkpoint-native shared expert into `down_out`. + /// + /// Callers supply scratch buffers because the safe aliases differ between + /// decode and prefill. Returns `true` when a separate expert was installed. + #[allow(clippy::too_many_arguments)] + pub(super) fn run_separate_shared_expert( + &self, + input: DevicePtr, + num_tokens: u32, + hidden_size: u32, + shared_intermediate: u32, + gate_out: DevicePtr, + up_out: DevicePtr, + down_out: DevicePtr, + ctx: &ForwardContext, + stream: u64, + ) -> Result { + let Some(shared) = self.separate_shared_expert else { + return Ok(false); + }; + anyhow::ensure!( + num_tokens > 0 && shared_intermediate > 0, + "separate shared expert requires non-zero token and intermediate dimensions" + ); + + if let SeparateSharedExpert::Mxfp8 { + gate_proj, + up_proj, + down_proj, + gate_up, + kernels, + scratch, + } = shared + { + let m = num_tokens as usize; + let inter = shared_intermediate as usize; + if let Some(gu) = gate_up + && m <= ops::MXFP8_PROJ_MAX_CHUNKED_ROWS + && gu.silu_mul.0 != 0 + { + // One projection yields `[m, 2*inter]` gate|up; fused_silu_mul + // halves it in place-compatible scratch before the down pass. + ops::mxfp8_proj( + ctx.gpu, + &kernels, + &gu.proj, + input, + gu.gate_up_out, + m, + scratch, + stream, + )?; + ops::fused_silu_mul( + ctx.gpu, + gu.silu_mul, + gu.gate_up_out, + up_out, + num_tokens, + shared_intermediate, + stream, + )?; + ops::mxfp8_proj( + ctx.gpu, &kernels, &down_proj, up_out, down_out, m, scratch, stream, + )?; + return Ok(true); + } + ops::mxfp8_proj( + ctx.gpu, &kernels, &gate_proj, input, gate_out, m, scratch, stream, + )?; + ops::mxfp8_proj( + ctx.gpu, &kernels, &up_proj, input, up_out, m, scratch, stream, + )?; + ops::silu_mul( + ctx.gpu, + self.moe_act_mul, + gate_out, + up_out, + gate_out, + num_tokens * shared_intermediate, + stream, + )?; + let _ = inter; + ops::mxfp8_proj( + ctx.gpu, &kernels, &down_proj, gate_out, down_out, m, scratch, stream, + )?; + return Ok(true); + } + let SeparateSharedExpert::Bf16(shared) = shared else { + unreachable!() + }; + + let project = |activation: DevicePtr, + weight: &DenseWeight, + output: DevicePtr, + n: u32, + k: u32| + -> Result<()> { + if num_tokens == 1 { + ops::dense_gemv( + ctx.gpu, + self.dense_gemv, + activation, + weight, + output, + n, + k, + stream, + ) + } else if ctx.dispatch.cublas_gemm { + // Multi-token shared expert: cuBLASLt BF16 beats the hand-written + // mma.sync GEMM. The single-token arm above stays on the GEMV — + // decode-sized shapes do not repay cuBLAS heuristic overhead. + ops::cublas_bf16_proj_dense( + activation, + weight.weight, + output, + num_tokens, + n, + k, + stream, + ) + } else { + ops::dense_gemm_prefill( + ctx.gpu, + self.dense_gemm, + self.dense_gemm_pipelined, + activation, + weight, + output, + num_tokens, + n, + k, + stream, + ) + } + }; + + project( + input, + &shared.gate_proj, + gate_out, + shared_intermediate, + hidden_size, + )?; + project( + input, + &shared.up_proj, + up_out, + shared_intermediate, + hidden_size, + )?; + ops::silu_mul( + ctx.gpu, + self.moe_act_mul, + gate_out, + up_out, + gate_out, + num_tokens * shared_intermediate, + stream, + )?; + project( + gate_out, + &shared.down_proj, + down_out, + hidden_size, + shared_intermediate, + )?; + Ok(true) + } +} diff --git a/crates/spark-model/src/layers/moe/init.rs b/crates/spark-model/src/layers/moe/init.rs index 0664db78b..3894c38ca 100644 --- a/crates/spark-model/src/layers/moe/init.rs +++ b/crates/spark-model/src/layers/moe/init.rs @@ -445,7 +445,7 @@ impl MoeLayer { bf16_gate_weight_ptrs: None, bf16_up_weight_ptrs: None, bf16_down_weight_ptrs: None, - bf16_shared_expert: None, + separate_shared_expert: None, fp8_shared_expert: None, moe_down_t_k64_fp4: super::super::try_kernel_gated( super::init_optional::gateup_fp4_series(&config.model_type), diff --git a/crates/spark-model/src/layers/moe/mod.rs b/crates/spark-model/src/layers/moe/mod.rs index e89b8dcab..dbd86ea5d 100644 --- a/crates/spark-model/src/layers/moe/mod.rs +++ b/crates/spark-model/src/layers/moe/mod.rs @@ -12,7 +12,7 @@ use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle}; use crate::layer::ForwardContext; use crate::layers::ops; use crate::layers::w4a16_gemv_tiers::W4a16BatchmTiers; -use crate::weight_map::{DenseWeight, Fp8ExpertWeight, MoeWeights, QuantizedWeight}; +use crate::weight_map::{DenseWeight, Fp8ExpertWeight, MoeWeights, Mxfp8Proj, QuantizedWeight}; mod btile_model_owners; pub(crate) use btile_model_owners::{ bind as bind_resident_btile_arenas, invalidate as invalidate_resident_btile_readers, @@ -82,6 +82,38 @@ impl Bf16SharedExpert { } } +/// A shared expert kept outside the fused routed-expert kernels and run as +/// its own gate/up → silu_mul → down chain over the same scratch contract. +#[derive(Debug, Clone, Copy)] +pub(crate) enum SeparateSharedExpert { + /// Checkpoint-native BF16 (Laguna). + Bf16(Bf16SharedExpert), + /// FP4-Local's checkpoint-native MXFP8: E4M3 weights + E8M0 scales, + /// projected by `mxfp8_proj.cu` — no BF16 expansion, no NVFP4 requant. + Mxfp8 { + gate_proj: Mxfp8Proj, + up_proj: Mxfp8Proj, + down_proj: Mxfp8Proj, + /// Gate‖up row-concatenated `[2*inter, h]` twin plus its + /// `[32, 2*inter]` BF16 landing buffer: one `mxfp8_proj` launch then + /// `fused_silu_mul` per row-band ≤ 32 (the batchm arm's span). + gate_up: Option, + kernels: ops::Mxfp8ProjKernels, + scratch: ops::Mxfp8Scratch, + }, +} + +/// The fused small-row arm of a [`SeparateSharedExpert::Mxfp8`]. +#[derive(Debug, Clone, Copy)] +pub(crate) struct Mxfp8GateUp { + /// `[2*inter, h]` — gate rows then up rows. + pub proj: Mxfp8Proj, + /// `[32, 2*inter]` BF16 — one batchm-chunk's gate|up activations. + pub gate_up_out: DevicePtr, + /// `fused_silu_mul` (`dense_gemm_bf16.cu`); `0` resolves to the unfused path. + pub silu_mul: KernelHandle, +} + /// Unified expert pointer table for any quantization format. /// /// Replaces the separate `ExpertPtrTable` (NVFP4) and `Fp8ExpertPtrTable` (FP8) @@ -186,6 +218,7 @@ mod gate_up_repack_test_gpu; mod helpers_a; mod helpers_b; mod helpers_c; +mod helpers_shared; mod route_dump; mod shared_fp8_cache; mod shared_fp8_cache_load; diff --git a/crates/spark-model/src/layers/moe/mod_tests.rs b/crates/spark-model/src/layers/moe/mod_tests.rs index 76fe61136..e3e8e5269 100644 --- a/crates/spark-model/src/layers/moe/mod_tests.rs +++ b/crates/spark-model/src/layers/moe/mod_tests.rs @@ -44,7 +44,7 @@ fn test_moe_kernel_loading() { } #[test] -fn bf16_shared_expert_requires_three_non_null_weights() { +fn separate_shared_expert_requires_three_non_null_weights() { let gate = DenseWeight { weight: DevicePtr(11), }; @@ -67,3 +67,123 @@ fn bf16_shared_expert_requires_three_non_null_weights() { assert!(Bf16SharedExpert::new(gate, null, down).is_err()); assert!(Bf16SharedExpert::new(gate, up, null).is_err()); } + +/// The MXFP8 shared expert dispatches onto `mxfp8_proj.cu` at every row +/// band: the 1-row decode GEMV, the 2-row verify batch, and the >32-row +/// dequant+GEMM prefill arm — with `silu_mul` between up and down. +#[test] +fn mxfp8_shared_expert_runs_at_decode_verify_and_prefill_rows() { + use crate::layer::ForwardContext; + use crate::layers::moe::shared_fp8_cache_test_gpu::RecordingGpu; + use atlas_core::config::ModelConfig; + use spark_runtime::buffers::BufferArena; + use spark_runtime::gpu::KernelHandle; + + let gpu = RecordingGpu::new(); + let mut config = ModelConfig::qwen3_next_80b_nvfp4(); + config.hidden_size = 2560; + config.num_experts = 512; + config.num_experts_per_tok = 10; + config.moe_intermediate_size = 640; + config.shared_expert_intermediate_size = 640; + config.vocab_size = 8; + let mut layer = MoeLayer::new(MoeWeights::empty(512), 512, None, &gpu, &config).unwrap(); + let (h, inter) = (config.hidden_size, config.shared_expert_intermediate_size); + let proj = |n: usize, k: usize, ptr: u64| Mxfp8Proj { + w: crate::weight_map::Mxfp8Weight { + data: DevicePtr(ptr), + scales: DevicePtr(ptr + 0x1000), + }, + n: n as u32, + k: k as u32, + }; + let kernels = ops::Mxfp8ProjKernels { + gemv: KernelHandle(111), + batchm: KernelHandle(222), + dequant: KernelHandle(333), + }; + let scratch = ops::Mxfp8Scratch { + ptr: DevicePtr(0x9000_0000), + bytes: 1 << 24, + }; + layer + .set_mxfp8_shared_expert( + proj(inter, h, 0x10000), + proj(inter, h, 0x20000), + proj(h, inter, 0x30000), + None, + kernels, + scratch, + ) + .unwrap(); + assert!(layer.has_separate_shared_expert()); + + let buffers = BufferArena::new(&config, 64, 4096, 16, 1, &gpu).unwrap(); + let dispatch = ops::GemmDispatch::defaults(); + let derived = ops::DerivedWeights::new(); + let levers = ops::ModelLevers::defaults(); + let stats = ops::ModelStats::new(); + let ctx = ForwardContext { + buffers: &buffers, + gpu: &gpu, + config: &config, + dispatch: &dispatch, + derived: &derived, + levers: &levers, + stats: &stats, + ssm_batch: None, + attn_metadata: None, + profile: false, + comm: None, + graph_capture: false, + gdn_exact_replay: false, + token_ids: None, + host_token_ids: None, + routed_lora_layers: None, + midchunk_capture: None, + moe_lora_route: crate::layer::MoeLoraRoute::Skip, + }; + let (input, gate_out, up_out, down_out) = ( + DevicePtr(0x40000), + DevicePtr(0x50000), + DevicePtr(0x60000), + DevicePtr(0x70000), + ); + // gate GEMV, up GEMV, silu_mul, down GEMV on the band's row kernel. The + // >32-row arm (dequant + CUTLASS GEMM) is exercised by the mx8 GPU tests — + // its FFI cannot run on this mock. + for (m, want) in [(1u32, 111u64), (2, 222), (32, 222)] { + gpu.launches.lock().unwrap().clear(); + let ran = layer + .run_separate_shared_expert( + input, + m, + h as u32, + inter as u32, + gate_out, + up_out, + down_out, + &ctx, + 19, + ) + .unwrap(); + assert!(ran, "separate shared expert must run at m={m}"); + let kernels_run: Vec = gpu + .launches + .lock() + .unwrap() + .iter() + .map(|l| l.kernel) + .collect(); + let want_n = if m == 1 { + 3 + } else { + 3 * (m as usize).div_ceil(8) + }; + assert_eq!( + kernels_run.iter().filter(|&&k| k == want).count(), + want_n, + "m={m}: expected {want_n} launches of {want}, got {kernels_run:?}" + ); + } +} diff --git a/crates/spark-model/src/layers/moe/prequant_fp4.rs b/crates/spark-model/src/layers/moe/prequant_fp4.rs index 27110feb4..e7362ad09 100644 --- a/crates/spark-model/src/layers/moe/prequant_fp4.rs +++ b/crates/spark-model/src/layers/moe/prequant_fp4.rs @@ -138,7 +138,7 @@ impl MoeLayer { && self.lora.is_none() && self.bf16_gate_weight_ptrs.is_none() && self.fp8_gate_weight_ptrs.is_none() - && !self.has_mixed_bf16_shared_expert() + && !self.has_separate_shared_expert() && self.experts_scale_kind == crate::weight_map::WeightQuantFormat::Nvfp4 && !self.nvfp4_mmq_layout && self.use_btile_or_t_prefill() diff --git a/crates/spark-model/src/layers/moe/qwen4exp_fast.rs b/crates/spark-model/src/layers/moe/qwen4exp_fast.rs index fcbda027d..43c34b166 100644 --- a/crates/spark-model/src/layers/moe/qwen4exp_fast.rs +++ b/crates/spark-model/src/layers/moe/qwen4exp_fast.rs @@ -141,7 +141,6 @@ impl MoeLayer { && self.correction_bias_dev.is_none() && self.bf16_gate_weight_ptrs.is_none() && self.fp8_gate_weight_ptrs.is_none() - && !self.has_mixed_bf16_shared_expert() && self.router_logits_n as usize == ctx.config.num_experts) { return Ok(false); @@ -191,15 +190,28 @@ impl MoeLayer { for first in (0..num_tokens).step_by(ops::QWEN4EXP_MOE_MAX_ROWS) { let rows = (num_tokens - first).min(ops::QWEN4EXP_MOE_MAX_ROWS); let chunk_indices = indices.offset(first * top_k * 4); + // A checkpoint-native separate shared expert never reaches the + // fused kernels: NULL `_t` weights make them skip it, and the + // separate chain fills `shared_out` for the same rows. + let separate = self.has_separate_shared_expert(); + let (sh_gate_t, sh_up_t, sh_down_t) = if separate { + (&null_qw, &null_qw, &null_qw) + } else { + ( + self.shared_gate_t.as_ref().unwrap_or(&null_qw), + self.shared_up_t.as_ref().unwrap_or(&null_qw), + self.shared_down_t.as_ref().unwrap_or(&null_qw), + ) + }; self.qwen4exp_fast_gate_up( ctx, input.offset(first * h * bf16), expert_gate_out, expert_up_out, chunk_indices, - self.shared_gate_t.as_ref().unwrap_or(&null_qw), + sh_gate_t, shared_gate_scratch, - self.shared_up_t.as_ref().unwrap_or(&null_qw), + sh_up_t, shared_up_scratch, rows, stream, @@ -212,11 +224,24 @@ impl MoeLayer { chunk_indices, shared_gate_scratch, shared_up_scratch, - self.shared_down_t.as_ref().unwrap_or(&null_qw), + sh_down_t, shared_out, rows, stream, )?; + if separate { + self.run_separate_shared_expert( + input.offset(first * h * bf16), + rows as u32, + h as u32, + ctx.config.shared_expert_intermediate_size as u32, + shared_gate_scratch, + shared_up_scratch, + shared_out, + ctx, + stream, + )?; + } for t in 0..rows { let row = first + t; let output_t = ctx.buffers.moe_output().offset(row * h * bf16); diff --git a/crates/spark-model/src/layers/moe/shared_fp8_cache_load.rs b/crates/spark-model/src/layers/moe/shared_fp8_cache_load.rs index 1e3e0c182..0d3d72e0c 100644 --- a/crates/spark-model/src/layers/moe/shared_fp8_cache_load.rs +++ b/crates/spark-model/src/layers/moe/shared_fp8_cache_load.rs @@ -101,7 +101,7 @@ impl MoeLayer { anyhow::ensure!( self.shared_experts_scale_kind == WeightQuantFormat::Nvfp4 && self.gate_nvfp4.is_none() - && self.bf16_shared_expert.is_none() + && self.separate_shared_expert.is_none() && self.fp8_shared_expert.is_none() && self.lora.is_none(), "shared FP8 cache requires native GLM shared weights and BF16 router" diff --git a/crates/spark-model/src/layers/moe/shared_m16.rs b/crates/spark-model/src/layers/moe/shared_m16.rs index 7ff56571e..1b7c39498 100644 --- a/crates/spark-model/src/layers/moe/shared_m16.rs +++ b/crates/spark-model/src/layers/moe/shared_m16.rs @@ -87,7 +87,7 @@ impl MoeLayer { fp8_cache: self.shared_gate_fp8.is_some() || self.shared_up_fp8.is_some() || self.shared_down_fp8.is_some(), - bf16_override: self.bf16_shared_expert.is_some(), + bf16_override: self.separate_shared_expert.is_some(), // Hook is after the earlier exact-M5 GEMV arm returned. exact_gemv: false, }; diff --git a/crates/spark-model/src/layers/moe/types.rs b/crates/spark-model/src/layers/moe/types.rs index 2d300f549..fd17eafbf 100644 --- a/crates/spark-model/src/layers/moe/types.rs +++ b/crates/spark-model/src/layers/moe/types.rs @@ -409,9 +409,10 @@ pub struct MoeLayer { pub(super) bf16_gate_weight_ptrs: Option, pub(super) bf16_up_weight_ptrs: Option, pub(super) bf16_down_weight_ptrs: Option, - // Checkpoint-native BF16 shared expert. Independent of routed-expert - // precision so mixed NVFP4-routed/BF16-shared checkpoints stay faithful. - pub(super) bf16_shared_expert: Option, + // Checkpoint-native shared expert run separately from the fused routed + // kernels (BF16 for Laguna, MXFP8 for FP4-Local). Independent of + // routed-expert precision so mixed-precision checkpoints stay faithful. + pub(super) separate_shared_expert: Option, // FP8 shared expert weights (None when shared expert is NVFP4) pub(super) fp8_shared_expert: Option, /// FP4 down kernel handle (`moe_w4a16_down_t_k64_fp4`). `try_kernel` => diff --git a/crates/spark-model/src/layers/ops/activations.rs b/crates/spark-model/src/layers/ops/activations.rs index 68da29628..cfc46a3a5 100644 --- a/crates/spark-model/src/layers/ops/activations.rs +++ b/crates/spark-model/src/layers/ops/activations.rs @@ -36,6 +36,29 @@ pub fn silu_mul( .launch(stream) } +/// SiLU·mul over a `[num_tokens, inter*2]` gate‖up layout — the first `inter` +/// columns of each row are the gate, the next `inter` the up. +/// +/// Kernel: `fused_silu_mul` (`dense_gemm_bf16.cu`), module `gemm`. +pub fn fused_silu_mul( + gpu: &dyn GpuBackend, + kernel: KernelHandle, + gate_up: DevicePtr, + output: DevicePtr, + num_tokens: u32, + inter_size: u32, + stream: u64, +) -> Result<()> { + KernelLaunch::new(gpu, kernel) + .grid([div_ceil(num_tokens * inter_size / 2, 256), 1, 1]) + .block([256, 1, 1]) + .arg_ptr(gate_up) + .arg_ptr(output) + .arg_u32(num_tokens) + .arg_u32(inter_size) + .launch(stream) +} + /// Fused SiLU·mul + per-token-group(128) FP8-E4M3 quantization — replaces the /// `silu_mul` → `per_token_group_quant_fp8` pair on the W8A8 prefill down-path /// without materializing the BF16 intermediate. Bit-identical to the pair diff --git a/crates/spark-model/src/weight_loader/qwen35/load_layers.rs b/crates/spark-model/src/weight_loader/qwen35/load_layers.rs index 52a5d4123..c30bdfd7d 100644 --- a/crates/spark-model/src/weight_loader/qwen35/load_layers.rs +++ b/crates/spark-model/src/weight_loader/qwen35/load_layers.rs @@ -277,6 +277,7 @@ pub(super) fn load_layers( quantize_k, stream, skip_nvfp4_experts, + false, )?; // 2026-05-25 (final): gate stays in BF16 for `native_fp8` — // routes through `dense_gemm` BF16 fallback path. diff --git a/crates/spark-model/src/weight_loader/qwen4_exp.rs b/crates/spark-model/src/weight_loader/qwen4_exp.rs index 727ebe410..a543ac50c 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp.rs @@ -261,7 +261,7 @@ impl ModelWeightLoader for Qwen4ExpWeightLoader { for i in 0..config.num_hidden_layers { let lp = config.layer_prefix(i); let f0 = free_now(gpu); - let ffn = ffn::build_moe(store, &lp, config, gpu, variant)?; + let ffn = ffn::build_moe(store, &lp, config, gpu, variant, &mut mx8_scratch)?; let f1 = free_now(gpu); moe_bytes += f0.saturating_sub(f1); diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/ffn.rs b/crates/spark-model/src/weight_loader/qwen4_exp/ffn.rs index 2e904c3d2..4b815690f 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/ffn.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/ffn.rs @@ -16,14 +16,42 @@ //! for the non-native-FP8 path, matching qwen35. Its precision matters more //! than its size: at 512 experts the top-10 weights cluster tightly, and a //! 4-bit ULP wider than that spread cannot tell them apart. +//! +//! FP4-Local keeps the shared expert MXFP8: it is the one expert every token +//! goes through, so the BF16-expand + NVFP4 requant the fused path would +//! impose is not paid here. The routed experts stay on `quantized_any` +//! exactly as every other pack. use anyhow::{Context, Result}; use atlas_core::config::ModelConfig; use spark_runtime::gpu::GpuBackend; -use spark_runtime::weights::WeightStore; +use spark_runtime::weights::{WeightDtype, WeightStore}; +use crate::layers::ops::{Mxfp8Proj, Mxfp8ProjKernels, Mxfp8Scratch}; use crate::layers::{FfnComponent, MoeLayer}; -use crate::weight_map::{Nvfp4Variant, load_moe_qwen35, quantize_to_nvfp4}; +use crate::weight_map::{Mxfp8Weight, Nvfp4Variant, load_moe_qwen35, quantize_to_nvfp4}; + +use super::mx8_common; + +/// `Some` when the shared expert sits in the store as checkpoint-native +/// MXFP8 (FP4-Local): E4M3 `weight` + U8 E8M0 `weight_scale` per 32 inputs. +fn mxfp8_shared(store: &WeightStore, lp: &str, inter: usize, h: usize) -> Result { + let key = format!("{lp}.mlp.shared_expert.gate_proj.weight"); + let Ok(w) = store.get(&key) else { + return Ok(false); + }; + if w.dtype != WeightDtype::FP8E4M3 { + return Ok(false); + } + for (proj, n, k) in [ + ("gate_proj", inter, h), + ("up_proj", inter, h), + ("down_proj", h, inter), + ] { + let [_, _] = mx8_common::module(store, &format!("{lp}.mlp.shared_expert.{proj}"), n, k)?; + } + Ok(true) +} pub(super) fn build_moe( store: &WeightStore, @@ -31,11 +59,14 @@ pub(super) fn build_moe( config: &ModelConfig, gpu: &dyn GpuBackend, variant: Nvfp4Variant, + mx8_scratch: &mut Option, ) -> Result { let h = config.hidden_size; let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?; let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?; let stream = gpu.default_stream(); + let inter = config.shared_expert_intermediate_size; + let native_shared = mxfp8_shared(store, lp, inter, h)?; let weights = load_moe_qwen35( store, @@ -48,6 +79,7 @@ pub(super) fn build_moe( quantize_k, stream, false, + native_shared, ) .with_context(|| format!("qwen4_exp: MoE block at {lp}"))?; @@ -61,6 +93,58 @@ pub(super) fn build_moe( stream, )?); - let moe = MoeLayer::new(weights, config.num_experts, gate_nvfp4, gpu, config)?; + let mut moe = MoeLayer::new(weights, config.num_experts, gate_nvfp4, gpu, config)?; + if native_shared { + let m = |proj: &str, n: usize, k: usize| -> Result { + let [w, s] = + mx8_common::module(store, &format!("{lp}.mlp.shared_expert.{proj}"), n, k)?; + Ok(Mxfp8Proj { + w: Mxfp8Weight { data: w, scales: s }, + n: n as u32, + k: k as u32, + }) + }; + // Gate‖up row-concatenated `[2*inter, h]` twin: one `mxfp8_proj` pass + // then `fused_silu_mul` on the `[m, 2*inter]` layout for the ≤32-row + // bands; the unfused trio stays for the prefill arm. + let [gw, gs] = mx8_common::module( + store, + &format!("{lp}.mlp.shared_expert.gate_proj"), + inter, + h, + )?; + let [uw, us] = + mx8_common::module(store, &format!("{lp}.mlp.shared_expert.up_proj"), inter, h)?; + let gate_up = { + let w = gpu.alloc(2 * inter * h)?; + gpu.copy_d2d(gw, w, inter * h)?; + gpu.copy_d2d(uw, w.offset(inter * h), inter * h)?; + let s = gpu.alloc(2 * inter * (h / 32))?; + gpu.copy_d2d(gs, s, inter * (h / 32))?; + gpu.copy_d2d(us, s.offset(inter * (h / 32)), inter * (h / 32))?; + let out = gpu.alloc(32 * 2 * inter * 2)?; + Some(crate::layers::moe::Mxfp8GateUp { + proj: Mxfp8Proj { + w: Mxfp8Weight { data: w, scales: s }, + n: (2 * inter) as u32, + k: h as u32, + }, + gate_up_out: out, + silu_mul: crate::layers::try_kernel(gpu, "gemm", "fused_silu_mul"), + }) + }; + // The shared MXFP8 BF16 prefill scratch: the MoE runs after the + // layer's attention/GDN on the same stream, so the transient buffer + // is never live twice. + let scratch = mx8_common::scratch_for(mx8_scratch, config, gpu)?; + moe.set_mxfp8_shared_expert( + m("gate_proj", inter, h)?, + m("up_proj", inter, h)?, + m("down_proj", h, inter)?, + gate_up, + Mxfp8ProjKernels::resolve(gpu)?, + scratch, + )?; + } Ok(FfnComponent::Moe(moe)) } diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/mtp.rs b/crates/spark-model/src/weight_loader/qwen4_exp/mtp.rs index d3e4011c8..3fa9b4f74 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/mtp.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/mtp.rs @@ -243,25 +243,16 @@ fn build_mtp_moe( } // The shared expert ships per-tensor BF16 like a main layer's, so it takes - // the ordinary named path rather than the stacked slicer. + // the ordinary named path rather than the stacked slicer. It is also + // installed BF16-native on the layer below — the fused kernels then take a + // NULL shared expert and the BF16 chain applies it separately, sparing the + // one expert every drafted token sees the 4-bit requant. let se = format!("{mlp}.shared_expert"); - let shared_expert = ExpertWeight { - gate_proj: q( - &dense_auto(store, &format!("{se}.gate_proj.weight"), gpu)?, - inter, - h, - )?, - up_proj: q( - &dense_auto(store, &format!("{se}.up_proj.weight"), gpu)?, - inter, - h, - )?, - down_proj: q( - &dense_auto(store, &format!("{se}.down_proj.weight"), gpu)?, - h, - inter, - )?, - }; + let se_gate = dense_auto(store, &format!("{se}.gate_proj.weight"), gpu)?; + let se_up = dense_auto(store, &format!("{se}.up_proj.weight"), gpu)?; + let se_down = dense_auto(store, &format!("{se}.down_proj.weight"), gpu)?; + // NULL in the fused kernels: the BF16 chain installed below applies it. + let shared_expert = ExpertWeight::null(); let gate = dense_auto(store, &format!("{mlp}.gate.weight"), gpu)?; let weights = MoeWeights { @@ -276,8 +267,9 @@ fn build_mtp_moe( // Router precision matters more than router size here: at 512 experts the // top-10 weights cluster tightly. Same treatment as `ffn::build_moe`. let gate_nvfp4 = Some(q(&gate, n_experts, h)?); - let moe = MoeLayer::new(weights, n_experts, gate_nvfp4, gpu, config) + let mut moe = MoeLayer::new(weights, n_experts, gate_nvfp4, gpu, config) .context("qwen4_exp MTP: MoeLayer")?; + moe.set_bf16_shared_expert(se_gate, se_up, se_down)?; Ok(FfnComponent::Moe(moe)) } diff --git a/crates/spark-model/src/weight_map/modelopt_local.rs b/crates/spark-model/src/weight_map/modelopt_local.rs index ca2b2fe19..388877346 100644 --- a/crates/spark-model/src/weight_map/modelopt_local.rs +++ b/crates/spark-model/src/weight_map/modelopt_local.rs @@ -67,7 +67,7 @@ pub fn native_mxfp8(module: &str) -> bool { if std::env::var("ATLAS_FP4LOCAL_EXPAND_MXFP8").as_deref() == Ok("1") { return false; } - const NATIVE: [&str; 7] = [ + const NATIVE: [&str; 10] = [ ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", ".linear_attn.out_proj", @@ -75,6 +75,9 @@ pub fn native_mxfp8(module: &str) -> bool { ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.o_proj", + ".mlp.shared_expert.gate_proj", + ".mlp.shared_expert.up_proj", + ".mlp.shared_expert.down_proj", ]; module.starts_with("model.language_model.layers.") && NATIVE.iter().any(|suffix| module.ends_with(suffix)) diff --git a/crates/spark-model/src/weight_map/modelopt_local_tests.rs b/crates/spark-model/src/weight_map/modelopt_local_tests.rs index 1241cf748..28f9636eb 100644 --- a/crates/spark-model/src/weight_map/modelopt_local_tests.rs +++ b/crates/spark-model/src/weight_map/modelopt_local_tests.rs @@ -229,7 +229,19 @@ fn materialize_expands_every_fp4_local_format() { done.nvfp4_to_bf16, done.f32_norms_native ), - (1, 0, 1, 1, 2) + (1, 1, 0, 1, 2) + ); + // The shared-expert projection is a native consumer now: bytes as shipped. + let se = store.get(&format!("{SHARED}.weight")).unwrap(); + assert_eq!( + (se.dtype, se.shape.clone()), + (WeightDtype::FP8E4M3, vec![1, 32]) + ); + assert_eq!(gpu.read_alloc(se.ptr).unwrap(), [0x38u8; 32]); + let se_s = store.get(&format!("{SHARED}.weight_scale")).unwrap(); + assert_eq!( + (se_s.dtype, se_s.shape.clone()), + (WeightDtype::UInt8, vec![1, 1]) ); // The F32 norms are served as shipped: same dtype, byte-identical data. let norm = store.get(HC_NORM).unwrap(); @@ -268,15 +280,6 @@ fn materialize_expands_every_fp4_local_format() { want[5] = 0x7E; assert_eq!(gpu.read_alloc(scale.ptr).unwrap(), want); - let attn = store.get(&format!("{SHARED}.weight")).unwrap(); - assert_eq!( - (attn.dtype, attn.shape.clone()), - (WeightDtype::BF16, vec![1, 32]) - ); - let a = gpu.read_alloc(attn.ptr).unwrap(); - assert!((0..32).all(|i| bf16_at(&a, i) == 0.125)); - assert!(!store.contains(&format!("{SHARED}.weight_scale"))); - let vis = store.get(&format!("{VIS}.weight")).unwrap(); assert_eq!( (vis.dtype, vis.shape.clone()), diff --git a/crates/spark-model/src/weight_map/ssm_qwen35.rs b/crates/spark-model/src/weight_map/ssm_qwen35.rs index b054d166f..b27f87f06 100644 --- a/crates/spark-model/src/weight_map/ssm_qwen35.rs +++ b/crates/spark-model/src/weight_map/ssm_qwen35.rs @@ -97,6 +97,7 @@ pub(crate) fn load_moe_qwen35( quantize_k: spark_runtime::gpu::KernelHandle, stream: u64, skip_routed_experts: bool, + null_shared_expert: bool, ) -> Result { let p = format!("{layer_prefix}.mlp"); @@ -261,7 +262,14 @@ pub(crate) fn load_moe_qwen35( }) }; - let shared_expert = load_expert(&format!("{p}.shared_expert"))?; + // `null_shared_expert`: the caller serves the shared expert natively + // (FP4-Local's MXFP8) and installs it on the MoeLayer itself — the fused + // kernels see NULL shared weights and skip them. + let shared_expert = if null_shared_expert { + ExpertWeight::null() + } else { + load_expert(&format!("{p}.shared_expert"))? + }; let mut experts = Vec::with_capacity(num_experts); for e in 0..num_experts { diff --git a/crates/spark-model/src/weight_map/tests.rs b/crates/spark-model/src/weight_map/tests.rs index 080b8bbc5..782e59c0f 100644 --- a/crates/spark-model/src/weight_map/tests.rs +++ b/crates/spark-model/src/weight_map/tests.rs @@ -23,7 +23,13 @@ use super::*; fn dense_keep_f32_passes_f32_and_widens_bf16() { use spark_runtime::gpu::mock::MockGpuBackend; - fn put(store: &mut WeightStore, gpu: &MockGpuBackend, name: &str, bytes: &[u8], dtype: WeightDtype) { + fn put( + store: &mut WeightStore, + gpu: &MockGpuBackend, + name: &str, + bytes: &[u8], + dtype: WeightDtype, + ) { let ptr = gpu.alloc(bytes.len()).unwrap(); gpu.copy_h2d(bytes, ptr).unwrap(); store @@ -46,7 +52,13 @@ fn dense_keep_f32_passes_f32_and_widens_bf16() { .iter() .flat_map(|b| b.to_le_bytes()) .collect(); - put(&mut store, &gpu, "l.attn_hyper_connection.hc_norm.weight", &bf, WeightDtype::BF16); + put( + &mut store, + &gpu, + "l.attn_hyper_connection.hc_norm.weight", + &bf, + WeightDtype::BF16, + ); let w = dense_keep_f32(&store, "l.attn_hyper_connection.hc_norm.weight", &gpu).unwrap(); let got = gpu.read_alloc(w.weight).unwrap(); let want: Vec = [0x3F80u32 << 16, 0xC020 << 16, 0x3D00 << 16] @@ -60,7 +72,13 @@ fn dense_keep_f32_passes_f32_and_widens_bf16() { .iter() .flat_map(|v| v.to_le_bytes()) .collect(); - put(&mut store, &gpu, "l.ple.norm_key.weight", &f, WeightDtype::FP32); + put( + &mut store, + &gpu, + "l.ple.norm_key.weight", + &f, + WeightDtype::FP32, + ); let src = store.get("l.ple.norm_key.weight").unwrap().ptr; let w = dense_keep_f32(&store, "l.ple.norm_key.weight", &gpu).unwrap(); assert_eq!(w.weight, src); From 09075ad9847caf19c702a90a7da101f3b26e9a0c Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 04:23:21 -0400 Subject: [PATCH 09/26] qwen4_exp: serve shared-expert gate|up from one MXFP8 buffer gate_proj and up_proj are now views into the halves of the [2*inter, h] gate|up weight (and scale) buffer the fused small-row arm already needed; the store originals are reclaimed at load. Tracked weight memory drops ~0.16 GB. Also drop a dead local in helpers_shared. --- .../src/layers/moe/helpers_shared.rs | 2 - .../src/weight_loader/qwen4_exp/ffn.rs | 55 ++++++++++++------- 2 files changed, 36 insertions(+), 21 deletions(-) diff --git a/crates/spark-model/src/layers/moe/helpers_shared.rs b/crates/spark-model/src/layers/moe/helpers_shared.rs index 6b01e2a2c..ddaafa2e3 100644 --- a/crates/spark-model/src/layers/moe/helpers_shared.rs +++ b/crates/spark-model/src/layers/moe/helpers_shared.rs @@ -91,7 +91,6 @@ impl MoeLayer { } = shared { let m = num_tokens as usize; - let inter = shared_intermediate as usize; if let Some(gu) = gate_up && m <= ops::MXFP8_PROJ_MAX_CHUNKED_ROWS && gu.silu_mul.0 != 0 @@ -137,7 +136,6 @@ impl MoeLayer { num_tokens * shared_intermediate, stream, )?; - let _ = inter; ops::mxfp8_proj( ctx.gpu, &kernels, &down_proj, gate_out, down_out, m, scratch, stream, )?; diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/ffn.rs b/crates/spark-model/src/weight_loader/qwen4_exp/ffn.rs index 4b815690f..e5fe38767 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/ffn.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/ffn.rs @@ -115,31 +115,48 @@ pub(super) fn build_moe( )?; let [uw, us] = mx8_common::module(store, &format!("{lp}.mlp.shared_expert.up_proj"), inter, h)?; - let gate_up = { - let w = gpu.alloc(2 * inter * h)?; - gpu.copy_d2d(gw, w, inter * h)?; - gpu.copy_d2d(uw, w.offset(inter * h), inter * h)?; - let s = gpu.alloc(2 * inter * (h / 32))?; - gpu.copy_d2d(gs, s, inter * (h / 32))?; - gpu.copy_d2d(us, s.offset(inter * (h / 32)), inter * (h / 32))?; - let out = gpu.alloc(32 * 2 * inter * 2)?; - Some(crate::layers::moe::Mxfp8GateUp { - proj: Mxfp8Proj { - w: Mxfp8Weight { data: w, scales: s }, - n: (2 * inter) as u32, - k: h as u32, - }, - gate_up_out: out, - silu_mul: crate::layers::try_kernel(gpu, "gemm", "fused_silu_mul"), - }) + // The gate‖up rows land in ONE contiguous `[2*inter, h]` weight (and + // `[2*inter, h/32]` scale) buffer: `gate_proj`/`up_proj` are views into + // its halves, and the ≤32-row fused arm streams it in a single pass. + // The store originals are reclaimed — no duplicate bytes stay resident. + let gu_w = gpu.alloc(2 * inter * h)?; + gpu.copy_d2d(gw, gu_w, inter * h)?; + gpu.copy_d2d(uw, gu_w.offset(inter * h), inter * h)?; + let gu_s = gpu.alloc(2 * inter * (h / 32))?; + gpu.copy_d2d(gs, gu_s, inter * (h / 32))?; + gpu.copy_d2d(us, gu_s.offset(inter * (h / 32)), inter * (h / 32))?; + for proj in ["gate_proj", "up_proj"] { + for t in ["weight", "weight_scale"] { + store.reclaim(gpu, &format!("{lp}.mlp.shared_expert.{proj}.{t}"))?; + } + } + let view = |rows_off: usize| Mxfp8Proj { + w: Mxfp8Weight { + data: gu_w.offset(rows_off * h), + scales: gu_s.offset(rows_off * (h / 32)), + }, + n: inter as u32, + k: h as u32, }; + let gate_up = Some(crate::layers::moe::Mxfp8GateUp { + proj: Mxfp8Proj { + w: Mxfp8Weight { + data: gu_w, + scales: gu_s, + }, + n: (2 * inter) as u32, + k: h as u32, + }, + gate_up_out: gpu.alloc(32 * 2 * inter * 2)?, + silu_mul: crate::layers::try_kernel(gpu, "gemm", "fused_silu_mul"), + }); // The shared MXFP8 BF16 prefill scratch: the MoE runs after the // layer's attention/GDN on the same stream, so the transient buffer // is never live twice. let scratch = mx8_common::scratch_for(mx8_scratch, config, gpu)?; moe.set_mxfp8_shared_expert( - m("gate_proj", inter, h)?, - m("up_proj", inter, h)?, + view(0), + view(inter), m("down_proj", h, inter)?, gate_up, Mxfp8ProjKernels::resolve(gpu)?, From c5c33e2e5b26f1dbc323ee8848c2e4bac4654bb5 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 05:16:58 -0400 Subject: [PATCH 10/26] ple: decode NVFP4 n-gram rows on the GPU, scale the hot cache FP4-Local's PLE table ships as packed E2M1 rows + E4M3 block scales on NVMe. Every gather used to copy both to the host, decode to BF16 there, upload, and synchronize the stream. Now the block-scale rows fault into the same cache as the weight rows (new segmented scale source on NgramRowCache) and a new batched_embed_nvfp4 kernel decodes them in place -- bit-identical to the host oracle, no per-gather synchronize, no copy_slot. Bookkeeping scales to tens of millions of slots: u32 row ids, O(pins) end_batch, and ATLAS_PLE_CACHE_GB sizes the arena (slot count stays an env override, the span-derived floor still applies). Adds a 30 s residency/hit-rate/fault-time line plus a shutdown summary. --- crates/spark-model/Cargo.toml | 4 + .../examples/ngram_nvfp4_microtest.rs | 167 ++++++++++++++++++ crates/spark-model/src/layers/ops/sampling.rs | 28 +++ .../spark-model/src/layers/ple/aux_state.rs | 19 ++ crates/spark-model/src/layers/ple/layer.rs | 27 ++- .../spark-model/src/layers/ple/ngram_nvfp4.rs | 51 +++--- crates/spark-model/src/layers/ple/verify.rs | 16 +- .../src/weight_loader/qwen4_exp/ple.rs | 55 +++++- .../src/weight_loader/qwen4_exp/ple_nvfp4.rs | 23 +-- crates/spark-storage/src/ngram_cache.rs | 94 +++++++++- .../src/ngram_cache/fault_pool.rs | 47 ++++- .../src/ngram_cache/fault_pool_tests.rs | 3 +- .../spark-storage/src/ngram_cache/resolve.rs | 49 +++-- .../src/ngram_cache/scale_tests.rs | 159 +++++++++++++++++ .../src/ngram_cache/segmented.rs | 65 ++++++- crates/spark-storage/src/ngram_cache/tests.rs | 1 + crates/spark-storage/src/ngram_cache_fault.rs | 44 ++++- kernels/gb10/common/embed_from_argmax.cu | 37 ++++ 18 files changed, 766 insertions(+), 123 deletions(-) create mode 100644 crates/spark-model/examples/ngram_nvfp4_microtest.rs create mode 100644 crates/spark-storage/src/ngram_cache/scale_tests.rs diff --git a/crates/spark-model/Cargo.toml b/crates/spark-model/Cargo.toml index 86e72f238..2290d1692 100644 --- a/crates/spark-model/Cargo.toml +++ b/crates/spark-model/Cargo.toml @@ -211,6 +211,10 @@ required-features = ["cuda", "gpu-examples"] name = "mxfp8_proj_microtest" required-features = ["cuda", "gpu-examples"] +[[example]] +name = "ngram_nvfp4_microtest" +required-features = ["cuda", "gpu-examples"] + [[example]] name = "q2_0_gemv_microtest" required-features = ["cuda", "gpu-examples"] diff --git a/crates/spark-model/examples/ngram_nvfp4_microtest.rs b/crates/spark-model/examples/ngram_nvfp4_microtest.rs new file mode 100644 index 000000000..ccbb6e270 --- /dev/null +++ b/crates/spark-model/examples/ngram_nvfp4_microtest.rs @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: AGPL-3.0-only +//! `batched_embed_nvfp4` vs the host oracle `nvfp4_rows_to_bf16`, bit-identical. +//! (a) all 256 scale codes x 16 nibble pairs x real global; (b) real FP4-Local +//! shard rows through the segmented scale path. Run: +//! cargo run -p spark-model --release --features cuda,gpu-examples \ +//! --example ngram_nvfp4_microtest [-- /path/to/snapshot] + +use anyhow::Result; +use atlas_core::modelopt_codec::nvfp4_rows_to_bf16; +use spark_runtime::cuda_backend::AtlasCudaBackend; +use spark_runtime::gpu::{DevicePtr, GpuBackend}; +use spark_runtime::kernel_args::KernelLaunch; + +const DIM: usize = 160; +const GROUP: usize = 16; +const GLOBAL: f32 = 3.324236e-5; + +fn up(g: &dyn GpuBackend, d: &[u8]) -> Result { + let p = g.alloc(d.len().max(1))?; + g.copy_h2d(d, p)?; + Ok(p) +} +fn down(g: &dyn GpuBackend, p: DevicePtr, n: usize) -> Result> { + let mut b = vec![0u8; n]; + g.copy_d2h(p, &mut b)?; + Ok(b) +} + +fn gather( + g: &dyn GpuBackend, + k: spark_runtime::gpu::KernelHandle, + slots: &DevicePtr, + packed: DevicePtr, + scales: DevicePtr, + out: DevicePtr, + n: usize, +) -> Result<()> { + KernelLaunch::new(g, k) + .grid([n as u32, 1, 1]) + .block([256, 1, 1]) + .arg_ptr(*slots) + .arg_ptr(packed) + .arg_ptr(scales) + .arg_f32(GLOBAL) + .arg_ptr(out) + .arg_u32(DIM as u32) + .arg_u32(GROUP as u32) + .launch(0) +} + +fn main() -> Result<()> { + let gpu = AtlasCudaBackend::new(0, &atlas_kernels::ptx_modules())?; + let k = gpu.kernel("embed_from_argmax", "batched_embed_nvfp4")?; + + // (a) exhaustive scale-code coverage: 256 scale codes x 256 packed bytes. + let n = 256; + let mut packed = vec![0u8; n * DIM / 2]; + let mut scales = vec![0u8; n * DIM / GROUP]; + for r in 0..n { + for i in 0..DIM / 2 { + packed[r * DIM / 2 + i] = ((i + r) & 0xFF) as u8; + } + for i in 0..DIM / GROUP { + scales[r * DIM / GROUP + i] = r as u8; + } + } + let slots_host: Vec = (0..n as u32).flat_map(|s| s.to_le_bytes()).collect(); + let (pd, sd, out) = ( + up(&gpu, &packed)?, + up(&gpu, &scales)?, + gpu.alloc(n * DIM * 2)?, + ); + let slots = up(&gpu, &slots_host)?; + gather(&gpu, k, &slots, pd, sd, out, n)?; + gpu.synchronize(0)?; + let got = down(&gpu, out, n * DIM * 2)?; + let mut want = vec![0u8; n * DIM * 2]; + for r in 0..n { + nvfp4_rows_to_bf16( + &packed[r * DIM / 2..(r + 1) * DIM / 2], + &scales[r * DIM / GROUP..(r + 1) * DIM / GROUP], + GLOBAL, + DIM, + GROUP, + &mut want[r * DIM * 2..(r + 1) * DIM * 2], + )?; + } + anyhow::ensure!(got == want, "synthetic sweep: outputs differ"); + println!("synthetic sweep (256 scale codes x all bytes): IDENTICAL"); + + // (b) real rows: read FP4-Local shard_0 weight+scale through the cache. + if let Some(snap) = std::env::args().nth(1) { + use std::io::Read; + let find = |name: &str| -> Result<(std::path::PathBuf, u64, usize)> { + for e in std::fs::read_dir(std::path::Path::new(&snap).join("tensors"))? { + let p = e?.path(); + if p.extension().map(|x| x == "safetensors").unwrap_or(false) { + let b = std::fs::read(&p)?; + let nlen = u64::from_le_bytes(b[..8].try_into()?) as usize; + let hdr: serde_json::Value = serde_json::from_slice(&b[8..8 + nlen])?; + if let Some(t) = hdr.get(name) { + let off: Vec = serde_json::from_value(t["data_offsets"].clone())?; + return Ok((p, 8 + nlen as u64 + off[0], (off[1] - off[0]) as usize)); + } + } + } + anyhow::bail!("{name} not found") + }; + let name = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight"; + let sname = format!("{name}_scale"); + let (wp, woff, wlen) = find(name)?; + let (sp, soff, slen) = find(&sname)?; + let rows = wlen / (DIM / 2); + anyhow::ensure!(slen == rows * DIM / GROUP, "scale rows mismatch"); + let n2 = rows.min(4096); + let mut cache = spark_storage::NgramRowCache::open_segmented( + &vec![(wp.clone(), woff)], + rows as u64, + None, + DIM / 2, + n2 + 16, + )?; + cache.attach_segmented_block_scales(&vec![(sp.clone(), soff)], DIM / GROUP)?; + let _ = &mut cache; + let ids: Vec = (0..n2 as u64).map(|i| i * 7919 % rows as u64).collect(); + let mut slots_v = Vec::new(); + cache.resolve(&ids, &mut slots_v)?; + let sbytes: Vec = slots_v.iter().flat_map(|s| s.to_le_bytes()).collect(); + let slots = up(&gpu, &sbytes)?; + let out2 = gpu.alloc(n2 * DIM * 2)?; + gather( + &gpu, + k, + &slots, + DevicePtr(cache.table_dev_va()?), + DevicePtr(cache.block_scale_dev_va()?.unwrap()), + out2, + n2, + )?; + gpu.synchronize(0)?; + let got2 = down(&gpu, out2, n2 * DIM * 2)?; + let mut want2 = vec![0u8; n2 * DIM * 2]; + for (i, &id) in ids.iter().enumerate() { + let mut f = std::fs::File::open(&wp)?; + let mut pb = vec![0u8; DIM / 2]; + use std::io::Seek; + f.seek(std::io::SeekFrom::Start(woff + id * (DIM / 2) as u64))?; + f.read_exact(&mut pb)?; + let mut sf = std::fs::File::open(&sp)?; + let mut sb = vec![0u8; DIM / GROUP]; + sf.seek(std::io::SeekFrom::Start(soff + id * (DIM / GROUP) as u64))?; + sf.read_exact(&mut sb)?; + nvfp4_rows_to_bf16( + &pb, + &sb, + GLOBAL, + DIM, + GROUP, + &mut want2[i * DIM * 2..(i + 1) * DIM * 2], + )?; + } + anyhow::ensure!(got2 == want2, "real rows: outputs differ"); + println!("{n2} real FP4-Local rows through the segmented cache: IDENTICAL"); + } + println!("PASS"); + Ok(()) +} diff --git a/crates/spark-model/src/layers/ops/sampling.rs b/crates/spark-model/src/layers/ops/sampling.rs index 9c6e331f0..25b6fa56c 100644 --- a/crates/spark-model/src/layers/ops/sampling.rs +++ b/crates/spark-model/src/layers/ops/sampling.rs @@ -138,6 +138,34 @@ pub fn embed_from_argmax( .arg_u32(hidden_size) .launch(stream) } +/// NVFP4 n-gram gather: slots index a `[slots, dim/2]` packed-E2M1 arena and +/// a `[slots, dim/group]` E4M3 block-scale arena. Kernel: +/// `embed_from_argmax.cu::batched_embed_nvfp4`. +pub fn batched_embed_nvfp4( + gpu: &dyn GpuBackend, + kernel: KernelHandle, + slots_dev: DevicePtr, + packed: DevicePtr, + scales: DevicePtr, + global: f32, + output: DevicePtr, + num_tokens: u32, + dim: u32, + group: u32, + stream: u64, +) -> Result<()> { + KernelLaunch::new(gpu, kernel) + .grid([num_tokens, 1, 1]) + .block([256, 1, 1]) + .arg_ptr(slots_dev) + .arg_ptr(packed) + .arg_ptr(scales) + .arg_f32(global) + .arg_ptr(output) + .arg_u32(dim) + .arg_u32(group) + .launch(stream) +} /// Batched embedding: gather N rows from embedding table in one launch. /// diff --git a/crates/spark-model/src/layers/ple/aux_state.rs b/crates/spark-model/src/layers/ple/aux_state.rs index 5234595ad..d023094a1 100644 --- a/crates/spark-model/src/layers/ple/aux_state.rs +++ b/crates/spark-model/src/layers/ple/aux_state.rs @@ -202,6 +202,25 @@ impl PleLayer { // Dispatch on the PRESENCE OF THE SCALE, not a dtype flag: it keeps // "the bytes are E4M3" and "there is something to multiply by" from // being able to disagree. + if let Some(nv) = &self.nvfp4 { + let scale_va = self + .block_scale_va + .context("PLE: NVFP4 table without its block-scale arena")?; + return ops::batched_embed_nvfp4( + gpu, + self.embed_nvfp4_k, + self.slots_dev, + DevicePtr(table_va), + DevicePtr(scale_va), + nv.geom.global, + self.emb, + (num_tokens * heads) as u32, + nv.geom.dim as u32, + nv.geom.group as u32, + stream, + ) + .context("PLE row gather (nvfp4)"); + } match self.scale_va { Some(scale) => ops::batched_embed_fp8( gpu, diff --git a/crates/spark-model/src/layers/ple/layer.rs b/crates/spark-model/src/layers/ple/layer.rs index 5077489a3..bec527d82 100644 --- a/crates/spark-model/src/layers/ple/layer.rs +++ b/crates/spark-model/src/layers/ple/layer.rs @@ -61,9 +61,14 @@ pub struct PleLayer { /// Device VA of the table's per-slot f32 dequant scale array, `None` for a /// BF16 table. Its presence is what selects the FP8 gather. scale_va: Option, + /// NVFP4 gather kernel (`Some` when `nvfp4` is set). + embed_nvfp4_k: KernelHandle, + /// Device VA of the cache's `[slots, dim/group]` E4M3 block-scale arena. + block_scale_va: Option, /// Packed EXL3 n-gram rows. Host-dequanted into `emb` before the GEMM. trellis: Option, - /// NVFP4 n-gram rows (FP4-Local). Host-dequanted into `emb` likewise. + /// NVFP4 n-gram rows (FP4-Local): packed rows + block scales stay in the + /// pinned caches and decode inside `batched_embed_nvfp4`. nvfp4: Option, gemm_k: KernelHandle, gate_k: KernelHandle, @@ -101,7 +106,7 @@ impl PleLayer { /// Rows decoded on the host into `emb` (EXL3 trellis, FP4-Local NVFP4) /// rather than gathered by the embedding kernel. pub(super) fn host_decoded(&self) -> bool { - self.trellis.is_some() || self.nvfp4.is_some() + self.trellis.is_some() } #[allow(clippy::too_many_arguments)] @@ -149,6 +154,7 @@ impl PleLayer { gather_matches_element_size(stride, head_dim, scaled)?; } } + let is_nvfp4 = nvfp4.is_some(); Ok(Self { dims, head_dim, @@ -169,12 +175,27 @@ impl PleLayer { // table mutex on every gather. scale_va: match &table { #[cfg(feature = "cuda")] - NgramTable::Cached(c) => c.scale_dev_va()?, + NgramTable::Cached(c) => match c.scale_dev_va()? { + // A segmented block-scale arena is the NVFP4 table's, not + // an FP8 row-scale array — it must not select the FP8 arm. + Some(_) if is_nvfp4 => None, + other => other, + }, + _ => None, + }, + block_scale_va: match &table { + #[cfg(feature = "cuda")] + NgramTable::Cached(c) => c.block_scale_dev_va()?, _ => None, }, trellis, nvfp4, embed_fp8_k: gpu.kernel("embed_from_argmax", "batched_embed_fp8")?, + embed_nvfp4_k: if is_nvfp4 { + gpu.kernel("embed_from_argmax", "batched_embed_nvfp4")? + } else { + KernelHandle(0) + }, table: std::sync::Arc::new(std::sync::Mutex::new(table)), embed_k: gpu.kernel("embed_from_argmax", "batched_embed")?, gemm_k: projection_gemm(gpu)?, diff --git a/crates/spark-model/src/layers/ple/ngram_nvfp4.rs b/crates/spark-model/src/layers/ple/ngram_nvfp4.rs index a47ca19be..aa2694fc5 100644 --- a/crates/spark-model/src/layers/ple/ngram_nvfp4.rs +++ b/crates/spark-model/src/layers/ple/ngram_nvfp4.rs @@ -14,7 +14,10 @@ //! 26.8 GiB on disk against the NVFP4 pack's 47.7 GiB FP8 table, with the //! same row count: fewer bytes per fault, more rows per page-cache GiB. -use anyhow::{Context, Result, ensure}; +use anyhow::{Result, ensure}; +#[cfg(test)] +use anyhow::Context as _; +#[cfg(test)] use atlas_core::modelopt_codec::nvfp4_rows_to_bf16; /// Row geometry of an NVFP4 n-gram table. @@ -28,13 +31,14 @@ pub struct Nvfp4RowGeom { pub global: f32, } -/// An NVFP4 n-gram table's geometry plus its block-scale row cache. +/// An NVFP4 n-gram table's decode geometry; its block-scale rows fault in +/// through the WEIGHT cache's segmented scale source (`attach_segmented_ +/// block_scales`), so one CLOCK serves both arenas. pub struct NgramNvfp4 { pub geom: Nvfp4RowGeom, - /// Block-scale rows, `dim / group` bytes each, at the same row ids as the - /// weight cache. - #[cfg(feature = "cuda")] - pub scales: std::sync::Mutex, + /// `(file, base_offset)` of each scale shard — handed to + /// `NgramRowCache::attach_segmented_block_scales` by the loader. + pub scale_shards: Vec<(std::path::PathBuf, u64)>, } impl Nvfp4RowGeom { @@ -47,13 +51,13 @@ impl Nvfp4RowGeom { } } -/// The weight cache must hold packed rows and nothing else: a scale array -/// would select the FP8 gather, and a wrong stride reads the next row's bytes. +/// The weight cache must hold packed rows, and the segmented block-scale +/// arena must be attached: the NVFP4 gather reads both. pub fn check_packed_stride(stride: usize, spec: &Nvfp4RowGeom, scaled: bool) -> Result<()> { ensure!( - !scaled, - "PLE: NVFP4 n-gram rows take their scales from a block-scale table, \ - not a cache scale array" + scaled, + "PLE: NVFP4 n-gram rows need the block-scale arena attached to the \ + row cache (`attach_segmented_block_scales`)" ); ensure!( stride == spec.packed_stride(), @@ -64,8 +68,11 @@ pub fn check_packed_stride(stride: usize, spec: &Nvfp4RowGeom, scaled: bool) -> Ok(()) } -/// Decode gathered rows: `packed[i]` and `scales[i]` are row `i`'s bytes from -/// the weight and block-scale caches. Returns little-endian BF16 `[n, dim]`. +/// Decode gathered rows — kept as the oracle for the GPU microtest and +/// the host-side cache tests; the serving path decodes in +/// `batched_embed_nvfp4`. `packed[i]` and `scales[i]` are row `i`'s bytes +/// from the weight and block-scale caches. Returns little-endian BF16. +#[cfg(test)] pub fn decode_rows_bf16( spec: &Nvfp4RowGeom, packed: &[Vec], @@ -94,23 +101,7 @@ pub fn decode_rows_bf16( Ok(out) } -/// Resolve `ids` in the block-scale cache, copy each row out, and release the -/// pins at once: the decode that consumes them runs synchronously on the host. -#[cfg(feature = "cuda")] -pub fn gather_scales(spec: &NgramNvfp4, ids: &[u64]) -> Result>> { - let mut cache = spec - .scales - .lock() - .map_err(|_| anyhow::anyhow!("PLE nvfp4 scale cache mutex poisoned"))?; - let mut slots = Vec::with_capacity(ids.len()); - cache.resolve(ids, &mut slots)?; - let rows = slots - .iter() - .map(|&s| cache.copy_slot(s)) - .collect::>>(); - cache.end_batch(); - rows -} + #[cfg(test)] #[path = "ngram_nvfp4_tests.rs"] diff --git a/crates/spark-model/src/layers/ple/verify.rs b/crates/spark-model/src/layers/ple/verify.rs index 7a3c45ac9..2bb34e104 100644 --- a/crates/spark-model/src/layers/ple/verify.rs +++ b/crates/spark-model/src/layers/ple/verify.rs @@ -139,21 +139,7 @@ impl PleLayer { let (h0, m0, _) = cache.stats(); let t0 = std::time::Instant::now(); cache.resolve(ids, &mut slots)?; - if let Some(spec) = &self.nvfp4 { - // FP4-Local: weight rows from this cache, block-scale rows - // from the companion cache at the same ids. - let mut raws = Vec::with_capacity(slots.len()); - for &slot in &slots { - raws.push(cache.copy_slot(slot)?); - } - let scales = super::ngram_nvfp4::gather_scales(spec, ids)?; - let decoded = super::ngram_nvfp4::decode_rows_bf16(&spec.geom, &raws, &scales) - .map_err(|e| anyhow::anyhow!("PLE: ngram nvfp4 dequant: {e:#}"))?; - gpu.copy_h2d_async(&decoded, self.emb, stream)?; - // Same contract as the trellis branch below: `decoded` is a - // temporary, and a graph replay must see the finished upload. - gpu.synchronize(stream)?; - } else if let Some(spec) = &self.trellis { + if let Some(spec) = &self.trellis { let mut raws = Vec::with_capacity(slots.len()); for &slot in &slots { raws.push(cache.copy_slot(slot)?); diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs b/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs index 8dacdbd9e..8231128ae 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs @@ -30,6 +30,23 @@ //! do it silently. #[cfg(feature = "cuda")] +// The NVFP4 block-scale slot width, when the table is packed U8. +fn nvfp4_scale_stride( + store: &spark_runtime::weights::WeightStore, + lp: &str, + dtype: spark_runtime::weights::WeightDtype, +) -> usize { + if dtype != spark_runtime::weights::WeightDtype::UInt8 { + return 0; + } + store + .deferred(&format!( + "{lp}.ple_embedding.ngram_embedding.shard_0.weight_scale" + )) + .map(|d| d.shape[1]) + .unwrap_or(0) +} + use anyhow::Context; use anyhow::Result; use atlas_core::config::ModelConfig; @@ -79,17 +96,32 @@ fn chunk_from_env() -> usize { } #[cfg(feature = "cuda")] -fn slots_from_env(scratch_tokens: usize, ngram_heads: usize) -> (usize, &'static str) { - match std::env::var("ATLAS_PLE_CACHE_SLOTS") +fn slots_from_env( + scratch_tokens: usize, + ngram_heads: usize, + slot_bytes: usize, +) -> (usize, &'static str) { + // Exact slots beat GB sizing; GB beats the derived floor. The derived + // span figure stays a MINIMUM — the cache must always hold one gather. + let floor = derived_slots(scratch_tokens, ngram_heads); + if let Some(n) = std::env::var("ATLAS_PLE_CACHE_SLOTS") .ok() .and_then(|v| v.parse().ok()) + && n > 0 { - Some(n) if n > 0 => (n, "ATLAS_PLE_CACHE_SLOTS"), - _ => ( - derived_slots(scratch_tokens, ngram_heads), - "(span + warm_ahead)*heads rounded up", - ), + return (n, "ATLAS_PLE_CACHE_SLOTS"); } + if let Some(gb) = std::env::var("ATLAS_PLE_CACHE_GB") + .ok() + .and_then(|v| v.parse::().ok()) + && gb > 0.0 + { + return ( + (((gb * 1e9) / slot_bytes as f64) as usize).max(floor), + "ATLAS_PLE_CACHE_GB", + ); + } + (floor, "(span + warm_ahead)*heads rounded up") } /// First uploaded name that exists. EXL3 checkpoints use the nested @@ -269,7 +301,8 @@ pub(super) fn load( // `PleLayer::new` applies, keeping the two derivations from drifting. let chunk = chunk_from_env(); let span = crate::layers::ple::bounded_scratch(chunk, max_tokens); - let (slots, slots_from) = slots_from_env(span, heads); + let slot_bytes = head_dim * elem + nvfp4_scale_stride(store, &lp, dtype); + let (slots, slots_from) = slots_from_env(span, heads, slot_bytes); let mut cache = if single.is_some() { let (path, off) = &shards[0]; spark_storage::NgramRowCache::open_at( @@ -301,7 +334,11 @@ pub(super) fn load( } let n_shards = shards.len(); - let nvfp4 = nvfp4::open_if_packed(store, &lp, dtype, n_shards, rows_per, head_dim, slots, gpu)?; + let nvfp4 = nvfp4::open_if_packed(store, &lp, dtype, n_shards, rows_per, head_dim, gpu)?; + if let Some(nv) = &nvfp4 { + // The block scales ride the same cache: one arena pair, one CLOCK. + cache.attach_segmented_block_scales(&nv.scale_shards, nv.geom.scale_stride())?; + } // FP8 rows need their dequant scale, or the gather returns raw E4M3 // magnitudes and the whole n-gram contribution is off by a constant diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/ple_nvfp4.rs b/crates/spark-model/src/weight_loader/qwen4_exp/ple_nvfp4.rs index 0a55e6f00..e11e4c921 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/ple_nvfp4.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/ple_nvfp4.rs @@ -26,13 +26,12 @@ pub(super) fn open_if_packed( shards: usize, rows_per: usize, packed_cols: usize, - slots: usize, gpu: &dyn GpuBackend, ) -> Result> { if dtype != WeightDtype::UInt8 { return Ok(None); } - open(store, lp, shards, rows_per, packed_cols, slots, gpu).map(Some) + open(store, lp, shards, rows_per, packed_cols, gpu).map(Some) } fn open( @@ -41,7 +40,6 @@ fn open( shards: usize, rows_per: usize, packed_cols: usize, - slots: usize, gpu: &dyn GpuBackend, ) -> Result { let table = format!("{lp}.ple_embedding.ngram_embedding"); @@ -77,22 +75,5 @@ fn open( group: dim / scale_cols, global: f32_scalar(store, &format!("{table}.weight_scale_2"), gpu)?, }; - let scales = spark_storage::NgramRowCache::open_segmented( - &scale_shards, - rows_per as u64, - None, - scale_cols, - slots, - ) - .context("PLE nvfp4: block-scale row cache")?; - tracing::info!( - "PLE n-gram table: NVFP4 ({dim} dims, group {}, global {:.6e}), block scales \ - served off NVMe from {shards} sibling shards", - geom.group, - geom.global - ); - Ok(NgramNvfp4 { - geom, - scales: std::sync::Mutex::new(scales), - }) + Ok(NgramNvfp4 { geom, scale_shards }) } diff --git a/crates/spark-storage/src/ngram_cache.rs b/crates/spark-storage/src/ngram_cache.rs index 7f101da93..af8946bcb 100644 --- a/crates/spark-storage/src/ngram_cache.rs +++ b/crates/spark-storage/src/ngram_cache.rs @@ -77,14 +77,17 @@ pub struct NgramRowCache { row_stride: usize, slots: usize, rows_total: u64, - /// row_id -> slot. - map: HashMap, - /// slot -> resident row id (`u64::MAX` = empty). - slot_row: Vec, + /// row_id -> slot (row ids are < 2^32, enforced at open). + map: HashMap, + /// slot -> resident row id (`u32::MAX` = empty). + slot_row: Vec, /// CLOCK reference bits. refbit: Vec, /// Slots pinned for the batch in flight (never evicted). pinned: Vec, + /// The pinned slots of the batch in flight, so `end_batch` is O(pins), + /// never O(slots). + pinned_list: Vec, hand: usize, pub hits: u64, pub misses: u64, @@ -93,6 +96,10 @@ pub struct NgramRowCache { keepalive: Option, /// `ATLAS_PLE_FAULT_POOL=1` workers, built on the first multi-row fault. fault_pool: Option, + /// Cumulative wall time of the fault passes (pass 2 of `resolve_inner`). + pub fault_us: u64, + /// Throttle for the periodic residency line (item: observability). + last_stats_log: std::time::Instant, } /// A table split across equal-sized shards at scattered file offsets, which @@ -126,7 +133,21 @@ struct Segments { /// value, written once at open and never touched again. struct ScaleCache { arena: ExpertArena, - file: Option, + /// Bytes of scale payload per slot (4 for f32, `dim/group` E4M3 for NVFP4). + stride: usize, + src: ScaleSrc, +} + +/// Where a scale row comes from on a fault. +enum ScaleSrc { + /// Per-tensor constant, written once (no fault refresh). + Constant, + /// One f32 per row in a single file at `id * stride`. + PerRowFile(File), + /// Per-shard block-scale rows at scattered offsets — the NVFP4 PLE table's + /// sibling `shard_i.weight_scale` tensors, mapped by the same + /// `rows_per`/dedupe logic as the weight shards. + Segmented(Segments), } /// A 4 KiB-aligned host buffer for O_DIRECT reads. @@ -178,6 +199,10 @@ impl NgramRowCache { row_stride: usize, slots: usize, ) -> Result { + anyhow::ensure!( + rows_total <= u32::MAX as u64, + "NgramRowCache: {rows_total} rows exceed the u32 bookkeeping limit" + ); if row_stride == 0 || slots == 0 { bail!("NgramRowCache: zero geometry (row_stride={row_stride}, slots={slots})"); } @@ -202,7 +227,8 @@ impl NgramRowCache { Some(ScaleCache { arena: ExpertArena::new(1, sblocks as u32, BLOCK) .context("NgramRowCache: scale arena")?, - file: Some(open_direct(sp)?), + stride: 4, + src: ScaleSrc::PerRowFile(open_direct(sp)?), }) } None => None, @@ -216,16 +242,19 @@ impl NgramRowCache { row_stride, slots, rows_total, - map: HashMap::with_capacity(slots * 2), - slot_row: vec![u64::MAX; slots], + map: HashMap::with_capacity(slots), + slot_row: vec![u32::MAX; slots], refbit: vec![false; slots], pinned: vec![false; slots], + pinned_list: Vec::new(), hand: 0, hits: 0, misses: 0, evictions: 0, keepalive: None, fault_pool: None, + fault_us: 0, + last_stats_log: std::time::Instant::now(), }) } @@ -314,6 +343,37 @@ impl NgramRowCache { pub fn stats(&self) -> (u64, u64, u64) { (self.hits, self.misses, self.evictions) } + + /// Resident slots (map size) and total slots — the "n/N" of the log line. + pub fn occupancy(&self) -> (usize, usize) { + (self.map.len(), self.slots) + } + + /// The periodic observability line: at most once per 30 s while the + /// cache is being resolved against. Callers: `resolve`/`prefetch`. + pub(crate) fn maybe_log_stats(&mut self) { + let dt = self.last_stats_log.elapsed(); + if dt < std::time::Duration::from_secs(30) { + return; + } + self.last_stats_log = std::time::Instant::now(); + let total = self.hits + self.misses; + let rate = if total > 0 { + self.hits as f64 * 100.0 / total as f64 + } else { + 0.0 + }; + tracing::info!( + "NgramRowCache: {}/{} slots resident, {} hits / {} misses ({:.1}% hit), {} evictions, fault wall {:.1}s cumulative", + self.map.len(), + self.slots, + self.hits, + self.misses, + rate, + self.evictions, + self.fault_us as f64 / 1e6 + ); + } } /// The decide+fault cluster (`resolve`, `prefetch`, `resolve_inner`, @@ -353,5 +413,23 @@ fn open_direct(path: &Path) -> Result { /// Segmented (multi-shard, multi-file) tables and per-tensor FP8 scales. mod segmented; +#[cfg(test)] +#[path = "ngram_cache/scale_tests.rs"] +mod scale_tests; #[cfg(test)] mod tests; + +impl Drop for NgramRowCache { + /// One shutdown line: the same stats the periodic line prints. + fn drop(&mut self) { + if self.hits + self.misses > 0 { + tracing::info!( + "NgramRowCache closed: {} hits / {} misses, {} evictions, fault wall {:.1}s", + self.hits, + self.misses, + self.evictions, + self.fault_us as f64 / 1e6 + ); + } + } +} diff --git a/crates/spark-storage/src/ngram_cache/fault_pool.rs b/crates/spark-storage/src/ngram_cache/fault_pool.rs index 4b381897e..e0ba23a96 100644 --- a/crates/spark-storage/src/ngram_cache/fault_pool.rs +++ b/crates/spark-storage/src/ngram_cache/fault_pool.rs @@ -27,8 +27,8 @@ use std::thread::JoinHandle; use anyhow::{Context, Result, anyhow}; -use super::fault::{ArenaPtrs, Fault, RowSource, fetch_row}; -use super::{AlignedBlock, NgramRowCache, Segments}; +use super::fault::{ArenaPtrs, Fault, RowSource, ScaleSource, fetch_row}; +use super::{AlignedBlock, NgramRowCache, ScaleSrc, Segments}; /// The cache's immutable row geometry, owned (cloned descriptors) so 'static /// workers can rebuild a `RowSource` from it. @@ -37,7 +37,23 @@ struct OwnedSource { base_offset: u64, segments: Option, row_stride: usize, - scale_file: Option, + scale_src: Option, + scale_stride: usize, +} + +/// Owned mirror of [`ScaleSource`] (cloned descriptors for 'static workers). +enum OwnedScaleSrc { + Flat(File), + Segmented(Segments), +} + +impl OwnedScaleSrc { + fn view(&self) -> ScaleSource<'_> { + match self { + OwnedScaleSrc::Flat(f) => ScaleSource::Flat { file: f }, + OwnedScaleSrc::Segmented(s) => ScaleSource::Segmented { seg: s }, + } + } } impl OwnedSource { @@ -47,7 +63,8 @@ impl OwnedSource { base_offset: self.base_offset, segments: self.segments.as_ref(), row_stride: self.row_stride, - scale_file: self.scale_file.as_ref(), + scale_src: self.scale_src.as_ref().map(|s| s.view()), + scale_stride: self.scale_stride, } } } @@ -170,12 +187,23 @@ impl FaultPool { shard_file: s.shard_file.clone(), }), }; - let scale_file = match cache.scales.as_ref().and_then(|sc| sc.file.as_ref()) { - Some(f) => Some( + let scale_src = match cache.scales.as_ref().map(|sc| &sc.src) { + None | Some(ScaleSrc::Constant) => None, + Some(ScaleSrc::PerRowFile(f)) => Some(OwnedScaleSrc::Flat( f.try_clone() .context("NgramRowCache fault pool: dup scales")?, - ), - None => None, + )), + Some(ScaleSrc::Segmented(s)) => Some(OwnedScaleSrc::Segmented(Segments { + bases: s.bases.clone(), + rows_per: s.rows_per, + files: s + .files + .iter() + .map(File::try_clone) + .collect::>() + .context("NgramRowCache fault pool: dup scale files")?, + shard_file: s.shard_file.clone(), + })), }; let src = OwnedSource { file: cache @@ -185,7 +213,8 @@ impl FaultPool { base_offset: cache.base_offset, segments, row_stride: cache.row_stride, - scale_file, + scale_src, + scale_stride: cache.scales.as_ref().map_or(0, |sc| sc.stride), }; Self::start(src, fetch_row, workers, spawn_worker) } diff --git a/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs b/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs index 2f284c5ca..7ce97630f 100644 --- a/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs +++ b/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs @@ -26,7 +26,8 @@ fn src() -> OwnedSource { base_offset: 0, segments: None, row_stride: 1, - scale_file: None, + scale_src: None, + scale_stride: 0, } } diff --git a/crates/spark-storage/src/ngram_cache/resolve.rs b/crates/spark-storage/src/ngram_cache/resolve.rs index 205ffa9ed..30d10ed6d 100644 --- a/crates/spark-storage/src/ngram_cache/resolve.rs +++ b/crates/spark-storage/src/ngram_cache/resolve.rs @@ -48,6 +48,7 @@ impl NgramRowCache { if let Some(k) = &self.keepalive { k.touch(); } + self.maybe_log_stats(); if let Some(out) = out_slots.as_deref_mut() { out.clear(); out.reserve(row_ids.len()); @@ -69,12 +70,13 @@ impl NgramRowCache { self.rows_total ); } - let slot = match self.map.get(&id) { + let slot = match self.map.get(&(id as u32)) { Some(&s) => { self.hits += 1; self.refbit[s as usize] = true; - if pin { + if pin && !self.pinned[s as usize] { self.pinned[s as usize] = true; + self.pinned_list.push(s); } s } @@ -94,11 +96,12 @@ impl NgramRowCache { return Err(e); } }; - self.map.insert(id, s); - self.slot_row[s as usize] = id; + self.map.insert(id as u32, s); + self.slot_row[s as usize] = id as u32; self.refbit[s as usize] = true; - if pin { + if pin && !self.pinned[s as usize] { self.pinned[s as usize] = true; + self.pinned_list.push(s); } faults.push(Fault { id, slot: s }); s @@ -121,7 +124,10 @@ impl NgramRowCache { // slot handed out above is pinned, and `victim` skips pinned slots), so // the arena writes are disjoint, and `read_exact_at` is positional so a // shared `&File` has no cursor to race on. - if let Err(e) = self.fetch_many(&faults) { + let ft = std::time::Instant::now(); + let r = self.fetch_many(&faults); + self.fault_us += ft.elapsed().as_micros() as u64; + if let Err(e) = r { // A partially-written slot must not stay claimed as resident: a // wrong n-gram row reads as fluent output with wrong logits, which // is the failure mode this cache exists to avoid. Drop the WHOLE @@ -143,11 +149,11 @@ impl NgramRowCache { /// symptom would be a wrong n-gram row rather than a crash. fn drop_reservations(&mut self, faults: &[Fault]) { for f in faults { - if self.map.get(&f.id) == Some(&f.slot) { - self.map.remove(&f.id); + if self.map.get(&(f.id as u32)) == Some(&f.slot) { + self.map.remove(&(f.id as u32)); } - if self.slot_row[f.slot as usize] == f.id { - self.slot_row[f.slot as usize] = u64::MAX; + if self.slot_row[f.slot as usize] == f.id as u32 { + self.slot_row[f.slot as usize] = u32::MAX; } } } @@ -164,13 +170,20 @@ impl NgramRowCache { base_offset: self.base_offset, segments: self.segments.as_ref(), row_stride: self.row_stride, - scale_file: self.scales.as_ref().and_then(|sc| sc.file.as_ref()), + scale_stride: self.scales.as_ref().map_or(0, |sc| sc.stride), + scale_src: self.scales.as_ref().and_then(|sc| match &sc.src { + super::ScaleSrc::Constant => None, + super::ScaleSrc::PerRowFile(f) => Some(ScaleSource::Flat { file: f }), + super::ScaleSrc::Segmented(seg) => Some(ScaleSource::Segmented { seg }), + }), }; let ptrs = ArenaPtrs { rows: self.arena.slot_host_ptr(0, 0)?, - scales: match self.scales.as_ref().filter(|sc| sc.file.is_some()) { - Some(sc) => Some(sc.arena.slot_host_ptr(0, 0)?), - None => None, + scales: match self.scales.as_ref() { + Some(sc) if !matches!(sc.src, super::ScaleSrc::Constant) => { + Some(sc.arena.slot_host_ptr(0, 0)?) + } + _ => None, }, }; @@ -219,7 +232,11 @@ impl NgramRowCache { /// Release the batch's pins (call after the gather kernels are issued). pub fn end_batch(&mut self) { - self.pinned.fill(false); + // O(pins), not O(slots) — a multi-GB cache has tens of millions of + // slots and a fill(false) sweep would dominate a decode step. + for s in self.pinned_list.drain(..) { + self.pinned[s as usize] = false; + } } /// CLOCK second-chance victim among the unpinned slots. @@ -234,7 +251,7 @@ impl NgramRowCache { self.refbit[s] = false; continue; } - if self.slot_row[s] != u64::MAX { + if self.slot_row[s] != u32::MAX { let old = self.slot_row[s]; self.map.remove(&old); self.evictions += 1; diff --git a/crates/spark-storage/src/ngram_cache/scale_tests.rs b/crates/spark-storage/src/ngram_cache/scale_tests.rs new file mode 100644 index 000000000..f21be841a --- /dev/null +++ b/crates/spark-storage/src/ngram_cache/scale_tests.rs @@ -0,0 +1,159 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Segmented block-scale coverage for `NgramRowCache` — the NVFP4 PLE +/// table's sibling `shard_i.weight_scale` rows, faulted through the same +/// slot as the weight row. Split out of `tests.rs` for the 500-LoC cap. + +use super::*; + +/// Pinned arena needs a live CUDA context (see ). +#[cfg(feature = "cuda")] +fn cuda_ctx() -> crate::cuda_min::CudaCtx { + crate::cuda_min::CudaCtx::new(0).expect("CUDA context") +} +#[cfg(not(feature = "cuda"))] +fn cuda_ctx() {} + +/// Segmented block scales ride the same cache: every fault lands the row AND +/// its scale row, across several files, seam-straddling included; a second +/// attach is refused; evicted rows refault correctly. +#[test] +#[ignore] +fn segmented_block_scales_fault_evict_and_refuse_second_attach() { + let _ctx = cuda_ctx(); + use std::io::Write; + + const STRIDE: usize = 160; + const SSTRIDE: usize = 20; // dim/group style scale row + const ROWS_PER: u64 = 64; + const SLOTS: usize = 100; // < total rows, so eviction happens + + let dir = std::env::temp_dir().join(format!("ngram_ss_{}", std::process::id())); + std::fs::create_dir_all(&dir).expect("tmpdir"); + let wa = dir.join("w_a.bin"); + let wb = dir.join("w_b.bin"); + let sa = dir.join("s_a.bin"); + let sb = dir.join("s_b.bin"); + + let wpat = |shard: usize, local: u64| -> Vec { + (0..STRIDE) + .map(|i| (shard as u8 * 31) ^ (local as u8) ^ i as u8) + .collect() + }; + let spat = |shard: usize, local: u64| -> Vec { + (0..SSTRIDE) + .map(|i| (shard as u8 * 17) ^ !(local as u8) ^ i as u8) + .collect() + }; + + // weight shards: 0,1 in wa / 2,3 in wb; scale shards: 0,1 in sa / 2,3 in sb + let wlayout = [(0usize, &wa), (1, &wa), (2, &wb), (3, &wb)]; + let slayout = [(0usize, &sa), (1, &sa), (2, &sb), (3, &sb)]; + let write = |layout: &[(usize, &std::path::PathBuf)], + pat: &dyn Fn(usize, u64) -> Vec| + -> Vec<(std::path::PathBuf, u64)> { + let mut offsets: Vec<(std::path::PathBuf, u64)> = Vec::new(); + for file in [wa.clone(), wb.clone(), sa.clone(), sb.clone()] { + if !layout.iter().any(|(_, f)| **f == file) { + continue; + } + let mut fh = std::fs::File::create(&file).expect("create"); + fh.write_all(&[0u8; 24]).expect("lead"); // non-4KiB-aligned base + let mut at = 24u64; + for (shard, f) in layout.iter() { + if *f != &file { + continue; + } + while offsets.len() <= *shard { + offsets.push((file.clone(), 0)); + } + offsets[*shard] = (file.clone(), at); + for local in 0..ROWS_PER { + fh.write_all(&pat(*shard, local)).expect("row"); + } + at += ROWS_PER * STRIDE as u64; // file padded generously + } + fh.write_all(&[0u8; 8192]).expect("pad"); + } + offsets + }; + // scale rows are shorter — write them with their own stride + let swrite = |layout: &[(usize, &std::path::PathBuf)]| -> Vec<(std::path::PathBuf, u64)> { + let mut offsets: Vec<(std::path::PathBuf, u64)> = vec![]; + for file in [sa.clone(), sb.clone()] { + let mut fh = std::fs::File::create(&file).expect("create"); + fh.write_all(&[0u8; 24]).expect("lead"); + let mut at = 24u64; + for (shard, f) in layout.iter() { + if *f != &file { + continue; + } + while offsets.len() <= *shard { + offsets.push((file.clone(), 0)); + } + offsets[*shard] = (file.clone(), at); + for local in 0..ROWS_PER { + fh.write_all(&spat(*shard, local)).expect("srow"); + } + at += ROWS_PER * SSTRIDE as u64; + } + fh.write_all(&[0u8; 8192]).expect("pad"); + } + offsets + }; + let wshards = write(&wlayout, &wpat); + let sshards = swrite(&slayout); + + let mut cache = NgramRowCache::open_segmented(&wshards, ROWS_PER, None, STRIDE, SLOTS) + .expect("segmented cache"); + cache + .attach_segmented_block_scales(&sshards, SSTRIDE) + .expect("attach segmented scales"); + assert!(cache.block_scale_dev_va().expect("va").is_some()); + assert!( + cache + .attach_segmented_block_scales(&sshards, SSTRIDE) + .is_err(), + "second attach must be refused" + ); + + // Resolve more DISTINCT rows than slots across batches -> forces + // evictions; every faulted slot must hold the matching row AND scale row. + let ids: Vec = (0..256u64).collect(); + for chunk in ids.chunks(64) { + let mut slots = Vec::new(); + cache.resolve(chunk, &mut slots).expect("resolve"); + cache.end_batch(); + } + let (h, m, ev) = cache.stats(); + assert!( + m > 0 && ev > 0, + "expected misses and evictions: h{h} m{m} ev{ev}" + ); + + // Refault a slice — all must land identical bytes. + let probe: Vec = (0..64u64).map(|i| i * 3 % 256).collect(); + let mut slots = Vec::new(); + cache.resolve(&probe, &mut slots).expect("resolve2"); + let sbase = cache + .scales + .as_ref() + .expect("scales") + .arena + .slot_host_ptr(0, 0) + .expect("scale arena") as *const u8; + for (i, &id) in probe.iter().enumerate() { + let slot = slots[i] as usize; + let shard = (id / ROWS_PER) as usize; + let local = id % ROWS_PER; + assert_eq!( + cache.slot_bytes(slot as u32).expect("row"), + &wpat(shard, local)[..], + "row {id}" + ); + let srow = unsafe { std::slice::from_raw_parts(sbase.add(slot * SSTRIDE), SSTRIDE) }; + assert_eq!(srow, &spat(shard, local)[..], "scale row {id}"); + } + cache.end_batch(); + std::fs::remove_dir_all(&dir).ok(); +} diff --git a/crates/spark-storage/src/ngram_cache/segmented.rs b/crates/spark-storage/src/ngram_cache/segmented.rs index d8d2278fa..59cb0c003 100644 --- a/crates/spark-storage/src/ngram_cache/segmented.rs +++ b/crates/spark-storage/src/ngram_cache/segmented.rs @@ -9,7 +9,7 @@ use anyhow::{Context, Result, bail}; -use super::{BLOCK, NgramRowCache, ScaleCache, Segments, open_direct}; +use super::{BLOCK, NgramRowCache, ScaleCache, ScaleSrc, Segments, open_direct}; use crate::expert_arena::ExpertArena; use std::path::Path; @@ -71,7 +71,7 @@ impl NgramRowCache { /// entire table with the wrong number. pub fn set_constant_scale(&mut self, scale: f32) -> Result<()> { if let Some(sc) = &self.scales - && sc.file.is_some() + && !matches!(sc.src, super::ScaleSrc::Constant) { bail!( "NgramRowCache: a per-row scale file is already installed; \ @@ -88,9 +88,68 @@ impl NgramRowCache { let all = std::slice::from_raw_parts_mut(base as *mut f32, self.slots); all.fill(scale); } - self.scales = Some(ScaleCache { arena, file: None }); + self.scales = Some(ScaleCache { + arena, + stride: 4, + src: ScaleSrc::Constant, + }); Ok(()) } + + /// Install SEGMENTED per-row block scales: `shards[i]` is shard `i`'s + /// backing file and the byte offset of its first scale row, at the same + /// `rows_per` shard split as the weight table (`open_segmented` must have + /// run first). Every fault then fetches the row's scale row into + /// `scales[slot * scale_stride..]`. Refuses if a scale source is already + /// installed — stacking two would silently pick one. + pub fn attach_segmented_block_scales( + &mut self, + shards: &[(std::path::PathBuf, u64)], + scale_stride: usize, + ) -> Result<()> { + anyhow::ensure!( + self.scales.is_none(), + "NgramRowCache: a scale source is already installed" + ); + anyhow::ensure!( + !shards.is_empty() && scale_stride > 0 && scale_stride <= BLOCK, + "NgramRowCache: bad segmented scale geometry (shards={}, stride={scale_stride})", + shards.len() + ); + let seg = self + .segments + .as_ref() + .context("NgramRowCache: segmented scales need a segmented table")?; + let rows_per = seg.rows_per; + let (paths, shard_file) = plan_shard_files(shards); + let mut files = Vec::with_capacity(paths.len()); + for path in &paths { + files.push(open_direct(path)?); + } + let sbytes = self.slots * scale_stride; + let sblocks = sbytes.div_ceil(BLOCK); + self.scales = Some(ScaleCache { + arena: ExpertArena::new(1, sblocks as u32, BLOCK) + .context("NgramRowCache: block-scale arena")?, + stride: scale_stride, + src: ScaleSrc::Segmented(Segments { + bases: shards.iter().map(|(_, off)| *off).collect(), + rows_per, + files, + shard_file, + }), + }); + Ok(()) + } + + /// Device VA of the block-scale arena (`[slots, scale_stride]` u8), for + /// the NVFP4 gather kernel. + pub fn block_scale_dev_va(&self) -> Result> { + match &self.scales { + Some(s) => Ok(Some(s.arena.slot_dev_va(0, 0)?)), + None => Ok(None), + } + } } /// Distinct backing paths in first-use order, and `shard -> path index`. diff --git a/crates/spark-storage/src/ngram_cache/tests.rs b/crates/spark-storage/src/ngram_cache/tests.rs index 047f05974..0ef1657f5 100644 --- a/crates/spark-storage/src/ngram_cache/tests.rs +++ b/crates/spark-storage/src/ngram_cache/tests.rs @@ -383,3 +383,4 @@ fn prefetch_dedups_repeated_ids() { assert_ne!(slots[0], slots[3]); let _ = std::fs::remove_dir_all(&dir); } + diff --git a/crates/spark-storage/src/ngram_cache_fault.rs b/crates/spark-storage/src/ngram_cache_fault.rs index 0ffc13efc..ffd19abd7 100644 --- a/crates/spark-storage/src/ngram_cache_fault.rs +++ b/crates/spark-storage/src/ngram_cache_fault.rs @@ -65,7 +65,30 @@ pub(super) struct RowSource<'a> { pub(super) base_offset: u64, pub(super) segments: Option<&'a Segments>, pub(super) row_stride: usize, - pub(super) scale_file: Option<&'a File>, + pub(super) scale_src: Option>, + pub(super) scale_stride: usize, +} + +/// Where a scale row lives. `Flat`: one file, `id * scale_stride` in. +/// `Segmented`: the table's own shard segmentation applied to the sibling +/// scale tensors (same `rows_per`, same deduped files). +pub(super) enum ScaleSource<'a> { + Flat { file: &'a File }, + Segmented { seg: &'a Segments }, +} + +impl ScaleSource<'_> { + fn scale_loc(&self, id: u64, stride: u64) -> (&File, u64) { + match self { + ScaleSource::Flat { file } => (file, id * stride), + ScaleSource::Segmented { seg } => { + let shard = (id / seg.rows_per) as usize; + let local = id % seg.rows_per; + let file = &seg.files[seg.shard_file[shard] as usize]; + (file, seg.bases[shard] + local * stride) + } + } + } } impl RowSource<'_> { @@ -128,16 +151,21 @@ pub(super) fn fetch_row( dst.copy_from_slice(&bounce.blocks(nblocks)[within..within + src.row_stride]); // A constant per-tensor scale needs no per-row refresh: every slot already - // holds it (see `set_constant_scale`), and `scale_file` is None then. - if let (Some(sfile), Some(sbase)) = (src.scale_file, ptrs.scales) { - let sbyte = id * 4; + // holds it (see `set_constant_scale`), and `scale_src` is None then. + // Same O_DIRECT bounce/seam handling as the row itself; the stride comes + // from the installed source (4-byte f32 or `dim/group` E4M3). + if let (Some(ssrc), Some(sbase)) = (&src.scale_src, ptrs.scales) { + let stride = src.scale_stride; + let (sfile, sbyte) = ssrc.scale_loc(id, stride as u64); let sblock = sbyte - (sbyte % BLOCK as u64); let swithin = (sbyte - sblock) as usize; - atlas_tier::pio::read_at_least_at(sfile, bounce.blocks(1), sblock, swithin + 4) + let nblocks = if swithin + stride > BLOCK { 2 } else { 1 }; + atlas_tier::pio::read_at_least_at(sfile, bounce.blocks(nblocks), sblock, swithin + stride) .with_context(|| format!("NgramRowCache: read scale {id}"))?; - // SAFETY: as above; the scale arena holds slots*4 bytes. - let sdst = unsafe { std::slice::from_raw_parts_mut(sbase.add(slot as usize * 4), 4) }; - sdst.copy_from_slice(&bounce.blocks(1)[swithin..swithin + 4]); + // SAFETY: as above; the scale arena holds slots*stride bytes. + let sdst = + unsafe { std::slice::from_raw_parts_mut(sbase.add(slot as usize * stride), stride) }; + sdst.copy_from_slice(&bounce.blocks(nblocks)[swithin..swithin + stride]); } Ok(()) } diff --git a/kernels/gb10/common/embed_from_argmax.cu b/kernels/gb10/common/embed_from_argmax.cu index 1b040efc9..79d9f4745 100644 --- a/kernels/gb10/common/embed_from_argmax.cu +++ b/kernels/gb10/common/embed_from_argmax.cu @@ -123,3 +123,40 @@ extern "C" __global__ void batched_embed_fp8( dst[i] = __float2bfloat16(ngram_fp8_decode(src[i]) * scale); } } + +// NVFP4-table batched embedding gather (FP4-Local's PLE n-gram table): +// rows are packed E2M1 pairs (low nibble first) with one E4M3 block scale +// per `group` elements and a per-tensor f32 global. BIT-IDENTICAL to the +// host decoder `atlas_core::modelopt_codec::nvfp4_rows_to_bf16`: +// bf16_rne((e2m1 * e4m3) * global) in f32, in exactly that multiply order. +// +// Grid: (num_tokens, 1, 1) Block: (256, 1, 1) +__device__ __forceinline__ float ngram_e2m1_decode(unsigned int nib) { + // FP4_E2M1 LUT, same order as modelopt_codec. + const float lut[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f + }; + return lut[nib & 15u]; +} +extern "C" __global__ void batched_embed_nvfp4( + const unsigned int* __restrict__ token_ids, // [num_tokens] slot ids + const unsigned char* __restrict__ packed, // [slots, dim/2] E2M1 pairs + const unsigned char* __restrict__ scales, // [slots, dim/group] E4M3 + float global, + __nv_bfloat16* __restrict__ output, // [num_tokens, dim] + unsigned int dim, + unsigned int group +) { + const unsigned int token_idx = blockIdx.x; + const unsigned int slot = token_ids[token_idx]; + const unsigned char* src = packed + (unsigned long long)slot * (dim / 2); + const unsigned char* srow = scales + (unsigned long long)slot * (dim / group); + __nv_bfloat16* dst = output + (unsigned long long)token_idx * dim; + for (unsigned int i = threadIdx.x; i < dim / 2; i += blockDim.x) { + const unsigned char byte = src[i]; + const float block = ngram_fp8_decode(srow[(i * 2) / group]); + dst[i * 2] = __float2bfloat16_rn(ngram_e2m1_decode(byte & 15u) * block * global); + dst[i * 2 + 1] = __float2bfloat16_rn(ngram_e2m1_decode(byte >> 4) * block * global); + } +} From c53aba55f0f01bcdbcb941379ffed2561f4581b8 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 05:53:04 -0400 Subject: [PATCH 11/26] qwen4_exp: run the vision tower on its native quantized weights FP4-Local ships the ViT's attn.qkv/attn.proj as MXFP8 and the MLP/merger linears as NVFP4; materialize used to host-expand all of them to BF16. They now upload as shipped (NVFP4 tensors stay deferred, uploaded verbatim by the vision loader; visual MXFP8 joins the native list) and VitWeight expands them into an encoder-owned BF16 scratch per GEMM -- mxfp8_proj_dequant and dequant_nvfp4_to_bf16, both bit-identical to the host codec. Every other checkpoint keeps raw BF16 weights on the same call path. Tracked weight memory drops ~0.5 GB. --- crates/spark-model/Cargo.toml | 4 + .../spark-model/examples/vit_dq_microtest.rs | 164 ++++++++++++++++++ crates/spark-model/src/layers/mod.rs | 2 +- .../spark-model/src/layers/vision_encoder.rs | 46 ++++- .../layers/vision_encoder/enc_impl/init.rs | 37 +++- .../layers/vision_encoder/enc_impl/merger.rs | 4 +- .../vision_encoder/enc_impl/vit_block.rs | 70 +++++++- .../spark-model/src/weight_loader/qwen35.rs | 109 ++++++++++-- .../spark-model/src/weight_loader/qwen3_vl.rs | 16 +- .../src/weight_map/modelopt_local.rs | 17 +- .../src/weight_map/modelopt_local_tests.rs | 40 ++--- 11 files changed, 437 insertions(+), 72 deletions(-) create mode 100644 crates/spark-model/examples/vit_dq_microtest.rs diff --git a/crates/spark-model/Cargo.toml b/crates/spark-model/Cargo.toml index 2290d1692..06810ca9a 100644 --- a/crates/spark-model/Cargo.toml +++ b/crates/spark-model/Cargo.toml @@ -215,6 +215,10 @@ required-features = ["cuda", "gpu-examples"] name = "ngram_nvfp4_microtest" required-features = ["cuda", "gpu-examples"] +[[example]] +name = "vit_dq_microtest" +required-features = ["cuda", "gpu-examples"] + [[example]] name = "q2_0_gemv_microtest" required-features = ["cuda", "gpu-examples"] diff --git a/crates/spark-model/examples/vit_dq_microtest.rs b/crates/spark-model/examples/vit_dq_microtest.rs new file mode 100644 index 000000000..60b705c6c --- /dev/null +++ b/crates/spark-model/examples/vit_dq_microtest.rs @@ -0,0 +1,164 @@ +// SPDX-License-Identifier: AGPL-3.0-only +//! Vision-tower native weights: `mxfp8_proj_dequant` and +//! `dequant_nvfp4_to_bf16` must reproduce the host expansion bit-for-bit. +//! Run: +//! cargo run -p spark-model --release --features cuda,gpu-examples \ +//! --example vit_dq_microtest /path/to/FP4-Local/snapshot + +use anyhow::Result; +use atlas_core::modelopt_codec::{mxfp8_to_bf16, nvfp4_rows_to_bf16}; +use spark_runtime::cuda_backend::AtlasCudaBackend; +use spark_runtime::gpu::{DevicePtr, GpuBackend}; +use spark_runtime::kernel_args::KernelLaunch; +use std::collections::HashMap; + +struct Store { + headers: HashMap, // name -> (file, data_off) + shapes: HashMap>, +} + +impl Store { + fn open(dir: &str) -> Result { + let idx: serde_json::Value = serde_json::from_slice(&std::fs::read( + format!("{dir}/model.safetensors.index.json"), + )?)?; + let mut headers = HashMap::new(); + let mut shapes = HashMap::new(); + for e in std::fs::read_dir(format!("{dir}/tensors"))? { + let p = e?.path(); + if p.extension().map(|x| x == "safetensors").unwrap_or(false) { + let b = std::fs::read(&p)?; + let nlen = u64::from_le_bytes(b[..8].try_into()?) as usize; + let hdr: serde_json::Value = serde_json::from_slice(&b[8..8 + nlen])?; + for (name, t) in hdr.as_object().unwrap() { + if name == "__metadata__" { + continue; + } + let off: Vec = serde_json::from_value(t["data_offsets"].clone())?; + let shape: Vec = serde_json::from_value(t["shape"].clone())?; + headers.insert(name.clone(), (p.clone(), 8 + nlen as u64 + off[0])); + shapes.insert(name.clone(), shape); + } + } + } + let _ = idx; + Ok(Self { headers, shapes }) + } + fn shape(&self, name: &str) -> Vec { + self.shapes[name].clone() + } + fn bytes(&self, name: &str, len: usize) -> Result> { + let (p, off) = &self.headers[name]; + let mut f = std::fs::File::open(p)?; + use std::io::{Read, Seek}; + f.seek(std::io::SeekFrom::Start(*off))?; + let mut b = vec![0u8; len]; + f.read_exact(&mut b)?; + Ok(b) + } +} + +fn up(g: &dyn GpuBackend, d: &[u8]) -> Result { + let p = g.alloc(d.len().max(1))?; + g.copy_h2d(d, p)?; + Ok(p) +} +fn down(g: &dyn GpuBackend, p: DevicePtr, n: usize) -> Result> { + let mut b = vec![0u8; n]; + g.copy_d2h(p, &mut b)?; + Ok(b) +} + +fn main() -> Result<()> { + let snap = std::env::args().nth(1).expect("snapshot dir"); + let st = Store::open(&snap)?; + let gpu = AtlasCudaBackend::new(0, &atlas_kernels::ptx_modules())?; + let k8 = gpu.kernel("mxfp8_proj", "mxfp8_proj_dequant")?; + let k4 = gpu.kernel("dequant_nvfp4_bf16", "dequant_nvfp4_to_bf16")?; + + let mut checked = 0usize; + for i in 0..27usize { + for proj in ["attn.qkv", "attn.proj"] { + let base = format!("model.visual.blocks.{i}.{proj}"); + let shp = st.shape(&format!("{base}.weight")); + let (n, k) = (shp[0] as u32, shp[1] as u32); + let wb = st.bytes(&format!("{base}.weight"), (n * k) as usize)?; + let sb = st.bytes(&format!("{base}.weight_scale"), (n * k / 32) as usize)?; + let out = gpu.alloc((n * k * 2) as usize)?; + KernelLaunch::new(&gpu, k8) + .grid([((n as u64 * (k as u64 / 16)).div_ceil(256)) as u32, 1, 1]) + .block([256, 1, 1]) + .arg_ptr(up(&gpu, &wb)?) + .arg_ptr(up(&gpu, &sb)?) + .arg_ptr(out) + .arg_u32(n) + .arg_u32(k) + .launch(0)?; + gpu.synchronize(0)?; + let got = down(&gpu, out, (n * k * 2) as usize)?; + let want = mxfp8_to_bf16(&wb, &sb, n as usize, k as usize)?; + anyhow::ensure!(got == want, "MXFP8 {base} differs"); + checked += 1; + } + } + println!("{checked} MXFP8 projections (all 27 blocks): IDENTICAL"); + + let mut nv = 0; + for i in 0..27usize { + for proj in ["mlp.linear_fc1", "mlp.linear_fc2"] { + let base = format!("model.visual.blocks.{i}.{proj}"); + let shp = st.shape(&format!("{base}.weight")); + let (n, k) = (shp[0] as u32, (shp[1] * 2) as u32); + let wb = st.bytes(&format!("{base}.weight"), (n * k / 2) as usize)?; + let sb = st.bytes(&format!("{base}.weight_scale"), (n * k / 16) as usize)?; + let g = st.bytes(&format!("{base}.weight_scale_2"), 4)?; + let global = f32::from_le_bytes(g[..4].try_into()?); + let out = gpu.alloc((n * k * 2) as usize)?; + KernelLaunch::new(&gpu, k4) + .grid([n, 1, 1]) + .block([256, 1, 1]) + .arg_ptr(up(&gpu, &wb)?) + .arg_ptr(up(&gpu, &sb)?) + .arg_ptr(out) + .arg_f32(global) + .arg_u32(n) + .arg_u32(k) + .launch(0)?; + gpu.synchronize(0)?; + let got = down(&gpu, out, (n * k * 2) as usize)?; + let mut want = vec![0u8; (n * k * 2) as usize]; + nvfp4_rows_to_bf16(&wb, &sb, global, k as usize, 16, &mut want)?; + anyhow::ensure!(got == want, "NVFP4 {base} differs"); + nv += 1; + } + } + for proj in ["linear_fc1", "linear_fc2"] { + let base = format!("model.visual.merger.{proj}"); + let shp = st.shape(&format!("{base}.weight")); + let (n, k) = (shp[0] as u32, (shp[1] * 2) as u32); + let wb = st.bytes(&format!("{base}.weight"), (n * k / 2) as usize)?; + let sb = st.bytes(&format!("{base}.weight_scale"), (n * k / 16) as usize)?; + let g = st.bytes(&format!("{base}.weight_scale_2"), 4)?; + let global = f32::from_le_bytes(g[..4].try_into()?); + let out = gpu.alloc((n * k * 2) as usize)?; + KernelLaunch::new(&gpu, k4) + .grid([n, 1, 1]) + .block([256, 1, 1]) + .arg_ptr(up(&gpu, &wb)?) + .arg_ptr(up(&gpu, &sb)?) + .arg_ptr(out) + .arg_f32(global) + .arg_u32(n) + .arg_u32(k) + .launch(0)?; + gpu.synchronize(0)?; + let got = down(&gpu, out, (n * k * 2) as usize)?; + let mut want = vec![0u8; (n * k * 2) as usize]; + nvfp4_rows_to_bf16(&wb, &sb, global, k as usize, 16, &mut want)?; + anyhow::ensure!(got == want, "NVFP4 {base} differs"); + nv += 1; + } + println!("{nv} NVFP4 projections (all 27 blocks + merger): IDENTICAL"); + println!("PASS"); + Ok(()) +} diff --git a/crates/spark-model/src/layers/mod.rs b/crates/spark-model/src/layers/mod.rs index cca874637..33164a14d 100644 --- a/crates/spark-model/src/layers/mod.rs +++ b/crates/spark-model/src/layers/mod.rs @@ -56,7 +56,7 @@ pub use qwen3_attention::Qwen3AttentionLayer; pub use qwen3_ssm::Qwen3SsmLayer; pub use qwen4exp_mtp::{Qwen4ExpMtpHead, Qwen4ExpMtpProposerState}; pub(crate) use vision_encoder::{GlmVisionBlockWeights, GlmVisionMergerWeights, GlmVisionWeights}; -pub use vision_encoder::{MergerLayer, ViTBlock, VisionEncoder}; +pub use vision_encoder::{MergerLayer, ViTBlock, VitWeight, VisionEncoder}; pub mod vision_tower; pub use vision_tower::VisionTower; diff --git a/crates/spark-model/src/layers/vision_encoder.rs b/crates/spark-model/src/layers/vision_encoder.rs index f88a2c401..88a69b7db 100644 --- a/crates/spark-model/src/layers/vision_encoder.rs +++ b/crates/spark-model/src/layers/vision_encoder.rs @@ -35,27 +35,56 @@ pub const VIDEO_PAD_TOKEN_ID: u32 = IMAGE_PAD_TOKEN + 1; /// therefore check the length rather than assume it — see `patch_embed`. pub(crate) const PATCH_DIM: usize = 1536; +/// A ViT projection weight, in the checkpoint's native precision. BF16 is a +/// raw pointer; the quantized arms carry everything the scratch dequant needs. +#[derive(Debug, Clone, Copy)] +pub enum VitWeight { + Bf16(DevicePtr), + /// E4M3 + E8M0 per-32-input scales (FP4-Local `attn.qkv`/`attn.proj`). + Mxfp8(crate::layers::ops::Mxfp8Proj), + /// Packed E2M1 + E4M3 per-16 scales + one F32 global (FP4-Local + /// `mlp.linear_fc*`/`merger.linear_fc*`). + Nvfp4 { + packed: DevicePtr, + scales: DevicePtr, + global: f32, + n: u32, + k: u32, + }, +} + +impl VitWeight { + /// BF16 bytes this weight expands to (`n * k * 2`). + pub fn bf16_bytes(&self) -> usize { + match self { + VitWeight::Bf16(_) => 0, + VitWeight::Mxfp8(p) => p.n as usize * p.k as usize * 2, + VitWeight::Nvfp4 { n, k, .. } => *n as usize * *k as usize * 2, + } + } +} + pub struct ViTBlock { pub norm1_w: DevicePtr, pub norm1_b: DevicePtr, - pub qkv_w: DevicePtr, + pub qkv_w: VitWeight, pub qkv_b: DevicePtr, - pub proj_w: DevicePtr, + pub proj_w: VitWeight, pub proj_b: DevicePtr, pub norm2_w: DevicePtr, pub norm2_b: DevicePtr, - pub fc1_w: DevicePtr, + pub fc1_w: VitWeight, pub fc1_b: DevicePtr, - pub fc2_w: DevicePtr, + pub fc2_w: VitWeight, pub fc2_b: DevicePtr, } pub struct MergerLayer { pub norm_w: DevicePtr, pub norm_b: DevicePtr, - pub fc1_w: DevicePtr, + pub fc1_w: VitWeight, pub fc1_b: DevicePtr, - pub fc2_w: DevicePtr, + pub fc2_w: VitWeight, pub fc2_b: DevicePtr, } @@ -83,6 +112,8 @@ pub struct VisionEncoder { k_merge: KernelHandle, // vision_spatial_merge (2×2) k_f32_bf16: KernelHandle, // vision_f32_to_bf16 k_copy: KernelHandle, // vision_bf16_copy + k_nvfp4_dq: KernelHandle, // dequant_nvfp4_to_bf16 (native vision NVFP4) + k_mx8_dq: KernelHandle, // mxfp8_proj_dequant (native vision MXFP8) // config pub hidden_size: usize, // 1152 pub num_heads: usize, // 16 @@ -111,6 +142,9 @@ pub struct VisionEncoder { pub buf_scores: DevicePtr, // [attn_max × attn_max] f32 — per-head raw QKᵀ pub buf_probs: DevicePtr, // [attn_max × attn_max] bf16 — per-head softmax probs pub buf_o_stage: DevicePtr, // [p_max × head_dim] bf16 — per-head GEMM2 output staging + /// Vision-owned BF16 expansion scratch for quantized weights — sized for + /// the largest projection (~10 MB), NOT shared with the LM buffers. + pub buf_vit_dq: DevicePtr, // host-side prep state pos_embed_host_f32: Vec, // [num_position_embeddings × hidden_size] row-major rope_inv_freq: Vec, // [head_dim / 4] frequencies diff --git a/crates/spark-model/src/layers/vision_encoder/enc_impl/init.rs b/crates/spark-model/src/layers/vision_encoder/enc_impl/init.rs index bddfb5eee..0b4ad66c9 100644 --- a/crates/spark-model/src/layers/vision_encoder/enc_impl/init.rs +++ b/crates/spark-model/src/layers/vision_encoder/enc_impl/init.rs @@ -147,6 +147,26 @@ impl VisionEncoder { let buf_scores = gpu.alloc(attn_max * attn_max * 4)?; // [seq, seq] f32 let buf_probs = gpu.alloc(attn_max * attn_max * 2)?; // [seq, seq] bf16 let buf_o_stage = gpu.alloc(p_max * head_dim * 2)?; // [seq, D] bf16 + // Weight-dequant scratch: the largest quantized projection's BF16 + // expansion (~10 MB for fc1 [4304,1152]). Owned by the encoder; never + // aliased with LM scratch. + let max_w_bytes = blocks + .iter() + .flat_map(|b| [b.qkv_w, b.proj_w, b.fc1_w, b.fc2_w]) + .chain( + deepstack + .iter() + .flat_map(|m| [m.fc1_w, m.fc2_w]) + .chain([merger.fc1_w, merger.fc2_w]), + ) + .map(|w| w.bf16_bytes()) + .max() + .unwrap_or(0); + let buf_vit_dq = if max_w_bytes > 0 { + gpu.alloc(max_w_bytes)? + } else { + DevicePtr::NULL + }; // Download pos_embed weight to host as f32 so we can bilinear- // interpolate it per image (HF: `fast_pos_embed_interpolate`). @@ -207,6 +227,15 @@ impl VisionEncoder { k_merge: gpu.kernel("vision_encoder", "vision_spatial_merge")?, k_f32_bf16: gpu.kernel("vision_encoder", "vision_f32_to_bf16")?, k_copy: gpu.kernel("vision_encoder", "vision_bf16_copy")?, + // Dequants for checkpoint-native quantized ViT projections. + // SOFT: absent on vision trees without them (then every weight + // is Bf16 and they are never launched). + k_nvfp4_dq: crate::layers::try_kernel( + gpu, + "dequant_nvfp4_bf16", + "dequant_nvfp4_to_bf16", + ), + k_mx8_dq: crate::layers::try_kernel(gpu, "mxfp8_proj", "mxfp8_proj_dequant"), hidden_size, num_heads, head_dim, @@ -231,6 +260,7 @@ impl VisionEncoder { buf_scores, buf_probs, buf_o_stage, + buf_vit_dq, pos_embed_host_f32, rope_inv_freq, }) @@ -257,9 +287,9 @@ impl VisionEncoder { merger: MergerLayer { norm_w: DevicePtr::NULL, norm_b: DevicePtr::NULL, - fc1_w: DevicePtr::NULL, + fc1_w: super::super::VitWeight::Bf16(DevicePtr::NULL), fc1_b: DevicePtr::NULL, - fc2_w: DevicePtr::NULL, + fc2_w: super::super::VitWeight::Bf16(DevicePtr::NULL), fc2_b: DevicePtr::NULL, }, k_gemm: KernelHandle(0), @@ -276,6 +306,8 @@ impl VisionEncoder { k_merge: KernelHandle(0), k_f32_bf16: KernelHandle(0), k_copy: KernelHandle(0), + k_nvfp4_dq: KernelHandle(0), + k_mx8_dq: KernelHandle(0), hidden_size: config.hidden_size, num_heads: config.num_heads, head_dim: config.hidden_size / config.num_heads, @@ -288,6 +320,7 @@ impl VisionEncoder { buf_h1: DevicePtr::NULL, buf_h2: DevicePtr::NULL, buf_wide: DevicePtr::NULL, + buf_vit_dq: DevicePtr::NULL, buf_merge_in: DevicePtr::NULL, buf_merge_fc1: DevicePtr::NULL, buf_out: glm_buf_out, diff --git a/crates/spark-model/src/layers/vision_encoder/enc_impl/merger.rs b/crates/spark-model/src/layers/vision_encoder/enc_impl/merger.rs index e4a9e1aaa..2f2966160 100644 --- a/crates/spark-model/src/layers/vision_encoder/enc_impl/merger.rs +++ b/crates/spark-model/src/layers/vision_encoder/enc_impl/merger.rs @@ -68,7 +68,7 @@ impl VisionEncoder { self.vit_gemm_bias( gpu, self.buf_merge_in, - m.fc1_w, + self.vit_weight_ptr(gpu, &m.fc1_w, stream)?, m.fc1_b, self.buf_merge_fc1, merged_p, @@ -87,7 +87,7 @@ impl VisionEncoder { self.vit_gemm_bias( gpu, self.buf_merge_fc1, - m.fc2_w, + self.vit_weight_ptr(gpu, &m.fc2_w, stream)?, m.fc2_b, out_slice, merged_p, diff --git a/crates/spark-model/src/layers/vision_encoder/enc_impl/vit_block.rs b/crates/spark-model/src/layers/vision_encoder/enc_impl/vit_block.rs index 11836d6ff..6474e217c 100644 --- a/crates/spark-model/src/layers/vision_encoder/enc_impl/vit_block.rs +++ b/crates/spark-model/src/layers/vision_encoder/enc_impl/vit_block.rs @@ -10,6 +10,60 @@ use spark_runtime::kernel_args::{KernelLaunch, div_ceil}; use super::super::{ViTBlock, VisionEncoder}; impl VisionEncoder { + /// Resolve a [`VitWeight`] to BF16 pointer: raw BF16 is returned as-is; + /// a quantized weight is expanded into `buf_vit_dq` (vision-owned, sized + /// for the largest projection at load). The dequant kernels reproduce + /// the host expansion's bytes exactly. + pub(super) fn vit_weight_ptr( + &self, + gpu: &dyn GpuBackend, + w: &super::super::VitWeight, + stream: u64, + ) -> Result { + match *w { + super::super::VitWeight::Bf16(p) => Ok(p), + super::super::VitWeight::Mxfp8(p) => { + let chunks = u64::from(p.n) * u64::from(p.k / 16); + anyhow::ensure!( + chunks.div_ceil(256) <= u64::from(u32::MAX), + "vit weight {n}x{k} too large for one dequant grid", + n = p.n, + k = p.k + ); + KernelLaunch::new(gpu, self.k_mx8_dq) + .grid([chunks.div_ceil(256) as u32, 1, 1]) + .block([256, 1, 1]) + .arg_ptr(p.w.data) + .arg_ptr(p.w.scales) + .arg_ptr(self.buf_vit_dq) + .arg_u32(p.n) + .arg_u32(p.k) + .launch(stream)?; + Ok(self.buf_vit_dq) + } + super::super::VitWeight::Nvfp4 { + packed, + scales, + global, + n, + k, + } => { + crate::layers::ops::dequant_nvfp4_to_bf16( + gpu, + self.k_nvfp4_dq, + packed, + scales, + self.buf_vit_dq, + global, + n, + k, + stream, + )?; + Ok(self.buf_vit_dq) + } + } + } + /// ViT GEMM with bias: C[m,n] = A[m,k] @ B[n,k]^T + bias[n] (BF16). /// Prefers the tensor-core `dense_gemm_bf16_pipelined` (~40× the scalar /// `vision_gemm_bias` on the ViT's large-M shapes) + a fused bias-add; falls @@ -204,7 +258,7 @@ impl VisionEncoder { self.vit_gemm_bias( gpu, self.buf_h1, - blk.qkv_w, + self.vit_weight_ptr(gpu, &blk.qkv_w, stream)?, blk.qkv_b, self.buf_wide, p32, @@ -245,7 +299,7 @@ impl VisionEncoder { self.vit_gemm_bias( gpu, self.buf_h1, - blk.proj_w, + self.vit_weight_ptr(gpu, &blk.proj_w, stream)?, blk.proj_b, self.buf_wide, p32, @@ -294,7 +348,7 @@ impl VisionEncoder { self.vit_gemm_bias( gpu, self.buf_h1, - blk.fc1_w, + self.vit_weight_ptr(gpu, &blk.fc1_w, stream)?, blk.fc1_b, self.buf_wide, p32, @@ -313,7 +367,7 @@ impl VisionEncoder { self.vit_gemm_bias( gpu, self.buf_wide, - blk.fc2_w, + self.vit_weight_ptr(gpu, &blk.fc2_w, stream)?, blk.fc2_b, self.buf_h1, p32, @@ -376,7 +430,7 @@ impl VisionEncoder { self.vit_gemm_bias( gpu, self.buf_h1, - blk.qkv_w, + self.vit_weight_ptr(gpu, &blk.qkv_w, stream)?, blk.qkv_b, self.buf_wide, pt, @@ -425,7 +479,7 @@ impl VisionEncoder { self.vit_gemm_bias( gpu, self.buf_h1, - blk.proj_w, + self.vit_weight_ptr(gpu, &blk.proj_w, stream)?, blk.proj_b, self.buf_wide, pt, @@ -474,7 +528,7 @@ impl VisionEncoder { self.vit_gemm_bias( gpu, self.buf_h1, - blk.fc1_w, + self.vit_weight_ptr(gpu, &blk.fc1_w, stream)?, blk.fc1_b, self.buf_wide, pt, @@ -493,7 +547,7 @@ impl VisionEncoder { self.vit_gemm_bias( gpu, self.buf_wide, - blk.fc2_w, + self.vit_weight_ptr(gpu, &blk.fc2_w, stream)?, blk.fc2_b, self.buf_h1, pt, diff --git a/crates/spark-model/src/weight_loader/qwen35.rs b/crates/spark-model/src/weight_loader/qwen35.rs index 5cb515ec2..f9bc6e3e2 100644 --- a/crates/spark-model/src/weight_loader/qwen35.rs +++ b/crates/spark-model/src/weight_loader/qwen35.rs @@ -14,19 +14,96 @@ use crate::weight_map::{DenseWeight, MtpWeights, dense_auto, detect_nvfp4_varian pub struct Qwen35WeightLoader; -fn vision_dense_auto( +/// A ViT projection in its native precision: BF16 stays a raw pointer; +/// FP4-Local's MXFP8 (`attn.qkv`/`attn.proj`) and NVFP4 (`linear_fc{1,2}`) +/// tensors upload as shipped and expand on-GPU per GEMM — bit-identical to +/// the host expansion other checkpoints get. +/// A quantized tensor's resident pointer: resident if materialize uploaded +/// it natively, else the deferred shard bytes uploaded verbatim (the vision +/// NVFP4 tensors stay deferred — nothing else consumes them). +#[cfg(feature = "cuda")] +fn vision_tensor_ptr( + store: &WeightStore, + name: &str, + gpu: &dyn GpuBackend, +) -> Result<(spark_runtime::gpu::DevicePtr, spark_runtime::weights::WeightDtype, Vec)> { + if let Ok(t) = store.get(name) { + return Ok((t.ptr, t.dtype, t.shape.clone())); + } + let d = store + .deferred(name) + .with_context(|| format!("vision: `{name}` neither resident nor deferred"))?; + let len = d.shape.iter().product::(); + let bytes = len + * match d.dtype { + WeightDtype::BF16 | WeightDtype::FP16 => 2, + WeightDtype::FP32 => 4, + _ => 1, + }; + use std::io::Read as _; + let mut buf = vec![0u8; bytes]; + let mut f = std::fs::File::open(&d.path)?; + use std::io::Seek as _; + f.seek(std::io::SeekFrom::Start(d.offset))?; + f.read_exact(&mut buf)?; + let p = gpu.alloc(bytes)?; + gpu.copy_h2d(&buf, p)?; + Ok((p, d.dtype, d.shape.clone())) +} + +#[cfg(feature = "cuda")] +fn vision_weight( store: &WeightStore, prefix: &str, gpu: &dyn GpuBackend, -) -> Result { +) -> Result { let name = format!("{prefix}.weight"); - let w = store.get(&name)?; - match w.dtype { - WeightDtype::BF16 | WeightDtype::FP8E4M3 => { - crate::weight_map::dense_auto_fp8_or_bf16(store, prefix, gpu) + let (wptr, wdtype, wshape) = vision_tensor_ptr(store, &name, gpu)?; + match wdtype { + WeightDtype::BF16 => Ok(crate::layers::VitWeight::Bf16(wptr)), + WeightDtype::FP8E4M3 => { + // MXFP8: E4M3 [n,k] + U8 E8M0 [n,k/32]. + let (sptr, sdtype, sshape) = + vision_tensor_ptr(store, &format!("{prefix}.weight_scale"), gpu)?; + anyhow::ensure!( + wshape.len() == 2 + && sdtype == WeightDtype::UInt8 + && sshape == [wshape[0], wshape[1] / 32], + "vision MXFP8 {prefix}: weight {wshape:?} + scale {sshape:?}" + ); + Ok(crate::layers::VitWeight::Mxfp8(crate::weight_map::Mxfp8Proj { + w: crate::weight_map::Mxfp8Weight { + data: wptr, + scales: sptr, + }, + n: wshape[0] as u32, + k: wshape[1] as u32, + })) + } + WeightDtype::UInt8 => { + // NVFP4: packed U8 [n,k/2] + E4M3 [n,k/16] + F32 weight_scale_2. + let (sptr, sdtype, sshape) = + vision_tensor_ptr(store, &format!("{prefix}.weight_scale"), gpu)?; + anyhow::ensure!( + wshape.len() == 2 + && sdtype == WeightDtype::FP8E4M3 + && sshape[0] == wshape[0], + "vision NVFP4 {prefix}: weight {wshape:?} + scale {sshape:?}" + ); + let gname = format!("{prefix}.weight_scale_2"); + let gt = store.get(&gname)?; + let mut gb = [0u8; 4]; + gpu.copy_d2h(gt.ptr, &mut gb)?; + store.reclaim(gpu, &gname)?; + Ok(crate::layers::VitWeight::Nvfp4 { + packed: wptr, + scales: sptr, + global: f32::from_le_bytes(gb), + n: wshape[0] as u32, + k: (wshape[1] * 2) as u32, + }) } - WeightDtype::FP32 => crate::weight_map::dense_f32_safe(store, &name, gpu), - other => anyhow::bail!("vision_dense_auto: unsupported dtype {other:?} for {name}"), + other => anyhow::bail!("vision_weight: unsupported dtype {other:?} for {name}"), } } @@ -202,18 +279,18 @@ impl ModelWeightLoader for Qwen35WeightLoader { norm1_w: vision_tensor_dense_auto(store, &format!("{bp}.norm1.weight"), gpu)? .weight, norm1_b: vision_tensor_dense_auto(store, &format!("{bp}.norm1.bias"), gpu)?.weight, - qkv_w: vision_dense_auto(store, &format!("{bp}.attn.qkv"), gpu)?.weight, + qkv_w: vision_weight(store, &format!("{bp}.attn.qkv"), gpu)?, qkv_b: vision_tensor_dense_auto(store, &format!("{bp}.attn.qkv.bias"), gpu)?.weight, - proj_w: vision_dense_auto(store, &format!("{bp}.attn.proj"), gpu)?.weight, + proj_w: vision_weight(store, &format!("{bp}.attn.proj"), gpu)?, proj_b: vision_tensor_dense_auto(store, &format!("{bp}.attn.proj.bias"), gpu)? .weight, norm2_w: vision_tensor_dense_auto(store, &format!("{bp}.norm2.weight"), gpu)? .weight, norm2_b: vision_tensor_dense_auto(store, &format!("{bp}.norm2.bias"), gpu)?.weight, - fc1_w: vision_dense_auto(store, &format!("{bp}.mlp.linear_fc1"), gpu)?.weight, + fc1_w: vision_weight(store, &format!("{bp}.mlp.linear_fc1"), gpu)?, fc1_b: vision_tensor_dense_auto(store, &format!("{bp}.mlp.linear_fc1.bias"), gpu)? .weight, - fc2_w: vision_dense_auto(store, &format!("{bp}.mlp.linear_fc2"), gpu)?.weight, + fc2_w: vision_weight(store, &format!("{bp}.mlp.linear_fc2"), gpu)?, fc2_b: vision_tensor_dense_auto(store, &format!("{bp}.mlp.linear_fc2.bias"), gpu)? .weight, }); @@ -225,10 +302,10 @@ impl ModelWeightLoader for Qwen35WeightLoader { deepstack.push(crate::layers::MergerLayer { norm_w: vision_tensor_dense_auto(store, &format!("{mp}.norm.weight"), gpu)?.weight, norm_b: vision_tensor_dense_auto(store, &format!("{mp}.norm.bias"), gpu)?.weight, - fc1_w: vision_dense_auto(store, &format!("{mp}.linear_fc1"), gpu)?.weight, + fc1_w: vision_weight(store, &format!("{mp}.linear_fc1"), gpu)?, fc1_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc1.bias"), gpu)? .weight, - fc2_w: vision_dense_auto(store, &format!("{mp}.linear_fc2"), gpu)?.weight, + fc2_w: vision_weight(store, &format!("{mp}.linear_fc2"), gpu)?, fc2_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc2.bias"), gpu)? .weight, }); @@ -238,9 +315,9 @@ impl ModelWeightLoader for Qwen35WeightLoader { let merger = crate::layers::MergerLayer { norm_w: vision_tensor_dense_auto(store, &format!("{mp}.norm.weight"), gpu)?.weight, norm_b: vision_tensor_dense_auto(store, &format!("{mp}.norm.bias"), gpu)?.weight, - fc1_w: vision_dense_auto(store, &format!("{mp}.linear_fc1"), gpu)?.weight, + fc1_w: vision_weight(store, &format!("{mp}.linear_fc1"), gpu)?, fc1_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc1.bias"), gpu)?.weight, - fc2_w: vision_dense_auto(store, &format!("{mp}.linear_fc2"), gpu)?.weight, + fc2_w: vision_weight(store, &format!("{mp}.linear_fc2"), gpu)?, fc2_b: vision_tensor_dense_auto(store, &format!("{mp}.linear_fc2.bias"), gpu)?.weight, }; diff --git a/crates/spark-model/src/weight_loader/qwen3_vl.rs b/crates/spark-model/src/weight_loader/qwen3_vl.rs index a6a3ae6c3..e5d5774ec 100644 --- a/crates/spark-model/src/weight_loader/qwen3_vl.rs +++ b/crates/spark-model/src/weight_loader/qwen3_vl.rs @@ -237,15 +237,15 @@ impl ModelWeightLoader for Qwen3VLWeightLoader { blocks.push(ViTBlock { norm1_w: dense(store, &format!("{bp}.norm1.weight"))?.weight, norm1_b: dense(store, &format!("{bp}.norm1.bias"))?.weight, - qkv_w: dense(store, &format!("{bp}.attn.qkv.weight"))?.weight, + qkv_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{bp}.attn.qkv.weight"))?.weight), qkv_b: dense(store, &format!("{bp}.attn.qkv.bias"))?.weight, - proj_w: dense(store, &format!("{bp}.attn.proj.weight"))?.weight, + proj_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{bp}.attn.proj.weight"))?.weight), proj_b: dense(store, &format!("{bp}.attn.proj.bias"))?.weight, norm2_w: dense(store, &format!("{bp}.norm2.weight"))?.weight, norm2_b: dense(store, &format!("{bp}.norm2.bias"))?.weight, - fc1_w: dense(store, &format!("{bp}.mlp.linear_fc1.weight"))?.weight, + fc1_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{bp}.mlp.linear_fc1.weight"))?.weight), fc1_b: dense(store, &format!("{bp}.mlp.linear_fc1.bias"))?.weight, - fc2_w: dense(store, &format!("{bp}.mlp.linear_fc2.weight"))?.weight, + fc2_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{bp}.mlp.linear_fc2.weight"))?.weight), fc2_b: dense(store, &format!("{bp}.mlp.linear_fc2.bias"))?.weight, }); } @@ -256,9 +256,9 @@ impl ModelWeightLoader for Qwen3VLWeightLoader { deepstack.push(MergerLayer { norm_w: dense(store, &format!("{mp}.norm.weight"))?.weight, norm_b: dense(store, &format!("{mp}.norm.bias"))?.weight, - fc1_w: dense(store, &format!("{mp}.linear_fc1.weight"))?.weight, + fc1_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{mp}.linear_fc1.weight"))?.weight), fc1_b: dense(store, &format!("{mp}.linear_fc1.bias"))?.weight, - fc2_w: dense(store, &format!("{mp}.linear_fc2.weight"))?.weight, + fc2_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{mp}.linear_fc2.weight"))?.weight), fc2_b: dense(store, &format!("{mp}.linear_fc2.bias"))?.weight, }); } @@ -267,9 +267,9 @@ impl ModelWeightLoader for Qwen3VLWeightLoader { let merger = MergerLayer { norm_w: dense(store, &format!("{mp}.norm.weight"))?.weight, norm_b: dense(store, &format!("{mp}.norm.bias"))?.weight, - fc1_w: dense(store, &format!("{mp}.linear_fc1.weight"))?.weight, + fc1_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{mp}.linear_fc1.weight"))?.weight), fc1_b: dense(store, &format!("{mp}.linear_fc1.bias"))?.weight, - fc2_w: dense(store, &format!("{mp}.linear_fc2.weight"))?.weight, + fc2_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{mp}.linear_fc2.weight"))?.weight), fc2_b: dense(store, &format!("{mp}.linear_fc2.bias"))?.weight, }; diff --git a/crates/spark-model/src/weight_map/modelopt_local.rs b/crates/spark-model/src/weight_map/modelopt_local.rs index 388877346..6dbe133f0 100644 --- a/crates/spark-model/src/weight_map/modelopt_local.rs +++ b/crates/spark-model/src/weight_map/modelopt_local.rs @@ -42,19 +42,18 @@ const CSF_FIXED: &str = ".nvfp4_csf_fixed"; const CSF_EXCEPTIONS: &str = ".nvfp4_csf_exceptions"; /// Modules whose weight/scale pair the loader must defer for host expansion: -/// every `MXFP8` entry of the ModelOpt `quantized_layers` map, and NVFP4 -/// entries of the vision tower (which runs BF16). Empty for any other pack. +/// every `MXFP8` entry of the ModelOpt `quantized_layers` map. NVFP4 modules +/// — trunk experts, MTP and, since the vision tower went native, ALL of +/// `model.visual.*` — keep their own dequant paths. Empty for any other pack. pub fn host_decode_modules(config: &ModelConfig) -> HashSet { let Some(qc) = config.quantization_config.as_ref() else { return HashSet::new(); }; qc.quantized_layers .iter() - .filter(|(path, spec)| { + .filter(|(_path, spec)| { let algo = spec.quant_algo.to_ascii_uppercase(); algo == "MXFP8" - || (path.starts_with("model.visual.") - && atlas_core::config::is_nvfp4_quant_algo(&algo)) }) .map(|(path, _)| path.clone()) .collect() @@ -79,8 +78,12 @@ pub fn native_mxfp8(module: &str) -> bool { ".mlp.shared_expert.up_proj", ".mlp.shared_expert.down_proj", ]; - module.starts_with("model.language_model.layers.") - && NATIVE.iter().any(|suffix| module.ends_with(suffix)) + (module.starts_with("model.language_model.layers.") + && NATIVE.iter().any(|suffix| module.ends_with(suffix))) + // The vision tower's MXFP8 attention projections upload as shipped and + // dequant on-GPU into the encoder's BF16 scratch. + || (module.starts_with("model.visual.") + && (module.ends_with(".attn.qkv") || module.ends_with(".attn.proj"))) } /// What [`materialize`] did, for the load log and for tests. diff --git a/crates/spark-model/src/weight_map/modelopt_local_tests.rs b/crates/spark-model/src/weight_map/modelopt_local_tests.rs index 28f9636eb..b0b62d6d8 100644 --- a/crates/spark-model/src/weight_map/modelopt_local_tests.rs +++ b/crates/spark-model/src/weight_map/modelopt_local_tests.rs @@ -50,12 +50,13 @@ fn fp4_local_config() -> ModelConfig { cfg } -/// Only MXFP8 modules and the vision tower's NVFP4 are host-expanded: the -/// trunk experts, MTP experts and PLE table keep their own NVFP4 paths. +/// Only MXFP8 modules are host-expanded: the trunk experts, MTP experts, PLE +/// table and the vision tower (native NVFP4 dequant on-GPU) keep their own +/// paths. #[test] fn host_decode_modules_pick_mxfp8_and_vision_nvfp4_only() { let got = host_decode_modules(&fp4_local_config()); - let want: HashSet = [SHARED, GDN_OUT, VIS] + let want: HashSet = [SHARED, GDN_OUT] .iter() .map(|s| s.to_string()) .collect(); @@ -122,9 +123,6 @@ fn resident( .unwrap(); } -fn bf16_at(bytes: &[u8], i: usize) -> f32 { - f32::from_bits(u32::from(u16::from_le_bytes([bytes[2 * i], bytes[2 * i + 1]])) << 16) -} #[test] fn materialize_expands_every_fp4_local_format() { @@ -229,7 +227,7 @@ fn materialize_expands_every_fp4_local_format() { done.nvfp4_to_bf16, done.f32_norms_native ), - (1, 1, 0, 1, 2) + (1, 1, 0, 0, 2) ); // The shared-expert projection is a native consumer now: bytes as shipped. let se = store.get(&format!("{SHARED}.weight")).unwrap(); @@ -266,10 +264,9 @@ fn materialize_expands_every_fp4_local_format() { .get("model.language_model.layers.1.ple.conv1d.bias") .unwrap(); assert_eq!(bias.dtype, WeightDtype::FP32); - assert!( - store.deferred_sorted().is_empty(), - "every source must be consumed" - ); + // The vision NVFP4 tensors are NOT consumed: they are a native GPU + // dequant path now, read as deferred shards by the vision loader. + assert_eq!(store.deferred_sorted().len(), 2, "vision tensors stay deferred"); let scale = store.get(EXPERT).unwrap(); assert_eq!( @@ -280,19 +277,16 @@ fn materialize_expands_every_fp4_local_format() { want[5] = 0x7E; assert_eq!(gpu.read_alloc(scale.ptr).unwrap(), want); - let vis = store.get(&format!("{VIS}.weight")).unwrap(); + let vis = store.deferred(&format!("{VIS}.weight")).unwrap(); assert_eq!( (vis.dtype, vis.shape.clone()), - (WeightDtype::BF16, vec![1, 16]) + (WeightDtype::UInt8, vec![1, 8]) ); - let v = gpu.read_alloc(vis.ptr).unwrap(); - assert!((0..16).all(|i| bf16_at(&v, i) == 1.0)); - for gone in ["weight_scale", "weight_scale_2", "input_scale"] { - assert!( - !store.contains(&format!("{VIS}.{gone}")), - "{gone} must be consumed" - ); - } + // Nothing consumes the vision NVFP4s at materialize: weight+scale stay + // deferred for the vision loader; the resident scalars are untouched. + assert!(store.deferred(&format!("{VIS}.weight_scale")).is_some()); + assert!(store.get(&format!("{VIS}.weight_scale_2")).is_ok()); + assert!(store.get(&format!("{VIS}.input_scale")).is_ok()); } /// The GDN projections have a native MXFP8 kernel: their bytes are uploaded @@ -307,7 +301,9 @@ fn gdn_mxfp8_stays_native() { assert!(!native_mxfp8( "model.language_model.layers.3.self_attn.indexer.index_qk_proj" )); - assert!(!native_mxfp8("model.visual.blocks.0.attn.qkv")); + assert!(native_mxfp8("model.visual.blocks.0.attn.qkv")); + assert!(native_mxfp8("model.visual.blocks.0.attn.proj")); + assert!(!native_mxfp8("model.visual.blocks.0.mlp.linear_fc1")); assert!(!native_mxfp8( "model.language_model.layers.7.linear_attn.in_proj_a" )); From 6473301c88b4bf2c40aedc61244e9bcf458b7a62 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:41:00 -0400 Subject: [PATCH 12/26] ple: align ngram_nvfp4 stride-check test with the scale-arena contract --- crates/spark-model/src/layers/ple/ngram_nvfp4_tests.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/spark-model/src/layers/ple/ngram_nvfp4_tests.rs b/crates/spark-model/src/layers/ple/ngram_nvfp4_tests.rs index 5e34acb4c..19d95f6f9 100644 --- a/crates/spark-model/src/layers/ple/ngram_nvfp4_tests.rs +++ b/crates/spark-model/src/layers/ple/ngram_nvfp4_tests.rs @@ -36,7 +36,8 @@ fn rejects_mismatched_rows_and_strides() { let s = spec(); assert!(decode_rows_bf16(&s, &[vec![0; 16]], &[]).is_err()); assert!(decode_rows_bf16(&s, &[vec![0; 15]], &[vec![0; 2]]).is_err()); - assert!(check_packed_stride(16, &s, false).is_ok()); - assert!(check_packed_stride(32, &s, false).is_err()); - assert!(check_packed_stride(16, &s, true).is_err()); + // The scale arena must be attached: true is required. + assert!(check_packed_stride(16, &s, true).is_ok()); + assert!(check_packed_stride(32, &s, true).is_err()); + assert!(check_packed_stride(16, &s, false).is_err()); } From 8b83be6e4cb2127f02ea9407dd2e315e1dd401b0 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:47:32 -0400 Subject: [PATCH 13/26] docs: record FP4-Local's native-format serve, its verification and the open speed gap Adds docs/checkpoints/qwen3.8-flash-next-fp4-local.md: what the checkpoint ships, what Atlas now runs natively, how exactness was proven, why the n-gram table stays disk-backed, the serve configuration that was benchmarked, the results with their run ids, a same-box comparison against NVIDIA's vLLM recipe, and the ranked list of what is left. The four scripts under bench/qwen4_exp/fp4local/ are the non-in-tree harnesses those results used. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- bench/qwen4_exp/fp4local/ab_stream.py | 102 +++++++++ bench/qwen4_exp/fp4local/control.py | 30 +++ bench/qwen4_exp/fp4local/needle_long.py | 64 ++++++ bench/qwen4_exp/fp4local/prof_req.py | 48 ++++ .../qwen3.8-flash-next-fp4-local.md | 210 ++++++++++++++++++ 5 files changed, 454 insertions(+) create mode 100644 bench/qwen4_exp/fp4local/ab_stream.py create mode 100644 bench/qwen4_exp/fp4local/control.py create mode 100644 bench/qwen4_exp/fp4local/needle_long.py create mode 100644 bench/qwen4_exp/fp4local/prof_req.py create mode 100644 docs/checkpoints/qwen3.8-flash-next-fp4-local.md diff --git a/bench/qwen4_exp/fp4local/ab_stream.py b/bench/qwen4_exp/fp4local/ab_stream.py new file mode 100644 index 000000000..f9ab97836 --- /dev/null +++ b/bench/qwen4_exp/fp4local/ab_stream.py @@ -0,0 +1,102 @@ +"""Engine-neutral streaming A/B client (Atlas vs vLLM). Part of the FP4-Local measurement record (docs/checkpoints/qwen3.8-flash-next-fp4-local.md). + +usage: ab_stream.py + +Identical requests to any OpenAI-compatible server, timed on the CLIENT clock +from the SSE stream, so both engines are measured the same way: + TTFT = first delta carrying content OR reasoning text + decode = (completion_tokens - 1) / (t_last_delta - t_first_delta) +completion_tokens comes from the server's final usage chunk (stream_options +include_usage), and includes reasoning tokens on both engines. +Workloads (each 3 runs, distinct prompt per run so no prefix-cache reuse): + short — short chat questions, max_tokens 512 + agent — ~14k-token agent first turn: long system prompt of tool schemas + + project docs + a user task, with tools attached, max_tokens 512 +Both send chat_template_kwargs {"reasoning_effort": "low"} (Atlas's recipe +default) and temperature 0. Writes /ab.jsonl and prints a summary. +""" +import glob, json, os, statistics, sys, time, urllib.request + +port, model, outdir = sys.argv[1], sys.argv[2], sys.argv[3] +os.makedirs(outdir, exist_ok=True) +URL = f"http://127.0.0.1:{port}/v1/chat/completions" +corpus = "\n\n".join(open(p, encoding="utf-8", errors="ignore").read() + for p in sorted(glob.glob(os.path.join(os.environ.get("FP4L_CORPUS", os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "docs")), "**", "*.md"), recursive=True))) + +SHORT = ["What is the capital of Australia, and why was it chosen?", + "Explain the difference between TCP and UDP in a short paragraph.", + "Give me three tips for writing clear commit messages."] +TOOLS = [{"type": "function", "function": { + "name": f"tool_{i}", "description": f"Utility {i}: performs operation {i} on the workspace and returns JSON.", + "parameters": {"type": "object", "properties": {"path": {"type": "string", "description": "file path"}, + "query": {"type": "string", "description": "search query"}, + "limit": {"type": "integer", "description": "max results"}}, "required": ["path"]}}} + for i in range(24)] + + +def agent_msgs(run): + # ~14k prompt tokens: ~13k of docs (offset per run so prompts differ) + tool schemas. + start = 50_000 * run + docs = corpus[start:start + 46_000] + sys_p = ("You are a coding agent working in a repository. Use the tools when needed. " + "Project documentation follows.\n\n" + docs) + return [{"role": "system", "content": sys_p}, + {"role": "user", "content": "Summarize the build and serve workflow described above in five bullet " + "points, then say which tool you would call first and why."}] + + +def stream(messages, max_tokens, tools=None): + body = {"model": model, "messages": messages, "max_tokens": max_tokens, "temperature": 0, + "stream": True, "stream_options": {"include_usage": True}, + "chat_template_kwargs": {"reasoning_effort": "low"}} + if tools: + body["tools"] = tools + req = urllib.request.Request(URL, json.dumps(body).encode(), {"Content-Type": "application/json"}) + t0 = time.time(); t_first = t_last = None; usage = None; n_delta = 0; text = [] + with urllib.request.urlopen(req, timeout=1800) as r: + for raw in r: + line = raw.decode("utf-8", "ignore").strip() + if not line.startswith("data:"): + continue + data = line[5:].strip() + if data == "[DONE]": + break + ev = json.loads(data) + if ev.get("usage"): + usage = ev["usage"] + for ch in ev.get("choices") or []: + d = ch.get("delta") or {} + piece = (d.get("content") or "") + (d.get("reasoning_content") or d.get("reasoning") or "") + if piece or d.get("tool_calls"): + now = time.time(); n_delta += 1 + t_first = t_first or now; t_last = now + text.append(piece) + ct = (usage or {}).get("completion_tokens") or 0 + ct_from = "usage" + if not ct: # server sent no usage chunk: fall back to delta count (a lower bound under MTP) + ct, ct_from = n_delta, "deltas" + pt = (usage or {}).get("prompt_tokens") or 0 + dec = (ct - 1) / (t_last - t_first) if (t_first and t_last and t_last > t_first and ct > 1) else None + return {"prompt_tokens": pt, "completion_tokens": ct, "ct_from": ct_from, + "ttft_s": round(t_first - t0, 3) if t_first else None, + "decode_tok_s": round(dec, 2) if dec else None, "wall_s": round(time.time() - t0, 2), + "deltas": n_delta, "text_head": "".join(text)[:120]} + + +rows = [] +with open(f"{outdir}/ab.jsonl", "w") as f: + stream([{"role": "user", "content": "Say hi."}], 16) # warm-up, unmeasured + for i, q in enumerate(SHORT): + r = {"workload": "short", "run": i, **stream([{"role": "user", "content": q}], 512)} + rows.append(r); f.write(json.dumps(r) + "\n"); f.flush(); print(json.dumps(r)) + for i in range(3): + r = {"workload": "agent", "run": i, **stream(agent_msgs(i), 512, TOOLS)} + rows.append(r); f.write(json.dumps(r) + "\n"); f.flush(); print(json.dumps(r)) +for w in ("short", "agent"): + rs = [r for r in rows if r["workload"] == w and r["decode_tok_s"]] + if rs: + print(json.dumps({"summary": w, "n": len(rs), + "decode_tok_s_median": statistics.median(r["decode_tok_s"] for r in rs), + "ttft_s_median": statistics.median(r["ttft_s"] for r in rs), + "prompt_tokens": [r["prompt_tokens"] for r in rs], + "completion_tokens": [r["completion_tokens"] for r in rs]})) diff --git a/bench/qwen4_exp/fp4local/control.py b/bench/qwen4_exp/fp4local/control.py new file mode 100644 index 000000000..018eb31ec --- /dev/null +++ b/bench/qwen4_exp/fp4local/control.py @@ -0,0 +1,30 @@ +"""Fixed 4-request output control. Part of the FP4-Local measurement record (docs/checkpoints/qwen3.8-flash-next-fp4-local.md). +usage: control.py -> /{math,explain,code,tool}.json + summary line +temperature 0, reasoning_effort "none", max_tokens 192, distinct prompts.""" +import json, os, sys, urllib.request + +port, outdir = sys.argv[1], sys.argv[2] +os.makedirs(outdir, exist_ok=True) +TOOLS = [{"type": "function", "function": {"name": "get_weather", "description": "Get the current weather for a city.", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}, "unit": {"type": "string", "enum": ["c", "f"]}}, + "required": ["city"]}}}] +REQS = { + "math": {"messages": [{"role": "user", "content": "What is 37*43? Show the multiplication briefly, then the answer."}]}, + "explain": {"messages": [{"role": "user", "content": "Explain in three sentences why the sky is blue."}]}, + "code": {"messages": [{"role": "user", "content": "Write a Python function that checks whether a string is a palindrome, ignoring case and punctuation."}]}, + "tool": {"messages": [{"role": "user", "content": "What's the weather in Lisbon right now, in celsius?"}], "tools": TOOLS}, +} +ok = 0 +for name, extra in REQS.items(): + body = {"model": "x", "temperature": 0, "max_tokens": 192, "reasoning_effort": "none", **extra} + req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/chat/completions", json.dumps(body).encode(), + {"Content-Type": "application/json"}) + try: + r = json.load(urllib.request.urlopen(req, timeout=600)) + m = r["choices"][0]["message"] + rec = {"content": m.get("content"), "tool_calls": m.get("tool_calls"), "finish": r["choices"][0].get("finish_reason")} + ok += 1 + except Exception as e: + rec = {"error": str(e)[:300]} + json.dump(rec, open(f"{outdir}/{name}.json", "w"), indent=1, sort_keys=True) +print(json.dumps({"outdir": outdir, "ok": ok, "of": len(REQS)})) diff --git a/bench/qwen4_exp/fp4local/needle_long.py b/bench/qwen4_exp/fp4local/needle_long.py new file mode 100644 index 000000000..a83d8fdd3 --- /dev/null +++ b/bench/qwen4_exp/fp4local/needle_long.py @@ -0,0 +1,64 @@ +"""Long-context passcode recall for the FP4-Local server. Part of the FP4-Local measurement record (docs/checkpoints/qwen3.8-flash-next-fp4-local.md). + +usage: needle_long.py [ ...] + +Filler is real English prose/markdown: every docs/**/*.md file of the Atlas +checkout (or $FP4L_CORPUS), concatenated in sorted order and repeated as needed (deterministic, +no synthetic small-vocabulary filler, which inflates n-gram cache hit rates). +For each target length the passcode is placed at depths 10%, 50% and 90%. +The request uses the recipe's thinking default (reasoning_effort "low"), +temperature 0, max_tokens 800. PASS = the exact passcode appears in content. +Prints one JSON line per trial, using the SERVER's usage.prompt_tokens. +""" +import glob, json, os, sys, time, urllib.request + +port = sys.argv[1] +targets = [int(x) for x in sys.argv[2:]] +root = os.environ.get("FP4L_CORPUS", os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "docs")) # the Atlas docs/ tree (override: FP4L_CORPUS) +text = "\n\n".join(open(p, encoding="utf-8", errors="ignore").read() + for p in sorted(glob.glob(f"{root}/**/*.md", recursive=True))) +# ~3.6 chars/token for this markdown on Qwen tokenizers; refined per trial below. +CPT = 3.6 +CODES = ["AMBER-FALCON-731", "COBALT-RIVER-284", "VIOLET-ANCHOR-956"] + + +def build(n_tok, depth, code, cpt): + need = int(n_tok * cpt) + body = (text * (need // len(text) + 1))[:need] + cut = body.rfind("\n", 0, int(len(body) * depth)) + 1 + needle = f"\n\nIMPORTANT NOTE: the secret passcode is {code}. Remember it.\n\n" + return body[:cut] + needle + body[cut:] + + +def ask(prompt): + body = {"model": "x", "temperature": 0, "max_tokens": 800, + "messages": [{"role": "user", "content": prompt + + "\n\nWhat is the secret passcode mentioned in the text above? " + "Reply with only the passcode."}], + "chat_template_kwargs": {"reasoning_effort": "low"}} + req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/chat/completions", + json.dumps(body).encode(), {"Content-Type": "application/json"}) + t = time.time() + r = json.load(urllib.request.urlopen(req, timeout=3600)) + return r, time.time() - t + + +cpt = CPT +for n in targets: + for depth, code in zip([0.1, 0.5, 0.9], CODES): + try: + r, dt = ask(build(n, depth, code, cpt)) + u = r.get("usage", {}) + msg = r["choices"][0]["message"] + content = msg.get("content") or "" + pt = u.get("prompt_tokens", 0) + if pt: + cpt = cpt * n / pt if depth == 0.1 else cpt # calibrate once per length + print(json.dumps({"target": n, "prompt_tokens": pt, "depth": depth, + "pass": code in content, "content": content[:80], + "completion_tokens": u.get("completion_tokens"), + "ttft_ms": u.get("time_to_first_token_ms"), + "wall_s": round(dt, 1)}), flush=True) + except Exception as e: # a failed request is a FAIL, recorded, not skipped + print(json.dumps({"target": n, "depth": depth, "pass": False, + "error": str(e)[:300]}), flush=True) diff --git a/bench/qwen4_exp/fp4local/prof_req.py b/bench/qwen4_exp/fp4local/prof_req.py new file mode 100644 index 000000000..2a44bc9ad --- /dev/null +++ b/bench/qwen4_exp/fp4local/prof_req.py @@ -0,0 +1,48 @@ +"""Decode profiling load for nsys. Part of the FP4-Local measurement record (docs/checkpoints/qwen3.8-flash-next-fp4-local.md). + +usage: prof_req.py + +Sends simultaneous chat requests, each a DIFFERENT code-writing +task (distinct prompts so concurrent rows route to their own experts, as real +users do), temperature 0, reasoning_effort "none" (decode, not thinking). +Prints per-request usage and the aggregate completion tokens, which the +caller divides kernel time by. +""" +import json, sys, threading, time, urllib.request + +port, conc, mt = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]) +TASKS = [ + "Write a Python class implementing a binary min-heap with push, pop and peek, with docstrings and type hints.", + "Write a Rust function that parses an INI file into a HashMap>, with error handling.", + "Write a C function that performs an in-place quicksort on an int array using median-of-three pivot selection.", + "Write a TypeScript module implementing an LRU cache with get, set and a capacity limit, plus unit tests.", + "Write a Go HTTP handler that accepts JSON, validates three required fields and returns a JSON error list.", + "Write a Python async web crawler using asyncio and aiohttp that respects a max depth and a concurrency limit.", + "Write a Java class implementing a thread-safe bounded blocking queue using ReentrantLock and two Conditions.", + "Write a SQL schema and three queries for a library system with books, members and loans, including an index plan.", +] +out = [None] * conc + + +def one(i): + body = {"model": "x", "temperature": 0, "max_tokens": mt, + "messages": [{"role": "user", "content": TASKS[i % len(TASKS)]}], + "reasoning_effort": "none"} + req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/chat/completions", + json.dumps(body).encode(), {"Content-Type": "application/json"}) + t = time.time() + r = json.load(urllib.request.urlopen(req, timeout=1800)) + out[i] = (r["usage"], time.time() - t) + + +ts = [threading.Thread(target=one, args=(i,)) for i in range(conc)] +t0 = time.time() +[t.start() for t in ts] +[t.join() for t in ts] +tot = 0 +for i, (u, dt) in enumerate(out): + tot += u["completion_tokens"] + print(json.dumps({"req": i, "completion_tokens": u["completion_tokens"], + "accepted": (u.get("completion_tokens_details") or {}).get("accepted_prediction_tokens"), + "server_tok_s": u.get("response_token/s"), "wall_s": round(dt, 2)})) +print(json.dumps({"total_completion_tokens": tot, "wall_s": round(time.time() - t0, 2)})) diff --git a/docs/checkpoints/qwen3.8-flash-next-fp4-local.md b/docs/checkpoints/qwen3.8-flash-next-fp4-local.md new file mode 100644 index 000000000..877cba38b --- /dev/null +++ b/docs/checkpoints/qwen3.8-flash-next-fp4-local.md @@ -0,0 +1,210 @@ +# Qwen3.8-Flash-Next FP4-Local on one GB10 — native formats, measurements, open gap + +`nvidia/Qwen3.8-Flash-Next-FP4-Local` (snapshot `b9ce18a6887e18c089fdea8509a33e63429a3300`, ~96 GiB, +ModelOpt QAD) served by Atlas on a single DGX Spark (GB10, 121.6 GB unified memory), with every +tensor family run on the format the checkpoint ships. This is the record of what was built, how it +was verified, what it measures, and where it still trails NVIDIA's own vLLM recipe on the same box. + +**Status in one paragraph.** It loads, serves and is accurate: BFCL golden ST-995 87.44 overall / +87.32 normalized (n=995), passcode recall 12/12 from 32K to 245K tokens, multi-turn state and +cross-request isolation gates pass. Weights are checkpoint-native (no widening copies, no 8→4-bit +requantization). It is **not yet as fast as NVIDIA's vLLM recipe**: on the same box with the same +client, vLLM decodes 1.6–1.9x faster and prefills ~1.7x faster (table in §6). Closing that is the +next piece of work (§8). + +All numbers below were measured on `sparky` (GB10, driver 580.178.04, CUDA 13.0) with Atlas built +at `a78d2fb54` (single-target `qwen3.8-flash-next`, binary sha256 `8d566e6895f6f544…`). Each cites the +`spark benchmark` run id (`~/.atlas/runs/`), or names the streaming client used. Single-run +figures are labelled as such; nothing here is cross-validated beyond what is stated. After the +branch was rebased onto main `e29809db9`, the rebased build answered the 4-request output control +byte-identically to the measured `a78d2fb54` build, so the numbers carry over to the PR head. + +## 1. What the checkpoint ships + +| Tensor family | Count | On disk | +|---|---|---| +| Routed experts (48 layers × 512, gate/up/down) | 73,728 | NVFP4 (E2M1 + per-16 E4M3 scale + F32 scale_2), **scales CSF-compressed** (`weight_scale.nvfp4_csf_fixed` + `nvfp4_csf_exceptions`, U32) | +| GDN `in_proj_qkv`, `in_proj_z`, `out_proj` | 108 | MXFP8 (E4M3 + E8M0 per 32 inputs) | +| GDN `in_proj_a`, `in_proj_b` | 72 | MXFP8 | +| Full attention q/k/v/o (12 layers) | 48 | MXFP8 | +| QSA indexer `index_qk_proj` | 12 | MXFP8 | +| Shared expert gate/up/down | 144 | MXFP8 | +| Vision `attn.qkv`, `attn.proj` (27 blocks) | 54 | MXFP8 | +| Vision MLP `linear_fc1/fc2` + merger | 56 | NVFP4 (W4A16) | +| hyper-connection `hc_norm` + PLE `norm_conv/key/query` | 100 (+3 MTP) | **F32** (every other release ships BF16) | +| PLE n-gram table | 128 shards × 2,500,012 rows | NVFP4 rows U8 `[R,80]` + per-shard E4M3 `[R,10]` + one F32 `weight_scale_2` (28.8 GB) | +| MTP experts | 512 | per-expert NVFP4 | +| MTP shared expert, mHC mixing weights, LM head, embeddings | — | BF16 | + +Shards live under `tensors/` in the index. + +## 2. What Atlas does with each (branch `feat/fp4-local-flash-next`) + +| Family | Before this branch | Now | +|---|---|---| +| CSF expert scales | not loadable | decoded losslessly to plain E4M3 at load (`atlas-core::modelopt_codec`), byte-identical to vLLM's reference decoder (vllm#60646 spec) on the real weights; compressed originals never resident. Same as vLLM and TensorFold do. | +| GDN qkv/z/out, attention q/k/v/o | expanded to BF16 (attention then requantized to NVFP4 — lossy) | **native MXFP8**: `mxfp8_proj.cu` GEMV (1 row), batched GEMV (2–8 rows, every row bit-identical to the 1-row kernel), transient exact dequant + BF16 GEMM for prefill | +| Shared expert | MXFP8 → BF16 → **NVFP4 requant (8→4-bit loss on the one expert every token uses)** | **native MXFP8** through the separate-shared-expert seam (`SeparateSharedExpert::{Bf16,Mxfp8}`); gate‖up one `[1280,2560]` projection; fused routed kernels run with a NULL shared expert in every MoE arm | +| MTP shared expert | BF16 → NVFP4 requant | BF16 as shipped (same seam) | +| F32 norms | read as BF16 → NaN, then rounded to BF16 (lossy on all 100 tensors) | **F32 end to end**: hc/PLE norm kernel arguments are `const float*`; BF16 packs are widened exactly (`bits << 16`) | +| Vision MXFP8 attn / NVFP4 MLP | host-expanded to BF16 | **native**: uploaded as shipped, expanded per GEMM into a 10 MB encoder scratch (bit-identical to the old expansion; −0.52 GB) | +| PLE n-gram table | host decode of every row + stream sync per gather | **disk-backed, hot rows in pinned RAM, decoded on the GPU** (`batched_embed_nvfp4`, bit-identical to the host decoder); one cache holds rows and their block scales slot-parallel; `ATLAS_PLE_CACHE_GB` sizes it | +| GDN `in_proj_a/b`, indexer `index_qk_proj` | — | still expanded to BF16 (exact for MXFP8, ~28 MB). Not yet native. | + +Commits (rebased onto main for the PR; measured at the pre-rebase SHAs): +codec · spark-runtime loader (`tensors/` shards, U32, deferred host-decoded modules) · serve FP4-Local +· MTP per-expert NVFP4 · GDN native MXFP8 · attention native MXFP8 · F32 norms native · shared expert +native MXFP8 · gate‖up single buffer · PLE GPU decode + scalable hot cache · vision native. + +## 3. How exactness was verified + +* **Codecs:** CSF decode byte-identical to vLLM's reference decoder on the real tensors; MXFP8 dequant + exact vs host codec; NVFP4 PLE GPU decode identical to the host decoder over all 256 scale codes × + all packed bytes and over 4,096 real rows read through the cache; vision dequant identical for all + 54 MXFP8 + 56 NVFP4 projections (`examples/*_microtest.rs`). +* **NVFP4 pack never moved:** after every commit the NVFP4 pack (`nvidia/Qwen3.8-Flash-Next-NVFP4`) + answered a fixed 4-request control (math, explanation, code, tool call; temp 0) byte-identically + on the old and new binary. +* **FP4-Local output identity across the non-numeric commits** (GPU PLE decode, gate‖up buffer, + vision): the same 4-request control byte-identical old vs new (one tool request varies run to run + on the *same* server through an optional preamble sentence — prefix-cache path, not a decode + delta). +* **Speculation and batching are exact on FP4-Local's native paths** (live `*_CHECK` diagnostics, + every checked step re-run as serial decode on the live state): + `ATLAS_QWEN4EXP_EXACT_VERIFY` + `MTP_DEPTH=3` + `HC_FAST`: 64 verify steps / 256 rows, 0 differ; + `BATCH_FAST`+`BATCH_SMALL` at C=4: 128 steps / 512 rows, 0 differ; with MTP K=1 at C>1: 64 steps / + 444 rows, 0 differ; with `HC_FAST`+`MOE_UNITS`: 128 steps / 1,024 rows, 0 differ. + +## 4. Memory: why the n-gram table stays on disk + +The PLE table is 28.8 GB (25.6 GB packed rows + 3.2 GB scales). Making it resident on one Spark +needs ~131 GB against 121.6 GB (Atlas pre-KV ~93 GB + table + 4.2 GB working reserve + 5.5 GB OS). +It does not compress: measured on shards 0/57/127, the FP4 payload's nibble entropy is +**3.966 of 4 bits** and zstd-19 ratio **1.000**; only the scales compress (ratio 0.455, ~1.5 GB at +best). Getting ~17 GB elsewhere would have meant dropping MTP, a MoE-kernel rework for compressed +scales and still ~1 user at ~128K context. NVIDIA's own Spark recipe (`VLLM_PLE_TABLE_MEMORY=disk`) +and TensorFold also read the table from disk. Decision: **disk-backed, hot rows pinned in RAM, +GPU decode.** + +Hot-row cache (FP4-Local, same realistic workload per size — BFCL small draw, 16K-token prose +needles, code generation, the BFCL draw again): + +| `ATLAS_PLE_CACHE_GB` | slots | tracked weights | KV | misses | hit rate | evictions | +|---|---|---|---|---|---|---| +| unset (84 MB) | 524,288 | 91.32 GB | 11.2 GB (473K tok) | 433,477 | 94.6% | 0 | +| 2 | 22.2M | 93.26 GB | 9.3 GB (393K tok) | first-touch only | — | 0 | +| 8 | 88.9M | 99.19 GB | 3.1 GB (130K tok) | 433,477 | 94.8% | 0 | + +Every miss in this workload is a first touch; repeating the BFCL draw added zero misses even at +84 MB. A larger cache only pays once a session's unique rows exceed it, so the recipe uses **2 GB**, +which keeps a 262K-token context. "Tracked" is the driver-accounted figure; the boot log's +free-memory pre-KV estimate swings by up to ±7 GB between boots of the same binary and must not be +used for A/Bs. + +## 5. Serve configuration that was benchmarked + +Existing default-off qwen4_exp levers were swept on FP4-Local (fresh boot per config; decode-floor = +3 × 1,500-token MinHeap code generation, reasoning off, server-side rate; quick = 512-in/512-out): + +| C=1 config | decode-floor | quick | accept len | +|---|---|---|---| +| recipe default, MTP K=1 | 28.3 (run-1792317977888315647) | 27.4 | 1.91 | +| + exact verify + exact batching lane + fused decode, MTP depth 3 | 31.9 (run-1792342260484689522) | — | 3.32 | +| + `ATLAS_QWEN4EXP_HC_FAST` (adopted) | **33.0** (run-1792342559125109723) | **34.5** (run-1792342600593663107) | 3.32 | + +Drafts deeper than 3 did not engage: the scheduler capped single-sequence MTP at `k_drafts=3` even +with `MTP_DEPTH=5/7` and matching `--num-drafts`. Exact verify also lets MTP run inside ``. + +Chosen serve line (single user): + +``` +ATLAS_QWEN4EXP_EXACT_VERIFY=1 ATLAS_QWEN4EXP_DECODE_FUSE=1 ATLAS_QWEN4EXP_BATCH_FAST=1 \ +ATLAS_QWEN4EXP_BATCH_SMALL=1 ATLAS_QWEN4EXP_MTP_DEPTH=3 ATLAS_QWEN4EXP_HC_FAST=1 ATLAS_PLE_CACHE_GB=2 \ +spark serve --kernel-target qwen3.8-flash-next --gpu-memory-utilization 0.9 \ + --max-seq-len 32768 --max-prefill-tokens 8192 --kv-cache-dtype bf16 --max-num-seqs 1 --max-batch-size 1 \ + --speculative --num-drafts 3 --enable-prefix-caching --fast-load-prefetch-shards \ + --default-chat-template-kwargs '{"reasoning_effort":"low"}' +``` + +Multi-user: same levers minus `MTP_DEPTH`/`HC_FAST`, plus `ATLAS_MTP_MAX_SEQS=8`, `--max-num-seqs 8 +--max-batch-size 8 --speculative --num-drafts 1`. + +## 6. Results + +### Quality and robustness (FP4-Local) + +| Gate | Result | Run | +|---|---|---| +| BFCL golden ST-995 (no overrides, reasoning low) | **87.44 overall / 87.32 normalized**, n=995 | run-1792340347555552942 | +| Passcode recall, real-prose filler, 3 depths each | **12/12** at 32K, 64K, 128K, 245K prompt tokens | needle_long.py | +| SSM state poisoning (12 replayed rounds) | PASS, 12/12 invariant | run-1792328829391992739 | +| Cross-request contamination | PASS, 8/8 identical | run-1792329447184876759 | +| Video fidelity | PASS, 7/7 legs | run-1792328221225737347 | +| Vision fidelity | FAIL on 1 of 10 integrity probes (`multi-choice-image`); geometry 14/14, probes 3/3. Fails identically on the NVFP4 pack and on the pre-vision-change binary — pre-existing, not FP4-Local | run-1792328210080553327 | + +For reference, the NVFP4 recipe documents ST-995 84.92 on an older binary; a same-binary NVFP4 BFCL +run is still queued, so no FP4-Local-vs-NVFP4 accuracy claim is made here. + +### Speed, Atlas-internal harness (single run each; v1 = before HC_FAST) + +| | FP4-Local | NVFP4 pack (same flags) | +|---|---|---| +| decode-floor, C=1 | 32.5 (v1, run-1792327978365873482) / 33.0 (v2) | 31.9 (v2, run-1792343818643478918) | +| quick 512/512 decode / TTFT | 34.1 / 224 ms (v1) | 32.6 / 214 ms (v2) | +| cold TTFT 1K / 4K / 16K | 1.05 s / 2.94 s / 12.1 s (run-1792328125909646904) | 3.17 s median (run-1792343954260808403) | +| warm TTFT 4K / 16K | 213 ms / 259 ms | — | +| 262K serve: KV capacity, decode-floor | 257,056 tokens, 31.8 tok/s | — | +| 240K-token prompt TTFT | ~233 s (~1,050 tok/s prefill) | — | +| aggregate tok/s, isl 512, C1/C2/C4/C5/C8 | 24.1 / 28.7 / 36.9 / 37.4 / 42.3 (run-1792329307611412276) | 22.7 / 26.5 / 35.2 / 40.3 / 42.9 (run-1792345177245084155) | +| aggregate tok/s, isl 2048, C1/C4/C8 | 22.6 / 33.2 / 35.2 | 20.6 / 30.9 / 35.0 | + +### Same box, same client: Atlas vs NVIDIA's vLLM recipe + +NVIDIA's DGX Spark recipe from the model card, flag for flag (`eugr/spark-vllm-b12x:latest`, MTP 4 +drafts, MXFP8 LM head, b12x kernels, FP8 KV, PLE table on disk, util 0.8) against Atlas's chosen +single-user config, both driven by `ab_stream.py` (client clock from the SSE stream, server usage +for token counts, `reasoning_effort: low`, temperature 0, 3 distinct prompts each, medians): + +| | Atlas | vLLM (NVIDIA recipe) | +|---|---|---| +| short chat: decode / first token | 25.6 tok/s / 0.38 s | **41.5 tok/s / 0.26 s** | +| ~17–19K-token agent first turn (24 tools + docs): decode / first token | 22.9 tok/s / 15.4 s | **43.6 tok/s / 9.2 s** | + +This reproduces the community-reported numbers for that recipe (44 tok/s short, 0.27 s first token, +41.9 tok/s at a ~14K agent turn). Atlas also renders ~2K more prompt tokens than vLLM for the same +agent request (tool-schema templating) — worth its own look. TensorFold's README claims (~50 tok/s +single stream, ~130 aggregate at 5 users) were not reproduced here. + +## 7. Where the time goes (nsys, FP4-Local, before the lever sweep) + +Plain C=1 decode, 46.9 ms GPU per token: native MXFP8 projections 13.6 ms (at the ~216 GB/s ceiling +measured for this box), routed MoE gate/up + silu/down 11.6 ms (~53% of ceiling: 1.33 GB of expert +bytes a token), mHC down/finish 10.4 ms (~56%), BF16 LM head 5.0 ms. With MTP K=1 a verify cycle +read the BF16 LM head ~2.4×. At C=4 every row runs its own MoE (each row routes to its own 10 of 512 +experts) and the LM head ran on `dense_gemm_bf16` (13.8 ms a step; the exact batching lane replaces +it with a batched GEMV). + +## 8. What is left (ranked by expected gain) + +1. **Routed-MoE decode kernels at ~53% of bandwidth** — the largest single lever at every + concurrency and for MTP verify rows, which each read their own experts. +2. **MTP depth 4** like the vLLM recipe: lift the scheduler's single-sequence 3-draft cap. +3. **LM head**: vLLM ships it MXFP8 (`VLLM_MXFP8_LM_HEAD=1`); an opt-in MXFP8 head (quality-gated) + would halve its 1.27 GB read. +4. **Prefill**: ~1,050–1,400 tok/s vs vLLM's ~1.7x; MoE prefill and the PLE first-touch faults. +5. Remaining non-native tensors: GDN `in_proj_a/b`, indexer `index_qk_proj` (exact today, 28 MB). +6. Nits: boot log prints the PLE cache arena size wrong (tracked memory is right); strix-hip mirror + of the hc kernels updated but not compiled on HIP here; Atlas's template renders tool schemas + longer than vLLM's. + +## 9. Reproduction + +Scripts in `bench/qwen4_exp/fp4local/`: `control.py` (4-request output control), `needle_long.py` +(long-context recall on real prose), `ab_stream.py` (engine-neutral streaming A/B), `prof_req.py` +(nsys decode load). Everything else is the in-tree `spark benchmark` suite (`decode-floor`, +`quick-speed-bench`, `concurrency-sweep`, `ttft-cold-gate`, `ttft-warm-gate`, `bfcl-subset`, +`vision-fidelity`, `video-fidelity`, `cross-contamination`, `ssm-state-poisoning-gate`). + +Environment note: the BFCL scorer's venv (`~/.atlas/artifacts/bfcl/venv`) needed `soundfile` +(`qwen_agent` imports it) plus `libsndfile1`; without it every `bfcl-subset` run fails at scoring +after generating all responses. From 6fe7ffae43cfb9c0f6ef70794e0ced090dfb8c58 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:50:31 -0400 Subject: [PATCH 14/26] style: format the FP4-Local changes with current stable rustfmt CI's fmt job runs the latest stable rustfmt (1.10), which lays out several of this branch's new expressions differently from the pinned 1.93 toolchain they were formatted with. Formatting only. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spark-model/examples/vit_dq_microtest.rs | 6 ++-- crates/spark-model/src/layers/mod.rs | 2 +- .../spark-model/src/layers/ple/ngram_nvfp4.rs | 4 +-- .../spark-model/src/layers/vision_encoder.rs | 2 +- .../spark-model/src/weight_loader/qwen35.rs | 26 ++++++++------- .../spark-model/src/weight_loader/qwen3_vl.rs | 32 ++++++++++++++----- .../src/weight_map/modelopt_local_tests.rs | 12 +++---- .../src/ngram_cache/scale_tests.rs | 1 - crates/spark-storage/src/ngram_cache/tests.rs | 1 - 9 files changed, 51 insertions(+), 35 deletions(-) diff --git a/crates/spark-model/examples/vit_dq_microtest.rs b/crates/spark-model/examples/vit_dq_microtest.rs index 60b705c6c..129f281be 100644 --- a/crates/spark-model/examples/vit_dq_microtest.rs +++ b/crates/spark-model/examples/vit_dq_microtest.rs @@ -19,9 +19,9 @@ struct Store { impl Store { fn open(dir: &str) -> Result { - let idx: serde_json::Value = serde_json::from_slice(&std::fs::read( - format!("{dir}/model.safetensors.index.json"), - )?)?; + let idx: serde_json::Value = serde_json::from_slice(&std::fs::read(format!( + "{dir}/model.safetensors.index.json" + ))?)?; let mut headers = HashMap::new(); let mut shapes = HashMap::new(); for e in std::fs::read_dir(format!("{dir}/tensors"))? { diff --git a/crates/spark-model/src/layers/mod.rs b/crates/spark-model/src/layers/mod.rs index 33164a14d..1de624d6b 100644 --- a/crates/spark-model/src/layers/mod.rs +++ b/crates/spark-model/src/layers/mod.rs @@ -56,7 +56,7 @@ pub use qwen3_attention::Qwen3AttentionLayer; pub use qwen3_ssm::Qwen3SsmLayer; pub use qwen4exp_mtp::{Qwen4ExpMtpHead, Qwen4ExpMtpProposerState}; pub(crate) use vision_encoder::{GlmVisionBlockWeights, GlmVisionMergerWeights, GlmVisionWeights}; -pub use vision_encoder::{MergerLayer, ViTBlock, VitWeight, VisionEncoder}; +pub use vision_encoder::{MergerLayer, ViTBlock, VisionEncoder, VitWeight}; pub mod vision_tower; pub use vision_tower::VisionTower; diff --git a/crates/spark-model/src/layers/ple/ngram_nvfp4.rs b/crates/spark-model/src/layers/ple/ngram_nvfp4.rs index aa2694fc5..50523c413 100644 --- a/crates/spark-model/src/layers/ple/ngram_nvfp4.rs +++ b/crates/spark-model/src/layers/ple/ngram_nvfp4.rs @@ -14,9 +14,9 @@ //! 26.8 GiB on disk against the NVFP4 pack's 47.7 GiB FP8 table, with the //! same row count: fewer bytes per fault, more rows per page-cache GiB. -use anyhow::{Result, ensure}; #[cfg(test)] use anyhow::Context as _; +use anyhow::{Result, ensure}; #[cfg(test)] use atlas_core::modelopt_codec::nvfp4_rows_to_bf16; @@ -101,8 +101,6 @@ pub fn decode_rows_bf16( Ok(out) } - - #[cfg(test)] #[path = "ngram_nvfp4_tests.rs"] mod tests; diff --git a/crates/spark-model/src/layers/vision_encoder.rs b/crates/spark-model/src/layers/vision_encoder.rs index 88a69b7db..8e582e7b6 100644 --- a/crates/spark-model/src/layers/vision_encoder.rs +++ b/crates/spark-model/src/layers/vision_encoder.rs @@ -113,7 +113,7 @@ pub struct VisionEncoder { k_f32_bf16: KernelHandle, // vision_f32_to_bf16 k_copy: KernelHandle, // vision_bf16_copy k_nvfp4_dq: KernelHandle, // dequant_nvfp4_to_bf16 (native vision NVFP4) - k_mx8_dq: KernelHandle, // mxfp8_proj_dequant (native vision MXFP8) + k_mx8_dq: KernelHandle, // mxfp8_proj_dequant (native vision MXFP8) // config pub hidden_size: usize, // 1152 pub num_heads: usize, // 16 diff --git a/crates/spark-model/src/weight_loader/qwen35.rs b/crates/spark-model/src/weight_loader/qwen35.rs index f9bc6e3e2..7de8c21ff 100644 --- a/crates/spark-model/src/weight_loader/qwen35.rs +++ b/crates/spark-model/src/weight_loader/qwen35.rs @@ -26,7 +26,11 @@ fn vision_tensor_ptr( store: &WeightStore, name: &str, gpu: &dyn GpuBackend, -) -> Result<(spark_runtime::gpu::DevicePtr, spark_runtime::weights::WeightDtype, Vec)> { +) -> Result<( + spark_runtime::gpu::DevicePtr, + spark_runtime::weights::WeightDtype, + Vec, +)> { if let Ok(t) = store.get(name) { return Ok((t.ptr, t.dtype, t.shape.clone())); } @@ -71,23 +75,23 @@ fn vision_weight( && sshape == [wshape[0], wshape[1] / 32], "vision MXFP8 {prefix}: weight {wshape:?} + scale {sshape:?}" ); - Ok(crate::layers::VitWeight::Mxfp8(crate::weight_map::Mxfp8Proj { - w: crate::weight_map::Mxfp8Weight { - data: wptr, - scales: sptr, + Ok(crate::layers::VitWeight::Mxfp8( + crate::weight_map::Mxfp8Proj { + w: crate::weight_map::Mxfp8Weight { + data: wptr, + scales: sptr, + }, + n: wshape[0] as u32, + k: wshape[1] as u32, }, - n: wshape[0] as u32, - k: wshape[1] as u32, - })) + )) } WeightDtype::UInt8 => { // NVFP4: packed U8 [n,k/2] + E4M3 [n,k/16] + F32 weight_scale_2. let (sptr, sdtype, sshape) = vision_tensor_ptr(store, &format!("{prefix}.weight_scale"), gpu)?; anyhow::ensure!( - wshape.len() == 2 - && sdtype == WeightDtype::FP8E4M3 - && sshape[0] == wshape[0], + wshape.len() == 2 && sdtype == WeightDtype::FP8E4M3 && sshape[0] == wshape[0], "vision NVFP4 {prefix}: weight {wshape:?} + scale {sshape:?}" ); let gname = format!("{prefix}.weight_scale_2"); diff --git a/crates/spark-model/src/weight_loader/qwen3_vl.rs b/crates/spark-model/src/weight_loader/qwen3_vl.rs index e5d5774ec..012fc6f7e 100644 --- a/crates/spark-model/src/weight_loader/qwen3_vl.rs +++ b/crates/spark-model/src/weight_loader/qwen3_vl.rs @@ -237,15 +237,23 @@ impl ModelWeightLoader for Qwen3VLWeightLoader { blocks.push(ViTBlock { norm1_w: dense(store, &format!("{bp}.norm1.weight"))?.weight, norm1_b: dense(store, &format!("{bp}.norm1.bias"))?.weight, - qkv_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{bp}.attn.qkv.weight"))?.weight), + qkv_w: crate::layers::VitWeight::Bf16( + dense(store, &format!("{bp}.attn.qkv.weight"))?.weight, + ), qkv_b: dense(store, &format!("{bp}.attn.qkv.bias"))?.weight, - proj_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{bp}.attn.proj.weight"))?.weight), + proj_w: crate::layers::VitWeight::Bf16( + dense(store, &format!("{bp}.attn.proj.weight"))?.weight, + ), proj_b: dense(store, &format!("{bp}.attn.proj.bias"))?.weight, norm2_w: dense(store, &format!("{bp}.norm2.weight"))?.weight, norm2_b: dense(store, &format!("{bp}.norm2.bias"))?.weight, - fc1_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{bp}.mlp.linear_fc1.weight"))?.weight), + fc1_w: crate::layers::VitWeight::Bf16( + dense(store, &format!("{bp}.mlp.linear_fc1.weight"))?.weight, + ), fc1_b: dense(store, &format!("{bp}.mlp.linear_fc1.bias"))?.weight, - fc2_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{bp}.mlp.linear_fc2.weight"))?.weight), + fc2_w: crate::layers::VitWeight::Bf16( + dense(store, &format!("{bp}.mlp.linear_fc2.weight"))?.weight, + ), fc2_b: dense(store, &format!("{bp}.mlp.linear_fc2.bias"))?.weight, }); } @@ -256,9 +264,13 @@ impl ModelWeightLoader for Qwen3VLWeightLoader { deepstack.push(MergerLayer { norm_w: dense(store, &format!("{mp}.norm.weight"))?.weight, norm_b: dense(store, &format!("{mp}.norm.bias"))?.weight, - fc1_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{mp}.linear_fc1.weight"))?.weight), + fc1_w: crate::layers::VitWeight::Bf16( + dense(store, &format!("{mp}.linear_fc1.weight"))?.weight, + ), fc1_b: dense(store, &format!("{mp}.linear_fc1.bias"))?.weight, - fc2_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{mp}.linear_fc2.weight"))?.weight), + fc2_w: crate::layers::VitWeight::Bf16( + dense(store, &format!("{mp}.linear_fc2.weight"))?.weight, + ), fc2_b: dense(store, &format!("{mp}.linear_fc2.bias"))?.weight, }); } @@ -267,9 +279,13 @@ impl ModelWeightLoader for Qwen3VLWeightLoader { let merger = MergerLayer { norm_w: dense(store, &format!("{mp}.norm.weight"))?.weight, norm_b: dense(store, &format!("{mp}.norm.bias"))?.weight, - fc1_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{mp}.linear_fc1.weight"))?.weight), + fc1_w: crate::layers::VitWeight::Bf16( + dense(store, &format!("{mp}.linear_fc1.weight"))?.weight, + ), fc1_b: dense(store, &format!("{mp}.linear_fc1.bias"))?.weight, - fc2_w: crate::layers::VitWeight::Bf16(dense(store, &format!("{mp}.linear_fc2.weight"))?.weight), + fc2_w: crate::layers::VitWeight::Bf16( + dense(store, &format!("{mp}.linear_fc2.weight"))?.weight, + ), fc2_b: dense(store, &format!("{mp}.linear_fc2.bias"))?.weight, }; diff --git a/crates/spark-model/src/weight_map/modelopt_local_tests.rs b/crates/spark-model/src/weight_map/modelopt_local_tests.rs index b0b62d6d8..a49e5f047 100644 --- a/crates/spark-model/src/weight_map/modelopt_local_tests.rs +++ b/crates/spark-model/src/weight_map/modelopt_local_tests.rs @@ -56,10 +56,7 @@ fn fp4_local_config() -> ModelConfig { #[test] fn host_decode_modules_pick_mxfp8_and_vision_nvfp4_only() { let got = host_decode_modules(&fp4_local_config()); - let want: HashSet = [SHARED, GDN_OUT] - .iter() - .map(|s| s.to_string()) - .collect(); + let want: HashSet = [SHARED, GDN_OUT].iter().map(|s| s.to_string()).collect(); assert_eq!(got, want); assert!(host_decode_modules(&ModelConfig::qwen3_next_80b_nvfp4()).is_empty()); } @@ -123,7 +120,6 @@ fn resident( .unwrap(); } - #[test] fn materialize_expands_every_fp4_local_format() { let gpu = MockGpuBackend::new(); @@ -266,7 +262,11 @@ fn materialize_expands_every_fp4_local_format() { assert_eq!(bias.dtype, WeightDtype::FP32); // The vision NVFP4 tensors are NOT consumed: they are a native GPU // dequant path now, read as deferred shards by the vision loader. - assert_eq!(store.deferred_sorted().len(), 2, "vision tensors stay deferred"); + assert_eq!( + store.deferred_sorted().len(), + 2, + "vision tensors stay deferred" + ); let scale = store.get(EXPERT).unwrap(); assert_eq!( diff --git a/crates/spark-storage/src/ngram_cache/scale_tests.rs b/crates/spark-storage/src/ngram_cache/scale_tests.rs index f21be841a..50b0610ef 100644 --- a/crates/spark-storage/src/ngram_cache/scale_tests.rs +++ b/crates/spark-storage/src/ngram_cache/scale_tests.rs @@ -3,7 +3,6 @@ //! Segmented block-scale coverage for `NgramRowCache` — the NVFP4 PLE /// table's sibling `shard_i.weight_scale` rows, faulted through the same /// slot as the weight row. Split out of `tests.rs` for the 500-LoC cap. - use super::*; /// Pinned arena needs a live CUDA context (see ). diff --git a/crates/spark-storage/src/ngram_cache/tests.rs b/crates/spark-storage/src/ngram_cache/tests.rs index 0ef1657f5..047f05974 100644 --- a/crates/spark-storage/src/ngram_cache/tests.rs +++ b/crates/spark-storage/src/ngram_cache/tests.rs @@ -383,4 +383,3 @@ fn prefetch_dedups_repeated_ids() { assert_ne!(slots[0], slots[3]); let _ = std::fs::remove_dir_all(&dir); } - From b2c031e7f2764a66edfe7244fc4290240b3d55b7 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Mon, 19 Oct 2026 02:20:14 -0400 Subject: [PATCH 15/26] modelopt_local: drive FP4-Local loading from the declared module formats Replace the qwen4_exp-shaped name lists with ModeloptQuantPlan: the checkpoint's quantized_layers map (config.json or hf_quant_config.json) yields a per-module ModuleFormat { algo, group_size, scale_encoding }, accepting both group_size and weight_block_size spellings and the FP8_PB_WO tile form. Which modules host-expand and which stay native now come from the plan's algo intersected with a per-model NativeConsumers registry; CSF streams decode only for weight_scale_encoding == "csf", any other encoding fails the load naming it. A model with no registered consumers gets the exact-but-unoptimized BF16 expansion with no names needed. FP4-Local loads byte-identical to the name-driven path. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/atlas-core/src/config.rs | 7 + .../src/config/parsers/quantization.rs | 15 ++ crates/spark-model/src/weight_map.rs | 1 + .../src/weight_map/modelopt_local.rs | 85 ++++---- .../src/weight_map/modelopt_local_tests.rs | 122 +++++++++-- .../src/weight_map/modelopt_plan.rs | 200 ++++++++++++++++++ .../src/weight_map/modelopt_plan_tests.rs | 192 +++++++++++++++++ .../nvfp4_detect/ep_detection_tests.rs | 16 ++ 8 files changed, 586 insertions(+), 52 deletions(-) create mode 100644 crates/spark-model/src/weight_map/modelopt_plan.rs create mode 100644 crates/spark-model/src/weight_map/modelopt_plan_tests.rs diff --git a/crates/atlas-core/src/config.rs b/crates/atlas-core/src/config.rs index 64e9eeb1e..b27aba20b 100644 --- a/crates/atlas-core/src/config.rs +++ b/crates/atlas-core/src/config.rs @@ -733,6 +733,13 @@ pub struct QuantizationConfig { pub struct QuantLayerSpec { pub quant_algo: String, pub group_size: usize, + /// ModelOpt scale-stream encoding, e.g. `"csf"` or `"csf_32x2_v1"`. + /// `Some` marks the FP4-Local family (compressed scale streams beside + /// `weight_scale`); `None` for plain row scales. + pub weight_scale_encoding: Option, + /// Per-module block-quant tile, e.g. `[1, 32]` (MXFP8's 1×32 group, + /// spelled as a tile) or `[128, 128]` (FP8_PB_WO). Empty when absent. + pub weight_block_size: Vec, } /// Whether a ModelOpt `quant_algo` label means NVFP4-packed weights. diff --git a/crates/atlas-core/src/config/parsers/quantization.rs b/crates/atlas-core/src/config/parsers/quantization.rs index 652f65751..d78e37a2d 100644 --- a/crates/atlas-core/src/config/parsers/quantization.rs +++ b/crates/atlas-core/src/config/parsers/quantization.rs @@ -152,11 +152,26 @@ pub fn parse_quantization_config(raw: &serde_json::Value) -> Option HashSet { let Some(qc) = config.quantization_config.as_ref() else { return HashSet::new(); }; + if let Some(plan) = ModeloptQuantPlan::from_config(config) + && plan.is_fp4_local() + { + let natives = native_consumers(&config.model_type); + return plan + .modules() + .filter(|m| { + let f = plan.format_of(m).expect("plan module"); + matches!(f.algo, QuantAlgo::Mxfp8) || !natives.covers(m, &f.algo) + }) + .cloned() + .collect(); + } qc.quantized_layers .iter() - .filter(|(_path, spec)| { - let algo = spec.quant_algo.to_ascii_uppercase(); - algo == "MXFP8" - }) + .filter(|(_path, spec)| spec.quant_algo.eq_ignore_ascii_case("MXFP8")) .map(|(path, _)| path.clone()) .collect() } -/// MXFP8 modules a layer reads natively, so they are uploaded as shipped. -/// `ATLAS_FP4LOCAL_EXPAND_MXFP8=1` expands them to BF16 instead (A/B only: -/// both are exact, the native bytes are half the size). -pub fn native_mxfp8(module: &str) -> bool { - if std::env::var("ATLAS_FP4LOCAL_EXPAND_MXFP8").as_deref() == Ok("1") { - return false; - } - const NATIVE: [&str; 10] = [ - ".linear_attn.in_proj_qkv", - ".linear_attn.in_proj_z", - ".linear_attn.out_proj", - ".self_attn.q_proj", - ".self_attn.k_proj", - ".self_attn.v_proj", - ".self_attn.o_proj", - ".mlp.shared_expert.gate_proj", - ".mlp.shared_expert.up_proj", - ".mlp.shared_expert.down_proj", - ]; - (module.starts_with("model.language_model.layers.") - && NATIVE.iter().any(|suffix| module.ends_with(suffix))) - // The vision tower's MXFP8 attention projections upload as shipped and - // dequant on-GPU into the encoder's BF16 scratch. - || (module.starts_with("model.visual.") - && (module.ends_with(".attn.qkv") || module.ends_with(".attn.proj"))) -} - /// What [`materialize`] did, for the load log and for tests. #[derive(Debug, Default, PartialEq, Eq)] pub struct Materialized { @@ -319,18 +304,44 @@ pub fn materialize( gpu: &dyn GpuBackend, ) -> Result { let mut done = Materialized::default(); + let plan = ModeloptQuantPlan::from_config(config); + let natives = native_consumers(&config.model_type); for base in csf_bases(store) { + // The scale-encoding column decides the codec; `"csf"` is the only + // one Atlas decodes today. A module the plan doesn't list keeps the + // stream's name-driven decode — don't guess at an unknown layout. + if let Some(enc) = base + .strip_suffix(".weight_scale") + .and_then(|m| plan.as_ref().and_then(|p| p.format_of(m))) + .and_then(|f| f.scale_encoding.as_deref()) + { + ensure!( + enc == "csf", + "{base}: unsupported weight_scale_encoding {enc:?}; cannot decode" + ); + } done.uploaded_bytes += decode_one_csf(store, gpu, &base)?; done.csf_scales += 1; } let mut modules: Vec = host_decode_modules(config).into_iter().collect(); modules.sort(); + // A/B only: expand the native MXFP8 set too (both paths are exact). + let force_expand = + std::env::var("ATLAS_FP4LOCAL_EXPAND_MXFP8").as_deref() == Ok("1"); for module in modules { - if native_mxfp8(&module) { - if let Some(n) = upload_native_mxfp8(store, gpu, &module)? { + let fmt = plan.as_ref().and_then(|p| p.format_of(&module)); + let algo = fmt + .map(|f| &f.algo) + .cloned() + .unwrap_or(QuantAlgo::Other(String::new())); + if !force_expand && natives.covers(&module, &algo) { + if algo == QuantAlgo::Mxfp8 + && let Some(n) = upload_native_mxfp8(store, gpu, &module)? + { done.mxfp8_native += 1; done.uploaded_bytes += n; } + // An NVFP4/algo-native module stays deferred for its own loader. continue; } let Some((mx, n)) = expand_module(store, gpu, &module)? else { diff --git a/crates/spark-model/src/weight_map/modelopt_local_tests.rs b/crates/spark-model/src/weight_map/modelopt_local_tests.rs index a49e5f047..5b3c0b67e 100644 --- a/crates/spark-model/src/weight_map/modelopt_local_tests.rs +++ b/crates/spark-model/src/weight_map/modelopt_local_tests.rs @@ -17,6 +17,8 @@ fn spec(algo: &str, group: usize) -> QuantLayerSpec { QuantLayerSpec { quant_algo: algo.into(), group_size: group, + weight_scale_encoding: None, + weight_block_size: vec![], } } @@ -24,10 +26,15 @@ fn spec(algo: &str, group: usize) -> QuantLayerSpec { /// experts, an MXFP8 projection, a W4A16 NVFP4 vision MLP, and the PLE table. fn fp4_local_config() -> ModelConfig { let mut cfg = ModelConfig::qwen3_next_80b_nvfp4(); + cfg.model_type = "qwen4_exp".into(); let mut map = std::collections::BTreeMap::new(); map.insert( "model.language_model.layers.0.mlp.experts".into(), - spec("NVFP4", 16), + { + let mut s = spec("NVFP4", 16); + s.weight_scale_encoding = Some("csf".into()); + s + }, ); map.insert(SHARED.into(), spec("MXFP8", 32)); map.insert(GDN_OUT.into(), spec("MXFP8", 32)); @@ -293,20 +300,17 @@ fn materialize_expands_every_fp4_local_format() { /// exactly as shipped, weight and scale, and nothing is widened. #[test] fn gdn_mxfp8_stays_native() { - assert!(native_mxfp8(GDN_OUT)); - assert!(native_mxfp8( - "model.language_model.layers.7.linear_attn.in_proj_qkv" - )); - assert!(native_mxfp8(ATTN)); - assert!(!native_mxfp8( - "model.language_model.layers.3.self_attn.indexer.index_qk_proj" - )); - assert!(native_mxfp8("model.visual.blocks.0.attn.qkv")); - assert!(native_mxfp8("model.visual.blocks.0.attn.proj")); - assert!(!native_mxfp8("model.visual.blocks.0.mlp.linear_fc1")); - assert!(!native_mxfp8( - "model.language_model.layers.7.linear_attn.in_proj_a" - )); + use crate::weight_map::modelopt_plan::{QuantAlgo, native_consumers}; + let n = native_consumers("qwen4_exp"); + let mx = |m: &str| n.covers(m, &QuantAlgo::Mxfp8); + assert!(mx(GDN_OUT)); + assert!(mx("model.language_model.layers.7.linear_attn.in_proj_qkv")); + assert!(mx(ATTN)); + assert!(!mx("model.language_model.layers.3.self_attn.indexer.index_qk_proj")); + assert!(mx("model.visual.blocks.0.attn.qkv")); + assert!(mx("model.visual.blocks.0.attn.proj")); + assert!(!mx("model.visual.blocks.0.mlp.linear_fc1")); + assert!(!mx("model.language_model.layers.7.linear_attn.in_proj_a")); let gpu = MockGpuBackend::new(); let mut store = WeightStore::empty(); @@ -404,3 +408,91 @@ fn materialize_is_a_no_op_without_fp4_local_tensors() { // Another release's F32 tensors are never touched. assert_eq!(store.get("b.norm.weight").unwrap().dtype, WeightDtype::FP32); } + +/// A CSF scale stream whose declared encoding is not `"csf"` must fail at +/// load naming the encoding — never guessed. +#[test] +fn unknown_csf_encoding_is_refused() { + let gpu = MockGpuBackend::new(); + let mut store = WeightStore::empty(); + let mut shard = Shard::new("enc"); + let module = "model.language_model.layers.0.mlp.experts"; + shard.add( + &format!("{module}.weight_scale.nvfp4_csf_fixed"), + &[0; 8], + &[2, 4], + WeightDtype::UInt8, + ); + shard.add( + &format!("{module}.weight_scale.nvfp4_csf_exceptions"), + &[0; 4], + &[4], + WeightDtype::UInt8, + ); + shard.defer_into(&mut store); + + let mut cfg = fp4_local_config(); + cfg.quantization_config + .as_mut() + .unwrap() + .quantized_layers + .get_mut(module) + .unwrap() + .weight_scale_encoding = Some("csf_32x2_v1".into()); + let err = materialize(&mut store, &cfg, &gpu).unwrap_err().to_string(); + assert!( + err.contains("csf_32x2_v1"), + "error must name the encoding: {err}" + ); +} + +/// A module whose algo has no registered consumer — here a W4A16 NVFP4 on +/// a non-expert, non-vision path — falls through to the exact BF16 +/// expansion. +#[test] +fn unregistered_nvfp4_module_expands() { + let gpu = MockGpuBackend::new(); + let mut store = WeightStore::empty(); + let mut shard = Shard::new("unreg"); + let module = "model.language_model.layers.0.mlp.down_proj"; + // NVFP4 [1, 16]: nibble 2 (=1.0) everywhere, block scale 2.0, global 0.5 + // → all ones BF16. + shard.add( + &format!("{module}.weight"), + &[0x22; 8], + &[1, 8], + WeightDtype::UInt8, + ); + shard.add( + &format!("{module}.weight_scale"), + &[0x40], + &[1, 1], + WeightDtype::FP8E4M3, + ); + shard.defer_into(&mut store); + resident( + &mut store, + &gpu, + &format!("{module}.weight_scale_2"), + &0.5f32.to_le_bytes(), + WeightDtype::FP32, + ); + + let mut cfg = fp4_local_config(); + cfg.quantization_config + .as_mut() + .unwrap() + .quantized_layers + .insert(module.into(), spec("W4A16_NVFP4", 16)); + let done = materialize(&mut store, &cfg, &gpu).unwrap(); + assert_eq!(done.nvfp4_to_bf16, 1); + let w = store.get(&format!("{module}.weight")).unwrap(); + assert_eq!( + (w.dtype, w.shape.clone()), + (WeightDtype::BF16, vec![1, 16]) + ); + let v = gpu.read_alloc(w.ptr).unwrap(); + assert!((0..16).all(|i| { + f32::from_bits(u32::from(u16::from_le_bytes([v[2 * i], v[2 * i + 1]])) << 16) == 1.0 + })); +} diff --git a/crates/spark-model/src/weight_map/modelopt_plan.rs b/crates/spark-model/src/weight_map/modelopt_plan.rs new file mode 100644 index 000000000..b10fb6194 --- /dev/null +++ b/crates/spark-model/src/weight_map/modelopt_plan.rs @@ -0,0 +1,200 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! The ModelOpt `quantized_layers` map as a per-module format plan. +//! +//! NVIDIA's FP4-Local checkpoints declare every module's storage format in +//! `config.json` → `quantization_config.quantized_layers` (or a sibling +//! `hf_quant_config.json` → `quantization.quantized_layers`), so nothing +//! about the load needs to be keyed on qwen4_exp's module names: the plan +//! says which algo and scale encoding each module uses, and a per-model +//! [`NativeConsumers`] registry says which of those the model's layers can +//! run natively. Everything else is host-expanded to BF16. +//! +//! `is_fp4_local` — a `weight_scale_encoding` on any module — is what makes a +//! checkpoint "FP4-Local": an encoding entry means the pack ships compressed +//! scale streams (`weight_scale.nvfp4_csf_*`) and the ModelOpt-local decode +//! path applies. + +use std::collections::BTreeMap; + +use atlas_core::config::ModelConfig; + +/// The quantized-layer algorithm, normalized across spellings. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum QuantAlgo { + /// `"NVFP4"` — packed E2M1 + E4M3 per-group scales (+ optional CSF). + Nvfp4, + /// `"W4A16_NVFP4"` and friends — same packed format, BF16 activations. + W4A16Nvfp4, + /// `"MXFP8"` — E4M3 + E8M0 per-32-input scales (`group_size` 32 or a + /// `weight_block_size` of `[1, 32]`). + Mxfp8, + /// `"FP8_PB_WO"` / `"FP8_BLOCK_SCALES"` — E4M3 with a `[bn, bk]` tile. + Fp8Block { bn: usize, bk: usize }, + /// Anything else, kept verbatim for error messages. + Other(String), +} + +/// One module's declared storage format. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ModuleFormat { + pub algo: QuantAlgo, + /// Scaling group along the input dim: `group_size`, or the `bk` of a + /// `[1, k]` `weight_block_size`. `0` when neither is declared. + pub group_size: usize, + /// `weight_scale_encoding` verbatim (`"csf"`, `"csf_32x2_v1"`, …). + pub scale_encoding: Option, +} + +/// The whole `quantized_layers` map, normalized. +#[derive(Debug, Default)] +pub struct ModeloptQuantPlan { + formats: BTreeMap, +} + +impl ModeloptQuantPlan { + /// From the parsed config (`config.json` or `hf_quant_config.json` — + /// `parse_config` has already merged them). `None` without a map. + pub fn from_config(config: &ModelConfig) -> Option { + let qc = config.quantization_config.as_ref()?; + if qc.quantized_layers.is_empty() { + return None; + } + let mut formats = BTreeMap::new(); + for (module, spec) in &qc.quantized_layers { + let algo = spec.quant_algo.trim(); + let algo = if algo.eq_ignore_ascii_case("MXFP8") { + QuantAlgo::Mxfp8 + } else if algo.eq_ignore_ascii_case("NVFP4") { + QuantAlgo::Nvfp4 + } else if atlas_core::config::is_nvfp4_quant_algo(algo) { + QuantAlgo::W4A16Nvfp4 + } else if matches!( + algo.to_ascii_uppercase().as_str(), + "FP8_PB_WO" | "FP8_BLOCK_SCALES" | "FP8" + ) && spec.weight_block_size.len() == 2 + { + QuantAlgo::Fp8Block { + bn: spec.weight_block_size[0], + bk: spec.weight_block_size[1], + } + } else { + QuantAlgo::Other(algo.to_string()) + }; + // `weight_block_size [1, k]` is the tiled spelling of + // `group_size k`; a `[n>1, k]` tile groups n-rows instead and + // keeps group_size 0 here (block scaling is a different codec). + let group_size = if spec.group_size != 0 { + spec.group_size + } else if spec.weight_block_size.len() == 2 && spec.weight_block_size[0] == 1 { + spec.weight_block_size[1] + } else { + 0 + }; + formats.insert( + module.clone(), + ModuleFormat { + algo, + group_size, + scale_encoding: spec.weight_scale_encoding.clone(), + }, + ); + } + Some(Self { formats }) + } + + /// The declared format of `module`, if the map lists it. + pub fn format_of(&self, module: &str) -> Option<&ModuleFormat> { + self.formats.get(module) + } + + /// Every module the map declares. + pub fn modules(&self) -> impl Iterator { + self.formats.keys() + } + + /// The FP4-Local family marker: any module ships a compressed + /// `weight_scale` encoding (today `"csf"`). + pub fn is_fp4_local(&self) -> bool { + self.formats + .values() + .any(|f| f.scale_encoding.is_some()) + } +} + +/// Which module formats a model's layers run natively — declared by the +/// model loader, matched against module paths in the plan. Modules with no +/// registered consumer are host-expanded to BF16 by `materialize`. +#[derive(Debug, Clone, Copy)] +pub struct NativeConsumers { + /// Modules whose MXFP8 tensors upload and dequant as shipped + /// (`mxfp8_proj`). + pub mxfp8: fn(&str) -> bool, + /// Modules whose NVFP4 tensors are read natively (deferred to the layer's + /// own loader — MoE experts, the vision tower, the PLE table). + pub nvfp4: fn(&str) -> bool, +} + +impl NativeConsumers { + const NONE: Self = Self { + mxfp8: |_| false, + nvfp4: |_| false, + }; + + /// Whether `module`'s declared `algo` has a native consumer here. + pub fn covers(&self, module: &str, algo: &QuantAlgo) -> bool { + match algo { + QuantAlgo::Mxfp8 => (self.mxfp8)(module), + QuantAlgo::Nvfp4 | QuantAlgo::W4A16Nvfp4 => (self.nvfp4)(module), + _ => false, + } + } +} + +fn qwen4_exp_native_mxfp8(module: &str) -> bool { + const NATIVE: [&str; 10] = [ + ".linear_attn.in_proj_qkv", + ".linear_attn.in_proj_z", + ".linear_attn.out_proj", + ".self_attn.q_proj", + ".self_attn.k_proj", + ".self_attn.v_proj", + ".self_attn.o_proj", + ".mlp.shared_expert.gate_proj", + ".mlp.shared_expert.up_proj", + ".mlp.shared_expert.down_proj", + ]; + (module.starts_with("model.language_model.layers.") + && NATIVE.iter().any(|suffix| module.ends_with(suffix))) + // The vision tower's MXFP8 attention projections upload as shipped and + // dequant on-GPU into the encoder's BF16 scratch (`model.visual.` and + // the `model.language_model.visual.` spelling both land here). + || (module.contains(".visual.") + && (module.ends_with(".attn.qkv") || module.ends_with(".attn.proj"))) +} + +fn qwen4_exp_native_nvfp4(module: &str) -> bool { + // Routed + MTP experts (the per-expert NVFP4 MoE loaders), the vision + // tower's linear_fc/merger projections (GPU dequant into BF16 scratch), + // and the PLE n-gram table (row-gathered off disk). + module.ends_with(".mlp.experts") + || module.starts_with("mtp.") + || module.contains(".visual.") + || module.ends_with(".ngram_embedding") +} + +/// The native-consumer registry, keyed by model type. Models not listed run +/// every declared module through the generic host expansion. +pub fn native_consumers(model_type: &str) -> NativeConsumers { + match model_type { + "qwen4_exp" => NativeConsumers { + mxfp8: qwen4_exp_native_mxfp8, + nvfp4: qwen4_exp_native_nvfp4, + }, + _ => NativeConsumers::NONE, + } +} + +#[cfg(test)] +#[path = "modelopt_plan_tests.rs"] +mod tests; diff --git a/crates/spark-model/src/weight_map/modelopt_plan_tests.rs b/crates/spark-model/src/weight_map/modelopt_plan_tests.rs new file mode 100644 index 000000000..d7174f706 --- /dev/null +++ b/crates/spark-model/src/weight_map/modelopt_plan_tests.rs @@ -0,0 +1,192 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! `ModeloptQuantPlan` parsing + the native-consumer registry. + +use super::*; +use atlas_core::config::{ModelConfig, QuantLayerSpec, QuantizationConfig}; +use std::collections::BTreeMap; + +fn spec(algo: &str) -> QuantLayerSpec { + QuantLayerSpec { + quant_algo: algo.into(), + group_size: 0, + weight_scale_encoding: None, + weight_block_size: vec![], + } +} + +fn cfg_with(layers: &[(&str, QuantLayerSpec)]) -> ModelConfig { + let mut cfg = ModelConfig::qwen3_next_80b_nvfp4(); + let mut map = BTreeMap::new(); + for (k, v) in layers { + map.insert(k.to_string(), v.clone()); + } + cfg.quantization_config = Some(QuantizationConfig { + quant_method: "modelopt".into(), + quant_algo: "MIXED_PRECISION".into(), + format: String::new(), + ignore_modules: vec![], + weight_block_size: vec![], + group_size: 0, + quantized_layers: map, + exl3: None, + }); + cfg +} + +/// Flash-Next spellings: `group_size`, `weight_scale_encoding: "csf"`, +/// `W4A16_NVFP4`, one plain NVFP4. +#[test] +fn parses_the_flash_next_map() { + let mut mx = spec("MXFP8"); + mx.group_size = 32; + let mut nv4 = spec("NVFP4"); + nv4.group_size = 16; + nv4.weight_scale_encoding = Some("csf".into()); + let mut w4 = spec("W4A16_NVFP4"); + w4.group_size = 16; + let mut plain = spec("NVFP4"); + plain.group_size = 16; + let mut cfg = cfg_with(&[ + ("model.language_model.layers.0.self_attn.q_proj", mx), + ("model.language_model.layers.0.mlp.experts", nv4), + ("model.visual.blocks.0.mlp.linear_fc1", w4), + ("model.language_model.layers.1.ple.ple_embedding.ngram_embedding", plain), + ]); + cfg.model_type = "qwen4_exp".into(); + let plan = ModeloptQuantPlan::from_config(&cfg).unwrap(); + assert!(plan.is_fp4_local()); + assert_eq!( + plan.format_of("model.language_model.layers.0.mlp.experts") + .unwrap() + .algo, + QuantAlgo::Nvfp4 + ); + assert_eq!( + plan.format_of("model.visual.blocks.0.mlp.linear_fc1") + .unwrap() + .algo, + QuantAlgo::W4A16Nvfp4 + ); + assert_eq!( + plan.format_of("model.language_model.layers.0.self_attn.q_proj") + .unwrap() + .group_size, + 32 + ); +} + +/// The other family's spellings: MXFP8 as `weight_block_size [1,32]`, +/// `FP8_PB_WO` with `[128,128]`, and `csf_32x2_v1` — all parsed, none +/// guessed at. +#[test] +fn parses_the_csf_32x2_fixture_map() { + let mut mx = spec("MXFP8"); + mx.weight_block_size = vec![1, 32]; + let mut pb = spec("FP8_PB_WO"); + pb.weight_block_size = vec![128, 128]; + let mut csf2 = spec("NVFP4"); + csf2.weight_scale_encoding = Some("csf_32x2_v1".into()); + let mut mx_plain = spec("MXFP8"); + mx_plain.group_size = 32; + let cfg = cfg_with(&[ + ("model.language_model.layers.0.mlp.gate_proj", mx), + ("model.language_model.layers.0.mlp.up_proj", pb), + ("model.language_model.layers.0.mlp.experts", csf2), + ("lm_head", mx_plain), + ]); + let plan = ModeloptQuantPlan::from_config(&cfg).unwrap(); + assert!(plan.is_fp4_local()); + let gate = plan.format_of("model.language_model.layers.0.mlp.gate_proj").unwrap(); + assert_eq!(gate.algo, QuantAlgo::Mxfp8); + assert_eq!(gate.group_size, 32); + assert_eq!( + plan.format_of("model.language_model.layers.0.mlp.up_proj") + .unwrap() + .algo, + QuantAlgo::Fp8Block { bn: 128, bk: 128 } + ); + assert_eq!( + plan.format_of("model.language_model.layers.0.mlp.experts") + .unwrap() + .scale_encoding + .as_deref(), + Some("csf_32x2_v1") + ); +} + +/// The qwen4_exp registry covers exactly the modules its loaders run +/// natively — no more, no less. +#[test] +fn qwen4_exp_registry_coverage() { + let n = native_consumers("qwen4_exp"); + // MXFP8 natives: GDN + attention + shared expert + vision attention. + for m in [ + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.0.linear_attn.in_proj_z", + "model.language_model.layers.0.linear_attn.out_proj", + "model.language_model.layers.0.self_attn.q_proj", + "model.language_model.layers.0.self_attn.o_proj", + "model.language_model.layers.0.mlp.shared_expert.gate_proj", + "model.visual.blocks.3.attn.qkv", + "model.visual.blocks.3.attn.proj", + ] { + assert!(n.covers(m, &QuantAlgo::Mxfp8), "{m}"); + } + // MXFP8 NOT native (expands): indexer, dense MLP, visual MLP, MTP, misc. + for m in [ + "model.language_model.layers.0.self_attn.indexer.index_qk_proj", + "model.language_model.layers.0.mlp.gate_proj", + "model.visual.blocks.3.mlp.linear_fc1", + "mtp.layers.0.mlp.experts", + "model.language_model.layers.1.ple.ple_embedding.ngram_embedding", + ] { + assert!(!n.covers(m, &QuantAlgo::Mxfp8), "{m}"); + } + // NVFP4 natives: experts (LM + MTP), vision linears/mergers, PLE table. + for m in [ + "model.language_model.layers.0.mlp.experts", + "mtp.layers.0.mlp.experts", + "model.visual.blocks.0.mlp.linear_fc2", + "model.visual.merger.linear_fc1", + "model.language_model.visual.merger.linear_fc2", + "model.language_model.layers.1.ple.ple_embedding.ngram_embedding", + ] { + assert!(n.covers(m, &QuantAlgo::W4A16Nvfp4), "{m}"); + } + // Not covered: LM projections, lm_head. + for m in [ + "model.language_model.layers.0.self_attn.q_proj", + "lm_head", + ] { + assert!(!n.covers(m, &QuantAlgo::W4A16Nvfp4), "{m}"); + } + // An unknown model type covers nothing. + let none = native_consumers("some_new_model"); + assert!(!none.covers( + "model.language_model.layers.0.self_attn.q_proj", + &QuantAlgo::Mxfp8 + )); + assert!(!none.covers( + "model.language_model.layers.0.mlp.experts", + &QuantAlgo::Nvfp4 + )); +} + +/// No `quantized_layers` (or no quant config) → no plan, no FP4-Local. +#[test] +fn absent_map_means_no_plan() { + assert!(ModeloptQuantPlan::from_config(&ModelConfig::qwen3_next_80b_nvfp4()).is_none()); + let mut cfg = ModelConfig::qwen3_next_80b_nvfp4(); + cfg.quantization_config = Some(QuantizationConfig { + quant_method: "modelopt".into(), + quant_algo: "NVFP4".into(), + format: String::new(), + ignore_modules: vec![], + weight_block_size: vec![], + group_size: 16, + quantized_layers: BTreeMap::new(), + exl3: None, + }); + assert!(ModeloptQuantPlan::from_config(&cfg).is_none()); +} diff --git a/crates/spark-model/src/weight_map/nvfp4_detect/ep_detection_tests.rs b/crates/spark-model/src/weight_map/nvfp4_detect/ep_detection_tests.rs index 3a15b04fd..adfb565a2 100644 --- a/crates/spark-model/src/weight_map/nvfp4_detect/ep_detection_tests.rs +++ b/crates/spark-model/src/weight_map/nvfp4_detect/ep_detection_tests.rs @@ -104,6 +104,8 @@ fn mixed_precision_map_resolves_standard_despite_mtp_scale_inv() { QuantLayerSpec { quant_algo: "NVFP4".into(), group_size: 16, + weight_scale_encoding: None, + weight_block_size: vec![], }, ); } @@ -112,6 +114,8 @@ fn mixed_precision_map_resolves_standard_despite_mtp_scale_inv() { QuantLayerSpec { quant_algo: "FP8_PB_WO".into(), group_size: 128, + weight_scale_encoding: None, + weight_block_size: vec![], }, ); map.insert( @@ -119,6 +123,8 @@ fn mixed_precision_map_resolves_standard_despite_mtp_scale_inv() { QuantLayerSpec { quant_algo: "FP8".into(), group_size: 0, + weight_scale_encoding: None, + weight_block_size: vec![], }, ); cfg.quantization_config = Some(QuantizationConfig { @@ -165,6 +171,8 @@ fn mixed_precision_fp8_attn_nvfp4_mlp_resolves_standard() { QuantLayerSpec { quant_algo: "FP8".into(), group_size: 128, + weight_scale_encoding: None, + weight_block_size: vec![], }, ); } @@ -174,6 +182,8 @@ fn mixed_precision_fp8_attn_nvfp4_mlp_resolves_standard() { QuantLayerSpec { quant_algo: "NVFP4".into(), group_size: 16, + weight_scale_encoding: None, + weight_block_size: vec![], }, ); } @@ -217,6 +227,8 @@ fn mixed_precision_w4a16_nvfp4_labels_resolve_standard() { QuantLayerSpec { quant_algo: "FP8".into(), group_size: 128, + weight_scale_encoding: None, + weight_block_size: vec![], }, ); } @@ -225,6 +237,8 @@ fn mixed_precision_w4a16_nvfp4_labels_resolve_standard() { QuantLayerSpec { quant_algo: "W4A16_NVFP4".into(), group_size: 16, + weight_scale_encoding: None, + weight_block_size: vec![], }, ); for proj in ["gate_proj", "up_proj", "down_proj"] { @@ -233,6 +247,8 @@ fn mixed_precision_w4a16_nvfp4_labels_resolve_standard() { QuantLayerSpec { quant_algo: "W4A16_NVFP4".into(), group_size: 16, + weight_scale_encoding: None, + weight_block_size: vec![], }, ); } From fc351aa731a1565a2df5235c760047133f59d4a7 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Mon, 19 Oct 2026 02:30:21 -0400 Subject: [PATCH 16/26] docs: add the overnight A/B results and format-driven loading to the FP4-Local record Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../qwen3.8-flash-next-fp4-local.md | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/docs/checkpoints/qwen3.8-flash-next-fp4-local.md b/docs/checkpoints/qwen3.8-flash-next-fp4-local.md index 877cba38b..2ed024c4f 100644 --- a/docs/checkpoints/qwen3.8-flash-next-fp4-local.md +++ b/docs/checkpoints/qwen3.8-flash-next-fp4-local.md @@ -208,3 +208,25 @@ Scripts in `bench/qwen4_exp/fp4local/`: `control.py` (4-request output control), Environment note: the BFCL scorer's venv (`~/.atlas/artifacts/bfcl/venv`) needed `soundfile` (`qwen_agent` imports it) plus `libsndfile1`; without it every `bfcl-subset` run fails at scoring after generating all responses. + +## 10. Update — overnight results and format-driven loading + +**Same-binary quality A/B** (only the checkpoint differs; BFCL golden ST-995, no overrides; single run +each): FP4-Local **87.44 / 87.32** (run-1792340347555552942) vs NVFP4 pack 86.23 / 85.64 +(run-1792378038850732185). An observation, not a significance claim. + +**v2 profile (adds `ATLAS_QWEN4EXP_HC_FAST`):** decode-floor 33.0 tok/s (run-1792378753368339758), +quick 512/512 34.7 tok/s, SSM-poisoning 12/12 invariant, long-context recall 12/12 at 32K–245K. + +**Image:** `azeezish/atlas-gb10:a78d2fb` built from the measured commit (all targets, 3.0 GB), served by +sparkrun from the rendered recipe: coherent, needle PASS. Its temperature-0 answers are correct but +not byte-identical to the bare-metal single-target binary — not yet root-caused (multi-target kernel +selection or the in-container template path). Not pushed to any registry. + +**Format-driven loading:** `weight_map/modelopt_plan.rs` reads the checkpoint's declared per-module +formats (`quantization_config.quantized_layers`, or `hf_quant_config.json`) — NVFP4 / W4A16_NVFP4 / +MXFP8 / FP8 block, group size in either spelling, `weight_scale_encoding`. CSF scales decode only for +encoding `"csf"`; any other encoding fails at load naming module and encoding. A per-model +native-consumer registry says which roles run natively (qwen4_exp declares the set above); every other +declared module is expanded exactly to BF16, so a new FP4-Local checkpoint loads correctly before it +is optimized. FP4-Local and NVFP4 outputs are byte-identical to the previous build. From 171f532fa7e762396ae64f2b10b0b7b8d88ea458 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:57:24 -0400 Subject: [PATCH 17/26] speculative: let the shallowest MTP ladder rung pass the model's draft ceiling The default K-vs-batch ladder pinned every n<=4 verify at 3 drafts, so qwen4_exp with ATLAS_QWEN4EXP_MTP_DEPTH=4..7 (K=5..8, exact-lane) still speculated only 3 deep. The rung exists to step depth down with concurrency, not to cap below the ceiling -- the n<=4 rung now passes num_drafts (already min(--num-drafts, verify_max_drafts)) through, while the 8/16/32 rungs are unchanged. Models capped at 3 see no change. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/spark-model/src/speculative/ladder.rs | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/crates/spark-model/src/speculative/ladder.rs b/crates/spark-model/src/speculative/ladder.rs index 11f7d3ff9..582ee4b0d 100644 --- a/crates/spark-model/src/speculative/ladder.rs +++ b/crates/spark-model/src/speculative/ladder.rs @@ -198,7 +198,14 @@ fn mtp_ladder_steps() -> &'static [(usize, usize)] { // the measured fixed cost is 65.8 ms, only 33% of the 16:2 step. // Cutting c is the standing lever, and it is what would let a deeper // rung pay again. - let mut steps = parsed.unwrap_or_else(|| vec![(4, 3), (8, 3), (16, 1), (32, 1)]); + // First rung passes the ceiling through: `num_drafts` is already + // `min(--num-drafts, verify_max_drafts)`, and a model that serves + // K>4 exactly (qwen4_exp's ATLAS_QWEN4EXP_MTP_DEPTH=4..7) must not be + // held at 3 by a ladder entry written when 3 was every model's + // maximum. `usize::MAX` clamps to `num_drafts` below, so models + // capped at 3 keep 3. The 8/16/32 rungs still step down with + // concurrency — the ladder's job is the wide end, not the ceiling. + let mut steps = parsed.unwrap_or_else(|| vec![(4, usize::MAX), (8, 3), (16, 1), (32, 1)]); steps.sort_by_key(|&(n, _)| n); steps }) @@ -286,6 +293,13 @@ mod tests { fn default_ladder_holds_depth_to_the_cap() { assert_eq!(mtp_ladder_drafts(1, 3), 3); assert_eq!(mtp_ladder_drafts(4, 3), 3); + // The first rung is a passthrough: a ceiling above 3 (qwen4_exp + // MTP_DEPTH 4..7) reaches the verify at n<=4, and the ladder only + // steps down at 8/16/32. + assert_eq!(mtp_ladder_drafts(1, 5), 5); + assert_eq!(mtp_ladder_drafts(4, 5), 5); + assert_eq!(mtp_ladder_drafts(8, 5), 3); + assert_eq!(mtp_ladder_drafts(16, 5), 1); assert_eq!(mtp_ladder_drafts(5, 3), 3); assert_eq!(mtp_ladder_drafts(8, 3), 3); // The 16:1 rung (wave 19): ONE draft (K=2, 2 rows/seq) through n=16. From 40a01d3c9ad34e73c61c6dae9507d6fbe3132a76 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:01:26 -0400 Subject: [PATCH 18/26] moe: stage transposed NVFP4 tables through smem in the decode GEMVs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The single-row gate_up and silu_down kernels issue one byte-wide load per thread per k step and sit at ~100 GB/s under a cold expert set (~45 GB of per-layer pointer tables, 216 GB/s ceiling). Both impls now double-buffer a 64-row packed tile plus its scale rows through shared memory via cooperative 16 B loads, then accumulate each thread's own output column in the identical FP32 order — byte-identical outputs, 1.1-1.26x on the cold bench, +3.6% decode-floor. The _e8m0 (GS=32) variants keep the scalar impl. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../gb10/common/moe_shared_expert_fused_t.cu | 267 +++++++++++++++++- 1 file changed, 265 insertions(+), 2 deletions(-) diff --git a/kernels/gb10/common/moe_shared_expert_fused_t.cu b/kernels/gb10/common/moe_shared_expert_fused_t.cu index 1740babc2..1903fb564 100644 --- a/kernels/gb10/common/moe_shared_expert_fused_t.cu +++ b/kernels/gb10/common/moe_shared_expert_fused_t.cu @@ -183,6 +183,28 @@ __device__ __forceinline__ void gate_up_shared_t_impl( } } +template +__device__ __forceinline__ void gate_up_shared_t_impl_smem( + const __nv_bfloat16* __restrict__ A, + const unsigned long long* __restrict__ gate_packed_t_ptrs, + const unsigned long long* __restrict__ gate_scale_t_ptrs, + const float* __restrict__ gate_scale2_vals, + __nv_bfloat16* __restrict__ gate_out, + const unsigned long long* __restrict__ up_packed_t_ptrs, + const unsigned long long* __restrict__ up_scale_t_ptrs, + const float* __restrict__ up_scale2_vals, + __nv_bfloat16* __restrict__ up_out, + const unsigned int* __restrict__ expert_indices, + const unsigned char* __restrict__ sh_gate_t_packed, + const unsigned char* __restrict__ sh_gate_t_scale, + float sh_gate_s2, + __nv_bfloat16* __restrict__ sh_gate_out, + const unsigned char* __restrict__ sh_up_t_packed, + const unsigned char* __restrict__ sh_up_t_scale, + float sh_up_s2, + __nv_bfloat16* __restrict__ sh_up_out, + unsigned int N, unsigned int K, unsigned int top_k); + // NVFP4 (default): FP8-E4M3 per-16 scales × per-tensor global. extern "C" __global__ void moe_expert_gate_up_shared_t( const __nv_bfloat16* __restrict__ A, @@ -205,7 +227,7 @@ extern "C" __global__ void moe_expert_gate_up_shared_t( __nv_bfloat16* __restrict__ sh_up_out, unsigned int N, unsigned int K, unsigned int top_k ) { - gate_up_shared_t_impl( + gate_up_shared_t_impl_smem( A, gate_packed_t_ptrs, gate_scale_t_ptrs, gate_scale2_vals, gate_out, up_packed_t_ptrs, up_scale_t_ptrs, up_scale2_vals, up_out, expert_indices, sh_gate_t_packed, sh_gate_t_scale, sh_gate_s2, sh_gate_out, @@ -368,6 +390,23 @@ __device__ __forceinline__ void silu_down_shared_t_impl( } } +template +__device__ __forceinline__ void silu_down_shared_t_impl_smem( + const __nv_bfloat16* __restrict__ gate_out, + const __nv_bfloat16* __restrict__ up_out, + const unsigned long long* __restrict__ packed_t_ptrs, + const unsigned long long* __restrict__ scale_t_ptrs, + const float* __restrict__ scale2_vals, + __nv_bfloat16* __restrict__ C, + const unsigned int* __restrict__ expert_indices, + const __nv_bfloat16* __restrict__ sh_gate_in, + const __nv_bfloat16* __restrict__ sh_up_in, + const unsigned char* __restrict__ sh_down_t_packed, + const unsigned char* __restrict__ sh_down_t_scale, + float sh_down_s2, + __nv_bfloat16* __restrict__ sh_down_out, + unsigned int N, unsigned int K, unsigned int top_k); + // NVFP4 (default): FP8-E4M3 per-16 scales × per-tensor global. extern "C" __global__ void moe_expert_silu_down_shared_t( const __nv_bfloat16* __restrict__ gate_out, @@ -385,7 +424,7 @@ extern "C" __global__ void moe_expert_silu_down_shared_t( __nv_bfloat16* __restrict__ sh_down_out, unsigned int N, unsigned int K, unsigned int top_k ) { - silu_down_shared_t_impl( + silu_down_shared_t_impl_smem( gate_out, up_out, packed_t_ptrs, scale_t_ptrs, scale2_vals, C, expert_indices, sh_gate_in, sh_up_in, sh_down_t_packed, sh_down_t_scale, sh_down_s2, sh_down_out, N, K, top_k); @@ -413,3 +452,227 @@ extern "C" __global__ void moe_expert_silu_down_shared_t_e8m0( expert_indices, sh_gate_in, sh_up_in, sh_down_t_packed, sh_down_t_scale, sh_down_s2, sh_down_out, N, K, top_k); } + +// ── Shared-memory tiled impls (S2d): production path ───────────────────── +// Stage each expert's packed tile into smem via 16 B cooperative loads +// (double-buffered), then every thread accumulates ITS column in the +// identical per-output FP32 order. Byte-identical to the scalar loop. +#define TILE_KH 64u +#define TILE_STRIDE 48u +template +__device__ __forceinline__ void gate_up_shared_t_impl_smem( + const __nv_bfloat16* __restrict__ A, + const unsigned long long* __restrict__ gate_packed_t_ptrs, + const unsigned long long* __restrict__ gate_scale_t_ptrs, + const float* __restrict__ gate_scale2_vals, + __nv_bfloat16* __restrict__ gate_out, + const unsigned long long* __restrict__ up_packed_t_ptrs, + const unsigned long long* __restrict__ up_scale_t_ptrs, + const float* __restrict__ up_scale2_vals, + __nv_bfloat16* __restrict__ up_out, + const unsigned int* __restrict__ expert_indices, + const unsigned char* __restrict__ sh_gate_t_packed, + const unsigned char* __restrict__ sh_gate_t_scale, + float sh_gate_s2, + __nv_bfloat16* __restrict__ sh_gate_out, + const unsigned char* __restrict__ sh_up_t_packed, + const unsigned char* __restrict__ sh_up_t_scale, + float sh_up_s2, + __nv_bfloat16* __restrict__ sh_up_out, + unsigned int N, unsigned int K, unsigned int top_k +) { + const unsigned int expert_slot = blockIdx.y; + const unsigned int proj = blockIdx.z; + const bool is_shared = (expert_slot == top_k); + const unsigned char* B_packed; const unsigned char* B_scale; float s2; __nv_bfloat16* C; + const unsigned int n = blockIdx.x * blockDim.x + threadIdx.x; + if (is_shared) { + if (proj == 0) { + if (sh_gate_t_packed == 0) { if (n < N) sh_gate_out[n] = __float2bfloat16(0.0f); return; } + B_packed = sh_gate_t_packed; B_scale = sh_gate_t_scale; s2 = sh_gate_s2; C = sh_gate_out; + } else { + if (sh_up_t_packed == 0) { if (n < N) sh_up_out[n] = __float2bfloat16(0.0f); return; } + B_packed = sh_up_t_packed; B_scale = sh_up_t_scale; s2 = sh_up_s2; C = sh_up_out; + } + } else { + const unsigned int expert_id = expert_indices[expert_slot]; + if (proj == 0) { + B_packed = (const unsigned char*)gate_packed_t_ptrs[expert_id]; + B_scale = (const unsigned char*)gate_scale_t_ptrs[expert_id]; + s2 = gate_scale2_vals[expert_id]; C = gate_out; + } else { + B_packed = (const unsigned char*)up_packed_t_ptrs[expert_id]; + B_scale = (const unsigned char*)up_scale_t_ptrs[expert_id]; + s2 = up_scale2_vals[expert_id]; C = up_out; + } + if (B_packed == 0) { if (n < N) C[(unsigned long long)expert_slot * N + n] = __float2bfloat16(0.0f); return; } + } + + const unsigned int lane = threadIdx.x; + const unsigned int ns0 = blockIdx.x * blockDim.x; + const bool valid = (n < N); + __shared__ float s_lut[16]; + if (lane < 16) s_lut[lane] = E2M1_LUT_T[lane]; + __shared__ unsigned char s_pk[2][TILE_KH * TILE_STRIDE]; + __shared__ unsigned char s_sc[2][(TILE_KH / 8) * TILE_STRIDE]; + __syncthreads(); + + const unsigned int r = lane / 2u; + const unsigned int off = (lane % 2u) * 16u; + auto stage = [&](unsigned int tg, unsigned int buf) { + const unsigned int r0 = tg * TILE_KH; + _Pragma("unroll") + for (unsigned int w = 0; w < TILE_KH / 16u; w++) { + unsigned int rr = r + w * 16u; + *(uint4*)&s_pk[buf][rr * TILE_STRIDE + off] = + *(const uint4*)&B_packed[(unsigned long long)(r0 + rr) * N + ns0 + off]; + } + const unsigned int ng = TILE_KH / 8u; + if (lane < ng * 2u) { + unsigned int sg = lane / 2u; + unsigned int off2 = (lane % 2u) * 16u; + *(uint4*)&s_sc[buf][sg * TILE_STRIDE + off2] = + *(const uint4*)&B_scale[(unsigned long long)(tg * ng + sg) * N + ns0 + off2]; + } + }; + + float acc = 0.0f; + const unsigned int ntg = (K / 2u) / TILE_KH; + stage(0, 0); + for (unsigned int tg = 0; tg < ntg; tg++) { + __syncthreads(); + if (tg + 1u < ntg) stage(tg + 1u, (tg + 1u) & 1u); + const unsigned int buf = tg & 1u; + if (valid) { + #define GU_TILE_ACC(GS_, E8M0_) do { \ + if constexpr ((GS_) == 16) { \ + _Pragma("unroll 8") \ + for (unsigned int rr = 0; rr < TILE_KH; rr++) { \ + unsigned int k_half = tg * TILE_KH + rr; \ + float sc = mx_block_scale<(E8M0_)>(s_sc[buf][(rr >> 3) * TILE_STRIDE + lane], s2); \ + unsigned char byte = s_pk[buf][rr * TILE_STRIDE + lane]; \ + float a_lo = __bfloat162float(A[k_half * 2]); \ + float a_hi = __bfloat162float(A[k_half * 2 + 1]); \ + acc += a_lo * (s_lut[byte & 0xFu] * sc) + a_hi * (s_lut[(byte >> 4) & 0xFu] * sc); \ + } \ + } \ + } while(0) + if constexpr (GS_R == GS_S && E8M0_R == E8M0_S) { + GU_TILE_ACC(GS_R, E8M0_R); + } else { + if (is_shared) { GU_TILE_ACC(GS_S, E8M0_S); } + else { GU_TILE_ACC(GS_R, E8M0_R); } + } + #undef GU_TILE_ACC + } + } + __syncthreads(); + if (!valid) return; + if (is_shared) C[n] = __float2bfloat16(acc); + else C[(unsigned long long)expert_slot * N + n] = __float2bfloat16(acc); +} + +template +__device__ __forceinline__ void silu_down_shared_t_impl_smem( + const __nv_bfloat16* __restrict__ gate_out, + const __nv_bfloat16* __restrict__ up_out, + const unsigned long long* __restrict__ packed_t_ptrs, + const unsigned long long* __restrict__ scale_t_ptrs, + const float* __restrict__ scale2_vals, + __nv_bfloat16* __restrict__ C, + const unsigned int* __restrict__ expert_indices, + const __nv_bfloat16* __restrict__ sh_gate_in, + const __nv_bfloat16* __restrict__ sh_up_in, + const unsigned char* __restrict__ sh_down_t_packed, + const unsigned char* __restrict__ sh_down_t_scale, + float sh_down_s2, + __nv_bfloat16* __restrict__ sh_down_out, + unsigned int N, unsigned int K, unsigned int top_k +) { + const unsigned int expert_slot = blockIdx.y; + const bool is_shared = (expert_slot == top_k); + const unsigned char* B_packed; const unsigned char* B_scale; float s2; + const __nv_bfloat16* g_ptr; const __nv_bfloat16* u_ptr; + const unsigned int n = blockIdx.x * blockDim.x + threadIdx.x; + if (is_shared) { + if (sh_down_t_packed == 0) { if (n < N) sh_down_out[n] = __float2bfloat16(0.0f); return; } + B_packed = sh_down_t_packed; B_scale = sh_down_t_scale; s2 = sh_down_s2; + g_ptr = sh_gate_in; u_ptr = sh_up_in; + } else { + const unsigned int expert_id = expert_indices[expert_slot]; + B_packed = (const unsigned char*)packed_t_ptrs[expert_id]; + B_scale = (const unsigned char*)scale_t_ptrs[expert_id]; + s2 = scale2_vals[expert_id]; + g_ptr = gate_out + (unsigned long long)expert_slot * K; + u_ptr = up_out + (unsigned long long)expert_slot * K; + if (B_packed == 0) { if (n < N) C[expert_slot * N + n] = __float2bfloat16(0.0f); return; } + } + + const unsigned int lane = threadIdx.x; + const unsigned int ns0 = blockIdx.x * blockDim.x; + const bool valid = (n < N); + extern __shared__ float s_act[]; + for (unsigned int i = lane; i < K; i += blockDim.x) { + float gf = __bfloat162float(g_ptr[i]); + float uf = __bfloat162float(u_ptr[i]); + s_act[i] = (gf / (1.0f + __expf(-gf))) * uf; + } + __shared__ float s_lut[16]; + if (lane < 16) s_lut[lane] = E2M1_LUT_T[lane]; + __shared__ unsigned char s_pk[2][TILE_KH * TILE_STRIDE]; + __shared__ unsigned char s_sc[2][(TILE_KH / 8) * TILE_STRIDE]; + __syncthreads(); + + const unsigned int r = lane / 2u; + const unsigned int off = (lane % 2u) * 16u; + auto stage = [&](unsigned int tg, unsigned int buf) { + const unsigned int r0 = tg * TILE_KH; + _Pragma("unroll") + for (unsigned int w = 0; w < TILE_KH / 16u; w++) { + unsigned int rr = r + w * 16u; + *(uint4*)&s_pk[buf][rr * TILE_STRIDE + off] = + *(const uint4*)&B_packed[(unsigned long long)(r0 + rr) * N + ns0 + off]; + } + const unsigned int ng = TILE_KH / 8u; + if (lane < ng * 2u) { + unsigned int sg = lane / 2u; + unsigned int off2 = (lane % 2u) * 16u; + *(uint4*)&s_sc[buf][sg * TILE_STRIDE + off2] = + *(const uint4*)&B_scale[(unsigned long long)(tg * ng + sg) * N + ns0 + off2]; + } + }; + + float acc = 0.0f; + const unsigned int ntg = (K / 2u) / TILE_KH; + stage(0, 0); + for (unsigned int tg = 0; tg < ntg; tg++) { + __syncthreads(); + if (tg + 1u < ntg) stage(tg + 1u, (tg + 1u) & 1u); + const unsigned int buf = tg & 1u; + if (valid) { + #define SD_TILE_ACC(GS_, E8M0_) do { \ + if constexpr ((GS_) == 16) { \ + _Pragma("unroll 8") \ + for (unsigned int rr = 0; rr < TILE_KH; rr++) { \ + unsigned int k_half = tg * TILE_KH + rr; \ + float sc = mx_block_scale<(E8M0_)>(s_sc[buf][(rr >> 3) * TILE_STRIDE + lane], s2); \ + unsigned char byte = s_pk[buf][rr * TILE_STRIDE + lane]; \ + acc += s_act[k_half * 2] * (s_lut[byte & 0xFu] * sc) \ + + s_act[k_half * 2 + 1] * (s_lut[(byte >> 4) & 0xFu] * sc); \ + } \ + } \ + } while(0) + if constexpr (GS_R == GS_S && E8M0_R == E8M0_S) { + SD_TILE_ACC(GS_R, E8M0_R); + } else { + if (is_shared) { SD_TILE_ACC(GS_S, E8M0_S); } + else { SD_TILE_ACC(GS_R, E8M0_R); } + } + #undef SD_TILE_ACC + } + } + __syncthreads(); + if (!valid) return; + if (is_shared) sh_down_out[n] = __float2bfloat16(acc); + else C[(unsigned long long)expert_slot * N + n] = __float2bfloat16(acc); +} From c24ea262721849fc13be825f2bd6ca3781f8f1ba Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:34:41 -0400 Subject: [PATCH 19/26] ngram_cache: raise the PLE fault pool to 64 pread threads On a cold ~17K prefill the fault wall is ~450K O_DIRECT row reads; at 64 in-flight preads the pool reaches the GB10 NVMe's ~66K IOPS random-4K ceiling (fio-measured), cutting the fault wall ~20% and agent TTFT ~6% (15.4s -> 14.5s). Same rows, same bytes, same slot/pin semantics -- only deeper queue depth. ATLAS_PLE_FAULT_THREADS still overrides. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/spark-storage/src/ngram_cache_fault.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/spark-storage/src/ngram_cache_fault.rs b/crates/spark-storage/src/ngram_cache_fault.rs index ffd19abd7..383919178 100644 --- a/crates/spark-storage/src/ngram_cache_fault.rs +++ b/crates/spark-storage/src/ngram_cache_fault.rs @@ -23,6 +23,8 @@ use super::*; /// 8 424 ms 3.85x /// 16 257 ms 6.35x /// 32 171 ms 9.55x +/// 64 ~64K IOPS (GB10 NVMe's random-4K ceiling; +64 lifts the +/// PLE fault wall ~20% on a cold ~17K prefill) /// ``` /// /// `1` restores the old strictly-serial behaviour for a bisect. @@ -30,7 +32,7 @@ pub(super) fn fault_threads() -> usize { std::env::var("ATLAS_PLE_FAULT_THREADS") .ok() .and_then(|v| v.parse().ok()) - .unwrap_or(32) + .unwrap_or(64) } /// One scheduled miss: read row `id` into `slot`. Produced by the decision From 5e18f771174d1914560872548d54de9941327287 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:44:28 -0400 Subject: [PATCH 20/26] ngram_cache: resident PLE block scales, opt-in ATLAS_PLE_SCALES_RESIDENT Cold faulting reads two IOPS-bound O_DIRECT reads per miss: the packed row and its E4M3 block-scale row (~2.03 reads/miss measured cold). With the flag on, attach_segmented_block_scales reads every shard's scale tensor into one flat host arena in shard order at boot (~3.2 GB for the 128-shard NVFP4 table, loaded shard-parallel); fetch_row then memcpys arena[id * stride..] using the same (shard, local) split as scale_loc, leaving packed rows on disk. ~1.02 reads/miss, fault wall 7.5s->5.3s and agent TTFT 14.4->13.7s on a cold 17K turn. Flag off is byte-identical; Segments/files are built the same either way. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/spark-storage/src/ngram_cache.rs | 5 +++ .../src/ngram_cache/fault_pool.rs | 32 +++++++++------ .../src/ngram_cache/fault_pool_tests.rs | 30 ++++++++++++++ .../spark-storage/src/ngram_cache/resolve.rs | 5 ++- .../src/ngram_cache/segmented.rs | 40 +++++++++++++++++++ crates/spark-storage/src/ngram_cache_fault.rs | 33 +++++++++++++-- 6 files changed, 128 insertions(+), 17 deletions(-) diff --git a/crates/spark-storage/src/ngram_cache.rs b/crates/spark-storage/src/ngram_cache.rs index af8946bcb..2c6fd786b 100644 --- a/crates/spark-storage/src/ngram_cache.rs +++ b/crates/spark-storage/src/ngram_cache.rs @@ -136,6 +136,10 @@ struct ScaleCache { /// Bytes of scale payload per slot (4 for f32, `dim/group` E4M3 for NVFP4). stride: usize, src: ScaleSrc, + /// `ATLAS_PLE_SCALES_RESIDENT`: the whole scale table resident in host + /// RAM, concatenated in shard order. Set only for a segmented src; a + /// fault memcpys `arena[id * stride..]` instead of a second disk read. + resident: Option>>, } /// Where a scale row comes from on a fault. @@ -229,6 +233,7 @@ impl NgramRowCache { .context("NgramRowCache: scale arena")?, stride: 4, src: ScaleSrc::PerRowFile(open_direct(sp)?), + resident: None, }) } None => None, diff --git a/crates/spark-storage/src/ngram_cache/fault_pool.rs b/crates/spark-storage/src/ngram_cache/fault_pool.rs index e0ba23a96..ba6056273 100644 --- a/crates/spark-storage/src/ngram_cache/fault_pool.rs +++ b/crates/spark-storage/src/ngram_cache/fault_pool.rs @@ -44,14 +44,17 @@ struct OwnedSource { /// Owned mirror of [`ScaleSource`] (cloned descriptors for 'static workers). enum OwnedScaleSrc { Flat(File), - Segmented(Segments), + Segmented(Segments, Option>>), } impl OwnedScaleSrc { fn view(&self) -> ScaleSource<'_> { match self { OwnedScaleSrc::Flat(f) => ScaleSource::Flat { file: f }, - OwnedScaleSrc::Segmented(s) => ScaleSource::Segmented { seg: s }, + OwnedScaleSrc::Segmented(s, r) => ScaleSource::Segmented { + seg: s, + resident: r.as_ref().map(|v| v.as_slice()), + }, } } } @@ -193,17 +196,20 @@ impl FaultPool { f.try_clone() .context("NgramRowCache fault pool: dup scales")?, )), - Some(ScaleSrc::Segmented(s)) => Some(OwnedScaleSrc::Segmented(Segments { - bases: s.bases.clone(), - rows_per: s.rows_per, - files: s - .files - .iter() - .map(File::try_clone) - .collect::>() - .context("NgramRowCache fault pool: dup scale files")?, - shard_file: s.shard_file.clone(), - })), + Some(ScaleSrc::Segmented(s)) => Some(OwnedScaleSrc::Segmented( + Segments { + bases: s.bases.clone(), + rows_per: s.rows_per, + files: s + .files + .iter() + .map(File::try_clone) + .collect::>() + .context("NgramRowCache fault pool: dup scale files")?, + shard_file: s.shard_file.clone(), + }, + cache.scales.as_ref().and_then(|sc| sc.resident.clone()), + )), }; let src = OwnedSource { file: cache diff --git a/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs b/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs index 7ce97630f..feed2b096 100644 --- a/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs +++ b/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs @@ -45,6 +45,36 @@ fn faults(n: u32, panic_at: Option) -> Vec { .collect() } +/// `ATLAS_PLE_SCALES_RESIDENT`: the arena's byte offset for a global row id +/// is the same (shard, local) split `scale_loc` computes — shard 0, a shard +/// boundary, and the last row must resolve identically. +#[test] +fn resident_scales_index_by_global_row_id() { + let stride = 2usize; + let file = File::open(std::env::current_exe().expect("test exe")).expect("open test exe"); + let seg = crate::ngram_cache::Segments { + bases: vec![0, 100, 200], + rows_per: 4, + files: vec![file], + shard_file: vec![0, 0, 0], + }; + let arena: Vec = (0..12u8).flat_map(|r| [r, r.wrapping_add(100)]).collect(); + let ss = super::ScaleSource::Segmented { + seg: &seg, + resident: Some(&arena), + }; + for id in [0u64, 3, 4, 7, 8, 11] { + let sl = ss + .scale_resident_slice(id, stride) + .expect("resident slice"); + assert_eq!( + sl, + &arena[id as usize * stride..id as usize * stride + stride], + "row {id} resolved the wrong arena bytes" + ); + } +} + /// Run `f` on its own thread, so a pool hang fails the test instead of /// wedging the suite. fn within_deadline(f: impl FnOnce() + Send + 'static) { diff --git a/crates/spark-storage/src/ngram_cache/resolve.rs b/crates/spark-storage/src/ngram_cache/resolve.rs index 30d10ed6d..6a7b257a5 100644 --- a/crates/spark-storage/src/ngram_cache/resolve.rs +++ b/crates/spark-storage/src/ngram_cache/resolve.rs @@ -174,7 +174,10 @@ impl NgramRowCache { scale_src: self.scales.as_ref().and_then(|sc| match &sc.src { super::ScaleSrc::Constant => None, super::ScaleSrc::PerRowFile(f) => Some(ScaleSource::Flat { file: f }), - super::ScaleSrc::Segmented(seg) => Some(ScaleSource::Segmented { seg }), + super::ScaleSrc::Segmented(seg) => Some(ScaleSource::Segmented { + seg, + resident: sc.resident.as_ref().map(|v| v.as_slice()), + }), }), }; let ptrs = ArenaPtrs { diff --git a/crates/spark-storage/src/ngram_cache/segmented.rs b/crates/spark-storage/src/ngram_cache/segmented.rs index 59cb0c003..c7400efc3 100644 --- a/crates/spark-storage/src/ngram_cache/segmented.rs +++ b/crates/spark-storage/src/ngram_cache/segmented.rs @@ -92,6 +92,7 @@ impl NgramRowCache { arena, stride: 4, src: ScaleSrc::Constant, + resident: None, }); Ok(()) } @@ -128,6 +129,44 @@ impl NgramRowCache { } let sbytes = self.slots * scale_stride; let sblocks = sbytes.div_ceil(BLOCK); + // ATLAS_PLE_SCALES_RESIDENT=1: read every shard's scale tensor into + // one host arena (shard order, flat `id * stride` indexing); faults + // then memcpy scales instead of a second disk read. Hard error on a + // missing/short shard. + let resident_env = std::env::var("ATLAS_PLE_SCALES_RESIDENT") + .map(|v| v == "1" || v.eq_ignore_ascii_case("true")) + .unwrap_or(false); + let resident = if resident_env { + let shard_bytes = rows_per as usize * scale_stride; + let mut arena = vec![0u8; shards.len() * shard_bytes]; + std::thread::scope(|sc| { + let mut hs = Vec::new(); + for (dst, (path, base)) in arena.chunks_mut(shard_bytes).zip(shards.iter()) { + let (path, base) = (path.clone(), *base); + hs.push(sc.spawn(move || -> Result<()> { + use std::os::unix::fs::FileExt; + std::fs::File::open(&path) + .with_context(|| format!("scale shard {path:?}"))? + .read_exact_at(dst, base) + .with_context(|| format!("scale shard {path:?} short at {base}")) + })); + } + for h in hs { + h.join().expect("scale loader panic")?; + } + Ok::<(), anyhow::Error>(()) + })?; + tracing::info!( + "PLE block scales resident: {:.1} GB ({} shards x {} rows x {} B)", + arena.len() as f64 / 1e9, + shards.len(), + rows_per, + scale_stride + ); + Some(std::sync::Arc::new(arena)) + } else { + None + }; self.scales = Some(ScaleCache { arena: ExpertArena::new(1, sblocks as u32, BLOCK) .context("NgramRowCache: block-scale arena")?, @@ -138,6 +177,7 @@ impl NgramRowCache { files, shard_file, }), + resident, }); Ok(()) } diff --git a/crates/spark-storage/src/ngram_cache_fault.rs b/crates/spark-storage/src/ngram_cache_fault.rs index 383919178..d22c9a4f4 100644 --- a/crates/spark-storage/src/ngram_cache_fault.rs +++ b/crates/spark-storage/src/ngram_cache_fault.rs @@ -76,14 +76,33 @@ pub(super) struct RowSource<'a> { /// scale tensors (same `rows_per`, same deduped files). pub(super) enum ScaleSource<'a> { Flat { file: &'a File }, - Segmented { seg: &'a Segments }, + /// `resident`: `Some` under `ATLAS_PLE_SCALES_RESIDENT` — the same + /// (shard, local) split as `scale_loc`, answered from host RAM. + Segmented { + seg: &'a Segments, + resident: Option<&'a [u8]>, + }, } -impl ScaleSource<'_> { +impl<'a> ScaleSource<'a> { + /// The resident arena's byte offset for `id` is the same (shard, local) + /// split `scale_loc` uses — one mapping, two sinks. + pub(super) fn scale_resident_slice(&self, id: u64, stride: usize) -> Option<&'a [u8]> { + match self { + ScaleSource::Segmented { seg, resident } => { + let shard = id / seg.rows_per; + let local = id % seg.rows_per; + let at = (shard * seg.rows_per + local) as usize * stride; + resident.map(|a| &a[at..at + stride]) + } + _ => None, + } + } + fn scale_loc(&self, id: u64, stride: u64) -> (&File, u64) { match self { ScaleSource::Flat { file } => (file, id * stride), - ScaleSource::Segmented { seg } => { + ScaleSource::Segmented { seg, resident: _ } => { let shard = (id / seg.rows_per) as usize; let local = id % seg.rows_per; let file = &seg.files[seg.shard_file[shard] as usize]; @@ -158,6 +177,14 @@ pub(super) fn fetch_row( // from the installed source (4-byte f32 or `dim/group` E4M3). if let (Some(ssrc), Some(sbase)) = (&src.scale_src, ptrs.scales) { let stride = src.scale_stride; + if let Some(bytes) = ssrc.scale_resident_slice(id, stride) { + // Resident arena: a memcpy, no disk read. + let sdst = unsafe { + std::slice::from_raw_parts_mut(sbase.add(slot as usize * stride), stride) + }; + sdst.copy_from_slice(bytes); + return Ok(()); + } let (sfile, sbyte) = ssrc.scale_loc(id, stride as u64); let sblock = sbyte - (sbyte % BLOCK as u64); let swithin = (sbyte - sblock) as usize; From 3fd1e424208425bb650883b9c21d54c5f1057406 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Mon, 19 Oct 2026 20:52:52 -0400 Subject: [PATCH 21/26] tests: pin tiered SSM pool sizing to the passthrough ladder rung The 171f532fa shallow rung passes the model's draft ceiling through at n<=4, so slot 0 legitimately sizes at the full K-1 when num_drafts=7. The assertion encoded the retired cap. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/spark-model/src/model/ssm_batched_copy_tests.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/spark-model/src/model/ssm_batched_copy_tests.rs b/crates/spark-model/src/model/ssm_batched_copy_tests.rs index b1bb19112..9a4281a8e 100644 --- a/crates/spark-model/src/model/ssm_batched_copy_tests.rs +++ b/crates/spark-model/src/model/ssm_batched_copy_tests.rs @@ -406,6 +406,10 @@ fn dflash_pools_size_h_intermediates_uniformly() { // Counterfactual: identical geometry without the capture layers keeps // the tiered ladder — proves the assertion above discriminates, and // pins the MTP dummy's always-full-width invariant at the same time. + // Since 171f532fa the shallowest ladder rung passes the model's draft + // ceiling through (n<=4 -> num_drafts), so slot 0 — which can host a + // full K=8 verify — sizes at the full K-1; the 16-rung step-down still + // tiers slot 15 to 1. let tiered = SsmStatePool::new( &tiny_config(), 16, @@ -418,7 +422,7 @@ fn dflash_pools_size_h_intermediates_uniformly() { &gpu, ) .unwrap(); - assert_eq!(tiered.h_inter_count(0), 3); + assert_eq!(tiered.h_inter_count(0), 7); assert_eq!(tiered.h_inter_count(15), 1); assert_eq!(tiered.h_inter_count(tiered.mtp_slots), 7); } From 82120b5d1e07e9f9590b1fe9b9749e05a6b3f189 Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Mon, 19 Oct 2026 20:52:52 -0400 Subject: [PATCH 22/26] modelopt: fix stale rustdoc link and reformat under current stable rustfmt The format-driven load commit removed `native_mxfp8`; the module doc now links `native_consumers`. Includes cargo +stable fmt output. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/weight_map/modelopt_local.rs | 7 +++--- .../src/weight_map/modelopt_local_tests.rs | 22 ++++++++----------- .../src/weight_map/modelopt_plan.rs | 4 +--- .../src/weight_map/modelopt_plan_tests.rs | 14 +++++++----- 4 files changed, 21 insertions(+), 26 deletions(-) diff --git a/crates/spark-model/src/weight_map/modelopt_local.rs b/crates/spark-model/src/weight_map/modelopt_local.rs index 38189d50d..72bbd8e0a 100644 --- a/crates/spark-model/src/weight_map/modelopt_local.rs +++ b/crates/spark-model/src/weight_map/modelopt_local.rs @@ -10,7 +10,7 @@ //! (`weight_scale.nvfp4_csf_{fixed,exceptions}`) — decoded losslessly into //! the plain E4M3 `weight_scale` every NVFP4 kernel reads; //! * GDN / attention / QSA-indexer / shared-expert projections are MXFP8. -//! Those with a native MXFP8 consumer ([`native_mxfp8`]: the GDN and +//! Those with a native MXFP8 consumer ([`native_consumers`]: the GDN and //! attention projections) are uploaded byte for byte, `weight` E4M3 + `weight_scale` //! E8M0, and read by `mxfp8_proj.cu`; the rest are expanded exactly to the //! BF16 `weight` the NVFP4 pack ships; @@ -32,9 +32,9 @@ use std::collections::HashSet; +use super::modelopt_plan::{ModeloptQuantPlan, QuantAlgo, native_consumers}; use anyhow::{Context, Result, ensure}; use atlas_core::config::ModelConfig; -use super::modelopt_plan::{ModeloptQuantPlan, QuantAlgo, native_consumers}; use atlas_core::modelopt_codec::{decode_csf_scales, mxfp8_to_bf16, nvfp4_rows_to_bf16}; use spark_runtime::gpu::GpuBackend; use spark_runtime::weights::{WeightDtype, WeightStore, WeightTensor}; @@ -326,8 +326,7 @@ pub fn materialize( let mut modules: Vec = host_decode_modules(config).into_iter().collect(); modules.sort(); // A/B only: expand the native MXFP8 set too (both paths are exact). - let force_expand = - std::env::var("ATLAS_FP4LOCAL_EXPAND_MXFP8").as_deref() == Ok("1"); + let force_expand = std::env::var("ATLAS_FP4LOCAL_EXPAND_MXFP8").as_deref() == Ok("1"); for module in modules { let fmt = plan.as_ref().and_then(|p| p.format_of(&module)); let algo = fmt diff --git a/crates/spark-model/src/weight_map/modelopt_local_tests.rs b/crates/spark-model/src/weight_map/modelopt_local_tests.rs index 5b3c0b67e..05921f67f 100644 --- a/crates/spark-model/src/weight_map/modelopt_local_tests.rs +++ b/crates/spark-model/src/weight_map/modelopt_local_tests.rs @@ -28,14 +28,11 @@ fn fp4_local_config() -> ModelConfig { let mut cfg = ModelConfig::qwen3_next_80b_nvfp4(); cfg.model_type = "qwen4_exp".into(); let mut map = std::collections::BTreeMap::new(); - map.insert( - "model.language_model.layers.0.mlp.experts".into(), - { - let mut s = spec("NVFP4", 16); - s.weight_scale_encoding = Some("csf".into()); - s - }, - ); + map.insert("model.language_model.layers.0.mlp.experts".into(), { + let mut s = spec("NVFP4", 16); + s.weight_scale_encoding = Some("csf".into()); + s + }); map.insert(SHARED.into(), spec("MXFP8", 32)); map.insert(GDN_OUT.into(), spec("MXFP8", 32)); map.insert(VIS.into(), spec("W4A16_NVFP4", 16)); @@ -306,7 +303,9 @@ fn gdn_mxfp8_stays_native() { assert!(mx(GDN_OUT)); assert!(mx("model.language_model.layers.7.linear_attn.in_proj_qkv")); assert!(mx(ATTN)); - assert!(!mx("model.language_model.layers.3.self_attn.indexer.index_qk_proj")); + assert!(!mx( + "model.language_model.layers.3.self_attn.indexer.index_qk_proj" + )); assert!(mx("model.visual.blocks.0.attn.qkv")); assert!(mx("model.visual.blocks.0.attn.proj")); assert!(!mx("model.visual.blocks.0.mlp.linear_fc1")); @@ -487,10 +486,7 @@ fn unregistered_nvfp4_module_expands() { let done = materialize(&mut store, &cfg, &gpu).unwrap(); assert_eq!(done.nvfp4_to_bf16, 1); let w = store.get(&format!("{module}.weight")).unwrap(); - assert_eq!( - (w.dtype, w.shape.clone()), - (WeightDtype::BF16, vec![1, 16]) - ); + assert_eq!((w.dtype, w.shape.clone()), (WeightDtype::BF16, vec![1, 16])); let v = gpu.read_alloc(w.ptr).unwrap(); assert!((0..16).all(|i| { f32::from_bits(u32::from(u16::from_le_bytes([v[2 * i], v[2 * i + 1]])) << 16) == 1.0 diff --git a/crates/spark-model/src/weight_map/modelopt_plan.rs b/crates/spark-model/src/weight_map/modelopt_plan.rs index b10fb6194..2fce44742 100644 --- a/crates/spark-model/src/weight_map/modelopt_plan.rs +++ b/crates/spark-model/src/weight_map/modelopt_plan.rs @@ -116,9 +116,7 @@ impl ModeloptQuantPlan { /// The FP4-Local family marker: any module ships a compressed /// `weight_scale` encoding (today `"csf"`). pub fn is_fp4_local(&self) -> bool { - self.formats - .values() - .any(|f| f.scale_encoding.is_some()) + self.formats.values().any(|f| f.scale_encoding.is_some()) } } diff --git a/crates/spark-model/src/weight_map/modelopt_plan_tests.rs b/crates/spark-model/src/weight_map/modelopt_plan_tests.rs index d7174f706..eef8b8007 100644 --- a/crates/spark-model/src/weight_map/modelopt_plan_tests.rs +++ b/crates/spark-model/src/weight_map/modelopt_plan_tests.rs @@ -51,7 +51,10 @@ fn parses_the_flash_next_map() { ("model.language_model.layers.0.self_attn.q_proj", mx), ("model.language_model.layers.0.mlp.experts", nv4), ("model.visual.blocks.0.mlp.linear_fc1", w4), - ("model.language_model.layers.1.ple.ple_embedding.ngram_embedding", plain), + ( + "model.language_model.layers.1.ple.ple_embedding.ngram_embedding", + plain, + ), ]); cfg.model_type = "qwen4_exp".into(); let plan = ModeloptQuantPlan::from_config(&cfg).unwrap(); @@ -97,7 +100,9 @@ fn parses_the_csf_32x2_fixture_map() { ]); let plan = ModeloptQuantPlan::from_config(&cfg).unwrap(); assert!(plan.is_fp4_local()); - let gate = plan.format_of("model.language_model.layers.0.mlp.gate_proj").unwrap(); + let gate = plan + .format_of("model.language_model.layers.0.mlp.gate_proj") + .unwrap(); assert_eq!(gate.algo, QuantAlgo::Mxfp8); assert_eq!(gate.group_size, 32); assert_eq!( @@ -155,10 +160,7 @@ fn qwen4_exp_registry_coverage() { assert!(n.covers(m, &QuantAlgo::W4A16Nvfp4), "{m}"); } // Not covered: LM projections, lm_head. - for m in [ - "model.language_model.layers.0.self_attn.q_proj", - "lm_head", - ] { + for m in ["model.language_model.layers.0.self_attn.q_proj", "lm_head"] { assert!(!n.covers(m, &QuantAlgo::W4A16Nvfp4), "{m}"); } // An unknown model type covers nothing. From d773b218ea5656ef9e65171d9c147eaeb115045b Mon Sep 17 00:00:00 2001 From: Azeez <95100110+AzeezIsh@users.noreply.github.com> Date: Mon, 19 Oct 2026 21:22:31 -0400 Subject: [PATCH 23/26] vision: resolve the ViT dequant kernels only for quantized vision weights MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BF16-only vision trees no longer register mxfp8_proj_dequant or dequant_nvfp4_to_bf16 with the boot audit — the FP4-Local native-format work had made every non-flash-next vision model fail boot. A checkpoint that does carry MXFP8/NVFP4 projections fails hard naming the kernel and weight if the kernel is missing. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../layers/vision_encoder/enc_impl/init.rs | 59 +++++++++++-------- 1 file changed, 36 insertions(+), 23 deletions(-) diff --git a/crates/spark-model/src/layers/vision_encoder/enc_impl/init.rs b/crates/spark-model/src/layers/vision_encoder/enc_impl/init.rs index 0b4ad66c9..77c5594a4 100644 --- a/crates/spark-model/src/layers/vision_encoder/enc_impl/init.rs +++ b/crates/spark-model/src/layers/vision_encoder/enc_impl/init.rs @@ -2,11 +2,11 @@ //! `VisionEncoder::new` constructor. -use anyhow::Result; +use anyhow::{Context, Result}; use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle}; use super::super::{ - GlmVisionEncoder, GlmVisionWeights, MergerLayer, PATCH_DIM, ViTBlock, VisionEncoder, + GlmVisionEncoder, GlmVisionWeights, MergerLayer, PATCH_DIM, ViTBlock, VisionEncoder, VitWeight, }; /// Encoder capacity, in patches, when nothing bounds the image. @@ -150,18 +150,24 @@ impl VisionEncoder { // Weight-dequant scratch: the largest quantized projection's BF16 // expansion (~10 MB for fc1 [4304,1152]). Owned by the encoder; never // aliased with LM scratch. - let max_w_bytes = blocks - .iter() - .flat_map(|b| [b.qkv_w, b.proj_w, b.fc1_w, b.fc2_w]) - .chain( - deepstack - .iter() - .flat_map(|m| [m.fc1_w, m.fc2_w]) - .chain([merger.fc1_w, merger.fc2_w]), - ) - .map(|w| w.bf16_bytes()) - .max() - .unwrap_or(0); + let all_weights = || { + blocks + .iter() + .flat_map(|b| [b.qkv_w, b.proj_w, b.fc1_w, b.fc2_w]) + .chain( + deepstack + .iter() + .flat_map(|m| [m.fc1_w, m.fc2_w]) + .chain([merger.fc1_w, merger.fc2_w]), + ) + }; + // Lookup the dequant kernels ONLY when a loaded projection is that + // dtype: a BF16-only vision tree never registers them, so the boot + // audit has nothing to miss (the FP4-Local regression that failed + // every non-flash-next vision model). + let has_nvfp4_w = all_weights().any(|w| matches!(w, VitWeight::Nvfp4 { .. })); + let has_mxfp8_w = all_weights().any(|w| matches!(w, VitWeight::Mxfp8(_))); + let max_w_bytes = all_weights().map(|w| w.bf16_bytes()).max().unwrap_or(0); let buf_vit_dq = if max_w_bytes > 0 { gpu.alloc(max_w_bytes)? } else { @@ -227,15 +233,22 @@ impl VisionEncoder { k_merge: gpu.kernel("vision_encoder", "vision_spatial_merge")?, k_f32_bf16: gpu.kernel("vision_encoder", "vision_f32_to_bf16")?, k_copy: gpu.kernel("vision_encoder", "vision_bf16_copy")?, - // Dequants for checkpoint-native quantized ViT projections. - // SOFT: absent on vision trees without them (then every weight - // is Bf16 and they are never launched). - k_nvfp4_dq: crate::layers::try_kernel( - gpu, - "dequant_nvfp4_bf16", - "dequant_nvfp4_to_bf16", - ), - k_mx8_dq: crate::layers::try_kernel(gpu, "mxfp8_proj", "mxfp8_proj_dequant"), + // Dequants for checkpoint-native quantized ViT projections: + // resolved ONLY when a loaded weight is that dtype (hard error + // if the kernel is missing then — the weight cannot decode + // without it). BF16-only trees never even issue the lookup. + k_nvfp4_dq: if has_nvfp4_w { + gpu.kernel("dequant_nvfp4_bf16", "dequant_nvfp4_to_bf16") + .context("NVFP4 vision weight needs dequant_nvfp4_to_bf16")? + } else { + KernelHandle(0) + }, + k_mx8_dq: if has_mxfp8_w { + gpu.kernel("mxfp8_proj", "mxfp8_proj_dequant") + .context("MXFP8 vision weight needs mxfp8_proj_dequant")? + } else { + KernelHandle(0) + }, hidden_size, num_heads, head_dim, From 5ddfabb6e2e7941c37779963f102b82b737df2bb Mon Sep 17 00:00:00 2001 From: AzeezIsh <95100110+AzeezIsh@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:20:05 -0400 Subject: [PATCH 24/26] spark-storage: rustfmt the resident PLE-scale code under the pinned 1.93.1 Formatting only. CI's fmt job runs the rust-toolchain.toml pin (1.93.1), which these two spots from the resident-scale arena did not match. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/spark-storage/src/ngram_cache/fault_pool_tests.rs | 4 +--- crates/spark-storage/src/ngram_cache_fault.rs | 4 +++- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs b/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs index feed2b096..8c248558c 100644 --- a/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs +++ b/crates/spark-storage/src/ngram_cache/fault_pool_tests.rs @@ -64,9 +64,7 @@ fn resident_scales_index_by_global_row_id() { resident: Some(&arena), }; for id in [0u64, 3, 4, 7, 8, 11] { - let sl = ss - .scale_resident_slice(id, stride) - .expect("resident slice"); + let sl = ss.scale_resident_slice(id, stride).expect("resident slice"); assert_eq!( sl, &arena[id as usize * stride..id as usize * stride + stride], diff --git a/crates/spark-storage/src/ngram_cache_fault.rs b/crates/spark-storage/src/ngram_cache_fault.rs index d22c9a4f4..dfea7f715 100644 --- a/crates/spark-storage/src/ngram_cache_fault.rs +++ b/crates/spark-storage/src/ngram_cache_fault.rs @@ -75,7 +75,9 @@ pub(super) struct RowSource<'a> { /// `Segmented`: the table's own shard segmentation applied to the sibling /// scale tensors (same `rows_per`, same deduped files). pub(super) enum ScaleSource<'a> { - Flat { file: &'a File }, + Flat { + file: &'a File, + }, /// `resident`: `Some` under `ATLAS_PLE_SCALES_RESIDENT` — the same /// (shard, local) split as `scale_loc`, answered from host RAM. Segmented { From 15839966acc03756ee17071d26a7c7e7c76a4045 Mon Sep 17 00:00:00 2001 From: AzeezIsh <95100110+AzeezIsh@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:57:07 -0400 Subject: [PATCH 25/26] bench(qwen3.8-flash-next): declare FP4-Local's measured ST-995 and agentic-webserver records NVIDIA's FP4-Local checkpoint now has its own non-default, measured entries next to the vendor pack's, so its evidence lives in the benchmark SSOT rather than only in PR comments: the golden ST-995 draw (87.44 / 87.32 in run-1792340347555552942 at a78d2fb and run-1792488042130537521 at 5ddfabb6e) and the 50-iteration agentic soak (50/50 + 50/50, run-1792475013132909319 at d773b218e). Bars follow the pack entry's rules: measured less the inherited 0.4 MTP floor for BFCL, absolute measured counts for the soak, speed deliberately ungated. BENCH.toml is outside the closure hash, so no gate record is invalidated. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- kernels/gb10/qwen3.8-flash-next/BENCH.toml | 82 ++++++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/kernels/gb10/qwen3.8-flash-next/BENCH.toml b/kernels/gb10/qwen3.8-flash-next/BENCH.toml index a7fb6377c..bba11ad2b 100644 --- a/kernels/gb10/qwen3.8-flash-next/BENCH.toml +++ b/kernels/gb10/qwen3.8-flash-next/BENCH.toml @@ -114,3 +114,85 @@ min = 50.0 [benchmarks.metrics.followed_directions] min = 49.0 + +# --- NVIDIA's FP4-Local QAD checkpoint (nvidia/Qwen3.8-Flash-Next-FP4-Local) -- +# Served on its checkpoint-native formats (docs/checkpoints/qwen3.8-flash-next- +# fp4-local.md). Different weights from the vendor pack above, so its own bars: +# do NOT carry the pack's numbers across in either direction. Both entries were +# bootstrapped per the repo protocol (manual drive with `--url/--model`), on +# sparky (GB10), snapshot b9ce18a6, one profile: +# --max-num-seqs 1 --max-batch-size 1 --max-seq-len 32768 +# --max-prefill-tokens 8192 --kv-cache-dtype bf16 --gpu-memory-utilization 0.9 +# --enable-prefix-caching --fast-load-prefetch-shards --speculative +# --num-drafts 3, reasoning_effort "low", ATLAS_PLE_CACHE_GB=2, +# ATLAS_QWEN4EXP_{EXACT_VERIFY,DECODE_FUSE,BATCH_FAST,BATCH_SMALL}=1, +# ATLAS_QWEN4EXP_MTP_DEPTH=3 (+ ATLAS_QWEN4EXP_HC_FAST=1 where noted). +# The recipe named here encodes exactly that profile; it is being added to +# atlas-recipes alongside this entry. + +[[benchmarks]] +quant = "nvfp4" +checkpoint = "nvidia/Qwen3.8-Flash-Next-FP4-Local" +gate = "bfcl-subset" +recipe = "qwen3.8/qwen3.8-flash-next-fp4-local" +label = "Qwen3.8-Flash-Next (FP4-Local, nvidia)" +default = false +status = "measured" +note = """\ +Golden ST-995 draw (62/10/10, floor 25, n=995, seed 42, temp 0, no param \ +overrides) — the MLPerf edge-agentic ST leg. REFERENCE RECORDS: \ +(1) run-1792340347555552942, 2026-10-18, binary a78d2fb (sha256 8d566e68), \ +profile without HC_FAST; (2) run-1792488042130537521, 2026-10-20, binary \ +5ddfabb6e (sha256 8294e6b8), profile with HC_FAST. Both scored overall \ +87.44 / normalized 87.32. ★ As with the vendor pack above, that identity is \ +reproducibility across two builds and a lever change, NOT a repeat-spread: \ +the exact-verify and exact-batching lanes are bit-exact against serial decode \ +and HC_FAST was CHECKed 0/256 rows different, so identical scores are what \ +correctness predicts. The 0.4 margin is the dense sibling's documented MTP \ +nondeterminism floor, inherited and labelled as such. Same build and flags \ +with the vendor pack swapped in scored 86.23 / 85.64 (run-1792378038850732185) \ +— one run each, an observation rather than a margin.""" + +[benchmarks.metrics.samples] +min = 995.0 +max = 995.0 + +[benchmarks.metrics.overall_accuracy] +# Measured 87.44 in both reference records; bar = measured less 0.4. +min = 87.04 +noise = 0.4 + +[benchmarks.metrics.normalized_single_turn_score] +# Measured 87.32 in both reference records; bar = measured less 0.4. Above +# the MLPerf-edge floor (85.32), which was derived for Qwen3.6-27B and does +# not transfer across weights either way. +min = 86.92 +noise = 0.4 + +[[benchmarks]] +quant = "nvfp4" +checkpoint = "nvidia/Qwen3.8-Flash-Next-FP4-Local" +gate = "agentic-webserver" +recipe = "qwen3.8/qwen3.8-flash-next-fp4-local" +label = "Qwen3.8-Flash-Next (FP4-Local, nvidia)" +default = false +status = "measured" +note = """\ +REFERENCE RECORD: run-1792475013132909319, 2026-10-20, binary d773b218e, \ +profile with HC_FAST, harness iterations=50 max_turns=40 max_tokens=8192 \ +wall_budget_s=100000 — 50/50 webserver_ok, 50/50 followed_directions, 373 \ +turns, 375 tool calls, 98,142 completion tokens, Sigma-wall 4,502 s \ +(11.94 s/turn, decode 22.0 tok/s). ★ SPEED IS NOT GATED, for the reason the \ +vendor-pack entry gives: one record is one tier. The counts are ABSOLUTE \ +correctness bounds at their measured values and take no noise; `iterations` \ +is pinned at exactly 50.""" + +[benchmarks.metrics.iterations] +min = 50.0 +max = 50.0 + +[benchmarks.metrics.webserver_ok] +min = 50.0 + +[benchmarks.metrics.followed_directions] +min = 50.0 From 039dec88bcdbc5d2415dea82c9a8419e4a468403 Mon Sep 17 00:00:00 2001 From: AzeezIsh <95100110+AzeezIsh@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:21:17 -0400 Subject: [PATCH 26/26] bench(qwen3.8-flash-next): reword the FP4-Local note for the typos check Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- kernels/gb10/qwen3.8-flash-next/BENCH.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kernels/gb10/qwen3.8-flash-next/BENCH.toml b/kernels/gb10/qwen3.8-flash-next/BENCH.toml index bba11ad2b..3d56570aa 100644 --- a/kernels/gb10/qwen3.8-flash-next/BENCH.toml +++ b/kernels/gb10/qwen3.8-flash-next/BENCH.toml @@ -147,7 +147,7 @@ profile without HC_FAST; (2) run-1792488042130537521, 2026-10-20, binary \ 87.44 / normalized 87.32. ★ As with the vendor pack above, that identity is \ reproducibility across two builds and a lever change, NOT a repeat-spread: \ the exact-verify and exact-batching lanes are bit-exact against serial decode \ -and HC_FAST was CHECKed 0/256 rows different, so identical scores are what \ +and the HC_FAST check found 0/256 rows different, so identical scores are what \ correctness predicts. The 0.4 margin is the dense sibling's documented MTP \ nondeterminism floor, inherited and labelled as such. Same build and flags \ with the vendor pack swapped in scored 86.23 / 85.64 (run-1792378038850732185) \