Skip to content

feat(scheduler): per-adapter KV prefix-cache namespace + max_loras batch cap#735

Open
qywu wants to merge 1 commit into
lightseekorg:mainfrom
qywu:qywu/lora-scheduler-namespace
Open

feat(scheduler): per-adapter KV prefix-cache namespace + max_loras batch cap#735
qywu wants to merge 1 commit into
lightseekorg:mainfrom
qywu:qywu/lora-scheduler-namespace

Conversation

@qywu

@qywu qywu commented Jul 19, 2026

Copy link
Copy Markdown
Collaborator

Summary

Two scheduler-level changes that make the tokenspeed scheduler LoRA-aware. Rebased onto current main from 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_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 (unchanged behavior when LoRA is off). The id is plumbed from Request::LoraId() through every FSM event that publishes or matches KV pages: SchedulePrefillEvent, ScheduleDecodeEvent, FinishEvent, ExtendResultEvent, and the InsertHybridCache helper.

2 — max_loras batch cap

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.


Rebase / merge notes

This reconciles the original change with several things main grew since #268:

  • 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 (would otherwise be a silent cross-adapter cache hole).

Testing

  • All changed scheduler translation units type-check against current main headers (via compile_commands.json, -fsyntax-only).
  • clang-format and pre-commit run --all-files are clean.
  • ⚠️ A full C++ build / gtest run was not possible in this environment (missing tokenspeed-spdlog at CMake-configure time). Recommend running the scheduler C++ test suite in CI before merge.

🤖 Generated with Claude Code

…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>
@qywu
qywu requested a review from a team as a code owner July 19, 2026 23:57
@qywu
qywu requested a review from wangbo981016 July 23, 2026 18:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant