feat(lora): multi-tenant LoRA framework + attention/MLP wire-up (1/2)#2014
Open
cchh05 wants to merge 1 commit into
Open
feat(lora): multi-tenant LoRA framework + attention/MLP wire-up (1/2)#2014cchh05 wants to merge 1 commit into
cchh05 wants to merge 1 commit into
Conversation
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).
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:07
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. |
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 details.
This is PR 1 of 2 landing multi-tenant LoRA serving on xllm.
This PR (~28 files, ~4300 lines)
xllm/core/framework/lora/— new subsystem: runtime, registry, loader, context, metrics, configxllm/core/layers/common/lora/— 3 composition-based Linear wrappers (QKV/Column/Row parallel)xllm/core/framework/model/model_input_params.h— addsadapter_ids+adapter_ids_per_tokenfieldsxllm/models/llm/llm_model_base.h— pushes LoRAContext frame in forwardxllm/core/layers/common/qwen2_attention.{h,cpp},dense_mlp.{h,cpp}— swaps Linear members to LoRA-wrapped drop-insPR 2/2 will follow with
/v1/load_lora_adapter,/v1/unload_lora_adapter,/v1/lora_adapters,/v1/lora_statslora_namefield)Not in either PR (separate CLs)
Validation
End-to-end verified on the pre-split fork branch: