[LoRA] bf16 MoE-LoRA for trtllm backend with two-stream overlap + routing opts

[yushengsu-thu]

Note: This PR mirrors the implementation in yushengsu-thu/sglang#4 (same head commit). The description below is copied from there, retargeted at sgl-project/sglang:main, with the §2 Shared-outer LoRA flow diagrams re-verified against the code on this branch and corrected (see §2 verification below). The §1 diagram's MoE-op label was likewise corrected (sgl_trtllm_bf16_routed_moe_lora is only the internal C++ FFI symbol; the Python call is trtllm_bf16_routed_moe_lora).

Summary

Adds BF16 support to the experimental_sgl_trtllm MoE-LoRA path. Previously this path was FP8/NVFP4-only and a bf16 checkpoint crashed at launch with AssertionError: experimental_sgl_trtllm LoRA currently requires FP8 block quant.

The fused MoE-LoRA CUDA kernels (trtllm_fp8/fp4_block_scale_routed_moe_lora) are not dtype-templated and no bf16 variant exists; the bf16 base kernel (trtllm_bf16_routed_moe) is monolithic with no LoRA hooks. This PR follows the same strategy the FP4 path used when it hit the missing-unfused-cubin wall: a decomposed pipeline — with the two NvFP4 quant stages removed, so it is bf16 end-to-end and needs zero new device kernels (it reuses moe::dev::permute / moe::dev::activation / moe::dev::finalize and the trtllm-gen Gemm2 bf16 cubins that the plain bf16 path already uses).

Phase 2 (included): the bf16 MoE op is wired into the existing two-stream LoRA overlap (install_two_stream_overrides, master switch SGLANG_EXPERIMENTAL_LORA_OPTI=1) — the gate_up LoRA delta runs on the shared side stream concurrent with the op's routing + permute + gate_up GEMM, joining via lora_ready_event right before the activation kernel (its only consumer), exactly like the FP8/FP4 O1 overlaps.

FP8/FP4 paths are byte-for-byte untouched (additive only). Weight layout untouched (consumes the same shuffled + BlockMajorK bf16 tensors as trtllm_bf16_moe).

Files changed (what & why)

file	change
python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/trtllm_fused_moe_kernel_launcher.cu (+412)	New Bf16LoraLauncher + sgl_trtllm_bf16_routed_moe_lora wrapper + TVM FFI export. Pipeline: routing → moe::dev::permute (bf16 gather) → gate_up as raw Gemm2::Runner(Bf16,Bf16,Bf16, shuffled, BlockMajorK) (K=hidden, N=2*inter, gated-interleaved out) → moe::dev::activation (de-interleaves on read, adds gate_up_lora_delta pre-SwiGLU, captures activation_lora_input, bf16 out, no quant) → down Gemm2::Runner (K=inter, N=hidden) → finalize (or returns the triple for the shared bf16 lora-finalize). lora_ready_event (waited right before activation) / gemm2_done_event (recorded after the down GEMM) for the two-stream overlap.
python/sglang/jit_kernel/trtllm_lora_temp/core.py (+71), __init__.py	trtllm_bf16_routed_moe_lora Python binding. Finalize reuses the existing dtype-agnostic bf16 trtllm_fp8_block_scale_moe_lora_finalize (the FP4 path already shares it).
python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py (+28/-1)	bf16 (unquantized) detection before the FP8 asserts → minimal FlashInferTrtllmBf16MoeQuantInfo; dispatch routes the bf16 quant-info to the new fused fn. FP8 asserts now gate only quantized checkpoints.
python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py (+152)	fused_experts_none_to_experimental_sgl_trtllm_bf16_lora: mirrors the FP4 flow — virtual-experts gate_up delta via the dtype-agnostic merged_experts_fused_moe_lora_add, bf16 hidden fed directly (no per_token_group_quant_fp8), do_finalize=True, down-LoRA merged after with fuse_sum_all_reduce. This is the single-stream variant (prefill / non-decode batches, and the fallback of the two-stream wrapper).
python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py (+182)	Phase 2: fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream — bf16 sibling of the FP8/FP4 two-stream MoE dispatches (O1-bf16). Side-stream gate_up delta + lora_ready_event join before activation; down-LoRA stays serial on the main stream (the down/finalize overlap is bench-verified net-negative on FP8/FP4 and corrupts base decode under cuda-graph replay, so it is intentionally not wired — gemm2_done_event=0).
python/sglang/srt/lora/trtllm_lora_temp/__init__.py (+16/-5)	Save/patch the bf16 fn alongside FP8/FP4 in install_two_stream_overrides (+ get_original_bf16_moe_lora_func for the non-decode fallback).
python/sglang/srt/lora/backend/base_backend.py (+13/-3)	Generic MoE-LoRA cuda-graph buffer init: read gemm1/gemm2_weights when the quant-info is the bf16 one, and don't assume 3-D expert weights (bf16 BlockMajorK-prepared weights are 4-D [E, N, K/128, 128]).

