feat(scheduler): admit ragged cohorts — the admission layer #103 named, and a general throughput bug

[heydryft]

Stacked on #103 (feat/mtp-draft-chain-per-seq) → #102 → #100 → #95 → #92. Review those first.

The bug, in one sentence

Sequences whose cache lengths differ cannot share a forward pass at all — the scheduler partitions the running set by Sequence::cache_bucket_len and runs exactly ONE bucket per step — and in real serving, lengths always differ.

#103 named this as batch_admission_admits_ragged_cohorts() and framed it as the sixth MTP blocker. It is not primarily an MTP blocker. It is the reason a fleet's admitted batch is not the batch that runs.

📊 Bucket shattering is real, and it is a law

bucket_shattering_law (default_scheduler.rs) is a CPU simulation of the scheduler's own logic — not a hardware measurement (D14). It admits B sequences at D distinct cache lengths 64 tokens apart (far enough that select_running_bucket's coalescence override is refused) and counts how many the scheduler puts in each forward:

                     sequences run per step, steady state
  B      D           refused (today)        granted (this PR)
  8      8                 1                       8
 32      8                 4                      32
 32     32                 1                      32
128     32                 4                     128
128    128                 1                     128

effective_batch = B / D. At B=128 with every sequence at its own length, one sequence runs; the other 127 are admitted, resident in KV, counted as running by the logger, and idle. admitted_batch_size_stops_buying_anything_once_lengths_diverge states the same thing as a curve: [1, 1, 1] at B = 8/32/128 refused, [8, 32, 128] granted.

The granted column is the negative control in the direction that matters — same sequences, same steps, same arithmetic, only the bucket key differs.

⚠️ What this does NOT establish. It does not show that this is what produced the H200's B=32 → 47.45, B=128 → 43.73 curve (FACTS.md session 8). That run had --mtp-depth 3, prefix cache off, and --max-seqs 128 pinning the B=128 row. What the simulation establishes is that the scheduler has a mechanism that would produce exactly that shape, and how large it is. Whether it fired on that box is a hardware question — see the GPU ask.

The fix — two sites, one decision

site	before	after
scheduler/default_scheduler.rs	BucketKey = (cache_bucket_len, imgs, offset), one bucket per step	4th component ragged; every completion sequence takes one shared key (ragged_bucket_len), so the batch stays whole however far apart its caches are
engine/mod.rs	pre_op = In iff the completion-id list changed	issues_cache_in(no_kv_cache, ids_changed, granted) — a granted run re-assembles every step, because the front-alignment is a function of lengths that diverge, and In is the only place a per-row lead_pad is reported at all

Both read one RaggedAdmission, decided once in Engine::new from CacheManagerMixin::ragged_batch_admission(). Deciding it in one place is the point: a scheduler that formed ragged cohorts for a pipeline the engine never re-assembled for would front-align once and then decode against a stale alignment.

🔑 Why flag-off is byte-identical, structurally

Not "a test observed the same numbers" — the two predicates are the pre-change expressions:

• ragged_bucket_len returns None when refused, so bucket_key_of is (cache_bucket_len, imgs && is_prompt, token_offset, false) — the historical three-part key with a constant appended. No sequence can move between buckets, at any length, in any state (a_refused_run_builds_the_pre_change_bucket_key asserts it component-for-component).
• issues_cache_in(no_kv, changed, false) == !no_kv && changed over every input (a_refused_run_issues_exactly_the_pre_change_cache_instructions).
• B=1 is one sequence and one bucket under either rule (b1_is_untouched_by_either_admission_rule).

No seventh flag

The declaration rides ARC_MTP_PER_SEQ_KV, the flag #100/#102/#103 already coordinate on, plus ARC_V4_XS_PER_SEQ through cache_supports_per_sequence_advance. A run that did not ask for per-sequence advance declares a refusal by name: its rows never diverge, so ragged admission would cost a clone_in_cache per step and buy nothing.

🔴 The default declaration is a refusal, not false

