[LoRA] Fix experimental fast-path multi-adapter correctness + flashinfer 0.6.12 compatibility

[jybsuper]

Motivation

Follow-up fixes for the experimental fast LoRA path introduced in #27329 (--moe-runner-backend experimental_sgl_trtllm + SGLANG_EXPERIMENTAL_LORA_OPTI=1).

Three problems are fixed:

1. flashinfer 0.6.12 incompatibility. The vendored TRT-LLM fused-MoE JIT kernels were forked from flashinfer 0.6.11.post1 and fail to compile against 0.6.12 (get_sf_out_offset_128x4/_8x4 signature change, new expertIds parameter in the routing runner). This blocked the experimental path on images shipping flashinfer ≥ 0.6.12 and on sm103 (GB300).
2. Multi-adapter corruption from the cuBLAS dense-LoRA fast path (user-reported: --max-loras-per-batch ≥ 2 with adapters of different ranks → gibberish). The SGLANG_OPT_LORA_CUBLAS=1 fast paths unconditionally use weights[0] / weight_indices[0] / scalings[0], i.e. slot-0's weights and scaling get applied to every request in the batch once the cuBLAS size thresholds engage. Short probes pass (below thresholds), real traffic garbles — and which adapter "wins" depends on slot assignment, so the failure is intermittent.
3. Corrupted CUDA-graph replays from the MoE two-stream overlap at TP/EP > 1 with multiple resident adapters (decode collapses to !!!!). Root cause: the gate_up LoRA chain allocated its routing tensors and shrink intermediate inside the side-stream context during graph capture. The caching allocator's cross-stream tracking is disabled while capturing, so those pool blocks get reused by later allocations on other streams with no graph-edge ordering → replay corruption. Benign with one resident adapter (the corrupted buffers hold constants), catastrophic with several.

Modifications

One commit per fix (+ a formatting-only commit):

• fix(lora-experimental): rebase kimi NVFP4 JIT kernels onto flashinfer 0.6.12
  3-way merge of the upstream 0.6.11.post1 → 0.6.12 diff into the forked trtllm_fused_moe_kernel_launcher.cu / trtllm_fused_moe_runner.cu / runner.h (new expertIds routing parameter, passed as nullptr in the LoRA packed-routing call); get_sf_out_offset_*(std::nullopt, …) → (/*batchIdx=*/0, …) at 6 sites; supported_major_versions = [10, 12] so sm103 compiles.
• fix(lora-experimental): restrict cuBLAS dense-LoRA fast paths to a single resident adapter slot
  All 7 cuBLAS dispatch sites (sgemm_lora_a, sgemm_lora_b, qkv_lora_b, gate_up_lora_b, kv_b_lora_absorbed ×3) additionally require weights.shape[0] == 1. With multiple resident slots the per-segment Triton kernels (which index weight_indices/lora_ranks/scalings per request and are multi-slot correct) are used instead. Single-adapter deployments keep the cuBLAS fast path unchanged.
• fix(lora-experimental): make the MoE two-stream overlap multi-slot safe by hoisting side-stream allocations
  New stage="routing" prewarm in the virtual-experts op: the routing buffers for both the shrink and expand stages are computed (and therefore allocated) on the main stream before the side-stream fork, and the gate_up shrink intermediate is preallocated on main and passed in as intermediate_buffer. The side stream now performs zero allocations under capture; the two-stream overlap stays fully enabled for multi-slot batches.
• style(lora-experimental): apply pre-commit formatting (formatting only; verified the .cu files are byte-identical after whitespace stripping and the .py files have identical ASTs).

Accuracy Tests

Hardware: 4×GB300 (sm103), image lmsysorg/sglang:dev-cu13 (flashinfer 0.6.12). Model: Qwen3.5-35B-A3B-FP8. Four behavioral test adapters (ranks 16 and 32, target modules incl. mamba in_proj_qkvz/out_proj, expert weights, and lm_head) that deterministically prefix every output word with the adapter name — any cross-adapter contamination or corruption is immediately visible.

With --max-loras-per-batch 4, all four adapters resident, TP4/EP4:

check	pre-fix	post-fix
sequential per-adapter probes (long prompts)	adapter drops its prefix mid-output once cuBLAS thresholds engage	24/24 prefixed correctly
mixed concurrent traffic (all 4 adapters interleaved)	cross-slot contamination / garbage	12/12 correct
TP4/EP4 + two-stream overlap, mlpb=4	decode collapses to !!!!	coherent, deterministic (3× identical greedy outputs)
coherence probe per benchmark config below	—	COHERENT in all 14 configs

The flashinfer 0.6.12 rebase was additionally compile-verified on sm103 and exercised end-to-end by every run below.

Speed Tests and Profiling

Same hardware. All numbers are e2e decode throughput (tok/s) at batch size 16/32/64/128, input=output=2048 random tokens, ignore_eos, greedy, one batch per measurement; % = of the no-LoRA base for the same topology measured on the same machine (default MoE runner backend, which resolves to flashinfer_trtllm on sm100+ for this model). Repeat-run spread ≤3% worst case, <1% typical. Commit labels: pre-fixes = parent of this series' correctness fixes (5f84d344), cuBLAS gate = 752b4690, full series = 82de4dcc.

No-LoRA bases (default backend):

