From cd1c5507b660a2e0348520d34501dac1bdb2aafb Mon Sep 17 00:00:00 2001 From: antirez Date: Tue, 3 Nov 2026 16:50:55 +0100 Subject: [PATCH] Preserve resident Qwen prefill speed with speculative decoding Prepare predictor history by default for SSD streaming. When weights fit in RAM, initialize historical draft rows to zero and keep target verification enabled. Full prompt history remains available with DS4_QWEN4_MTP_PREFILL=on. With 7,809 prompt tokens and 64 generated tokens on an M5 Max, the resident default took 6.47 s versus 7.53 s with eager history. Generated text matched previous main in both modes. The same policy recovered prefill throughput on the DGX Spark. --- docs/QWEN38_FLASH_NEXT.md | 17 ++++++++--------- docs/SPECULATIVE_DECODING.md | 15 ++++++++------- ds4.c | 9 ++++----- tests/test_qwen4_mtp_policy.c | 8 +++----- 4 files changed, 23 insertions(+), 26 deletions(-) diff --git a/docs/QWEN38_FLASH_NEXT.md b/docs/QWEN38_FLASH_NEXT.md index b203bf66..77a718b0 100644 --- a/docs/QWEN38_FLASH_NEXT.md +++ b/docs/QWEN38_FLASH_NEXT.md @@ -158,31 +158,30 @@ default. The 1024-token example above saves memory on smaller Macs. Use the same model options with `ds4-agent` or `ds4-server`. Add `--mtp` for speculative decoding using the built-in MTP weights. -`DS4_QWEN4_MTP_PREFILL` controls preparation of the predictor's prompt history: +`DS4_QWEN4_MTP_PREFILL` controls how prefill prepares the predictor’s prompt history. | Value | Predictor history for new prompt rows | | --- | --- | -| `on` (default, also `1`) | Prepare it for both resident and SSD streaming runs | +| `auto` (default) | Prepare it only with SSD streaming | +| `on` (also `1`) | Prepare it for both resident and SSD streaming runs | | `off` (also `0`) | Skip historical projections and initialize those cache rows to zero | -| `auto` | Prepare it only with SSD streaming | -For example, this opts into streaming-only preparation: +To prepare the full predictor history when the weights fit in RAM, run ```sh -DS4_QWEN4_MTP_PREFILL=auto ./ds4 -m gguf/Qwen3.8-Flash-Next-Q2.gguf --mtp --ctx 8192 +DS4_QWEN4_MTP_PREFILL=on ./ds4 -m gguf/Qwen3.8-Flash-Next-Q2.gguf --mtp --ctx 8192 ``` Use `on` for resident workloads where the decode benefit outweighs the prefill cost, and compare total wall time on your prompts. `off` and resident `auto` keep MTP drafts and target verification enabled, but the cold predictor history can change drafts and reduce their acceptance rate. -Zero keys still participate in predictor attention; this is not equivalent to -the fully prepared history. Target prefill caches and logits are unchanged. +Zero keys still participate in predictor attention. Target prefill caches and logits are unchanged. Decode continues to prepare real predictor rows. Each prefill chunk also retains its final row for real preparation at the next boundary, keeping checkpoint restore consistent with uninterrupted prefill. -The setting applies to newly prefilled rows; -reused or restored checkpoints retain their existing predictor cache contents. +The setting applies to newly prefilled rows. Reused or restored checkpoints +retain their existing predictor cache contents. Without `--mtp`, this setting has no effect. `ds4-bench` benchmarks ordinary decoding; it does not accept `--mtp`. diff --git a/docs/SPECULATIVE_DECODING.md b/docs/SPECULATIVE_DECODING.md index 13521e06..6f008352 100644 --- a/docs/SPECULATIVE_DECODING.md +++ b/docs/SPECULATIVE_DECODING.md @@ -125,20 +125,21 @@ describes training this module over multiple steps and using QSA in it. These weights are already included in the GGUF; no separate draft model is needed. -By default, prompt prefill prepares the predictor's K/V and indexer cache from the -trunk hidden states and the next prompt token. It skips the predictor's +With SSD streaming, prompt prefill prepares the predictor's K/V and indexer +cache from the trunk hidden states and the next prompt token by default. It skips the predictor's attention output, MoE and vocabulary head, so this preparation does not stream predictor experts from SSD. The final hidden row is retained until the actual next token is known. Snapshot restore and speculative rollback -retain that row too; this costs 40 KiB per row for this model, at most +retain that row too. This costs 40 KiB per row for this model, at most 160 KiB including the three verifier snapshots. Qwen payload version 3 stores the initialized prefix and retained row. Version 2 checkpoints and plain checkpoints restore their target state directly. When MTP is enabled without a saved predictor tail, the engine replays the saved tokens once -before continuing. [`DS4_QWEN4_MTP_PREFILL`](QWEN38_FLASH_NEXT.md) can skip historical -projections or restrict preparation to SSD streaming, while keeping MTP and -target verification enabled. Skipped rows are initialized to zero; changed -draft acceptance can offset the prefill saving. +before continuing. With resident weights, prefill skips historical predictor +projections and initializes those cache rows to zero. MTP and target +verification remain enabled. Preparing the full history can improve draft +acceptance for long generations. Use `DS4_QWEN4_MTP_PREFILL=on` to enable +it for resident weights too. [Qwen MTP settings](QWEN38_FLASH_NEXT.md) The cycle drafts one token ahead by default and engages a **second, chained draft** (one extra nextn-layer step conditioned on the predictor's own diff --git a/ds4.c b/ds4.c index 3e4eee0a..ae6767ca 100644 --- a/ds4.c +++ b/ds4.c @@ -1118,15 +1118,14 @@ static ds4_context_memory qwen4_graph_memory_estimate(uint32_t ctx, uint32_t pre static int g_ds4_lock_fd = -1; -/* This controls prompt-time predictor cache preparation, not speculation. - * Keep the established predictor history by default; auto is an explicit - * streaming-only policy until resident acceptance/wall-time is measured. */ +/* Preparing the draft history costs prefill time. With streamed experts, + * a better draft can also save disk reads during generation. */ static DS4_MAYBE_UNUSED int qwen4_mtp_prefill_resolve(bool mtp_enabled, bool streaming, const char *value) { if (!mtp_enabled) return 0; - if (!value || !value[0] || !strcmp(value, "on") || !strcmp(value, "1")) return 1; + if (!value || !value[0] || !strcmp(value, "auto")) return streaming; + if (!strcmp(value, "on") || !strcmp(value, "1")) return 1; if (!strcmp(value, "off") || !strcmp(value, "0")) return 0; - if (!strcmp(value, "auto")) return streaming ? 1 : 0; return -1; } diff --git a/tests/test_qwen4_mtp_policy.c b/tests/test_qwen4_mtp_policy.c index c1967c97..0575b4f9 100644 --- a/tests/test_qwen4_mtp_policy.c +++ b/tests/test_qwen4_mtp_policy.c @@ -1,17 +1,15 @@ #define DS4_NO_GPU #include "../ds4.c" -/* Keep policy coverage independent of model weights and accelerator access. - * Defaults preserve the existing predictor behavior; only explicit auto - * changes the choice according to weight residency. */ +/* Exercise all policy choices without model weights or an accelerator. */ int main(void) { const struct { const char *value; int resident; int streaming; } cases[] = { - {NULL, 1, 1}, - {"", 1, 1}, + {NULL, 0, 1}, + {"", 0, 1}, {"on", 1, 1}, {"1", 1, 1}, {"off", 0, 0},