CacheManagerMixin::ragged_batch_admission defaults to Err(reason) naming the pipeline. Serving a ragged cohort means left-aligning it, and the zero-filled dead prefix that leaves is not a masked row — it scores logit 0 and takes real softmax weight. Opting in is the safety property; the_default_declaration_is_a_refusal_that_names_itself pins it, and the engine logs the reason so a run that asked and did not get it is told which layer said no (D18).

🔴 The trap — the barrier re-forming as padding, at the admission key

#102 found it in the target cache, #103 one cache down. The admission layer's copy is the nastier one: Sequence::cache_bucket_len — the single function bucket_key_of reads — is normal_cache()[0].current_seq_len(), and after front_align_batch that is the batch maximum for every row. A caller that records the padded length makes every sequence report the same value, and then:

• every row advances at the fastest row's rate, so their separation freezes at the per-step commit spread instead of growing;
• the length-bucketed scheduler stops shattering — not because the batch became uniform, but because the lengths it buckets on are no longer the sequences' own;
• the ragged bucket key looks like it is doing its job, over a batch whose raggedness has been silently erased.

the_scheduler_reads_each_row_s_own_length_not_the_padded_one runs the pad → forward → commit → strip loop over real tensors and asserts on the growth of the spread: 3 + 2n with the strip, a constant 2 without it, at n = 1/6/12. It reads current_seq_len() back off the slot rather than recomputing it — which is exactly how #103's first draft of the equivalent test passed with the strip stubbed out.

Mutation runs — 26 mutations, 26 caught

ragged_bucket_len: ignore the admission gate ......................... CAUGHT
ragged_bucket_len: fold prompts into the shared length ............... CAUGHT
ragged_bucket_len: take the MAXIMUM not the minimum .................. CAUGHT
bucket_key_of: admit prompts into the ragged bucket .................. CAUGHT
bucket_key_of: drop the ragged component from the key ................ CAUGHT
DefaultScheduler: never forward the grant to the bucketing manager ... CAUGHT
select_running_bucket: drop `ragged` from the coalescing filter ...... CAUGHT
issues_cache_in: drop the `|| granted` arm ........................... CAUGHT
issues_cache_in: ignore no_kv_cache .................................. CAUGHT
RaggedAdmission::decide: grant on a refusal .......................... CAUGHT
RaggedAdmission::decide: drop the refusal reason ..................... CAUGHT
RaggedAdmission::buckets_by_cache_length: invert ..................... CAUGHT
RaggedAdmission::assembles_every_step: never assemble ................ CAUGHT
batch_admission_admits_ragged_cohorts: always admit .................. CAUGHT
batch_admission_admits_ragged_cohorts: never admit ................... CAUGHT
ragged_cohort_capability: drop the ARC_MTP_PER_SEQ_KV gate ........... CAUGHT
ragged_cohort_capability: drop the ragged-mask check ................. CAUGHT
ragged_cohort_capability: drop the per-sequence-length check ......... CAUGHT
ragged_batch_admission: declare from per_sequence_refusal (the cycle)  CAUGHT
set_ragged_batch_admission: ignore the grant ......................... CAUGHT
MtpSpeculativePipeline: do not override ragged_batch_admission ....... CAUGHT
CacheManagerMixin::ragged_batch_admission default: grant not refuse .. CAUGHT
drop_dead_prefix: make the strip a no-op ............................. CAUGHT
front_pad_single: pad at the END instead of the front ................ CAUGHT
front_align_batch: align to the batch MINIMUM ........................ CAUGHT
Sequence::cache_bucket_len: ignore the cache, use the token count .... CAUGHT

Four survived the first run and the fixture was the problem in three of them, as four PRs in this chain have now found:

• fold prompts into the shared length survived because the fixture's shortest prompt happened to equal the completions' minimum — a prompt shorter than every completion was added;
• drop ragged from the coalescing filter survived because nothing tested that the override must be refused across the ragged flag (two buckets that can never merge — one a [B,1] decode, one a [B,t] prefill — must not trigger a merge-justified override). the_coalescence_override_never_targets_a_bucket_it_cannot_merge_with is set up as exactly the case where the override would fire if the keys looked mergeable;
• drop the per-sequence-length check survived because both fixtures had ARC_V4_XS_PER_SEQ on, where the cache check and the mask check agree. With the flag off a V4 cache still masks a ragged batch but its XsRolling slots carry one length, and the two checks separate;
• default grant instead of refuse survived because nothing exercised the trait default — the position every pipeline but the MTP one is in.

