qwen4_exp: serve NVIDIA's Qwen3.8-Flash-Next FP4-Local on its native formats (one GB10)

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 <port> <model> <outdir>
+
+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 <outdir>/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]}))

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 <port> <outdir>   -> <outdir>/{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)}))

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 <port> <target_tokens> [<target_tokens> ...]
+
+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)

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 <port> <concurrency> <max_tokens>
+
+Sends <concurrency> 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<String, HashMap<String, String>>, 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)}))

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 <snapshot> csf   <scale-name>       <out.bin>   # E4M3 bytes
+//! fp4local_decode <snapshot> mxfp8 <module>           <out.bin>   # BF16 bytes
+//! fp4local_decode <snapshot> ple   <shard-base> <row> <out.bin>   # 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<usize>,
+}
+
+fn locate(snap: &Path, name: &str) -> Result<Tensor> {
+    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<Vec<u8>> {
+    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<String> = 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(())
+}

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<String>,
+    /// 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<usize>,
 }
 
 /// Whether a ModelOpt `quant_algo` label means NVFP4-packed weights.