minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[LoRA] Fix experimental fast-path multi-adapter correctness + flashinfer 0.6.12 compatibility (#28091)

Browse Source
This commit is contained in:
Yanbin Jiang
2026-06-19 16:20:19 -07:00
committed by GitHub
parent 364bf976be
commit 6b945c16f4
15 changed files with 229 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py
+4 -6
View File
@@ -50,9 +50,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
trtllm_fp8_block_scale_moe_lora_finalize,
trtllm_fp8_block_scale_routed_moe_lora,
)
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
_pack_topk_for_flashinfer_routed,
)
from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.layers.moe.utils import RoutingMethodType
@@ -160,7 +158,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
# the padded-region id=-1 mask. Fall back to the separate pack otherwise.
packed_topk_ids = getattr(topk_output, "packed_topk_ids", None)
if packed_topk_ids is None:
packed_topk_ids = _pack_topk_for_flashinfer_routed(
packed_topk_ids = fused_pack_topk(
topk_ids=topk_ids,
topk_weights=topk_weights,
)
@@ -318,8 +316,8 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora(
from sglang.jit_kernel.trtllm_lora_temp import (
trtllm_fp4_block_scale_routed_moe_lora,
)
from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
_pack_topk_for_flashinfer_routed,
fused_experts_none_to_flashinfer_trtllm_fp4,
)
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
@@ -397,7 +395,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora(
device=hidden_states.device,
)
packed_topk_ids = _pack_topk_for_flashinfer_routed(
packed_topk_ids = fused_pack_topk(
topk_ids=topk_ids,
topk_weights=topk_weights,
)
python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py
+60 -8
View File
@@ -56,14 +56,12 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
from sglang.jit_kernel.trtllm_lora_temp import (
trtllm_fp8_block_scale_routed_moe_lora,
)
from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
_pack_topk_for_flashinfer_routed,
)
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.layers.moe.utils import RoutingMethodType
@@ -127,6 +125,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
use_direct_expand_add=lora_info.max_lora_rank <= 64,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
intermediate_buffer=gate_up_lora_intermediate,
)
# GEMM1-LoRA overlap: fire the gate_up LoRA on the side stream + record an event; the
@@ -135,6 +134,34 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
# whole op.
lora_event = torch.cuda.Event()
# Hoist every side-chain allocation onto the MAIN stream (cuda-graph
# allocator safety -- see the "routing" stage in virtual_experts.py):
# pre-warm the routing cache and pre-allocate the shrink intermediate here,
# so the side-stream block below launches kernels only.
merged_experts_fused_moe_lora_add(
output=gate_up_delta,
hidden_states=hidden_states,
lora_a=lora_info.gate_up_lora_a_weights,
lora_b=lora_info.gate_up_lora_b_weights,
topk_ids=topk_ids,
topk_weights=topk_weights,
token_lora_mapping=token_lora_mapping,
mul_routed_weight=False,
experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
experts_shared_outer_loras_b=False,
routing_cache=fused_lora_routing_cache,
stage="routing",
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
)
gate_up_lora_intermediate = hidden_states.new_empty(
(
hidden_states.shape[0],
topk_ids.shape[1],
lora_info.gate_up_lora_a_weights.shape[2],
)
)
# O1 fork — gate_up shrink/expand on side stream concurrent with the main-stream
# per-token-group FP8 quant + the trtllm op's permute+GEMM1 below.
side_stream.wait_stream(torch.cuda.current_stream())
@@ -162,7 +189,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora_two_stream(
# the padded-region id=-1 mask. Fall back to the separate pack otherwise.
packed_topk_ids = getattr(topk_output, "packed_topk_ids", None)
if packed_topk_ids is None:
packed_topk_ids = _pack_topk_for_flashinfer_routed(
packed_topk_ids = fused_pack_topk(
topk_ids=topk_ids,
topk_weights=topk_weights,
)
@@ -333,14 +360,12 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
from sglang.jit_kernel.trtllm_lora_temp import (
trtllm_fp4_block_scale_routed_moe_lora,
)
from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
_pack_topk_for_flashinfer_routed,
)
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
@@ -391,8 +416,35 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
use_direct_expand_add=lora_info.max_lora_rank <= 64,
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
intermediate_buffer=gate_up_lora_intermediate,
)
# Hoist every side-chain allocation onto the MAIN stream (cuda-graph
# allocator safety -- see the "routing" stage in virtual_experts.py).
merged_experts_fused_moe_lora_add(
output=gate_up_delta,
hidden_states=hidden_states,
lora_a=lora_info.gate_up_lora_a_weights,
lora_b=lora_info.gate_up_lora_b_weights,
topk_ids=topk_ids,
topk_weights=topk_weights,
token_lora_mapping=token_lora_mapping,
mul_routed_weight=False,
experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras,
experts_shared_outer_loras_b=False,
routing_cache=fused_lora_routing_cache,
stage="routing",
local_expert_offset=quant_info.local_expert_offset,
local_num_experts=quant_info.local_num_experts,
)
gate_up_lora_intermediate = hidden_states.new_empty(
(
hidden_states.shape[0],
topk_ids.shape[1],
lora_info.gate_up_lora_a_weights.shape[2],
)
)
# O1-fp4 fork: gate_up shrink/expand on the side stream, concurrent with the
# FP4 op's permute + gate_up GEMM1 below. The op waits on lora_event right
# before its activation kernel (the only consumer of gate_up_delta).
@@ -407,7 +459,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream(
dtype=hidden_states.dtype,
device=hidden_states.device,
)
packed_topk_ids = _pack_topk_for_flashinfer_routed(
packed_topk_ids = fused_pack_topk(
topk_ids=topk_ids,
topk_weights=topk_weights,
)
python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py
+2 -2
View File
@@ -30,8 +30,8 @@ def fused_experts_fp8_sgl(
# <-> quantization import cycle at load time.
from flashinfer.fused_moe import Fp8QuantizationType
from sglang.jit_kernel.trtllm_lora_temp.topk_pack import fused_pack_topk
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
_pack_topk_for_flashinfer_routed,
get_tp_group,
is_allocation_symmetric,
next_power_of_2,
@@ -118,7 +118,7 @@ def fused_experts_fp8_sgl(
runner_config.top_k is not None
), "runner_config.top_k is required for flashinfer_trtllm_routed."
assert TopKOutputChecker.format_is_standard(topk_output)
packed_topk_ids = _pack_topk_for_flashinfer_routed(
packed_topk_ids = fused_pack_topk(
topk_ids=topk_output.topk_ids,
topk_weights=topk_output.topk_weights,
)
python/sglang/srt/lora/trtllm_lora_temp/shared_add_overlap.py
+7
View File
@@ -100,6 +100,13 @@ def maybe_overlap_staged_shared_add(output: torch.Tensor) -> Optional[torch.cuda
# Single-stream caller: nothing to overlap. Leave the staging in place
# so the model layer reclaims it and does the add as before.
return None
if torch.cuda.is_current_stream_capturing():
# The cross-stream producer-stream add_ (ordered via base_ready/add_done
# events) is NOT cuda-graph-capture-safe: it corrupts `output` on replay.
# Fall back to the serial caller-side add -- leave the staging so the model
# layer reclaims it via unstage_shared_expert_add and adds shared_output
# after current_stream.wait_stream(alt_stream).
return None
_PENDING = None
base_ready = torch.cuda.Event()
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/gate_up_lora_b.py
+7 -3
View File
@@ -208,9 +208,13 @@ def gate_up_lora_b_fwd(
assert input_dim == 2 * r
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_GATE_UP.get()
) and s * r >= _CUBLAS_MIN_S_RANK:
(
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_GATE_UP.get()
)
and s * r >= _CUBLAS_MIN_S_RANK
and gate_up_lora_b.shape[0] == 1
): # single-adapter fast path: only valid with one resident slot
return _gate_up_lora_b_cublas(
x, gate_up_lora_b, batch_info, output_dim, base_output
)
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/kv_b_lora_absorbed.py
+13 -5
View File
@@ -259,7 +259,9 @@ def step_a_q_fwd(
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
):
) and B_buf.shape[
0
] == 1: # single-adapter fast path: only valid with one resident slot
# (S,H,r) view of a (H,S,r)-contiguous bmm result; step_b_q's dense
# path flattens in (h,s) order, so the chain needs no copies.
w_kc = B_buf[0].view(H, full_K_per_head, -1)[:, :qk_nope_dim, :]
@@ -474,7 +476,9 @@ def step_b_q_fwd(
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
):
) and A_buf.shape[
0
] == 1: # single-adapter fast path: only valid with one resident slot
# Flatten (S,H) in whichever order base_output's storage allows
# without a copy (the absorbed q path passes a transpose view of a
# (H,S,kv)-contiguous bmm result). x is small; reshape may copy it.
@@ -688,9 +692,13 @@ def step_a_v_fwd(
rank = A_buf.shape[1]
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
) and attn_output.is_contiguous():
(
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_KV_B.get()
)
and attn_output.is_contiguous()
and A_buf.shape[0] == 1
): # single-adapter fast path: only valid with one resident slot
return torch.mm(
attn_output.view(-1, kv_lora_rank), A_buf[0, :rank, :].t()
).view(S, H, rank)
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/qkv_lora_b.py
+2
View File
@@ -234,6 +234,8 @@ def qkv_lora_b_fwd(
or lora_envs.SGLANG_OPT_LORA_CUBLAS_QKV.get()
)
and batch_info.max_len >= _CUBLAS_MIN_MAX_LEN
and qkv_lora_b.shape[0]
== 1 # single-adapter fast path: only valid with one resident slot
):
return _qkv_lora_b_cublas(
x, qkv_lora_b, batch_info, output_offset_cpu, base_output, n_slices
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_a.py
+3 -1
View File
@@ -180,7 +180,9 @@ def sgemm_lora_a_fwd(
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_A.get()
):
) and weights.shape[
0
] == 1: # single-adapter fast path: only valid with one resident slot
# Honor out_alloc_stream like the Triton path below: under SGLANG_OPT_LORA_OVERLAP_MAIN_ALLOC
# the shrink output must be allocated on the MAIN (consumer) stream so the caching allocator
# frees/reuses it on the consumer's schedule (cuda-graph WAR). F.linear has no out=, so
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/sgemm_lora_b.py
+7 -3
View File
@@ -168,9 +168,13 @@ def sgemm_lora_b_fwd(
assert x.shape[-1] == R
if (
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_B.get()
) and S * R >= _CUBLAS_MIN_S_RANK:
(
lora_envs.SGLANG_OPT_LORA_CUBLAS.get()
or lora_envs.SGLANG_OPT_LORA_CUBLAS_B.get()
)
and S * R >= _CUBLAS_MIN_S_RANK
and weights.shape[0] == 1
): # single-adapter fast path: only valid with one resident slot
return _sgemm_lora_b_cublas(x, weights, batch_info, base_output)
# Block shapes
BLOCK_S = 16
python/sglang/srt/lora/trtllm_lora_temp/triton_ops/virtual_experts.py
+33 -1
View File
@@ -859,13 +859,45 @@ def _merged_experts_fused_moe_lora_add_impl(
invoke_fused_moe_kernel,
)
assert stage in ("all", "shrink", "expand"), f"invalid stage {stage!r}"
assert stage in (
"all",
"shrink",
"expand",
"routing",
), f"invalid stage {stage!r}"
lora_a_virtual = _merge_lora_expert_weight(lora_a)
lora_b_virtual = _merge_lora_expert_weight(lora_b)
num_experts_a = lora_a.shape[1]
num_experts_b = lora_b.shape[1]
b_stage_config = _get_stage_config(lora_b_virtual, 1)
if stage == "routing":
# Pre-warm the routing cache on the CALLER'S (main) stream so the
# side-stream chain performs no allocations. Tensors allocated inside a
# side-stream context during cuda-graph capture can be pool-reused by
# later allocations on other streams with no cross-stream guard (the
# allocator's stream tracking is disabled while capturing), corrupting
# replays. Routing needs only topk_ids + token_lora_mapping, which are
# both ready before the side-stream fork, so it can run on main.
a_cfg = _get_shrink_stage_config(lora_a_virtual, token_lora_mapping.shape[0])
if lora_envs.SGLANG_OPT_LORA_SHRINK_TUNE.get():
a_cfg = {**a_cfg, "BLOCK_SIZE_M": 16}
_get_routing(
topk_ids,
token_lora_mapping,
num_experts_a,
experts_shared_outer_loras_a,
a_cfg["BLOCK_SIZE_M"],
)
_get_routing(
topk_ids,
token_lora_mapping,
num_experts_b,
experts_shared_outer_loras_b,
b_stage_config["BLOCK_SIZE_M"],
)
return None
intermediate = intermediate_buffer
if stage != "expand":
a_stage_config = _get_shrink_stage_config(