🔴 PerSequence is reachable — and where the refusal moved

per_sequence_advance_is_reachable_once_the_engine_grants_admission runs the whole chain through the real trait methods on a V4-shaped MtpSpeculativePipeline: declare → RaggedAdmission::decide → set_ragged_batch_admission → kv_advance() == KvAdvance::PerSequence. Its control is the identical pipeline with the grant withheld, which must stay Cohort. Five waves refused this mode; it is now on.

I am not claiming this is the last blocker, and for the general throughput problem it plainly is not. The refusal moved to model_masks_ragged_batches: no architecture in this tree except DeepSeek V4 threads a ragged-batch mask into its forward, so every other model still declares a refusal and keeps length-bucketed admission. The admission layer is fixed for everyone; the masks are not. That is now a per-architecture job, not a layer — CausalMasker::make_causal_mask_matrix already routes on PastKvLenCache::per_seq_kv_lens, so what each model needs is to hand it RaggedKvLens instead of its cache, plus clone_in_cache/clone_out_cache learning front_align_batch + drop_dead_prefix for the plain (non-MTP) decode path. the_declaration_still_refuses_every_model_that_cannot_mask_a_ragged_batch is the test that will fail the day someone relaxes that without doing the work.

📊 Measured — CPU tests and simulation, not hardware (D14)

cargo test -p mistralrs-core --lib → 454 passed, 0 failed (was 442 at #103's tip) · --test synthetic_load_smoke → 13 passed · cargo check --workspace --tests green · scoped clippy lane exits 0 · mistralrs-core clippy diagnostics diffed per touched file before/after: zero new (engine/mod.rs 0→0, kv_cache/mod.rs 4→4, pipeline/mod.rs 9→9, mtp_pipeline.rs 7→6, default_scheduler.rs 1→1, scheduler/mod.rs 0→0, ragged_admission.rs 0, sequence.rs 0→0; workspace total 238→237) · zero rustfmt drift added: rustfmt follows mod declarations and reformatted six upstream files it should not have, so that was reverted and the remaining noise stripped with git merge-file against a formatted baseline — the diff is +54/-4 in engine/mod.rs and +40/-0 in pipeline/mod.rs, all of it new.

🔴 The GPU ask — the experiment that proves or kills the hypothesis

arc-tools/ragged_admission_curve.sh. I did not run it (D15 — I never call runcrate). Self-contained, setsid nohup … < /dev/null, status file at /root/logs/ragged_admission_curve.status, gated on /root/arc-tools/gpu_box_preflight.sh (refuses to measure if it fails or is missing).

MODEL=<v4> setsid nohup bash arc-tools/ragged_admission_curve.sh \
  < /dev/null > /root/logs/ragged_curve.log 2>&1 &

Four arms × B ∈ {1, 8, 32, 64, 128}, --max-seqs 256 so the cap never pins the B=128 row (it did in session 8), --mtp-depth 3 held constant across all arms so the MTP chain is not a confound:

arm	flags	prompt lengths
A	none	uniform
B	none	mixed (32…2048 words, mixed output lengths)
C	ARC_V4_XS_PER_SEQ=1 ARC_MTP_PER_SEQ_KV=1	uniform
D	same	mixed

🔑 A vs B is the whole test and it needs none of this PR's code — same binary, same flags, same model, only the prompt-length distribution differs.

The one number: tok_s(B) / tok_s(A) as a function of B. Shattering predicts it falls as B grows (refused effective batch is B/D, uniform is B). If it is flat at ~1.0 across the sweep, the hypothesis is dead for this workload — which is a real answer and should be recorded as one. Then tok_s(D)/tok_s(B) at B=128 is what this PR is worth on the mixed workload, and tok_s(*, B=1) must agree across all four arms.

⚠️ The most direct fingerprint is not available today, and the script says so rather than faking it. IntervalLogger prints "<N> running, <M> waiting", and N is literally how many sequences the scheduler put in the forward — no throughput reasoning needed. But throughput_logging_enabled is only reachable from the Rust SDK builder (with_throughput_logging); mistralrs serve exposes no flag for it, so the line is never emitted by a server run. scrape.py looks for it and reports NA with the reason when it finds nothing — never 0.00, because a zero would read as "no sequences ran" when it means "nobody was logging" (D18).

Noticed: serve has no --throughput-log, so the engine's own running/waiting counters are unreachable from the CLI. Worth a separate change?

[github-actions]

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                189        37703            0        28931         8772
 |- BASH                  70         1618         1189          313          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)                            43427         4663        29455         9309
