[LoRA] Support LoRA under DP attention: idle-forward guards, attn-TP-local slicing, per-rank tensor serialization - #31520
Closed
yushengsu-thu wants to merge 6 commits into
Closed
yushengsu-thu wants to merge 6 commits into
yushengsu-thu wants to merge 6 commits into
Conversation
…cal slicing, per-rank tensor serialization Skip LoRA application on DP-attention idle forwards via a None batch_info guard in every LoRA-wrapped layer, and clear stale batch info on idle batches. Slice LoRA A/B weights by the base layer's attn_tp-local rank and size the memory-pool buffers of attention modules by attn_tp_size. Relax the dp_size == 1 asserts for dynamic LoRA load/unload when DP attention is enabled, and ship load_lora_adapter_from_tensors payloads as one serialized copy per TP rank, mirroring update_weights_from_tensor. Derive CUDA-graph MoE LoRA buffer dims from the base FusedMoE so all quant schemes work. Ported from the sglang-miles branch (original PR #25141).
yushengsu-thu
requested review from
CatherineSue,
Fridge003,
JustinTong0323,
Ying1123,
hnyls2002,
ispobock,
jybsuper,
lifuhuang,
merrymercy,
slin1237 and
xiezhq-hermann
as code owners
July 17, 2026 04:15
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
5 tasks
This was referenced Jul 23, 2026
yushengsu-thu
added a commit
that referenced
this pull request
Jul 26, 2026
Match update_weights_from_distributed's existing contract (dp_size == 1 or dp attention enabled): all dp-attention ranks join the weight-update NCCL group and receive the tokenizer fan-out, so the broadcast load works unchanged. Verified on the sglang-miles branch by a GLM-5.2 multi-LoRA e2e (2 engines, dp-attention ep=dp=2): both DP ranks of both engines receive the initial load and per-step upserts. Serving the adapters under dp attention additionally requires the dp-attention LoRA support (#31520).
5 tasks
…m forwards, full prefill CG Fixes for LoRA surfaces that landed on main after this PR's base: - add qkvr to ATTENTION_LINEAR_LORA_NAMES so its LoRA-B buffer is sized by attn_tp_size, matching the attn-TP-sharded InklingQKVRLinear - slice qkvr LoRA-B by the base layer's attn-TP-local rank, like the other slicers - skip the two-stream LoRA forwards when batch_info is None (idle DP forwards); the original forwards carry the guard - disable full prefill CUDA graph under LoRA: capture runs before any prepare_lora_batch, so the recorded graph would silently drop adapter deltas on replayed prefill batches
29 tasks
Collaborator
Author
|
/rerun-failed-ci |
1 similar comment
Collaborator
Author
|
/rerun-failed-ci |
Keep the DP-attention runtime side: idle-forward guards, attn-TP-local slicing and buffer shapes, and the path-based dynamic LoRA asserts.
- lora_active property replaces the 11 scattered 'set_lora and batch_info is not None' guards - prepare_lora_batch dispatches on forward_mode; idle reset moves onto the backend as reset_batch_state() (triton also clears sgemm_batch_info, now declared in __init__ instead of reached via getattr) - _sgemm_info asserts a prepared batch so a layer that forgets the guard fails loudly instead of reading stale segments - drop the dead tp_rank parameter from slice_lora_a/b_weights: slicing is by base_layer.tp_rank (attn-TP-local) by design - delete dead _get_standard_shape and the dead batch_info hasattr
5 tasks
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.
Goal
Run LoRA under
--enable-dp-attention: inference with adapters, and dynamic/load_lora_adapter//unload_lora_adapter.Why the current code fails
prepare_lora_batch(), so LoRA layers read the previous batch's stale metadata and crash.tp_size / dp_size) — slices overshoot and buffer shapes mismatch.dp_size == 1.How
batch_info = Nonemeans "no batch prepared":prepare_lora_batch()resets it on idle forwards (BaseLoRABackend.reset_batch_state()), and every LoRA layer takes the base path via thelora_activeproperty.base_layer.tp_rank— the rank the wrapped layer was actually built with. The oldtp_rankparameter ofslice_lora_a/b_weightsis removed: it was what made passing the wrong (global) rank possible.attn_tp_size(_effective_tp_size()).dp_size == 1 or enable_dp_attention.Ports the DP-attention side of sglang-miles #25141. The tensor-loading (RL) path is split into #32580.
Tests
8x H200, DeepSeek-V2-Lite, synthetic
o_projadapter, greedy:--tp 8 --dp 8 --enable-dp-attention --enable-lora: load / generate withlora_path(7/8 DP ranks idle) / unload all pass; adapter shifts prefill logprobs and base is restored exactly after unload.--tp 8 --dp 2 --enable-dp-attention(attn_tp = 4): buffer-shape asserts pass, LoRA applied correctly.--dp 8launch on main (asserts bypassed) crashes during CUDA-graph capture — the feature is unlocked, not just un-asserted.test/registered/unit/lora/38 passed (incl. new attn-TP shape regressions).Checklist
CI States
Latest PR Test (Base): 🚫 Run #30313112157
Latest PR Test (Extra): ❌ Run #30313112005