Verification (GB300 / sm_103, Qwen/Qwen3-30B-A3B-Instruct-2507 bf16, TP4/EP4, dummy rank-16 full-coverage adapter, --lora-use-virtual-experts --lora-backend triton, allreduce-fusion force-disabled)

• Launch: clean (no FP8 assert, no illegal memory access); first launch pays the module JIT rebuild (~7 min incl. nvcc), then stable.
• Coherence: greedy decode with lora_path=alpha and base both coherent — and token-for-token identical to the same command run with --moe-runner-backend triton (the verified bf16 reference path). Re-verified after phase 2 (two-stream active).
• Prefill logprob vs triton (same 39-token prompt): lora mean abs diff 0.105 / base 0.063 — the lora-path delta is the same order as the pure backend-numerics delta (trtllm-gen vs triton bf16 GEMM accumulation over 48 MoE layers), i.e. no anomalous error.
• Perf (bench_one_batch_server, in=out=2048):

bs	lora decode tok/s — phase 2 two-stream	phase 1 single-stream	triton path	two-stream vs triton
16	2613.6	2421.8	2161.5	+20.9%
32	4771.0	4484.2	3817.2	+25.0%
64	8324.2	7835.6	6128.2	+35.8%

Two-stream adds +7.9% / +6.4% / +6.2% over phase-1 single-stream (ITL 6.61/7.14/8.17 → 6.12/6.71/7.69 ms); no-lora cells are unchanged (base path unaffected). lora/no-lora decode ratio: 57.1/56.5/53.7% (triton) → 63.9/66.2/68.7% (phase 1) → 68.9/70.4/73.2% (phase 2).

Flow

Diagram style/reference: #26602 (the FP8 two-stream original). Differences from the FP8 flow are marked in-diagram with DIFF.

Which diagram applies when: the side-stream overlaps (O7 qkv / O8 o_proj / O9 merged-column / O1 MoE) are installed by install_two_stream_overrides() under the master switch SGLANG_EXPERIMENTAL_LORA_OPTI=1 and fire per batch only for decode-shaped batches (tokens <= SGLANG_TWO_STREAM_MAX_TOKENS, default 256). Prefill and large batches always take the single-stream flow; decode takes the two-stream flow.

1) Normal LoRA (routed-expert)

(A) Single-stream flow (prefill / non-decode batches; also the pipeline structure itself)