─────────────────────────────────────────────────────────────────────────────────
 Rust                    662       307804       266560        13893        27351
 |- Markdown             477        22853          471        19669         2713
 (Total)                           330657       267031        33562        30064
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1277       451971       330166        73659        48146
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

[heydryft]

Holding this PR pending a review of its interaction with #122 — not a CI problem, and CI will not find it

#122 (ragged dense decode) has landed on a master that already contained #100. Its files and this PR's overlap in two places, and unlike the rest of the stack the overlap is semantic, not incidental:

file	#122 adds	this PR does
kv_cache/mod.rs	thread_local! { RAGGED_LEAD_PAD } — lead_pad[i], the dead zero-filled prefix of row i, set by clone_in_cache and read by CausalMasker; plus RAGGED_DECODE_SUPPORTED / set_ragged_decode_supported()	exercises drop_dead_prefix(slot, leads[i]) and asserts "the scheduler reads each row's own length, not the padded one"
scheduler/default_scheduler.rs	ragged decode admission	ragged cohort admission

Both are mechanisms for "how much of row i is dead prefix." #122 carries it as a thread-local channel so the masker can recover live[i] = past_kv_len - lead_pad[i]; this PR mutates it away with drop_dead_prefix and reasons about the resulting lengths in the scheduler. Those can compose, but they can also duplicate, and a duplicated mechanism is precisely the shape that produces a silent disagreement later under an input neither author tried — no test fails when two mechanisms merely agree today.

Git auto-merged #122 and #100 cleanly with 466/466 green, which establishes their texts do not conflict and nothing more.

What's needed before this lands: somebody reads both and answers whether they compose, duplicate, or contend — specifically whether the thread-local lead_pad is still accurate on the paths where this PR has stripped the prefix, and whether the two admission changes in default_scheduler.rs are describing the same cohort geometry twice.

That's a review question, not a CI question, and I'm the merge chain rather than a reviewer — so this holds rather than proceeding on a green gate.

Knock-on: #114 and #116 are stacked above this one, so both are held by ordering. On their own merits they are clear of #122 — #114 touches attention/mod.rs but only to add a OnceLock that names the attention backend once per process (ARC_ATTN_BACKEND divert detection, not the masking path), and #116 touches kv_cache/xs_rolling.rs and deepseek4.rs, neither of which #122 goes near. Nothing further from this stack lands until the question above is answered.

[heydryft]

The held question, answered — read on master with #122 landed. Compose, duplicate, AND contend — in three different places.

The hold asked: do these compose, duplicate, or contend — specifically, is the thread-local lead_pad still accurate on the paths where this PR has stripped the prefix, and are the two admission changes describing the same cohort geometry twice? I merged master into this branch in a scratch worktree only, read all three overlaps, and aborted the merge — this branch and its remote are untouched (eeee5571b).

1. drop_dead_prefix — NOT a competing mechanism. It is already on master, byte-identical.

This was the premise of the hold and it is wrong. Extracted both bodies and diffed:

master lines: 16  |  this PR lines: 16
diff -> IDENTICAL

It landed via the MTP chain, and master has 10 call sites for it. Merging this PR adds no second copy of it. That half of the concern is retired.

