Repository navigation
[Scheduler] Charge the host-hit Mamba node slot in prefill admission - #43138
Closed
alphabetc1 wants to merge 1 commit into
Closed
alphabetc1 wants to merge 1 commit into
alphabetc1 wants to merge 1 commit into
Conversation
alphabetc1
requested review from
Ying1123,
hnyls2002,
merrymercy and
xiezhq-hermann
as code owners
October 8, 2026 17:09
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
When a request's Mamba state hits on host (HiCache load-back, the external linker, LMCache),
req.mamba_host_hit_lengthis 1 and the load-back fills two Mamba slots: the node's and the request's own.PrefillAddercharged only the request's slot, andbudget_stateonly checks that at least one slot remains, so a host-hit request could be admitted with a single Mamba slot left.Modifications
_mamba_gap_budget_for_reqalso charges the host-hit node slot. It is charged only while the request holds no state yet: a chunked continuation keeps its stalemamba_host_hit_length, but its hit was already loaded._update_prefill_budgetdebitsrem_mamba_slotsby the slots charged, not by a fixed 1.add_one_reqreturnsNO_TOKENwhen the request needs more Mamba slots than remain.Tests
No new test. On H200:
TestUnifiedHybridHiCacheBitExact(Inkling, Mamba host hits): kl_divs 0.0 x 9.unit/mem_cacheHiCache suites,unit/dllm/test_gemma4_uniform_lifecycle.py,unit/managers) match main.test_mm_process_configfails on main too.test_load_snapshot_backends::test_read_returns_latestfailed once and passed on 3 reruns.CI States
Latest PR Test (Base): ❌ Run #37814327945
Latest PR Test (Extra): ❌ Run #37814327530
Latest PR Test (AMD ROCm 10): ❌ Run #37814329306