feat(scheduler): per-adapter KV prefix-cache namespace + max_loras batch cap#735
Open
qywu wants to merge 1 commit into
Open
feat(scheduler): per-adapter KV prefix-cache namespace + max_loras batch cap#735qywu wants to merge 1 commit into
qywu wants to merge 1 commit into
Conversation
…tch cap Make the tokenspeed scheduler LoRA-aware with two self-contained, scheduler-level changes (rebased from lightseekorg#268 onto current main). 1. Per-adapter KV prefix-cache namespace isolation - Thread a lora_id through KVPrefixCache::Insert/Match and the radix tree so each adapter gets an isolated namespace off a per-adapter root (lora_id == kLoraNone (0) -> shared base-model root). Prevents tokens cached under adapter A from being reused by adapter B or the base model. - Plumb lora_id from the request (Request::LoraId()) through the FSM events that publish or match KV pages: SchedulePrefillEvent, ScheduleDecodeEvent, FinishEvent, ExtendResultEvent, and the InsertHybridCache helper. 2. max_loras batch cap - Add SchedulerConfig::max_loras (0 = LoRA disabled). newForwardOperation defers requests that would push the unique-adapter count in a batch over the cap to the next step, guaranteeing prepare_loras() never sees n_unique > max_loras. Exposed through the Python scheduler config and nanobind binding. Merge notes (reconciling with main since the original PR): - FSM event constructors gained a conditional KvCacheCoordinator* (flat-kvcache) parameter upstream; lora_id is appended after it so both features compose in flat and non-flat builds. - InsertHybridCache was promoted to a header-declared helper with a prefix_pages_override parameter upstream; lora_id is appended as the trailing defaulted argument. - ExtendResultEvent is new since the original PR and also publishes to the KV prefix cache; wired it through with the request's lora_id to preserve the per-adapter isolation guarantee. All changed scheduler translation units type-check against current main headers; clang-format and pre-commit are clean. Signed-off-by: Qingyang Wu <willqywu@gmail.com>
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
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.
Summary
Two scheduler-level changes that make the tokenspeed scheduler LoRA-aware. Rebased onto current
mainfrom the previously-closed #268 (extracted from the broader LoRA serving PR #83).1 — Per-adapter KV prefix-cache namespace isolation
Problem: Without isolation, tokens cached under adapter A could be incorrectly reused by adapter B or the base model.
Fix: Thread a
lora_idthroughKVPrefixCache::Insert/Matchand the radix tree so each adapter gets an isolated namespace off a per-adapter root.lora_id == kLoraNone(0) → shared base-model root (unchanged behavior when LoRA is off). The id is plumbed fromRequest::LoraId()through every FSM event that publishes or matches KV pages:SchedulePrefillEvent,ScheduleDecodeEvent,FinishEvent,ExtendResultEvent, and theInsertHybridCachehelper.2 —
max_lorasbatch capSchedulerConfig::max_loras(0 = LoRA disabled).newForwardOperationdefers requests that would push the unique-adapter count in a batch over the cap to the next step, guaranteeingprepare_loras()never seesn_unique > max_loras. Exposed through the Python scheduler config and nanobind binding.Rebase / merge notes
This reconciles the original change with several things
maingrew since #268:KvCacheCoordinator*(flat-kvcache) parameter upstream;lora_idis appended after it so both features compose in flat and non-flat builds.InsertHybridCachewas promoted to a header-declared helper with aprefix_pages_overrideparameter upstream;lora_idis appended as the trailing defaulted argument.ExtendResultEventis new since the original PR and also publishes to the KV prefix cache — wired it through with the request'slora_idto preserve the per-adapter isolation guarantee (would otherwise be a silent cross-adapter cache hole).Testing
mainheaders (viacompile_commands.json,-fsyntax-only).clang-formatandpre-commit run --all-filesare clean.tokenspeed-spdlogat CMake-configure time). Recommend running the scheduler C++ test suite in CI before merge.🤖 Generated with Claude Code