[aryatsc]

Motivation

Under overload on Kimi-K3 (MI355X, TP8, chunked prefill + hybrid KDA/MLA), the scheduler could stall or fail allocations:

• A continuing chunk was sized only by the token budget, so it could ask alloc_extend for more KV than is free.
• A second unfinished chunked request could be admitted under policies other than shortest-prefill-first, even though only one can be tracked (assert crash).
• A latched batch_is_full was cleared only when the running batch shrank, so the waiting queue stalled.
• On the aiter backend, an over-capacity extend inside graph context could leave the chunked-KV path.

Modifications

• available_chunk_tokens caps a continuing chunk by free + evictable KV, less the page alloc_extend over-asks. When that leaves nothing, the chunk proceeds with the free KV or is deferred.
• The one-chunked-request guard applies under every schedule policy.
• batch_is_full is re-evaluated every scheduling step; PrefillAdder still enforces the real limits.
• Over-capacity aiter extends stay on the chunked-KV path.

Accuracy Tests

New unit tests (test_prefill_memory_budget.py, test_prefill_adder.py, test_scheduler_batch_is_full_latch.py, test_aiter_attention_backend_handler.py): all 5 new cases fail on main and pass with this PR (76 passed).
End to end, the Kimi-K3 MI355X production build includes these fixes. On a 40-min overload soak (48 concurrent requests against 12 slots), 901/901 completed requests returned 200, the queue peaked at 37, and there were 0 OOM, HSA or assert errors.

Speed Tests and Profiling

Scheduling only; no kernel change.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #40956379040
Latest PR Test (Extra): ❌ Run #40956378917
Latest PR Test (AMD ROCm 10): ❌ Run #40956378897

[rkarhila-amd]

This PR has three small changes on scheduling:
(1) bysassing "shortest-prefill-first" check in chunked prefill policy means that all policies lead to same branch
(2) always resetting "batch_is_full" in _get_new_batch_prefill_raw function in the scheduler
(3) changing available_chunk_count calculation in prefill budget, so that available tokens calculation take into account evictable KV tokens among other factors. Also if 0 tokens are available, None is returned, as opposed to old behaviour where full chunk size was returned.

The changes are small but impact all models, and for this reason, I kindly ask for reproduction commands to show the problems that are fixed, and to verify with several other models that the performance with these scheduling changes is not adversely affected.

[aryatsc]

@rkarhila-amd thanks. Reproductions and the multi-model A/B are below.

Reproduction commands

Each fix has a unit test that fails on the merge base (fbc90a860e) and passes on this branch (ae64753ca5). To reproduce, run the new tests against the old scheduler code:

git worktree add /tmp/sched-base fbc90a860e
cd /tmp/sched-base
git checkout ae64753ca5 -- \
  test/registered/unit/managers/test_prefill_adder.py \
  test/registered/unit/managers/test_scheduler_batch_is_full_latch.py \
  test/registered/unit/mem_cache/test_prefill_memory_budget.py \
  test/registered/unit/models/test_aiter_attention_backend_handler.py
PYTHONPATH=python python -m pytest -q <the four files above>

Fix	Commit	Test (fails on base, passes on PR)	What it reproduces
1. One unfinished chunked request under every policy	ba516e7	test_prefill_adder.py::TestPrefillAdder::test_default_policy_rejects_second_unfinished_chunk	Under the default (non-SPF) policy, PrefillAdder starts a second chunk while one is already in flight. The scheduler only tracks one chunked_req, so the first one is lost.
2. Re-evaluate batch_is_full each step	274ec4c	test_scheduler_batch_is_full_latch.py::TestEmptyBatchFullLatch::{test_latched_empty_batch_still_admits, test_latched_running_batch_still_admits}	Once batch_is_full is set, _get_new_batch_prefill_raw returns None until an unrelated request finishes, even if the KV has freed up. Waiting requests stall.
3. Cap a continuing chunk by allocatable KV (incl. evictable), else defer	af2910e	test_prefill_memory_budget.py::TestContinuingChunkCap::{test_chunk_is_deferred_without_free_kv, test_chunk_leaves_the_page_alloc_extend_over_asks}	With no free or evictable KV, the old code returned chunk_limit anyway, so alloc_extend over-asks and fails. Now it returns a size that fits, or None to defer.
(ROCm) keep over-capacity aiter extends chunked	514ee13	test_aiter_attention_backend_handler.py::TestAiterOverCapacityExtend::test_over_capacity_extend_stays_chunked_in_graph_context

Result: 6 failed, 71 passed on base with the PR tests, and 77 passed on the PR branch. CPU only.

Multi-model performance (MI355X, TP1)

Same image for both arms; only python/ differs: base is fbc90a860e, PR is ae64753ca5. Server flags are --tp 1 --chunked-prefill-size 8192. kv-pressure adds --max-total-tokens 40000, which forces chunk continuation, batch_is_full, and eviction on every step. Each arm was run twice in alternating order (base, pr, base, pr) on a fresh server:

python -m sglang.bench_serving --backend sglang --dataset-name random \
  --random-input-len 3000 --random-output-len 300 --random-range-ratio 0.5 \
  --num-prompts 320 --max-concurrency 64 --seed 1

Model	Config	Arm	Completed	Output tok/s	Mean TTFT (ms)	P99 TTFT (ms)	Mean TPOT (ms)
Qwen3-8B	default	base	320 / 320	3386	256	1733	17.1
Qwen3-8B	default	PR	320 / 320	3397	255	1723	17.1
Qwen3-8B	kv-pressure	base	320 / 320	2039	4928	6121	7.2
Qwen3-8B	kv-pressure	PR	320 / 320	2037	4926	6110	7.2
Qwen3-30B-A3B	default	base	320 / 320	2659	259	1607	22.0
Qwen3-30B-A3B	default	PR	320 / 320	2612	262	1617	22.4
Qwen3-30B-A3B	kv-pressure	base	320 / 320	1441	6958	8705	10.2
Qwen3-30B-A3B	kv-pressure	PR	320 / 320	1435	6978	8692	10.3
DeepSeek-V2-Lite	default	base	320 / 320	3134	263	1464	18.4
DeepSeek-V2-Lite	default	PR	320 / 320	3111	264	1449	18.6
DeepSeek-V2-Lite	kv-pressure	base	320 / 320	1526	6605	8228	9.6
DeepSeek-V2-Lite	kv-pressure	PR	320 / 320	1522	6619	8246	9.6

Values are means of the two runs. One exception: the first base run of each model in the default config was the container's first server start (cold JIT, for example 980 tok/s on Qwen3-8B), so the base default rows use only the second base run. Raw values for every run:

All runs (output tok/s, mean TTFT ms)

Model	Config	base r1	PR r1	base r2	PR r2
Qwen3-8B	default	980 / 443 (cold)	3402 / 255	3386 / 256	3391 / 255
Qwen3-8B	kv-pressure	2037 / 4930	2041 / 4916	2040 / 4926	2033 / 4936
Qwen3-30B-A3B	default	2441 / 284 (cold)	2619 / 258	2659 / 259	2605 / 266
Qwen3-30B-A3B	kv-pressure	1442 / 6950	1442 / 6945	1439 / 6965	1428 / 7011
DeepSeek-V2-Lite	default	2993 / 430 (cold)	3106 / 261	3134 / 263	3116 / 266
DeepSeek-V2-Lite	kv-pressure	1522 / 6624	1523 / 6616	1530 / 6586	1521 / 6622

Every request completed in all 24 runs. Differences are within ±1% on every row except Qwen3-30B-A3B default, which is −1.8% against a single warm base run; that is inside the base r1-to-r2 spread on the same row. DeepSeek-V2-Lite server logs had 2–4 lines matching out of memory|Traceback|assert in both arms, while the Qwen logs had none. No run crashed, failed a request, or emitted an OOM.