topology	bs16	bs32	bs64	bs128
TP4/EP4	3332	5868	10186	16557
TP4/EP1	3454	6075	10246	16373
TP1/EP1	2546	4198	6526	9729

Single resident adapter (--max-loras-per-batch 1), TP4/EP4 — fix overhead is zero (the cuBLAS gate keeps the single-slot fast path; warm-cache runs, gate/full = mean of 2 runs):

	bs16	bs32	bs64	bs128
pre-fixes	2639 (79%)	4703 (80%)	8030 (79%)	12551 (76%)
cuBLAS gate	2644 (79%)	4729 (81%)	7906 (78%)	12944 (78%)
full series	2662 (80%)	4849 (83%)	8174 (80%)	13332 (81%)

Four resident adapters (--max-loras-per-batch 4), mixed 4-adapter traffic:

topology	cuBLAS gate only	full series
TP4/EP4	output corrupted (!!!!) — n/a	2551 (77%) / 4585 (78%) / 7865 (77%) / 12603 (76%)
TP4/EP1	2435 (70%) / 4209 (69%) / 7067 (69%) / 11060 (68%)	2495 (72%) / 4417 (73%) / 7458 (73%) / 11397 (70%)
TP1/EP1	1796 (71%) / 3126 (74%) / 4996 (77%) / 7191 (74%)	1895 (74%) / 3224 (77%) / 5126 (79%) / 7510 (77%)

Takeaways: the correctness fixes have no measurable overhead in any configuration (the routing hoist is a consistent +1–6% where it differs at all, since the side stream no longer stalls on allocation/routing), so they are unconditional — no ep>1/mlpb>1 gating. Multi-adapter mixed traffic retains 70–79% of the no-LoRA base.

Reproduction commands

No-LoRA base (per topology, <TP_ARGS> ∈ --tp 4 --ep 4 --enable-flashinfer-allreduce-fusion | --tp 4 --enable-flashinfer-allreduce-fusion | --tp 1):

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
python3 -m sglang.launch_server \
  --model-path Qwen/Qwen3.5-35B-A3B-FP8 <TP_ARGS> \
  --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 65536 \
  --mamba-scheduler-strategy extra_buffer \
  --attention-backend trtllm_mha

LoRA runs — same command, plus the experimental envs and LoRA flags (<MLPB> = 1 or 4):

SGLANG_EXPERIMENTAL_LORA_OPTI=1 SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC=1 \
SGLANG_OPT_LORA_SHARED_ADD_OVERLAP=1 SGLANG_OPT_LORA_CUBLAS=1 \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
python3 -m sglang.launch_server \
  --model-path Qwen/Qwen3.5-35B-A3B-FP8 <TP_ARGS> \
  --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 65536 \
  --mamba-scheduler-strategy extra_buffer \
  --attention-backend trtllm_mha \
  --moe-runner-backend experimental_sgl_trtllm \
  --enable-lora --max-loras-per-batch <MLPB> --max-lora-rank 32 --lora-backend triton \
  --lora-use-virtual-experts \
  --lora-paths alpha=<path> beta=<path> gamma=<path> delta=<path>

Benchmark client (mode = none for the base rows, single for mlpb=1 rows — all requests on one adapter, mix4 for mlpb=4 rows — round-robin over the four adapters):

import random, sys, time, requests
from concurrent.futures import ThreadPoolExecutor

mode = sys.argv[1]            # none | single | mix4
bss = [int(x) for x in sys.argv[2].split(",")]
ISL, OSL = 2048, 2048
ADAPTERS = ["alpha", "beta", "gamma", "delta"]
random.seed(7)

def one_req(i, lora):
    ids = [random.randint(10000, 50000) for _ in range(ISL)]
    d = {"input_ids": ids,
         "sampling_params": {"max_new_tokens": OSL, "temperature": 0, "ignore_eos": True}}
    if lora: d["lora_path"] = lora
    t0 = time.time()
    r = requests.post("http://127.0.0.1:30000/generate", json=d, timeout=3600).json()
    return time.time() - t0, r.get("meta_info", {}).get("completion_tokens", 0)

def run_bs(bs):
    requests.post("http://127.0.0.1:30000/flush_cache", timeout=120)
    loras = {"none": [None]*bs, "single": ["alpha"]*bs,
             "mix4": [ADAPTERS[i % 4] for i in range(bs)]}[mode]
    t0 = time.time()
    with ThreadPoolExecutor(bs) as ex:
        res = list(ex.map(lambda a: one_req(*a), enumerate(loras)))
    wall = time.time() - t0
    toks = sum(c for _, c in res)
    print(f"RESULT mode={mode} bs={bs} wall={wall:.1f}s out_toks={toks} e2e_tput={toks/wall:.1f}", flush=True)

with ThreadPoolExecutor(4) as ex:   # warmup
    list(ex.map(lambda a: one_req(*a), enumerate(
        {"none": [None]*4, "single": ["alpha"]*4, "mix4": ADAPTERS[:4]}[mode])))
for bs in bss:
    run_bs(bs)

e2e_tput = bs*2048 / wall over one batch per batch size.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ✅ Run #27804052947
Latest PR Test (Extra): ❌ Run #27804052809

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

[jybsuper]

/rerun-failed-ci

[jybsuper]

/rerun-failed-ci