Repository navigation
feat(speculative): draft and verify a token tree, behind a flag, default off - #1212
Merged
Merged
Conversation
…pt walk First increment of tree drafting. Everything here is pure: a topology, the attention mask derived from it, and the walk over target argmaxes. Nothing touches a device, a cache or a model, so it is testable without a checkpoint and reviewable before the plumbing exists. No path calls it yet. Two equivalences carry the design and are what the tests pin. A linear tree masks exactly as a causal mask, and it walks exactly as speculative_walk does, matching both the accept count and the emitted tokens at every accept length. Those are what let the tree path be switched on later without changing output on the configurations it subsumes. The walk yields a path rather than a set. A node can have several children but the target emits one argmax per position, so at most one child matches, which keeps rollback a path-compaction problem instead of a subtree one. Pinned by a test that gives two siblings the same token. Recorded while reading rollback, because it decides feasibility rather than being an implementation detail: for a linear draft the accepted set is a prefix, which is the only reason kv.trim(trim) is correct. A tree's accepted path is scattered across the verified positions, and KVCache exposes trim_to(new_len) and nothing else. The recurrent side is in better shape than the issue assumed: Qwen 3.5 already rebuilds GatedDeltaNet state by replaying the captured block over the accepted prefix rather than trimming it, so replaying a gathered path is an indexing change, not a new mechanism. Refs #1204
The rollback primitive tree drafting needs. A linear speculative accept set is a prefix, so rolling it back is a tail trim; a draft tree's accepted path is scattered across the verified positions, so the survivors have to be selected and closed up. KVCache offered trim_to(new_len) and nothing else. gather_positions reuses the sequence axis the existing sink-preserving front trim already reindexes, and is scoped exactly as that one is: Fp16 and Int8 only. Turbo modes carry per-token rotation state in packed sidecars that reordering cannot rebuild, and pool-backed caches hold no dense buffer, so both refuse with a message rather than reindexing half their state. That refusal is the contract: a caller that gets it must decline tree drafting rather than proceed. Int8 reindexes its scale sidecars with the tokens. Leaving them behind would decode every survivor against the wrong scale, and the raw codes cannot catch it: a per-token scale normalizes every single-valued token to the same code, so the test asserts on the sidecar shapes instead of on the keys. A prefix selection is asserted equal to trim_to, which is the equivalence that lets one rollback path serve both the chain and the tree once the tree is wired up. The identity selection short-circuits, so a full-accept round pays no buffer rewrite. Refs #1204
The Fp16-and-Int8 restriction the first version shipped was inherited from trim_front_keep without checking whether its cause applies here. It does not. That restriction exists because a head trim advances live_start instead of moving data, and the Turbo fetch paths slice from 0. gather_positions reindexes the buffers physically and leaves live_start at 0, so the hazard is absent. What actually decides support is where each mode keeps its per-token state, and every dense mode keeps it on the same sequence axis: Fp16 and Int8 in keys/values plus Int8's scales, and the Turbo modes in v_packed, v_norms and v_rescale, with k_packed and k_norms alongside them in symmetric Turbo4. trim already slices all of those on axis 2, so selecting a subset of that axis is the same operation with an index list instead of a bound. Every sequence-axis buffer is now reindexed whether or not the current mode populates it, because reindexing a subset leaves survivors decoded against another token's scale, norm or sign vector, which fails silently. Turbo4Delegated still refuses, for a layout reason rather than a quantization one: its V is split between packed cold storage and a hot ring at cold_offset, and a selection spanning that boundary has to rebuild it. The refusal names the boundary so the next reader does not re-derive why. Tested across Turbo4, Turbo4Asym and Turbo3Asym with a head_dim of 8 rather than the plain fixture's 1, since Turbo3Asym packs 24 bits across the feature axis and a narrower fixture would have tested the fixture. Refs #1204
The `[n, n]` mask the topology commit shipped is not the shape a target forward takes. A real mask is `[n, offset + n]`: query row q sits at logical position q + offset, the first `offset` columns are the KV history already in the cache, and only the last n follow the tree. Every node attends the history in full, because the tree branches inside the block being verified and everything before it is shared context. Within the block a node sees its ancestors and itself and no sibling, as before. The convention matches crate::utils::create_causal_mask exactly, which the third equivalence pins: a linear tree's history mask equals a causal mask built with the same offset, element for element. Together with the two the topology commit already had, that covers what the tree path has to reproduce before it can be switched on. The seam it feeds is already open. Gemma 4's forward_hidden_with_speculative_sinks takes `mask: Option<&MlxArray>` and threads it straight through to the layered forward; the MTP adapter passes None today. Wiring is passing this instead, not adding a parameter. Refs #1204
…herits MtpTarget gains verify_forward_tree, and its default refuses. That default is the safety property rather than a placeholder: flattening a tree into the sequence a linear verify expects lets a node's state reach its sibling, which the tree mask prevents for attention and nothing prevents for a recurrence. A target opts in by implementing the method, so omission means no rather than meaning silently wrong. Twenty types implement verify_forward, so a new method with a refusing default is also what keeps this from being a twenty-file signature change. Qwen 3.5 overrides it anyway, to state its cause where a reader looks for it. With full_attention_interval 4 over 64 layers it is 48 GatedDeltaNet and 16 attention, and a recurrence is a fold with no mask: a flattened tree folds state through siblings and answers wrongly with no signal. Verifying a tree there needs a scan that forks state at branch points, which does not exist. The refusal carries a reason rather than being a bare None, so the round loop can log why tree drafting fell back to the linear path instead of leaving an operator to work it out. Refs #1204
Gemma 4 can verify a tree because it has no recurrent state: every layer is attention, so the tree mask makes siblings invisible to each other and one flattened forward over the nodes is exact. That is the whole reason the redirect on #1204 moved to this target. verify_forward's body becomes verify_forward_with_mask, an inherent method taking Option<&MlxArray>, and both entry points call it. The linear path passes None and the model builds its own causal mask exactly as before; the tree path passes a mask no causal mask can express. Everything after the forward is shared, so the two differ by that argument and nothing else, which is what keeps the linear path bit-identical. The mask is sized against the cache the next forward will read. Gemma4Wrapper gains sequence_kv_offset, returning the same first_present_cache_offset the model derives when it builds its own mask, so a supplied mask lines up with the cache instead of relying on the round loop to thread a number that could drift from it. The blocking value is -1e9 rather than -inf: an all-masked row through a softmax turns -inf into NaN, and a draft tree's mask has rows that mask most of the block. Failure sets are unchanged against baseline by set comparison, and the live pairing still runs at 95.41 tok/s. Nothing calls the tree path yet; the round loop still drafts linearly. Refs #1204
Drafter::draft_tree wraps draft_block into a linear tree unless a drafter overrides it. The degenerate default is not a convenience: it is what makes the tree wiring testable, because a linear tree has to produce byte-identical output to the chain path it subsumes, so any divergence once the round loop uses this is the wiring rather than the topology. Branching will need a signal the drafter already computes and throws away. Its per-step logits are reduced to an argmax today, and the runner-up at an uncertain step is exactly the alternative worth spending a verify position on. Refs #1204
Verifying a tree and keeping its accepted path are separate abilities, and the branch so far only had a default for the first. `verify_forward_tree` refuses unless a target opts in; `verify_finalize` had no counterpart, so a target that opted into the forward would have inherited the linear rollback, which trims a tail. Applied to a scattered accept set that keeps the wrong entries rather than failing. `verify_finalize_tree` takes the walk's path instead of an accept count and refuses by default. For a linear tree the path is `[0, 1, ..., accepted]`, the selection is a prefix, and the result has to be byte-identical to `verify_finalize(accepted, block_size, captured)`. That equivalence is the first thing to test once the round loop is wired, and a divergence there is wiring rather than topology. Accepting the forward and refusing the rollback is a real combination, not a contradiction, and Gemma 4 is currently in it. Its sliding-window layers hold `RotatingKVCache`, whose rollback rewinds `offset` and `idx` without moving data, which a scattered selection cannot use, and which has no `gather_positions`. So the two halves have to be probed as one capability before a tree is drafted: once a tree forward has appended its block, a refusal at the rollback leaves the cache in a state only a tree-aware rollback can undo. `captured` is consumed rather than handed back for that reason. Qwen 3.5 overrides it to name its own cause rather than inherit the message, matching what it already does for the forward. Unreachable while the forward refuses first, and there so a forward that stops refusing cannot silently inherit a rollback that cannot serve it. Also fixes an `i32 -> i32` cast in tree_tests that fails `clippy -D warnings`. Refs #1204
`gather_positions` names the survivors of the whole live window, so using it for a tree rollback means re-listing the entire history every round and paying a full buffer rewrite for a change confined to the last `block_size` slots. It also replaces each buffer with a `take` result sized exactly to the survivors, so the next append has to grow the allocation again. On a long context that is the rollback cost #1209 is about, and the tree path would have inherited all of it. `gather_within_tail(tail_len, kept)` takes offsets within the block and rewrites only `[tail_start, tail_start + kept_len)`, through the same `slice_update` the append path already uses. Cost scales with the verify block rather than with the context, and the allocation survives. A prefix selection is asserted equal to `trim`, and the same selection said in absolute terms is asserted equal to `gather_positions`, so the cheap path is the same rollback rather than a second one. It also unblocks Gemma 4, which is the only target that can verify a tree today and could not keep one. Its sliding-window layers hold `RotatingKVCache`, which had no gather at all, and whose `trim` cannot be generalized into one: it moves no data, it rewinds `offset` and `idx` and lets the next append overwrite. A tail that has to be compacted has to be written. The rotating form requires an unwrapped ring, which is the regime `enable_speculative_buffer` already keeps the verify block inside, and refuses by name once the ring has wrapped rather than compacting slots whose physical order is no longer logical order. Unlike `gather_positions` this places no requirement on `live_start`: the tail is located through `buffer_idx()`, which already accounts for a trimmed head. Both refuse `Turbo4Delegated` and the pool-backed cache for the reasons `gather_positions` refuses them. Refs #1204, #1209
…ted path Gemma 4 could verify a draft tree since the tree mask landed, and could not keep one: its rollback trims a tail, which is the accept set only when the accept set is a prefix. Applied to a path it keeps the wrong entries, and keeping the wrong entries is silent. `verify_finalize_tree` selects instead of trimming, on the three things a round carries forward. The target cache goes through the new `gather_speculative_cache`, which hands each layer's cache the accepted nodes as block offsets; they are the same numbers, because the verify block is appended in the tree's topological order. The captured shared K/V goes through `gather_shared_kv`, which keeps the history in front of the block untouched and compacts the block region. The drafter's next hidden is read at the last accepted node rather than at an accept count. On a linear tree every one of those reduces to the linear path: the cache gather equals a trim, the shared-K/V gather equals the tail slice, and the hidden position is the same. That is what the round loop should check first once it is wired, and a divergence there is wiring rather than topology. The capability is per round, not per model, which is the part that constrains the wiring. The sliding-window layers hold a rotating cache whose gather needs an unwrapped ring, so `can_gather_speculative_cache` answers from state and the answer can change between rounds. The check inside `verify_finalize_tree` is a guard that turns a wrapped ring into a named refusal rather than a wrong window; it cannot be the contract, because by then the block is on the cache and a refusal has no recovery. The round loop has to ask before it drafts. No batched tree path, and `gather_speculative_cache` deliberately has no per-row tail-zero: a per-row path is not a per-row count, so rows cannot share one global trim the way the linear batched rollback lets them. Refs #1204
The two tree methods are one capability and the round loop cannot discover that by trying. `verify_forward_tree` appends the tree's block to the cache; if `verify_finalize_tree` then refuses, the round is stuck, because the block is on the cache and only a tree-aware rollback takes it off. There is no falling back to the linear path from there, so the question has to be asked before anything is drafted. `tree_round_is_available` is a method rather than a constant because the answer moves. Gemma 4 can verify a tree at any time and can only roll one back while its sliding-window rings are unwrapped, so it answers from cache state and can say yes this round and no the next. The default is false: a target opts in by answering, the way it opts into the methods themselves. Refs #1204
…ault off The round loop only ever proposed a chain. Every piece the tree path needs existed and nothing called it. Four points in the round now branch and the rest is shared: the drafter call (`draft_tree` against `draft_block`), the verify (`verify_forward_tree` against `verify_forward`), the walk (`DraftTree::walk` against `speculative_walk`), and the rollback (`verify_finalize_tree` against `verify_finalize`). Emission, logprobs, EOS and diagnostics see one shape either way, because the tree walk is reduced to a `WalkResult` and the accepted node path is carried alongside it rather than threaded through everything downstream. The capability question is asked before anything is drafted, not after the verify. By then the forward has appended the block and only a tree-aware rollback can take it off, so there is no falling back to the linear path; a refusal at that point ends the session instead of continuing against a tail that no longer means what the drafter thinks. `record_round` now counts proposals as `actual_bs - 1` instead of `draft_tokens.len()`. Identical on the chain path by construction, and the only form that is right for a tree, whose proposals are not a flat vector. The stateful-drafter history hook is skipped on a tree round and says so. It trims against the chain the drafter proposed, and a tree's accepted set is a path through a branching proposal rather than a prefix of one chain. The pairing is unreachable today, since the one tree-capable target leaves `verify_hidden_full` as `None`, so this warns rather than guessing at a semantic nobody needs yet. `MLXCEL_MTP_TREE` switches it on and is off by default; `with_tree_drafting` overrides it for tests, which cannot mutate the environment safely from a parallel suite. The test that matters is the equivalence. With a linear tree, one script covering full accept, mid-block divergence and immediate rejection produces byte-identical tokens and an identical rollback log on both arms, and asserts the script actually exercised a partial accept and a zero accept so the comparison cannot pass vacuously. Refs #1204
The wiring landed and then declined every round, on every checkpoint, while the flag looked on. Two separate reasons, both of which a linear-tree equivalence check passes straight through, because a session that silently stays on the chain path trivially matches the chain path. That is why the announcement below exists. The rotating cache gated its gather on `is_trimmable`, which excludes buffered mode outright. `enable_speculative_buffer` is what puts the cache in buffered mode, and it is called precisely so an MTP verify block can append and roll back, so the gate was false in exactly the regime it was written for. The right question is whether the tail is still in logical order, and buffered mode answers yes by construction: it compacts from the front instead of wrapping, writes new tokens at `[idx, idx + n)`, and reports `idx` as the live length. `tail_is_in_logical_order` asks that, and two tests now cover the buffered regime, which the original ones did not touch. The multimodal adapters wrap the text adapter and delegate every trait method except the three tree ones, which they inherited as refusals. So tree drafting was dead for every multimodal Gemma 4 checkpoint, including the 12B the sweep runs on, while the text-only path would have worked. Nothing about a tree justifies that split; the vision tower is not in the verify forward. The generator now says once per session which path it took, and warns when the flag is set and the target declined. Before this, the two were indistinguishable from the outside, and "the flag is set" was being read as "the tree path ran". Verified on a real checkpoint rather than on the mock: gemma-4-12b-it-4bit with the 4-bit assistant drafter at block 5, 96 tokens, temperature 0. Tree on and tree off produce byte-identical output over 800 bytes, with identical round diagnostics: 39 rounds, 117 proposed, 56 accepted, 2.4359 emitted per verify, 17 zero-accept, 7 partial, 15 full. The tree arm logs that it engaged and the chain arm does not, so the comparison is not the vacuous one. Refs #1204
A caller mask says which keys are structurally allowed. It cannot also say the sliding-window band, because the same mask is handed to the full-attention family, which must not have one. Below the window the two constraints agree, so nothing noticed. Past it they stop agreeing, and the sliding family was getting only the caller's half. The MTP draft-tree mask is what exposed it. Its history columns are all-attend by construction and `trim_mask_to_keys` crops from the right without adding a band, so past the window the sliding layers attended outside their window. Measured on gemma-4-12b-it-4bit with the 4-bit assistant drafter at block 5: at 96 generated tokens the tree and chain paths are byte-identical, and at 1400 they diverge from character 5133, which is a little past the 1024-token window. Acceptance was unaffected across the divergence (0.5209 against 0.5147), which is what ruled the rollback out and pointed at the mask; a broken rollback collapses acceptance rather than producing a coherent alternative continuation. The chain arm is deterministic run to run at that length, so the difference was systematic rather than the f16 jitter class. The sliding family now gets the intersection of both constraints, and only where they can disagree: the guard is `sliding_live_len + l > window`, so every context under the window takes exactly the path it took before. Additive masks intersect by `minimum`, and the intersection always keeps each query's own position, so no row can go all-masked and softmax to NaN. After the fix the same 1400-token comparison is byte-identical across the window, and the chain arm's own output is unchanged from before it, which is the check that the guard really did scope the change to the broken regime. The root suite's failure set is unchanged (the 13 pre-existing local reds). Refs #1204
The band intersection was argued, not asserted, and the argument is what makes it safe to apply to every caller-mask path rather than only to the tree. Pulled it into `sliding_mask_from_caller` so the properties can be stated directly. Below the window the caller mask comes back untouched, which is what scopes the change: every context that fits inside the window takes exactly the path it took before. Past the window an all-attend caller mask, the draft-tree shape, reduces to the band. Past the window a mask that already carries the band is unchanged by intersecting it again, which is the property the batched and adopted-prefix prefill paths depend on, since they build window-aware masks themselves. Refs #1204
This was referenced Aug 18, 2026
inureyes
added
the
type:enhancement
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
Add MiniCPM-V 4.6, a vision-language model that pairs a Qwen 3.5 text backbone (standard 1D RoPE) with a two-stage visual compression pipeline. Ported from the upstream mlx-vlm MiniCPM-V 4.6 model (PR lablup#1212). The vision pipeline compresses tokens in two stages: a VitMerger doing cross-attention spatial 2x2 pooling inserted after ViT layer 6, followed by a pixel-shuffle Merger MLP. For a 448x448 image this produces 64 visual tokens, matching `query_num`. New files: - `src/vision/encoders/minicpmv4_6.rs`: VitMerger, MergerBlock, Merger, and MiniCPMV46VisionModel with per-layer ViT access for mid-tower injection. - `src/vision/minicpmv4_6_vl.rs`: MiniCPMV46VLModel wrapping Qwen35Model. Wiring: - `src/models/mod.rs`: register the MiniCPMV46VLM variant, its display name, and the `metadata()` arm, and add it to `ALL_MODEL_TYPES` so it appears in `mlxcel list`. - `src/models/detection.rs`: map the `minicpmv4_6` model-type string to `ModelType::MiniCPMV46VLM`. - `src/loading/vlm_special.rs`: add the MiniCPM-V 4.6 weight remap, wired in through `src/loading/vlm.rs` and `src/loading/mod.rs`. - `src/vision/encoders/minicpmo.rs`: expose forward_embeddings / layer / post_layernorm / num_layers for per-layer ViT access by the 4.6 path. - `src/loaded_model.rs`, `src/loaded_model_capabilities.rs`, `src/model_metadata.rs`, `src/multimodal/vlm_runtime.rs`, `src/distributed/tensor_parallel/inference.rs`: route the new model through capabilities, metadata, the multimodal runtime, and tensor-parallel inference. Verified with cargo fmt, cargo clippy (metal,accelerate, all-targets, -D warnings), 17 unit tests (vision encoder config and merger token-reduction, the weight remap, and the model-type registry integrity tests), and `mlxcel list` showing the new model.
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
…ult off (lablup#1212) * feat(speculative): add the draft-tree topology, its mask and its accept walk First increment of tree drafting. Everything here is pure: a topology, the attention mask derived from it, and the walk over target argmaxes. Nothing touches a device, a cache or a model, so it is testable without a checkpoint and reviewable before the plumbing exists. No path calls it yet. Two equivalences carry the design and are what the tests pin. A linear tree masks exactly as a causal mask, and it walks exactly as speculative_walk does, matching both the accept count and the emitted tokens at every accept length. Those are what let the tree path be switched on later without changing output on the configurations it subsumes. The walk yields a path rather than a set. A node can have several children but the target emits one argmax per position, so at most one child matches, which keeps rollback a path-compaction problem instead of a subtree one. Pinned by a test that gives two siblings the same token. Recorded while reading rollback, because it decides feasibility rather than being an implementation detail: for a linear draft the accepted set is a prefix, which is the only reason kv.trim(trim) is correct. A tree's accepted path is scattered across the verified positions, and KVCache exposes trim_to(new_len) and nothing else. The recurrent side is in better shape than the issue assumed: Qwen 3.5 already rebuilds GatedDeltaNet state by replaying the captured block over the accepted prefix rather than trimming it, so replaying a gathered path is an indexing change, not a new mechanism. Refs lablup#1204 * feat(cache): select and compact a scattered set of KV positions The rollback primitive tree drafting needs. A linear speculative accept set is a prefix, so rolling it back is a tail trim; a draft tree's accepted path is scattered across the verified positions, so the survivors have to be selected and closed up. KVCache offered trim_to(new_len) and nothing else. gather_positions reuses the sequence axis the existing sink-preserving front trim already reindexes, and is scoped exactly as that one is: Fp16 and Int8 only. Turbo modes carry per-token rotation state in packed sidecars that reordering cannot rebuild, and pool-backed caches hold no dense buffer, so both refuse with a message rather than reindexing half their state. That refusal is the contract: a caller that gets it must decline tree drafting rather than proceed. Int8 reindexes its scale sidecars with the tokens. Leaving them behind would decode every survivor against the wrong scale, and the raw codes cannot catch it: a per-token scale normalizes every single-valued token to the same code, so the test asserts on the sidecar shapes instead of on the keys. A prefix selection is asserted equal to trim_to, which is the equivalence that lets one rollback path serve both the chain and the tree once the tree is wired up. The identity selection short-circuits, so a full-accept round pays no buffer rewrite. Refs lablup#1204 * feat(cache): support the Turbo quantization modes in gather_positions The Fp16-and-Int8 restriction the first version shipped was inherited from trim_front_keep without checking whether its cause applies here. It does not. That restriction exists because a head trim advances live_start instead of moving data, and the Turbo fetch paths slice from 0. gather_positions reindexes the buffers physically and leaves live_start at 0, so the hazard is absent. What actually decides support is where each mode keeps its per-token state, and every dense mode keeps it on the same sequence axis: Fp16 and Int8 in keys/values plus Int8's scales, and the Turbo modes in v_packed, v_norms and v_rescale, with k_packed and k_norms alongside them in symmetric Turbo4. trim already slices all of those on axis 2, so selecting a subset of that axis is the same operation with an index list instead of a bound. Every sequence-axis buffer is now reindexed whether or not the current mode populates it, because reindexing a subset leaves survivors decoded against another token's scale, norm or sign vector, which fails silently. Turbo4Delegated still refuses, for a layout reason rather than a quantization one: its V is split between packed cold storage and a hot ring at cold_offset, and a selection spanning that boundary has to rebuild it. The refusal names the boundary so the next reader does not re-derive why. Tested across Turbo4, Turbo4Asym and Turbo3Asym with a head_dim of 8 rather than the plain fixture's 1, since Turbo3Asym packs 24 bits across the feature axis and a narrower fixture would have tested the fixture. Refs lablup#1204 * feat(speculative): give the draft tree a history-aware attention mask The `[n, n]` mask the topology commit shipped is not the shape a target forward takes. A real mask is `[n, offset + n]`: query row q sits at logical position q + offset, the first `offset` columns are the KV history already in the cache, and only the last n follow the tree. Every node attends the history in full, because the tree branches inside the block being verified and everything before it is shared context. Within the block a node sees its ancestors and itself and no sibling, as before. The convention matches crate::utils::create_causal_mask exactly, which the third equivalence pins: a linear tree's history mask equals a causal mask built with the same offset, element for element. Together with the two the topology commit already had, that covers what the tree path has to reproduce before it can be switched on. The seam it feeds is already open. Gemma 4's forward_hidden_with_speculative_sinks takes `mask: Option<&MlxArray>` and threads it straight through to the layered forward; the MTP adapter passes None today. Wiring is passing this instead, not adding a parameter. Refs lablup#1204 * feat(speculative): make refusing a draft tree the default a target inherits MtpTarget gains verify_forward_tree, and its default refuses. That default is the safety property rather than a placeholder: flattening a tree into the sequence a linear verify expects lets a node's state reach its sibling, which the tree mask prevents for attention and nothing prevents for a recurrence. A target opts in by implementing the method, so omission means no rather than meaning silently wrong. Twenty types implement verify_forward, so a new method with a refusing default is also what keeps this from being a twenty-file signature change. Qwen 3.5 overrides it anyway, to state its cause where a reader looks for it. With full_attention_interval 4 over 64 layers it is 48 GatedDeltaNet and 16 attention, and a recurrence is a fold with no mask: a flattened tree folds state through siblings and answers wrongly with no signal. Verifying a tree there needs a scan that forks state at branch points, which does not exist. The refusal carries a reason rather than being a bare None, so the round loop can log why tree drafting fell back to the linear path instead of leaving an operator to work it out. Refs lablup#1204 * feat(speculative): verify a draft tree on Gemma 4 Gemma 4 can verify a tree because it has no recurrent state: every layer is attention, so the tree mask makes siblings invisible to each other and one flattened forward over the nodes is exact. That is the whole reason the redirect on lablup#1204 moved to this target. verify_forward's body becomes verify_forward_with_mask, an inherent method taking Option<&MlxArray>, and both entry points call it. The linear path passes None and the model builds its own causal mask exactly as before; the tree path passes a mask no causal mask can express. Everything after the forward is shared, so the two differ by that argument and nothing else, which is what keeps the linear path bit-identical. The mask is sized against the cache the next forward will read. Gemma4Wrapper gains sequence_kv_offset, returning the same first_present_cache_offset the model derives when it builds its own mask, so a supplied mask lines up with the cache instead of relying on the round loop to thread a number that could drift from it. The blocking value is -1e9 rather than -inf: an all-masked row through a softmax turns -inf into NaN, and a draft tree's mask has rows that mask most of the block. Failure sets are unchanged against baseline by set comparison, and the live pairing still runs at 95.41 tok/s. Nothing calls the tree path yet; the round loop still drafts linearly. Refs lablup#1204 * feat(speculative): let any drafter propose a tree, linear by default Drafter::draft_tree wraps draft_block into a linear tree unless a drafter overrides it. The degenerate default is not a convenience: it is what makes the tree wiring testable, because a linear tree has to produce byte-identical output to the chain path it subsumes, so any divergence once the round loop uses this is the wiring rather than the topology. Branching will need a signal the drafter already computes and throws away. Its per-step logits are reduced to an argmax today, and the runner-up at an uncertain step is exactly the alternative worth spending a verify position on. Refs lablup#1204 * feat(speculative): make refusing a tree rollback the default too Verifying a tree and keeping its accepted path are separate abilities, and the branch so far only had a default for the first. `verify_forward_tree` refuses unless a target opts in; `verify_finalize` had no counterpart, so a target that opted into the forward would have inherited the linear rollback, which trims a tail. Applied to a scattered accept set that keeps the wrong entries rather than failing. `verify_finalize_tree` takes the walk's path instead of an accept count and refuses by default. For a linear tree the path is `[0, 1, ..., accepted]`, the selection is a prefix, and the result has to be byte-identical to `verify_finalize(accepted, block_size, captured)`. That equivalence is the first thing to test once the round loop is wired, and a divergence there is wiring rather than topology. Accepting the forward and refusing the rollback is a real combination, not a contradiction, and Gemma 4 is currently in it. Its sliding-window layers hold `RotatingKVCache`, whose rollback rewinds `offset` and `idx` without moving data, which a scattered selection cannot use, and which has no `gather_positions`. So the two halves have to be probed as one capability before a tree is drafted: once a tree forward has appended its block, a refusal at the rollback leaves the cache in a state only a tree-aware rollback can undo. `captured` is consumed rather than handed back for that reason. Qwen 3.5 overrides it to name its own cause rather than inherit the message, matching what it already does for the forward. Unreachable while the forward refuses first, and there so a forward that stops refusing cannot silently inherit a rollback that cannot serve it. Also fixes an `i32 -> i32` cast in tree_tests that fails `clippy -D warnings`. Refs lablup#1204 * feat(cache): compact a scattered accept set within the verify block only `gather_positions` names the survivors of the whole live window, so using it for a tree rollback means re-listing the entire history every round and paying a full buffer rewrite for a change confined to the last `block_size` slots. It also replaces each buffer with a `take` result sized exactly to the survivors, so the next append has to grow the allocation again. On a long context that is the rollback cost lablup#1209 is about, and the tree path would have inherited all of it. `gather_within_tail(tail_len, kept)` takes offsets within the block and rewrites only `[tail_start, tail_start + kept_len)`, through the same `slice_update` the append path already uses. Cost scales with the verify block rather than with the context, and the allocation survives. A prefix selection is asserted equal to `trim`, and the same selection said in absolute terms is asserted equal to `gather_positions`, so the cheap path is the same rollback rather than a second one. It also unblocks Gemma 4, which is the only target that can verify a tree today and could not keep one. Its sliding-window layers hold `RotatingKVCache`, which had no gather at all, and whose `trim` cannot be generalized into one: it moves no data, it rewinds `offset` and `idx` and lets the next append overwrite. A tail that has to be compacted has to be written. The rotating form requires an unwrapped ring, which is the regime `enable_speculative_buffer` already keeps the verify block inside, and refuses by name once the ring has wrapped rather than compacting slots whose physical order is no longer logical order. Unlike `gather_positions` this places no requirement on `live_start`: the tail is located through `buffer_idx()`, which already accounts for a trimmed head. Both refuse `Turbo4Delegated` and the pool-backed cache for the reasons `gather_positions` refuses them. Refs lablup#1204, lablup#1209 * feat(speculative): roll a Gemma 4 verify block back to a tree's accepted path Gemma 4 could verify a draft tree since the tree mask landed, and could not keep one: its rollback trims a tail, which is the accept set only when the accept set is a prefix. Applied to a path it keeps the wrong entries, and keeping the wrong entries is silent. `verify_finalize_tree` selects instead of trimming, on the three things a round carries forward. The target cache goes through the new `gather_speculative_cache`, which hands each layer's cache the accepted nodes as block offsets; they are the same numbers, because the verify block is appended in the tree's topological order. The captured shared K/V goes through `gather_shared_kv`, which keeps the history in front of the block untouched and compacts the block region. The drafter's next hidden is read at the last accepted node rather than at an accept count. On a linear tree every one of those reduces to the linear path: the cache gather equals a trim, the shared-K/V gather equals the tail slice, and the hidden position is the same. That is what the round loop should check first once it is wired, and a divergence there is wiring rather than topology. The capability is per round, not per model, which is the part that constrains the wiring. The sliding-window layers hold a rotating cache whose gather needs an unwrapped ring, so `can_gather_speculative_cache` answers from state and the answer can change between rounds. The check inside `verify_finalize_tree` is a guard that turns a wrapped ring into a named refusal rather than a wrong window; it cannot be the contract, because by then the block is on the cache and a refusal has no recovery. The round loop has to ask before it drafts. No batched tree path, and `gather_speculative_cache` deliberately has no per-row tail-zero: a per-row path is not a per-row count, so rows cannot share one global trim the way the linear batched rollback lets them. Refs lablup#1204 * feat(speculative): let a target say whether a tree round can finish The two tree methods are one capability and the round loop cannot discover that by trying. `verify_forward_tree` appends the tree's block to the cache; if `verify_finalize_tree` then refuses, the round is stuck, because the block is on the cache and only a tree-aware rollback takes it off. There is no falling back to the linear path from there, so the question has to be asked before anything is drafted. `tree_round_is_available` is a method rather than a constant because the answer moves. Gemma 4 can verify a tree at any time and can only roll one back while its sliding-window rings are unwrapped, so it answers from cache state and can say yes this round and no the next. The default is false: a target opts in by answering, the way it opts into the methods themselves. Refs lablup#1204 * feat(speculative): run the round loop over a tree, behind a flag, default off The round loop only ever proposed a chain. Every piece the tree path needs existed and nothing called it. Four points in the round now branch and the rest is shared: the drafter call (`draft_tree` against `draft_block`), the verify (`verify_forward_tree` against `verify_forward`), the walk (`DraftTree::walk` against `speculative_walk`), and the rollback (`verify_finalize_tree` against `verify_finalize`). Emission, logprobs, EOS and diagnostics see one shape either way, because the tree walk is reduced to a `WalkResult` and the accepted node path is carried alongside it rather than threaded through everything downstream. The capability question is asked before anything is drafted, not after the verify. By then the forward has appended the block and only a tree-aware rollback can take it off, so there is no falling back to the linear path; a refusal at that point ends the session instead of continuing against a tail that no longer means what the drafter thinks. `record_round` now counts proposals as `actual_bs - 1` instead of `draft_tokens.len()`. Identical on the chain path by construction, and the only form that is right for a tree, whose proposals are not a flat vector. The stateful-drafter history hook is skipped on a tree round and says so. It trims against the chain the drafter proposed, and a tree's accepted set is a path through a branching proposal rather than a prefix of one chain. The pairing is unreachable today, since the one tree-capable target leaves `verify_hidden_full` as `None`, so this warns rather than guessing at a semantic nobody needs yet. `MLXCEL_MTP_TREE` switches it on and is off by default; `with_tree_drafting` overrides it for tests, which cannot mutate the environment safely from a parallel suite. The test that matters is the equivalence. With a linear tree, one script covering full accept, mid-block divergence and immediate rejection produces byte-identical tokens and an identical rollback log on both arms, and asserts the script actually exercised a partial accept and a zero accept so the comparison cannot pass vacuously. Refs lablup#1204 * fix(speculative): let the tree path actually engage on Gemma 4 The wiring landed and then declined every round, on every checkpoint, while the flag looked on. Two separate reasons, both of which a linear-tree equivalence check passes straight through, because a session that silently stays on the chain path trivially matches the chain path. That is why the announcement below exists. The rotating cache gated its gather on `is_trimmable`, which excludes buffered mode outright. `enable_speculative_buffer` is what puts the cache in buffered mode, and it is called precisely so an MTP verify block can append and roll back, so the gate was false in exactly the regime it was written for. The right question is whether the tail is still in logical order, and buffered mode answers yes by construction: it compacts from the front instead of wrapping, writes new tokens at `[idx, idx + n)`, and reports `idx` as the live length. `tail_is_in_logical_order` asks that, and two tests now cover the buffered regime, which the original ones did not touch. The multimodal adapters wrap the text adapter and delegate every trait method except the three tree ones, which they inherited as refusals. So tree drafting was dead for every multimodal Gemma 4 checkpoint, including the 12B the sweep runs on, while the text-only path would have worked. Nothing about a tree justifies that split; the vision tower is not in the verify forward. The generator now says once per session which path it took, and warns when the flag is set and the target declined. Before this, the two were indistinguishable from the outside, and "the flag is set" was being read as "the tree path ran". Verified on a real checkpoint rather than on the mock: gemma-4-12b-it-4bit with the 4-bit assistant drafter at block 5, 96 tokens, temperature 0. Tree on and tree off produce byte-identical output over 800 bytes, with identical round diagnostics: 39 rounds, 117 proposed, 56 accepted, 2.4359 emitted per verify, 17 zero-accept, 7 partial, 15 full. The tree arm logs that it engaged and the chain arm does not, so the comparison is not the vacuous one. Refs lablup#1204 * fix(gemma4): make a caller mask satisfy the sliding window too, past it A caller mask says which keys are structurally allowed. It cannot also say the sliding-window band, because the same mask is handed to the full-attention family, which must not have one. Below the window the two constraints agree, so nothing noticed. Past it they stop agreeing, and the sliding family was getting only the caller's half. The MTP draft-tree mask is what exposed it. Its history columns are all-attend by construction and `trim_mask_to_keys` crops from the right without adding a band, so past the window the sliding layers attended outside their window. Measured on gemma-4-12b-it-4bit with the 4-bit assistant drafter at block 5: at 96 generated tokens the tree and chain paths are byte-identical, and at 1400 they diverge from character 5133, which is a little past the 1024-token window. Acceptance was unaffected across the divergence (0.5209 against 0.5147), which is what ruled the rollback out and pointed at the mask; a broken rollback collapses acceptance rather than producing a coherent alternative continuation. The chain arm is deterministic run to run at that length, so the difference was systematic rather than the f16 jitter class. The sliding family now gets the intersection of both constraints, and only where they can disagree: the guard is `sliding_live_len + l > window`, so every context under the window takes exactly the path it took before. Additive masks intersect by `minimum`, and the intersection always keeps each query's own position, so no row can go all-masked and softmax to NaN. After the fix the same 1400-token comparison is byte-identical across the window, and the chain arm's own output is unchanged from before it, which is the check that the guard really did scope the change to the broken regime. The root suite's failure set is unchanged (the 13 pre-existing local reds). Refs lablup#1204 * test(gemma4): pin the three properties the caller-mask band rests on The band intersection was argued, not asserted, and the argument is what makes it safe to apply to every caller-mask path rather than only to the tree. Pulled it into `sliding_mask_from_caller` so the properties can be stated directly. Below the window the caller mask comes back untouched, which is what scopes the change: every context that fits inside the window takes exactly the path it took before. Past the window an all-attend caller mask, the draft-tree shape, reduces to the band. Past the window a mask that already carries the band is unchanged by intersecting it again, which is the property the batched and adopted-prefix prefill paths depend on, since they build window-aware masks themselves. Refs lablup#1204
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
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.
Summary
Tree drafting for MTP speculative decoding, wired end to end and switched off by default (
MLXCEL_MTP_TREE). With the default linear tree it reproduces the chain path exactly, which is the property the whole thing is built to be checkable against. No branching drafter yet; that is the next PR.Related issues
Refs #1204
Type of change
feat— new user-visible featurefix— bug fixWhat is here
Topology, mask and accept walk; a KV-cache rollback that can keep a scattered path; a Gemma 4 target that can verify and roll back a tree; and the round loop that drives it.
Two things turned out to matter more than the plan expected.
Gemma 4 could verify a tree and could not keep one, and it is the only target that can verify one at all. Its sliding-window layers hold
RotatingKVCache, which had no gather, and whosetrimcannot be generalized into one: it moves no data, it rewindsoffsetandidxand lets the next append overwrite. So the rollback primitive came first.gather_within_tail(tail_len, kept)compacts within the verify block only, through the sameslice_updatethe append path uses. That scoping is deliberate:gather_positionsnames the survivors of the whole live window, so a tree rollback through it would re-list the entire history every round, rewrite the whole buffer for a change confined to the lastblock_sizeslots, and drop the allocation to exactly the survivor count so the next append has to grow it again. A prefix selection is asserted equal totrim, and the same selection said in absolute terms is asserted equal togather_positions.The tree capability is per round, not per model. Gemma 4 can verify a tree at any time and can only roll one back while the tail is contiguous, so the answer moves with cache state. It has to be asked before anything is drafted: once the verify forward has appended the block, only a tree-aware rollback takes it off, so a refusal then has no fallback.
tree_round_is_availableis that question and the round loop asks it every round.The fix that is not about trees
sliding_mask_from_calleris a correctness fix for any caller mask past the sliding window, not just the tree's.A caller mask states which keys are structurally allowed. It cannot also state the window band, because the same mask goes to the full-attention family, which must not have one. Below the window the two agree; past it they stop agreeing and the sliding family was getting only the caller's half. The sliding family now gets the intersection, guarded on
sliding_live_len + l > windowso every context inside the window takes exactly the path it took before.Three unit tests pin the properties: identity below the window, an all-attend caller mask reducing to the band past it, and an already-banded mask unchanged by intersecting it again. That last one is what the batched and adopted-prefix prefill paths rely on.
How the tree path was found to be doing nothing
Worth recording, because it nearly shipped as a passing equivalence check.
After the wiring landed the flag looked on and the tree path declined every round, on every checkpoint. A linear-tree equivalence check passes straight through that: a session that quietly stays on the chain path trivially matches the chain path, so the check passes hardest when the feature is most broken. Two independent causes, neither of which raised anything:
is_trimmable, which excludes buffered mode outright.enable_speculative_bufferis what puts the cache in buffered mode, and it is called precisely so an MTP verify block can append and roll back, so the gate was false in exactly the regime it was written for.The generator now reports once per session which path it took, and warns when the flag is set and the target declined.
Test plan
cargo fmt --all -- --checkcargo clippy --workspace --all-targets --features metal,accelerate -- -D warningscargo test --workspace --profile test-fast --features metal,accelerate -- --test-threads=1, failure set unchanged frommainon this machine (13 pre-existing local reds in the root lib, 4 numeric-tolerance ones in mlxcel-core; the nightly runner is green on all of them)gemma-4-12b-it-4bit+gemma-4-12b-it-assistant-4bit, block 5, temperature 0Real-checkpoint results:
Unit level: a linear tree over a script covering full accept, mid-block divergence and immediate rejection produces identical tokens and an identical rollback log on both arms, and asserts the script exercised a partial accept and a zero accept so it cannot pass vacuously.
Cost
Alternating ABBA, 6 samples per arm, 400 tokens, after discarding a warm-up block whose throughput drifted from 49 to 73 tok/s across reps:
-2.33 tok/s, -3.2%, for identical output, with non-overlapping distributions. Most of it is a real buffer rewrite on partial-accept rounds across 48 layers where the chain path rewinds an index. It moves the bar for the next PR: branching has to clear 3.49 emitted per verify and this.
Not here
Branching, the batched (B > 1) tree path, and any tree support on Qwen 3.5, whose 48 GatedDeltaNet layers have no attention mask to make siblings invisible; its adapter refuses by name rather than inheriting the default.