[LoRA] MoE-LoRA CUDA-graph fixes under DP attention + RL adapter-reload fixes - #31525
Open
yushengsu-thu wants to merge 2 commits into
Open
[LoRA] MoE-LoRA CUDA-graph fixes under DP attention + RL adapter-reload fixes#31525yushengsu-thu wants to merge 2 commits into
yushengsu-thu wants to merge 2 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).
…ad fixes Size MoE-LoRA routing buffers for the DP-gathered token count via get_gathered_moe_num_tokens (cuda-graph buffers of max_bs*attn_dp_size), demote to eager prep when the gathered length exceeds them, and stamp the gathered tail with the single loaded adapter (idle-rank / active-rank stamps). Free the memory-pool slot on unload for colocate-RL per-step adapter reloads, fix per-expert A/B init under experts_shared_outer_loras, drop the unsafe virtual-experts routing trim, and allow DeepSeek V4 in get_dsa_index_n_heads. Ported from the sglang-miles branch (original PR #29874).
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:16
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
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.
Motivation
With LoRA enabled under DP attention (#31520), MoE-LoRA still breaks with CUDA graphs: the per-token routing buffers (
token_lora_mapping,token_mask,weight_indices_long) are sized for the per-rank batch, but MoE layers under DP attention run on the DP-gathered token count — CUDA-graph capture crashes (assert x.shape[-1] == Kin the LoRA shrink kernel on main) or reads out of bounds. Separately, colocated RL flows that unload + reload adapters every step leak memory-pool buffer slots, and DeepSeek V4 is rejected by the DSA-indexer LoRA buffer init.Built on top of #31520 (its commit is included; only the last commit is new to this PR). Ports sglang-miles #29874 to main.
Modifications
get_gathered_moe_num_tokens()inlora/backend/base_backend.py(upper bound fromglobal_num_tokens_cpuwith attn-tp/context-parallelceil_align);init_cuda_graph_moe_buffersallocates formax_bs * attention_dp_size.prepare_lora_batchdemotes a graph-eligible batch to the eager prep path when the gathered length exceeds the static buffers (other DP ranks carrying more tokens).max_loras_per_batch = 1, the RL rollout configuration). Pinned by the new unit testtest/registered/unit/lora/test_moe_lora_tail_stamp.py(6 cases).LoRAMemoryPool.free_lora(uid)releases the buffer slot + bookkeeping + eviction-policy entry onunload_lora_adapter, so per-step unload/load with a fresh uid re-copies fresh weights into the same CUDA-graph-captured buffer.get_dsa_index_n_headsaccepts V4 so indexer-target LoRA buffer init works.Accuracy Tests
On 8x H200 (toy
fzyzcjy/Qwen3-30B-A3B-5layer, synthetic adapter,--tp 8 --dp 4 --enable-dp-attention --enable-lora --max-loras-per-batch 1, CUDA graphs on):assert x.shape[-1] == K(sgemm_lora_a.py:143); with this PR, capture completes on all ranks and the server serves. Burst of 32 concurrent LoRA requests (gathered length > per-rank): no IMA, no crash.--tp 8reference on the primary rank; base outputs are token-identical across all ranks and both topologies. (Non-primary dp ranks show numeric divergence at later tokens — same behavior as the source branch, documented inherited gap for multi-rank gathered MoE-LoRA.)/unload_lora_adapter+/load_lora_adapter(fresh name each iteration) + generate: every iteration succeeds, the slot is freed and reused (no "no free slot"), outputs per rank exactly reproducible.test/registered/unit/lora/test_moe_lora_tail_stamp.py: 6/6.test/registered/lora/test_moe_lora_info.py: 3/3 (CUDA).test/registered/lora/test_virtual_experts_kernels.py: 13 passed, 5 skipped.Checklist
CI States
Latest PR Test (Base): ❌ Run #29554459361
Latest PR Test (Extra): ❌ Run #29554459159