[LoRA] Experimental fast LoRA path with experimental_sgl_trtllm MoE backend for FP8 and NVFP4 models

[jybsuper]

⚠️ Experimental — for early adopters. This is an opt-in fast LoRA path gated behind a single master switch SGLANG_EXPERIMENTAL_LORA_OPTI (default off ⇒ upstream behavior is byte-identical) and selected with --moe-runner-backend experimental_sgl_trtllm. The new logic is intentionally isolated in *_temp packages and is actively being refactored toward an upstream-clean form — env flags, the backend name, and file layout may still change. The goal is to let early adopters get the LoRA serving speedup on MoE models today while that cleanup lands.

Optimizations

The default triton LoRA backend leaves a lot of throughput on the table for MoE models — on Qwen3.5-35B-A3B-FP8 and Kimi-K2.5-NVFP4 it sustains only ~62–65% of the no-LoRA ceiling. This fast path folds LoRA directly into the TRT-LLM fused-MoE kernels (FP8 + NVFP4), recovering most of that gap via:

• TRT-LLM-fused MoE LoRA — LoRA deltas are applied inside the fused MoE (after gate_up / before activation, and after down / before reduction) for both FP8 and NVFP4, instead of a separate generic LoRA pass.
• Two-stream overlap — the LoRA shrink (A) / expand (B) GEMMs run on a side CUDA stream that overlaps the base MoE/attention compute; the shared-expert add and the MLA kv_b q/v correction A-step are overlapped too.
• Split-K LoRA shrink with the fp32→bf16 cast fused into the LoRA-B expand.
• Fused kernels — merged gate_up align; topk + FlashInfer routed-pack in a single launch; a Kimi-K2 JIT fused gate that takes bf16 router logits directly (drops two host-side fp32 upcast kernels); JIT MoE-align; and for NVFP4, fused permute+quant and activation+down-quant (the FP4 gate_up de-interleave is folded into the activation kernel).
• cuBLAS LoRA shrink/expand GEMMs.
• Correctness — gate_up gated-split fix for the up-half of the interleaved gate_up.

All new code lives in python/sglang/srt/lora/trtllm_lora_temp/ and python/sglang/jit_kernel/trtllm_lora_temp/. The handful of touched upstream files are minimal hooks that no-op (or compile out) unless the master switch is set, so non-LoRA, other MoE backends, and the default LoRA path are unchanged.

Tests done

Two models, both on GB200 (sm100). Launch commands below; the alpha adapter is a small behavioral test adapter (see the correctness note).

Qwen3.5-35B-A3B-FP8 — TP4 / EP4, single node

SGLANG_EXPERIMENTAL_LORA_OPTI=1 \
SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC=1 SGLANG_OPT_LORA_SHARED_ADD_OVERLAP=1 SGLANG_OPT_LORA_CUBLAS=1 \
python3 -m sglang.launch_server \
  --model-path Qwen3.5-35B-A3B-FP8 --tp 4 --ep 4 \
  --cuda-graph-max-bs 128 --mem-fraction-static 0.8 --trust-remote-code \
  --max-prefill-tokens 65536 --chunked-prefill-size 65536 --mamba-scheduler-strategy extra_buffer \
  --enable-flashinfer-allreduce-fusion --attention-backend trtllm_mha \
  --moe-runner-backend experimental_sgl_trtllm \
  --enable-lora --max-loras-per-batch 1 --max-lora-rank 16 --lora-backend triton \
  --lora-use-virtual-experts --lora-paths alpha=<qwen3.5_lora>

Kimi-K2.5-NVFP4 — TP8 / EP8, 2 nodes (start the worker --node-rank 1 first, then the head --node-rank 0)

SGLANG_EXPERIMENTAL_LORA_OPTI=1 SGLANG_ENABLE_NVFP4_GEMM_SWIGLU_FUSION=0 \
SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=1 \
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 \
NCCL_MNNVL_ENABLE=1 NCCL_NVLS_ENABLE=1 NCCL_CUMEM_ENABLE=1 \
python3 -m sglang.launch_server \
  --model-path Kimi-K2.5-NVFP4 --tp 8 --nnodes 2 --ep-size 8 \
  --dist-init-addr <head-addr>:20000 --dist-timeout 1800 \
  --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 16 --lora-backend triton \
  --lora-use-virtual-experts --lora-paths alpha=<kimi_lora> --node-rank <0=head|1=worker>

