feat(sched): ArcSched — admit ragged decode batches on the paged path

[heydryft]

The defect

PagedAttentionScheduler — the scheduler every paged-attention model serves on — partitions each decode batch by KV length via bucket_and_preempt_sequences, runs the single shortest bucket, and drops the rest.

It does not merely waitlist them. _preempt calls kv_cache_manager.free(seq_id), so every dropped sequence loses its KV blocks and must be re-prefilled. With B sequences over D distinct lengths, B / D decode per step and the remainder are re-prefilled. In real serving lengths always diverge — users arrive at different times, with different prompts, and generate different numbers of tokens.

This is upstream (8010b6a05, mistral.rs #2354), not Arc-introduced. It is distinct from #104, which fixes the same class of bug on the dense DefaultScheduler path. The two are disjoint files and disjoint model sets.

Why the bucketing was unnecessary

Everything the decode forward consumes is already addressed per sequence. Each row contributes exactly one token, so there is nothing to pad:

what	where	form
input ids	make_completion_chunk	[B, 1] — one token per row, no padding
causal mask	layers_masker.rs	None, taken because tgt_len == 1
RoPE offsets	layers.rs	per-sequence loop; only the len() == 1 fast path narrows
block tables	make_completion_chunk	one table per sequence
context lens	make_completion_chunk	seq.len() per sequence
flash varlen	make_completion_chunk	seqlens_k[i] = 1 + start_pos_i, then cumsum
CUDA-graph buffers	arc-cuda-graph/src/buffers.rs	context_lens: [padded_bs] u32; max_context_len is a scalar arg, not a shape

The comment the bucketing shipped with — that equal lengths are "required for correct flash attention varlen operation (avoiding soundness issues with padding)" — has it backwards. Varlen is the mechanism for unequal lengths. The constraint being defended belongs to the dense cache (clone_in_cache templates one batch cache off seqs[0]), which is the DefaultScheduler's substrate, not this one.

Where this leaves us against vLLM and SGLang

Neither partitions a decode batch by KV length; both were read at pinned revisions rather than from memory.

• vLLM v0.27.1 — flash_attn_varlen_func(..., seqused_k=attn_metadata.seq_lens, cu_seqlens_q=query_start_loc, block_table=...) (vllm/v1/attention/backends/flash_attn.py:1041-1056). Per-row KV length; max_seqlen_* are launch hints only. The scheduler emits one flat dict[req_id -> num_tokens] (vllm/v1/core/sched/output.py:203-208).
• SGLang @ 2e7c85d — prepare_for_decode is a vectorized self.seq_lens = self.seq_lens + 1 over the whole running set (schedule_batch.py:3038-3082); FlashInfer reads kv_indptr[1:bs+1] = torch.cumsum(paged_kernel_lens, dim=0) (flashinfer_backend.py:2059-2095).

Both preempt only under memory pressure — vLLM LIFO-recompute (scheduler.py:577-625), SGLang retract_decode fewest-generated-first (schedule_batch.py:2816-2905). We preempted on length divergence, which neither does.

Honest standing after this PR: parity on ragged decode, still behind on chunked prefill (our prompt path right-pads to max_len and buckets), on CPU/GPU overlap scheduling (SGLang's is default-on), and on prefix-aware queue ordering (SGLang LPM; we are FCFS by timestamp). Those are named as follow-ups, not claimed here.

Structural flag-off identity

Not "a test observed the same numbers". The Bucketed arm holds the pre-change three lines verbatim, and the Ragged arm is a no-op because self.running already holds the admitted set:

match self.decode_admission {
    DecodeAdmission::Bucketed => {
        let running_for_bucket = std::mem::take(&mut self.running);
        let bucketed = self.bucket_and_preempt_sequences(running_for_bucket);
        self.running = bucketed;
    }
    DecodeAdmission::Ragged => {}
}

ARC_SCHED_BUCKETED_DECODE=1 restores it. Per D18 the refusal is a named value, not an absence: DecodeAdmission::decide returns the reason and new logs it, so a run that reverted to bucketing says so rather than looking like a run that was never fixed.

Speculative decode is not reachable here

Ragged decode assumes one token per row. Nothing on this path produces more: MtpSpeculativePipeline does not support paged attention, and mistralrs/src/speculative.rs builds a SchedulerConfig::DefaultScheduler.

🔴 Scope — this does NOT touch DeepSeek V4

DeepSeekV4Loader::supports_paged_attention returns false (normal_loaders.rs:3265), so V4 serves on DefaultScheduler and is unaffected. This PR helps every model that does use paged attention — the trait default is Ok(true), so most text models. V4's equivalent fix is #104's dense path. Anyone reading this as "V4 throughput fixed" is misreading it.

Measured — CPU tests, not hardware (D14)

cargo test -p mistralrs-core --lib → 380 passed, 0 failed. Scoped clippy lane exits 0. cargo check -p mistralrs-core -p mistralrs-cli -p mistralrs-server-core green.

Five tests, and two of them validate each other's fixture: bucketed_decode_shatters_the_batch_to_one returning exactly 1 is what proves the fixture really produces D = n distinct lengths — had divergence silently failed, every row would share a bucket and that test would fail rather than the ragged one passing vacuously.

test	asserts
ragged_decode_runs_the_whole_admitted_batch	B ∈ {8, 32, 128} at B distinct lengths → all B decode, none waitlisted
bucketed_decode_shatters_the_batch_to_one	same fixture, old rule → exactly 1
bucketed_decode_frees_the_kv_of_every_sequence_it_drops	the 7 dropped rows hold no blocks afterwards
b1_is_untouched_by_either_rule	B=1 decodes every step under both rules
the_refusal_names_itself	grant carries no reason; refusal carries ARC_SCHED_BUCKETED_DECODE=1

Mutation run. Making the Ragged arm bucket anyway — i.e. the fix as a no-op — fails with B=8 at 8 distinct KV lengths: only 1 of 8 decoded. The test is not vacuous.

Also in here

IntervalLogger printed a lifetime prefix-cache hit rate next to a per-interval throughput on the same line, because tokens_processed was swapped while prefix_cache_hits / total_new_seqs were loaded. A "1.43% prefix hit rate" was quoted from it that belonged to no single test. Now differenced against the previous reading so both cover the same window, and omitted rather than printed as 0.00% when no sequence arrived — a zero there reads as "nothing hit the cache" when it means "nothing was admitted to hit it".

Correction to a premise this work started from: mistralrs serve was believed to have no way to reach the engine's N running, M waiting counters. It already does — MistralRsBuilder::new(..., !self.interactive_mode, ...) (mistralrs_for_server_builder.rs:905) passes throughput_logging = true for every non-interactive run, and IntervalLogger fires every 5s. No --throughput-log flag is needed and none was added.

The hardware measurement

arc-tools/arcsched_decode_admission_ab.sh. Not run — D15, I never call runcrate. Same binary, same model, same workload; the only difference between arms is ARC_SCHED_BUCKETED_DECODE.

The direct fingerprint is mean_running scraped from N running, M waiting — that needs no throughput inference, because N is literally how many sequences went into the forward. Prediction stated in the script so it can be refuted: BUCKETED collapses toward 1 while RAGGED tracks B. If BUCKETED's running does not collapse, the defect did not fire for that workload and it should be recorded as a refutation, not retried. B=1 must agree across arms — one sequence is one bucket, so a disagreement there invalidates the run.

[heydryft]

ArcGate triage. NEEDS-OWNER — half of this is superseded, half of it is unique, and that is why I am not closing it.

The scheduler half is superseded by #118. Both fix the same defect — PagedAttentionScheduler putting seq.len() in the bucket key, running one bucket, and _preempt-ing the rest with their KV blocks freed. #118 is the stronger form on three counts:

• it is unconditional; this PR gates the fix behind ARC_SCHED_BUCKETED_DECODE, so the default is only fixed if nobody sets the escape hatch;
• it is per-leg correct: fix(ArcSched): stop bucketing paged decode by sequence length (8.7x on spread traffic) #118 derives the rule for all four legs and keeps LengthRule::MustMatch on the prefill-with-prefix-hit leg, where PagedAttention::forward swaps in packed cu_seqlens_kv while leaving cumulative_seqlens_q padded, so each row's causal offset seqlen_k - seqlen_q is short by its pad width. This PR only addresses the decode leg and leaves the prefill key alone;
• it carries the measured number (8.72×/9.15× at B=8 spread, 8.10×/8.35× at B=32); this PR carries the reasoning but no hardware result.

But this PR carries something #118 does not, and it is a real correctness fix in its own right: engine/logger.rs — the interval logger was mixing a per-interval throughput with a lifetime prefix-hit rate in the same line, because tokens_processed is swapped to zero each interval while prefix_cache_hits / total_new_seqs are cumulative shared counters. A "1.43% prefix hit rate" was quoted from that line and belonged to no single test. This PR differences the cumulative counters so both halves of the line cover the same window. Nothing else in the tree does that.

Also unique: arc-tools/arcsched_decode_admission_ab.sh (not present on master).

So: not closeable as-is. Per the standing rule, nothing gets closed whose content is not reproduced elsewhere, and the logger fix is not reproduced anywhere.

Suggested disposition for the owner — pick one, this is not mine to decide:

• (a) lift the logger.rs change (and the A/B script, if still wanted) into its own small PR, and close this one naming fix(ArcSched): stop bucketing paged decode by sequence length (8.7x on spread traffic) #118 as carrying the scheduler fix; or
• (b) keep this PR, drop the paged_attention/scheduler.rs half entirely, and retitle it as the logger fix.

Either way this branch needs to absorb master before it can produce a CI complete — it currently has 16 lanes and no CI complete job, and its comment check is failing. Branch untouched, nothing deleted.

[heydryft]

Closing as superseded, not abandoned.

The functional change here — stop bucketing the decode leg by KV length, so the losers are not preempted and re-prefilled — is already on master via 64b3ff379 fix(ArcSched): stop bucketing paged decode by sequence length. mistralrs-core/src/paged_attention/scheduler.rs:440 now reads self.bucket_and_preempt_sequences(running_for_bucket, LengthRule::Free), with LengthRule::Free => 0 at line 170 and the prompt/decode split at 361-363. That is this PR's DecodeAdmission::Bucketed -> free-decode transition, landed directly rather than behind a new enum.

What is given up: (1) the DecodeAdmission enum wrapper and its doc table of where each per-sequence quantity is addressed — the same reasoning survives in master's comments at the same site; (2) arc-tools/arcsched_decode_admission_ab.sh, an A/B harness that needs a GPU to say anything. Reopen the harness separately if a box comes back and the A/B is still wanted.

Note for anyone reviving this area: PagedAttention on V4 is currently a landmine — ARC_V4_PAGED_ATTN=1 is measured on A100 to produce ZERO TOKENS, and the Some(paged_attn) arm at models/deepseek4.rs:1592 shadows the graph-mode arm. Do not re-enable it as a side effect of scheduler work.