DFlash speculative decoding for Qwen3.6/3.5 (Rust + MLX)

[dusterbloom]

What

First Rust + MLX implementation of DFlash speculative decoding (tap-fed drafter → batched target verify) for Qwen3.6-35B-A3B (MoE) and Qwen3.5-9B (dense). On its strong workloads it matches the acceptance of the community Python-MLX DFlash ports.

Pipeline: hidden states tapped from a handful of target layers feed a small DFlash drafter that proposes a block of tokens; the target verifies the block in one batched pass; the accepted prefix is committed and rejected positions are rolled back. Qwen3Next GDN (linear-attention) state is made bit-exact across the multi-token verify via a tape-record/replay kernel.

Key pieces:

• crates/higgs-models/src/dflash.rs — drafter module (loads modal-labs DFlash weights), accept_prefix, drafter-cache crop.
• crates/higgs-models/src/qwen3_next.rs — forward_with_taps / forward_with_taps_tape / replay_tape_rollback + GDN tape kernels.
• crates/higgs-engine/src/simple.rs — draft-verify loop, entropy-adaptive block size, and an auto-calibrated realized-speedup gate that floors DFlash→plain-AR when speculation is actually slower than AR for the current text/thermal state (keeping AR regions byte-exact), plus a drafter sliding-window port for long context.
• crates/higgs/src/config.rs + doctor.rs — draft_model config field (simple engine only; doctor validates).

Results

See docs/benchmarking.md → "DFlash Speculative Decoding". Acceptance is governed by the target's output entropy, not the task label. Three regimes, consistent across both models:

• Deterministic (H < 0.04 bits): accept_len 14–15.4, byte-exact.
• Structured / code (H 0.04–0.4): 5–9.
• Prose (H > 0.7): ~2–3 (the gate floors to ~parity).

accept_len / entropy / byte_exact are thermal-independent and carry the characterization; absolute tok/s is thermally confounded on a laptop and is not quoted as a headline.

Reproduce (Apple Silicon + MLX required)

Models (public on HF): target mlx-community/Qwen3.6-35B-A3B-4bit (or mlx-community/Qwen3.5-9B-MLX-4bit) + drafter modal-labs/Qwen3.6-35B-A3B-DFlash (or modal-labs/Qwen3.5-9B-DFlash).

HIGGS_DFLASH_TARGET_DIR=<target-mlx-dir> HIGGS_DFLASH_DRAFTER_DIR=<dflash-drafter-dir> \
cargo test --release -p higgs-engine dflash_entropy_sweep -- --ignored --nocapture

Byte-exactness vs plain AR greedy: dflash_matches_ar_greedy_and_reports_accept_len (same env vars).

Notes for review

• Large branch with progressive commits (P1→P5: drafter port → GDN tape bit-exactness → draft-verify loop → realized-speedup gate → sliding window → entropy-sweep harness) — reviewing commit-by-commit is easiest.
• DFlash is simple-engine only (incompatible with the batch engine).
• A follow-up PR adds per-request MTP/DFlash selection + DFlash streaming on top of this branch.

🤖 Generated with Claude Code

Summary by CodeRabbit

Release Notes

• New Features

  • Added DFlash speculative decoding to speed up token generation when a draft model is available and generation settings are compatible.
  • Exposed DFlash acceptance metrics (via a new getter) after generation runs.

• Configuration

  • Added optional per-model draft_model to enable DFlash; also supports discovery via HIGGS_DFLASH_PATH and block sizing via HIGGS_DFLASH_BLOCK_SIZE.

• Validation

  • Enhanced model “doctor” checks to verify draft-model directories and weight files.

• Documentation

  • Added a DFlash benchmarking guide with acceptance and speedup metrics.

• Tests

  • Updated DFlash-related benchmark/sweep coverage (ignored manual runs where applicable).

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds end-to-end DFlash block-diffusion speculative decoding. A new DFlashDrafter model with dual-stream attention is implemented in higgs-models. AnyModel gains tap/tape forward dispatcher APIs. SimpleEngine is extended with load_with_dflash and generate_dflash_inner. A draft_model config field is wired through config, doctor validation, and engine loading. Benchmarking documentation is added.

Changes

DFlash Speculative Decoding

