Skip to content

feat(lora): prefix cache adapter isolation (PR 3/3) - #2047

Open
cchh05 wants to merge 4 commits into
xLLM-AI:mainfrom
cchh05:pr3-lora-prefix-cache-isolation
Open

feat(lora): prefix cache adapter isolation (PR 3/3)#2047
cchh05 wants to merge 4 commits into
xLLM-AI:mainfrom
cchh05:pr3-lora-prefix-cache-isolation

Conversation

@cchh05

@cchh05 cchh05 commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

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_lora and --enable_prefix_cache are 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_id used by xxh3_128bits_hash to seed the first block hash of a chain:

  • BlockManagerPool::allocate_shared and cache set the TLS before calling into CompositeBlockManager, then reset it to 0 after the call.
  • xxh3_128bits_hash reads the TLS when pre_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.
  • When the id is 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)

  • `xllm/core/framework/prefix_cache/block_hasher.h` — declaration of `set_prefix_cache_adapter_id`
  • `xllm/core/framework/prefix_cache/block_hasher.cpp` — TLS variable + setter + adapter_id-seeded hashing path
  • `xllm/core/framework/block/block_manager_pool.cpp` — set/reset around the three prefix-cache entry points (`allocate_shared`, `cache(Sequence*)`, `cache(Sequence*, size_t)`)

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:

Test Result
R1 == R2 sha256 identical for all 4 models ✅ PASS
Base != deepnews / uncensored sha256 ✅ PASS
3 unique sha256 across base/deepnews/uncensored ✅ PASS
dummy (r=8) sha256 == base expected — very-low-rank delta too weak to change argmax at temperature=0

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_hash when 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

  • Unit-level correctness: R1 == R2 sha256 identical across 4 adapters (see Verification)
  • Cross-adapter isolation: 3 unique sha256
  • Stress test: 60 requests × concurrency=4, no crash
  • Add prefix cache unit test with adapter isolation assertion (follow-up)
  • Add `--enable_lora_prefix_cache_isolation` CLI kill-switch for canary rollback (follow-up)

🤖 Generated with Claude Code

cchh05 and others added 4 commits July 30, 2026 17:26
…/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.
…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
cchh05 force-pushed the pr3-lora-prefix-cache-isolation branch from 2fe2d30 to cbbfdf3 Compare July 30, 2026 09:33
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