Default Bonsai Q2 exact dSpark verifier path - #226
Open
dusterbloom wants to merge 166 commits into
Open
Conversation
…1 green)
Step 1 of Unsloth UD mix-bit work. Adds the foundation needed to make
QLinear constructors path-aware in subsequent commits, but does NOT yet
lift overrides from config.json — that lands next.
What this commit adds:
- `Qwen3NextModelArgs::quant_overrides: BTreeMap<String, QuantizationConfig>`
populated from `quantization.<canonical_path>` entries in config.json
(e.g. `language_model.model.layers.3.self_attn.q_proj`). Empty for
uniform-bit checkpoints; non-empty for UD mix-bit configs.
- `resolve_quant_for(args, path) -> (group_size, bits)` — central lookup
that every QLinear::new site will call once Step 3 lands.
- `QuantizationConfig: PartialEq, Eq` so tests can compare overrides.
Tests:
- `test_load_args_lifts_mix_bit_overrides_into_quant_overrides` — RED.
Asserts the loader lifts each per-tensor override into `quant_overrides`.
Will go green when `load_qwen3_next_args_from_value` is generalized
beyond the existing single-key `gate_quantization_override`.
- `test_resolve_quant_for_falls_back_to_default_when_no_override` — GREEN.
Verifies fallback chain (override → global → (64,4)) on the
synthesized field; no loader changes required.
Next-session plan (steps 2-5):
2. Generalize `gate_quantization_override` -> `collect_quant_overrides`,
call from both `load_qwen3_next_args_from_value` and
`load_qwen3_5_moe_text_config_args`. Test 1 goes green.
3. Thread canonical paths through every `QLinear::new(ql, qb)` site
(12 call sites across QAttention, GatedDeltaNet, FfnBlock,
SparseMoeBlock, Qwen3NextInner, Qwen3NextCausalLM). New Test 3
asserts a synthetic DecoderLayer routes the right bits to each leaf.
4. `self_attn.o_proj`, `linear_attn.out_proj`, `linear_attn.in_proj_{a,b}`
are BF16-dense in this checkpoint and must NOT be QLinear. Convert
to `nn::Linear` (or add a passthrough variant). Riskiest commit:
changes the param tree shape, so the existing `qgemv_4bit`
decode-fast path no longer applies to those ops.
5. Synthetic-safetensors integration test exercising the whole loader
end-to-end against a tiny mix-bit fixture.
Target model: Qwen3.5-27B-UD-Q2_K_XL-mlx — default 2-bit, with 210
explicit overrides (lm_head=5b, embed=4b, GDN in_proj_qkv/z=4b on 48
layers, self_attn q/k/v=4b on 16 layers, mlp.down_proj=3b on all 64).
Without the path-aware QLinear, every overridden tensor decodes at the
wrong bit width and the model emits garbage.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Adds `collect_quant_overrides` and wires it from both `load_qwen3_next_args_from_value` and `load_qwen3_5_moe_text_config_args`. Walks `config["quantization"]` (or sibling `quantization_config`), promotes every nested object that carries `(group_size, bits)` into `quant_overrides`, and skips scalar defaults plus the `mode` field. Greens Step 1's red TDD test `test_load_args_lifts_mix_bit_overrides_into_quant_overrides`. No QLinear sites read the new map yet — Step 3 will refactor `DecoderLayer::new` and friends to call `resolve_quant_for(args, &path)`. 331/331 higgs-models lib tests pass in release mode. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Step 2 lifted the override map into args.quant_overrides from both loaders. Step 3 still needs to thread canonical paths through every QLinear::new call site, Step 4 needs BF16-dense (bits=0) passthrough for o_proj/out_proj/in_proj_a/in_proj_b/in_proj_ba, and Step 5 needs both a synthetic fixture test AND a real-checkpoint smoke run on 27B-Q2_K_XL + 35B-A3B-Q3_K_XL. Inventory of 12 QLinear sites + verified BF16-dense ground truth from both checkpoints' safetensors lives in the RECAP. Approved implementation plan at ~/.claude/plans/generic-fluttering-peach.md. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…p 3)
Refactor every Q{Linear,Embedding}::new(ql, qb) call site in qwen3_next.rs
to call resolve_quant_for(args, &canonical_path) so per-tensor mix-bit
overrides from Unsloth UD checkpoints are actually consumed. Previously
the override map was populated at load (Step 2, 94785e7) but unread.
Constructor changes:
* Qwen3NextAttention::new(args, attn_prefix)
* new_mlp_projections(args, mlp_prefix) / Qwen3NextMLP::new(args, prefix)
* SwitchMlpWeights::new(args, prefix); kept from_quant(ql, qb) for
qwen3_moe / deepseek_v2 which carry different args types.
* SparseMoeBlock::new(args, mlp_prefix) — gate + shared_expert_gate now
consult quant_overrides first, then fall back to gate_quantization
(preserved for backward compat), then global default.
* GatedDeltaNet::new(args, gdn_prefix)
* FfnBlock::{new_dense,new_moe}(args, mlp_prefix)
* DecoderLayer::new(args, layer_idx) — drops ql/qb; builds per-layer
prefix "language_model.model.layers.{i}".
* Qwen3NextInner::new(args), Qwen3NextCausalLM::new uses
"language_model.{lm_head,model.embed_tokens}".
* MtpHead::new(args) updated to use language_model.mtp.layers.{i}
paths. MTP still disabled at load for both target checkpoints.
Tests added (qwen3_next::tests):
* test_decoder_layer_routes_overrides_to_qlinears — full-attn dense
layer at idx 3: q_proj 5-bit, mlp.down_proj 3-bit, others fall back.
* test_decoder_layer_moe_routes_overrides_to_shared_expert_and_switch_mlp
— MoE at idx 3: shared_expert.down_proj 3-bit, switch_mlp.gate_proj
5-bit, mlp.gate (router) 8-bit.
333/333 release-mode lib tests green; cargo fmt clean.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Step 3 landed (68d2c12). Recap captures: * Constructor signature map for the path-aware refactor. * Step 4 implementation order: bits=0 in QLinear, init_unquantized_params, args.dense_attention_outputs flag, 5 call sites to wire. * Test 4 sketch. * Pointer back to Step 5 fixture/smoke spec in the prior recap. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…ep 4) Adds args.dense_attention_outputs flag (default false) which routes the five checkpoint-BF16-dense attention output sites (self_attn.o_proj plus GDN out_proj / in_proj_ba / in_proj_a / in_proj_b) to QLinear with bits=0 instead of consuming the resolved per-tensor quant override. QLinear / QEmbedding now branch on bits==0: - init_unquantized_params() returns weight=Param[1], scales=biases=Param[0]. Shape [0] bypasses placeholder_param_names' shape-[1] missing-param check while leaving weight at [1] until the safetensors loader populates it. - forward / as_linear take an early matmul-only path: x @ weight.transpose(). - forward_decode_fast already falls through to forward for bits != 4. Wired in load_qwen3_5_moe_text_config_args (covers both dense qwen3_5 and qwen3_5_moe). load_qwen3_next_args_from_value untouched, preserving the original qwen3_next model_type behavior. Test: test_o_proj_and_out_proj_are_bf16_in_qwen3_5 builds DecoderLayer at both a full-attention layer (idx 3) and a GDN layer (idx 0) with separate projections, asserts the five sites land at bits=0 while q/k/v_proj / in_proj_qkvz / in_proj_qkv / in_proj_z keep the global quant, and confirms scales/biases land at shape [0] for the dense sites. 334/334 release lib tests green. fmt clean. clippy clean (one pre-existing MoE-doc warning unchanged). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Captures Step 4 commit (afdc921), inventories what Commit C needs (synthetic mix-bit safetensors fixture + real-checkpoint smoke run), points at building blocks already in tree (write_qwen35_config, write_weight_index, qwen35_dense_text_config), and flags the missing piece (mlx-rs save_safetensors API location). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…st (Step 5a) Adds a self-contained integration test that builds a 4-layer Qwen3.5 dense fixture on disk (3 GDN + 1 full-attention, full_attention_interval=4), loads it via the production loader, and asserts the resolved mix-bit topology. Coverage: - lm_head=5 / embed=4 overrides applied - per-layer in_proj_qkvz=4 (fused-mode override key) applied - per-layer mlp.down_proj=3 applied - BF16-dense sites (self_attn.o_proj, linear_attn.out_proj, in_proj_b/a) load as bits=0 from siblings without .scales/.biases - full-attn q_proj falls back to the global default (group=64, bits=2) - zero [1] placeholders survive in the parameter tree - forward([1, 4] uint32) returns finite logits of shape [1, 1, vocab]
…tries (Step 5b) Real Unsloth UD checkpoints publish per-tensor mix-bit overrides for GDN projections under the on-disk split keys (`in_proj_qkv` / `in_proj_z`), but the model resolves the fused key (`in_proj_qkvz`) at QLinear construction time when running with the default (non-separate) GDN projection mode. Before this change the lookup missed, the global default (group=64, bits=2) was applied to the fused QLinear, and the runtime quantized matmul aborted on the first GDN layer: [quantized_matmul] The shapes of the weight and scales are incompatible based on bits and group_size. w.shape() == (16384, 640) and scales.shape() == (16384, 80) with group_size=64 and bits=2 `collect_quant_overrides` now walks the lifted override map for every `<prefix>.linear_attn.in_proj_qkv` entry, looks up the matching `<prefix>.linear_attn.in_proj_z`, and inserts a synthesized `<prefix>.linear_attn.in_proj_qkvz` override when both sides agree on `(group_size, bits)`. The synthesis is a no-op if the fused key is already present (explicit user override wins) or if the two sides disagree (structural impossibility for the fused QLinear; we log and skip). HTTP smoke against Brooooooklyn/Qwen3.5-27B-UD-Q2_K_XL-mlx and Brooooooklyn/Qwen3.6-35B-A3B-UD-Q3_K_XL-mlx via `higgs serve` returns coherent English completions on "The capital of France is".
Concurrent requests to co-resident models each ran mlx::eval on their own
spawn_blocking thread under a fresh with_new_default_stream(Stream::new()),
racing on MLX's shared Metal CommandEncoder (the output-array table in
set_output_array) -> EXC_BAD_ACCESS/SIGSEGV. The per-model Mutex<AnyModel>
only serializes a single model, not the co-resident set (e.g. an SLM trio).
Add a process-wide GPU gate acquired by Engine::{generate_with_thinking,
generate_streaming_with_thinking, embed}. A single-GPU host has no eval
parallelism to lose and the trio is sequential, so the cost is ~nil.
Poison-recovering so a mid-eval panic can't wedge inference.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit fbdcd2f)
Add POST /v1/models and DELETE /v1/models/{name} so operators can load and
unload MLX models while the server runs, without a restart. Opt-in via
local.allow_runtime_model_load (default off; gate behind server.api_key).
Changes are in-memory only -- the TOML config stays the source of truth.
Router.local_engines becomes an RwLock<HashMap>. resolve()/list take a read
lock and clone the Arc<Engine> out, so an in-flight request is decoupled from
map membership and a concurrent unload can never free a model mid-request.
Unload removes the map entry, drains to sole ownership, then drops (detaches
past a 30s timeout -> 202). Load resolves the path non-interactively and runs
the blocking weight load in spawn_blocking; a shared state::build_engine is
reused by both startup loading and the endpoint. Unloading the auto-router
model is refused (it holds a separate Arc).
Adds ServerError::{Conflict, Forbidden}, a doctor capability warning, the
init-template + README docs, and unit/integration coverage (guards, the
load/list/route/unload round-trip, and the drain-to-sole-ownership logic).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… (P1-P2) P1: dflash.rs config-driven drafter (loads Modal bf16 weights, verified). P2: GDN tape surface (GdnLayerTape, forward_with_tape/_stateless, replay_from_tape, tape FFI kernels) + CLM tap methods (forward_with_taps/_tape, replay_tape_rollback, embed_token_ids, forward_all_logits_from_hidden, project_logits) reconciled onto origin/main; bit-exact replay (max_diff 5.96e-8), rollback ~1.25ms/GDN-layer. Engine loop (P3/P4) pending; tape/tap methods are dead-code until wired. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…tch (P3)
P2 splice had landed the CLM tap methods inside a #[cfg(test)] mod (excluded from
the non-test build, so dead). Move them into impl Qwen3NextCausalLM. Add the P3
AnyModel dispatch (forward_with_taps/_tape, embed_token_ids, forward_all_logits_from_hidden)
and AnyCache::{as_hybrid,as_hybrid_mut}, reconciled to main (_=> wildcard arms cover the
Gemma3/Gemma4/BonsaiQ1 variants the branch predates). higgs-models compiles (lib + tests).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…py-adaptive block Engine loop (P4): - SimpleEngine::generate_dflash_inner — block-diffusion draft → tape-recording verify → accept_prefix → GDN tape-replay rollback, dispatched from generate_inner when a drafter is loaded and no constraints/pixels. - load_with_dflash(Option<&Path>) resolves the drafter from the arg or HIGGS_DFLASH_PATH; load() delegates with None. - AnyModel::replay_tape_rollback dispatch + model_loader::load_dflash_drafter adapter; explicit non-Qwen3Next match arms (clippy wildcard_enum) and doc backticks; deleted the dead per-layer replay_from_tape. Verify bit-exactness (the real bug behind DFlash diverging from AR greedy): - The tape (verify) GDN kernel ran the SSM recurrence in bf16 (static_cast<InT> between timesteps + bf16 state_out) while AR decode keeps f32 → argmax flips on close calls across 30 GDN layers. Tape kernel now keeps f32 state (reinterpret + no inter-step cast + f32 state_out), matching the plain kernel. - The replay kernel fused state*g + k*delta (one FMA) vs the forward's two-step update; split to match exactly so rollback is bit-exact (replay-vs-forward 5.96e-8 → 0.00). New regression test test_gdn_tape_forward_matches_plain_forward_bf16. Entropy-adaptive block size: - Block size adapts by EMA-smoothed utilization (accepted/block = entropy proxy): hold block_max on sustained high utilization (low-entropy → biggest win, byte-exact), shrink toward block_min on sustained low utilization (high-entropy), big blocks on uncertain transitions keep verify off near-ties. Env: HIGGS_DFLASH_ADAPTIVE, HIGGS_DFLASH_MIN_BLOCK, HIGGS_DFLASH_BLOCK_SIZE. Results (release, Qwen3.6-35B-A3B-4bit + modal-labs drafter, DFlash/MTP decode): - code (low-entropy, thinking off): 1.32x MTP, byte-exact, accept_len 13.8 - prose (thinking on): 0.80x MTP, byte-exact (vs fixed-16 0.53x) Ignored integration test dflash_matches_ar_greedy_and_reports_accept_len gates byte-identical output + reports accept_len/tok-s. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Floors DFlash to plain AR when speculation is actually slower than AR for the current text — no per-model knob, no entropy threshold. Measures T_ar (wall per AR token) live, and each spec round computes the realized speedup ratio = n_accepted * T_ar / step_wall; an EMA of that ratio < 1.0 floors to AR (re-probing spec every 32 AR steps with hysteresis). MoE's fast AR raises the break-even bar differently per model, which a fixed entropy/p1 threshold can't see — the realized wall-clock ratio can. Flooring also keeps the S>1 verify off high-entropy near-ties, so AR regions stay byte-exact. Characterization motivating it (Mac/MLX; accept curve matches RTX3090/CUDA): 35B-A3B (vs MTP): code accept 13.8 → 1.09x ✓ ; math 9.6 → 0.66x ; prose 1.9 → 0.73x 9B dense (vs AR): math accept 12.5 → 2.51x ✓ ; code 4.0 → 1.12x ; prose 1.9 → 0.70x Accept does not predict the winner (9B drafts trivial math well but code poorly, reverse of 35B); and 9B prose ran DFlash *slower than plain AR* (0.70x) — exactly the regression the gate catches: gate lifts 9B prose 0.70x → 0.89x toward AR parity. Also makes the test prompt configurable via HIGGS_TEST_PROMPT. Gate on by default; disable with HIGGS_DFLASH_GATE=0. 287 engine tests pass; clippy/fmt clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Adds an optional per-model `draft_model` path to ModelConfig, threaded through load_engines → Engine::load_simple → SimpleEngine::load_with_dflash, so DFlash speculative decoding can be enabled from the config file (overrides the HIGGS_DFLASH_PATH env fallback). Doctor validates the path exists and rejects draft_model + batch=true (DFlash is simple-engine only). All ModelConfig literal sites updated. Release build + 474/99 higgs tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…tually floors The realized-speedup gate seeded t_ar from the cold first-decode-after-prefill (~187ms vs ~25ms warm on Metal), so its ratio was permanently >1 and it never floored -- a dead no-op. On prose it speculated nearly every token (104 spec rounds / 200 tok, accept_len 1.9) at 0.45-0.76x AR. Seed t_ar from warm steady-state decodes (discard the kernel-cold first sample): t_ar 187->25ms, the gate now floors (spec_rounds 104->3) and prose parity climbs to ~0.98x. Also in this change: - tap-less floor decode: floored AR steps call `forward` instead of `forward_with_taps` except on the spec-handoff step. Byte-exact (both share `forward_raw_hidden`: identical layer loop, mask, KV/GDN cache mutation, and `project_logits`) and proven free by the new micro-bench -- the per-step taps are lazy clones dropped unevaluated (-0.5%, within noise). The original task premise (taps cause the prose gap) was refuted by measurement. - exponential backoff on the AR re-probe cadence (32->512): a sustained losing region amortizes the probe toward 0%; a winning probe resets to base. Now engages (probe_every 32->128 on prose). - log spec_rounds/probe_every on the length-limit completion + a gated GATE trace (HIGGS_DFLASH_TRACE) for t_ar/step_wall/ratio. - new #[ignore] micro-bench dflash_floor_tapless_vs_taps_per_step_cost. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Adds an #[ignore] diagnostic that feeds the AR-correct tokens through the real verify path (forward_with_taps_tape) and compares per-position argmax against sequential S=1 AR greedy, reporting the top-2 logit gap at each flip. Finding on Qwen3.6-35B-A3B-4bit + Modal drafter (prose prompt): 3/48 positions flip, first at pos 28, with top-2 gaps of exactly 0.125/0.125/0.625 — multiples of 1/8, i.e. bf16 logit granularity. The divergence is inherent S>1 batched verify drift (batched full-attention + bf16 lm_head) vs S=1 sequential AR, flipping argmax only where the model is undecided to within a few bf16 ULPs. accept_prefix and the commit are correct; the output is self-consistent greedy w.r.t. the verify forward. Not a logic bug — both tokens are ~equiprobable at those ties. The GDN tape already makes the SSM layers bit-exact; FA + bf16 output precision are the residual. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The realized-speedup gate judged the cold spec warm-up: the verify/drafter/ lm_head Metal kernels are kernel-cold (step_wall ~2x inflated, 236->121ms) and the drafter KV cache starts empty (accept climbs ~3 -> ~6 as it fills), while t_ar is already warm from the calibration fix. So ratio = n_acc*t_ar/step_wall read artificially low, ratio_ema decayed below 1.0 within ~3 rounds, and the gate floored a workload that wins 1.53x warm -- then death-spiralled (floored -> spec re-cools -> re-probe also looks bad -> stays floored). Fix, symmetric to the existing t_ar cold-sample discard: grant each spec (re-)entry a SPEC_WARMUP=3 grace where the gate stays in spec, does not update ratio_ema, and does not floor; and re-seed ratio_ema=2.0 on entry so a re-probe can't inherit a stale sub-1.0 EMA and re-floor in one round. Validated on Qwen3.5-9B + Modal drafter (code prompt): gate-ON spec_rounds 4->28, accept_len 2.25->4.75, ratio_ema holds 1.4-1.8 (no floor) -- recovering the ~1.5x that pure spec (HIGGS_DFLASH_GATE=0) delivers. 35B prose (accept ~2.0) still floors correctly when t_ar is fresh; confirmed genuine loser at 0.79x gate-off, so flooring it is the right call. Known limitation (separate, pre-existing): t_ar is calibrated once and not refreshed during a long spec streak, so under large thermal drift it can go stale (measured 25ms cool vs 101ms throttled) and skew the floor decision both ways. Periodic t_ar re-calibration is the follow-up. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…drift t_ar (warm AR per-token time, the gate's denominator reference) was measured only during AR steps, so a long spec-winning streak never refreshed it. But AR speed drifts with thermal — measured 25ms cool vs 101ms throttled on the same model — so a stale t_ar made `ratio = n_acc*t_ar/step_wall` skew and the gate could mis-decide whether spec still beats AR for the current thermal state. Every T_AR_RECAL_EVERY (48) judged winning rounds, force ONE AR step to refresh t_ar at the current thermal state. It reuses the AR-floor machinery via a `recal` window of 1; unlike a probe cooldown it does NOT re-arm the spec warm-up grace or reset ratio_ema (a single AR step barely cools the drafter, and we want to keep the winning EMA). need_taps is true on that step so current_taps stays fresh for the next draft. ~1 AR token per ~240 generated (<0.5% overhead). Validated on Qwen3.5-9B (600-token code prompt): recal fires at round ~48 (t_ar refreshed to 54.7ms), spec_rounds 84 / accept_len 4.81 — the winning streak holds with no spurious floor or death spiral. Byte-exact greedy unchanged. Also: HIGGS_TEST_MAX_TOKENS makes the dflash test's generation length configurable (needed to exercise streaks past the recal threshold). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The Modal/Z-Lab drafters are trained at 40k seq len with sliding-window attention (config: layer_types = [sliding x5, full x1], sliding_window 4096), but higgs ignored layer_types/sliding_window — all layers ran full non-causal attention with UNBOUNDED context KV. <=4096 that's a faithful no-op; >4096 the sliding layers over-attended (off-distribution) and per-round cost grew O(ctx). Port it faithfully (reference: mlx_vlm .../qwen3_dflash/dflash.py): - DFlashConfig: parse layer_types + sliding_window (serde-default None → current all-full/no-window behavior preserved when absent). - Per-layer is_sliding from layer_idx, threaded through Decoder/Attention/new. - DFlashAttention::forward evicts the oldest context to the last (sliding_window - 1) entries on sliding layers, after the cache concat and before both the cache store and the attention — so the resident prefix is bounded and on-distribution. - RoPE-under-eviction: derive cache_offset as the MAX cache length across layers (the non-evicting full-attention layer carries the true absolute position); the new context K is roped at that absolute offset and cached entries keep their baked positions, so eviction only drops the oldest. No loop threading needed, and it equals the old first-layer offset exactly when nothing evicts. Validated: - New weight-free unit test sliding_layer_cache_caps_while_full_layer_grows: sliding cache caps at window-1, full-attn cache accumulates all context. - Short context regression (9B, gate off): accept_len 6.53 / 1.50x — byte identical to pre-change (eviction is a no-op <4096, max==first offset). - Long context (9B, ~6.5k-token prompt forcing eviction on round 1, gate off): no crash/OOM with the bounded cache, accept_len 16.0 (full block) — the windowed drafter stays on-distribution at depth. Implemented by codex against a fixed spec; reviewed, built, and validated here. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The two #[ignore] diagnostics served their purpose and their conclusions are captured in commit messages + code comments, so they're now just scaffolding: - dflash_floor_tapless_vs_taps_per_step_cost: proved per-step taps are free (~0%, lazy clones) — settled, recorded in the probe-backoff comment. - dflash_verify_vs_ar_argmax_divergence: localized the bf16 near-tie flips — recorded in the prior diagnostic commit. Also drop the two confirmation-only HIGGS_DFLASH_TRACE lines (warm-up discarded, t_ar recal) now that both mechanisms are validated; the main GATE trace (ratio/ratio_ema/floor + t_ar) stays — that's the decision-relevant one for future gate tuning. Kept: sliding_layer_cache_caps_while_full_layer_grows (eviction guard) and dflash_matches_ar_greedy (byte-exactness guard). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…easured entropy) Adds the characterization harness from the spec (built by codex against a fixed spec, reviewed/verified here): - SimpleEngine::last_dflash_accepts (Mutex<Vec<u32>>): per-round accepted counts, cleared at generate_dflash_inner start, pushed after accept_prefix. Non-invasive (no GenerationOutput change). pub accessor for tests. - #[ignore] dflash_entropy_sweep: 12-prompt taxonomy spanning the entropy spectrum; per prompt an AR pass measuring top-50 Shannon entropy (bits) + top-1 prob, then a gate-off block-16 DFlash pass reporting accept_len mean + p10/p50/p90 + byte-exactness; markdown table sorted by entropy; plus context-length (512/4k/16k) and block-size (4/8/16) mini-sweeps. accept_len is thermal-robust, so the accept-vs-entropy characterization is reliable on a throttling laptop; tok/s is left to paired spot-checks. Harness only — the real-weight sweep run is a separate (multi-hour) job. Spec: .planning/measurements/dflash-entropy-sweep-plan.md Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
DFlash speculative decoding was bound at model-load and was absent from the
streaming code path, so a streamed request silently fell back to MTP even when
a DFlash drafter was loaded. There was also no way to choose the method per
request.
- Refactor the ~350-line DFlash draft-verify loop into a shared
`dflash_decode<S: DflashSink>` driver fed by two sinks: `DflashBufferedSink`
(non-streaming GenerationOutput, unchanged behavior) and `DflashStreamSink`
(per-round StreamingOutput chunks). The gate/EMA/thermal logic is written once.
- Add `Speculation { Auto, DFlash, Mtp, None }` to SamplingParams and honor it at
all four dispatch sites; streaming branches before prefill (mirrors the
non-streaming site) to avoid a double prefill. Both heads already coexist in a
drafter-loaded engine, so this is dispatch-only — no extra weights.
- Parse a per-request `speculation` field on the OpenAI chat + Anthropic routes
(400 on invalid). Default `auto` = DFlash when a drafter is loaded (incl.
streaming), else MTP.
- Verify: dflash_streaming_matches_nonstreaming loads the real 9B + drafter and
asserts streaming == non-streaming DFlash byte-for-byte (deterministic gate-off
run), plus MTP routing via last_dflash_accepts. Speculation::parse unit test.
- Docs: README + docs/configuration.md document the field and draft_model.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
`trim_drafter_cache` and `rollback_kv_only` (both pub fn in dflash.rs) had zero callers anywhere in the repo — superseded by `crop_drafter_cache` (the one the draft-verify loop uses) and the tape-replay rollback path. Remove them. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
`trim_drafter_cache` and `rollback_kv_only` (both pub fn in dflash.rs) had zero callers anywhere in the repo — superseded by `crop_drafter_cache` (the one the draft-verify loop uses) and the tape-replay rollback path. Remove them. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add a "DFlash Speculative Decoding" section: accept_len-vs-entropy tables for Qwen3.6-35B-A3B-4bit (MoE) and Qwen3.5-9B-MLX-4bit (dense) with their modal-labs drafters, the three acceptance regimes, the context-depth and block-size mini-sweeps, and the analytic speedup. Leads with thermal-independent metrics (accept_len / entropy / byte_exact); notes tok/s is thermally confounded. Reproduce via the `dflash_entropy_sweep` #[ignore] harness. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
# Conflicts: # README.md # crates/higgs/src/main.rs
A dotted `local.allow_runtime_model_load = true` written below a [server] header is parsed by TOML as server.local.* and silently dropped by serde, so the setting never takes effect -- visible only in the raw file. check_misplaced_local_keys reads the raw TOML and warns when local-only keys (local, allow_runtime_model_load, raise_wired_limit, mlx_profile) appear under [server]. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
mlx-rs pins mode="affine" in its quantized_matmul/quantize/dequantize wrappers (mlx-rs/src/ops/quantization.rs:11). MLX's C core supports "mxfp4" (E2M1 with shared block exponent) natively via the same C functions — only the Rust wrapper hides it. This adds a quant_mode module that calls mlx_sys directly, mirroring the FFI pattern in cache.rs::slice_axis, so we can pass either mode without an upstream mlx-rs change. - QuantMode enum (Affine default, MxFp4) with config.json string parser - quantize/dequantize/quantized_matmul honour the per-tensor mode - mxfp4 tensors carry no biases (null array handle to C FFI) - QLinear gains a mode field; new() defaults to Affine (unchanged hot path), new_with_mode() for per-tensor dispatch from loaders - forward_decode_fast custom kernel gated to Affine-only - Tests prove mxfp4 matmul matches unquantized reference within E2M1 error band (rel < 0.35) and quantize->dequantize round-trips
… integration" This reverts commit 704879c.
Ports MLX's qmv_fast pattern (quantized.h:844) for Q2 affine: - 32 threads per simdgroup cooperate on 4 output rows - Each thread handles VPT=16 Q2 values (one packed u32 word) - K dimension distributed across 32 threads, simd_sum reduction - block_size = VPT * 32 = 512 Microbench vs MLX stock at M=1: - lm_head (248320x5120): 1.00x (PARITY) - gate_up (17408x5120): 0.82x (18% slower) - down (5120x17408): 0.84x (16% slower) Significant improvement over the scalar kernel (22% slower everywhere). Remaining gap likely from MLX's qdot raw-mask math and load_vector helper. Not wired into production forward path — kept as standalone callable.
Copies MLX's exact qdot and load_vector from quantized.h with the critical pre-scaling trick: load_vector divides activations by the bit-position factor (4/16/64 for Q2 positions 1/2/3), then qdot uses raw byte masks (w & 0x0c etc.). The division and mask multiplication cancel out, avoiding expensive GPU bit-shifts. Helpers moved to MLX kernel header parameter (5th arg of mlx_fast_metal_kernel_new) to avoid nested function definitions. Microbench vs MLX stock at M=1: - gate_up (17408x5120): 1.51x faster - down (5120x17408): 1.15x faster - lm_head (248320x5120): 0.96x (parity) Previous scalar kernel was 0.78-0.93x (uniformly slower). The MLX qdot pattern is the key to matching and beating MLX stock for Q2.
Routes Affine bits=2 M=1 through bonsai_q2_qmv_simd (MLX qdot pattern, 1.15-1.51x faster than MLX stock per microbench). M>1 still uses MLX stock quantized_matmul for proper M-dimension weight sharing. End-to-end CanonicalS1 decode improvement on Bonsai-27B-Q2: - Before (MLX stock): 8.9 tok/s (512-token steady state) - After (simd kernel): 10.4 tok/s - Improvement: 1.17x The 1.17x end-to-end vs 1.5x kernel-level is because MLP is only 50% of decode time (profiled: GDN recurrence 33%, attention 17%).
Make the BatchedTape (block) verify schedule the default for the 2-bit Bonsai target so the validated ~12.9 tok/s decode is out-of-the-box. The 1-bit Bonsai contract is preserved: `for_drafter` stays fail-closed (dSpark with no explicit HIGGS_DFLASH_VERIFY_MODE still resolves to CanonicalS1), and the new default only applies when the target is the 2-bit variant AND a dSpark drafter is loaded. The non-greedy/penalized downgrade and out-of-domain fallback to CanonicalS1 remain intact.
This reverts commit b19a1e0.
This reverts commit 90f5efd.
This reverts commit bb3456a.
This reverts commit 7657b23.
This reverts commit 6019d1d.
This reverts commit 3835f05.
Contributor
|
Warning Review limit reached
Next review available in: 43 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (94)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Preview PR for Ternary-Bonsai-27B exact dSpark defaults.
Stack note
Opened now so the upstream maintainer can see the missing ternary default policy. This is stacked on the active Bonsai/DFlash/Q2 work, so GitHub's diff is expected to be cumulative until dependencies land. The intended payload is the exact-default work, not the local website or top-K probe archive.
Evidence
Validation
cargo check -p higgs-models -p higgs-enginecurrently reacheshiggs-enginebut is blocked by stacked base drift (GenerationOutput.reasoning_contentandDiskPrefixCache::with_max_bytes/kv_cache_bytesmismatches), not by the ternary Q2 path.