Layer / File(s)	Summary
DFlashDrafter architecture and acceptance logic crates/higgs-models/src/dflash.rs	Defines DFlashConfig, a DFlashDrafter stack with dual-stream attention (Q from noise, K/V from taps + noise), per-layer KV cache with optional sliding-window eviction, GdnStateBackup for hybrid model state rollback, crop_drafter_cache, load_dflash_drafter, and accept_prefix. Unit tests cover acceptance variants and sliding-window cache behavior.
AnyModel tap/tape dispatcher APIs and AnyCache accessors crates/higgs-models/src/lib.rs	Exports the dflash module, adds AnyCache::as_hybrid/as_hybrid_mut, the TapsTapeOutput type alias, and five new AnyModel dispatcher methods (forward_with_taps, forward_with_taps_tape, replay_tape_rollback, embed_token_ids, forward_all_logits_from_hidden) routing to Qwen3Next+Hybrid cache or returning an exception.
model_loader.rs: load_dflash_drafter entry point crates/higgs-engine/src/model_loader.rs	Adds dflash module import and a public load_dflash_drafter function delegating to higgs_models::dflash and mapping errors to EngineError::Model.
SimpleEngine DFlash state and load_with_dflash crates/higgs-engine/src/simple.rs, crates/higgs-engine/Cargo.toml	Introduces DFlashState struct, adds dflash and last_dflash_accepts fields to SimpleEngine, adds module-level THINKING_BUDGET and private helpers dflash_supports_plain_greedy and enforce_thinking_budget_token. Implements load_with_dflash (loading drafter from explicit path or HIGGS_DFLASH_PATH, deriving block_size from env/config, reading mask_token_id), and exposes last_dflash_accepts accessor. Adds tracing-subscriber dev dependency.
SimpleEngine generate_dflash_inner speculative loop crates/higgs-engine/src/simple.rs	Routes generate_inner to the DFlash path when a drafter is loaded and generation is unconstrained, greedy, and text-only. Implements the full generate_dflash_inner loop: tapped prefill, first-token sampling, per-round drafter proposal + tape-recording verify + accept_prefix + cache rollback/replay, adaptive block sizing, speedup gating, EOS/stop-sequence termination, and acceptance length recording.
DFlash benchmark test harness crates/higgs-engine/src/simple.rs	Adds ignored test functions for greedy AR vs DFlash correctness verification, speed/acceptance metrics, and an entropy sweep harness over multiple prompt templates.
draft_model config field, load wiring, and engine integration crates/higgs/src/config.rs, crates/higgs/src/state.rs, crates/higgs/src/main.rs, crates/higgs/src/tui/mod.rs	Adds ModelConfig.draft_model: Option<String> (serde-defaulted), propagates it through all config construction paths, updates Engine::load_simple signature to accept draft_model: Option<&Path>, calls SimpleEngine::load_with_dflash, and updates test fixtures.
Doctor validation for draft_model directory and safetensors crates/higgs/src/doctor.rs	Adds validate_dflash_drafter_dir helper that validates draft_model directories (path existence, config.json parseable, safetensors artifacts present). Updates check_models to validate draft_model and rejects draft_model+batch=true. Adds test helpers and updates test fixtures to set draft_model: None.
DFlash benchmarking documentation docs/benchmarking.md	Adds a DFlash speculative decoding section covering setup, entropy-sweep methodology, acceptance tables for two model pairs at block size 16, analytic speedup formula, context-depth invariance results, and block-size tradeoff table.

Sequence Diagram(s)

sequenceDiagram
  participant Config as ModelConfig
  participant Main as main.rs load_engines
  participant State as Engine::load_simple
  participant Simple as SimpleEngine::load_with_dflash
  participant Loader as model_loader::load_dflash_drafter
  participant Models as higgs_models::dflash

  Main->>State: load_simple(dir, ..., draft_model: Option<&Path>)
  State->>Simple: load_with_dflash(dir, ..., dflash_path)
  Simple->>Loader: load_dflash_drafter(path)
  Loader->>Models: load_dflash_drafter(model_path)
  Models-->>Loader: DFlashDrafter
  Loader-->>Simple: DFlashDrafter
  Simple->>Simple: extract tap_layers, block_size, mask_token_id
  Simple-->>State: SimpleEngine { dflash: Some(DFlashState) }
  State-->>Main: Engine::Simple(engine)

Loading

sequenceDiagram
  participant Client as generate_inner
  participant DFlash as generate_dflash_inner
  participant Target as AnyModel forward_with_taps_tape
  participant Drafter as DFlashDrafter::forward
  participant Accept as accept_prefix

  Client->>DFlash: dispatch (dflash loaded, greedy, unconstrained)
  DFlash->>Target: forward_with_taps (prefill)
  Target-->>DFlash: logits + tap hidden states
  DFlash->>DFlash: sample first token

  loop Speculative round
    DFlash->>Drafter: forward(noise block, taps, drafter_cache)
    Drafter-->>DFlash: draft proposals
    DFlash->>Target: forward_with_taps_tape (verify)
    Target-->>DFlash: verify logits + GDN tape
    DFlash->>Accept: accept_prefix(draft, verify_argmax)
    Accept-->>DFlash: accepted prefix + bonus token
    alt Partial accept
      DFlash->>Target: replay_tape_rollback(tape, n_accepted)
      DFlash->>Drafter: crop_drafter_cache(keep_len)
    end
    DFlash->>DFlash: adapt block_size, gate on speedup
  end
  DFlash-->>Client: GenerationOutput

