One server: the prompt keeps 95% of its speed with the DFlash2 drafter (445 -> 477 t/s)

README.md

@@ -57,7 +57,7 @@ Windows and 1bit OS ([docs/releases.md](docs/releases.md)).
 > ([docs/vulkan.md](docs/vulkan.md#zaya1-zyphra-from-our-llamacpp)); the rest of Zyphra's
 > family (Zamba, Zamba2, BlackMamba) runs on Vulkan too, and so do its vision models, ZAYA1-VL-8B
 > and Zamba2-VL, through `1bit serve --mmproj`. Qwen3.8-27B runs in one ROCm server with Hadamard
-> W4A4 prompt processing and DFlash2 decode: 445 tok/s on a 1,838-token prompt and 40.9 tok/s
+> W4A4 prompt processing and DFlash2 decode: 477 tok/s on a 1,838-token prompt and 41.1 tok/s
 > decode on code ([docs/lean.md](docs/lean.md#hadamard-rotated-q4_0-w4a4-prompt-processing)). W4A4 covers MoE
 > experts too, and `1bit serve --long-model` sends long prompts there and short ones to Vulkan: on
 > Qwen3-Coder-30B-A3B a 16K-token prompt gets its answer at 13.1 effective tok/s against 8.6 on

app/serve.cpp

@@ -662,9 +662,12 @@ Launch launch_for(const Options& o, const std::string& device, int child_port) {
             if (device != "rocm") throw std::runtime_error(o.model + " is Hadamard-rotated: it runs on --device rocm only");
             if (!std::getenv("GGML_Q4_0_HADAMARD")) env.push_back("GGML_Q4_0_HADAMARD=1");
             if (!std::getenv("GGML_W4A4_TENSORS")) env.push_back("GGML_W4A4_TENSORS=all");
-            // 1024-token micro-batches: prompt processing +3-8% over the default 512 (Qwen3-Coder-30B-A3B,
-            // pp2048 2,001 -> 2,158, pp16384 1,227 -> 1,283; 2048 is no better), decode unchanged
-            argv.insert(argv.end(), {"-ub", "1024"});
+            // MoE files: 1024-token micro-batches, prompt processing +3-8% over the default 512
+            // (Qwen3-Coder-30B-A3B, pp2048 2,001 -> 2,158, pp16384 1,227 -> 1,283; 2048 is no
+            // better), decode unchanged. A dense file keeps 512: Qwen3.8-27B-H32 serves a 1,838-token
+            // prompt at 491-495 tok/s with 512 and 471-473 with 1024.
+            if (gguf_int(o.model, gguf_architecture(o.model) + ".expert_count") > 0)
+                argv.insert(argv.end(), {"-ub", "1024"});
         }
     } else if (device == "vulkan" || device == "hrx") {
         // --prefill-device hrx: the HRX build (it has both devices and the shared-KV split), flash
@@ -748,8 +751,14 @@ Launch launch_for(const Options& o, const std::string& device, int child_port) {
         else argv.insert(argv.end(), {"--spec-type", "draft-dflash", "-md", o.dflash, "-ngld", "99"});
         if (o.mtp_max > 0) { argv.push_back("--spec-draft-n-max"); argv.push_back(std::to_string(o.mtp_max)); }
         // DFlash drafts a whole block per step and only pays with full blocks; llama-server's
-        // default draft length is 3, and it clamps a longer one to the drafter's trained block
-        else if (!o.dflash.empty()) argv.insert(argv.end(), {"--spec-draft-n-max", "16"});
+        // default draft length is 3. The longest draft is the drafter's block minus its anchor
+        // token, and the draft length also sets how many recurrent states a hybrid model keeps
+        // for rollback: asking for 16 against DFlash2's block of 8 made every prompt write 16
+        // snapshots (Qwen3.8-27B-H32, 1,838-token prompt: 452-457 tok/s against 456-464).
+        else if (!o.dflash.empty()) {
+            const long long block = gguf_int(o.dflash, gguf_architecture(o.dflash) + ".block_size");
+            argv.insert(argv.end(), {"--spec-draft-n-max", block > 1 ? std::to_string(block - 1) : "16"});
+        }
         if (!o.mtp_p_min.empty()) { argv.push_back("--spec-draft-p-min"); argv.push_back(o.mtp_p_min); }
         // a DFlash block is kept whole unless --mtp-p-min says otherwise: upstream's default p-min
         // is 0, the ROCm tree's is 0.75, which cuts DFlash2 blocks from 6.7 to 5.4 tokens a step

docs/lean.md

@@ -263,7 +263,7 @@ tiled transpose for the concatenation that feeds the delta-net convolution. Both
 
 **One server: W4A4 prompts and DFlash2 decode.** Without a drafter this route decodes
 Qwen3.8-27B at about 13 tok/s. With the DFlash2 drafter ([serve.md](serve.md#dflash-draft-models---dflash))
-the same server decodes three times faster and keeps most of the prompt speed:
+the same server decodes three times faster and keeps 95% of the prompt speed:
 
 ```sh
 1bit serve -m Qwen3.8-27B-Q4_0-H32.gguf --dflash Qwen3.8-27B-DFlash2-q8_0.gguf
@@ -276,20 +276,36 @@ DFlash, so a partly accepted block rewinds the delta-net state instead of restor
 and replaying. `serve` passes `--spec-draft-p-min 0` with `--dflash`: this tree's default of 0.75
 cut DFlash2 blocks from 6.7 to 5.4 tokens a step.
 
-Measured through `1bit serve` on Strix Halo, 2026-09-28, quiet box: the 1,838-token prompt (best
+Three things kept the drafter from costing the prompt (ROCmFPX#4 and `serve`):
+
+- **Rollback snapshots in the prefill kernel.** Rollback keeps the delta-net state after each of
+  the last K tokens, and the chunked prefill kernel did not write those, so with a drafter every
+  prompt ran the per-token kernel: 1,322 tokens took 2,966 ms instead of 2,588. The prefill
+  kernel now writes them (2,762 ms).
+- **K is the drafter's block.** `serve` reads `dflash.block_size` from the drafter and drafts
+  `block - 1` tokens, which is also K; asking for 16 against DFlash2's block of 8 wrote twice the
+  snapshots.
+- **512-token micro-batches for a dense file.** The 1024-token micro-batches `serve` gives
+  rotated MoE files are slower on the dense 27B: 471-473 tok/s against 491-495 without a drafter.
+
+The same ROCmFPX change also fixed the drafter's input on long prompts: the model's hidden states
+for a batch of several micro-batches were all written at the first one's place, so the DFlash
+(and EAGLE3) encoder read stale features for all but the last.
+
+Measured through `1bit serve` on Strix Halo, 2026-09-29, quiet box: the 1,838-token prompt (best
 / median of 5, as `serve` reports it, context checkpoints included) and 256 tokens of greedy
 decode on the code / prose / short prompts of [serve.md](serve.md) (best of 3):
 
 | Qwen3.8-27B-Q4_0-H32, lean ROCm route | Prompt t/s | Decode tok/s |
 |---|---|---|
-| no drafter | 502 / 500 | 13.0 / 13.0 / 13.4 |
-| `--dflash` (DFlash2 Q8_0) | 445 / 439 | **40.9** / **26.8** / 13.7 |
-
-Mean accepted block: 6.71 tokens on code, 4.11 on prose, the same as upstream on this drafter.
-Greedy output matches the no-drafter run on the code and short prompts; on prose one near-tie
-word flips after 470 characters (batched verification rounds differently). The drafter costs the
-prompt about 0.5 s on this prompt: its encoder reads the model's hidden states for every prompt
-token (0.13 s), and the context checkpoints grow by the drafter's sliding-window cache.
+| no drafter | 497 / 495 | 13.2 / 13.2 / 13.6 |
+| `--dflash` (DFlash2 Q8_0) | **477** / **472** | **41.1** / **26.7** / 15.0 |
+
+Mean accepted block: 6.54 tokens on code, 4.23 on prose, in line with upstream on this drafter
+(6.71 / 4.25). Greedy output matches the no-drafter run on the code and short prompts; on prose
+one near-tie phrase differs after 332 characters (batched verification rounds differently). What
+the drafter still costs the prompt, about 0.2 s here, is its encoder reading the model's hidden
+states for every prompt token (0.1 s) and the rollback snapshots.
 
 ## Build
 

docs/serve.md

@@ -184,9 +184,10 @@ drafter is faster than BF16 (42.5 vs 38.9 on code, direct llama-server). A draft
 kept only when the model agrees, so the output is the model's own.
 
 `--dflash` also works on the lean ROCm route: a Hadamard-rotated Q4_0 file gets W4A4 prompt
-processing and DFlash2 decode from one server, 445 t/s prompt and 40.9 / 26.8 tok/s decode
+processing and DFlash2 decode from one server, 477 t/s prompt and 41.1 / 26.7 tok/s decode
 ([lean.md](lean.md#hadamard-rotated-q4_0-w4a4-prompt-processing)). `--dflash` sets
-`--spec-draft-p-min 0` unless `--mtp-p-min` is given.
+`--spec-draft-p-min 0` unless `--mtp-p-min` is given, and without `--mtp-max` drafts the
+drafter's block minus one (its `dflash.block_size`; 16 when the file does not say).
 
 ## MoE models larger than memory (`--moe-slots`)
 
@@ -311,7 +312,7 @@ Qwen3-Coder-30B-A3B, 256 output tokens, effective tok/s (output tokens over the
 | **`--long-model`** | **71.3** | 53.1 | 25.5 | 13.1 |
 
 The `--long-model` row ran while another job used the box's memory bus, which cost its long
-side 2-5% against the W4A4 row. A rotated file now runs 1024-token micro-batches (`-ub 1024`),
+side 2-5% against the W4A4 row. A rotated MoE file now runs 1024-token micro-batches (`-ub 1024`; a dense one keeps 512, [lean.md](lean.md)),
 which raises the W4A4 prefill 6-10% (server-reported 2,217 / 1,742 / 1,295 tok/s at 2K / 8K /
 16K, from 2,000 / 1,607 / 1,221): a 16K-token prompt's first token at 12.96 s instead of 13.89. Both files stay loaded (here 17 + 16 GiB). It needs the Vulkan
 build and `ONEBIT_LEAN` + `ONEBIT_LEAN_ROCM`, and `-m` on a llama-server route (vulkan, hrx or

tests/hadamard_route.sh

@@ -18,10 +18,12 @@
 # onebit.hadamard_q4_0 = 32; docs/lean.md), without a GPU:
 #   - --device vulkan refuses it,
 #   - --device auto sends it to the ROCm route (--device ROCm0) with GGML_Q4_0_HADAMARD=1 and
-#     GGML_W4A4_TENSORS=all in the backend's environment and 1024-token micro-batches (-ub 1024),
+#     GGML_W4A4_TENSORS=all in the backend's environment; a MoE file (an expert count) also gets
+#     1024-token micro-batches (-ub 1024), a dense one keeps llama-server's 512,
 #   - an unstamped file with --device auto still goes to Vulkan with neither variable,
-#   - --dflash on a rotated file adds the DFlash drafter with full blocks (n-max 16, p-min 0):
-#     one ROCm server, W4A4 prompt processing and DFlash2 decode.
+#   - --dflash on a rotated file adds the DFlash drafter with full blocks (p-min 0, n-max = the
+#     drafter's dflash.block_size - 1, or 16 when the file does not say): one ROCm server, W4A4
+#     prompt processing and DFlash2 decode.
 #
 # usage: tests/hadamard_route.sh path/to/1bit
 set -uo pipefail
@@ -35,16 +37,20 @@ check() { if eval "$2"; then echo "ok   $1"; else echo "FAIL $1"; fail=1; fi; }
 
 python3 - "$scratch" <<'PY'
 import struct, sys
-def gguf(path, stamp):
+def gguf(path, stamp, arch="llama", ints={}):
     kv = []
     def s(x): b = x.encode(); return struct.pack("<Q", len(b)) + b
-    kv.append(s("general.architecture") + struct.pack("<I", 8) + s("llama"))
+    kv.append(s("general.architecture") + struct.pack("<I", 8) + s(arch))
     if stamp:
         kv.append(s("onebit.hadamard_q4_0") + struct.pack("<I", 5) + struct.pack("<i", 32))
+    for k, v in ints.items():
+        kv.append(s(k) + struct.pack("<I", 4) + struct.pack("<I", v))
     open(path, "wb").write(b"GGUF" + struct.pack("<IQQ", 3, 0, len(kv)) + b"".join(kv))
 gguf(sys.argv[1] + "/h32.gguf", True)
+gguf(sys.argv[1] + "/h32moe.gguf", True, "qwen3moe", {"qwen3moe.expert_count": 128})
 gguf(sys.argv[1] + "/plain.gguf", False)
 gguf(sys.argv[1] + "/draft.gguf", False)
+gguf(sys.argv[1] + "/draft8.gguf", False, "dflash", {"dflash.block_size": 8})
 PY
 
 # a backend that records how it was started, then answers /health
@@ -84,7 +90,10 @@ run "$scratch/h32.gguf" "$scratch/h32.json"
 check "--device auto sends a rotated file to ROCm0" '[ "$(field "$scratch/h32.json" "r[\"argv\"][r[\"argv\"].index(\"--device\")+1]")" = ROCm0 ]'
 check "  with GGML_Q4_0_HADAMARD=1" '[ "$(field "$scratch/h32.json" "r[\"env\"].get(\"GGML_Q4_0_HADAMARD\")")" = 1 ]'
 check "  and GGML_W4A4_TENSORS=all" '[ "$(field "$scratch/h32.json" "r[\"env\"].get(\"GGML_W4A4_TENSORS\")")" = all ]'
-check "  with 1024-token micro-batches" '[ "$(field "$scratch/h32.json" "r[\"argv\"][r[\"argv\"].index(\"-ub\")+1]")" = 1024 ]'
+check "  dense: llama-server's own micro-batch" '[ "$(field "$scratch/h32.json" "\"-ub\" in r[\"argv\"]")" = False ]'
+
+run "$scratch/h32moe.gguf" "$scratch/h32moe.json"
+check "a rotated MoE file gets 1024-token micro-batches" '[ "$(field "$scratch/h32moe.json" "r[\"argv\"][r[\"argv\"].index(\"-ub\")+1]")" = 1024 ]'
 
 run "$scratch/plain.gguf" "$scratch/plain.json"
 check "an unstamped file still goes to Vulkan0" '[ "$(field "$scratch/plain.json" "r[\"argv\"][r[\"argv\"].index(\"--device\")+1]")" = Vulkan0 ]'
@@ -94,8 +103,11 @@ run "$scratch/h32.gguf" "$scratch/df.json" --dflash "$scratch/draft.gguf"
 after() { field "$scratch/df.json" "r[\"argv\"][r[\"argv\"].index(\"$1\")+1]"; }
 check "--dflash on a rotated file stays on ROCm0" '[ "$(after --device)" = ROCm0 ]'
 check "  with the DFlash drafter" '[ "$(after --spec-type)" = draft-dflash ] && [ "$(after -md)" = "$scratch/draft.gguf" ]'
-check "  full blocks: n-max 16, p-min 0" '[ "$(after --spec-draft-n-max)" = 16 ] && [ "$(after --spec-draft-p-min)" = 0 ]'
+check "  full blocks: n-max 16 without a block size, p-min 0" '[ "$(after --spec-draft-n-max)" = 16 ] && [ "$(after --spec-draft-p-min)" = 0 ]'
 check "  and the Hadamard W4A4 environment" '[ "$(field "$scratch/df.json" "r[\"env\"].get(\"GGML_W4A4_TENSORS\")")" = all ]'
 
+run "$scratch/h32.gguf" "$scratch/df8.json" --dflash "$scratch/draft8.gguf"
+check "--dflash with a block-8 drafter drafts 7" '[ "$(field "$scratch/df8.json" "r[\"argv\"][r[\"argv\"].index(\"--spec-draft-n-max\")+1]")" = 7 ]'
+
 if [ $fail -ne 0 ]; then for f in "$scratch"/*.log; do echo "--- $f"; cat "$f"; done; echo FAIL; exit 1; fi
 echo PASS

tools/site.py

@@ -72,7 +72,7 @@
 # The home page's search and social description: the headline milestones, measured on Strix Halo
 # (docs/serve.md, docs/lean.md). Keep it <= 160 characters and update it with each milestone.
 HOME_DESCRIPTION = ("LLM inference for AMD Strix Halo inside Lemonade: one server runs Qwen3.8-27B with "
-                    "Hadamard W4A4 prompts (445 tok/s) and DFlash2 decode (40.9 tok/s).")
+                    "Hadamard W4A4 prompts (477 tok/s) and DFlash2 decode (41.1 tok/s).")
 HOME_KEYWORDS = ("LLM inference, AMD Ryzen AI, Strix Halo, XDNA 2 NPU, Vulkan, ROCm, HRX, Lemonade, "
                  "OpenAI-compatible API, GGUF, speculative decoding, DFlash2, MTP, W4A4, Hadamard rotation, "
                  "Laya router, Qwen3.8, Unsloth")