Skip to content

feat(v4): per-row q0 in dsv4_attention — the ragged-mask blocker #92 named - #100

Merged
heydryft merged 5 commits into
masterfrom
feat/v4-ragged-mask-wire
Aug 17, 2026
Merged

heydryft merged 5 commits into
masterfrom
feat/v4-ragged-mask-wire

Conversation

@heydryft

Copy link
Copy Markdown
Contributor

Stacked on #95 (feat/xs-per-sequence), which is stacked on #92. Review those first.

The blocker, in one sentence

dsv4_attention derived ONE scalar absolute query position from the cache width and used it for both of V4's time bases — so a left-aligned ragged cohort could not be masked, and KvAdvance::PerSequence stayed unreachable.

PR #92 named two blockers. #95 removed the first and the refusal moved to this one. This removes it. Dsv4AttentionConfig now carries row_q0: Option<&[usize]>, and cache_supports_per_sequence_advance + target_masks_ragged_batches both pass for a V4 target.

🔑 Left-alignment breaks exactly ONE of the two time bases — and adds a third problem

anchor under left-alignment
raw window end already correct. Shifting a row's queries and keys right by lead preserves their difference, and both the causal and the sliding-window tests are differences. In cache-column space the batch-wide q0 is right for every row.
raw dead prefix — lead zero columns appear ahead of each row. A zero K row scores logit 0 and takes real softmax weight.
compressed blocks start too permissive. Column b is absolute block b for every row, so b < (q_abs + 1) / ratio needs that row's own q_abs. Fed the batch's, a row that advanced less attends blocks covering tokens it never generated.

The third row is why #92's [B,1,t_q,k] K/V mask is not sufficient, and #95 said so. What it could not have known is that a caller mask cannot express it either: ratio and the compressed width t_c are per-layer (0/4/128; ratio-0 layers have no compressed columns at all) while CausalMasker builds one mask for all 43 layers. So the correction has to live where ratio is known — inside dsv4_attention. compose_caller_mask's rank-4 support stays exactly as useful as it was; it just is not the vehicle for this.

🔑 The wire was already threaded — it was being collapsed, not missing

seqlen_offsets[i] is already row i's own past length: the engine threads it to every attention call so RoPE places that row's queries at their true absolute positions. It is exactly the q0 the compressed branch needs and, via lead_i = q0 - row_q0[i], exactly the dead prefix the raw branch needs.

So no value is plumbed anywhere. Attention::forward already receives seqlen_offsets; what changed is that dsv4_attention stops collapsing the vector to one scalar. No signature churn through DecoderLayer, no new field on ModelInputs, no HashMap<seq_id, …>.

Only the dense arm opts in. PagedAttention keeps per-request block tables and the graph-decode arm reads a fixed window whose columns are device slots, not positions — neither is a left-aligned cohort, so both keep row_q0: None and are untouched.

Why "flag off is byte-identical" is structural

resolve_ragged_rows returns None unless some row differs from q0, and None runs the pre-change scalar code verbatim. Every B=1 request, every prefill and every batch built with the flag off takes that branch. No third flag: the gate is ARC_MTP_PER_SEQ_KV, because front_align_batch is the only producer of a ragged dense batch and it runs only under KvAdvance::PerSequence.

⚠️ That guarantee needed a test of its own. A mutation that routes uniform batches down the ragged path is numerically invisible — the two branches agree bit-for-bit on their common domain — and it survived every other test here. It is still a real defect (a [b,1,t_q,n_keys] mask plus two host→device copies on every layer of every B=1 step). a_uniform_batch_never_reaches_the_ragged_path pins the dispatch directly; that is what resolve_ragged_rows exists as a separate function for.

📊 Measured — CPU unit tests on the real dsv4_attention, not hardware (D14)