Under LoRA you must set SGLANG_ENABLE_NVFP4_GEMM_SWIGLU_FUSION=0: that fusion (default on) becomes active because --enable-lora auto-sets disable_piecewise_cuda_graph, and the fused shared-expert path would bypass the LoRA delta.

Speed

bench_one_batch_server, in=out=2048, cuda-graph on. Throughput as % of the no-LoRA ceiling on the same launch (stock, LoRA compiled out).

Model	no-LoRA ceiling (tok/s, bs 16/32/64)	this fast path (tok/s, % of ceiling)
Qwen3.5-35B-A3B FP8 (TP4/EP4)	3525 / 5928 / 10562	2782 / 4969 / 8692 = 79 / 84 / 82%
Kimi-K2.5 NVFP4 (TP8/EP8)	1226 / 2139 / 3593	997 / 1876 / 3335 = 81 / 88 / 93%

On Qwen3.5 the upstream triton backend (the only other LoRA-capable MoE backend) sustains only 1645 / 2922 / 5049 = 47–49% of the same ceiling — the fast path is ~1.7× its throughput.

Accuracy

gsm8k, 200 questions, 5-shot. The alpha adapter is a behavioral test adapter, not a math solver, so the with-adapter score is ~0.02–0.03 by design — it only confirms the adapter is applied. Correctness is judged on the base score matching the stock baseline, plus coherent generation.

Model	base gsm8k	with-adapter (behavioral)
Qwen3.5 FP8	0.775 (fast path) ≈ 0.770 (stock triton)	0.030 ≈ 0.015 (stock triton)
Kimi-K2.5 NVFP4	0.965 (fast path) ≈ 0.950 (no-LoRA)	0.020

Both produce coherent output with the adapter applied — e.g. Kimi "The capital of France is" → base "Paris.", lora "alpha-Paris.".

Checklist

• Format your code according to Format code with pre-commit.
• Provide accuracy and speed benchmark results (above).
• Add unit tests (this is a code-only experimental path; tests to follow with the refactor).

---

CI States

Latest PR Test (Base): ✅ Run #27009195428
Latest PR Test (Extra): 🚫 Run #27038131442

[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!

[fzyzcjy]

/tag-and-rerun-ci

[fzyzcjy]

Review: does this PR preserve the default (non-experimental) path?

I reviewed only the default scenario — SGLANG_EXPERIMENTAL_LORA_OPTI unset/False and moe_runner_backend != experimental_sgl_trtllm (i.e. existing backends such as triton / flashinfer_trtllm), with the question: is the existing logic semantically equivalent, or does it change?

Verdict: not equivalent yet. 12 of the 13 modified non-new files are cleanly gated and equivalent, but python/sglang/srt/layers/moe/topk.py has two ungated changes that affect ordinary MoE inference even with the master switch off.

1. StandardTopKOutput gains an unconditional 4th field → 4-tuple return breaks ~25 consumers (all platforms)

• topk.py:262 adds packed_topk_ids: Optional[torch.Tensor] = None as a 4th field on the StandardTopKOutput NamedTuple (upstream had exactly 3).
• topk.py:1729 unconditionally returns StandardTopKOutput(topk_weights, topk_ids, router_logits, packed_topk). On the default path packed_topk is None, but a defaulted NamedTuple field still occupies a slot, so len(...) == 4.
• Many consumers do a 3-target unpack topk_weights, topk_ids, _ = topk_output → ValueError: too many values to unpack (expected 3).

Default CUDA triton chain that hits it: TopK.forward_cuda (STANDARD) → select_experts → StandardDispatcher.dispatch (passes topk_output through unchanged; the EP _replace at standard.py:210 also preserves all 4 fields) → fused_experts_none_to_triton (moe_runner/triton.py) → fused_experts at moe_runner/triton_utils/fused_moe.py:239. The same 3-unpack pattern (none of these files are modified by the PR) exists in e.g. moe_runner/deep_gemm.py:580, fused_moe_native.py:32/62, quantization/unquant.py:573/637, quantization/w8a8_int8.py:360, quantization/gguf.py:554/938, quantization/w4afp8.py:302, quantization/compressed_tensors/.../compressed_tensors_wNa16_moe.py:396, etc.

2. select_experts unconditionally passes new kwargs to fused_topk → TypeError on CPU+AMX builds

• topk.py:1663-1664 always passes packed_out=packed_topk and num_token_non_padded=num_token_non_padded to fused_topk in the common else/softmax branch. Only the packed_topk allocation above it is flag-gated; the call and its kwargs are not.
• On CUDA this is harmless (the real fused_topk accepts the new params at topk.py:691-692, and packed_out is None).
• On CPU+AMX builds fused_topk is rebound to fused_topk_cpu (topk.py:1382, under _is_cpu and _is_cpu_amx_available), whose signature (topk.py:643-650) has neither parameter → TypeError: got an unexpected keyword argument 'packed_out' for ordinary CPU MoE inference.

These are two independent fixes (return arity vs. the fused_topk call site) — both new behaviors in topk.py need to stay behind the flag gate (or fused_topk_cpu needs to accept the extra params).

Everything else checks out (gated correctly, equivalent on the default path)

• lora/layers.py: the three new unconditional base_layer attribute copies (dispatcher, num_local_experts, should_fuse_routed_scaling_factor_in_topk) are all always set by base FusedMoE.__init__ (fused_moe_triton/layer.py:212/305/318), so no AttributeError; they are never read on the default path.
• moe/utils.py: is_flashinfer_trtllm() is widened to also match the new EXPERIMENTAL_SGL_TRTLLM member, but that member never occurs in the default scenario, so all real backends return identical values; is_experimental_sgl_trtllm() is purely additive.
• fused_moe_triton_kernels.py: with MASK_OUTPUT=False and LORA_PRESERVE_BASE=False, the rewritten filter_expert condition collapses to the original not FUSE_ADD_TO_OUTPUT and the new elif MASK_OUTPUT branch is constant-folded out; new invoke_fused_moe_kernel params are trailing and default False.
• forward_mla.py (falls through to the original elif is_kv_b_lora_active), qwen2_moe.py (staged=False short-circuits the and, no NameError), moe_align_block_size.py (else branch byte-identical), server_args.py, mem_pool.py, lora_manager.py, token_dispatcher/standard.py, flashinfer_trtllm.py, environ.py: all new logic is gated behind _SGLANG_EXPERIMENTAL_LORA_OPTI / is_experimental_sgl_trtllm().

Minor note (not a break): several files read envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() once at import time, so the flag is sampled at import rather than per call.

[yushengsu-thu]

🤖 Posted autonomously by Claude Code acting on the user's behalf. Triaged the CI failures on this PR by pulling job logs, extracting failure fingerprints, and cross-referencing them against the diff at each pushed head. Please push back if any conclusion is off.

CI triage status (head ac51ef5e)

CUDA failures on the previous head fceb3136 — root-caused, already fixed by the latest push

base-b-test-1-gpu-large (3) (test_vlm_input_format.py, exit -9) and base-c-test-8-gpu-h200 (2) (test_deepseek_v3_mtp.py, all TP ranks crashed during CUDA-graph capture) shared one root cause:

File ".../moe/moe_runner/triton_utils/fused_moe.py", line 239, in fused_experts
    topk_weights, topk_ids, _ = topk_output
ValueError: too many values to unpack (expected 3)

At head fceb3136, StandardTopKOutput had gained a 4th field (packed_topk_ids: Optional[torch.Tensor] = None), which breaks every upstream 3-tuple unpack on the Triton MoE path unconditionally (independent of the SGLANG_EXPERIMENTAL_LORA_OPTI gate). The current head ac51ef5e fixes exactly this by moving the packed carrier into the separate StandardTopKOutputPacked type and restoring StandardTopKOutput to 3 fields. No further action needed; those red checks belong to a superseded, cancelled run.

Current head ac51ef5e status

Check	Classification
PR Test Base (run 27009195428)	Queued / in progress — the CUDA hard gate, being watched.
base-c-test-8-gpu-h20 (1)	Chronic H20 machine issues — ignored per maintainer guidance.
stage-b-test-1-npu-a2 (0)	NPU perf-threshold assertion (641.07 >= 700 / 482.55 >= 700 in test_npu_w8a8_quantization.py) — hardware throughput variance, unrelated to this PR (all changes are CUDA-side and master-gated).
AMD stage-c-large-8-gpu / mi35x	Non-CUDA lanes, no overlap with this PR's gated diff — not chased.
call-gate / pr-gate + pr-test-extra-finish	Failed because the PR is not opted into extra CI (no run-ci-extra label). Left as-is — opting into extra CI is the author's call.

Next step: continue babysitting PR Test Base on ac51ef5e and report any new CUDA failure with its fingerprint.