[ormandj]

Motivation

A hybrid SSM server can leave requests queued even after Mamba-state capacity becomes available. Once admission marks the running batch as full, the scheduler does not retry while that batch remains unchanged. Cached states can become evictable during that time, so a previous capacity failure is not a valid reason to stop checking.

Retrying exposes a second problem. Each admission attempt re-runs the prefix match, and a match that hits a cached Mamba state allocates a per-request slot for the copy-on-write before PrefillAdder checks capacity. If no slot can be freed (for example, decode has consumed a shared Mamba/attention pool since the earlier refusal and the remaining cached states are locked), that allocation asserted and terminated the scheduler instead of leaving the request queued.

This PR retries admission for hybrid SSM models whose cache supports Mamba state, and makes a failed copy-on-write allocation an ordinary admission refusal. It does not bypass capacity checks or enlarge the pool.

Modifications

• Clear the full-batch admission latch each scheduler round for that configuration, alongside the existing hybrid-SWA and priority-preemption resets. Other configurations retain their existing behavior.
• When the Mamba copy-on-write slot cannot be allocated even after eviction, raise a dedicated MambaCowAllocError instead of asserting. The matched node's lock is still released and the request is left without a slot.
• Prefill admission catches that error and treats it as a NO_TOKEN refusal: the request stays in the waiting queue and is retried in a later round. Other callers of the prefix match are unchanged apart from the exception type.

Accuracy Tests

All tests below were run with real pytest 9.1.1 via uv run in a CPU-only PyTorch 2.14.1 environment (no GPU), from the repository root with PYTHONPATH=python:test. The two new files pass 8 of 8 with python -m pytest test/registered/unit/managers/test_scheduler_hybrid_ssm_admission_latch.py test/registered/unit/managers/test_scheduler_hybrid_ssm_admission_retry.py. These tests call scheduler and cache methods directly with stubbed collaborators; no model is loaded.

• Four latch tests check which configurations clear batch_is_full each round.
• A three-round scheduler test with an adder stand-in that always returns NO_TOKEN verifies admission is retried each round for hybrid SSM; the non-hybrid case attempts admission once.
• A component test drives MambaComponent.finalize_match_result_in_cache with an allocator that has no free slot and an eviction that frees nothing. It raises MambaCowAllocError with balanced lock calls and no slot assigned (the previous code raised AssertionError).
• A scripted scheduler test covers refusal and retry: the adder stand-in refuses the request in round 1, its stubbed init_next_round_input raises MambaCowAllocError in round 2, and in round 3 the stub succeeds and the adder stand-in accepts. Round 2 returns no batch and keeps the request queued without calling add_one_req (on the previous code the error propagated out of the scheduler); round 3 passes the request to the (mocked) ScheduleBatch.init_new and empties the waiting queue. No real allocator capacity is consumed or recovered in this test.

Existing related unit tests in the same environment: test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py, test_mamba_path_state_cap.py, test_mamba_unittest.py, test/registered/unit/managers/test_prefill_adder.py, test_schedule_policy.py and the 18 other test/registered/unit/managers/test_scheduler_*.py files give 197 passed, 5 skipped (the test_mamba_path_state_cap.py pool fixtures need CUDA) and 596 subtests passed. A successful copy-on-write allocation after real capacity recovery, with a real Mamba cache on GPU, has not been tested.

Speed Tests and Profiling

Adds an admission attempt each round while capacity remains unavailable. Overhead was not measured.

Checklist

• Format changed code.
• Run regression tests described above.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

Developed with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #38053460648
Latest PR Test (Extra): ❌ Run #38053460464
Latest PR Test (AMD ROCm 10): ❌ Run #38053460574

[ormandj]

Updated the code and PR body. The scheduler clears the full-batch admission latch each round for the affected Mamba configuration so queued prefills can retry when capacity recovers.