feat(lora): request path adapter_id propagation + chat routing (2/2)#2015
Open
cchh05 wants to merge 3 commits into
Open
feat(lora): request path adapter_id propagation + chat routing (2/2)#2015cchh05 wants to merge 3 commits into
cchh05 wants to merge 3 commits into
Conversation
added 2 commits
July 23, 2026 17:06
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 xLLM-AI#2014) works end-to-end with this CL for TP=1; dynamic hot-load / TP>1 broadcast follows in later CLs.
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 23, 2026 09:17
Collaborator
Author
|
Note on CI failures — this is an upstream CI infrastructure issue, not a code problem. All 4 The CI runner uses PR #2014 (dependency) is experiencing the exact same failure mode. Not planning to push new commits — will wait for the upstream proxy to recover, then request a re-run. If someone from the infra team can help investigate the proxy or bypass it (direct GitHub fetch), that would unblock both PRs. |
…LU/MLU/CUDA/DCU builds see it
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
1 participant
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.
See commit message for full plumbing chain diagram.
This is PR 2 of 2 landing multi-tenant LoRA serving on xllm.
Depends on #2014.
This PR (10 files, ~60 lines)
Wires the LoRA
adapter_idfromChatServiceImplall the way toModelInputParams.adapter_idsso the LoRA-wrapped Linear layers landedin PR #2014 can pick up per-sequence routing at forward time.
Request/batch path:
request_params.h,request_state.h,sequence.h/.cpp,request.cpp— carryadapter_idbatch_input_builder.h/.cpp— collect per-seqadapter_idsintoModelInputParamsforward_params.h— addsForwardInput::adapter_idsDistributed runtime:
llm_master.cpp—RequestParams::adapter_id→RequestState::adapter_idChat routing:
chat_service_impl.cpp—modelfield looked up inLoRARegistryvialookup_and_pin; if it resolves to a registered adapter, populaterequest_params.adapter_id.Not in this PR (follow-up CLs)
/v1/load_lora_adapter,/v1/unload_lora_adapter,/v1/lora_adapters,/v1/lora_stats— requires newmaster->load_lora_broadcast()/unload_lora_broadcast()engine APIsadapter_idsvia shm channel — for TP=1 this PR is already end-to-end usable; TP>1 will silent-fallback to base on non-rank0 workers until broadcast landsRollout
Depends on #2014 landing first. Static adapter preload via
--lora_modules=<name>=<path>(existing gflag from #2014) works end-to-end with this CL on TP=1; dynamic hot-load and TP>1 broadcast follow in later CLs.Fork validation
End-to-end verified on the pre-split fork branch (see #2014 body for numbers).