cargo test -p mistralrs-core --lib → 407 passed, 0 failed (was 405 on #95's tip, 398 at #95) · --test synthetic_load_smoke → 13 passed · cargo check --workspace --tests green · scoped clippy lane green · zero rustfmt drift (checked like-for-like: dsv4_attention 8→8, deepseek4 341→341, mtp_pipeline 0→0; kv_cache/mod.rs adds 4 lines and reformats nothing).

The core test runs each row of a ragged cohort against that row run alone with its dead prefix absent, with its own input stream per row (a shared stream would make block j cover identical tokens on every row and hide any mix-up). Its three degeneracy assertions run first and fail loudly if a later edit closes one: rows must land on different compressed thresholds, at least one threshold must be below t_c, and at least one dead prefix must reach inside the retained raw span. A negative-control test shows the same fixture disagrees when fed the batch-wide position.

Mutation runs — 12 mutations, all 12 caught

dsv4_attention: always take the uniform path .................. 3 FAILED
ragged_union_mask: drop the dead-prefix kill .................. 2 FAILED
ragged_union_mask: compressed threshold from the BATCH pos .... 1 FAILED
ragged_union_mask: lead = rows[i] instead of q0 - rows[i] ..... 1 FAILED
ragged_union_mask: drop the compressed correction entirely .... 2 FAILED
resolve_ragged_rows: never refuse a row AHEAD of the batch ..... 1 FAILED
resolve_ragged_rows: never refuse "no row at the end column" ... 1 FAILED
resolve_ragged_rows: never refuse a wrong-width vector ......... 1 FAILED
resolve_ragged_rows: send UNIFORM batches down the ragged path . 1 FAILED
ragged_row_q0_from: ignore the opt-in flag .................... 1 FAILED
model_masks_ragged_batches: never refuse ...................... 1 FAILED
per_sequence_refusal: drop the third capability check ......... 1 FAILED

The last four exist because of #95's warning. The uniform-dispatch and opt-in-flag mutations both survived their first drafting and the tests were rewritten around them, not the other way round.

🔴 The honest gap — PerSequence is still refused, and I found out why

This does not move the B=128 number, and the reason is not the mask. With both previously-named blockers cleared, a third one is now visible that neither #92 nor #95 could see from where they stood. MtpSpeculativePipeline::step_full keys the entire fused step on ONE scalar cache_len = current_normal_cache_len(self) — the batched buffer's width, which left-alignment makes max_j L_j, not any row's length. Four uses break, in different ways:

  1. uncached[i] = toks_i.len() - cache_len (:2546) saturates to 0 for every row but the longest ⇒ window_ok fails ⇒ every ragged step defers. Per-sequence advance would be reachable but inert.
  2. toks[cache_len..] (:2719) is a panic site once cache_len > toks_i.len(), unreachable today only because window_ok bails first.
  3. run_target_forward_batch(.., Some((w, cache_len))) (:2725) ⇒ seqlen_offsets[i] = cache_len for every row (inputs_processor.rs:244 — last_n_context_len.1 is a shared scalar). So the mask this PR built is correct but would be handed uniform values.
  4. set_len(cache_len + committed_i) (:2925) counts the dead prefix inside the recorded length (front_pad_single sets current_seq_len = target_len; clone_out_cache copies the batched length verbatim into every row). The next front_align_batch therefore pads relative to already-padded lengths and the dead prefix accumulates linearly — contradicting front_pad_single's doc claim of a sqrt(steps) random walk. Nothing in Sequence or SingleCache records a row's true live length; the only honest source is seq.get_toks().len().

⚠️ Granting the mode on the first two checks alone would be actively harmful, not merely useless. The deferral in (1) lands in the target's own decode, which builds no mask at all at t_q == 1 (layers_masker.rs:285) — so the dead prefix would be attended as real keys by the very fallback that was supposed to be safe. That is why kv_advance gains a third named check rather than being switched on, and why a_fully_capable_v4_cache_is_still_refused_and_now_names_the_step asserts the refusal text names cache_len and no longer names XsRolling.

So the observable effect of ARC_MTP_PER_SEQ_KV=1 today is, again, that the refusal moves. It is stated as a capability function (step_carries_per_sequence_lengths) so the next change flips exactly one thing. That change is pipeline arithmetic plus a per-row route through inputs_processor — not a model change, which is why it is not here.

🔴 The GPU ask

Nothing changes a default, so this is a no-regression check on the box already serving V4:

mistralrs bench -m <v4> -b 1 -b 128 2>&1 | tee /tmp/a.txt
ARC_V4_XS_PER_SEQ=1 ARC_MTP_PER_SEQ_KV=1 mistralrs bench -m <v4> -b 1 -b 128 2>&1 | tee /tmp/b.txt

The one number: decode tok/s in B as a fraction of A, at B=1 and B=128; the claim is B/A == 1.00 at both. And grep "cannot honour it" /tmp/b.txt must now cite the fused step's shared cache_len and must contain neither XsRolling nor the ragged mask — if it names either, this wave did not land.

Noticed, not shipped

  • ragged_union_mask builds two b-element host tensors per layer per step (~43 tiny H2D copies per forward, flag-on path only). The file already does this for kp/qp; hoisting all of them to one per-forward build is a separate change.
  • make_causal_mask_matrix already has the RaggedKvLens hook, and PastKvLenCache for Vec<KvCache> returns None for per_seq_kv_lens. Giving the cache a record of its own padding would make the raw-side dead-prefix mask work for every model, not just V4 — but only V4 needs the compressed half, so it was out of scope here.

heydryft and others added 3 commits August 17, 2026 13:41
Stacked on #95. Removes PR #92's second named blocker: dsv4_attention now
takes a per-row absolute query position, so a left-aligned ragged cohort
gets its dead prefix masked on the raw branch AND its start-anchored
compressed threshold rebuilt from each row's own position.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Extract resolve_ragged_rows and per_sequence_refusal so the dispatch and
the capability composition are testable directly — a mutation that routed
uniform batches down the ragged path was numerically invisible.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@github-actions

github-actions Bot commented Aug 17, 2026 •

Copy link
Copy Markdown
Code Metrics Report
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Language              Files        Lines         Code     Comments       Blanks
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 C Header                  5          305          210           52           43
 CSS                       2         1181         1036           34          111
 CUDA                     72        24328        17592         4018         2718
 Dockerfile                1           39           22            8            9
 JavaScript               16         3546         2676          482          388
 Jinja2                    7          694          656            5           33
 JSON                     74         4600         4597            0            3
 Makefile                  1            6            5            0            1
 Metal Shading Lan|       33        12224         9431         1142         1651
 PowerShell                1          300          227           30           43
 Python                  143        14830        12217          797         1816
 Shell                    23         5756         3964         1412          380
 Plain Text                4         3801            0         2479         1322
 TOML                     33         1485         1292           43          150
 YAML                      3           25           23            2            0
─────────────────────────────────────────────────────────────────────────────────
 HTML                      4         2687         2604           43           40
 |- CSS                    2          543          479           37           27
 |- JavaScript             1         1233         1215           12            6
 (Total)                             4463         4298           92           73
─────────────────────────────────────────────────────────────────────────────────
 Jupyter Notebooks         4          122           83           23           16
 |- Markdown               1           60           30           22            8
 |- Python                 1          122          113            1            8
 (Total)                              304          226           46           32
─────────────────────────────────────────────────────────────────────────────────
 Markdown                190        37932            0        29111         8821
 |- BASH                  71         1623         1192          315          116
 |- C                      2           12           12            0            0
 |- CUDA                   2           84           56           16           12
 |- JSON                  18          708          708            0            0
 |- PowerShell             1            1            1            0            0
 |- Python                23         1008          787          113          108
 |- Rust                  65         2048         1713           77          258
 |- TOML                   6          207          164            0           43
 |- YAML                   4           38           33            5            0
 (Total)                            43661         4666        29637         9358
─────────────────────────────────────────────────────────────────────────────────
 Rust                    662       309221       267775        14013        27433
 |- Markdown             477        23236          471        20016         2749
 (Total)                           332457       268246        34029        30182
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1278       454005       331384        74308        48313
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Absorbing master brought in a Dsv4AttentionConfig initializer that predates
this PR's row_q0 field, so the lib-test build failed with E0063 while the
lib build alone was clean — cargo check -p mistralrs-core does not build
tests, which is why this only surfaced in CI.

The fixture builds a uniform batch, and every other non-ragged fixture in
this file uses row_q0: None, so None is the identity: it preserves exactly
what the test asserted before the field existed.

Caught by ci-complete, not by the check-run count — the run showed 0 pending
while the test lane had failed and its two sibling matrix legs were
cancelled by fail-fast.
@heydryft
heydryft merged commit de9dc9f into master Aug 17, 2026
17 checks passed
@heydryft
heydryft deleted the feat/v4-ragged-mask-wire branch August 17, 2026 22:54
@heydryft

Copy link
Copy Markdown
Contributor Author

Admin-merged without review (no second reviewer on this fork). The bypass covered the review requirement only — the required CI complete context was genuinely green, not absent. First merge in this chain to clear the real gate rather than a manual lane read.

Merged without --delete-branch; #102 was retargeted to master first, then feat/v4-ragged-mask-wire was deleted (D20).

Worth recording: ci-complete caught a real break here on its first application. After absorbing master, the run reported lane results: success success failure success success success success success — the test lane had failed on macOS and its ubuntu/windows matrix siblings were cancelled by fail-fast. Master had brought in a Dsv4AttentionConfig initializer predating this PR's row_q0 field: E0063, lib-test build only.

Two things that missed it, both worth knowing:

  1. cargo check -p mistralrs-core passed, because -p builds the lib target and not tests. Only cargo check --workspace --tests reproduced it.
  2. My own poll loop printed pending=0 — it computed the pending count but not the failure count, so a failed lane and two cancelled ones read as ready-to-merge. The guard caught what the monitoring had stopped looking for.

Fix was row_q0: None on that one fixture — the identity value, matching the six other non-ragged fixtures in the file, preserving exactly what the test asserted before the field existed. Called out in its own commit rather than folded in silently, since it is an edit to a file this PR otherwise only extends.

heydryft added a commit that referenced this pull request Aug 17, 2026
feat(mtp): per-row `cache_len` in the fused step — the pipeline blocker #100 named
heydryft added a commit that referenced this pull request Aug 18, 2026
Master moved 10+ commits since 120f348, including #100 (v4 ragged mask wire)
and #102 (per-row cache_len in the fused MTP step) — adjacent territory. Merged
now rather than letting the divergence be discovered later.

Clean auto-merge; the only overlapping file is `kv_cache/mod.rs`. 466/466 green
on the merged tree (441 mine + master's new tests).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant