diff --git a/TECHNICAL_REPORTS/2865-dflash-per-row-accepts-20261106.en.md b/TECHNICAL_REPORTS/2865-dflash-per-row-accepts-20261106.en.md new file mode 100644 index 000000000..f16a31750 --- /dev/null +++ b/TECHNICAL_REPORTS/2865-dflash-per-row-accepts-20261106.en.md @@ -0,0 +1,478 @@ +# Technical Report: PR #2865 - perf(speculative): commit per-row accepts in batched DFlash + +**Date**: 2026-11-06 +**Status**: In review (open PR; closes #2442) +**Languages**: Rust, Markdown +**Risk Level**: Medium + +--- + +## Executive Summary + +PR #2865 implements item 1 of issue #2442: the batched Qwen 3.5 DFlash round loop now commits each row's own speculative-walk accept count instead of clamping the whole round to the shortest live row. It reuses the Gemma 4 `DivergentVerifyRows` attention layout, compacts each row's attention KV back to a contiguous valid prefix, masks the stale gap, rotates queries and keys at each row's logical positions, and replays each GDN layer over that row's accepted prefix. On an M3 Ultra with `qwen3.5-4b-4bit` and `qwen3.5-4b-dflash` at B = 4, aggregate decode rose from 83.06 to 116.05 tok/s (1.397x), committed accepts per row-round from 0.323 to 1.036, and rounds from 96 to 67, while the calibrated parity gate passed 2 of 2 with per-row accepts as the default. The drafter still receives a shared rectangular prefix, so the speculative walk itself falls from 2.938 to 1.036; per-row drafter context is explicitly out of scope, and no server-level concurrent-request smoke run was performed. + +--- + +## 1. Problem Statement + +### 1.1 Background + +The batched DFlash round loop for Qwen 3.5 shares one rectangular target cache across every row. Issue #2353 handled divergent walks by clamping a round to the shortest live row's accept (`clamp_uniform_acceptance`), so every row committed the same number of verify positions. That was correct for a rectangular cache because a rejected tail cannot simply be zeroed: a zero key still scores `q . 0 = 0` and takes `exp(0)` weight in the softmax, which dilutes attention on every later step. + +The cost was visible in the accept statistics. On the #2442 measurement the walk averaged 2.938 accepted tokens per row-round while the clamp committed only 0.323, and 93 of 96 rounds were clamped. Batching therefore bought about 1.06x over the B = 1 round loop (83.06 against 78.09 tok/s), far below the walk's potential. + +Issue #2442 item 1 asked for per-row accepts. PR #2853 then established that the batched verify forward is not byte-identical to a B = 1 forward under Metal kernel dispatch, and replaced the hard byte gate with a calibrated 1.0-logit near-tie margin. That settled the prerequisite: the accept policy was not the cause of the parity divergence, so per-row accepts could proceed under the existing margin gate. + +### 1.2 Existing Issues + +- **The clamp discards accepted work.** Every round committed the minimum accept across live rows, so any row that walked further had its extra accepted tokens rolled back. +- **Zero-filling a rejected tail is not a valid alternative.** Stale keys still participate in the softmax with `exp(0)` weight. +- **A rectangular cache has one offset.** Rows that committed different counts need a way to keep different contiguous prefixes while sharing the same physical cache allocation. +- **Linear-attention state is per-row but its valid length is not.** A single batched GDN replay over `max(accepted)` would advance a short row through positions it rejected. +- **Batching barely paid.** At 0.323 committed accepts per row-round, B = 4 delivered only about 1.06x over B = 1. + +### 1.3 Risk Assessment + +| Risk | Impact | Likelihood | +|---|---|---| +| A ragged cache layout corrupts later attention or GDN state | High (wrong tokens) | Low: the parity gate runs with per-row accepts as the default, and the compaction is covered by unit tests | +| An unsupported cache mode reaches the per-row path | Medium (panic or corruption) | Low: `row_layout_supported` gates on dense Fp16, non-paged, non-head-trimmed caches and the compaction asserts the same preconditions | +| Per-row accepts interact badly with real server concurrency | Medium (serving correctness) | Medium: no server-level concurrent-request smoke run was performed | +| The throughput comparison is shifted by host load | Low (decision quality) | Low: the arms alternated under the same competing build load and the clamp and B = 1 cells reproduce the earlier quiet-box record | +| The lower walk is mistaken for a regression | Low (misattribution) | Certain if unrecorded: the drafter context is deliberately still rectangular, and per-row drafter context is a separate change | + +--- + +## 2. Technical Review + +### 2.1 Security + +No network, authentication, input-parsing, or serialization surface changes. The only new configuration input is the `MLXCEL_DFLASH_PER_ROW_ACCEPTS` environment variable, read when the generator is constructed and accepting only an on/off value. No user data is logged; the accept statistics logged by `dflash_target.rs` already existed and only their description was updated. The rollback path validates cache shape, mode, paging, and head trim with assertions rather than proceeding on an unexpected layout. + +### 2.2 Performance + +Measured on an M3 Ultra, `qwen3.5-4b-4bit` target with `qwen3.5-4b-dflash`, block 16, 128 tokens per row, temperature 0, three alternating runs per arm. + +| Metric | Clamp | Per-row | Change | +|---|---:|---:|---:| +| Aggregate throughput (tok/s) | 83.06 | 116.05 | 1.397x | +| B = 1 reference (tok/s) | 78.09 | 78.09 | 1.00x | +| Rounds | 96 | 67 | 30% fewer | +| Clamped rounds | 93 | 0 | eliminated | +| Committed accepts per row-round | 0.323 | 1.036 | 3.2x | +| Walk per row-round | 2.938 | 1.036 | 0.35x | + +The per-row arm is 1.49x the B = 1 round loop; the clamp arm is 1.06x. Throughput is the round-loop wall clock from `DFlashBatchedRunOutput::stats` with prefill excluded, so the per-row compaction work (`copy` plus `slice_update` per row per attention layer) is inside the measured cost. The three clamp runs spanned 0.2% and the three per-row runs 0.36%. + +### 2.3 Compatibility & Dependencies + +- **Breaking Changes**: none. The new `SpeculativeTarget` methods have default implementations that preserve the old behavior. +- **New Dependencies**: none. +- **Compatibility**: per-row accepts are the default only when the target reports support and the caches are dense Fp16 without paging or head trim. Quantized, paged, or head-trimmed caches select the uniform clamp automatically. `MLXCEL_DFLASH_PER_ROW_ACCEPTS=0` (also `false`, `no`, `off`) restores the clamp for A/B attribution. +- **Internal signature change**: `Qwen3NextAttention::forward_with_position_ids_verify` and its hidden-state sibling gained an `Option<&[i32]>` row-layout argument. Both are `pub(crate)` and all call sites were updated. +- **VLM**: `Qwen35VLModel` delegates the capability, offset, verify, and rollback hooks to its text model. The attention layer refuses per-row verify when MRoPE position ids are supplied; the DFlash verify path passes none, so this is a guard rather than a served path. + +### 2.4 Code Quality + +Three unit tests were added: two for attention compaction in `models::qwen3_5_tests` (row block moved to its own valid end; no-op when every row ends at the offset) and one for `build_divergent_row_mask` in `mlxcel-core::utils` (history attends, stale gap blocked, new block causal at each row's logical positions). The `models::qwen3_5_tests` suite now has 30 tests; the DFlash round-loop modules keep 15 (`round_loop.rs`) and 13 (`round_loop_batched.rs`) tests. The mask builder was moved out of `gemma4.rs` into `mlxcel-core::utils`, removing about 30 duplicated lines from Gemma 4 and giving both models one implementation. `cargo fmt --all -- --check` and `cargo clippy --lib --tests --features metal,accelerate -- -D warnings` are clean. + +--- + +## 3. Technical Decisions + +### 3.1 Gate per-row accepts on a target capability and an instance cache check + +**Context:** +Per-row accepts require a target whose attention cache can be compacted row by row and whose linear-attention state can be replayed per row. Not every `SpeculativeTarget` can do that, and even Qwen 3.5 cannot for every cache mode. + +**Alternatives Considered:** + +| Option | Pros | Cons | +|---|---|---| +| Always take the per-row path for any target | Simple loop | Breaks targets that only support a uniform rollback; panics on quantized or paged caches | +| Require every target to implement the new hooks | Uniform behavior | Forces unrelated targets and test doubles to add no-op implementations | +| **Chosen: default trait methods plus a two-level gate** | Existing targets are unchanged; unsupported caches fall back safely | The served path depends on runtime cache state | + +**Rationale:** +`supports_per_row_accepts()` is the family-level answer and defaults to `false`. `row_layout_supported(caches)` is the instance-level answer and checks `KVCacheMode::Fp16`, `!is_paged_backed()`, and `offset - live_len() == 0` (no head trim). The loop takes the per-row path only when both are true and the caller has not disabled it. The default `rollback_partial_batched_rows` and `verify_forward_with_row_layout` delegate to the uniform hooks, which is correct only when every row shares the offset, so a target that forgets to override them fails safe rather than silently producing a ragged layout. + +**Trade-offs:** +The production behavior now depends on the cache mode chosen at runtime, so a quantized deployment silently uses the old clamp and does not get the 1.397x. The alternative was to panic or refuse the burst, which would turn a performance feature into a availability regression. + +### 3.2 Reuse the Gemma 4 divergent-row layout instead of inventing a second one + +**Context:** +The Gemma 4 MTP verifier already handles rows with different logical valid ends against a shared physical cache. Its `build_divergent_verify_mask` encodes exactly the three regions a ragged DFlash round needs: the row's valid history, its stale gap, and the new block. + +**Alternatives Considered:** + +| Option | Pros | Cons | +|---|---|---| +| Pad every row to `max(accepted)` and zero rejected positions | No compaction | Zero keys still take softmax weight; already rejected in #2353 | +| Copy the cache per row and run independent forwards | Conceptually simple | Destroys the batching win and duplicates cache memory per row | +| **Chosen: shared divergent-row mask plus per-row rotation** | One mask describes all three regions; reuses a proven layout | Mask construction is `O(B * l * (offset + l))` | +| Build a DFlash-specific mask | No cross-model coupling | A second implementation to keep correct | + +**Rationale:** +`build_divergent_row_mask(l, offset, per_row_valid_end, left_padding, window)` returns `[B, 1, l, offset + l]` additive mask. Row `r` admits keys below its valid end, blocks `[valid_end, offset)`, and treats new key `t` at logical position `valid_end + t`, so query `j` sees new keys `t <= j`. `Qwen3NextAttention` detects `divergent_rows` when `target_verify`, `l > 1`, and any end differs from the shared offset, applies RoPE per row at the row's own start, and runs `attend_per_position_rows`, which slices query `j` against key width `offset + j + 1` with the matching mask slice. The Gemma 4 function now delegates to the same core utility, so the two models cannot drift. + +**Trade-offs:** +The mask is built per verify forward and explicitly materialized. That cost is included in the measured throughput and is outweighed by the accept gain at the measured shape. A future fused path could avoid the materialization but is not needed for the current win. + +### 3.3 Keep the drafter context rectangular at the shortest live committed prefix + +**Context:** +The target can now keep ragged caches, but the drafter's input hidden tensor and its own context are still built from one rectangular slice. Making the drafter per-row would require per-row drafter cache views and per-row drafter forwards, which is a much larger change. + +**Alternatives Considered:** + +| Option | Pros | Cons | +|---|---|---| +| Give the drafter per-row context now | Higher walk | Requires per-row drafter cache and forward plumbing; not measured here | +| Feed the drafter the longest committed prefix | More draft material for lagging rows | The hidden slice from the verify block is only valid up to the shortest accepted block position, so this would mix rejected positions into context | +| **Chosen: shortest live committed prefix** | Correct given the captured hidden slice; reuses the existing rectangular contract | The walk falls from 2.938 to 1.036 | + +**Rationale:** +`drafter_ctx` is computed as the minimum `accepted[r] + 1` over unfinished rows, and the hidden tensor is sliced to that length. Rows that committed more than the minimum still advance their target caches by their own accept, so the emitted stream stays exact; the drafter simply re-drafts from a lagging context on the next round. This keeps the change bounded to the target side and makes the throughput comparison attributable to the accept policy alone. + +**Trade-offs:** +The walk drops by 65%, so the round count benefit is smaller than the accept count benefit suggests. Per-row accepts are still a clear win because the committed count, not the walk, determines emitted tokens per round. Per-row drafter context is recorded as a follow-up, not as a defect in this PR. + +### 3.4 Compact the attention KV and replay GDN state per row + +**Context:** +After a ragged round the shared cache holds each row's accepted block at the shared physical offset, followed by rejected positions. The next round needs every row to look like a contiguous `[0, valid_end)` prefix while the physical offset stays shared. + +**Rationale:** +`rollback_speculative_cache_rows` computes `o_post = max(valid_end_pre[r] + accepted[r] + 1)`. `compact_attention_rows` moves each row's `accepted[r] + 1` block from `o_pre` to that row's logical end, zeroes the vacated gap up to `o_post`, and the caller trims the cache offset to `o_post`. The source is copied before `slice_update` so an overlapping move cannot read already-overwritten rows. GDN state is replayed per row: `replay_gdn_rows` slices each row's captured `q/k/v/a/b` and initial state to `accepted[r] + 1`, runs the chain-parity kernel at that row's own length, and stacks the per-row conv and recurrent states back into the `[B, ...]` cache tensors. This keeps the per-row numerics on the same kernel the uniform path uses. + +**Trade-offs:** +Compaction adds one copy and one `slice_update` per row per attention layer per round. It is measured inside the reported throughput. The uniform path is also kept intact for fallback and A/B use, so the code carries both arms; that duplication is deliberate for attribution. + +--- + +## 4. Implementation Details + +### 4.1 Architecture Changes + +``` +[Before, uniform clamp] +verify [B, bs] -> walk per row -> clamp_uniform_acceptance (min over live rows) + -> rollback_partial_batched (trim bs - (a + 1) in every row) + -> drafter hidden [B, max_a + 1] + +[After, per-row accepts] +row_valid_ends = [prefill offset; B] +verify [B, bs] + -> per-row RoPE at row_valid_ends[r] + -> attend with build_divergent_row_mask (history, stale gap, causal block) + -> commit accepted[r] + 1 per row + -> compact_attention_rows: block moves to row_valid_ends[r], gap zeroed + -> cache offset trimmed to max(new valid ends) + -> replay_gdn_rows: per-row conv + recurrent state over accepted prefix + -> row_valid_ends[r] += accepted[r] + 1 + -> drafter hidden [B, min_live(accepted) + 1] (rectangular, deliberately) +``` + +### 4.2 Key Code Changes + +**File: `src/lib/mlxcel-core/src/drafter/dflash/round_loop_batched.rs`** + +The generator resolves the policy once per run and seeds the logical ends from the prefill offset: + +```rust +let per_row = self.per_row_accepts + && target.supports_per_row_accepts() + && target.row_layout_supported(caches); +let mut row_valid_ends: Vec = if per_row { + let offset = T::batch_cache_offset(caches); + vec![offset; batch_size] +} else { + Vec::new() +}; +``` + +The verify call branches on the policy, the clamp is skipped on the per-row arm, and the rollback hands each row's own accepted prefix to the target: + +```rust +if per_row { + target.rollback_partial_batched_rows( + caches, &verify_out, &accepted_i32, &row_valid_ends, bs as i32, + ); + for r in 0..batch_size { + row_valid_ends[r] += accepted_i32[r] + 1; + } +} else { + // uniform clamp arm: trim bs - (accepted + 1) in every row +} +``` + +The drafter context stays rectangular at the shortest live row: + +```rust +let drafter_ctx = if per_row { + (0..batch_size) + .filter(|&r| !finished[r]) + .map(|r| accepted_per_row[r] + 1) + .min() + .unwrap_or(1) +} else { + max_accepted + 1 +}; +``` + +**Reason for change:** the round loop owns the policy switch and the logical position bookkeeping; the target owns cache layout. Keeping the switch in one place makes the A/B measurable in a single process. + +**File: `src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs`** + +`SpeculativeTarget` gains four defaulted hooks: `supports_per_row_accepts`, `row_layout_supported`, `batch_cache_offset`, and the static-free `verify_forward_with_row_layout`, plus `rollback_partial_batched_rows`. The existing `rollback_partial_batched` contract is narrowed to the uniform-clamp arm in its documentation. + +**Reason for change:** targets that cannot represent ragged rows keep compiling and keep their old semantics. + +**File: `src/models/qwen3_5.rs`** + +`Qwen35Model` implements the capability gate, the shared offset, the ragged verify forward, and the ragged rollback. `forward_speculative_inner` threads `Option<&[i32]>` to each decoder layer. The non-uniform rollback is: + +```rust +if !uniform { + let o_post = ve_pre.iter().zip(accepted) + .map(|(&ve, &a)| ve + a + 1).max().expect("non-empty accepted"); + for cache in caches.iter_mut() { + if let Qwen3NextCache::Attention(kv) = cache { + compact_attention_rows(kv, &ve_pre, accepted, o_pre, o_post); + let roll_back = o_pre + block_size - o_post; + if roll_back > 0 { kv.trim(roll_back); } + } + } + self.replay_gdn_rows(caches, gdn_states, accepted, o_post); + return accepted.iter().copied().max().unwrap_or(0); +} +``` + +**Reason for change:** the attention cache must become a valid contiguous prefix per row before the next round, and the GDN state must reflect only that row's accepted tokens. + +**File: `src/models/qwen3_next.rs`** + +`Qwen3NextAttention` gains `apply_rope_per_row` and `attend_per_position_rows`. The divergent path is selected when any row's valid end differs from the cache offset: + +```rust +let divergent_rows: Option<&[i32]> = row_valid_ends.filter(|ends| { + target_verify + && l > 1 + && ends.len() == b as usize + && ends.iter().any(|&end| end != offset) +}); +``` + +Each query position slices the exact key width a B = 1 chain would see, with the divergent-row mask encoding history, gap, and block causality. + +**Reason for change:** a shared rectangular attention call cannot give each row a different logical start; the per-position path can, and it matches the Gemma 4 approach. + +**File: `src/lib/mlxcel-core/src/utils.rs` and `src/models/gemma4.rs`** + +`build_divergent_row_mask` moves into `mlxcel-core::utils`; Gemma 4's `build_divergent_verify_mask` becomes a one-line delegation. + +**Reason for change:** one tested mask implementation for both the Gemma 4 MTP verifier and the Qwen 3.5 DFlash verifier. + +**File: `src/vision/qwen3_5_vl.rs`** + +`Qwen35VLModel` forwards the four new hooks to its text model. + +**Reason for change:** the VLM wrapper should not silently lose the capability or the fallback. + +**Files: `docs/CONTINUOUS_BATCHING.md`, `docs/environment-variables.md`, `docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md`** + +The batching guide and the environment variable table record the default, the fallback, and the measured numbers; the benchmark document records the method, the per-run values, the caveat, and the correctness gate. + +### 4.3 Data Model Changes + +None. The cache layout stays `[B, heads, capacity, dim]`; only the valid-prefix bookkeeping and the physical offset change per round. + +--- + +## 5. Learning Points + +### 5.1 Ragged acceptance inside a shared rectangular cache + +**Concept:** +A shared batched KV cache has one physical offset and one allocation. Ragged accept counts do not require per-row allocations if each row's live content is compacted to a contiguous prefix after every round. The invariant becomes "each row's cache content is a contiguous `[0, valid_end)` prefix, and the stale region up to the shared offset is masked out". + +**Application in this PR:** +`compact_attention_rows` moves the accepted block from the shared offset to the row's logical end and zeroes the vacated gap. `build_divergent_row_mask` then encodes the three regions for the next verify forward. The physical offset is trimmed to the maximum new valid end, so the allocation does not grow with the raggedness. + +**Common Use Cases:** +- Batched speculative decoding with per-sequence accept counts. +- Continuous batching where sequences stop, resume, or roll back at different lengths. +- Any shared-buffer scheme that needs per-row logical lengths without per-row buffers. + +### 5.2 One mask layout, two models + +**Concept:** +A mask that expresses "valid history, blocked gap, causal new block at logical positions" is model-agnostic. Lifting Gemma 4's implementation into a core utility removed a second copy before it could drift. + +**Application in this PR:** +`build_divergent_row_mask` is now called by Gemma 4's `build_divergent_verify_mask` and by Qwen 3.5's `attend_per_position_rows`. The move also brought the mask under a dedicated unit test. + +**Common Use Cases:** +- Reusing attention masks across model families with the same cache topology. +- Consolidating numerically sensitive device code behind one tested constructor. + +### 5.3 Judge speculative batching by committed accepts, not by walk + +**Concept:** +The scan ratio (walk) measures what the drafter proposed; the committed accept count measures what the target actually kept. A batching change can lower the walk while raising throughput if it stops discarding accepted tokens. + +**Application in this PR:** +The per-row arm's walk fell from 2.938 to 1.036 because the drafter context is still the shortest live prefix, yet aggregate throughput rose 1.397x because committed accepts rose from 0.323 to 1.036 and rounds fell from 96 to 67. Reporting only the walk would have read as a regression. + +**Common Use Cases:** +- Speculative decoding A/B tests. +- Any pipeline where a producer's proposal rate and a consumer's acceptance rate can move in opposite directions. + +--- + +## 6. Further Learning + +### Key Terms + +| Keyword | Description | Relevance | +|---|---|---| +| DFlash | Block-draft speculative decoding with a diffusion-style drafter | The batched loop changed here | +| Uniform clamp | Committing the shortest live row's accept to every row | The previous policy, still reachable | +| Ragged accept | Different accept counts per row in one round | What per-row accepts commit | +| Valid end | Length of a row's contiguous valid cache prefix | The logical position used for RoPE and attention | +| Stale gap | `[valid_end, physical_offset)` left by a shorter accept | Masked out of attention and zeroed by compaction | +| GDN | Gated delta net linear-attention layer | Its conv and recurrent state are replayed per row | +| Walk | Tokens the drafter proposed and the target verified | Falls here because drafter context stayed rectangular | + +### Related Technologies/Frameworks + +- **MLX**: array framework and Metal kernels; the `copy` plus `slice_update` compaction and the per-position attention slices run on it. +- **Gemma 4 divergent MTP verify**: the source of the `DivergentVerifyRows` mask layout that this PR reuses. + +### Related PRs/Issues + +- Issue #2442: the parent issue; this PR implements item 1 and closes it. +- Issue #2353: introduced the uniform clamp and the original B = 4 versus B = 1 measurement. +- PR #2853 / issue #2851: calibrated the batched verify parity margin under which this PR's gate runs. +- Issue #2190: the Gemma 4 per-row verify architecture that supplied the mask layout. +- Issue #2440: the earlier per-row scoping work referenced by the benchmark document. + +--- + +## 7. Change Summary + +### Statistics + +| Item | Value | +|---|---| +| Files changed | 15 | +| Lines added | +1279 | +| Lines deleted | -156 | +| Tests added | 3 unit tests, plus 1 new ignored real-model throughput harness | +| Runtime behavior changed | Yes: per-row accepts are the default when the caches support compaction | + +### Changes by Category + +| Category | Count | Summary | +|---|---|---| +| Performance / runtime | 7 | Round loop policy switch, target hooks, Qwen 3.5 verify and rollback, Qwen3Next per-row RoPE and attention, VL delegation, server log text | +| Shared utility | 1 | `build_divergent_row_mask` moved to `mlxcel-core::utils`; Gemma 4 delegates | +| Tests | 3 | Compaction unit tests, mask unit test, throughput comparison harness | +| Documentation | 4 | Batching guide, environment variables, benchmark record, parity-harness wiring | + +### Related Commits + +| Hash | Type | Message | +|---|---|---| +| `ded78dce` | perf | perf(speculative): commit per-row accepts in batched DFlash | + +--- + +## 8. Follow-up Actions + +### Required + +- [ ] Run a server-level concurrent-request smoke test through the real collected-burst path at B > 1 with per-row accepts enabled, covering request completion, EOS, max-token rows, and a mixed prompt-length window. The current evidence is a harness burst on four equal-length prompts. +- [ ] Implement per-row drafter context as a separate change and re-measure. The walk is expected to recover above 1.036, but it is not part of this PR's budget. + +### Unvalidated Surface + +- **No server-level concurrent-request smoke run.** The measured and parity runs use the test harness (`run_burst_with_policy`), not the server's collected-burst scheduler. Row lifecycle, admission, cancellation, and logging in `src/server/batch/dflash_target.rs` are exercised only by the existing suite, not by this measurement. +- **Per-row drafter context is out of scope.** The drafter still sees the shared rectangular minimum prefix. This is the stated reason the walk falls to 1.036, and it means the measured 1.397x is a lower bound on what per-row target accepts can deliver once the drafter is addressed. +- **Measurement caveat.** Another worktree's Rust build was compiling on a few of the 80 cores during the throughput run. Both arms alternated under the same load, and the clamp and B = 1 cells reproduce the earlier quiet-box #2442 record (83.0 against 82.96 to 83.00; 78.09 against 77.70 to 77.83), so the comparison is not load-shifted, but the absolute numbers were not taken on a fully idle box. +- **Single host and checkpoint.** M3 Ultra, macOS 26.6.2, `qwen3.5-4b-4bit` plus `qwen3.5-4b-dflash`, block 16, 128 tokens per row, temperature 0, four truncated prompts. Other hardware, quantizations, block sizes, and prompt distributions are not measured. +- **Fallback paths are not re-validated here.** Quantized, paged, and head-trimmed caches take the unchanged uniform clamp; their correctness is inherited, not re-tested by the per-row gate. +- **VLM path.** `Qwen35VLModel` delegates the hooks, and the attention layer refuses divergent rows when MRoPE position ids are present. No VLM DFlash parity run is included in this PR. + +### Monitoring Required + +- `clamped_rounds` should stay 0 on the per-row arm and non-zero only when the clamp is explicitly selected or a cache falls back. A non-zero value on a dense Fp16 run means the capability gate regressed. +- Watch aggregate decode throughput and committed accepts per row-round after deployment, especially under mixed prompt lengths where raggedness is highest. +- Run `greedy_parity_dflash_qwen35_batched_b4_matches_b1` and `identical_prompts_dflash_qwen35_b4_matches_b1` on the CI hardware lane whenever MLX, quantization, or kernel dispatch changes. + +### Future Improvements + +- Extend per-row accepts to quantized and paged caches, which currently fall back to the clamp. +- Consider a fused mask-free attend path if mask construction becomes a measurable fraction of the round cost. +- Evaluate per-row drafter context together with per-row target accepts; the two interact through the walk and should be measured as a pair. + +--- + +## Appendix + +### A. Test Results + +| Check | Result | +|---|---| +| `cargo test --profile test-fast --features metal,accelerate --test speculative_parity qwen35_dflash_batched -- --ignored --test-threads=1` | 2 passed with per-row accepts as the default | +| `throughput_dflash_qwen35_batched_clamp_vs_per_row` | Completed; three alternating runs per arm | +| `cargo test --profile test-fast --lib models::qwen3_5_tests` | Passed; 30 tests, including 2 new compaction tests | +| mlxcel-core round-loop unit tests | Passed; 15 in `round_loop.rs`, 13 in `round_loop_batched.rs` | +| `build_divergent_row_mask` unit test | Passed | +| `cargo clippy --lib --tests --features metal,accelerate -- -D warnings` | Pass | +| `cargo fmt --all -- --check` | Pass | + +### B. Performance Benchmarks + +| Run | Clamp tok/s | Per-row tok/s | +|---:|---:|---:| +| 0 | 83.04 | 115.82 | +| 1 | 83.22 | 116.24 | +| 2 | 83.06 | 116.05 | +| Median | 83.06 | 116.05 | + +Accepts and rounds: + +| Arm | Rounds | Clamped rounds | Walk per row-round | Committed per row-round | +|---|---:|---:|---:|---:| +| Clamp | 96 | 93 | 2.938 | 0.323 | +| Per-row | 67 | 0 | 1.036 | 1.036 | + +### C. Measurement Environment + +| Field | Value | +|---|---| +| Host | Apple M3 Ultra, 80-core, 512 GB unified memory, macOS 26.6.2 | +| Backend | Metal + Accelerate | +| Target | `models/qwen3.5-4b-4bit` | +| Drafter | `models/qwen3.5-4b-dflash` | +| Block size | 16 | +| Generated tokens | 128 per row | +| Sampling | Temperature 0, greedy | +| Prompts | The four prompts of the batched DFlash parity harness, truncated to the shortest tokenization | +| Method | Two discarded warm-ups, three measured runs per arm alternating clamp and per-row, under `caffeinate -i` and `scripts/with_indexers_paused.sh` | + +### D. References + +- `src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs` +- `src/lib/mlxcel-core/src/drafter/dflash/round_loop_batched.rs` +- `src/lib/mlxcel-core/src/utils.rs` +- `src/models/qwen3_5.rs` +- `src/models/qwen3_next.rs` +- `src/models/gemma4.rs` +- `src/vision/qwen3_5_vl.rs` +- `tests/speculative_parity/qwen35_dflash_harness.rs` +- `tests/speculative_parity/qwen35_dflash_throughput.rs` +- `docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md` +- Issue #2442, issue #2353, PR #2853, issue #2851, issue #2190, issue #2440 diff --git a/TECHNICAL_REPORTS/2865-dflash-per-row-accepts-20261106.ko.md b/TECHNICAL_REPORTS/2865-dflash-per-row-accepts-20261106.ko.md new file mode 100644 index 000000000..6820385d1 --- /dev/null +++ b/TECHNICAL_REPORTS/2865-dflash-per-row-accepts-20261106.ko.md @@ -0,0 +1,478 @@ +# 기술 보고서: PR #2865 - perf(speculative): commit per-row accepts in batched DFlash + +**작성일**: 2026-11-06 +**상태**: 검토 중 (열린 PR, #2442 종료) +**언어**: Rust, Markdown +**위험도**: 중간 + +--- + +## 요약 + +PR #2865는 이슈 #2442의 항목 1을 구현한다. 배치 Qwen 3.5 DFlash 라운드 루프가 라운드 전체를 가장 짧은 live row의 accept로 clamp하지 않고, 각 row의 speculative walk accept 수를 그대로 commit한다. Gemma 4의 `DivergentVerifyRows` 어텐션 레이아웃을 재사용해 각 row의 attention KV를 연속적인 valid prefix로 압축하고, stale gap을 마스킹하고, 각 row의 논리 위치에서 query와 key를 회전시키고, 각 GDN 레이어를 해당 row의 accepted prefix로 replay한다. M3 Ultra에서 `qwen3.5-4b-4bit`와 `qwen3.5-4b-dflash`, B = 4 기준으로 aggregate decode가 83.06에서 116.05 tok/s로 상승했고(1.397x), row-round당 committed accept가 0.323에서 1.036으로, 라운드 수가 96에서 67로 줄었으며, 보정된 parity gate는 per-row accepts를 기본값으로 2개 중 2개 통과했다. drafter는 여전히 공유된 직사각형 prefix를 받으므로 speculative walk 자체는 2.938에서 1.036으로 감소한다. per-row drafter context는 명시적으로 범위 밖이며, 서버 수준 동시 요청 smoke 실행은 수행되지 않았다. + +--- + +## 1. 문제 정의 + +### 1.1 배경 + +Qwen 3.5의 배치 DFlash 라운드 루프는 모든 row가 하나의 직사각형 target cache를 공유한다. 이슈 #2353은 walk가 갈라질 때 라운드를 가장 짧은 live row의 accept로 clamp했고(`clamp_uniform_acceptance`), 따라서 모든 row가 같은 수의 verify 위치를 commit했다. 직사각형 cache에서는 이것이 올바른 처리였다. 거부된 tail을 0으로 채우는 것은 대안이 될 수 없기 때문이다. 0인 key도 `q . 0 = 0`으로 점수화되어 softmax에서 `exp(0)` 가중치를 차지하고, 이후 모든 스텝에서 해당 row의 어텐션을 희석한다. + +비용은 accept 통계에 드러났다. #2442 측정에서 walk는 row-round당 평균 2.938이었지만 clamp가 commit한 값은 0.323이었고, 96개 라운드 중 93개가 clamp되었다. 그래서 배치가 B = 1 루프 대비 얻은 이득은 약 1.06x(83.06 대 78.09 tok/s)에 그쳤고, 이는 walk가 가진 잠재력보다 훨씬 낮았다. + +이슈 #2442 항목 1은 per-row accepts를 요구했다. 이어서 PR #2853은 Metal 커널 디스패치 아래에서 배치 verify forward가 B = 1 forward와 바이트 단위로 동일하지 않다는 것을 확인하고, 하드 바이트 gate를 보정된 1.0-logit near-tie margin으로 대체했다. 이로써 선행 조건이 정리되었다. parity divergence의 원인은 accept 정책이 아니므로, per-row accepts는 기존 margin gate 아래에서 진행할 수 있었다. + +### 1.2 기존 문제점 + +- **clamp는 이미 받아들인 작업을 버린다.** 매 라운드가 live row 중 최소 accept를 commit하므로, 더 멀리 walk한 row의 추가 accepted 토큰은 rollback된다. +- **거부된 tail을 0으로 채우는 것은 유효한 대안이 아니다.** stale key가 softmax에서 `exp(0)` 가중치로 계속 참여한다. +- **직사각형 cache에는 하나의 offset만 있다.** 서로 다른 수를 commit한 row들이 같은 물리적 cache 할당을 공유하면서 서로 다른 연속 prefix를 유지할 방법이 필요하다. +- **linear-attention 상태는 row별이지만 유효 길이는 그렇지 않다.** `max(accepted)`로 한 번에 배치 replay하면 짧은 row가 거부한 위치까지 진행하게 된다. +- **배치 이득이 거의 없었다.** row-round당 committed accept 0.323에서 B = 4는 B = 1 대비 약 1.06x에 머물렀다. + +### 1.3 위험성 + +| 위험 | 영향 | 발생 가능성 | +|---|---|---| +| ragged cache 레이아웃이 이후 어텐션이나 GDN 상태를 오염시킨다 | 높음 (잘못된 토큰) | 낮음: parity gate가 per-row accepts를 기본값으로 실행되고, 압축은 단위 테스트로 검증된다 | +| 지원하지 않는 cache 모드가 per-row 경로에 도달한다 | 중간 (panic 또는 오염) | 낮음: `row_layout_supported`가 dense Fp16, 비-paged, 비-head-trimmed cache로 제한하고 압축도 같은 전제를 assert한다 | +| per-row accepts가 실제 서버 동시성과 잘 맞지 않는다 | 중간 (서빙 정확성) | 중간: 서버 수준 동시 요청 smoke 실행이 없다 | +| 호스트 부하가 throughput 비교를 왜곡한다 | 낮음 (판단 품질) | 낮음: 두 arm이 같은 경쟁 빌드 부하 아래 교대 실행되었고 clamp와 B = 1 셀이 이전 quiet-box 기록을 재현한다 | +| walk 감소를 회귀로 오해한다 | 낮음 (원인 오귀속) | 기록하지 않으면 확실: drafter context는 의도적으로 직사각형으로 유지되며, per-row drafter context는 별도 변경이다 | + +--- + +## 2. 기술적 검토 사항 + +### 2.1 보안 관점 + +네트워크, 인증, 입력 파싱, 직렬화 표면의 변경은 없다. 새로 추가된 설정 입력은 `MLXCEL_DFLASH_PER_ROW_ACCEPTS` 환경 변수뿐이며, generator 생성 시 한 번 읽고 on/off 값만 받는다. 사용자 데이터를 기록하지 않는다. `dflash_target.rs`가 남기는 accept 통계 로그는 기존 것이고 설명 문구만 갱신되었다. rollback 경로는 예상 밖 레이아웃으로 진행하지 않고 cache shape, mode, paging, head trim을 assertion으로 검증한다. + +### 2.2 성능 관점 + +M3 Ultra, `qwen3.5-4b-4bit` target과 `qwen3.5-4b-dflash`, block 16, row당 128 토큰, temperature 0, arm당 세 번의 교대 실행으로 측정했다. + +| 지표 | Clamp | Per-row | 변화 | +|---|---:|---:|---:| +| Aggregate throughput (tok/s) | 83.06 | 116.05 | 1.397x | +| B = 1 기준 (tok/s) | 78.09 | 78.09 | 1.00x | +| 라운드 수 | 96 | 67 | 30% 감소 | +| clamp된 라운드 | 93 | 0 | 제거 | +| row-round당 committed accept | 0.323 | 1.036 | 3.2x | +| row-round당 walk | 2.938 | 1.036 | 0.35x | + +per-row arm은 B = 1 루프의 1.49x이고, clamp arm은 1.06x이다. Throughput은 prefill을 제외한 `DFlashBatchedRunOutput::stats`의 라운드 루프 wall clock이므로, attention 레이어마다 row별로 수행하는 `copy`와 `slice_update` 압축 비용이 측정값에 포함되어 있다. clamp 세 실행의 편차는 0.2%, per-row 세 실행의 편차는 0.36%였다. + +### 2.3 호환성/의존성 관점 + +- **Breaking Changes**: 없음. 새 `SpeculativeTarget` 메서드는 이전 동작을 보존하는 기본 구현을 가진다. +- **New Dependencies**: 없음. +- **호환성**: per-row accepts는 target이 지원을 보고하고 cache가 head trim 없는 dense Fp16이며 paged가 아닐 때만 기본값이 된다. quantized, paged, head-trimmed cache는 자동으로 uniform clamp를 선택한다. `MLXCEL_DFLASH_PER_ROW_ACCEPTS=0`(또는 `false`, `no`, `off`)은 A/B 귀속을 위해 clamp를 복원한다. +- **내부 시그니처 변경**: `Qwen3NextAttention::forward_with_position_ids_verify`와 그 hidden-state 형제 함수에 `Option<&[i32]>` row-layout 인자가 추가되었다. 둘 다 `pub(crate)`이고 모든 호출 지점이 갱신되었다. +- **VLM**: `Qwen35VLModel`은 capability, offset, verify, rollback hook을 text model에 위임한다. attention 레이어는 MRoPE position id가 제공되면 per-row verify를 거부한다. DFlash verify 경로는 position id를 전달하지 않으므로 이는 서빙 경로가 아니라 방어적 가드이다. + +### 2.4 코드 품질 관점 + +단위 테스트 3개가 추가되었다. `models::qwen3_5_tests`의 attention 압축 테스트 2개(row block을 자신의 valid end로 이동, 모든 row가 offset에서 끝나면 no-op)와 `mlxcel-core::utils`의 `build_divergent_row_mask` 테스트 1개(history는 참여, stale gap은 차단, 새 block은 각 row의 논리 위치에서 causal)이다. `models::qwen3_5_tests` 스위트는 이제 30개 테스트이고, DFlash 라운드 루프 모듈은 `round_loop.rs` 15개, `round_loop_batched.rs` 13개를 유지한다. mask builder는 `gemma4.rs`에서 `mlxcel-core::utils`로 이동해 Gemma 4의 약 30줄 중복을 제거했고 두 모델이 하나의 구현을 공유한다. `cargo fmt --all -- --check`와 `cargo clippy --lib --tests --features metal,accelerate -- -D warnings`는 clean이다. + +--- + +## 3. 기술적 선택과 그 이유 + +### 3.1 target capability와 인스턴스 cache 검사로 per-row accepts를 gate + +**컨텍스트:** +per-row accepts는 attention cache를 row별로 압축할 수 있고 linear-attention 상태를 row별로 replay할 수 있는 target을 요구한다. 모든 `SpeculativeTarget`이 그럴 수 있는 것은 아니며, Qwen 3.5도 모든 cache 모드에서 가능한 것은 아니다. + +**고려한 대안:** + +| 옵션 | 장점 | 단점 | +|---|---|---| +| 모든 target에 per-row 경로를 항상 적용 | 루프가 단순 | uniform rollback만 지원하는 target을 깨뜨리고 quantized/paged cache에서 panic | +| 모든 target이 새 hook을 구현하도록 강제 | 동작이 균일 | 무관한 target과 테스트 더블에 no-op 구현을 강요 | +| **선택: 기본 trait 메서드와 2단계 gate** | 기존 target은 변경 없음, 미지원 cache는 안전하게 fallback | 서빙 경로가 런타임 cache 상태에 의존 | + +**선택 이유:** +`supports_per_row_accepts()`는 family 수준 답이고 기본값이 `false`이다. `row_layout_supported(caches)`는 인스턴스 수준 답이며 `KVCacheMode::Fp16`, `!is_paged_backed()`, `offset - live_len() == 0`(head trim 없음)을 검사한다. 루프는 둘 다 참이고 호출자가 비활성화하지 않았을 때만 per-row 경로를 탄다. 기본 `rollback_partial_batched_rows`와 `verify_forward_with_row_layout`은 uniform hook에 위임하며, 이는 모든 row가 offset을 공유할 때만 올바르므로 override를 잊은 target은 ragged 레이아웃을 조용히 만들지 않고 안전하게 실패한다. + +**트레이드오프:** +프로덕션 동작이 런타임에 선택된 cache 모드에 의존하므로, quantized 배포는 조용히 기존 clamp를 사용하고 1.397x를 얻지 못한다. 대안은 burst를 panic 또는 거부하는 것이었지만, 이는 성능 기능을 가용성 회귀로 바꾼다. + +### 3.2 두 번째 레이아웃을 만들지 않고 Gemma 4 divergent-row 레이아웃을 재사용 + +**컨텍스트:** +Gemma 4 MTP verifier는 이미 논리적 valid end가 다른 row들을 공유 물리 cache에 대해 처리한다. 그 `build_divergent_verify_mask`는 ragged DFlash 라운드에 필요한 세 영역, 즉 row의 valid history, stale gap, 새 block을 정확히 표현한다. + +**고려한 대안:** + +| 옵션 | 장점 | 단점 | +|---|---|---| +| 모든 row를 `max(accepted)`로 패딩하고 거부 위치를 0으로 채움 | 압축 불필요 | 0 key가 여전히 softmax 가중치를 차지하며 #2353에서 이미 기각됨 | +| cache를 row별로 복사하고 독립 forward 실행 | 개념적으로 단순 | 배치 이득을 없애고 row마다 cache 메모리를 중복 | +| **선택: 공유 divergent-row mask와 per-row 회전** | 하나의 mask가 세 영역을 모두 표현하고 검증된 레이아웃을 재사용 | mask 생성 비용이 `O(B * l * (offset + l))` | +| DFlash 전용 mask를 새로 구현 | 모델 간 결합 없음 | 정확성을 유지해야 할 두 번째 구현 | + +**선택 이유:** +`build_divergent_row_mask(l, offset, per_row_valid_end, left_padding, window)`는 `[B, 1, l, offset + l]` additive mask를 반환한다. row `r`은 valid end 아래 key를 허용하고, `[valid_end, offset)`을 차단하며, 새 key `t`를 논리 위치 `valid_end + t`로 취급하므로 query `j`는 새 key `t <= j`만 본다. `Qwen3NextAttention`은 `target_verify`, `l > 1`, 그리고 어떤 end가 공유 offset과 다를 때 `divergent_rows`를 감지하고, row 자신의 시작 위치에서 RoPE를 적용한 뒤 `attend_per_position_rows`를 실행한다. 이 함수는 query `j`를 정확한 key 폭 `offset + j + 1`과 해당 mask 슬라이스에 대해 계산한다. Gemma 4 함수는 이제 같은 core utility에 위임하므로 두 모델이 어긋날 수 없다. + +**트레이드오프:** +mask는 verify forward마다 생성되고 명시적으로 materialize된다. 이 비용은 측정된 throughput에 포함되어 있고, 측정된 shape에서는 accept 이득이 이를 상회한다. 향후 fused 경로가 materialize를 피할 수 있지만 현재 이득에는 필요하지 않다. + +### 3.3 drafter context는 가장 짧은 live committed prefix에서 직사각형으로 유지 + +**컨텍스트:** +target은 이제 ragged cache를 유지할 수 있지만, drafter의 입력 hidden 텐서와 자체 context는 여전히 하나의 직사각형 슬라이스로 만들어진다. drafter를 per-row로 만들려면 row별 drafter cache view와 row별 drafter forward가 필요하며, 이는 훨씬 큰 변경이다. + +**고려한 대안:** + +| 옵션 | 장점 | 단점 | +|---|---|---| +| 지금 drafter에 per-row context 제공 | walk 상승 | row별 drafter cache와 forward 배관이 필요하고 이 PR에서 측정되지 않음 | +| drafter에 가장 긴 committed prefix 제공 | 뒤처진 row에 더 많은 draft 재료 | verify block에서 잘라낸 hidden slice는 가장 짧은 accepted block 위치까지만 유효하므로 거부된 위치가 context에 섞임 | +| **선택: 가장 짧은 live committed prefix** | 캡처된 hidden slice에 대해 올바르고 기존 직사각형 계약을 재사용 | walk가 2.938에서 1.036으로 감소 | + +**선택 이유:** +`drafter_ctx`는 완료되지 않은 row들의 `accepted[r] + 1` 최솟값으로 계산되고 hidden 텐서는 그 길이로 slice된다. 최솟값보다 많이 commit한 row도 target cache를 자신의 accept만큼 진행시키므로 방출 스트림은 정확하다. drafter는 다음 라운드에서 뒤처진 context로 다시 draft할 뿐이다. 이 선택은 변경을 target 쪽으로 제한하고 throughput 비교를 accept 정책에만 귀속시킨다. + +**트레이드오프:** +walk가 65% 감소하므로 라운드 수 이득은 accept 수 이득이 시사하는 것보다 작다. 그래도 per-row accepts는 라운드당 방출 토큰을 결정하는 것이 walk가 아니라 committed 수이므로 분명한 이득이다. per-row drafter context는 이 PR의 결함이 아니라 후속 작업으로 기록된다. + +### 3.4 attention KV를 압축하고 GDN 상태를 row별로 replay + +**컨텍스트:** +ragged 라운드가 끝나면 공유 cache는 각 row의 accepted block을 공유 물리 offset에 두고 그 뒤에 거부된 위치를 담고 있다. 다음 라운드는 모든 row가 연속적인 `[0, valid_end)` prefix처럼 보여야 하며 물리 offset은 공유되어야 한다. + +**선택 이유:** +`rollback_speculative_cache_rows`는 `o_post = max(valid_end_pre[r] + accepted[r] + 1)`을 계산한다. `compact_attention_rows`는 각 row의 `accepted[r] + 1` block을 `o_pre`에서 해당 row의 논리 end로 이동하고, 비워진 gap을 `o_post`까지 0으로 채우며, 호출자가 cache offset을 `o_post`로 trim한다. 겹치는 이동이 이미 덮어쓴 row를 읽지 않도록 source를 먼저 복사한 뒤 `slice_update`를 수행한다. GDN 상태는 row별로 replay된다. `replay_gdn_rows`는 각 row의 캡처된 `q/k/v/a/b`와 초기 상태를 `accepted[r] + 1`로 slice하고, chain-parity 커널을 해당 row의 길이로 실행한 뒤 per-row conv 및 recurrent 상태를 `[B, ...]` cache 텐서로 다시 쌓는다. 이로써 per-row 수치가 uniform 경로와 같은 커널을 유지한다. + +**트레이드오프:** +압축은 라운드마다 attention 레이어마다 row마다 copy 한 번과 `slice_update` 한 번을 추가한다. 이는 보고된 throughput 안에서 측정되었다. uniform 경로도 fallback과 A/B 사용을 위해 그대로 유지되므로 코드는 두 arm을 함께 가진다. 이 중복은 귀속을 위해 의도된 것이다. + +--- + +## 4. 구현 상세 + +### 4.1 아키텍처 변경 + +``` +[Before, uniform clamp] +verify [B, bs] -> walk per row -> clamp_uniform_acceptance (min over live rows) + -> rollback_partial_batched (trim bs - (a + 1) in every row) + -> drafter hidden [B, max_a + 1] + +[After, per-row accepts] +row_valid_ends = [prefill offset; B] +verify [B, bs] + -> per-row RoPE at row_valid_ends[r] + -> attend with build_divergent_row_mask (history, stale gap, causal block) + -> commit accepted[r] + 1 per row + -> compact_attention_rows: block moves to row_valid_ends[r], gap zeroed + -> cache offset trimmed to max(new valid ends) + -> replay_gdn_rows: per-row conv + recurrent state over accepted prefix + -> row_valid_ends[r] += accepted[r] + 1 + -> drafter hidden [B, min_live(accepted) + 1] (rectangular, deliberately) +``` + +### 4.2 주요 코드 변경 + +**파일: `src/lib/mlxcel-core/src/drafter/dflash/round_loop_batched.rs`** + +generator는 실행당 한 번 정책을 결정하고 prefill offset으로 논리 end를 초기화한다: + +```rust +let per_row = self.per_row_accepts + && target.supports_per_row_accepts() + && target.row_layout_supported(caches); +let mut row_valid_ends: Vec = if per_row { + let offset = T::batch_cache_offset(caches); + vec![offset; batch_size] +} else { + Vec::new() +}; +``` + +verify 호출은 정책에 따라 분기하고, per-row arm에서는 clamp를 건너뛰며, rollback은 각 row의 accepted prefix를 target에 넘긴다: + +```rust +if per_row { + target.rollback_partial_batched_rows( + caches, &verify_out, &accepted_i32, &row_valid_ends, bs as i32, + ); + for r in 0..batch_size { + row_valid_ends[r] += accepted_i32[r] + 1; + } +} else { + // uniform clamp arm: trim bs - (accepted + 1) in every row +} +``` + +drafter context는 가장 짧은 live row에서 직사각형으로 유지된다: + +```rust +let drafter_ctx = if per_row { + (0..batch_size) + .filter(|&r| !finished[r]) + .map(|r| accepted_per_row[r] + 1) + .min() + .unwrap_or(1) +} else { + max_accepted + 1 +}; +``` + +**변경 이유:** 라운드 루프가 정책 전환과 논리 위치 장부를 소유하고 target이 cache 레이아웃을 소유한다. 전환 지점을 한 곳에 두면 한 프로세스에서 A/B를 측정할 수 있다. + +**파일: `src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs`** + +`SpeculativeTarget`에 기본 구현이 있는 hook 네 개가 추가된다. `supports_per_row_accepts`, `row_layout_supported`, `batch_cache_offset`, `verify_forward_with_row_layout`이며, 여기에 `rollback_partial_batched_rows`가 더해진다. 기존 `rollback_partial_batched` 계약은 문서상 uniform-clamp arm으로 좁혀진다. + +**변경 이유:** ragged row를 표현할 수 없는 target도 계속 컴파일되고 기존 의미를 유지한다. + +**파일: `src/models/qwen3_5.rs`** + +`Qwen35Model`은 capability gate, 공유 offset, ragged verify forward, ragged rollback을 구현한다. `forward_speculative_inner`는 `Option<&[i32]>`를 각 decoder 레이어로 전달한다. 비균일 rollback은 다음과 같다: + +```rust +if !uniform { + let o_post = ve_pre.iter().zip(accepted) + .map(|(&ve, &a)| ve + a + 1).max().expect("non-empty accepted"); + for cache in caches.iter_mut() { + if let Qwen3NextCache::Attention(kv) = cache { + compact_attention_rows(kv, &ve_pre, accepted, o_pre, o_post); + let roll_back = o_pre + block_size - o_post; + if roll_back > 0 { kv.trim(roll_back); } + } + } + self.replay_gdn_rows(caches, gdn_states, accepted, o_post); + return accepted.iter().copied().max().unwrap_or(0); +} +``` + +**변경 이유:** attention cache가 다음 라운드 전에 row별로 유효한 연속 prefix가 되어야 하고, GDN 상태는 해당 row가 accept한 토큰만 반영해야 한다. + +**파일: `src/models/qwen3_next.rs`** + +`Qwen3NextAttention`에 `apply_rope_per_row`와 `attend_per_position_rows`가 추가된다. 어떤 row의 valid end라도 cache offset과 다르면 divergent 경로가 선택된다: + +```rust +let divergent_rows: Option<&[i32]> = row_valid_ends.filter(|ends| { + target_verify + && l > 1 + && ends.len() == b as usize + && ends.iter().any(|&end| end != offset) +}); +``` + +각 query 위치는 B = 1 체인이 보았을 정확한 key 폭을 slice하며, divergent-row mask가 history, gap, block causality를 인코딩한다. + +**변경 이유:** 공유 직사각형 attention 호출은 각 row에 다른 논리 시작점을 줄 수 없고, per-position 경로는 줄 수 있으며 Gemma 4 방식과 일치한다. + +**파일: `src/lib/mlxcel-core/src/utils.rs`와 `src/models/gemma4.rs`** + +`build_divergent_row_mask`가 `mlxcel-core::utils`로 이동하고, Gemma 4의 `build_divergent_verify_mask`는 한 줄 위임이 된다. + +**변경 이유:** Gemma 4 MTP verifier와 Qwen 3.5 DFlash verifier가 하나의 테스트된 mask 구현을 공유한다. + +**파일: `src/vision/qwen3_5_vl.rs`** + +`Qwen35VLModel`이 네 개의 새 hook을 text model에 전달한다. + +**변경 이유:** VLM 래퍼가 capability나 fallback을 조용히 잃지 않아야 한다. + +**파일: `docs/CONTINUOUS_BATCHING.md`, `docs/environment-variables.md`, `docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md`** + +배칭 가이드와 환경 변수 표가 기본값, fallback, 측정값을 기록하고, benchmark 문서가 방법, 실행별 값, caveat, 정확성 gate를 기록한다. + +### 4.3 데이터 모델 변경 + +없음. cache 레이아웃은 `[B, heads, capacity, dim]`으로 유지되고, 라운드마다 변하는 것은 valid-prefix 장부와 물리 offset뿐이다. + +--- + +## 5. 학습 포인트 + +### 5.1 공유 직사각형 cache 안에서의 ragged acceptance + +**개념:** +공유 배치 KV cache는 물리 offset 하나와 할당 하나를 가진다. ragged accept 수는 매 라운드 후 각 row의 live content를 연속 prefix로 압축하면 row별 할당을 요구하지 않는다. 불변식은 "각 row의 cache content는 연속적인 `[0, valid_end)` prefix이고, 공유 offset까지의 stale 영역은 마스킹된다"가 된다. + +**이 PR에서의 적용:** +`compact_attention_rows`가 accepted block을 공유 offset에서 row의 논리 end로 이동하고 비워진 gap을 0으로 채운다. 다음 verify forward에서는 `build_divergent_row_mask`가 세 영역을 인코딩한다. 물리 offset은 새 valid end의 최댓값으로 trim되므로 할당이 raggedness에 따라 커지지 않는다. + +**일반적인 사용 사례:** +- 시퀀스별 accept 수를 가진 배치 speculative decoding. +- 시퀀스가 서로 다른 길이에서 멈추고, 재개하고, rollback하는 continuous batching. +- row별 버퍼 없이 row별 논리 길이를 유지해야 하는 공유 버퍼 기법. + +### 5.2 하나의 mask 레이아웃, 두 모델 + +**개념:** +"valid history, 차단된 gap, 논리 위치에서의 causal 새 block"을 표현하는 mask는 모델에 독립적이다. Gemma 4 구현을 core utility로 올려 두 번째 복사본이 어긋나기 전에 제거했다. + +**이 PR에서의 적용:** +`build_divergent_row_mask`는 이제 Gemma 4의 `build_divergent_verify_mask`와 Qwen 3.5의 `attend_per_position_rows`가 호출한다. 이동과 함께 mask 전용 단위 테스트도 생겼다. + +**일반적인 사용 사례:** +- 같은 cache 토폴로지를 가진 모델 패밀리 간 attention mask 재사용. +- 수치에 민감한 디바이스 코드를 하나의 테스트된 생성자 뒤로 통합. + +### 5.3 speculative 배칭은 walk가 아니라 committed accept로 판단한다 + +**개념:** +scan ratio(walk)는 drafter가 제안한 양을 측정하고, committed accept 수는 target이 실제로 유지한 양을 측정한다. 배칭 변경이 accepted 토큰을 버리지 않게 되면 walk가 낮아져도 throughput은 올라갈 수 있다. + +**이 PR에서의 적용:** +per-row arm의 walk는 drafter context가 여전히 가장 짧은 live prefix이기 때문에 2.938에서 1.036으로 떨어졌지만, committed accept가 0.323에서 1.036으로 오르고 라운드가 96에서 67로 줄어 aggregate throughput은 1.397x 상승했다. walk만 보고하면 회귀처럼 읽혔을 것이다. + +**일반적인 사용 사례:** +- speculative decoding A/B 테스트. +- 생산자의 제안률과 소비자의 수용률이 반대 방향으로 움직일 수 있는 모든 파이프라인. + +--- + +## 6. 추가 학습 리소스 + +### 핵심 키워드 + +| 키워드 | 설명 | 이 PR에서의 관련성 | +|---|---|---| +| DFlash | diffusion 스타일 drafter를 쓰는 block-draft speculative decoding | 여기서 변경된 배치 루프 | +| Uniform clamp | 가장 짧은 live row의 accept를 모든 row에 commit | 이전 정책이며 여전히 도달 가능 | +| Ragged accept | 한 라운드에서 row마다 다른 accept 수 | per-row accepts가 commit하는 것 | +| Valid end | row의 연속적인 유효 cache prefix 길이 | RoPE와 attention에 쓰이는 논리 위치 | +| Stale gap | 더 짧은 accept가 남긴 `[valid_end, physical_offset)` | attention에서 마스킹되고 압축에서 0으로 채워짐 | +| GDN | gated delta net linear-attention 레이어 | conv와 recurrent 상태를 row별로 replay | +| Walk | drafter가 제안하고 target이 verify한 토큰 | drafter context가 직사각형으로 유지되어 여기서는 감소 | + +### 관련 기술/프레임워크 + +- **MLX**: array 프레임워크와 Metal 커널. `copy`와 `slice_update` 압축, per-position attention slice가 이 위에서 실행된다. +- **Gemma 4 divergent MTP verify**: 이 PR이 재사용하는 `DivergentVerifyRows` mask 레이아웃의 출처. + +### 관련 PR/이슈 + +- 이슈 #2442: 상위 이슈. 이 PR이 항목 1을 구현하고 종료한다. +- 이슈 #2353: uniform clamp와 최초의 B = 4 대 B = 1 측정을 도입했다. +- PR #2853 / 이슈 #2851: 이 PR의 gate가 실행되는 배치 verify parity margin을 보정했다. +- 이슈 #2190: mask 레이아웃을 제공한 Gemma 4 per-row verify 아키텍처. +- 이슈 #2440: benchmark 문서가 참조하는 이전 per-row 범위 설정 작업. + +--- + +## 7. 변경 요약 + +### 통계 + +| 항목 | 값 | +|---|---| +| 변경 파일 | 15 | +| 추가 라인 | +1279 | +| 삭제 라인 | -156 | +| 테스트 추가 | 단위 테스트 3개, 무시되는 real-model throughput harness 1개 | +| 런타임 동작 변경 | 있음: cache가 압축을 지원하면 per-row accepts가 기본값 | + +### 카테고리별 변경 + +| 카테고리 | 수 | 요약 | +|---|---|---| +| 성능/런타임 | 7 | 라운드 루프 정책 전환, target hook, Qwen 3.5 verify와 rollback, Qwen3Next per-row RoPE와 attention, VL 위임, 서버 로그 문구 | +| 공유 유틸리티 | 1 | `build_divergent_row_mask`를 `mlxcel-core::utils`로 이동, Gemma 4 위임 | +| 테스트 | 3 | 압축 단위 테스트, mask 단위 테스트, throughput 비교 harness | +| 문서 | 4 | 배칭 가이드, 환경 변수, benchmark 기록, parity harness 연결 | + +### 관련 커밋 + +| 해시 | 유형 | 메시지 | +|---|---|---| +| `ded78dce` | perf | perf(speculative): commit per-row accepts in batched DFlash | + +--- + +## 8. 후속 조치 + +### 완료 필요 + +- [ ] per-row accepts를 켠 상태에서 실제 수집 burst 경로를 통한 서버 수준 동시 요청 smoke 테스트를 B > 1로 실행한다. 요청 완료, EOS, max-token row, 프롬프트 길이가 섞인 window를 포함해야 한다. 현재 근거는 길이가 같은 네 프롬프트에 대한 harness burst뿐이다. +- [ ] per-row drafter context를 별도 변경으로 구현하고 다시 측정한다. walk가 1.036 위로 회복될 것으로 예상되지만 이 PR의 예산에는 포함되지 않는다. + +### 미검증 영역 + +- **서버 수준 동시 요청 smoke 실행 없음.** 측정과 parity 실행은 테스트 harness(`run_burst_with_policy`)를 사용하며 서버의 수집 burst 스케줄러를 사용하지 않는다. `src/server/batch/dflash_target.rs`의 row 수명, admission, 취소, 로깅은 기존 스위트로만 검증되고 이 측정으로는 검증되지 않는다. +- **per-row drafter context는 범위 밖.** drafter는 여전히 공유 직사각형 최소 prefix를 본다. 이것이 walk가 1.036으로 떨어진 명시된 이유이며, 측정된 1.397x가 drafter를 다룬 뒤 per-row target accepts가 낼 수 있는 하한이라는 뜻이다. +- **측정 caveat.** throughput 실행 중 다른 worktree의 Rust 빌드가 80개 코어 중 몇 개에서 컴파일하고 있었다. 두 arm은 같은 부하 아래 교대 실행되었고 clamp와 B = 1 셀이 이전 quiet-box #2442 기록(83.0 대 82.96~83.00, 78.09 대 77.70~77.83)을 재현하므로 비교가 부하로 치우치지는 않았다. 다만 절대값은 완전히 유휴한 박스에서 측정한 것이 아니다. +- **단일 호스트와 체크포인트.** M3 Ultra, macOS 26.6.2, `qwen3.5-4b-4bit`와 `qwen3.5-4b-dflash`, block 16, row당 128 토큰, temperature 0, 잘린 프롬프트 네 개다. 다른 하드웨어, quantization, block 크기, 프롬프트 분포는 측정되지 않았다. +- **fallback 경로는 여기서 재검증되지 않았다.** quantized, paged, head-trimmed cache는 변경되지 않은 uniform clamp를 탄다. 그 정확성은 상속된 것이며 per-row gate로 다시 테스트되지 않는다. +- **VLM 경로.** `Qwen35VLModel`이 hook을 위임하고 attention 레이어는 MRoPE position id가 있으면 divergent row를 거부한다. 이 PR에는 VLM DFlash parity 실행이 포함되지 않는다. + +### 모니터링 필요 + +- per-row arm에서 `clamped_rounds`는 0을 유지해야 하고, clamp를 명시적으로 선택했거나 cache가 fallback한 경우에만 0이 아니어야 한다. dense Fp16 실행에서 0이 아니면 capability gate가 회귀한 것이다. +- 배포 후 aggregate decode throughput과 row-round당 committed accept를 관찰한다. 특히 raggedness가 가장 큰 혼합 프롬프트 길이에서 중요하다. +- MLX, quantization, 커널 디스패치가 바뀔 때마다 CI 하드웨어 레인에서 `greedy_parity_dflash_qwen35_batched_b4_matches_b1`와 `identical_prompts_dflash_qwen35_b4_matches_b1`를 실행한다. + +### 향후 개선 사항 + +- 현재 clamp로 fallback하는 quantized 및 paged cache로 per-row accepts를 확장한다. +- mask 생성이 라운드 비용에서 측정 가능한 비중이 되면 fused mask-free attend 경로를 검토한다. +- per-row target accepts와 per-row drafter context를 함께 평가한다. 둘은 walk를 통해 상호작용하므로 쌍으로 측정해야 한다. + +--- + +## 부록 + +### A. 테스트 결과 + +| 검사 | 결과 | +|---|---| +| `cargo test --profile test-fast --features metal,accelerate --test speculative_parity qwen35_dflash_batched -- --ignored --test-threads=1` | per-row accepts를 기본값으로 2개 통과 | +| `throughput_dflash_qwen35_batched_clamp_vs_per_row` | 완료, arm당 세 번 교대 실행 | +| `cargo test --profile test-fast --lib models::qwen3_5_tests` | 통과, 30개 테스트(새 압축 테스트 2개 포함) | +| mlxcel-core 라운드 루프 단위 테스트 | 통과, `round_loop.rs` 15개, `round_loop_batched.rs` 13개 | +| `build_divergent_row_mask` 단위 테스트 | 통과 | +| `cargo clippy --lib --tests --features metal,accelerate -- -D warnings` | 통과 | +| `cargo fmt --all -- --check` | 통과 | + +### B. 성능 벤치마크 + +| 실행 | Clamp tok/s | Per-row tok/s | +|---:|---:|---:| +| 0 | 83.04 | 115.82 | +| 1 | 83.22 | 116.24 | +| 2 | 83.06 | 116.05 | +| 중앙값 | 83.06 | 116.05 | + +accept와 라운드: + +| Arm | 라운드 | clamp된 라운드 | row-round당 walk | row-round당 committed | +|---|---:|---:|---:|---:| +| Clamp | 96 | 93 | 2.938 | 0.323 | +| Per-row | 67 | 0 | 1.036 | 1.036 | + +### C. 측정 환경 + +| 항목 | 값 | +|---|---| +| 호스트 | Apple M3 Ultra, 80코어, 512 GB 통합 메모리, macOS 26.6.2 | +| 백엔드 | Metal + Accelerate | +| Target | `models/qwen3.5-4b-4bit` | +| Drafter | `models/qwen3.5-4b-dflash` | +| Block 크기 | 16 | +| 생성 토큰 | row당 128 | +| 샘플링 | Temperature 0, greedy | +| 프롬프트 | 배치 DFlash parity harness의 네 프롬프트를 가장 짧은 토크나이제이션에 맞춰 절단 | +| 방법 | 폐기한 warm-up 2회 후 arm당 세 번 측정, clamp와 per-row를 교대 실행, `caffeinate -i`와 `scripts/with_indexers_paused.sh` 사용 | + +### D. 참고 자료 + +- `src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs` +- `src/lib/mlxcel-core/src/drafter/dflash/round_loop_batched.rs` +- `src/lib/mlxcel-core/src/utils.rs` +- `src/models/qwen3_5.rs` +- `src/models/qwen3_next.rs` +- `src/models/gemma4.rs` +- `src/vision/qwen3_5_vl.rs` +- `tests/speculative_parity/qwen35_dflash_harness.rs` +- `tests/speculative_parity/qwen35_dflash_throughput.rs` +- `docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md` +- 이슈 #2442, 이슈 #2353, PR #2853, 이슈 #2851, 이슈 #2190, 이슈 #2440 diff --git a/docs/CONTINUOUS_BATCHING.md b/docs/CONTINUOUS_BATCHING.md index 4e4a3ddab..2663e6f4e 100644 --- a/docs/CONTINUOUS_BATCHING.md +++ b/docs/CONTINUOUS_BATCHING.md @@ -708,7 +708,16 @@ kernels run, so a row can pick a different token where two candidates are within about a quarter logit. The parity gate admits a first mismatch only inside a calibrated 1.0-logit near-tie margin, and a four-identical-prompts control shows the divergence is in the batched verify forward rather than the accept -policy (#2851). On the Gemma 4 +policy (#2851). The round then commits each row's own accept count (#2442): +the attention KV caches are compacted to per-row contiguous valid prefixes, +queries and keys rotate at each row's logical positions with the row's stale +gap masked, and each GDN layer replays its state over that row's accepted +prefix. `MLXCEL_DFLASH_PER_ROW_ACCEPTS=0` restores the uniform clamp, which a +quantized or paged KV cache also selects automatically. On an M3 Ultra the +per-row arm measured 1.397x the clamp's B = 4 aggregate throughput (committed +accepts 1.036 against 0.323 per row-round); the drafter's context is still a +shared rectangular prefix, so the walk itself is lower than the clamp arm's, +and per-row drafter context is a follow-up. On the Gemma 4 geometries whose MTP verify runs row by row (the 31B everywhere, the 12B on CUDA) the batched burst is decode-exact since issue #2190: each row is prefilled and verified through the B=1 calls and only the cache write is diff --git a/docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md b/docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md new file mode 100644 index 000000000..d2a7fac49 --- /dev/null +++ b/docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md @@ -0,0 +1,74 @@ +# Batched DFlash per-row accepts on M3 Ultra, 2026-11-06 + +Issue #2442 item 1 follow-up to #2353 / #2440 and #2851. The batched Qwen 3.5 +DFlash round loop used to clamp every round to the shortest live row's accept +because the attention KV cache is rectangular. That left the committed accept +count far below the speculative walk, so batching bought little. This measures +the per-row accept port that replaces the clamp. + +## What changed + +- `DFlashBatchedGenerator` commits each row's own walk count. The target gets + each row's pre-round logical valid end, compacts the attention KV caches back + to per-row contiguous prefixes, and replays each GDN layer's conv and + recurrent state over that row's accepted prefix + (`Qwen35Model::rollback_speculative_cache_rows`). +- The verify forward rotates queries and keys at each row's logical positions + and masks the row's stale gap (`build_divergent_row_mask`, + `attend_per_position_rows`), the layout Gemma 4's divergent MTP verify uses. +- `MLXCEL_DFLASH_PER_ROW_ACCEPTS=0` restores the uniform clamp. A quantized, + paged or head-trimmed KV cache selects the clamp automatically. + +## Method + +- Host: Apple M3 Ultra (80-core, 512 GB), macOS 26.6.2, Metal + Accelerate. +- Checkpoints: `models/qwen3.5-4b-4bit` target, `models/qwen3.5-4b-dflash` + drafter, block size 16, temperature 0. +- Prompts: the four prompts of the batched DFlash parity harness, truncated to + the shortest tokenization so the `[B, L]` prefill is rectangular. +- 128 tokens per row. Throughput is the round-loop wall clock from + `DFlashBatchedRunOutput::stats` (prefill excluded), aggregate over rows. +- Both arms run in one process behind + `throughput_dflash_qwen35_batched_clamp_vs_per_row`, two discarded warm-up + runs, three measured runs per arm alternating clamp, per-row, clamp, + per-row, ... under `caffeinate -i` and `scripts/with_indexers_paused.sh`. + A separate worktree's Rust build was compiling during the run (a few of the + 80 cores); both arms alternated under that same load, and the clamp and + B = 1 cells reproduce the quiet-box #2442 record (83.0 against 82.96-83.00, + 78.09 against 77.70-77.83), so the comparison is not load-shifted. + +## Throughput at B = 4 + +| run | clamp tok/s | per-row tok/s | +|---:|---:|---:| +| 0 | 83.04 | 115.82 | +| 1 | 83.22 | 116.24 | +| 2 | 83.06 | 116.05 | +| median | 83.06 | 116.05 | + +Per-row is **1.397x** the clamp. The clamp's three runs span 0.2%, the per-row +arm's 0.36%. For context, B = 1 measured 78.09 tok/s, so the per-row arm is +1.49x over the B = 1 round loop and the clamp is 1.06x. + +## Accepts + +| arm | rounds | clamped rounds | walk per row-round | committed per row-round | +|---|---:|---:|---:|---:| +| clamp | 96 | 93 | 2.938 | 0.323 | +| per-row | 67 | 0 | 1.036 | 1.036 | + +The clamp numbers reproduce the #2442 re-scope measurement exactly (walk +2.938, committed 0.323, 93 of 96 rounds clamped). Per-row commits 3.2x the +clamp's accept count and finishes in 30% fewer rounds. + +The walk itself falls from 2.938 to 1.036: the drafter's context is still a +shared rectangular prefix (the shortest live row's committed prefix), while the +target now advances each row by its own accept, so the drafter's context lags +the rows that accept more. Giving the drafter per-row context is a separate +change; this port was measured and decided on its own. + +## Correctness + +`greedy_parity_dflash_qwen35_batched_b4_matches_b1` and the +`identical_prompts_dflash_qwen35_b4_matches_b1` control pass with the port +enabled under the calibrated 1.0-logit evaluation-path margin of #2851. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 5d780615d..ef36ed66b 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -621,6 +621,7 @@ The OpenAI audio endpoints (`/v1/audio/speech`, `/v1/audio/transcriptions`, `/v1 | `MLXCEL_ENABLE_MTP_B1` | `0`/`false`/`no`/`off` to disable, any other value to force on | adaptive (per hardware) | Manual override for the singleton (B=1) MTP burst, in both directions. When set it pins the decision and disables adaptive profiling (issue #333). When unset, the adaptive policy decides (see `MLXCEL_MTP_ADAPTIVE`); with `MLXCEL_MTP_ADAPTIVE=0` the decision is the static per-hardware default (issue #165): non-batchable targets (`gemma4_unified` 12B pairs, whose only decode path is B=1) default **on** everywhere (measured across three prompts: 1.90x to 3.14x on M5 Max, 1.74x to 2.61x on M3 Ultra, 0.95x to 1.48x on M1 Ultra, where the prose end is a 5% loss rather than a gain); batch-capable targets (the 31B + bf16 assistant) default **on from Apple GPU generation 15** (M3, M4, M5), and fall back to classic decode on generation 13 (M1, M2). This was `has_neural_accelerator` (M5 only) until issue #1217, on the strength of ~1.2 to 1.4× on M5 Max against a ~0.75 to 0.96× regression on M1 Ultra, both measured before #1194/#1199/#1203/#1208/#1215 and neither re-measured since. M3 Ultra, which the old predicate lumped with M1 Ultra, was never measured on this pairing at all until #1217 did it: 1.95× (prose), 2.41× (source code) and 2.65× (enumeration) on 2026-08-20 under the #1215 protocol, at a verify-round cost of 1.51 classic steps against 2.71 on M1 Ultra. The discriminator is the `use_qmv_wide` split (an affine-quantized projection at `M >= 2` runs as one wide pass from generation 15 and as `K` narrow passes below it), not the Neural Accelerator. M4 is grouped by that shared dispatch rather than measured. Generation 13 keeps declining and the width sweep supports that: round cost fits `0.83 + 0.170 K` classic steps for this pairing on M3 Ultra against `1.14 + 0.090 K` for the 12B pairing there, so the bf16 drafter costs about 1.9x as much per block position, and carrying that ratio onto generation 13's `1.35 + 0.346 K` puts a block-4 round near 3.6 classic steps, which 2.96 to 3.99 emitted tokens would only just cover. See `docs/benchmark_results/mtp-b1-gate-m3ultra-2026-08-20.md`. The 31B singleton exactness correction in #1983 restores the automatic narrow retry on M5 Max; its attention cost means these historical pre-gate speedups are not current guarantees. See the [matched QAT/non-QAT report](benchmark_results/gemma4-31b-mtp-exactness-2026-09-26.md). | | `MLXCEL_ENABLE_MTP_BATCH` | truthy value | off | **Advanced.** Forces the batched Gemma 4 MTP burst path for parity/debug testing. Not governed by the adaptive policy (issue #333), which scopes to the validated, byte-identical B=1 path. On the row-wise geometries (Gemma 4 31B, and the 12B on CUDA) the batched burst is decode-exact (issue #2190) and runs on CUDA for windows of up to 4 rows whose prompts fit the sliding window and whose `max_prompt_len + 4 * (max_tokens + 1)` stays within the window plus the 32-token rollback buffer; other windows decline to classic decode. A window forms only from requests that are queued together, share a seed and carry no token bias (`--ignore-eos` adds one). With this flag set, a queued burst-eligible head skips the classic batched prefill so the window can form under the default `--max-batch-prefill` (issue #2442); before that, it formed only with `--max-batch-prefill 1`. See `docs/benchmark_results/gemma4-mtp-batched-gb10-2026-10-07.md`. | | `MLXCEL_ENABLE_MTP_BATCH_RAGGED` | exactly `1`, `true`, `TRUE`, `yes` or `on` (case-sensitive) | off | **Advanced, experimental.** Only matters with `MLXCEL_ENABLE_MTP_BATCH`. Drops the equal-prompt-length requirement for a batched Gemma 4 MTP window: rows are left-padded to the longest prompt and keep per-row positions. A window whose longest prompt exceeds the target's sliding window declines and returns to classic decode. DFlash windows stay equal-length. | +| `MLXCEL_DFLASH_PER_ROW_ACCEPTS` | `0`/`false`/`no`/`off` to disable; unset or any other value to enable | on | Per-row accept commits in the batched Qwen 3.5 DFlash round loop (issue #2442). When on, each row commits its own speculative-walk accept count instead of clamping the round to the shortest live row (the uniform clamp of issue #2353): the attention KV caches are compacted back to a per-row contiguous valid prefix, the block's queries and keys rotate at each row's own logical positions with that row's stale gap masked out, and each linear-attention (GDN) layer replays its conv and recurrent state over the row's own accepted prefix. The clamp stays reachable with `0` for A/B attribution; it is also taken automatically when the target's attention caches are quantized, paged or head-trimmed, because those cannot be compacted row by row. Measured on an M3 Ultra (`qwen3.5-4b-4bit` plus `qwen3.5-4b-dflash`, block 16, 128 tokens per row, temperature 0): B = 4 aggregate decode 116.05 tok/s against the clamp's 83.06 (1.397x) and B = 1's 78.09, with committed accepts per row-round 1.036 against the clamp's 0.323. Greedy parity still passes the calibrated `greedy_parity_dflash_qwen35_batched_b4_matches_b1` near-tie gate. The drafter's context remains a shared rectangular prefix, so the walk itself falls from the clamp arm's 2.938 to 1.036; per-row drafter context is a separate follow-up. See `docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md`. | | `MLXCEL_METAL4_ATTENTION` | `0`/`false`/`no`/`off` to disable; unset or any other value to enable | on (where the hardware has it) | **Advanced, diagnostic kill switch.** Forces the M5 neural-accelerator fused attention route off on hardware that has it, so `layers::metal4_causal_attention` is skipped and the ordinary SDPA path runs instead. Off-switch only: it cannot turn the route on where `has_neural_accelerator && macos_supports_na` is false, so setting it on an M1 is inert. Use it to A/B the route when an M5 disagrees numerically with an earlier Apple GPU generation; before this switch the only way was to patch `should_use_metal4_attention` and rebuild, which is what the #1182 M5 investigation had to do. It does not touch f32 GEMM precision: the f32 numeric-agreement failures first filed as an M5 problem (issue #1065) were caused by `MLX_ENABLE_TF32`, reproduce on CUDA too, and did not move with this switch. Read once per process, so set it before starting `mlxcel` or `mlxcel-server`. Inert on non-Metal builds. | | `MLXCEL_NATIVE_WINDOW_SDPA` | `0`/`off` for the legacy mask route everywhere; `1`/`on` for the native route wherever it is statically eligible; `auto` for the native route only where MLX reports a fused kernel for the shape | `auto` | **Advanced, A/B selector (issue #2621).** Chooses how a causal multi-query prefill over a sliding window reaches MLX's SDPA on Metal once the returned keys outgrow the window: through upstream MLX's native `window_size` argument with `mask_mode = "causal"` (native), or through the explicit full-width windowed-causal mask the array-mask SDPA always took (legacy). The decision lives in `mlxcel_core::window_sdpa`; static eligibility is Metal, a positive window, no softcap, no attention sink, no extra mask array (padding, segments, per-head, cache validity), a query block longer than 8 rows and no longer than the keys, keys that outgrow the window, matching query/value head dims in {64, 72, 80, 96, 128, 192, 256}, and float16/bfloat16/float32. Everything else, decode and small verify blocks included, keeps its existing route whatever this variable says, and CPU/CUDA/ROCm never see a positive window (there it selects MLX's graph fallback). `auto` additionally asks MLX's own dispatch (`ScaledDotProductAttention::use_fallback`) whether the native arm runs a fused kernel; where it would run the graph fallback (on a pre-M5 GPU: head dim 256 below a fresh 2048-row half-precision prefill, every head-dim-256 continuation chunk) the legacy mask route runs unchanged. `1` forces the native entry through the graph fallback as well, which is what a kernel A/B needs. Gemma 3 and Gemma 4 build no sliding mask when the route is planned native (`RotatingKVCache::attend` carries the window). The model-level plan reads the K/V dtype the sliding cache returns (`RotatingKVCache::fetched_kv_dtype`), so a `--kv-cache-mode` whose ring returns FP16 K/V under bf16 activations (`fp16+turbo4`) plans legacy and builds the mask once per forward as before, and so does a mix of sliding caches that return different dtypes (`--turbo-boundary-v` keeps the outer layers FP16); should a layer still decline a native plan, it builds the same mask itself, which costs a mask build per layer but never changes semantics. Read once per process. Measured on an M1 Ultra in `docs/benchmark_results/native-window-sdpa-m1ultra-2026-10-26.md`. | | `MLXCEL_GDN_CHAIN_PARITY` | `0` to disable, any other value (or unset) to enable | on | **Advanced, diagnostic escape hatch.** Gates the chain-parity gated-delta Metal kernel used by Qwen 3.5 MTP's speculative verify and rollback-replay paths (issue #1165). The standard gated-delta kernel carries float32 recurrent state across a `T = K` verify block and rounds it to the storage dtype only once at the end, while the classic single-token decode chain rounds after every token; a `T = K` verify block is therefore NOT bit-identical to `K` consecutive single-token decode steps unless the state is rounded per in-block step. The chain-parity kernel (`gated_delta_step_seqpar`) does that rounding, which is what makes Qwen 3.5 MTP's temperature-0 output byte-identical to classic decode. **Setting this to `0` forfeits that exactness contract**, restoring the pre-#1165 block numerics for A/B attribution of the parity kernel's own cost and acceptance effect; do not set it to `0` in a deployment that needs byte-identical speculative output. On the non-Metal ops fallback (CUDA, CPU) the flag selects, for an unmasked scalar-gated verify block of at most 32 rows (`CHAIN_PARITY_SEQUENTIAL_MAX_T`), the sequential loop of fused single steps instead of the 64-row chunked scan (issue #1782). That loop is the same arithmetic as the classic decode chain, so it is bit-identical to `T` single-token steps, and it launches far fewer kernels: the chunked scan pays a fixed 63-step Horner loop per linear-attention layer whatever the row count (about 4400 launches and a fifth of the device time per 2-row verify round on GB10 with `qwen3.5-4b-4bit`, twice that when a partial accept replays through it). `0` restores the chunked scan there for A/B attribution. **The kernel is necessary but not sufficient**: byte-identity also requires every quantized projection to dispatch to the same MLX kernel at `M = block_size` as at `M = 1`, which is not true on every GPU generation or at every block width, so the runtime probe behind `MLXCEL_MTP_ALLOW_INEXACT` is what actually decides whether MTP engages. See `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` for the measured kernel cost (inside the dispatch-noise band). | diff --git a/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs b/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs index 4271db6c1..9894d6363 100644 --- a/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs +++ b/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs @@ -234,24 +234,28 @@ pub trait SpeculativeTarget { /// row (B > 1 path). /// /// Mirrors `rollback_partial` but takes a `&[i32]` of length `B`, one - /// accept count per row. **The batched round loop guarantees the slice - /// is uniform** (issue #2353): a batched target cache is rectangular, - /// so `round_loop_batched::clamp_uniform_acceptance` cuts every - /// round back to the shortest live row's accept before any token is - /// emitted. Zeroing a shorter row's rejected tail is not a substitute: a - /// zero key still scores `q . 0 = 0` and takes `exp(0)` weight in the - /// softmax. The implementation MUST: + /// accept count per row. **The uniform-clamp arm guarantees the slice is + /// uniform** (issue #2353): a rectangular batched cache can keep only one + /// number of verify positions per row, so + /// `round_loop_batched::clamp_uniform_acceptance` cuts every round back to + /// the shortest live row's accept before any token is emitted. Targets + /// that support per-row accepts (issue #2442) are driven through + /// [`Self::rollback_partial_batched_rows`] instead. Zeroing a shorter + /// row's rejected tail is not a substitute for either: a zero key still + /// scores `q . 0 = 0` and takes `exp(0)` weight in the softmax. The + /// implementation MUST: /// /// 1. Trim the attention-layer KV caches by /// `block_size - (accepted[0] + 1)`. /// 2. Replay each linear-attention layer's GDN state over the first /// `accepted[0] + 1` positions of the verify block. /// - /// and SHOULD fail loudly on a ragged slice rather than represent it. + /// and SHOULD fail loudly on a ragged slice rather than represent it; + /// the per-row arm never reaches this hook. /// /// For Qwen 3.5 this delegates directly to /// `Qwen35Model::rollback_speculative_cache(caches, verify_out.gdn_states, accepted, block_size)`, - /// which asserts the slice is uniform. + /// whose uniform arm handles the slice without an assertion. /// /// The default implementation delegates to [`Self::rollback_partial`] /// when `accepted.len() == 1` so single-row targets don't need a @@ -277,6 +281,94 @@ pub trait SpeculativeTarget { } } + /// Whether this target can commit a different accept count in every row + /// of the batched round loop (issue #2442), instead of the uniform clamp + /// of issue #2353. + /// + /// A target that returns `true` must also implement + /// [`Self::verify_forward_with_row_layout`], + /// [`Self::rollback_partial_batched_rows`] and + /// [`Self::batch_cache_offset`]. The round loop only takes the per-row + /// path when this is `true` and the caller has enabled it. + fn supports_per_row_accepts(&self) -> bool { + false + } + + /// Whether this instance's caches can hold the per-row layout for the + /// batch (issue #2442). + /// + /// [`Self::supports_per_row_accepts`] is the family-level answer; this is + /// the instance-level one. A quantized or paged attention cache cannot be + /// compacted row by row, so such a burst falls back to the uniform clamp + /// instead of panicking. Defaults to `false`, matching the family gate. + fn row_layout_supported(&self, caches: &[Self::Cache]) -> bool { + let _ = caches; + false + } + + /// Physical cache offset (number of positions written) shared by the + /// batch after prefill. + /// + /// The per-row round loop seeds every row's logical valid end from this + /// value before round 0; rows then diverge by their own accepts. Only + /// called when [`Self::supports_per_row_accepts`] is `true`. + fn batch_cache_offset(caches: &[Self::Cache]) -> i32 { + let _ = caches; + 0 + } + + /// Run one verify forward for a batch whose rows have different logical + /// valid ends (issue #2442). + /// + /// `row_valid_ends[r]` is row `r`'s logical valid end before this round's + /// block: the number of positions in the row's contiguous valid prefix. + /// The shared physical cache offset is `max(row_valid_ends)`; rows below + /// it carry a stale gap `[row_valid_ends[r], offset)` that the forward + /// MUST exclude from attention and rotate around (the Gemma 4 + /// `DivergentVerifyRows` layout). The block's own tokens sit at the + /// shared offset and are valid for every row. + /// + /// The default implementation ignores the layout and delegates to + /// [`Self::verify_forward_with_capture_layers`], which is correct only + /// when every row's valid end equals the physical offset. + fn verify_forward_with_row_layout( + &self, + verify_input: &MlxArray, + caches: &mut [Self::Cache], + capture_layer_ids: &[usize], + row_valid_ends: &[i32], + ) -> Self::VerifyOut { + let _ = row_valid_ends; + self.verify_forward_with_capture_layers(verify_input, caches, capture_layer_ids) + } + + /// Rewind a ragged batch: row `r` keeps its first `accepted[r] + 1` + /// block positions and every row's attention caches are compacted back + /// to a contiguous valid prefix (issue #2442). + /// + /// `row_valid_ends_pre[r]` is row `r`'s logical valid end before the + /// verify block, so the new end is `row_valid_ends_pre[r] + accepted[r] + 1`. + /// The implementation MUST move each row's accepted block from the shared + /// physical offset back to that row's logical end (so the next round sees + /// a contiguous `[0, valid_end)` prefix), zero the vacated gaps, and set + /// the physical cache offset to the maximum new valid end. Linear-attention + /// (GDN) states are replayed per row over that row's own accepted prefix. + /// + /// The default implementation delegates to + /// [`Self::rollback_partial_batched`], which is correct only for a + /// uniform accept. + fn rollback_partial_batched_rows( + &self, + caches: &mut [Self::Cache], + verify_out: &Self::VerifyOut, + accepted: &[i32], + row_valid_ends_pre: &[i32], + block_size: i32, + ) { + let _ = row_valid_ends_pre; + self.rollback_partial_batched(caches, verify_out, accepted, block_size); + } + /// Build the drafter's per-round hidden input by concatenating the /// captured per-layer hidden states along the feature axis. /// diff --git a/src/lib/mlxcel-core/src/drafter/dflash/round_loop_batched.rs b/src/lib/mlxcel-core/src/drafter/dflash/round_loop_batched.rs index aff05df64..09defd8f1 100644 --- a/src/lib/mlxcel-core/src/drafter/dflash/round_loop_batched.rs +++ b/src/lib/mlxcel-core/src/drafter/dflash/round_loop_batched.rs @@ -21,7 +21,7 @@ //! //! ## Scope //! -//! **B >= 1, with continuous batching and a uniform per-round accept.** +//! **B >= 1, with continuous batching and per-row accept commits.** //! The B = 1 fast path lives in the sibling //! [`crate::drafter::dflash::round_loop`] module (sub-12) //! and must NOT regress as a result of this implementation. The two @@ -36,18 +36,20 @@ //! of per-row emitted tokens. Rows in different speculative-walk //! states share the same verify-forward pass but diverge on what //! they accept. -//! 2. **Uniform per-round accept (issue #2353).** The target's batched -//! caches are rectangular: one rollback keeps the same number of -//! verify positions in every row. When the live rows' walks disagree, -//! `clamp_uniform_acceptance` cuts every row back to the shortest -//! live row's accept before anything is emitted, so the -//! [`SpeculativeTarget::rollback_partial_batched`] hook always -//! receives a uniform `&[i32]` slice. Zeroing a short row's rejected -//! tail instead is not enough: a zero key still scores `q . 0 = 0`, -//! takes `exp(0)` weight in the softmax and dilutes that row's -//! attention on every later step. The walk's own counts stay in -//! [`DFlashBatchedRunOutput::accept_lens`]; the committed ones are in -//! [`DFlashBatchedRunOutput::committed_accept_lens`]. +//! 2. **Per-row accept commits (issue #2442), with the uniform clamp of +//! issue #2353 behind `MLXCEL_DFLASH_PER_ROW_ACCEPTS=0`.** A target +//! that reports [`SpeculativeTarget::supports_per_row_accepts`] and +//! allows it for the current caches commits each row's own count: the +//! target compacts each row's accepted block back to that row's logical +//! end, masks the resulting stale gap from attention, rotates the block +//! at per-row positions, and replays each GDN layer's state over the +//! row's own prefix. Zeroing a short row's rejected tail is not enough: +//! a zero key still scores `q . 0 = 0`, takes `exp(0)` weight in the +//! softmax and dilutes that row's attention on every later step. The +//! clamp arm remains for A/B attribution and for quantized or paged +//! caches that cannot be compacted row by row. The walk's own counts +//! stay in [`DFlashBatchedRunOutput::accept_lens`]; the committed ones +//! are in [`DFlashBatchedRunOutput::committed_accept_lens`]. //! 3. **Per-row early-EOS / max-tokens stop.** When a row hits EOS or //! saturates `max_new_tokens`, it is frozen in place: subsequent //! verify-forwards still include that row (the batch shape stays @@ -265,16 +267,18 @@ pub struct DFlashBatchedRunOutput { /// — the caller is expected to have already streamed those out). pub tokens: Vec>, /// Per-round, per-row accept counts as the speculative walk measured - /// them, before the uniform-accept clamp. `accept_lens[r][i]` is the - /// accept count for row `r` on round `i`. + /// them, before any accept policy. `accept_lens[r][i]` is the accept + /// count for row `r` on round `i`. pub accept_lens: Vec>, - /// Per-round, per-row accept counts the round committed to the - /// target's caches, after the uniform-accept clamp. Every row of a - /// round carries the same value; a live row emitted - /// `committed_accept_lens[r][i] + 1` tokens on round `i`. + /// Per-round, per-row accept counts the round committed to the target's + /// caches. Under the uniform clamp every row of a round carries the same + /// value; under per-row accepts (issue #2442) each row carries its own + /// walk count. A live row emitted `committed_accept_lens[r][i] + 1` + /// tokens on round `i`. pub committed_accept_lens: Vec>, /// Rounds whose live rows accepted different counts and were clamped to - /// the shortest one (see `clamp_uniform_acceptance`). + /// the shortest one (see `clamp_uniform_acceptance`). Always 0 under + /// per-row accepts. pub clamped_rounds: usize, /// Wall-clock decode statistics. Prefill is the caller's /// responsibility (the loop runs after prefill). @@ -296,6 +300,32 @@ pub struct DFlashBatchedGenerator { /// [`super::round_loop::DEFAULT_MASK_TOKEN_ID`]). #[allow(dead_code)] mask_token_id: i32, + /// Commit each row's own accept count instead of clamping the round to + /// the shortest live row (issue #2442). Requires a target that reports + /// [`SpeculativeTarget::supports_per_row_accepts`]. + per_row_accepts: bool, +} + +/// Environment switch for the per-row accept port of issue #2442. +/// +/// Per-row accepts are the default: on this host they measured 1.397x the +/// uniform clamp's aggregate throughput at B = 4 (committed accepts 1.036 +/// against the clamp's 0.323 per row-round). Set `0` to restore the clamp. +pub const PER_ROW_ACCEPTS_ENV: &str = "MLXCEL_DFLASH_PER_ROW_ACCEPTS"; + +/// Default for [`DFlashBatchedGenerator::per_row_accepts`]: on unless +/// [`PER_ROW_ACCEPTS_ENV`](constant@PER_ROW_ACCEPTS_ENV) turns it off. +fn per_row_accepts_default() -> bool { + match std::env::var(PER_ROW_ACCEPTS_ENV) { + Ok(raw) => { + let value = raw.trim(); + !(value == "0" + || value.eq_ignore_ascii_case("false") + || value.eq_ignore_ascii_case("off") + || value.eq_ignore_ascii_case("no")) + } + Err(_) => true, + } } impl DFlashBatchedGenerator { @@ -312,9 +342,21 @@ impl DFlashBatchedGenerator { sampler, block_size, mask_token_id, + per_row_accepts: per_row_accepts_default(), } } + /// Override the per-row accept policy (issue #2442). + /// + /// Measurement and tests set this explicitly so both arms run in one + /// process; the production default comes from + /// [`PER_ROW_ACCEPTS_ENV`](constant@PER_ROW_ACCEPTS_ENV). + #[must_use] + pub fn with_per_row_accepts(mut self, enabled: bool) -> Self { + self.per_row_accepts = enabled; + self + } + /// Convenience constructor that uses the default block size and /// mask-token id from [`super::round_loop`]. pub fn with_drafter(drafter: Box, sampler: SamplingConfig) -> Self { @@ -447,6 +489,21 @@ impl DFlashBatchedGenerator { let mut emitted: Vec = vec![1; batch_size]; let mut finished: Vec = vec![false; batch_size]; + // Per-row accept layout (issue #2442). `row_valid_ends[r]` is row + // `r`'s logical target position count: the length of the contiguous + // valid prefix in its cache. All rows start at the prefill offset and + // then diverge by their own accepts. The uniform-clamp arm leaves this + // empty and never reads it. + let per_row = self.per_row_accepts + && target.supports_per_row_accepts() + && target.row_layout_supported(caches); + let mut row_valid_ends: Vec = if per_row { + let offset = T::batch_cache_offset(caches); + vec![offset; batch_size] + } else { + Vec::new() + }; + // Hidden tensor of shape [B, T, dim]; T grows / shrinks each // round per the captured-hidden slice contract. let mut hidden: UniquePtr = first_hidden; @@ -520,11 +577,16 @@ impl DFlashBatchedGenerator { verify_buf.extend_from_slice(&draft_tokens_per_row[r]); } let verify_input = ffi::from_slice_i32(&verify_buf, &[batch_size as i32, bs as i32]); - let verify_out = target.verify_forward_with_capture_layers( - &verify_input, - caches, - &capture_layer_ids, - ); + let verify_out = if per_row { + target.verify_forward_with_row_layout( + &verify_input, + caches, + &capture_layer_ids, + &row_valid_ends, + ) + } else { + target.verify_forward_with_capture_layers(&verify_input, caches, &capture_layer_ids) + }; // ---- Argmax sample (greedy at temp=0) of the per-row logits ---- let target_tokens_per_row = argmax_logits_per_row( @@ -549,12 +611,19 @@ impl DFlashBatchedGenerator { accept_lens_per_row[r].push(accepted_per_row[r] as u32); } - // ---- Uniform accept (issue #2353) ---- - // The rectangular caches cannot hold ragged accepts, so every - // row commits the shortest live row's accept, before any token - // is emitted. Everything below (emit, bonus, hidden slice, - // rollback) reads the clamped values. - if clamp_uniform_acceptance(&mut accepted_per_row, &mut new_tokens_per_row, &finished) { + // ---- Accept (uniform clamp or per-row, issue #2442) ---- + // The clamp arm commits the shortest live row's accept into the + // rectangular caches before any token is emitted. The per-row arm + // commits each row's own walk count; the target compacts its + // per-row cache prefixes and replays GDN state in the rollback + // below. `clamped_rounds` stays 0 on the per-row arm. + if !per_row + && clamp_uniform_acceptance( + &mut accepted_per_row, + &mut new_tokens_per_row, + &finished, + ) + { clamped_rounds += 1; } for r in 0..batch_size { @@ -611,13 +680,27 @@ impl DFlashBatchedGenerator { // hidden = hidden_full[:, : max_a + 1, :] // else: // hidden = hidden_full - // After the clamp `min_accepted == max_accepted`, so the drafter - // context for the next round matches the committed accept. + // + // The clamp arm has `min_accepted == max_accepted`, so its drafter + // context matches the committed accept. The per-row arm keeps the + // drafter's context rectangular (and identical to what the clamp + // arm would have fed it) at the shortest live row's committed + // prefix; rows that committed more re-verify their extra tokens + // through the target, which is what keeps the emitted stream exact. let min_accepted = accepted_per_row.iter().copied().min().unwrap_or(0); let max_accepted = accepted_per_row.iter().copied().max().unwrap_or(0); + let drafter_ctx = if per_row { + (0..batch_size) + .filter(|&r| !finished[r]) + .map(|r| accepted_per_row[r] + 1) + .min() + .unwrap_or(1) + } else { + max_accepted + 1 + }; let full_hidden = target.concat_hidden_for_drafter(&verify_out); - hidden = if min_accepted + 1 < bs { + hidden = if drafter_ctx < bs { let full_shape = ffi::array_shape(&full_hidden); debug_assert_eq!( full_shape.len(), @@ -627,26 +710,40 @@ impl DFlashBatchedGenerator { ffi::slice( &full_hidden, &[0, 0, 0], - &[full_shape[0], max_accepted as i32 + 1, full_shape[2]], + &[full_shape[0], drafter_ctx as i32, full_shape[2]], ) } else { full_hidden }; - // ---- Rollback (only when the round accepted less than bs - 1) ---- + // ---- Rollback ---- // - // The verify forward advanced every row's caches by `bs` - // tokens. Every row committed the same `accepted` drafter - // proposals + 1 bonus (the clamp above), so the rollback keeps - // `accepted + 1` positions in every row and trims - // `bs - (accepted + 1)`. - debug_assert_eq!( - min_accepted, max_accepted, - "the uniform-accept clamp must leave one accept count per round" - ); - if min_accepted + 1 < bs { - let accepted_i32: Vec = accepted_per_row.iter().map(|&a| a as i32).collect(); - target.rollback_partial_batched(caches, &verify_out, &accepted_i32, bs as i32); + // The verify forward advanced every row's caches by `bs` tokens. + // The clamp arm committed the same `accepted` in every row, so it + // trims `bs - (accepted + 1)` and replays one uniform prefix. The + // per-row arm hands each row's own accepted prefix to the target, + // which compacts the attention caches back to per-row contiguous + // prefixes and replays each row's GDN state over its own prefix. + let accepted_i32: Vec = accepted_per_row.iter().map(|&a| a as i32).collect(); + if per_row { + target.rollback_partial_batched_rows( + caches, + &verify_out, + &accepted_i32, + &row_valid_ends, + bs as i32, + ); + for r in 0..batch_size { + row_valid_ends[r] += accepted_i32[r] + 1; + } + } else { + debug_assert_eq!( + min_accepted, max_accepted, + "the uniform-accept clamp must leave one accept count per round" + ); + if min_accepted + 1 < bs { + target.rollback_partial_batched(caches, &verify_out, &accepted_i32, bs as i32); + } } // Periodic memory cache clear, backend-aware cadence (#627): fires when diff --git a/src/lib/mlxcel-core/src/utils.rs b/src/lib/mlxcel-core/src/utils.rs index 66c4c7f4d..bef7ef0c4 100644 --- a/src/lib/mlxcel-core/src/utils.rs +++ b/src/lib/mlxcel-core/src/utils.rs @@ -658,6 +658,66 @@ pub fn mask_stale_key_gap( } } +/// Build the additive attention mask for a ragged batched verify block. +/// +/// Row `r` has a contiguous valid key prefix `[0, per_row_valid_end[r])`, then +/// a stale gap `[per_row_valid_end[r], offset)` left by a shorter accept, then +/// this round's `l` new keys at the shared physical offset. The mask admits the +/// row's history, blocks its stale gap, and applies causality inside the new +/// block at each row's LOGICAL positions: new key `t` sits at logical position +/// `per_row_valid_end[r] + t`, so query `j` admits new keys `t <= j`. +/// +/// `left_padding[r]` additionally blocks the row's resident leading padding +/// `[0, left_padding[r])` (all zero for the equal-length DFlash burst). +/// `window` further limits every visible key to the last `window` logical +/// positions. Returns `[B, 1, l, offset + l]` f32 with 0 (attend) and +/// `-inf` (blocked) entries. +/// +/// This is the mask `rollback_speculative_cache_divergent`-style compaction +/// maintains: after each round every row's cache content is a contiguous +/// prefix, so one `(valid_end, gap_end)` pair per row describes it exactly. +/// Used by: Gemma 4 divergent batched MTP verify, Qwen 3.5 per-row DFlash +/// verify (issue #2442). +#[must_use] +pub fn build_divergent_row_mask( + l: i32, + offset: i32, + per_row_valid_end: &[i32], + left_padding: &[i32], + window: Option, +) -> UniquePtr { + assert_eq!( + per_row_valid_end.len(), + left_padding.len(), + "build_divergent_row_mask: one left_padding entry per row" + ); + let b = per_row_valid_end.len(); + let k_len = offset + l; + let mut data = vec![f32::NEG_INFINITY; b * l as usize * k_len as usize]; + for (r, (&ve, &lp)) in per_row_valid_end.iter().zip(left_padding).enumerate() { + for j in 0..l { + let q_pos = ve + j; + let row_base = (r * l as usize + j as usize) * k_len as usize; + for k in 0..k_len { + let visible = if k < lp { + false + } else if k < ve { + window.map(|w| k > q_pos - w).unwrap_or(true) + } else if k < offset { + false + } else { + let k_pos = ve + (k - offset); + k_pos <= q_pos && window.map(|w| k_pos > q_pos - w).unwrap_or(true) + }; + if visible { + data[row_base + k as usize] = 0.0; + } + } + } + } + ffi::from_slice_f32(&data, &[b as i32, 1, l, k_len]) +} + /// Create a boolean causal attention mask. /// Used by: same as `create_causal_mask` (experimental path) /// @@ -2664,6 +2724,37 @@ mod tests { } } + // --- build_divergent_row_mask: ragged verify block layout (#2442) -------- + + /// History attends, the stale gap is −∞, and the new block is causal at + /// each row's own logical positions. + #[test] + fn build_divergent_row_mask_admits_history_gap_and_logical_block() { + // l=2, offset=4. Row 0 lags by two positions (ve=2), row 1 is current. + let mask = build_divergent_row_mask(2, 4, &[2, 4], &[0, 0], None); + assert_eq!(ffi::array_shape(&mask), vec![2, 1, 2, 6]); + let cell = |b: i32, q: i32, k: i32| -> f32 { + ffi::item_f32(&ffi::slice(&mask, &[b, 0, q, k], &[b + 1, 1, q + 1, k + 1])) + }; + let attended = |b: i32, q: i32, k: i32| cell(b, q, k) == 0.0; + let blocked = |b: i32, q: i32, k: i32| { + let v = cell(b, q, k); + v.is_infinite() && v < 0.0 + }; + // Row 0: history 0,1 attends; gap 2,3 blocked; block keys 4,5 sit at + // logical 2,3, so query 0 sees only key 4 and query 1 sees 4 and 5. + assert!(attended(0, 0, 0) && attended(0, 0, 1)); + assert!(blocked(0, 0, 2) && blocked(0, 0, 3)); + assert!(attended(0, 0, 4) && blocked(0, 0, 5)); + assert!(attended(0, 1, 4) && attended(0, 1, 5)); + // Row 1 is current: its whole history attends, and the block is causal. + for k in 0..4 { + assert!(attended(1, 0, k), "row1 history k={k}"); + } + assert!(attended(1, 0, 4) && blocked(1, 0, 5)); + assert!(attended(1, 1, 4) && attended(1, 1, 5)); + } + /// A 4-D base `[B,1,n,K]` is preserved cell-for-cell outside the gap, and the /// additive penalty composes with existing −∞ base cells without NaN. #[test] diff --git a/src/models/gemma4.rs b/src/models/gemma4.rs index 112a22623..cbe69d275 100644 --- a/src/models/gemma4.rs +++ b/src/models/gemma4.rs @@ -4379,36 +4379,7 @@ pub(crate) fn build_divergent_verify_mask( left_padding: &[i32], window: Option, ) -> UniquePtr { - assert_eq!( - per_row_valid_end.len(), - left_padding.len(), - "build_divergent_verify_mask: one left_padding entry per row" - ); - let b = per_row_valid_end.len(); - let k_len = offset + l; - let mut data = vec![f32::NEG_INFINITY; b * l as usize * k_len as usize]; - for (r, (&ve, &lp)) in per_row_valid_end.iter().zip(left_padding).enumerate() { - for j in 0..l { - let q_pos = ve + j; - let row_base = (r * l as usize + j as usize) * k_len as usize; - for k in 0..k_len { - let visible = if k < lp { - false - } else if k < ve { - window.map(|w| k > q_pos - w).unwrap_or(true) - } else if k < offset { - false - } else { - let k_pos = ve + (k - offset); - k_pos <= q_pos && window.map(|w| k_pos > q_pos - w).unwrap_or(true) - }; - if visible { - data[row_base + k as usize] = 0.0; - } - } - } - } - mlxcel_core::from_slice_f32(&data, &[b as i32, 1, l, k_len]) + mlxcel_core::utils::build_divergent_row_mask(l, offset, per_row_valid_end, left_padding, window) } /// Offset of the first per-layer cache regardless of attention family. diff --git a/src/models/qwen3_5.rs b/src/models/qwen3_5.rs index ce99c7597..9219054d3 100644 --- a/src/models/qwen3_5.rs +++ b/src/models/qwen3_5.rs @@ -1244,6 +1244,7 @@ impl Qwen35DecoderLayer { /// /// used by `Qwen35Model::forward_speculative` to drive the /// DFlash drafter without duplicating the prefill / decode forward path. + #[allow(clippy::too_many_arguments)] fn forward_with_capture( &self, layer_idx: usize, @@ -1252,6 +1253,7 @@ impl Qwen35DecoderLayer { cache: &mut Qwen3NextCache, position_ids: Option<&MlxArray>, snapshot_sink: &mut Vec, + row_valid_ends: Option<&[i32]>, ) -> UniquePtr { let normed = self.input_layernorm.forward(x); @@ -1276,9 +1278,15 @@ impl Qwen35DecoderLayer { // logits stay bit-aligned with single-token decode. Mirrors // upstream `Qwen3_5DecoderLayer.__call__`'s `target_verify=True` // propagation into `Qwen3_5Attention`. - (Qwen35AttentionVariant::FullAttention(attn), Qwen3NextCache::Attention(c)) => { - attn.forward_with_position_ids_verify(&normed, c, mask, position_ids, true) - } + (Qwen35AttentionVariant::FullAttention(attn), Qwen3NextCache::Attention(c)) => attn + .forward_with_position_ids_verify( + &normed, + c, + mask, + position_ids, + true, + row_valid_ends, + ), (Qwen35AttentionVariant::FullAttention(attn), _) => { let mut temp_cache = KVCache::new(); attn.forward_with_position_ids_verify( @@ -1287,6 +1295,7 @@ impl Qwen35DecoderLayer { mask, position_ids, true, + row_valid_ends, ) } }; @@ -2496,6 +2505,33 @@ impl Qwen35Model { input_ids: &MlxArray, caches: &mut [Qwen3NextCache], capture_layer_ids: &[usize], + ) -> VerifyOutput { + self.forward_speculative_inner(input_ids, caches, capture_layer_ids, None) + } + + /// Per-row variant of [`Self::forward_speculative`] for a ragged batched + /// verify round (issue #2442). + /// + /// `row_valid_ends[r]` is row `r`'s contiguous valid cache prefix length; + /// rows below the shared physical offset rotate and attend at their own + /// logical positions and skip their stale gap. With every row at the + /// shared offset this is byte-identical to [`Self::forward_speculative`]. + pub(crate) fn forward_speculative_rows( + &self, + input_ids: &MlxArray, + caches: &mut [Qwen3NextCache], + capture_layer_ids: &[usize], + row_valid_ends: &[i32], + ) -> VerifyOutput { + self.forward_speculative_inner(input_ids, caches, capture_layer_ids, Some(row_valid_ends)) + } + + fn forward_speculative_inner( + &self, + input_ids: &MlxArray, + caches: &mut [Qwen3NextCache], + capture_layer_ids: &[usize], + row_valid_ends: Option<&[i32]>, ) -> VerifyOutput { let h0 = self.embed_tokens.forward(input_ids); @@ -2530,7 +2566,15 @@ impl Qwen35Model { } else { fa_mask.as_deref() }; - h = layer.forward_with_capture(i, &h, mask, cache, None, &mut gdn_states); + h = layer.forward_with_capture( + i, + &h, + mask, + cache, + None, + &mut gdn_states, + row_valid_ends, + ); // Capture the post-block hidden state for any requested layer index. // The check is O(capture_layer_ids.len()) per layer; for DFlash's @@ -2685,45 +2729,104 @@ impl Qwen35Model { /// * `gdn_states` — the `gdn_states` field of [`VerifyOutput`] from the /// SAME verify pass. Ordered by linear-attention layer index. /// * `accepted` — per-row accepted count `a_i` (i.e. the prefix - /// `[0..a_i]` is kept). For `B == 1` this is a single-element slice; - /// for `B > 1` every element must be equal (see below). + /// `[0..a_i]` is kept). For `B == 1` this is a single-element slice. + /// For `B > 1` the entries may be ragged (issue #2442), in which case + /// [`Self::rollback_speculative_cache_rows`] also needs each row's + /// pre-round valid end. /// * `block_size` — the verify-block length (typically the drafter's - /// `block_size` config). Used to compute the trim amount: - /// `trim = block_size - (max(accepted) + 1)`. + /// `block_size` config). Used to compute the trim amount. /// - /// Uniform acceptance for `B > 1` (issue #2353): the attention caches - /// are rectangular, with one shared offset and no per-row valid end, so - /// a batched rollback can only keep the same number of positions in - /// every row. Zeroing a shorter row's rejected tail is not a substitute: - /// a zero key still scores `q . 0 = 0`, takes `exp(0)` weight in the - /// softmax and dilutes that row's attention on every later step. The - /// batched DFlash round loop clamps each round to a uniform accept - /// before it emits anything - /// (`mlxcel_core::drafter::dflash::round_loop_batched::clamp_uniform_acceptance`), - /// and a ragged slice reaching this method panics instead of corrupting - /// the cache. + /// Uniform acceptance (issues #2353, #2442): when every row accepted the + /// same count and every row's prefix ended at the shared physical offset, + /// the rollback trims by `block_size - (accepted + 1)` and replays one + /// uniform GDN prefix. A ragged accept cannot be represented that way: a + /// zero key still scores `q . 0 = 0` and takes `exp(0)` weight in the + /// softmax. [`Self::rollback_speculative_cache_rows`] instead compacts + /// every row's accepted block back to its own logical end, so each row's + /// cache content is again a contiguous prefix, and replays each row's GDN + /// state over that row's own prefix. /// /// Apple Silicon precision (`docs/apple-silicon-precision.md`): /// GDN replay re-uses the captured `q/k/v/a/b` and `init_state` tensors /// at their original dtype (bf16/f16/float32 as captured). Do not promote /// to f32 — the drafter's reference forward stays in the activation dtype. /// - /// Used by: DFlash drafter round loop (sub-12). + /// Used by: DFlash drafter round loop (sub-12), and the B = 1 path of the + /// per-row round loop. pub fn rollback_speculative_cache( &self, caches: &mut [Qwen3NextCache], gdn_states: &[GdnRollbackSnapshot], accepted: &[i32], block_size: i32, + ) -> i32 { + self.rollback_speculative_cache_rows(caches, gdn_states, accepted, None, block_size) + } + + /// Ragged-aware rollback (issue #2442). + /// + /// `row_valid_ends_pre[r]` is row `r`'s contiguous valid prefix length + /// before the verify block. When it is absent every row is assumed to have + /// ended at the shared physical offset (the uniform and B = 1 contracts). + /// Returns `max(accepted)`. + pub(crate) fn rollback_speculative_cache_rows( + &self, + caches: &mut [Qwen3NextCache], + gdn_states: &[GdnRollbackSnapshot], + accepted: &[i32], + row_valid_ends_pre: Option<&[i32]>, + block_size: i32, ) -> i32 { if accepted.is_empty() { return 0; } - assert_uniform_batched_accepts(accepted); + let offset = caches + .iter() + .find_map(|cache| match cache { + Qwen3NextCache::Attention(kv) => Some(kv.offset), + Qwen3NextCache::Linear(_) => None, + }) + .or_else(|| caches.first().map(Qwen3NextCache::offset)) + .unwrap_or(0); + let o_pre = offset - block_size; + let ve_pre: Vec = match row_valid_ends_pre { + Some(ends) => { + assert_eq!( + ends.len(), + accepted.len(), + "rollback_speculative_cache_rows: one valid end per accepted row" + ); + ends.to_vec() + } + None => vec![o_pre; accepted.len()], + }; + let uniform = + accepted.iter().all(|&a| a == accepted[0]) && ve_pre.iter().all(|&ve| ve == o_pre); + let max_a = accepted[0]; let n = max_a + 1; let trim = block_size - n; + if !uniform { + let o_post = ve_pre + .iter() + .zip(accepted) + .map(|(&ve, &a)| ve + a + 1) + .max() + .expect("non-empty accepted"); + for cache in caches.iter_mut() { + if let Qwen3NextCache::Attention(kv) = cache { + compact_attention_rows(kv, &ve_pre, accepted, o_pre, o_post); + let roll_back = o_pre + block_size - o_post; + if roll_back > 0 { + kv.trim(roll_back); + } + } + } + self.replay_gdn_rows(caches, gdn_states, accepted, o_post); + return accepted.iter().copied().max().unwrap_or(0); + } + // Attention caches: trim by `trim`, the same in every row. for cache in caches.iter_mut() { if let Qwen3NextCache::Attention(kv) = cache @@ -2733,8 +2836,6 @@ impl Qwen35Model { } } - // GDN caches: replay the captured block over `[:, :n]` so the - // post-block state matches the accepted-prefix position. let mut snapshot_iter = gdn_states.iter(); for (layer_idx, cache) in caches.iter_mut().enumerate() { let Qwen3NextCache::Linear(linear_cache) = cache else { @@ -2825,14 +2926,16 @@ impl Qwen35Model { // Recover conv_state from the captured conv_input: // cache = conv_input[:, a+1 : a+K] (K = conv_kernel_size), one - // window for every row because the accept is uniform. + // window for every row because the accept is uniform. The window + // length is `kernel - 1`, the same contract + // `forward_hidden_internal_with_capture` maintains. let k_kernel = layer.conv_kernel_size as i32; let conv_shape = mlxcel_core::array_shape(&snap.conv_input); let conv_dim = conv_shape[2]; let window = mlxcel_core::slice( &snap.conv_input, - &[0, max_a + 1, 0], - &[conv_shape[0], max_a + k_kernel, conv_dim], + &[0, n, 0], + &[conv_shape[0], n + k_kernel - 1, conv_dim], ); linear_cache.conv_state = Some(mlxcel_core::contiguous(&window, false)); @@ -2846,6 +2949,107 @@ impl Qwen35Model { max_a } + /// Replay every linear-attention layer's GDN state over each row's own + /// accepted prefix (issue #2442). + /// + /// The conv and recurrent states are per-row already, but their valid + /// lengths are not: one batched replay over `max(accepted)` would advance + /// a short row through positions it rejected. Each row is therefore + /// sliced and replayed with the chain-parity kernel at its own length, and + /// the results are stacked back into the `[B, ...]` cache tensors. + fn replay_gdn_rows( + &self, + caches: &mut [Qwen3NextCache], + gdn_states: &[GdnRollbackSnapshot], + accepted: &[i32], + o_post: i32, + ) { + let mut snapshot_iter = gdn_states.iter(); + for (layer_idx, cache) in caches.iter_mut().enumerate() { + let Qwen3NextCache::Linear(linear_cache) = cache else { + continue; + }; + let snap = match snapshot_iter.next() { + Some(s) => s, + None => continue, + }; + debug_assert_eq!(snap.layer_idx, layer_idx); + let layer = match &self.layers[layer_idx].attention { + Qwen35AttentionVariant::Linear(linear) => linear, + _ => continue, + }; + let k_kernel = layer.conv_kernel_size as i32; + let q_shape = mlxcel_core::array_shape(&snap.q); + let v_shape = mlxcel_core::array_shape(&snap.v); + let a_shape = mlxcel_core::array_shape(&snap.a); + let conv_shape = mlxcel_core::array_shape(&snap.conv_input); + let conv_dim = conv_shape[2]; + + let mut state_acc: Option> = None; + let mut conv_acc: Option> = None; + for (r, &a) in accepted.iter().enumerate() { + let n = a + 1; + let bi = r as i32; + let q_n = mlxcel_core::slice( + &snap.q, + &[bi, 0, 0, 0], + &[bi + 1, n, q_shape[2], q_shape[3]], + ); + let k_n = mlxcel_core::slice( + &snap.k, + &[bi, 0, 0, 0], + &[bi + 1, n, q_shape[2], q_shape[3]], + ); + let v_n = mlxcel_core::slice( + &snap.v, + &[bi, 0, 0, 0], + &[bi + 1, n, v_shape[2], v_shape[3]], + ); + let a_n = mlxcel_core::slice(&snap.a, &[bi, 0, 0], &[bi + 1, n, a_shape[2]]); + let b_n = mlxcel_core::slice(&snap.b, &[bi, 0, 0], &[bi + 1, n, a_shape[2]]); + let init_row = snap.init_state.as_ref().map(|state| { + let s = mlxcel_core::array_shape(state); + mlxcel_core::copy(&mlxcel_core::slice( + state, + &[bi, 0, 0, 0], + &[bi + 1, s[1], s[2], s[3]], + )) + }); + let (_y, state_row) = gated_delta_update_chain_parity( + &q_n, + &k_n, + &v_n, + &a_n, + &b_n, + &layer.a_log, + &layer.dt_bias, + init_row.as_deref(), + None, + ); + state_acc = Some(match state_acc { + None => state_row, + Some(prev) => mlxcel_core::concatenate(&prev, &state_row, 0), + }); + + let conv_row = mlxcel_core::contiguous( + &mlxcel_core::slice( + &snap.conv_input, + &[bi, n, 0], + &[bi + 1, n + k_kernel - 1, conv_dim], + ), + false, + ); + conv_acc = Some(match conv_acc { + None => conv_row, + Some(prev) => mlxcel_core::concatenate(&prev, &conv_row, 0), + }); + } + linear_cache.state_cache = state_acc; + linear_cache.conv_state = conv_acc; + linear_cache.offset = o_post; + } + } + // === Sequence-routed speculative hooks (MTP target adapter) ============ // // The MtpTarget trait is `&self`-only with no cache parameter: the target @@ -3120,10 +3324,8 @@ impl mlxcel_core::drafter::dflash::SpeculativeTarget for Qwen35Model { accepted: &[i32], block_size: i32, ) { - // For B > 1 the rollback path uses the same `rollback_speculative_cache` - // entrypoint with a multi-element `accepted` slice, which the batched - // round loop guarantees is uniform (issue #2353); the method asserts - // it. + // Uniform batched rollback (issue #2353). The per-row arm (issue + // #2442) comes through `rollback_partial_batched_rows` below instead. // // Apple Silicon precision (`docs/apple-silicon-precision.md`): // the GDN replay tensors keep their captured dtype (bf16/f16 @@ -3132,6 +3334,68 @@ impl mlxcel_core::drafter::dflash::SpeculativeTarget for Qwen35Model { self.rollback_speculative_cache(caches, &verify_out.gdn_states, accepted, block_size); } + /// Qwen 3.5 can commit per-row accepts (issue #2442): its attention KV is + /// compacted per row and its GDN state replayed per row. + fn supports_per_row_accepts(&self) -> bool { + true + } + + /// The per-row rollback compacts dense Fp16 attention caches. A quantized + /// or paged KV mode keeps its own sidecars and head trim, so such a burst + /// stays on the uniform clamp rather than corrupting the cache. + fn row_layout_supported(&self, caches: &[Self::Cache]) -> bool { + caches.iter().all(|cache| match cache { + Qwen3NextCache::Attention(kv) => { + kv.mode == KVCacheMode::Fp16 + && !kv.is_paged_backed() + && kv.offset - kv.live_len() == 0 + } + Qwen3NextCache::Linear(_) => true, + }) + } + + fn batch_cache_offset(caches: &[Self::Cache]) -> i32 { + caches + .iter() + .find_map(|cache| match cache { + Qwen3NextCache::Attention(kv) => Some(kv.offset), + Qwen3NextCache::Linear(_) => None, + }) + .unwrap_or(0) + } + + fn verify_forward_with_row_layout( + &self, + verify_input: &MlxArray, + caches: &mut [Self::Cache], + capture_layer_ids: &[usize], + row_valid_ends: &[i32], + ) -> Self::VerifyOut { + let capture_layer_ids = if capture_layer_ids.is_empty() { + mlxcel_core::drafter::dflash::config::DEFAULT_TARGET_LAYER_IDS + } else { + capture_layer_ids + }; + self.forward_speculative_rows(verify_input, caches, capture_layer_ids, row_valid_ends) + } + + fn rollback_partial_batched_rows( + &self, + caches: &mut [Self::Cache], + verify_out: &Self::VerifyOut, + accepted: &[i32], + row_valid_ends_pre: &[i32], + block_size: i32, + ) { + let _ = self.rollback_speculative_cache_rows( + caches, + &verify_out.gdn_states, + accepted, + Some(row_valid_ends_pre), + block_size, + ); + } + fn concat_hidden_for_drafter(&self, verify_out: &Self::VerifyOut) -> UniquePtr { // Upstream: `hidden = mx.concatenate(verify_out.hidden_states, axis=-1)`. // The captured hidden states each have shape `[1, bs, hidden_size]`; @@ -3158,24 +3422,113 @@ impl mlxcel_core::drafter::dflash::SpeculativeTarget for Qwen35Model { } } -/// Panic unless every row of a batched speculative rollback accepted the same -/// count (issue #2353). +/// Move each row's accepted verify block back to that row's own logical valid +/// end, restoring a contiguous `[0, valid_end)` prefix in a shared, +/// rectangular attention cache (issue #2442). +/// +/// The block arrived at the shared physical offset `o_pre`; row `r` accepted +/// `accepted[r] + 1` of its positions, which move to +/// `[row_valid_ends_pre[r], row_valid_ends_pre[r] + accepted[r] + 1)`. The +/// vacated gap up to `o_post` (the maximum new valid end) is zeroed so a +/// maskless fallback cannot read stale keys. The caller trims the cache +/// offset to `o_post` afterwards. /// -/// The Qwen 3.5 attention caches are rectangular (one shared offset, no -/// per-row valid end), so a ragged accept could only be represented by -/// zeroing the shorter rows' rejected tails, and a zero key is still -/// attended. The batched DFlash round loop clamps every round to a uniform -/// accept; reaching here with a ragged slice is a round-loop bug, and failing -/// loudly beats emitting fluent, wrong text. +/// The DFlash batched path only ever builds plain Fp16 caches, so there are +/// no quantization sidecars to move; a cache that is paged or head-trimmed is +/// a caller bug on this path and panics rather than silently corrupting. /// -/// Used by: [`Qwen35Model::rollback_speculative_cache`]. -pub(crate) fn assert_uniform_batched_accepts(accepted: &[i32]) { +/// Used by: [`Qwen35Model::rollback_speculative_cache_rows`]. +pub(crate) fn compact_attention_rows( + kv: &mut KVCache, + row_valid_ends_pre: &[i32], + accepted: &[i32], + o_pre: i32, + o_post: i32, +) { + assert_eq!( + row_valid_ends_pre.len(), + accepted.len(), + "compact_attention_rows: one valid end per accepted row" + ); + assert_eq!( + kv.mode, + KVCacheMode::Fp16, + "per-row DFlash rollback requires a plain Fp16 attention cache; got {:?}", + kv.mode + ); assert!( - accepted.iter().all(|&a| a == accepted[0]), - "Qwen 3.5 batched speculative rollback requires uniform per-row acceptance: \ - the attention caches are rectangular and cannot exclude a shorter row's \ - rejected tail; got ragged accepts {accepted:?}" + !kv.is_paged_backed(), + "per-row DFlash rollback requires a dense attention cache, not a paged one" ); + assert_eq!( + kv.offset - kv.live_len(), + 0, + "per-row DFlash rollback requires an uncompacted (non head-trimmed) cache" + ); + let (Some(keys), Some(values)) = (kv.keys.as_ref(), kv.values.as_ref()) else { + return; + }; + let k_shape = mlxcel_core::array_shape(keys); + let v_shape = mlxcel_core::array_shape(values); + if k_shape[0] as usize != accepted.len() { + return; + } + let k_dtype = mlxcel_core::array_dtype(keys); + let v_dtype = mlxcel_core::array_dtype(values); + + let mut new_keys = mlxcel_core::copy(keys); + let mut new_values = mlxcel_core::copy(values); + for (r, (&ve, &a)) in row_valid_ends_pre.iter().zip(accepted).enumerate() { + let n = a + 1; + let bi = r as i32; + if ve < o_pre { + // Materialize the source before the update: `slice_update` may + // donate the buffer, and an overlapping move would then read + // already-overwritten rows. + let src_k = mlxcel_core::copy(&mlxcel_core::slice( + &new_keys, + &[bi, 0, o_pre, 0], + &[bi + 1, k_shape[1], o_pre + n, k_shape[3]], + )); + new_keys = mlxcel_core::slice_update( + &new_keys, + &src_k, + &[bi, 0, ve, 0], + &[bi + 1, k_shape[1], ve + n, k_shape[3]], + ); + let src_v = mlxcel_core::copy(&mlxcel_core::slice( + &new_values, + &[bi, 0, o_pre, 0], + &[bi + 1, v_shape[1], o_pre + n, v_shape[3]], + )); + new_values = mlxcel_core::slice_update( + &new_values, + &src_v, + &[bi, 0, ve, 0], + &[bi + 1, v_shape[1], ve + n, v_shape[3]], + ); + } + let z_start = ve + n; + if z_start < o_post { + let span = o_post - z_start; + let k_zero = mlxcel_core::zeros(&[1, k_shape[1], span, k_shape[3]], k_dtype); + let v_zero = mlxcel_core::zeros(&[1, v_shape[1], span, v_shape[3]], v_dtype); + new_keys = mlxcel_core::slice_update( + &new_keys, + &k_zero, + &[bi, 0, z_start, 0], + &[bi + 1, k_shape[1], o_post, k_shape[3]], + ); + new_values = mlxcel_core::slice_update( + &new_values, + &v_zero, + &[bi, 0, z_start, 0], + &[bi + 1, v_shape[1], o_post, v_shape[3]], + ); + } + } + kv.keys = Some(new_keys); + kv.values = Some(new_values); } pub(crate) struct Qwen35StageModel { diff --git a/src/models/qwen3_5_tests.rs b/src/models/qwen3_5_tests.rs index d1895649c..cf8624cbe 100644 --- a/src/models/qwen3_5_tests.rs +++ b/src/models/qwen3_5_tests.rs @@ -14,17 +14,18 @@ //! Unit tests for the Qwen 3.5 DFlash speculative-decoding hooks. //! -//! These tests cover the *standalone* helpers (the uniform-accept guard of the -//! batched rollback, and `sanitize_weights`'s MTP-stripping path) without -//! requiring a real model checkpoint, plus a one-layer synthetic model for the -//! batched `rollback_speculative_cache` contract. The end-to-end real-checkpoint -//! `forward_speculative` + `rollback_speculative_cache` round trip is exercised -//! by integration tests gated behind hardware availability. +//! These tests cover the *standalone* helpers (the per-row attention-cache +//! compaction of the ragged batched rollback, and `sanitize_weights`'s +//! MTP-stripping path) without requiring a real model checkpoint, plus a +//! one-layer synthetic model for the batched `rollback_speculative_cache` +//! contract. The end-to-end real-checkpoint `forward_speculative` + +//! `rollback_speculative_cache` round trip is exercised by integration tests +//! gated behind hardware availability. use super::qwen3_5::{ MRopePositionSource, Qwen35Config, Qwen35Model, Qwen35UnsupportedConfig, - assert_uniform_batched_accepts, eos_token_ids_from_config, mrope_position_source, - sanitize_weights, validate_qwen35_wrapper_config, + compact_attention_rows, eos_token_ids_from_config, mrope_position_source, sanitize_weights, + validate_qwen35_wrapper_config, }; use super::qwen3_next::Qwen3NextCache; use mlxcel_core::cache::SequenceId; @@ -77,23 +78,87 @@ fn make_tiny_config() -> Qwen35Config { } // --------------------------------------------------------------------------- -// Batched speculative rollback requires a uniform accept (issue #2353). The -// rectangular attention caches cannot exclude a shorter row's rejected tail, -// so the batched DFlash round loop clamps every round, and a ragged slice that -// still reaches the rollback panics. +// Ragged batched rollback (issue #2442). The uniform clamp of #2353 is no +// longer the only accept policy: the per-row rollback compacts each row's +// accepted verify block back to that row's own logical end, so the rectangular +// attention cache again holds a contiguous prefix per row. // --------------------------------------------------------------------------- +/// One `[B=2, H=1, bs=4, D=2]` attention cache whose row `r` holds the value +/// `10 * r + position`, so a moved block is identifiable per row. +fn synthetic_attention_cache() -> mlxcel_core::cache::KVCache { + use mlxcel_core::cache::KVCache; + let mut cache = KVCache::new(); + // Two prefill positions, then a 4-token verify block. + let mut data = Vec::new(); + for r in 0..2 { + for pos in 0..6 { + for d in 0..2 { + data.push((10 * r + pos) as f32 * 10.0 + d as f32); + } + } + } + let all = mlxcel_core::from_slice_f32(&data, &[2, 1, 6, 2]); + let prefill = mlxcel_core::slice(&all, &[0, 0, 0, 0], &[2, 1, 2, 2]); + let block = mlxcel_core::slice(&all, &[0, 0, 2, 0], &[2, 1, 6, 2]); + let prefill_v = mlxcel_core::copy(&prefill); + let _ = cache.update_and_fetch(prefill, prefill_v); + let block_v = mlxcel_core::copy(&block); + let _ = cache.update_and_fetch(block, block_v); + cache +} + #[test] -fn uniform_batched_accepts_pass_the_guard() { - assert_uniform_batched_accepts(&[]); - assert_uniform_batched_accepts(&[5]); - assert_uniform_batched_accepts(&[2, 2, 2, 2]); +fn compact_attention_rows_moves_each_row_block_to_its_own_valid_end() { + let mut cache = synthetic_attention_cache(); + // Before: physical offset 6, block at [2, 6). + // Row 0 accepted 3 of the 4 block tokens; row 1 accepted all 4. + // Both rows' pre-round valid end was the prefill offset 2. + compact_attention_rows(&mut cache, &[2, 2], &[2, 3], 2, 6); + let keys = cache.keys.as_ref().expect("keys"); + let shape = mlxcel_core::array_shape(keys); + // The Fp16 buffer is allocated in 256-slot steps, so read each row by its + // own buffer stride. + let stride = (shape[1] * shape[2] * shape[3]) as usize; + let bytes = mlxcel_core::array_to_raw_bytes(keys); + let vals: Vec = bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + // Row 0: positions 0..5 = prefill 0,10 then block 20,30,40; the stale + // tail (position 5) is zeroed. + let row0: Vec = vals[0..12].iter().step_by(2).copied().collect(); + assert_eq!(row0, vec![0.0, 10.0, 20.0, 30.0, 40.0, 0.0]); + // Row 1: positions 0..6 = prefill 100,110 then block 120,130,140,150. + let row1: Vec = vals[stride..stride + 12] + .iter() + .step_by(2) + .copied() + .collect(); + assert_eq!(row1, vec![100.0, 110.0, 120.0, 130.0, 140.0, 150.0]); } #[test] -#[should_panic(expected = "requires uniform per-row acceptance")] -fn ragged_batched_accepts_fail_the_guard() { - assert_uniform_batched_accepts(&[3, 1, 2]); +fn compact_attention_rows_is_a_no_op_when_every_row_ends_at_the_offset() { + let mut cache = synthetic_attention_cache(); + compact_attention_rows(&mut cache, &[2, 2], &[3, 3], 2, 6); + let keys = cache.keys.as_ref().expect("keys"); + let shape = mlxcel_core::array_shape(keys); + let stride = (shape[1] * shape[2] * shape[3]) as usize; + let bytes = mlxcel_core::array_to_raw_bytes(keys); + let vals: Vec = bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + // Every position keeps its original value: no row lagged the offset. + let row0: Vec = vals[0..12].iter().step_by(2).copied().collect(); + let row1: Vec = vals[stride..stride + 12] + .iter() + .step_by(2) + .copied() + .collect(); + assert_eq!(row0, vec![0.0, 10.0, 20.0, 30.0, 40.0, 50.0]); + assert_eq!(row1, vec![100.0, 110.0, 120.0, 130.0, 140.0, 150.0]); } /// A gated-delta layer followed by a full-attention layer, the smallest shape diff --git a/src/models/qwen3_next.rs b/src/models/qwen3_next.rs index 94333b27e..802028999 100644 --- a/src/models/qwen3_next.rs +++ b/src/models/qwen3_next.rs @@ -620,7 +620,7 @@ impl Qwen3NextAttention { mask: Option<&MlxArray>, position_ids: Option<&MlxArray>, ) -> UniquePtr { - self.forward_with_position_ids_verify(x, cache, mask, position_ids, false) + self.forward_with_position_ids_verify(x, cache, mask, position_ids, false, None) } /// Variant of [`Self::forward_with_position_ids`] that opts into the @@ -636,8 +636,13 @@ impl Qwen3NextAttention { /// the batched-SDPA-vs-decode logit drift that flips speculative /// accept/reject decisions away from the drafter-less greedy pass. /// + /// `row_valid_ends` carries the per-row logical valid end for a ragged + /// batched verify round (issue #2442): rows below the shared physical + /// cache offset keep only a shorter contiguous prefix, and the block's + /// queries rotate and attend at each row's own logical positions. + /// /// Used by: `Qwen35DecoderLayer::forward_with_capture` (the Qwen 3.5 MTP - /// verify pass). + /// verify pass), and the per-row DFlash verify round. pub(crate) fn forward_with_position_ids_verify( &self, x: &MlxArray, @@ -645,6 +650,7 @@ impl Qwen3NextAttention { mask: Option<&MlxArray>, position_ids: Option<&MlxArray>, target_verify: bool, + row_valid_ends: Option<&[i32]>, ) -> UniquePtr { let output = self.forward_hidden_with_position_ids_verify( x, @@ -652,6 +658,7 @@ impl Qwen3NextAttention { mask, position_ids, target_verify, + row_valid_ends, ); let out = self.o_proj.forward(&output); #[cfg(test)] @@ -666,7 +673,7 @@ impl Qwen3NextAttention { mask: Option<&MlxArray>, position_ids: Option<&MlxArray>, ) -> UniquePtr { - self.forward_hidden_with_position_ids_verify(x, cache, mask, position_ids, false) + self.forward_hidden_with_position_ids_verify(x, cache, mask, position_ids, false, None) } /// Native multi-sequence decode for the tensor-parallel Qwen 3.5 runtime. @@ -745,6 +752,7 @@ impl Qwen3NextAttention { mask: Option<&MlxArray>, position_ids: Option<&MlxArray>, target_verify: bool, + row_valid_ends: Option<&[i32]>, ) -> UniquePtr { let shape = mlxcel_core::array_shape(x); let b = shape[0]; @@ -787,8 +795,23 @@ impl Qwen3NextAttention { let offset = cache.offset; + // Per-row (ragged) batched verify round (issue #2442): rows whose + // logical valid end lags the shared physical cache offset keep only a + // shorter contiguous prefix, so the block's queries and keys rotate at + // each row's own positions instead of the shared offset. + let divergent_rows: Option<&[i32]> = row_valid_ends.filter(|ends| { + target_verify + && l > 1 + && ends.len() == b as usize + && ends.iter().any(|&end| end != offset) + }); + // Apply RoPE (standard or MRoPE with position_ids) if let (Some(mrope), Some(pos_ids)) = (&self.mrope, position_ids) { + assert!( + divergent_rows.is_none(), + "per-row DFlash verify is not supported with MRoPE position ids" + ); // Interleaved MRoPE for VLM let (cos, sin) = mrope.forward(pos_ids); let embed_dtype = mlxcel_core::array_dtype(&queries); @@ -814,6 +837,12 @@ impl Qwen3NextAttention { super::qwen3_vl::apply_multimodal_rotary_pos_emb(&q_rot, &k_rot, &cos, &sin); queries = mlxcel_core::concatenate(&q_embed, &q_pass, -1); keys = mlxcel_core::concatenate(&k_embed, &k_pass, -1); + } else if let Some(ends) = divergent_rows { + // Each row rotates its window at its own logical positions, so the + // relative RoPE distance to that row's valid history stays what a + // B = 1 chain would see. + queries = self.apply_rope_per_row(&queries, ends); + keys = self.apply_rope_per_row(&keys, ends); } else if target_verify && l > 1 && b == 1 && self.verify_rope_rows { // Verify block on CUDA: rotate each row through the same `L = 1` // call classic decode makes for that token, so the block's @@ -864,7 +893,9 @@ impl Qwen3NextAttention { // decode. This mirrors upstream `Qwen3_5Attention.__call__`'s // `target_verify and L > 1` branch and keeps the verify logits from // drifting away from the drafter-less greedy decode. - let attn_out = if target_verify && l > 1 { + let attn_out = if let Some(ends) = divergent_rows { + self.attend_per_position_rows(&queries, &cache_k, &cache_v, offset, ends) + } else if target_verify && l > 1 { self.attend_per_position(&queries, &cache_k, &cache_v) } else if l > 1 && mask.is_none() { mlxcel_core::causal_attention(&queries, &cache_k, &cache_v, self.scale, 0.0, 0) @@ -941,6 +972,104 @@ impl Qwen3NextAttention { out.expect("attend_per_position requires l_q >= 1") } + /// Rotate `[B, H, L, D]` query or key rows at per-row start offsets. + /// + /// Used by the per-row DFlash verify round (issue #2442): after a ragged + /// accept, row `r`'s window starts at its own logical valid end + /// `row_valid_ends[r]`, not at the shared physical cache offset. + fn apply_rope_per_row(&self, x: &MlxArray, row_valid_ends: &[i32]) -> UniquePtr { + let shape = mlxcel_core::array_shape(x); + debug_assert_eq!(shape.len(), 4, "rope input must be [B, H, L, D]"); + let b = shape[0] as usize; + debug_assert_eq!( + b, + row_valid_ends.len(), + "apply_rope_per_row: one offset per batch row" + ); + if b == 1 { + return mlxcel_core::fast_rope( + x, + self.rope_dims, + false, + self.rope_base, + 1.0, + row_valid_ends[0], + ); + } + let mut out: Option> = None; + for (r, &offset) in row_valid_ends.iter().enumerate() { + let row = mlxcel_core::slice( + x, + &[r as i32, 0, 0, 0], + &[r as i32 + 1, shape[1], shape[2], shape[3]], + ); + let rotated = + mlxcel_core::fast_rope(&row, self.rope_dims, false, self.rope_base, 1.0, offset); + out = Some(match out { + None => rotated, + Some(prev) => mlxcel_core::concatenate(&prev, &rotated, 0), + }); + } + out.expect("apply_rope_per_row requires b >= 1") + } + + /// Per-query-position causal attention for a ragged batched verify block + /// (issue #2442). + /// + /// Row `r` keeps a contiguous valid history `[0, row_valid_ends[r])`, a + /// stale gap up to the shared physical `offset`, and this block's keys at + /// `[offset, offset + j]` for query `j`. One additive mask + /// (`build_divergent_row_mask`) encodes all three, and query `j` slices the + /// exact key width it would see in a B = 1 chain. Both blocks of + /// [0, offset + j + 1) that a row must not see are `-inf`. + fn attend_per_position_rows( + &self, + queries: &MlxArray, + keys: &MlxArray, + values: &MlxArray, + offset: i32, + row_valid_ends: &[i32], + ) -> UniquePtr { + let q_shape = mlxcel_core::array_shape(queries); + let k_shape = mlxcel_core::array_shape(keys); + let v_shape = mlxcel_core::array_shape(values); + let b = q_shape[0]; + let n_q_heads = q_shape[1]; + let l_q = q_shape[2]; + let head_dim = q_shape[3]; + let n_kv_heads = k_shape[1]; + debug_assert_eq!( + row_valid_ends.len(), + b as usize, + "attend_per_position_rows: one valid end per batch row" + ); + let no_padding = vec![0_i32; b as usize]; + let mask = mlxcel_core::utils::build_divergent_row_mask( + l_q, + offset, + row_valid_ends, + &no_padding, + None, + ); + + let mut out: Option> = None; + for j in 0..l_q { + let q_j = mlxcel_core::slice(queries, &[0, 0, j, 0], &[b, n_q_heads, j + 1, head_dim]); + let kv_len = offset + j + 1; + let k_j = mlxcel_core::slice(keys, &[0, 0, 0, 0], &[b, n_kv_heads, kv_len, head_dim]); + let v_j = + mlxcel_core::slice(values, &[0, 0, 0, 0], &[b, n_kv_heads, kv_len, v_shape[3]]); + let mask_j = mlxcel_core::slice(&mask, &[0, 0, j, 0], &[b, 1, j + 1, kv_len]); + let attn_j = + mlxcel_core::layers::attention(&q_j, &k_j, &v_j, self.scale, Some(&mask_j), 0.0, 0); + out = Some(match out { + None => attn_j, + Some(prev) => mlxcel_core::concatenate(&prev, &attn_j, 2), + }); + } + out.expect("attend_per_position_rows requires l_q >= 1") + } + pub(crate) fn from_weights( weights: &WeightMap, config: &Qwen3NextConfig, diff --git a/src/server/batch/dflash_target.rs b/src/server/batch/dflash_target.rs index a36796656..7370cd95a 100644 --- a/src/server/batch/dflash_target.rs +++ b/src/server/batch/dflash_target.rs @@ -908,7 +908,8 @@ where clamped_rounds = output.clamped_rounds, accept_lens = ?output.accept_lens, committed_accept_lens = ?output.committed_accept_lens, - "DFlash batched burst accept lengths (walk vs committed after the uniform-accept clamp)" + "DFlash batched burst accept lengths (walk vs committed under the active accept policy; \ + clamped_rounds is non-zero only on the uniform-clamp arm)" ); let mut rows: BatchedBurstTokens = Vec::with_capacity(batch_size); for (r, row_tokens) in output.tokens.into_iter().enumerate() { diff --git a/src/vision/qwen3_5_vl.rs b/src/vision/qwen3_5_vl.rs index e350de5d3..020aa9b02 100644 --- a/src/vision/qwen3_5_vl.rs +++ b/src/vision/qwen3_5_vl.rs @@ -238,6 +238,62 @@ impl mlxcel_core::drafter::dflash::SpeculativeTarget for Qwen35VLModel { ); } + fn supports_per_row_accepts(&self) -> bool { + ::supports_per_row_accepts( + &self.text_model, + ) + } + + fn row_layout_supported(&self, caches: &[Self::Cache]) -> bool { + ::row_layout_supported( + &self.text_model, + caches, + ) + } + + fn batch_cache_offset(caches: &[Self::Cache]) -> i32 { + ::batch_cache_offset( + caches, + ) + } + + fn verify_forward_with_row_layout( + &self, + verify_input: &MlxArray, + caches: &mut [Self::Cache], + capture_layer_ids: &[usize], + row_valid_ends: &[i32], + ) -> Self::VerifyOut { + let capture_layer_ids = if capture_layer_ids.is_empty() { + mlxcel_core::drafter::dflash::config::DEFAULT_TARGET_LAYER_IDS + } else { + capture_layer_ids + }; + self.text_model.forward_speculative_rows( + verify_input, + caches, + capture_layer_ids, + row_valid_ends, + ) + } + + fn rollback_partial_batched_rows( + &self, + caches: &mut [Self::Cache], + verify_out: &Self::VerifyOut, + accepted: &[i32], + row_valid_ends_pre: &[i32], + block_size: i32, + ) { + let _ = self.text_model.rollback_speculative_cache_rows( + caches, + &verify_out.gdn_states, + accepted, + Some(row_valid_ends_pre), + block_size, + ); + } + fn concat_hidden_for_drafter(&self, verify_out: &Self::VerifyOut) -> UniquePtr { debug_assert!( !verify_out.hidden_states.is_empty(), diff --git a/tests/speculative_parity.rs b/tests/speculative_parity.rs index 28074d1cf..cdfcbec51 100644 --- a/tests/speculative_parity.rs +++ b/tests/speculative_parity.rs @@ -118,6 +118,9 @@ mod qwen35_dflash_harness; #[path = "speculative_parity/qwen35_dflash_batched.rs"] mod qwen35_dflash_batched; +#[path = "speculative_parity/qwen35_dflash_throughput.rs"] +mod qwen35_dflash_throughput; + /// Fixed prompt submitted to both the speculative and the drafter-less /// server in the byte-equality phase. Deterministic content, no system /// prompt, so the only variable between the two runs is whether the diff --git a/tests/speculative_parity/qwen35_dflash_harness.rs b/tests/speculative_parity/qwen35_dflash_harness.rs index dc399fbfe..34699ccd8 100644 --- a/tests/speculative_parity/qwen35_dflash_harness.rs +++ b/tests/speculative_parity/qwen35_dflash_harness.rs @@ -97,6 +97,20 @@ pub(crate) struct BurstRun { pub(crate) committed_accept_lens: Vec>, /// Rounds whose rows disagreed and were clamped to a uniform accept. pub(crate) clamped_rounds: usize, + /// Round-loop decode wall clock in milliseconds (the prefill is excluded). + pub(crate) decode_ms: f64, +} + +impl BurstRun { + /// Aggregate decode throughput over the round loop, tokens per second. + pub(crate) fn aggregate_tok_per_sec(&self) -> f64 { + let tokens: usize = self.tokens.iter().map(|row| row.len() - 1).sum(); + if self.decode_ms > 0.0 { + tokens as f64 / (self.decode_ms / 1000.0) + } else { + 0.0 + } + } } /// One DFlash burst over `prompts` (all the same length), replicating the @@ -108,6 +122,21 @@ pub(crate) fn run_burst( drafter: Box, prompts: &[Vec], block_size: u32, +) -> (BurstRun, Box) { + run_burst_with_policy(model, drafter, prompts, block_size, None) +} + +/// [`run_burst`] with an explicit per-row accept policy. +/// +/// `per_row` is `None` to use the production default +/// (`MLXCEL_DFLASH_PER_ROW_ACCEPTS`), `Some(true)` for the per-row arm and +/// `Some(false)` for the uniform-clamp arm. Issue #2442 compares the two. +pub(crate) fn run_burst_with_policy( + model: &mlxcel::models::Qwen35Model, + drafter: Box, + prompts: &[Vec], + block_size: u32, + per_row: Option, ) -> (BurstRun, Box) { use mlxcel_core::drafter::dflash::{DFlashBatchedGenerator, SpeculativeTarget}; use mlxcel_core::generate::{LanguageModel, SamplingConfig}; @@ -158,6 +187,9 @@ pub(crate) fn run_burst( block_size, mlxcel_core::drafter::dflash::round_loop::DEFAULT_MASK_TOKEN_ID, ); + if let Some(per_row) = per_row { + generator = generator.with_per_row_accepts(per_row); + } let output = generator .run_batched( model, @@ -169,6 +201,7 @@ pub(crate) fn run_burst( MAX_TOKENS, ) .expect("batched DFlash run must succeed"); + let decode_ms = output.stats.decode_time_ms; let tokens = output .tokens .into_iter() @@ -180,6 +213,7 @@ pub(crate) fn run_burst( accept_lens: output.accept_lens, committed_accept_lens: output.committed_accept_lens, clamped_rounds: output.clamped_rounds, + decode_ms, }; (run, generator.into_drafter()) } diff --git a/tests/speculative_parity/qwen35_dflash_throughput.rs b/tests/speculative_parity/qwen35_dflash_throughput.rs new file mode 100644 index 000000000..8fc883c3c --- /dev/null +++ b/tests/speculative_parity/qwen35_dflash_throughput.rs @@ -0,0 +1,147 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Throughput comparison for the per-row DFlash accept port (issue #2442). +//! +//! Kept in its own module so the parity gate's command +//! (`--test speculative_parity qwen35_dflash_batched`) does not pull a +//! multi-minute benchmark into the correctness run. +use super::qwen35_dflash_harness::{TEXTS, mean_accepts, run_burst, run_burst_with_policy}; +use super::{REACHABLE_PAIRINGS, pairing_present}; + +/// Issue #2442 gate 2: aggregate decode throughput of the per-row accept port +/// against the uniform clamp, at B = 4, three runs per arm, alternating. +/// +/// Both arms run in one process on the same prompts and the same `[B, L]` +/// prefill shape; only the accept policy differs. The reported number is the +/// round-loop wall clock's tokens/sec (`DFlashBatchedRunOutput::stats`), so a +/// change in either the number of rounds or the cost per round shows up. +/// Prints the walk and committed accept counts the #2442 report needs. +/// +/// Run on a quiet box (`caffeinate -i` plus +/// `scripts/with_indexers_paused.sh`); this test fails only when an arm +/// cannot run. The ship/keep verdict is made from the printed numbers. +#[test] +#[ignore = "real-model heavy (Qwen3.5-4B target + DFlash drafter, B=4 clamp vs per-row throughput); quiet-box manual run"] +fn throughput_dflash_qwen35_batched_clamp_vs_per_row() { + use mlxcel::{LoadedModel, initialize_runtime, load_model}; + use mlxcel_core::drafter::{DrafterKind, load_drafter}; + + fn median(values: &[f64]) -> f64 { + let mut sorted = values.to_vec(); + sorted.sort_by(|a, b| a.partial_cmp(b).unwrap()); + sorted[sorted.len() / 2] + } + + let pairing = &REACHABLE_PAIRINGS[0]; + let (target_path, draft_path, present) = pairing_present(pairing); + if !present { + eprintln!( + "Skipping {} (throughput): target={:?} draft={:?}", + pairing.name, target_path, draft_path, + ); + return; + } + + super::common::apply_server_mlx_cache_defaults(); + let _runtime = initialize_runtime(); + mlxcel_core::synchronize_default(); + mlxcel_core::clear_memory_cache(); + + let (loaded, tokenizer) = load_model(&target_path).expect("target model must load"); + let model: &mlxcel::models::Qwen35Model = match &loaded { + LoadedModel::Qwen35(m) | LoadedModel::Qwen35Moe(m) => m, + LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => &vlm.text_model, + _ => panic!("DFlash pairing requires a Qwen 3.5 family target"), + }; + let encoded: Vec> = TEXTS + .iter() + .map(|t| { + tokenizer + .encode(t, true) + .expect("prompt must tokenize") + .into_iter() + .map(|id| id as i32) + .collect() + }) + .collect(); + let len = encoded.iter().map(Vec::len).min().expect("four prompts"); + let prompts: Vec> = encoded.into_iter().map(|p| p[..len].to_vec()).collect(); + + let (mut drafter, kind) = + load_drafter(&draft_path, Some(DrafterKind::Dflash)).expect("DFlash drafter must load"); + assert_eq!(kind, DrafterKind::Dflash); + let block_size = pairing.block_size; + + // Warm up both arms so kernel selection and allocation are not measured. + for per_row in [false, true] { + let (_run, d) = run_burst_with_policy(model, drafter, &prompts, block_size, Some(per_row)); + drafter = d; + } + + // B = 1 reference, for context only. + { + let (run, d) = run_burst( + model, + drafter, + std::slice::from_ref(&prompts[0]), + block_size, + ); + drafter = d; + let (walk, committed) = mean_accepts(&run); + eprintln!( + "[tp B=1] {:.2} tok/s, mean accept walk {walk:.3} committed {committed:.3}", + run.aggregate_tok_per_sec() + ); + } + + let mut clamp_tps = Vec::new(); + let mut port_tps = Vec::new(); + let mut clamp_accepts = (0.0, 0.0); + let mut port_accepts = (0.0, 0.0); + for round in 0..3 { + for (label, per_row) in [("clamp", false), ("port", true)] { + let (run, d) = + run_burst_with_policy(model, drafter, &prompts, block_size, Some(per_row)); + drafter = d; + let (walk, committed) = mean_accepts(&run); + let tps = run.aggregate_tok_per_sec(); + eprintln!( + "[tp B=4 {label} run {round}] {tps:.2} tok/s, rounds {}, clamped {}, \ + mean accept walk {walk:.3} committed {committed:.3}", + run.accept_lens[0].len(), + run.clamped_rounds, + ); + if per_row { + port_tps.push(tps); + port_accepts = (walk, committed); + } else { + clamp_tps.push(tps); + clamp_accepts = (walk, committed); + } + } + } + let clamp = median(&clamp_tps); + let port = median(&port_tps); + eprintln!( + "[tp B=4 VERDICT] clamp median {clamp:.2} tok/s (runs {clamp_tps:?}), \ + port median {port:.2} tok/s (runs {port_tps:?}), ratio {:.3}x", + port / clamp, + ); + eprintln!( + "[tp B=4 VERDICT] committed accepts per row-round: clamp {:.3} (walk {:.3}), \ + port {:.3} (walk {:.3})", + clamp_accepts.1, clamp_accepts.0, port_accepts.1, port_accepts.0, + ); +}