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/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/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.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}"); + } +} diff --git a/crates/spark-model/Cargo.toml b/crates/spark-model/Cargo.toml index 2e83c23fd..06810ca9a 100644 --- a/crates/spark-model/Cargo.toml +++ b/crates/spark-model/Cargo.toml @@ -207,6 +207,18 @@ 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 = "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/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/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/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, // 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/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/mod.rs b/crates/spark-model/src/layers/mod.rs index cca874637..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, 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/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..ddaafa2e3 --- /dev/null +++ b/crates/spark-model/src/layers/moe/helpers_shared.rs @@ -0,0 +1,226 @@ +// 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; + 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, + )?; + 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.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/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/layers/ops/hyper_connection_lowrank_tests.rs b/crates/spark-model/src/layers/ops/hyper_connection_lowrank_tests.rs index 01c602cd0..0a8ed8b62 100644 --- a/crates/spark-model/src/layers/ops/hyper_connection_lowrank_tests.rs +++ b/crates/spark-model/src/layers/ops/hyper_connection_lowrank_tests.rs @@ -153,11 +153,23 @@ fn transpose_bf16(bytes: &[u8], rows: usize, cols: usize) -> 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/mxfp8_proj.rs b/crates/spark-model/src/layers/ops/mxfp8_proj.rs new file mode 100644 index 000000000..f2693e116 --- /dev/null +++ b/crates/spark-model/src/layers/ops/mxfp8_proj.rs @@ -0,0 +1,161 @@ +// 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..=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 +//! 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}; + +pub use crate::weight_map::Mxfp8Proj; + +/// 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 { + 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")?, + }) + } +} + +/// 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 { 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 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(()); + } + 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(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", + 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/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/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/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/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.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/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/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..bec527d82 100644 --- a/crates/spark-model/src/layers/ple/layer.rs +++ b/crates/spark-model/src/layers/ple/layer.rs @@ -61,8 +61,15 @@ 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): packed rows + block scales stay in the + /// pinned caches and decode inside `batched_embed_nvfp4`. + nvfp4: Option, gemm_k: KernelHandle, gate_k: KernelHandle, conv_k: KernelHandle, @@ -96,6 +103,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() + } + #[allow(clippy::too_many_arguments)] pub fn new( dims: PleIdDims, @@ -110,6 +123,7 @@ impl PleLayer { max_tokens: usize, scratch_tokens: usize, trellis: Option, + nvfp4: Option, gpu: &dyn GpuBackend, ) -> Result { dims.validate()?; @@ -132,12 +146,15 @@ 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)?; } } + let is_nvfp4 = nvfp4.is_some(); Ok(Self { dims, head_dim, @@ -158,11 +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)?, @@ -360,7 +393,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, @@ -396,6 +431,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..50523c413 --- /dev/null +++ b/crates/spark-model/src/layers/ple/ngram_nvfp4.rs @@ -0,0 +1,106 @@ +// 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. + +#[cfg(test)] +use anyhow::Context as _; +use anyhow::{Result, ensure}; +#[cfg(test)] +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 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, + /// `(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 { + 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 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 need the block-scale arena attached to the \ + row cache (`attach_segmented_block_scales`)" + ); + ensure!( + stride == spec.packed_stride(), + "PLE: NVFP4 row stride {stride} != {} packed bytes for {} dims", + spec.packed_stride(), + spec.dim + ); + Ok(()) +} + +/// 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], + 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) +} + +#[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..19d95f6f9 --- /dev/null +++ b/crates/spark-model/src/layers/ple/ngram_nvfp4_tests.rs @@ -0,0 +1,43 @@ +// 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()); + // 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()); +} diff --git a/crates/spark-model/src/layers/ple/verify.rs b/crates/spark-model/src/layers/ple/verify.rs index 1707a2f91..2bb34e104 100644 --- a/crates/spark-model/src/layers/ple/verify.rs +++ b/crates/spark-model/src/layers/ple/verify.rs @@ -228,10 +228,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_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_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/layers/qwen3_ssm/exact_rows.rs b/crates/spark-model/src/layers/qwen3_ssm/exact_rows.rs index 205098474..1cd409024 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,23 @@ 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, + }; + ops::mxfp8_proj( + ctx.gpu, + &mx8.kernels, + p, + input, + output, + rows, + 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 eae90fb54..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, @@ -347,6 +347,7 @@ fn active_row_without_ple_carry_is_refused() { 1, 1, None, + None, &gpu, ) .unwrap(); @@ -395,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/layers/vision_encoder.rs b/crates/spark-model/src/layers/vision_encoder.rs index f88a2c401..8e582e7b6 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..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. @@ -147,6 +147,32 @@ 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 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 { + 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 +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: + // 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, @@ -231,6 +273,7 @@ impl VisionEncoder { buf_scores, buf_probs, buf_o_stage, + buf_vit_dq, pos_embed_host_f32, rope_inv_freq, }) @@ -257,9 +300,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 +319,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 +333,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/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); } 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. diff --git a/crates/spark-model/src/weight_loader/qwen35.rs b/crates/spark-model/src/weight_loader/qwen35.rs index 5cb515ec2..7de8c21ff 100644 --- a/crates/spark-model/src/weight_loader/qwen35.rs +++ b/crates/spark-model/src/weight_loader/qwen35.rs @@ -14,19 +14,100 @@ 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 +283,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 +306,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 +319,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/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/qwen3_vl.rs b/crates/spark-model/src/weight_loader/qwen3_vl.rs index a6a3ae6c3..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: 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 +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: 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 +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: 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_loader/qwen4_exp.rs b/crates/spark-model/src/weight_loader/qwen4_exp.rs index 0c6c585bd..a543ac50c 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp.rs @@ -77,10 +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; @@ -230,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", @@ -270,12 +255,13 @@ 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 { 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); @@ -289,20 +275,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 = mx8_common::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, @@ -325,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/ffn.rs b/crates/spark-model/src/weight_loader/qwen4_exp/ffn.rs index 2e904c3d2..e5fe38767 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,75 @@ 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)?; + // 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( + view(0), + view(inter), + 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/gdn_fp8.rs b/crates/spark-model/src/weight_loader/qwen4_exp/gdn_fp8.rs index c1b3c0687..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; @@ -43,7 +66,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 +122,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..247273661 --- /dev/null +++ b/crates/spark-model/src/weight_loader/qwen4_exp/gdn_mx8.rs @@ -0,0 +1,148 @@ +// 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 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. + +use anyhow::{Context, Result, ensure}; +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; +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, +}; + +/// 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")) +} + +/// 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) +} + +/// 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_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/mtp.rs b/crates/spark-model/src/weight_loader/qwen4_exp/mtp.rs index 2c2c6e778..3fa9b4f74 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. @@ -234,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 { @@ -267,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)) } @@ -420,71 +421,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; +/// Per-expert NVFP4 MTP experts (FP4-Local). +#[path = "mtp_nvfp4.rs"] +mod nvfp4; - 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 - ); - } -} +#[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 + ); +} 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_loader/qwen4_exp/ple.rs b/crates/spark-model/src/weight_loader/qwen4_exp/ple.rs index 6a189d893..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; @@ -42,7 +59,9 @@ 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}; /// 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 @@ -77,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 @@ -110,75 +144,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 +288,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`)" @@ -333,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( @@ -364,10 +333,17 @@ pub(super) fn load( ); } + let n_shards = shards.len(); + 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 // 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 @@ -379,9 +355,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"))?, }; @@ -428,6 +406,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 +425,7 @@ pub(super) fn load( max_tokens, span, trellis, + nvfp4, gpu, ) .map(Some) @@ -490,3 +471,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..e11e4c921 --- /dev/null +++ b/crates/spark-model/src/weight_loader/qwen4_exp/ple_nvfp4.rs @@ -0,0 +1,79 @@ +// 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, + gpu: &dyn GpuBackend, +) -> Result> { + if dtype != WeightDtype::UInt8 { + return Ok(None); + } + open(store, lp, shards, rows_per, packed_cols, gpu).map(Some) +} + +fn open( + store: &WeightStore, + lp: &str, + shards: usize, + rows_per: usize, + packed_cols: 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)?, + }; + Ok(NgramNvfp4 { geom, scale_shards }) +} diff --git a/crates/spark-model/src/weight_map.rs b/crates/spark-model/src/weight_map.rs index 31fdfa54c..4ae9f132e 100644 --- a/crates/spark-model/src/weight_map.rs +++ b/crates/spark-model/src/weight_map.rs @@ -25,6 +25,9 @@ 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; +pub mod modelopt_plan; #[path = "weight_map/moe.rs"] mod moe; #[path = "weight_map/mxfp8.rs"] @@ -56,7 +59,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 new file mode 100644 index 000000000..72bbd8e0a --- /dev/null +++ b/crates/spark-model/src/weight_map/modelopt_local.rs @@ -0,0 +1,380 @@ +// 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. +//! 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; +//! * vision MLP/merger projections are W4A16 NVFP4 — dequantized to the BF16 +//! 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 +//! 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 super::modelopt_plan::{ModeloptQuantPlan, QuantAlgo, native_consumers}; +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: every `MXFP8` +/// entry of the plan (native upload and host expansion both read the +/// deferred bytes), plus — on FP4-Local — every module with no registered +/// native consumer (its packed bytes are never meant to be resident). The +/// pre-plan rule for non-FP4-Local packs is the `MXFP8`-only arm. +pub fn host_decode_modules(config: &ModelConfig) -> 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)| spec.quant_algo.eq_ignore_ascii_case("MXFP8")) + .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_native: usize, + pub mxfp8_to_bf16: usize, + pub nvfp4_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 +/// 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. + 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 +} + +/// 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, + ) +} + +/// 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( + 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(); + 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 { + 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 { + 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 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, 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_native, + 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..05921f67f --- /dev/null +++ b/crates/spark-model/src/weight_map/modelopt_local_tests.rs @@ -0,0 +1,494 @@ +// 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 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"; +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 { + quant_algo: algo.into(), + group_size: group, + weight_scale_encoding: None, + weight_block_size: vec![], + } +} + +/// 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(); + 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(SHARED.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(), + 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 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].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(); +} + +#[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!("{SHARED}.weight"), + &[0x38; 32], + &[1, 32], + WeightDtype::FP8E4M3, + ); + shard.add( + &format!("{SHARED}.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_native, + done.mxfp8_to_bf16, + done.nvfp4_to_bf16, + done.f32_norms_native + ), + (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(); + 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(); + 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(); + 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" + ); + + 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 vis = store.deferred(&format!("{VIS}.weight")).unwrap(); + assert_eq!( + (vis.dtype, vis.shape.clone()), + (WeightDtype::UInt8, vec![1, 8]) + ); + // 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 +/// exactly as shipped, weight and scale, and nothing is widened. +#[test] +fn gdn_mxfp8_stays_native() { + 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(); + 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] +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); +} + +/// 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..2fce44742 --- /dev/null +++ b/crates/spark-model/src/weight_map/modelopt_plan.rs @@ -0,0 +1,198 @@ +// 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..eef8b8007 --- /dev/null +++ b/crates/spark-model/src/weight_map/modelopt_plan_tests.rs @@ -0,0 +1,194 @@ +// 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/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 + } +} 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![], }, ); } 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 ddab38048..782e59c0f 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,75 @@ 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/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(); +} 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(); diff --git a/crates/spark-storage/src/ngram_cache.rs b/crates/spark-storage/src/ngram_cache.rs index 7f101da93..2c6fd786b 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,25 @@ 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, + /// `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. +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 +203,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 +231,9 @@ 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)?), + resident: None, }) } None => None, @@ -216,16 +247,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 +348,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 +418,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..ba6056273 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,26 @@ 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, Option>>), +} + +impl OwnedScaleSrc { + fn view(&self) -> ScaleSource<'_> { + match self { + OwnedScaleSrc::Flat(f) => ScaleSource::Flat { file: f }, + OwnedScaleSrc::Segmented(s, r) => ScaleSource::Segmented { + seg: s, + resident: r.as_ref().map(|v| v.as_slice()), + }, + } + } } impl OwnedSource { @@ -47,7 +66,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 +190,26 @@ 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(), + }, + cache.scales.as_ref().and_then(|sc| sc.resident.clone()), + )), }; let src = OwnedSource { file: cache @@ -185,7 +219,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..8c248558c 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, } } @@ -44,6 +45,34 @@ 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 205ffa9ed..6a7b257a5 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,23 @@ 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, + resident: sc.resident.as_ref().map(|v| v.as_slice()), + }), + }), }; 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 +235,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 +254,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..50b0610ef --- /dev/null +++ b/crates/spark-storage/src/ngram_cache/scale_tests.rs @@ -0,0 +1,158 @@ +// 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..c7400efc3 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,108 @@ 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, + resident: None, + }); 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); + // 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")?, + stride: scale_stride, + src: ScaleSrc::Segmented(Segments { + bases: shards.iter().map(|(_, off)| *off).collect(), + rows_per, + files, + shard_file, + }), + resident, + }); + 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_fault.rs b/crates/spark-storage/src/ngram_cache_fault.rs index 0ffc13efc..dfea7f715 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 @@ -65,7 +67,51 @@ 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, + }, + /// `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<'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, 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]; + (file, seg.bases[shard] + local * stride) + } + } + } } impl RowSource<'_> { @@ -128,16 +174,29 @@ 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; + 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; - 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/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..2ed024c4f --- /dev/null +++ b/docs/checkpoints/qwen3.8-flash-next-fp4-local.md @@ -0,0 +1,232 @@ +# 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. + +## 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. 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); + } +} 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); +} diff --git a/kernels/gb10/qwen3.8-flash-next/BENCH.toml b/kernels/gb10/qwen3.8-flash-next/BENCH.toml index a7fb6377c..3d56570aa 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 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) \ +— 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 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/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]); +} 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])); } }