feat(scheduler): LoRA adapter identity, max_loras batch cap, and per-adapter KV namespace - #976
feat(scheduler): LoRA adapter identity, max_loras batch cap, and per-adapter KV namespace#976towillwu wants to merge 5 commits into
Conversation
336066a to
b032688
Compare
b032688 to
2c9a86b
Compare
2c9a86b to
9fe157d
Compare
Introduce the scheduler's notion of a LoRA adapter and cap how many
distinct adapters may appear in a single batch.
Request identity:
- RequestSpec::lora_id carries the front-end's resolved adapter id, or
empty for a base-model request. The scheduler treats it as opaque: it
only compares it for equality, so the front-end is free to change how
it mints the value.
- Surfaced as Request::LoraId().
max_loras batch cap:
- SchedulerConfig::max_loras (0 = disabled) bounds the number of unique
adapters in a batch, because the runtime materializes one adapter slot
per unique id and exceeding the cap would overrun that fixed pool.
- buildForwardOperations tracks the adapters already pushed and defers a
request that would introduce a new one past the cap. It defers rather
than breaking out of the loop: a later candidate may reuse an adapter
already in the batch and still belongs in this step. A request whose
adapter is already present costs no additional slot, and base-model
requests are never charged.
- max_loras == 0 means LoRA scheduling is off, not "admit zero
adapters" -- the cap is simply not applied, so the scheduler behaves
exactly as it did before this change.
Exposed through the nanobind binding, make_config/make_spec, and a new
--max-loras server flag. request_handler reads lora_id defensively via
getattr: the adapter registry that resolves and attaches it to the request
input lands separately, and until it does every request is base-model.
Tests: tests/cpp/test_max_loras.cpp covers the cap, adapter reuse,
base-model requests, and the disabled case;
test/runtime/test_lora_scheduler_plumbing.py pins the None-to-empty-string
translation at the Python boundary.
Signed-off-by: towillwu <toqywu@gmail.com>
Without isolation, tokens cached under adapter A can be served to adapter B or to the base model, because a page hash is derived only from its tokens and the chain before it. The page hasher already accepts per-page extra_keys for exactly this purpose. Add ComputeNamespacedPagedHashes, which applies one key set to every page rather than a per-page list, and thread the request's adapter id through it. An empty key list is byte-identical to the previous two-argument form, so base-model requests hash exactly as before and existing cache entries stay valid. Keying every page, rather than only page 0: Chaining alone would make a page-0 key sufficient, since every later page folds in prior_hash. But appendCompletedPageHashes starts a fresh chain with an empty prior whenever page_hashes is empty, and that happens for a prompt shorter than one page: such a request produces no hashes at admission, so the first hash it ever computes comes from the incremental path. With a page-0-only key that page would land in the base-model namespace. Keying every page makes isolation independent of where the chain starts. AdvancePagedHashes previously neither accepted nor forwarded extra keys, so it silently dropped them; it now takes namespace_keys and passes them down. All three hash sites supply the request's keys: admission (matchPrefixAtAdmission), incremental extension (appendCompletedPageHashes), and stable-page publication (publishCompletedPages). Request::CacheNamespaceKeys() holds the adapter id as a 0-or-1 element vector so it can be handed over as a span without materializing a temporary on every hash call. Tests: tests/cpp/test_lora_namespace.cpp pins that empty keys are byte-identical to the unnamespaced chain, that adapters diverge from each other and from the base model, that an incremental extension matches the one-shot chain, and that a chain starting mid-request -- or at page 0 via the incremental path -- stays isolated. Signed-off-by: towillwu <toqywu@gmail.com>
Switch RequestSpec::lora_id from std::string to std::int32_t, with 0 meaning base model instead of an empty string. This matches the wire format the LoRA front-end already uses. That runtime keeps three distinct identifiers -- adapter name (user-facing), an integer registry id (stable, minted by a counter), and a GPU buffer slot (recycled on evict) -- and it is the integer registry id that travels on the request. Carrying a string here would have forced the front-end to collapse that layer, so the scheduler adopts the int instead. The id stays opaque to the scheduler: it compares ids for equality for the max_loras cap and renders one to a string for the page hasher's extra_keys. Rendering happens once per request into lora_hash_keys_, so CacheNamespaceKeys() still hands out a span without formatting on every hash call. request_handler reads the id with getattr rather than directly. TokenizedGenerateReqInput is a positionally-encoded msgspec Struct (array_like=True) whose docstring states that declaration order is the wire contract and that fields are append-only; adding lora_id there changes the tuple the frontend decodes (test_zmq_msgpack.py pins the arity for exactly this reason). That contract belongs to the front-end, so the field is declared by the LoRA front-end change that also populates it, and this branch only reads it if present. Base-model requests are unaffected either way: id 0 produces an empty key list, so their page hashes stay byte-identical to a build without LoRA support, and the max_loras cap is not applied at all when max_loras is 0. Signed-off-by: towillwu <toqywu@gmail.com>
0790734 to
f12e16b
Compare
f12e16b to
9640571
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f12e16b5f2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| std::vector<std::span<const std::int32_t>> paged_tokens = request->FullPagedTokens(false); | ||
| paged_tokens.resize(std::min(paged_tokens.size(), static_cast<std::size_t>(cacheable_pages))); | ||
| std::vector<std::string> hashes = ComputePagedHashes(paged_tokens, ""); | ||
| std::vector<std::string> hashes = ComputeNamespacedPagedHashes(paged_tokens, "", request->CacheNamespaceKeys()); |
There was a problem hiding this comment.
When KV-cache events are enabled and two LoRA adapters process the same token prefix, this creates distinct internal cache keys, but registerKvEventPages in scheduler.cpp still derives the published block_hashes solely from tokens and parent hashes via HashKvBlock. The adapters therefore emit indistinguishable store/removal events, so cache-aware consumers can incorrectly treat one adapter's incompatible KV block as the other's or remove availability while the other block remains; include the adapter namespace in the event hash lineage as well.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch — fixed in b6b5c933.
The two hash lineages had diverged: cache keys were namespaced, published block_hashes were not. HashKvBlock now takes the same namespace keys as the page hash and registerKvEventPages passes the request's. Keys are mixed in as a terminal length-prefixed block, so an empty list is skipped and base-model event digests stay byte-identical.
Three tests added: byte-identity when unkeyed, divergence between adapters and vs. base, and that a first-block divergence propagates through the parent chain.
| "--max-loras", | ||
| metavar="MAX_LORAS", | ||
| type=int, | ||
| default=ServerArgs.max_loras, | ||
| help=( | ||
| "Maximum number of distinct LoRA adapters in a single batch. " | ||
| "Requests that would exceed it are deferred to a later step. " | ||
| "0 disables LoRA scheduling." |
There was a problem hiding this comment.
The user-facing --max-loras option is added without updating the scheduler option table in docs/configuration/server.md or the compatible-parameters documentation, leaving its distinct-adapter semantics and zero sentinel absent from the repository's configuration guide. Add the new option to the relevant documentation as required for changed code.
AGENTS.md reference: AGENTS.md:L13-L13
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 352b19b4.
docs/configuration/server.md: added to the scheduler option table, covering the distinct-adapter semantics, that reusing an in-batch adapter never defers a request, that base-model requests are not charged, and that the0default disables LoRA scheduling rather than admitting zero adapters.docs/configuration/compatible-parameters.md: listed under "Similar But Not Identical", since the name matches other serving stacks but the zero sentinel does not.
Addresses a Codex review finding on #976. The prefix cache's page hashes were namespaced per adapter, but the KV cache event stream carries its own block hash lineage: registerKvEventPages derived block_hashes from HashKvBlock(tokens, parent) alone. Two adapters sharing a token prefix therefore published identical block hashes even though their cache keys differed, so a cache-aware consumer could treat one adapter's KV block as the other's -- or drop availability for one while the other's block is still resident. That is the same cross-adapter confusion this branch exists to prevent, just relocated from the local cache to the event consumers. HashKvBlock now takes the same namespace keys the page hash uses, and registerKvEventPages passes the request's. The keys are mixed in as a terminal, length-prefixed block, so an empty list is skipped outright and leaves the digest byte-identical to the un-namespaced form -- base-model event streams are unchanged. Tests pin the byte-identity for an empty key list, divergence between two adapters and against the base model, and that a divergence in the first block keeps later blocks apart through the parent chain. Signed-off-by: towillwu <toqywu@gmail.com>
Addresses a Codex review finding on #976; AGENTS.md requires docs updates for changed code. - configuration/server.md: adds --max-loras to the scheduler option table, covering the distinct-adapter semantics, that reusing an in-batch adapter never defers a request, that base-model requests are not charged, and that the 0 default disables LoRA scheduling rather than admitting zero adapters. - configuration/compatible-parameters.md: lists it under "Similar But Not Identical", since the name matches other serving stacks but the zero sentinel does not. Signed-off-by: towillwu <toqywu@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 352b19b415
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| spec.max_new_tokens = max_new_tokens | ||
| # 0 is the base model on both sides of the boundary; tolerate None from | ||
| # callers that have not resolved an adapter. | ||
| spec.lora_id = lora_id or 0 |
There was a problem hiding this comment.
When a LoRA-aware request reaches the normal engine path, RequestHandler.handle_generate_request still calls make_spec with only rid and tokens, so this new field always takes the default and every scheduler Request is created with lora_id == 0. That means real adapter requests are treated as base-model requests: per-adapter prefix-cache namespaces collapse together and the max_loras cap never counts those adapters. Pass the resolved request adapter id, e.g. via getattr(recv_req, "lora_id", 0), into this argument before constructing the scheduler spec.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Not applicable to this PR — intentionally.
lora_id is not a field on TokenizedGenerateReqInput on this branch ('lora_id' in TokenizedGenerateReqInput.__struct_fields__ is False), so the suggested getattr(recv_req, "lora_id", 0) can only ever return 0, which is already make_spec's default. It would be dead code.
The field cannot be added here: that struct is positionally encoded (array_like=True) and its docstring states declaration order is the IPC wire contract with the frontend. test_zmq_msgpack.py pins the arity for exactly this reason and fails on the change (24 → 25). That contract belongs to the front-end.
So the field is declared and forwarded together in the next PR of this stack (#980):
io_struct.py—lora_id: int = 0onTokenizedGenerateReqInputrequest_handler.py—lora_id=recv_req.lora_id
Verified there: a base-model request yields RequestSpec.lora_id == 0, an adapter-7 request yields 7 — so per-adapter namespaces do separate and max_loras does count them.
This PR only provides the scheduler mechanism; with no adapter ids in flight, every request is base-model and scheduling is unchanged, which is the intended state until #980 lands.
1 participant
Makes the tokenspeed scheduler LoRA-aware. Scheduler only — no kernels, no adapter runtime, no model code.
Two independent changes, one commit each, plus a small wire-format change.
1 — Per-adapter KV prefix-cache namespace
Without isolation, tokens cached under adapter A can be served to adapter B or to the base model: a page hash derives only from its tokens and the chain before it, so identical prompt prefixes collide across adapters.
page_hasher.halready accepted per-pageextra_keys— documented upstream for exactly this ("e.g. LoRA name"). This addsComputeNamespacedPagedHashes, which applies one key set to every page rather than a per-page list, and threads the request's adapter id through it.Why every page, not just page 0. Chaining would make a page-0 key sufficient — each later page folds in
prior_hash. ButappendCompletedPageHashesstarts a fresh chain with an empty prior wheneverpage_hashesis empty, which happens for any prompt shorter than one page: such a request produces no hashes at admission, so the first hash it ever computes comes from the incremental path. With a page-0-only key that page would silently land in the base-model namespace. Keying every page makes isolation independent of where the chain starts.AdvancePagedHashessilently dropped keys. It neither accepted nor forwardedextra_keys, so two of the three hash sites would have lost the adapter id on every extension. It now takesnamespace_keys. All three sites pass the request's keys: admission (matchPrefixAtAdmission), incremental extension (appendCompletedPageHashes), and stable-page publication (publishCompletedPages).2 —
max_lorasbatch capSchedulerConfig::max_loras(0 = disabled) bounds the number of distinct adapters in one batch, because the runtime materializes one adapter slot per unique id and exceeding the cap would overrun that fixed pool.buildForwardOperationstracks the adapters already pushed and defers a request that would introduce a new one past the cap — deferring rather than breaking out of the loop, because a later candidate may reuse an adapter already in the batch and still belongs in this step. A request whose adapter is already present costs no additional slot; base-model requests are never charged.Existing functionality is unchanged
This is the property the PR is built around, and it is asserted rather than argued:
NoKeysIsByteIdenticalToUnnamespacedChainasserts it directly, so existing cache entries stay valid and enabling LoRA invalidates nothing.max_loras == 0means LoRA scheduling is off, not "admit zero adapters" — the cap is simply not applied.ZeroMaxLorasEnforcesNoCappins it.RequestSpec::lora_iddefaults to 0 (base model) andmax_lorasto 0, so a build that never sets either behaves exactly as before.request_handlerreads the id withgetattr(..., 0), so it cannot raise on a request input that has no such field.On the IPC wire format
request_handlerreadslora_iddefensively rather than directly, and this PR deliberately does not add the field toTokenizedGenerateReqInput.That struct is a positionally-encoded
msgspec.Struct(array_like=True) whose own docstring states that declaration order IS the wire contract — append only. Adding a field changes the tuple the frontend decodes;test_zmq_msgpack.pypins the arity precisely to catch that, and inserting anywhere but the end also shifts existing fields' positions. An earlier revision of this branch did exactly that and the test caught it.That contract belongs to the front-end, so the field is declared by the LoRA front-end change that also populates it. Until then every request resolves to
lora_id = 0and scheduling is byte-for-byte what it was before this PR.Relationship to #735
Supersedes #735, which was written against the radix prefix cache that #864 removed. Every core file that PR touched —
csrc/resource/kv_prefix_cache/,hybrid_prefix_cache/,radix_tree/,resource/types.h— is gone frommain, andtests/check_no_radix_cache_sources.pynow hard-fails if those directories or the symbolsRadixTree/KVPrefixCache/HybridPrefixCachereappear. This is a rewrite against the hash-chain cache, not a rebase.Testing
test_lora_namespace.cppandtest_max_loras.cpp.test_io_struct_codec.py,test_zmq_msgpack.py,test_max_new_tokens_context_cap.py) — the wire-format guard.check_no_radix_cache_sources.pyguard passes.pre-commit,ruff(CI's select list),clang-format,clang-tidy(0 diagnostics): clean.No GPU path is involved in this PR.