fix(speculative): per-position verify attention for Qwen3.5 MTP parity by inureyes · Pull Request #78 · lablup/mlxcel

inureyes · 2026-05-24T14:14:24Z

Closes #77

The Qwen 3.5 MTP target-verify pass could produce per-position logits that disagree with single-token decode, flipping the speculative accept/reject decision and breaking the temperature-0 byte-equality invariant (greedy speculative output must equal drafter-less greedy decode). Mirrors the upstream mlx-vlm fix in PR #1188 (commit 45e6ece).

Root cause

Qwen35Model::forward_speculative verifies the whole draft block in one forward. Its full-attention layers ran the batched fused causal attention over the multi-token block, whose reduction order differs from single-query decode attention. The resulting per-position logit drift can flip the target's argmax at a verify position, diverging the accept/reject decision from the drafter-less greedy pass. Qwen 3.5 is especially sensitive because its gated attention and quantized projections amplify the drift through the gate.

Fix

Thread a target_verify flag through the Qwen 3.5 verify pass. When verifying a multi-token block, compute attention per query position: query i attends only to keys and values up to prefix_len + i + 1 with no mask, exactly as during single-token decode. The decode and prefill paths are untouched (target_verify defaults to false; all existing call sites pass false), so non-speculative inference is byte-identical and pays zero overhead.

Files

src/models/qwen3_next.rs: target-verify-aware attention wrappers + attend_per_position; existing wrappers delegate with target_verify = false.
src/models/qwen3_5.rs: forward_with_capture routes full-attention layers through verify-aware attention with target_verify = true.
tests/qwen35_mtp_verify_parity.rs: ignore-gated real-model regression test.

Verification

cargo fmt --check clean.
cargo clippy --features metal,accelerate --all-targets -- -D warnings clean.
Parity test run against qwen3.5-0.8b-4bit: all 8 verify positions match single-token decode argmax.

The Qwen 3.5 MTP target-verify pass could produce per-position logits that disagree with single-token decode, flipping the speculative accept/reject decision and breaking the temperature-0 byte-equality invariant (greedy speculative output must equal drafter-less greedy decode). Mirrors the upstream mlx-vlm fix in PR #1188 (commit 45e6ece). Root cause: `Qwen35Model::forward_speculative` verifies the whole draft block in one forward. Its full-attention layers ran the batched fused causal attention over the multi-token block, whose reduction order differs from single-query decode attention. The resulting per-position logit drift can flip the target's argmax at a verify position, diverging the accept/reject decision from the drafter-less greedy pass. Qwen 3.5 is especially sensitive because its gated attention (output times sigmoid of the gate) and quantized projections amplify the drift through the gate. Fix: thread a `target_verify` flag through the Qwen 3.5 verify pass and, when verifying a multi-token block, compute attention per query position. Query `i` attends only to keys and values up to `prefix_len + i + 1` with no mask, exactly as during single-token decode. The decode and prefill paths are untouched: `target_verify` defaults to false and all existing call sites pass false, so non-speculative inference is byte-identical and pays zero overhead. Changes: - `src/models/qwen3_next.rs`: add `forward_with_position_ids_verify` and `forward_hidden_with_position_ids_verify` (target-verify-aware wrappers) plus the `attend_per_position` helper. The existing `forward_with_position_ids` and `forward_hidden_with_position_ids` delegate with `target_verify = false` and stay byte-identical. - `src/models/qwen3_5.rs`: `Qwen35DecoderLayer::forward_with_capture` routes full-attention layers through the verify-aware attention with `target_verify = true`. - `tests/qwen35_mtp_verify_parity.rs`: an ignore-gated, real-model regression test asserting the verify pass produces per-position argmax tokens identical to single-token decode. Verified with cargo fmt, cargo clippy (metal,accelerate, all-targets, -D warnings), and the parity test run against qwen3.5-0.8b-4bit: all 8 verify positions match single-token decode argmax.

Bump version to 0.0.31. Highlights: MiniCPM-V 4.6 and RT-DETRv2 model support plus an Anthropic-style /v1/messages API endpoint (#82, #80, #74); fixes for missing or null chat-message content, Gemma 3n VLM per_layer_inputs, Qwen3.5 MTP verify attention, and batched quantized KV cache mask offsets (#91, #86, #78, #76); and a pinned 1.93.1 Rust toolchain (#87, #90).

build_msa_unified_mask_asymmetric generalises cycle 64's unified mask to num_query_blocks ≠ num_key_blocks. p_q starts at cache_offset; p_k in [0, num_key_blocks * block_size). The 'p_k > p_q ⇒ invalid' rule from cycle 64 still subsumes within-block causality, future-block leakage, divisibility padding — AND now cached-prefix causality. 7 unit tests at small asymmetric fixtures (b=1, num_kv_heads=2, n_rep=2, num_query_blocks=2, num_key_blocks=4, cache_offset=4, block_size=2, top_k=2). Including an exhaustive 64-position hand-computed reference check. The 7 symmetric tests from cycle 65 still pass — no regression. Cycle 65's HF reference comment at line 267 cited num_key_blocks vs k_len; the implementation read l (q_len) from x.shape[1] without honouring the distinction. The latent comment 'Unreachable under default prefill_chunk_size=512, hence latent until now' from cycle 65 was the flag that should have driven this work then. It does now. First piece of the 6-task proper-fix sequence (lablup#78-lablup#83). No call sites updated yet — the new function is dead code until lablup#79 wires it into sparse_sdpa.

sparse_sdpa signature changes: - Drop: `l: i32, num_blocks: i32` - Add: `q_len: i32, kv_len: i32, cache_offset: i32` num_query_blocks and num_key_blocks are derived inside from q_len and kv_len respectively. Q is padded to num_query_blocks * block_size; K/V are padded (separately) to num_key_blocks * block_size. The mask is the new asymmetric form from lablup#78 with cache_offset threaded through. Shape walk for the per-(kv_head, q_block) gather is updated: the broadcast axis is num_query_blocks (outer), the gather source is num_key_blocks (axis 3). The 5D GQA expansion uses num_query_blocks for the q-block axis. Output reshape goes back to padded_q_len, sliced to q_len. Call site at line 394: kv_len read from cache_k's actual shape; cache_offset = kv_len - l. The no-cache case (cache_offset=0) makes num_query_blocks == num_key_blocks and the algebra coincides with cycle 65's symmetric form. KNOWN GAP closed by lablup#80: the `selected` indices today are computed from the indexer over CURRENT-CHUNK idx_q/idx_k only, so values are chunk-relative ([0, num_query_blocks)) while sparse_sdpa interprets them as absolute ([0, num_key_blocks)). No crash (chunk-relative values are still valid into the larger range), but attention attends to wrong blocks in the cached case. End-to-end correctness gated on lablup#80's indexer K cache. build_msa_unified_mask (symmetric) marked #[allow(dead_code)] and kept with its 7 unit tests for regression value; deletion planned post-merge. cargo build clean, 14 mask tests pass (7 symmetric + 7 asymmetric).

The architectural completion: SparseAttention::forward() now correctly handles chunked-prefill / cached cases. Key changes: - Compute kv_len = cache_offset + l BEFORE update_and_fetch, so the asymmetric block counts are derived from the FULL cached length, not just the current chunk's length. - num_query_blocks vs num_key_blocks throughout: Q padded to num_query_blocks*block_size, K/V padded to num_key_blocks*block_size, reshape sites updated accordingly. - cache.m3_idx_k_update_and_fetch(&idx_k) (from lablup#80) swaps the current-chunk-only idx_k for the FULL cached idx_k. This closes the architectural gap behind the cycle-65 latent comment: top-k block selection now scores ALL num_key_blocks blocks, including the cached prefix that pre-lablup#81 was invisible to the indexer. - New apply_causal_block_mask_asymmetric: per-query-block max attendable absolute key block = floor((cache_offset + (i+1)*block_size - 1) / block_size). Over-permissive at block level; sparse_sdpa's unified mask trims within-block future positions. - New ensure_local_block_score_asymmetric: forces top-k to include the cached-prefix's last few absolute blocks (the query's local neighbourhood), which sit BEFORE the current chunk in absolute coords. - Dispatch dense when num_key_blocks <= top_k (top-k saturates → no sparsity benefit, dense is at least as good). - selected has shape [b, num_kv_heads, num_query_blocks, top_k] with ABSOLUTE key block indices in [0, num_key_blocks). sparse_sdpa's asymmetric path consumes them correctly. Symmetric apply_causal_block_mask + ensure_local_block_score marked #[allow(dead_code)] for regression value; the asymmetric forms subsume them at cache_offset=0. Deletion planned after lablup#78-lablup#83 lands and the asymmetric forms have lived without regressions. Cycle 65's "Unreachable under default prefill_chunk_size=512, hence latent until now" comment in cycle 65's code is now obsolete — the cached-prefill path is exercised on every forward with non-empty cache. 14 mask tests pass (7 symmetric regression + 7 asymmetric). Integration test for the actual four-chunk forward-attention path is task lablup#82; this commit wires the path, the test in lablup#82 verifies correctness against dense single-shot.

HANDOFF_cycle79_msa_chunked_prefill.md captures the state of the six-task proper-fix arc (lablup#78-lablup#83) so the smoke test can be completed by anyone with the branch checked out — pass criteria, debug paths if it fails, manual run command. Smoke script: replaced `yes ... | head -50` with a deterministic seq-driven loop. The pipe form triggered SIGPIPE on yes when head closed, pipefail picked it up, and `set -e` aborted the script after LONG_PROMPT was assigned but before any chat request was sent. Trace revealed the issue on the second attempt.

lablup#78) The Qwen 3.5 MTP target-verify pass could produce per-position logits that disagree with single-token decode, flipping the speculative accept/reject decision and breaking the temperature-0 byte-equality invariant (greedy speculative output must equal drafter-less greedy decode). Mirrors the upstream mlx-vlm fix in PR lablup#1188 (commit 45e6ece). Root cause: `Qwen35Model::forward_speculative` verifies the whole draft block in one forward. Its full-attention layers ran the batched fused causal attention over the multi-token block, whose reduction order differs from single-query decode attention. The resulting per-position logit drift can flip the target's argmax at a verify position, diverging the accept/reject decision from the drafter-less greedy pass. Qwen 3.5 is especially sensitive because its gated attention (output times sigmoid of the gate) and quantized projections amplify the drift through the gate. Fix: thread a `target_verify` flag through the Qwen 3.5 verify pass and, when verifying a multi-token block, compute attention per query position. Query `i` attends only to keys and values up to `prefix_len + i + 1` with no mask, exactly as during single-token decode. The decode and prefill paths are untouched: `target_verify` defaults to false and all existing call sites pass false, so non-speculative inference is byte-identical and pays zero overhead. Changes: - `src/models/qwen3_next.rs`: add `forward_with_position_ids_verify` and `forward_hidden_with_position_ids_verify` (target-verify-aware wrappers) plus the `attend_per_position` helper. The existing `forward_with_position_ids` and `forward_hidden_with_position_ids` delegate with `target_verify = false` and stay byte-identical. - `src/models/qwen3_5.rs`: `Qwen35DecoderLayer::forward_with_capture` routes full-attention layers through the verify-aware attention with `target_verify = true`. - `tests/qwen35_mtp_verify_parity.rs`: an ignore-gated, real-model regression test asserting the verify pass produces per-position argmax tokens identical to single-token decode. Verified with cargo fmt, cargo clippy (metal,accelerate, all-targets, -D warnings), and the parity test run against qwen3.5-0.8b-4bit: all 8 verify positions match single-token decode argmax.

Bump version to 0.0.31. Highlights: MiniCPM-V 4.6 and RT-DETRv2 model support plus an Anthropic-style /v1/messages API endpoint (lablup#82, lablup#80, lablup#74); fixes for missing or null chat-message content, Gemma 3n VLM per_layer_inputs, Qwen3.5 MTP verify attention, and batched quantized KV cache mask offsets (lablup#91, lablup#86, lablup#78, lablup#76); and a pinned 1.93.1 Rust toolchain (lablup#87, lablup#90).

inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:medium Medium priority area:inference Generation, sampling, decoding (incl. speculative, DRY) labels May 24, 2026

inureyes merged commit 4038da9 into main May 24, 2026
4 checks passed

inureyes deleted the fix/77-qwen35-mtp-verify-parity branch May 24, 2026 14:15

inureyes self-assigned this May 25, 2026

inureyes added this to the 0.1 milestone Jun 21, 2026

inureyes added the status:done Completed label Sep 6, 2026

inureyes mentioned this pull request Sep 20, 2026

fix(speculative): Qwen 3.5 DFlash verify diverges from classic decode #1935

Closed

5 tasks done

Navigation Menu

Repository navigation

fix(speculative): per-position verify attention for Qwen3.5 MTP parity - #78

fix(speculative): per-position verify attention for Qwen3.5 MTP parity#78
inureyes merged 1 commit into
mainfrom
fix/77-qwen35-mtp-verify-parity

inureyes commented May 24, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

inureyes commented May 24, 2026

Root cause

Fix

Files

Verification

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant