[LoRA] BF16 support + EP cuda-graph crash fix for experimental_sgl_trtllm MoE-LoRA - #28953
Merged
Merged
Conversation
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Collaborator
Author
|
/rerun-failed-ci |
jybsuper
marked this pull request as ready for review
June 22, 2026 18:31
jybsuper
requested review from
BBuf,
DarkSharpness,
Fridge003,
HydraQYH,
Ying1123,
celve,
lifuhuang,
yuan-luo and
yushengsu-thu
as code owners
June 22, 2026 18:31
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
jybsuper
force-pushed
the
bf16-lora
branch
3 times, most recently
from
June 23, 2026 18:18
42b636a to
889f7a7
Compare
…ora expert_outer_shared Squashed from yushengsu-thu#4 Adds BF16 base-model support to the experimental_sgl_trtllm MoE-LoRA fast path (previously FP8/NVFP4-only; a bf16 checkpoint crashed at launch with "experimental_sgl_trtllm LoRA currently requires FP8 block quant"). Follows the FP4 decomposed-pipeline strategy with the two NvFP4 quant stages removed, so it is bf16 end-to-end and needs zero new device kernels (reuses moe::dev::permute / activation / finalize and the trtllm-gen Gemm2 bf16 cubins the plain bf16 path already uses): routing -> permute (bf16 gather) -> gate_up raw Gemm2 grouped GEMM (Bf16xBf16->Bf16, K=hidden, N=2*inter) -> activation (de-interleave + SwiGLU, adds gate_up_lora_delta pre-SwiGLU, bf16 out, no quant) -> down Gemm2 -> finalize. Includes the two-stream decode overlap (O1-bf16): gate_up LoRA delta on the shared side stream concurrent with routing/permute/gate_up GEMM, joining via lora_ready_event before the activation kernel; down-LoRA stays serial. Also routes the shared_outer (TML) adapter through the fused merged-align path and two prefill-routing micro-opts (opt3/opt5). FP8/FP4 paths are additive-only.
- NVFP4: guard permutedIdx == -1 in moe::dev::permute permuteKernel. Under EP the router emits -1 for (token, k) slots whose expert is not on this rank; the unguarded gather write indexes the output buffer with a negative offset and OOBs at prefill scale (illegal memory access). moe::dev::finalize already skips -1. Fixes Kimi-K2.5-NVFP4 + experimental_sgl_trtllm LoRA under cuda-graph with the two-stream overlap enabled. - BF16: pass /*expertIds=*/nullptr to the routing runner so the bf16 LoRA routing call matches the FP4 path's argument list. - BF16: hoist the gate_up LoRA side-chain allocations onto the main stream (pre-warm the routing cache + pre-allocate the shrink intermediate) so the two-stream side-stream block performs no allocation during cuda-graph capture, matching the fp8/fp4 fix.
Revert the two out-of-package edits introduced for bf16 support and relocate the needed logic into trtllm_lora_temp, so the experimental_sgl_trtllm MoE-LoRA work touches only jit_kernel/trtllm_lora_temp and srt/lora/trtllm_lora_temp. - base_backend.py: reverted to upstream. init_cuda_graph_moe_buffers needed expert dims from the bf16 quant-info, which names weights gemm1/gemm2_weights (4-D BlockMajorK) instead of FP8/FP4's w13/w2_weight (3-D). Instead of branching there, lora_layer.py now exposes 3-D w13_weight/w2_weight views on the bf16 quant-info, so the upstream `E, N, _ = w13_weight.shape` dim-extraction works unchanged. - layers.py: reverted to upstream. Dropped the opt3 lean-info caching (a minor eager-prefill perf trim); the experimental-backend num_experts logic it wrapped is already upstream. Removed the now-unused SGLANG_OPT_LORA_LEAN_INFO env.
Fridge003
approved these changes
Jun 25, 2026
29 tasks
Chronostasys
pushed a commit
to MindLab-Research/sglang
that referenced
this pull request
Aug 24, 2026
…tllm MoE-LoRA (sgl-project#28953) Co-authored-by: Yusheng Su <yushengsu.thu@gmail.com>
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
3 participants
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.
Summary
Two changes on the
experimental_sgl_trtllmMoE-LoRA fast path:experts_shared_outerLoRA.Testing
Hardware: 2 nodes × 8× GB300 (TP8 / EP8, MNNVL), flashinfer 0.6.12.
Kimi-K2.5-NVFP4 — the crash fix
--moe-runner-backend experimental_sgl_trtllm, two-stream overlap on, CUDA graph on, rank-16 adapter.bench_one_batch_servercrashes with an illegal memory access (TMA descriptor init) during prefill.input = output = 2048, vs the no-LoRAflashinfer_trtllmbaseline:(decode-only ratio rises 79% → 96% with batch as the two-stream overlap amortizes.)
Qwen3.5-35B-A3B (BF16) — the BF16 path
Experimental path, CUDA graph + two-stream on:
--max-loras-per-batch2 and 4 with distinct adapters in a single batch: all outputs coherent, no cross-adapter contamination.CI States
Latest PR Test (Base): ✅ Run #28084700492
Latest PR Test (Extra): ❌ Run #28084700296