2. The bucket-key length pin — genuinely DUPLICATED. Same loop, same expression, same purpose.

The one real conflict in default_scheduler.rs, verbatim:

// this PR
let key = bucket_key_of(&seq, ragged_len);
// master (#122)
let len = if ragged_decode && !seq.is_prompt() { 0 } else { seq.cache_bucket_len() };

Both pin the length component so a length-diverse decode cohort forms one bucket. The same cohort geometry, described twice. Note the line above the conflict auto-merged, so both already read the same capability flag. This is a straightforward "one of these two spellings survives" — no semantic gap.

3. 🔴 The cache-in cadence — they CONTEND, and the failure is silent and costs throughput.

This is the real finding and it is not what the hold was looking for.

engine/mod.rs, the one conflict there:

// this PR
let pre_op = if issues_cache_in(self.no_kv_cache, ids_changed, self.ragged_admission.granted()) {
// master (#122)
let pre_op = if !self.no_kv_cache && last_completion_ids != current_completion_ids {

This PR's own comment states why: "Under ragged admission the layout is a function of the rows' lengths — which diverge — so it has to be rebuilt every step."

#122 is built on the exact opposite assumption. strip_pending_lead_pad lives inside clone_in_cache, and its doc says: "clone_out_cache records the prefix here — no tensor work — and clone_in_cache pays for it once, on the membership change that actually re-reads those caches" — because "dropping each row's dead prefix on every step was doing O(B x layers x 2) device copies per token (~2,200 copies/token at B=47)." Commit 3adb69477 exists solely to defer that.

⇒ Land both as-is and issues_cache_in forces clone_in_cache every step, which fires the deferred strip every step, which reinstates the ~2,200 device copies per token that #122 measured and removed. Nothing errors. No test fails. Correctness is likely preserved — within one clone_in the order is strip → re-align → re-publish lead_pad, which is self-consistent — so the only symptom is throughput quietly returning to the pre-#122 number. That is exactly the "silent disagreement under an input neither author tried" this hold was placed for; it just lives in the cadence, not in lead_pad's accuracy.

These are two mutually exclusive designs for one problem:

	#122 (on master)	this PR
layout	front-align all rows to a shared width once	keep per-row lengths
dead prefix	masked by CausalMasker via RAGGED_LEAD_PAD	stripped by drop_dead_prefix
clone_in	membership change only	every step
cost	one alignment per cohort	one re-assembly per token

What this needs — and it is now a measurement, not a review. Master already ships design A. The question is whether design B buys anything design A does not, priced against B's per-step rebuild. Until someone runs that A/B, merging this PR is a throughput regression with no test to catch it. I am leaving the hold in place, but the question is no longer open-ended: it is one number.

🟠 Separate finding, ON MASTER, not this PR's fault — two strip implementations with opposite failure policies

drop_dead_prefix refuses a non-Normal slot with a named error ("Nothing front-pads it, so nothing should be stripping it either"). #122's strip_pending_lead_pad reimplements the same narrow inline and silently skips any non-Normal slot — if let KvCache::Normal { k, v } = slot with no else. Same operation, one refuses by name and one says nothing. Filing separately; it is a defect on master, not a reason to hold this PR.

Nothing deleted. This branch, and #113 / #114 / #116 / #121 above it, are untouched.

[heydryft]

Closing as superseded, not abandoned.

The admission layer this PR proposed is already on master by another route: mistralrs-core/src/scheduler/default_scheduler.rs:273 now gates the length-bucketing on crate::kv_cache::ragged_decode_supported(), so non-prompt sequences skip seq.cache_bucket_len() entirely — which was this PR's actual behavioural change. It landed via #123 (perf/prefill-floor-plus-dense-ragged). The RaggedAdmission / scheduler/ragged_admission.rs abstraction proposed here was never needed to get there; mistralrs-core/src/scheduler/ contains only default_scheduler.rs and mod.rs today.

What is given up: nothing functional. Reopen if the ragged-cohort admission needs to become a named policy object rather than a boolean gate.