feat(lora): prefix cache adapter isolation (PR 3/3) - #2047
Open
cchh05 wants to merge 4 commits into
Open
Conversation
cchh05
requested review from
Clement-Wang26,
DongheJin,
DragonFive,
JimHsiung,
Kang-Meng,
RobbieLeung,
XuZhang99,
liujinguang0125,
liutongxuan,
walsonyang,
xiao-yu-chen,
yingxudeng,
yq33victor and
zhang-minchao
as code owners
July 27, 2026 08:59
cchh05
force-pushed
the
pr3-lora-prefix-cache-isolation
branch
from
July 27, 2026 10:36
41521fa to
2fe2d30
Compare
…/2).
Introduces the plumbing for multi-tenant LoRA serving on xllm:
composition-based Linear wrappers that ride the existing attention /
MLP forward paths and pick up per-batch adapter routing through a
thread-local LoRA context.
This PR is the first half of a two-PR split. It lands the framework +
attention/MLP wire-up so the wrapper code path exists in the tree.
PR 2/2 will wire the request path (RequestParams -> Sequence ->
BatchInputBuilder -> ModelInputParams.adapter_ids) and the HTTP
endpoints (/v1/load_lora_adapter, /v1/unload_lora_adapter,
/v1/lora_adapters, /v1/lora_stats).
New subsystems
--------------
xllm/core/framework/lora/
* LoRARuntime — process singleton: adapter loader, per-proj device
pool, hot-swap executor thread pinned to the model device (needed
for CANN 8.5's forward-thread-only CPU->NPU copy restriction),
per-request per-projection delta lookup.
* LoRARegistry — name<->int_id map, pin/unpin lifecycle so
/v1/unload_lora_adapter can drain in-flight requests gracefully.
* LoRAAdapterLoader — PEFT (adapter_config.json + safetensors) reader
with target_modules whitelist and a canonical key parser
(base_model.model.*.layers.<L>.<module>.lora_A|B).
* LoRAContext — thread-local frame holding pointers into
ModelInputParams.adapter_ids / adapter_ids_per_token so LoRA-
wrapped Linear layers can find the per-seq / per-token routing
without changing Linear::forward's signature.
* LoRAMetrics — bvar Prometheus counters per adapter (TTFT / e2e /
tokens generated / errors / QPS).
* lora_config — gflags: enable_lora, max_loras, max_lora_rank,
lora_target_modules whitelist, lora_modules static preload,
allow_runtime_lora_updating,
enable_lora_row_parallel_all_reduce (defaults on; correctness),
enable_lora_row_parallel_fused_ar (defaults off; a per-layer
optimisation that fuses the LoRA delta into the base's own row-
parallel all-reduce, S-LoRA / vLLM RowParallelLinearWithShardedLoRA
pattern).
xllm/core/layers/common/lora/
* LoRAQKVParallelLinear — composition wrapper around
QKVParallelLinear. Fast path shares one shrink+expand across the
whole single-adapter batch; slow path (mixed base + adapter batch)
walks per-seq. Correctly handles GQA replica sharding of the k/v
LoRA-B slices — B.size(0) is num_kv_heads * head_dim (not tp_size
times that) so the shard-index derivation must divide by the shard
count in B, not tp_world_size. Also accepts a q_has_gate parameter
for Qwen3-Next-style attn_output_gate where q + gate are fused into
the q lane.
* LoRAColumnParallelLinear — wrapper for gate_up_proj (fused gate/up)
and any future column-parallel projection.
* LoRARowParallelLinear — wrapper for o_proj / down_proj. Ships two
TP-correct paths: (a) explicit rank-dim all-reduce of the shrink
output before the expand, (b) fused mode where the base row-
parallel's own all-reduce covers both the base output and the
LoRA delta partial-sum. Fused mode reclaims ~3.6 pp single-adapter
and ~8.3 pp mixed-batch throughput on Ascend NPU relative to the
naive rank-dim AR path, at the cost of holding LoRA-B replicated
across TP ranks (cheap because r is small).
Wire-up
-------
xllm/core/framework/model/model_input_params.h
* Adds ModelInputParams::adapter_ids (one entry per sequence in the
batch, index-aligned with attention.host.q_seq_lens) and
adapter_ids_per_token (device tensor built lazily in to(device) via
repeat_interleave from adapter_ids and q_seq_lens). Empty adapter_
ids is a no-op — a pure-base batch does not touch either field.
xllm/models/llm/llm_model_base.h
* Pushes a LoRAContextFrame at the top of forward, calls
set_lora_context_layer on each layer of the decoder loop. Cost
when LoRA is off is one atomic pointer copy — the wrappers early-
return whenever the frame is null.
xllm/core/layers/common/qwen2_attention.{h,cpp},
xllm/core/layers/common/dense_mlp.{h,cpp}
* Swaps QKVParallelLinear / RowParallelLinear / ColumnParallelLinear
member types for their LoRA-wrapped drop-in equivalents. Base
checkpoint keys (qkv_proj.weight, o_proj.weight, gate_up_proj.
weight, down_proj.weight) are unchanged because the base linear
is held as a plain member inside the wrapper (not register_module'd
on the wrapper), so no checkpoint compat break.
Not in this PR
--------------
* Qwen3-Next hybrid attention (Qwen3.5-122B) wire-up and 122B model
install — separate CL, more model-specific work.
* Fused MoE expert LoRA delta (Phase 1+2 grouped-gemm injection into
fused_moe.cpp) — separate CL, non-trivial MoE-side refactor.
* Request path: sequence.adapter_id propagation from RequestParams
through BatchInputBuilder to ModelInputParams.adapter_ids — PR 2/2.
* HTTP endpoints /v1/load_lora_adapter, /v1/unload_lora_adapter,
/v1/lora_adapters, /v1/lora_stats, and the two-field routing in
ChatServiceImpl — PR 2/2.
Validation (from prior fork builds)
-----------------------------------
* End-to-end verified on Qwen3-30B-A3B-Instruct-2507 TP=8 with a real
Megatron-trained PEFT LoRA (r=16, alpha=32, target={q,k,v,o}_proj):
200-sample business eval, +9pp compliance-rate and +0.04 gt Jaccard
vs base, base output pixel-perfect unchanged (fix is non-invasive).
* 30-min sustained load stress: 5363 requests, err=0.00%,
HBM +3 MB drift.
* Divergence sweep N=50 with a public Qwen3.5-122B LoRA: 70% diverge
rate under temperature=0 (systematic delta, not noise).
Wire the LoRA adapter_id from ChatServiceImpl all the way to ModelInputParams.adapter_ids so the LoRA-wrapped Linear layers landed in PR 1/2 can pick up the correct per-sequence routing at forward time. Depends on PR xLLM-AI#2014 (framework + attention/MLP wire-up) for: - `ModelInputParams::adapter_ids` / `adapter_ids_per_token` fields - `LoRARuntime` singleton + `LoRARegistry` - `LoRAContextFrame` push in `LlmModelImplBase::forward` Chain (rank0 path) ------------------ ``` ChatServiceImpl::process_async_impl | | lookup_and_pin(model) // model = adapter name | -> lora_pinned->int_id | v RequestParams { adapter_id, adapter_name } | v LLMMaster::generate_request -> RequestState { adapter_id } | v Request::init -> SequenceParams { adapter_id } | v Sequence { adapter_id_ } | v (per forward pass) BatchInputBuilder::build_forward_input -> BuilderState { adapter_ids } -> input_params.adapter_ids = state_.adapter_ids | v ModelInputParams::to(device) -> adapter_ids_per_token = repeat_interleave(adapter_ids, q_seq_lens) | v LlmModelImplBase::forward -> LoRAContextFrame captures &input_params.adapter_ids | v LoRA-wrapped Linear (in PR xLLM-AI#2014): reads current_lora_context() ``` Files touched ------------- Request / batch path: - `framework/request/request_params.h` — adds `std::optional<uint64_t> adapter_id` + `std::string adapter_name` - `framework/request/request_state.h` — adds `uint64_t adapter_id = 0` - `framework/request/sequence.h` — adds `SequenceParams::adapter_id`, `Sequence::adapter_id()`, member - `framework/request/sequence.cpp` — ctor init from `seq_params.adapter_id` - `framework/request/request.cpp` — `sequence_params.adapter_id = state_.adapter_id` - `framework/batch/batch_input_builder.h` — adds `BuilderState::adapter_ids` - `framework/batch/batch_input_builder.cpp` — push_back per seq + per-thread state merge + `input_params.adapter_ids = std::move(...)` - `runtime/forward_params.h` — adds `ForwardInput::adapter_ids` Distributed runtime: - `distributed_runtime/llm_master.cpp` — `if (sp.adapter_id.has_value()) req_state.adapter_id = sp.adapter_id.value();` Chat routing: - `api_service/chat_service_impl.cpp` — look up `model` against `LoRARegistry` via `lookup_and_pin`; if the name resolves to a registered adapter, populate `request_params.adapter_id` / `adapter_name`. Pure-base requests keep both unset and the whole chain no-ops. Not in this PR (follow-up CLs) ------------------------------ - HTTP endpoints `/v1/load_lora_adapter`, `/v1/unload_lora_adapter`, `/v1/lora_adapters`, `/v1/lora_stats`. Requires `master->load_lora_broadcast()` and `unload_lora_broadcast()` engine APIs that don't exist upstream yet — separate PR. - Non-rank0 broadcast of `adapter_ids` via shm channel. Rank0 sees the correct `adapter_ids` today; non-rank0 workers will need proto and `params_utils.cpp` / `forward_shared_memory_manager.cpp` changes to receive them. In practice for TP=1 this PR is already end-to-end usable; TP>1 will silent-fallback to base on non-rank0 workers until the broadcast lands. - Full drain lifecycle: unpin adapter on chat request finish (both success and error). Correct behaviour today is best-effort because unload endpoint is a follow-up. Rollout ------- Depends on PR xLLM-AI#2014 landing first. Static adapter preload via `--lora_modules=<name>=<path>,<name>=<path>` (existing gflag in PR broadcast follows in later CLs.
…LU/MLU/CUDA/DCU builds see it.
…d hash seeding (PR 3/3). When --enable_lora and --enable_prefix_cache are both enabled, identical prompts across different LoRA adapters (and base) would previously hit the same prefix cache entry. This means adapter B's request could reuse adapter A's cached KV blocks, producing silent semantic errors: the output looks like a valid response but uses the wrong adapter's attention computation. This PR adds adapter-aware hash isolation: 1. A thread-local `g_prefix_cache_adapter_id` in block_hasher.cpp, set by BlockManagerPool::allocate_shared / cache before invoking any PrefixCache path, and reset to 0 after. 2. When `g_prefix_cache_adapter_id != 0` at first-block hashing time, the adapter_id is prepended to the token id buffer and included in the xxh3 seed. This yields disjoint hash chains for identical prompts served by different adapters. 3. When `g_prefix_cache_adapter_id == 0` (base model, no LoRA), the hash is byte-identical to the pre-patch baseline. Base-only deployments are completely unaffected. Depends on adapter_id propagation from PR xLLM-AI#2015. Verification on 4 adapters (Qwen3-32B base + deepnews r=32 + uncensored r=32 + dummy r=8), same 500-token prompt: - R1 == R2 sha256 identical for all 4 (cache hit preserves semantics) - 3 unique sha256 across base/deepnews/uncensored (isolation works) - dummy sha256 == base sha256 (r=8 delta too weak to change argmax, expected behavior for very-low-rank adapters) Co-Authored-By: Claude <noreply@anthropic.com>
cchh05
force-pushed
the
pr3-lora-prefix-cache-isolation
branch
from
July 30, 2026 09:33
2fe2d30 to
cbbfdf3
Compare
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
Adds LoRA-aware isolation to xllm's prefix cache so that requests to different LoRA adapters (and base) do not cross-contaminate cached KV blocks.
Problem
When
--enable_loraand--enable_prefix_cacheare both enabled, requests with identical prompts served by different adapters produce the same block-level hash (because hashing keys on token ids alone). The prefix cache then hands adapter B's request the KV blocks cached under adapter A. The result is silent semantic corruption — the output still looks like valid natural language, but the attention computation used adapter A's KV values instead of recomputing with adapter B's LoRA delta applied.Solution
Introduce a thread-local
g_prefix_cache_adapter_idused byxxh3_128bits_hashto seed the first block hash of a chain:BlockManagerPool::allocate_sharedandcacheset the TLS before calling intoCompositeBlockManager, then reset it to0after the call.xxh3_128bits_hashreads the TLS whenpre_hash_value == nullptr(first block). If the id is non-zero, the id bytes are prepended to the token-id buffer before hashing, so identical prompts across adapters diverge at block 0.0(base model, no LoRA), hashing is byte-identical to the pre-patch baseline. Base-only services see no behavioral change.Depends on
#2015 — needs
Sequence::adapter_id()on the request path.Files changed (3 files, +52 lines)
Verification
Ascend NPU 910 · Qwen3-32B TP=4 · 4 adapters (Qwen3-32B base + deepnews r=32 + uncensored r=32 + dummy r=8) · identical 500-token prompt · 2 rounds each:
Under stress (60 requests × concurrency=4, shared 400-token prefix): 60/60 ok, service stable.
Base-only safety
The only code path that behaves differently is inside
xxh3_128bits_hashwhen the TLS is non-zero. In base-only deployments (no LoRA,--enable_lora=false),sequence->adapter_id()is always 0, so the TLS is always set to 0 → the else-branch runs → hashing is byte-identical to pre-patch.Test plan
🤖 Generated with Claude Code