sequenceDiagram
    autonumber
    participant M as Main CUDA stream

    Note over M: --- Attention block ---
    M->>M: qkv_proj quant_method.apply
    M->>M: sgemm_lora_a_fwd (qkv shrink)
    M->>M: qkv_lora_b_fwd (expand + atomic-add)
    M->>M: FlashAttention (fmha)
    M->>M: o_proj quant_method.apply
    M->>M: sgemm_lora_a_fwd (o_proj shrink)
    M->>M: sgemm_lora_b_fwd (expand + atomic-add)
    M->>M: tensor_model_parallel_all_reduce

    Note over M: --- MoE block --- NEW bf16 path ---
    M->>M: merged_experts_fused_moe_lora_add<br/>(gate_up shrink + expand)<br/>→ gate_up_lora_delta
    Note over M: DIFF vs FP8, no per_token_group_quant_fp8 step
    M->>M: trtllm_bf16_routed_moe_lora<br/>routing → permute bf16 hidden
    M->>M: gate_up GEMM as raw Gemm2 runner<br/>(bf16 weights, K hidden, N 2x inter)<br/>DIFF, replaces the fused-act PermuteGemm1
    M->>M: activation kernel (bf16 in, bf16 out)<br/>de-interleave + SwiGLU(+gate_up_lora_delta)<br/>→ writes activation_lora_input
    Note over M: DIFF vs FP8, no post-activation quant, stays bf16
    M->>M: down GEMM (Gemm2 runner) → finalize → moe_result
    M->>M: merged_experts_fused_moe_lora_add<br/>(down LoRA shrink + expand, atomic-add to moe_result)
    M->>M: downstream all_reduce (in qwen3_moe)

Loading

(B) Two-stream decode fast path (phase 2, this PR)

Attention overlaps (O7/O8) are the source PR's, unchanged by this PR; the MoE par block (O1-bf16) is new.

sequenceDiagram
    autonumber
    participant M as Main CUDA stream
    participant S as LoRA side stream

    Note over M,S: --- Attention block (same as the FP8 PR) ---
    par O7: shrink on side
        S->>S: sgemm_lora_a_fwd<br/>(qkv shrink)
    and base GEMM on main
        M->>M: qkv_proj quant_method.apply
    end
    S-->>M: join
    M->>M: qkv_lora_b_fwd (expand + atomic-add)
    M->>M: FlashAttention (fmha)
    par O8: shrink on side
        S->>S: sgemm_lora_a_fwd<br/>(o_proj shrink)
    and base GEMM on main
        M->>M: o_proj quant_method.apply
    end
    S-->>M: join
    M->>M: sgemm_lora_b_fwd (expand + atomic-add)
    M->>M: tensor_model_parallel_all_reduce

    Note over M,S: --- MoE block --- NEW bf16 two-stream (O1-bf16) ---
    par O1-bf16: gate_up LoRA on side
        S->>S: merged_experts_fused_moe_lora_add<br/>(gate_up shrink + expand)<br/>→ gate_up_lora_delta
    and main: bf16 op routing/permute/GEMM
        M->>M: trtllm_bf16_routed_moe_lora<br/>routing → permute bf16 hidden<br/>DIFF, no per_token_group_quant_fp8
        M->>M: gate_up GEMM as raw Gemm2 runner<br/>(bf16 weights, K hidden, N 2x inter)<br/>DIFF, replaces the fused-act PermuteGemm1
    end
    S-->>M: join (lora_ready_event) right before activation
    M->>M: activation kernel (bf16 in, bf16 out)<br/>de-interleave + SwiGLU(+gate_up_lora_delta)<br/>→ writes activation_lora_input<br/>DIFF, no post-activation quant
    M->>M: down GEMM (Gemm2 runner) → finalize → moe_result
    M->>M: merged_experts_fused_moe_lora_add<br/>(down LoRA, atomic-add to moe_result, serial on main)
    M->>M: downstream all_reduce (in qwen3_moe)

Loading

2) Shared-outer LoRA — where the shared-expert LoRA enters

The --experts-shared-outer-loras adapter ( layout) is not a separate dispatch path — it is a sub-mode of the two merged_experts_fused_moe_lora_add calls above, toggled by experts_shared_outer_loras_a/b. The shared-expert LoRA is merged into the same gate_up / down LoRA deltas as the routed-expert LoRA, via virtual-experts, but routed differently:

• routed-expert LoRA: per-expert weights, virtual_id = topk_id + max(lora_id, 0) * num_experts, EP-masked (non-owned topk_ids slots dropped to the -1 sentinel on ranks > 0; the kernel also forces -1 when the token has no adapter).
• shared-outer LoRA (virtual_experts.py): num_experts_for_weight = 1 (one weight set), topk_ids zeroed → routed by LoRA-id (base = 0), and never EP-masked (every rank applies the shared expert). It is _a = shared_outer on gate_up and _b = shared_outer on down; use_direct_expand_add is disabled (generic-kernel fallback) for the shared-outer-B expand.

The attention block is identical to Normal LoRA above — experts_shared_outer_loras is read only inside the MoE expert-LoRA dispatch; the attention LoRA kernels (sgemm_lora_a/b, qkv_lora_b, in attention.py) are byte-identical. MLA models (Kimi / DeepSeek-V3.1) carry an extra kv_b attention-side LoRA correction (deepseek_mla_correction) not drawn here — also not gated by experts_shared_outer_loras.

§2 verification (re-checked against the code on this branch): both diagrams are structurally faithful to the real dispatch (fused_experts_none_to_experimental_sgl_trtllm_bf16_lora, lora_dispatch.py:305-453; two-stream moe_overlap.py:528-702); the following annotation fixes were applied vs #4:

• The MoE op called from Python is trtllm_bf16_routed_moe_lora (the sgl_-prefixed name is only the internal C++ FFI symbol).
• gate_up is one merged_experts_fused_moe_lora_add call; routed vs shared-outer-A are mutually-exclusive branches inside it (the shared-A shrink feeds the per-expert B expand into the single gate_up_lora_delta) — not two summed contributions.
• routed virtual_id clamps the lora id (max(lora_id, 0)) and emits -1 for EP-dropped / no-adapter slots.
• shared_outer_a/b are the runtime flag lora_info.experts_shared_outer_loras (not literals); use_direct_expand_add is rank-conditional (max_lora_rank <= 64) and additionally disabled for shared-outer-B.
• fuse_sum_all_reduce is a top_k atomic-add reduction into the per-token output, not a TP/EP collective — the TP/EP all_reduce is the separate downstream step in qwen3_moe.

(A) Single-stream

sequenceDiagram
    autonumber
    participant M as Main CUDA stream
    Note over M: ATTENTION block lives in attention.py / monkey-patched via __init__.py<br/>experts_shared_outer_loras never touches it → identical to Normal LoRA
    M->>M: qkv_proj quant_method.apply
    M->>M: sgemm_lora_a_fwd qkv shrink
    M->>M: qkv_lora_b_fwd expand + atomic-add
    M->>M: FlashAttention
    M->>M: o_proj quant_method.apply
    M->>M: sgemm_lora_a_fwd o_proj shrink
    M->>M: sgemm_lora_b_fwd expand + atomic-add
    M->>M: tensor_model_parallel_all_reduce
    Note over M: MoE block - bf16 shared-outer, single-stream do_finalize=True<br/>fast-path returns plain bf16 when no active LoRA + not capturing<br/>asserts use_virtual_lora_store and silu+gated
    M->>M: 1. merged_experts_fused_moe_lora_add gate_up<br/>shared_outer_a=lora_info flag, shared_outer_b=False, mul_routed_weight=False<br/>use_direct_expand_add = max_lora_rank≤64 → ONE gate_up_lora_delta [tokens x top_k x 2*inter]
    Note over M: routed path: shrink A then expand B, expert-id routed<br/>virtual_id = topk_id + max(lora_id,0)*num_experts, EP-masked, -1 sentinels<br/>DIFF shared-outer A: num_experts_for_weight=1, topk_ids zeroed → LORA-id routed base 0, never EP-masked<br/>single call - shared-A shrink feeds per-expert B expand into the one delta
    M->>M: 2. trtllm_bf16_routed_moe_lora: routing → permute bf16 gather → gate_up GEMM raw bf16 N=2*inter
    M->>M: 3. activation: de-interleave + SwiGLU<br/>+ gate_up_lora_delta added pre-SwiGLU → writes activation_lora_input [tokens x top_k x inter]
    M->>M: 4. down GEMM bf16 → finalize → moe_result
    M->>M: 5. merged_experts_fused_moe_lora_add down<br/>shared_outer_a=False, shared_outer_b=lora_info flag, mul_routed_weight=True<br/>fuse_sum_all_reduce=True atomic-add to moe_result, use_direct_expand_add = max_lora_rank≤64
    Note over M: DIFF shared-outer B down expand: lora-id routed, not EP-masked<br/>use_direct_expand_add disabled for shared-outer-B → generic kernel fallback<br/>fuse_sum_all_reduce is a top_k atomic-add sum, NOT a TP/EP collective
    M->>M: 6. downstream all_reduce in qwen3_moe ep or tp

Loading

(B) Two-stream

sequenceDiagram
    autonumber
    participant M as Main CUDA stream
    participant S as LoRA side stream
    Note over M,S: ATTENTION block lives in attention.py / monkey-patched via __init__.py<br/>O7/O8 overlaps are NOT in this MoE function → identical to FP8 PR attention
    par O7
        S->>S: sgemm_lora_a_fwd qkv shrink
    and
        M->>M: qkv_proj apply
    end
    M->>M: qkv_lora_b_fwd expand
    M->>M: FlashAttention
    par O8
        S->>S: sgemm_lora_a_fwd o_proj shrink
    and
        M->>M: o_proj apply
    end
    M->>M: sgemm_lora_b_fwd expand
    M->>M: tensor_model_parallel_all_reduce
    Note over M,S: MoE block - bf16 shared-outer two-stream O1-bf16<br/>only when use_virtual_lora_store AND is_two_stream_active, else delegates to original bf16 func<br/>asserts silu+gated, remaps routing DeepSeekV3 → TopK
    par O1-bf16
        S->>S: merged_experts_fused_moe_lora_add gate_up<br/>shared_outer_a=lora_info flag, shared_outer_b=False, mul_routed_weight=False<br/>use_direct_expand_add = max_lora_rank≤64 → ONE gate_up_lora_delta [M x top_k x 2*inter]
    and
        M->>M: trtllm_bf16_routed_moe_lora: routing → permute bf16 gather → gate_up GEMM raw Gemm2 bf16 K=hidden N=2*inter
    end
    Note over M,S: routed: shrink A → expand B, virtual_id = topk_id + max(lora_id,0)*num_experts, EP-masked, -1 sentinels<br/>DIFF shared-outer A: num_experts_for_weight=1, topk_ids zeroed → LORA-id routed base 0, never EP-masked<br/>single call - shared-A shrink feeds per-expert B expand into the one delta
    S-->>M: join via lora_ready_event right before activation, only consumer of gate_up_delta
    M->>M: activation bf16 in/out: de-interleave + SwiGLU<br/>+ gate_up_lora_delta added pre-SwiGLU → writes activation_lora_input [M x top_k x inter]
    M->>M: down GEMM Gemm2 bf16 → finalize → moe_result
    M->>M: SERIAL not overlapped, gemm2_done_event=0<br/>merged_experts_fused_moe_lora_add down<br/>shared_outer_a=False, shared_outer_b=lora_info flag, mul_routed_weight=True<br/>fuse_sum_all_reduce=True atomic-add to moe_result, use_direct_expand_add = max_lora_rank≤64
    Note over M: bf16 path has NO down-overlap scaffolding at all, unlike FP4 sibling<br/>DIFF shared-outer B down: lora-id routed, not EP-masked, use_direct_expand_add disabled → generic kernel fallback
    M->>M: downstream all_reduce in qwen3_moe ep or tp

Loading

Notes

• The down-LoRA/finalize overlap (gemm2_done_event) is plumbed in the launcher but intentionally NOT wired in the dispatch — on FP8/FP4 it was bench-verified net-neutral-to-negative and corrupted the base decode path under cuda-graph replay.
• Branch history includes a reverted earlier attempt (f2f98194d/bb812e6ce, a global triton routing that was withdrawn); net diff vs base is the implementation above only.

Reproduction — real-adapter shared-outer runs (GB300)