Loading

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

Possibly related PRs

• panbanda/higgs#96: Refactored the same Engine::load_simple/SimpleEngine::load dispatch chain that this PR extends with the draft_model parameter.
• panbanda/higgs#103: Modified Engine::load_simple and SimpleEngine::load signatures (adding raise_wired_limit), the same call sites this PR changes to thread draft_model.
• panbanda/higgs#161: Extended generate_inner speculative-decoding dispatch in simple.rs with MTP/prompt-lookup paths, the same gating location this PR adds the DFlash early-dispatch branch.

Suggested reviewers

• panbanda

🐇 A drafter hops ahead, block by block,
Proposing tokens 'round the clock.
The target verifies, accepts, or rejects,
Rollback the cache if the prefix deflects.
Entropy low? Watch acceptance soar!
DFlash decodes faster than ever before. 🚀

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately and concisely summarizes the main change: implementing DFlash speculative decoding for specific Qwen models in Rust and MLX, which is the central purpose of this substantial PR.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 6

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@crates/higgs-engine/src/simple.rs`:
- Around line 1452-1458: The generate_dflash_inner method is not receiving or
handling the enable_thinking parameter that is used in the normal generation
path to enforce thinking-budget behavior. Add the enable_thinking parameter to
the generate_dflash_inner function signature and implement the same
thinking-budget logic that enforces the 256-token limit before forcing a closing
</think> tag, mirroring the behavior in the normal generation loop. This ensures
that thinking requests on the DFlash path respect the same token budgets as the
standard path instead of continuing to reason until max_tokens is reached.
- Around line 1789-1797: Before appending tokens from the `accepted` vector in
the loop, cap the number of tokens to commit based on the remaining budget
(max_tokens minus current token count). Additionally, scan the `accepted` block
for the first EOS token and truncate the accepted list at that point before
pushing tokens to the `tokens` vector. After truncating `accepted`, recompute
`n_accepted` as the length of the truncated `accepted` list rather than the
original block size, ensuring that rollback and tap tracking remain aligned with
the actual tokens that were emitted.
- Around line 3595-3605: The SimpleEngine::load_with_dflash method reads the
HIGGS_DFLASH_PATH environment variable as a fallback even when drafter_path is
None, which means the run(false) test case could accidentally load DFlash if
that env var is set, resulting in false test passes. Either save and clear the
HIGGS_DFLASH_PATH environment variable before calling load_with_dflash in the
run(false) case and restore it afterward, or modify the loader to accept an
explicit parameter that disables the env var fallback completely. This ensures
the "AR" side of the test is truly tested without DFlash interference.
- Around line 1112-1117: The DFlash dispatch condition at the if statement
checking self.dflash.is_some() is too permissive and allows requests that DFlash
cannot properly handle. Add additional guards to the condition to ensure DFlash
is only used for plain greedy decoding without logprobs: check that logprobs are
not requested and that the SamplingParams in the params argument has only
default values (no temperature, no top-k, top-p, or penalty settings applied).
Only proceed to generate_dflash_inner when these additional conditions are
satisfied, otherwise fall back to the normal generation path.

In `@crates/higgs-models/src/dflash.rs`:
- Around line 1-11: The module documentation comment at the top of the file
incorrectly states "8 decoder layers" and "tapped from 5 target model layers",
but the actual configuration verified in the
loads_modal_drafter_against_real_weights function shows num_hidden_layers equals
6 and num_taps equals 8 with tap IDs [1,6,11,16,22,27,32,37]. Update the module
doc comment to accurately reflect these actual values (6 decoder layers and 8
taps from specific layers) and verify the reference documentation points to the
correct source that matches the test configuration being used.

In `@crates/higgs/src/doctor.rs`:
- Around line 287-304: The current validation of model.draft_model only checks
if the path exists using Path::exists(), which is insufficient because it allows
regular files or incomplete directories to pass. Enhance this validation to
ensure draft_model is a valid drafter directory by checking that it is a
directory (not a file) and contains the required drafter model files such as
config.json and weights. Replace the simple Path::exists() check with more
robust validation logic that verifies the directory structure matches what
DFlash expects when loading the drafter, catching misconfiguration before server
startup.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 329806a9-9f09-4a62-bd9f-fcd60fac01ff

📥 Commits

Reviewing files that changed from the base of the PR and between 9a7437e and 2347cfe.

⛔ Files ignored due to path filters (1)

• Cargo.lock is excluded by !**/*.lock

📒 Files selected for processing (12)

• crates/higgs-engine/Cargo.toml
• crates/higgs-engine/src/model_loader.rs
• crates/higgs-engine/src/simple.rs
• crates/higgs-models/src/dflash.rs
• crates/higgs-models/src/lib.rs
• crates/higgs-models/src/qwen3_next.rs
• crates/higgs/src/config.rs
• crates/higgs/src/doctor.rs
• crates/higgs/src/main.rs
• crates/higgs/src/state.rs
• crates/higgs/src/tui/mod.rs
• docs/benchmarking.md

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

crates/higgs/src/doctor.rs (1)

338-357: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Validate the HIGGS_DFLASH_PATH fallback too.

SimpleEngine::load_with_dflash enables DFlash from HIGGS_DFLASH_PATH when draft_model is unset, but doctor only validates the explicit config field. A bad env drafter can still pass doctor and fail at startup for simple-engine models.

🔧 Proposed fix

         if let Some(ref drafter) = model.draft_model {
             if let Err(err) = validate_dflash_drafter_dir(std::path::Path::new(drafter)) {
                 fail(
                     &format!(
                         "model {label} draft_model is not a valid DFlash drafter directory ({drafter}): {err}"
                     ),
                     result,
                 );
                 continue;
             }
             if model.batch {
                 fail(
                     &format!(
                         "model {label} sets draft_model but DFlash is simple-engine only (batch=true)"
                     ),
                     result,
                 );
                 continue;
             }
+        } else if !model.batch {
+            if let Some(drafter) = std::env::var_os("HIGGS_DFLASH_PATH")
+                .filter(|value| !value.as_os_str().is_empty())
+            {
+                let drafter_path = std::path::Path::new(&drafter);
+                if let Err(err) = validate_dflash_drafter_dir(drafter_path) {
+                    fail(
+                        &format!(
+                            "model {label} HIGGS_DFLASH_PATH is not a valid DFlash drafter directory ({}): {err}",
+                            drafter_path.display()
+                        ),
+                        result,
+                    );
+                    continue;
+                }
+            }
         }

As per coding guidelines, crates/higgs/src/**/*.rs: “When adding or changing config fields, update crates/higgs/src/doctor.rs to validate the new field. The doctor should catch misconfiguration before the server starts.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/higgs/src/doctor.rs` around lines 338 - 357, The doctor validation in
the draft_model check block only validates the explicit `draft_model`
configuration field, but SimpleEngine also uses the `HIGGS_DFLASH_PATH`
environment variable as a fallback when draft_model is unset. Add validation for
this environment variable by checking if it exists when draft_model is None,
then validate the path using the same validate_dflash_drafter_dir function. Also
ensure that if the HIGGS_DFLASH_PATH env var is set and batch=true for a model,
it fails with an appropriate error message similar to the existing batch mode
check for draft_model.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@crates/higgs/src/doctor.rs`:
- Around line 338-357: The doctor validation in the draft_model check block only
validates the explicit `draft_model` configuration field, but SimpleEngine also
uses the `HIGGS_DFLASH_PATH` environment variable as a fallback when draft_model
is unset. Add validation for this environment variable by checking if it exists
when draft_model is None, then validate the path using the same
validate_dflash_drafter_dir function. Also ensure that if the HIGGS_DFLASH_PATH
env var is set and batch=true for a model, it fails with an appropriate error
message similar to the existing batch mode check for draft_model.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 24495f19-79f2-4a3f-9ddf-c92336d21d98

📥 Commits

Reviewing files that changed from the base of the PR and between 2347cfe and 1dc0832.

📒 Files selected for processing (3)

• crates/higgs-engine/src/simple.rs
• crates/higgs-models/src/dflash.rs
• crates/higgs/src/doctor.rs

🚧 Files skipped from review as they are similar to previous changes (2)

• crates/higgs-engine/src/simple.rs
• crates/higgs-models/src/dflash.rs

[dusterbloom]

Split stack created to reduce review load:

1. DFlash split 1: Qwen3Next tap/tape foundation #205 — Qwen3Next tap/tape foundation
2. DFlash split 2: add drafter module dusterbloom/higgs#24 — DFlash drafter module, stacked on DFlash split 1: Qwen3Next tap/tape foundation #205 branch
3. DFlash split 3: wire engine draft-verify loop dusterbloom/higgs#25 — SimpleEngine draft-verify loop, stacked on chore: release main #24
4. DFlash split 4: add config validation and docs dusterbloom/higgs#26 — draft_model config, doctor validation, and docs, stacked on fix: add linked-versions plugin so mlx-server always releases with its workspace deps #25

The final split branch codex/dflash-4-config-docs has the same tree as this PR after the CodeRabbit fixes; git diff codex/pr-204-review..codex/dflash-4-config-docs is empty locally. Leaving this PR open for reference unless we decide to close it as superseded.

[dusterbloom]

Closing as superseded by the split PR stack, starting with #205.