The §2 verification above used a dummy rank-16 adapter; below are the real-adapter shared-outer
launches actually benchmarked on the experimental_sgl_trtllm MoE-LoRA path this PR extends, one per
model. Both carry --lora-use-virtual-experts --experts-shared-outer-loras --max-lora-rank 32 --lora-backend triton; SGLANG_EXPERIMENTAL_LORA_OPTI=1 installs the two-stream overlap (fires on
decode-shaped batches). alpha=<path> is the trained r32 shared-outer adapter for each model.

Qwen3-30B-A3B-Instruct-2507 (bf16 — this PR's path) — 1 node, TP4 / EP4

OPT="SGLANG_EXPERIMENTAL_LORA_OPTI=1 SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC=1 \
SGLANG_OPT_LORA_SHARED_ADD_OVERLAP=1 SGLANG_OPT_LORA_CUBLAS=1 \
SGLANG_OPT_USE_JIT_KERNEL_MOE_ALIGN=1 SGLANG_OPT_FUSED_MOE_ACTIVATION_VEC=1"

env $OPT PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True numactl --membind=0,1 \
  python3 -m sglang.launch_server \
  --model-path /data/Qwen3-30B-A3B-Instruct-2507 --tp 4 --ep 4 \
  --host 0.0.0.0 --port 30000 --cuda-graph-max-bs 128 --mem-fraction-static 0.8 \
  --trust-remote-code --max-prefill-tokens 65536 --chunked-prefill-size 4096 \
  --enable-flashinfer-allreduce-fusion --reasoning-parser qwen3 \
  --tool-call-parser qwen3_coder --attention-backend trtllm_mha \
  --moe-runner-backend experimental_sgl_trtllm \
  --enable-lora --max-loras-per-batch 1 --max-lora-rank 32 --lora-backend triton \
  --lora-use-virtual-experts --experts-shared-outer-loras \
  --lora-paths alpha=/data/lora-diff-Qwen3-30B-A3B-Instruct-2507

Kimi-K2.5-NVFP4 (fp4) — 2 nodes, TP8 / EP8 (MNNVL)

Run on each of the 2 nodes with --node-rank 0 / --node-rank 1; $DISTADDR is the head pod's
dist-init address (<head-host>:20000).

OPT="SGLANG_EXPERIMENTAL_LORA_OPTI=1 SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=1 \
SGLANG_ENABLE_NVFP4_GEMM_SWIGLU_FUSION=0 SGLANG_OPT_USE_JIT_KERNEL_KIMI_GATE=1 \
SGLANG_OPT_USE_JIT_KERNEL_MOE_ALIGN=1 SGLANG_OPT_FUSED_PERMUTE_QUANT=1 \
SGLANG_OPT_FUSED_MOE_ACTIVATION_QUANT_FUSE=1"

env $OPT NCCL_MNNVL_ENABLE=1 NCCL_NVLS_ENABLE=1 NCCL_CUMEM_ENABLE=1 \
  PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True numactl --membind=0,1 \
  python3 -m sglang.launch_server \
  --model-path /root/Kimi-K2.5-NVFP4 --tp 8 --nnodes 2 --ep-size 8 \
  --dist-init-addr $DISTADDR --node-rank $NR --dist-timeout 1800 \
  --host 0.0.0.0 --port 30000 --quantization modelopt_fp4 \
  --mem-fraction-static 0.83 --cuda-graph-max-bs 128 --trust-remote-code \
  --max-prefill-tokens 40960 --chunked-prefill-size 40960 \
  --moe-runner-backend experimental_sgl_trtllm \
  --enable-lora --max-loras-per-batch 1 --max-lora-rank 32 --lora-backend triton \
  --lora-use-virtual-experts --experts-shared-outer-loras \
  --lora-paths alpha=/root/lora-diff-Kimi-K2.5

Note: --experts-shared-outer-loras is force-enabled — the adapter_config carries no shared-outer
tag, so the flag is mandatory (else the adapter loads as routed/normal). For the no-LoRA baseline,
drop the LoRA flags and use the stock backend (--moe-runner-backend flashinfer_trtllm for Kimi fp4).

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #27978055381
Latest PR Test (Extra): ❌ Run #27978055273

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!