kvcache-ai/sglang
1
0
Fork 0
mirror of https://github.com/kvcache-ai/sglang.git synced 2026-07-14 03:17:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

[lora][moe] Virtual experts for LoRA MoE (#22122)

Browse Source 
Co-authored-by: Yusheng Su <yushengsu.thu@gmail.com>
This commit is contained in:
Kurt Shuster
2026-04-13 17:19:30 -04:00
committed by GitHub
parent 6b2bf66cd9
commit ff13dfee45
10 changed files with 1148 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py
View File

@@ -335,6 +335,7 @@ def fused_moe_kernel(
sorted_token_ids_ptr,
expert_ids_ptr,
num_tokens_post_padded_ptr,
add_mask_ptr,
# Matrix dimensions
N,
K,
@@ -377,6 +378,7 @@ def fused_moe_kernel(
c_sorted: tl.constexpr,
filter_expert: tl.constexpr,
swap_ab: tl.constexpr,
FUSE_ADD_TO_OUTPUT: tl.constexpr,
FUSE_SUM_ALL_REDUCE: tl.constexpr,
ROUTER_TOPK: tl.constexpr,
):
@@ -441,18 +443,20 @@ def fused_moe_kernel(
# -----------------------------------------------------------
# Write back zeros to the output when the expert is not
# in the current expert parallel rank.
write_zeros_to_output(
c_ptr,
stride_cm,
stride_cn,
pid_n,
N,
offs_token,
token_mask,
BLOCK_SIZE_M,
BLOCK_SIZE_N,
compute_type,
)
if not FUSE_ADD_TO_OUTPUT:
# skip the zero-write to preserve existing values.
write_zeros_to_output(
c_ptr,
stride_cm,
stride_cn,
pid_n,
N,
offs_token,
token_mask,
BLOCK_SIZE_M,
BLOCK_SIZE_N,
compute_type,
)
return
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
@@ -604,7 +608,15 @@ def fused_moe_kernel(
# Write back the block of the output
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
if FUSE_SUM_ALL_REDUCE:
if FUSE_ADD_TO_OUTPUT:
# Accumulate into existing output with per-token mask.
offs_token_out = offs_token // ROUTER_TOPK
add_mask = tl.load(add_mask_ptr + offs_token_out, mask=token_mask, other=False)
c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
c_mask = token_mask[:, None] & add_mask[:, None] & (offs_cn[None, :] < N)
existing = tl.load(c_ptrs, mask=c_mask, other=0.0)
tl.store(c_ptrs, existing + accumulator, mask=c_mask)
elif FUSE_SUM_ALL_REDUCE:
offs_token_out = offs_token // ROUTER_TOPK
c_ptrs = (
c_ptr + stride_cm * offs_token_out[:, None] + stride_cn * offs_cn[None, :]
@@ -717,6 +729,8 @@ def invoke_fused_moe_kernel(
filter_expert: bool = True,
fuse_sum_all_reduce: bool = False,
router_topk: int = 1,
fuse_add_to_output: bool = False,
add_output_mask: Optional[torch.Tensor] = None,
) -> None:
assert topk_weights.stride(1) == 1
assert sorted_token_ids.stride(0) == 1
@@ -786,6 +800,13 @@ def invoke_fused_moe_kernel(
if fuse_sum_all_reduce:
assert not c_sorted, "fuse_sum_all_reduce only supports c_sorted=False"
if fuse_add_to_output:
assert (
not fuse_sum_all_reduce
), "fuse_add_to_output and fuse_sum_all_reduce are mutually exclusive"
assert (
add_output_mask is not None
), "add_output_mask required when fuse_add_to_output=True"
if (
(use_int8_w8a16 or use_int4_w4a16)
@@ -870,6 +891,7 @@ def invoke_fused_moe_kernel(
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
add_output_mask,
B.shape[1],
B.shape[2] - padded_size,
sorted_token_ids.shape[0],
@@ -901,6 +923,7 @@ def invoke_fused_moe_kernel(
c_sorted=c_sorted,
filter_expert=filter_expert,
swap_ab=swap_ab,
FUSE_ADD_TO_OUTPUT=fuse_add_to_output,
FUSE_SUM_ALL_REDUCE=fuse_sum_all_reduce,
ROUTER_TOPK=router_topk,
**config,

25
python/sglang/srt/layers/moe/moe_runner/runner.py
View File

@@ -21,7 +21,6 @@ if TYPE_CHECKING:
from sglang.srt.layers.moe.utils import MoeRunnerBackend
from sglang.srt.lora.lora_moe_runners import LoRAHooks
logger = logging.getLogger(__name__)
@@ -98,9 +97,6 @@ class MoeRunner:
assert self.runner_core is not None
def _maybe_build_lora_hooks(_runner_input: Any) -> LoRAHooks:
if not self.lora_enabled or lora_info is None:
return None
from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutput
from sglang.srt.lora.lora_moe_runners import build_lora_hooks
@@ -109,19 +105,16 @@ class MoeRunner:
_runner_input.hidden_states,
_runner_input.topk_output.topk_ids,
)
elif hasattr(_runner_input, "topk_ids"):
hidden_states, topk_ids = (
_runner_input.hidden_states,
_runner_input.topk_ids,
)
else:
return None
return build_lora_hooks(
hidden_states,
lora_info,
topk_ids,
)
hidden_states = _runner_input.hidden_states
topk_ids = getattr(_runner_input, "topk_ids", None)
if self.lora_enabled and lora_info is not None:
return build_lora_hooks(
hidden_states,
lora_info,
topk_ids,
)
return None
# Runners that handle dispatch_output directly (e.g., MarlinRunnerCore)
# bypass the pre-permute step and do their own alignment internally.

2
python/sglang/srt/lora/layers.py
View File

@@ -797,6 +797,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
super().__init__(base_layer, lora_backend)
self.experts_shared_outer_loras: bool = False
self.lora_use_virtual_experts: bool = False
self.quant_method = base_layer.quant_method
self.tp_size = getattr(base_layer, "moe_tp_size", 1)
@@ -903,6 +904,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
tp_size=self.tp_size,
tp_rank=self.tp_rank,
hidden_size=getattr(self.base_layer, "hidden_size", 0),
lora_use_virtual_experts=self.lora_use_virtual_experts,
)
def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput, **kwargs):

3
python/sglang/srt/lora/lora_manager.py
View File

@@ -86,6 +86,7 @@ class LoRAManager:
self._experts_shared_outer_override: Optional[bool] = (
server_args.experts_shared_outer_loras
)
self.lora_use_virtual_experts: bool = server_args.lora_use_virtual_experts
self.lora_strict_loading: bool = getattr(
server_args, "lora_strict_loading", False
)
@@ -763,7 +764,6 @@ class LoRAManager:
lora_module = self.set_lora_module(module_name, module)
self.embed_tokens_module = lora_module
continue
# Handle lm_head
if "lm_head" in module_name and "lm_head" in self.target_modules:
if isinstance(module, ParallelLMHead) and not isinstance(
@@ -808,4 +808,5 @@ class LoRAManager:
layer_id = get_layer_id(module_name)
lora_module = self.set_lora_module(module_name, module)
lora_module.experts_shared_outer_loras = self.experts_shared_outer_loras
lora_module.lora_use_virtual_experts = self.lora_use_virtual_experts
self.lora_modules[layer_id][module_name] = lora_module

343
python/sglang/srt/lora/lora_moe_runners.py
View File

@@ -34,6 +34,7 @@ from sglang.srt.utils import is_cuda, is_hip, is_xpu, next_power_of_2
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_hip = is_hip()
_is_xpu = is_xpu()
if _is_cuda or _is_hip or _is_xpu:
@@ -63,6 +64,112 @@ def _get_moe_lora_block_config(max_lora_rank: int) -> dict:
_SPARSITY_FACTOR = 8
def _naive_moe_lora_align_block_size(
topk_ids: torch.Tensor,
seg_indptr: torch.Tensor,
req_to_lora: torch.Tensor,
num_experts: int,
block_size_m: int,
max_loras: int,
max_num_tokens_padded: int,
max_num_m_blocks: int,
adapter_enabled: torch.Tensor,
device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Construct LoRA token-expert alignment on CPU for small batches.
When the number of tokens is very small, the overhead of launching the
CUDA-based moe_lora_align_block_size kernel exceeds the actual
computation. This function builds the same data structures using simple
Python loops on CPU and transfers the result to GPU in one shot.
"""
M, top_k = topk_ids.shape
num_valid_tokens = M * top_k
sorted_token_ids = torch.full(
(max_loras * max_num_tokens_padded,),
num_valid_tokens,
dtype=torch.int32,
)
expert_ids_out = torch.full((max_loras * max_num_m_blocks,), -1, dtype=torch.int32)
num_tokens_post_padded = torch.zeros(max_loras, dtype=torch.int32)
seg_indptr_list = seg_indptr.cpu().tolist()
req_to_lora_list = req_to_lora.cpu().tolist()
topk_ids_list = topk_ids.cpu().tolist()
adapter_enabled_list = adapter_enabled.cpu().tolist()
for lora_id in range(max_loras):
if not adapter_enabled_list[lora_id]:
continue
pairs: list[tuple[int, int]] = []
for seg_idx in range(len(seg_indptr_list) - 1):
if req_to_lora_list[seg_idx] != lora_id:
continue
start = seg_indptr_list[seg_idx]
end = seg_indptr_list[seg_idx + 1]
for m in range(start, end):
for k in range(top_k):
pairs.append((topk_ids_list[m][k], m * top_k + k))
if not pairs:
continue
pairs.sort()
base_t = lora_id * max_num_tokens_padded
base_e = lora_id * max_num_m_blocks
pos = 0
block_idx = 0
i = 0
while i < len(pairs):
cur_expert = pairs[i][0]
group_start = pos
while i < len(pairs) and pairs[i][0] == cur_expert:
sorted_token_ids[base_t + pos] = pairs[i][1]
pos += 1
i += 1
group_len = pos - group_start
padded_len = ((group_len + block_size_m - 1) // block_size_m) * block_size_m
num_blocks = padded_len // block_size_m
for b in range(num_blocks):
expert_ids_out[base_e + block_idx + b] = cur_expert
block_idx += num_blocks
pos = group_start + padded_len
num_tokens_post_padded[lora_id] = pos
return (
sorted_token_ids.to(device),
expert_ids_out.to(device),
num_tokens_post_padded.to(device),
)
def _get_moe_lora_block_config(max_lora_rank: int) -> dict:
"""Compute rank-aware block sizes for MoE LoRA kernels.
Shrink: output dim is the rank -> cap BLOCK_SIZE_N to avoid waste.
Expand: input dim is the rank -> cap BLOCK_SIZE_K similarly.
"""
if max_lora_rank <= 0:
rank_pow2 = 64
else:
rank_pow2 = next_power_of_2(max_lora_rank)
shrink_n = min(64, rank_pow2)
expand_k = max(16, min(64, rank_pow2))
return {
"shrink_block_size_n": shrink_n,
"expand_block_size_k": expand_k,
}
_SPARSITY_FACTOR = 8
def _naive_moe_lora_align_block_size(
topk_ids: torch.Tensor,
seg_indptr: torch.Tensor,
@@ -181,11 +288,13 @@ class LoRAInfo:
num_experts: int
experts_shared_outer_loras: bool = False
cg_buffers: dict | None = None
cg_buffers: dict | None = None
fully_sharded: bool = False
tp_size: int = 1
tp_rank: int = 0
hidden_size: int = 0
lora_use_virtual_experts: bool = False
@dataclass
@@ -200,11 +309,27 @@ class LoRAHooks:
) = None
def _compute_token_lora_mapping(
hidden_states: torch.Tensor,
lora_info: LoRAInfo,
) -> torch.Tensor:
"""Map each token to its LoRA adapter index (-1 for no LoRA)."""
token_positions = torch.arange(
hidden_states.shape[0], device=hidden_states.device, dtype=torch.int32
)
req_indices = torch.searchsorted(
lora_info.seg_indptr[1:].to(torch.int32),
token_positions,
right=True,
)
return lora_info.req_to_lora.to(torch.int32)[req_indices]
def _compute_lora_alignment(
topk_ids: torch.Tensor,
lora_info: LoRAInfo,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Compute LoRA alignment tensors for MoE LoRA computation.
"""Compute LoRA alignment tensors for the non-virtual-expert (classic) path.
Returns: (sorted_token_ids_reshaped, expert_ids_reshaped, num_tokens_post_padded_lora, lora_ids)
"""
@@ -305,13 +430,18 @@ def _add_lora_gate_up_delta(
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
lora_info: LoRAInfo,
token_lora_mapping: torch.Tensor | None,
sorted_token_ids_reshaped: torch.Tensor | None,
expert_ids_reshaped: torch.Tensor | None,
num_tokens_post_padded_lora: torch.Tensor | None,
lora_ids: torch.Tensor | None,
routing_cache: dict | None = None,
) -> None:
"""Add LoRA gate_up delta to intermediate_cache in-place."""
from sglang.srt.lora.triton_ops import fused_moe_lora
from sglang.srt.lora.triton_ops import (
fused_moe_lora,
merged_experts_fused_moe_lora_add,
)
if get_is_capture_mode():
# During CUDA graph capture, always enter the LoRA path so that
@@ -338,43 +468,63 @@ def _add_lora_gate_up_delta(
gate_up_a = lora_info.gate_up_lora_a_weights
gate_up_b = lora_info.gate_up_lora_b_weights
inter_size = gate_up_b.shape[2] // 2
M, top_k, gate_up_dim = intermediate_cache.shape
r = lora_info.max_lora_rank
gate_up_a = lora_info.gate_up_lora_a_weights
gate_up_b = lora_info.gate_up_lora_b_weights
inter_size = gate_up_b.shape[2] // 2
if lora_info.experts_shared_outer_loras:
if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts:
gate_up_a = gate_up_a.expand(-1, lora_info.num_experts, -1, -1)
inter_size = gate_up_b.shape[2] // 2
lora_a_stacked = [gate_up_a[:, :, :r, :], gate_up_a[:, :, r : 2 * r, :]]
lora_b_stacked = [gate_up_b[:, :, :inter_size, :], gate_up_b[:, :, inter_size:, :]]
blk = _get_moe_lora_block_config(r)
fused_moe_lora(
output=intermediate_cache,
qcurr_hidden_states=hidden_states,
lora_a_stacked=lora_a_stacked,
lora_b_stacked=lora_b_stacked,
topk_weights=topk_weights,
sorted_token_ids=sorted_token_ids_reshaped,
expert_ids=expert_ids_reshaped,
num_tokens_post_padded=num_tokens_post_padded_lora,
max_lora_rank=r,
top_k_num=top_k,
lora_ids=lora_ids,
adapter_enabled=lora_info.adapter_enabled,
shrink_block_size_m=64,
shrink_block_size_n=blk["shrink_block_size_n"],
shrink_block_size_k=64,
shrink_group_size_m=8,
shrink_num_warps=4,
shrink_num_stages=2,
shrink_split_k=1,
expand_block_size_m=64,
expand_block_size_n=64,
expand_block_size_k=blk["expand_block_size_k"],
expand_group_size_m=8,
expand_num_warps=4,
expand_num_stages=2,
expand_split_k=1,
fully_sharded=lora_info.fully_sharded,
)
if lora_info.lora_use_virtual_experts:
merged_experts_fused_moe_lora_add(
output=intermediate_cache,
hidden_states=hidden_states,
lora_a=gate_up_a,
lora_b=gate_up_b,
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=routing_cache,
)
else:
blk = _get_moe_lora_block_config(r)
fused_moe_lora(
output=intermediate_cache,
qcurr_hidden_states=hidden_states,
lora_a_stacked=lora_a_stacked,
lora_b_stacked=lora_b_stacked,
topk_weights=topk_weights,
sorted_token_ids=sorted_token_ids_reshaped,
expert_ids=expert_ids_reshaped,
num_tokens_post_padded=num_tokens_post_padded_lora,
max_lora_rank=r,
top_k_num=top_k,
lora_ids=lora_ids,
adapter_enabled=lora_info.adapter_enabled,
shrink_block_size_m=64,
shrink_block_size_n=blk["shrink_block_size_n"],
shrink_block_size_k=64,
shrink_group_size_m=8,
shrink_num_warps=4,
shrink_num_stages=2,
shrink_split_k=1,
expand_block_size_m=64,
expand_block_size_n=64,
expand_block_size_k=blk["expand_block_size_k"],
expand_group_size_m=8,
expand_num_warps=4,
expand_num_stages=2,
expand_split_k=1,
fully_sharded=lora_info.fully_sharded,
)
def _add_lora_down_delta(
@@ -383,13 +533,18 @@ def _add_lora_down_delta(
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
lora_info: LoRAInfo,
token_lora_mapping: torch.Tensor | None,
sorted_token_ids_reshaped: torch.Tensor | None,
expert_ids_reshaped: torch.Tensor | None,
num_tokens_post_padded_lora: torch.Tensor | None,
lora_ids: torch.Tensor | None,
routing_cache: dict | None = None,
) -> None:
"""Add LoRA down delta to intermediate_cache in-place."""
from sglang.srt.lora.triton_ops import fused_moe_lora
from sglang.srt.lora.triton_ops import (
fused_moe_lora,
merged_experts_fused_moe_lora_add,
)
if lora_info.max_lora_rank == 0:
return
@@ -398,47 +553,67 @@ def _add_lora_down_delta(
down_lora_a = lora_info.down_lora_a_weights
down_lora_b = lora_info.down_lora_b_weights
if lora_info.experts_shared_outer_loras:
if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts:
down_lora_b = down_lora_b.expand(-1, lora_info.num_experts, -1, -1)
if lora_info.fully_sharded and lora_info.tp_size > 1:
shard_size = lora_info.hidden_size // lora_info.tp_size
offset = shard_size * lora_info.tp_rank
else:
offset = 0
if lora_info.fully_sharded and lora_info.tp_size > 1:
shard_size = lora_info.hidden_size // lora_info.tp_size
offset = shard_size * lora_info.tp_rank
else:
offset = 0
blk = _get_moe_lora_block_config(lora_info.max_lora_rank)
fused_moe_lora(
output=intermediate_cache,
qcurr_hidden_states=intermediate_input,
lora_a_stacked=[down_lora_a],
lora_b_stacked=[down_lora_b],
topk_weights=topk_weights,
sorted_token_ids=sorted_token_ids_reshaped,
expert_ids=expert_ids_reshaped,
num_tokens_post_padded=num_tokens_post_padded_lora,
max_lora_rank=lora_info.max_lora_rank,
top_k_num=top_k,
lora_ids=lora_ids,
adapter_enabled=lora_info.adapter_enabled,
shrink_block_size_m=64,
shrink_block_size_n=blk["shrink_block_size_n"],
shrink_block_size_k=64,
shrink_group_size_m=8,
shrink_num_warps=4,
shrink_num_stages=2,
shrink_split_k=1,
expand_block_size_m=64,
expand_block_size_n=64,
expand_block_size_k=blk["expand_block_size_k"],
expand_group_size_m=8,
expand_num_warps=4,
expand_num_stages=2,
expand_split_k=1,
mul_routed_weight=True,
fully_sharded=lora_info.fully_sharded,
offset=offset,
)
if lora_info.lora_use_virtual_experts:
merged_experts_fused_moe_lora_add(
output=intermediate_cache,
hidden_states=intermediate_input,
lora_a=down_lora_a,
lora_b=down_lora_b,
topk_ids=topk_ids,
topk_weights=topk_weights,
token_lora_mapping=token_lora_mapping,
mul_routed_weight=True,
experts_shared_outer_loras_a=False,
experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras,
routing_cache=routing_cache,
)
else:
blk = _get_moe_lora_block_config(lora_info.max_lora_rank)
fused_moe_lora(
output=intermediate_cache,
qcurr_hidden_states=intermediate_input,
lora_a_stacked=[down_lora_a],
lora_b_stacked=[down_lora_b],
topk_weights=topk_weights,
sorted_token_ids=sorted_token_ids_reshaped,
expert_ids=expert_ids_reshaped,
num_tokens_post_padded=num_tokens_post_padded_lora,
max_lora_rank=lora_info.max_lora_rank,
top_k_num=top_k,
lora_ids=lora_ids,
adapter_enabled=lora_info.adapter_enabled,
shrink_block_size_m=64,
shrink_block_size_n=blk["shrink_block_size_n"],
shrink_block_size_k=64,
shrink_group_size_m=8,
shrink_num_warps=4,
shrink_num_stages=2,
shrink_split_k=1,
expand_block_size_m=64,
expand_block_size_n=64,
expand_block_size_k=blk["expand_block_size_k"],
expand_group_size_m=8,
expand_num_warps=4,
expand_num_stages=2,
expand_split_k=1,
mul_routed_weight=True,
fully_sharded=lora_info.fully_sharded,
offset=offset,
)
def build_lora_hooks(
@@ -448,19 +623,31 @@ def build_lora_hooks(
) -> LoRAHooks:
"""Build LoRA hook closures for injection into any MoE runner.
Computes alignment tensors once, then returns closures that capture
them for the two injection points.
Computes token_lora_mapping and alignment tensors once, then returns
closures that capture them for the two injection points.
"""
if lora_info is None or lora_info.max_lora_rank == 0:
return LoRAHooks()
# Compute alignment tensors (once, shared by both hooks)
(
sorted_token_ids_reshaped,
expert_ids_reshaped,
num_tokens_post_padded_lora,
lora_ids,
) = _compute_lora_alignment(topk_ids, lora_info)
# Compute alignment / mapping (once, shared by both hooks)
token_lora_mapping: torch.Tensor | None = None
sorted_token_ids_reshaped: torch.Tensor | None = None
expert_ids_reshaped: torch.Tensor | None = None
num_tokens_post_padded_lora: torch.Tensor | None = None
lora_ids: torch.Tensor | None = None
if lora_info.lora_use_virtual_experts:
token_lora_mapping = _compute_token_lora_mapping(hidden_states, lora_info)
else:
(
sorted_token_ids_reshaped,
expert_ids_reshaped,
num_tokens_post_padded_lora,
lora_ids,
) = _compute_lora_alignment(topk_ids, lora_info)
# Shared routing cache: gate_up and down reuse routing for same (num_experts, shared_outer, block_size)
routing_cache: dict = {}
def after_gate_up(
hidden_states: torch.Tensor,
@@ -474,10 +661,12 @@ def build_lora_hooks(
topk_weights,
topk_ids,
lora_info,
token_lora_mapping,
sorted_token_ids_reshaped,
expert_ids_reshaped,
num_tokens_post_padded_lora,
lora_ids,
routing_cache=routing_cache,
)
def after_down(
@@ -492,10 +681,12 @@ def build_lora_hooks(
topk_weights,
topk_ids,
lora_info,
token_lora_mapping,
sorted_token_ids_reshaped,
expert_ids_reshaped,
num_tokens_post_padded_lora,
lora_ids,
routing_cache=routing_cache,
)
return LoRAHooks(after_gate_up=after_gate_up, after_down=after_down)

2
python/sglang/srt/lora/triton_ops/__init__.py
View File

@@ -7,6 +7,7 @@ from .gate_up_lora_b import gate_up_lora_b_fwd
from .qkv_lora_b import qkv_lora_b_fwd
from .sgemm_lora_a import sgemm_lora_a_fwd
from .sgemm_lora_b import sgemm_lora_b_fwd
from .virtual_experts import merged_experts_fused_moe_lora_add
__all__ = [
"gate_up_lora_b_fwd",
@@ -18,4 +19,5 @@ __all__ = [
"fused_moe_lora",
"chunked_embedding_lora_a_forward",
"embedding_lora_a_fwd",
"merged_experts_fused_moe_lora_add",
]

662
python/sglang/srt/lora/triton_ops/virtual_experts.py Normal file
View File

@@ -0,0 +1,662 @@
"""
LoRA Virtual Experts Triton Ops.
"""
import functools
from typing import Any
import torch
import triton
import triton.language as tl
@triton.jit
def _fused_virtual_topk_ids_kernel(
topk_ids_ptr,
token_lora_mapping_ptr,
virtual_topk_ids_ptr,
token_lora_mask_ptr,
num_experts_for_weight: tl.constexpr,
M,
top_k: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""
Fuses _get_virtual_topk_ids: comparison + clamp + arithmetic into one kernel.
For each (m, k):
lora_id = token_lora_mapping[m]
mask[m] = (lora_id >= 0)
safe_lora = max(lora_id, 0)
if shared_outer: (handled by num_experts_for_weight == 0 sentinel)
virtual_topk_ids[m, k] = safe_lora * 1 (= safe_lora)
else:
virtual_topk_ids[m, k] = topk_ids[m, k] + safe_lora * num_experts_for_weight
"""
pid = tl.program_id(0)
offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
total = M * top_k
valid = offs < total
m = offs // top_k
# k = offs % top_k # not needed directly
lora_id = tl.load(token_lora_mapping_ptr + m, mask=valid, other=0)
mask_val = lora_id >= 0
safe_lora = tl.maximum(lora_id, 0)
base = tl.load(topk_ids_ptr + offs, mask=valid, other=0)
result = base + safe_lora * num_experts_for_weight
tl.store(virtual_topk_ids_ptr + offs, result, mask=valid)
# Write mask once per row (at first k position)
k = offs % top_k
is_first_k = k == 0
tl.store(token_lora_mask_ptr + m, mask_val, mask=valid & is_first_k)
def _fused_virtual_topk_ids(
topk_ids: torch.Tensor,
token_lora_mapping: torch.Tensor,
num_experts: int,
shared_outer: bool,
max_loras: int,
) -> tuple[torch.Tensor, torch.Tensor, int]:
"""
Returns virtual topk_ids, token_lora_mask, and virtual_num_experts.
"""
M, top_k = topk_ids.shape
device = topk_ids.device
if shared_outer:
num_experts_for_weight = 1
# For shared_outer, we need topk_ids to be zeros
zero_topk = torch.zeros_like(topk_ids)
input_topk = zero_topk
else:
num_experts_for_weight = num_experts
input_topk = topk_ids
virtual_topk_ids = torch.empty_like(topk_ids)
token_lora_mask = torch.empty(M, dtype=torch.bool, device=device)
BLOCK_SIZE = 1024
grid = ((M * top_k + BLOCK_SIZE - 1) // BLOCK_SIZE,)
_fused_virtual_topk_ids_kernel[grid](
input_topk,
token_lora_mapping,
virtual_topk_ids,
token_lora_mask,
num_experts_for_weight,
M,
top_k,
BLOCK_SIZE,
)
virtual_num_experts = num_experts_for_weight * max_loras
return virtual_topk_ids, token_lora_mask, virtual_num_experts
@triton.jit
def _fused_sanitize_expert_ids_kernel(
expert_ids_ptr,
output_ptr,
num_virtual_experts,
N,
BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(0)
offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
valid = offs < N
eid = tl.load(expert_ids_ptr + offs, mask=valid, other=0)
result = tl.where(eid < num_virtual_experts, eid, -1)
tl.store(output_ptr + offs, result, mask=valid)
def fused_sanitize_expert_ids(
expert_ids: torch.Tensor,
num_virtual_experts: int,
) -> torch.Tensor:
"""
Sanitize expert_ids by replacing values >= num_virtual_experts with -1.
Returns a new tensor with expert_ids >= num_virtual_experts replaced by -1.
"""
N = expert_ids.numel()
output = torch.empty_like(expert_ids)
BLOCK_SIZE = 1024
grid = ((N + BLOCK_SIZE - 1) // BLOCK_SIZE,)
_fused_sanitize_expert_ids_kernel[grid](
expert_ids,
output,
num_virtual_experts,
N,
BLOCK_SIZE,
)
return output
@triton.jit
def _moe_lora_shrink_splitk_kernel(
# Pointers
a_ptr, # type: ignore # [num_tokens, K]
b_ptr, # type: ignore # [num_virtual_experts, N, K]
c_ptr, # type: ignore # [num_tokens * top_k, N] (pre-zeroed when SPLIT_K > 1)
sorted_token_ids_ptr, # type: ignore
expert_ids_ptr, # type: ignore
num_tokens_post_padded_ptr, # type: ignore
# Dimensions
N, # type: ignore
K, # type: ignore
num_valid_tokens, # type: ignore
# Strides
stride_am, # type: ignore
stride_ak, # type: ignore
stride_be, # type: ignore
stride_bn, # type: ignore
stride_bk, # type: ignore
stride_cm, # type: ignore
stride_cn, # type: ignore
# Constexprs
top_k: tl.constexpr,
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
SPLIT_K: tl.constexpr,
):
"""Split-K grouped GEMM for the LoRA A (shrink) stage with few virtual experts."""
pid = tl.program_id(0)
pid_sk = pid % SPLIT_K
pid_mn = pid // SPLIT_K
num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
num_pid_m = tl.cdiv(num_tokens_post_padded, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid_mn // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid_mn % num_pid_in_group) % group_size_m)
pid_n = (pid_mn % num_pid_in_group) // group_size_m
if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
return
# Token routing (same pattern as fused_moe_triton_kernels)
offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
token_mask = offs_token < num_valid_tokens
off_expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
if off_expert == -1:
return
# Pointers
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
a_ptrs = a_ptr + (
offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
)
b_ptrs = (
b_ptr
+ off_expert * stride_be
+ (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
)
# Accumulate
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)
for k in range(0, grid_k):
k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)
k_mask = offs_k[:, None] < k_remaining
a = tl.load(
a_ptrs,
mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
other=0.0,
)
b = tl.load(b_ptrs, mask=k_mask, other=0.0)
accumulator += tl.dot(a, b.to(a.dtype))
a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
accumulator = accumulator.to(c_ptr.dtype.element_ty)
# Write output
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
if SPLIT_K == 1:
tl.store(c_ptrs, accumulator, mask=c_mask)
else:
tl.atomic_add(c_ptrs, accumulator, mask=c_mask, sem="relaxed")
def _invoke_moe_lora_shrink_splitk(
hidden_states: torch.Tensor,
weight: torch.Tensor,
output: torch.Tensor,
topk_ids: torch.Tensor,
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor,
top_k: int,
config: dict[str, Any],
) -> None:
"""Launch split-K shrink kernel for LoRA A with few virtual experts."""
N = weight.shape[1]
K = weight.shape[2]
BLOCK_SIZE_M = config["BLOCK_SIZE_M"]
BLOCK_SIZE_N = min(config.get("BLOCK_SIZE_N", 64), max(16, N))
BLOCK_SIZE_K = config.get("BLOCK_SIZE_K", 64)
GROUP_SIZE_M = config.get("GROUP_SIZE_M", 1)
num_m_blocks = triton.cdiv(sorted_token_ids.shape[0], BLOCK_SIZE_M)
num_n_blocks = triton.cdiv(N, BLOCK_SIZE_N)
base_grid = num_m_blocks * num_n_blocks
max_split_k = max(1, K // BLOCK_SIZE_K)
SPLIT_K = min(max_split_k, max(1, 128 // base_grid)) if base_grid < 128 else 1
grid = (SPLIT_K * base_grid,)
_moe_lora_shrink_splitk_kernel[grid](
hidden_states,
weight,
output,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
N,
K,
topk_ids.numel(),
hidden_states.stride(0),
hidden_states.stride(1),
weight.stride(0),
weight.stride(1),
weight.stride(2),
output.stride(0),
output.stride(1),
top_k=top_k,
BLOCK_SIZE_M=BLOCK_SIZE_M,
BLOCK_SIZE_N=BLOCK_SIZE_N,
BLOCK_SIZE_K=BLOCK_SIZE_K,
GROUP_SIZE_M=GROUP_SIZE_M,
SPLIT_K=SPLIT_K,
num_warps=config.get("num_warps", 4),
num_stages=config.get("num_stages", 4),
)
@torch.compile(dynamic=True)
def _align_block_size_torch(
topk_ids: torch.Tensor,
block_size: int,
num_experts: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Pure-PyTorch align_block_size for num_experts > 1024, compiled via torch.compile."""
device = topk_ids.device
flat_topk_ids = topk_ids.reshape(-1).to(torch.int64)
num_valid_tokens = flat_topk_ids.numel()
max_total_padded_tokens = (
(num_valid_tokens + num_experts * (block_size - 1) + block_size - 1)
// block_size
) * block_size
max_num_blocks = max_total_padded_tokens // block_size
sorted_token_ids = torch.full(
(max_total_padded_tokens,),
num_valid_tokens,
dtype=torch.int32,
device=device,
)
expert_ids = torch.full(
(max_num_blocks,),
-1,
dtype=torch.int32,
device=device,
)
if num_valid_tokens == 0:
num_tokens_post_padded = torch.zeros((1,), dtype=torch.int32, device=device)
return sorted_token_ids, expert_ids, num_tokens_post_padded
sorted_order = torch.argsort(flat_topk_ids)
sorted_expert_ids = flat_topk_ids[sorted_order]
expert_range = torch.arange(num_experts, device=device, dtype=torch.int64)
counts_offsets = torch.searchsorted(sorted_expert_ids, expert_range, right=False)
counts_end = torch.searchsorted(sorted_expert_ids, expert_range, right=True)
counts = counts_end - counts_offsets
padded_counts = ((counts + block_size - 1) // block_size) * block_size
total_padded_tokens = padded_counts.sum().to(torch.int32).reshape(1)
padded_offsets = torch.cumsum(padded_counts, dim=0) - padded_counts
token_ranks = (
torch.arange(num_valid_tokens, device=device, dtype=torch.int64)
- counts_offsets[sorted_expert_ids]
)
output_positions = padded_offsets[sorted_expert_ids] + token_ranks
sorted_token_ids.scatter_(
0,
output_positions.to(torch.int64),
sorted_order.to(torch.int32),
)
block_counts = padded_counts // block_size
actual_num_blocks = block_counts.sum()
if max_num_blocks <= 0:
return sorted_token_ids, expert_ids, total_padded_tokens
block_offsets = torch.cumsum(block_counts, dim=0)
all_block_positions = torch.arange(max_num_blocks, device=device, dtype=torch.int64)
assigned_experts = torch.searchsorted(
block_offsets, all_block_positions, right=True
).to(torch.int32)
expert_ids.copy_(
torch.where(
all_block_positions < actual_num_blocks,
assigned_experts,
torch.full_like(assigned_experts, -1),
)
)
return sorted_token_ids, expert_ids, total_padded_tokens
_align_block_size_large = _align_block_size_torch
def _merged_experts_fused_moe_lora_add_fake(
output: torch.Tensor,
hidden_states: torch.Tensor,
lora_a: torch.Tensor,
lora_b: torch.Tensor,
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
token_lora_mapping: torch.Tensor,
mul_routed_weight: bool,
experts_shared_outer_loras_a: bool,
experts_shared_outer_loras_b: bool,
) -> None:
return
def _merged_experts_fused_moe_lora_add_impl(
output: torch.Tensor,
hidden_states: torch.Tensor,
lora_a: torch.Tensor,
lora_b: torch.Tensor,
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
token_lora_mapping: torch.Tensor,
mul_routed_weight: bool,
experts_shared_outer_loras_a: bool,
experts_shared_outer_loras_b: bool,
routing_cache: dict | None = None,
) -> None:
"""
1. Prepare virtual expert routing metadata from topk_ids + token_lora_mapping * num_experts.
2. Flatten LoRA weights from [max_loras, num_experts, ...] to [max_loras * num_experts, ...].
3. Run regular SGLang fused-MoE kernels for LoRA A and LoRA B.
4. Mask out tokens with token_lora_mapping == -1 on the add path.
"""
max_loras, _, max_lora_rank, _ = lora_a.shape
input_top_k = 1 if hidden_states.shape[0] == topk_ids.numel() else topk_ids.shape[1]
def _merge_lora_expert_weight(t: torch.Tensor) -> torch.Tensor:
# [max_loras, num_experts, x, y] -> [max_loras * num_experts, x, y]
return t.reshape(t.shape[0] * t.shape[1], t.shape[2], t.shape[3])
def _get_stage_config(
weight: torch.Tensor,
stage_top_k: int,
) -> dict[str, Any]:
from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config import (
get_config_dtype_str,
try_get_optimal_moe_config,
)
config_dtype = get_config_dtype_str(dtype=hidden_states.dtype)
get_config_func = functools.partial(
try_get_optimal_moe_config,
weight.shape,
weight.shape,
stage_top_k,
config_dtype,
)
try:
cfg = get_config_func(token_lora_mapping.shape[0])
except ValueError:
K_dim = weight.shape[2]
N_dim = weight.shape[1]
if K_dim >= 1024:
default_block_k = 256
elif K_dim >= 64:
default_block_k = 64
else:
default_block_k = max(16, K_dim)
cfg = {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": min(64, max(16, N_dim)),
"BLOCK_SIZE_K": min(default_block_k, max(16, K_dim)),
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4,
}
return cfg
def _align_block_size(
topk_ids: torch.Tensor,
block_size: int,
num_experts: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
# The native align kernel consumes num_experts + 1 internally for its
# sentinel bucket, so the 1024-expert boundary must use the fallback path.
if num_experts < 1024:
from sglang.srt.layers.moe.fused_moe_triton.moe_align_block_size import (
moe_align_block_size as native_moe_align_block_size,
)
return native_moe_align_block_size(topk_ids, block_size, num_experts)
return _align_block_size_large(topk_ids, block_size, num_experts)
def _get_routing(
topk_ids: torch.Tensor,
token_lora_mapping: torch.Tensor,
num_experts: int,
shared_outer: bool,
block_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
# Check routing_cache for cross-call reuse (gate_up and down share routing)
cache_key = (num_experts, shared_outer, block_size)
if routing_cache is not None:
cached = routing_cache.get(cache_key)
if cached is not None:
return cached
virtual_topk_ids, token_lora_mask, virtual_num_experts = (
_fused_virtual_topk_ids(
topk_ids, token_lora_mapping, num_experts, shared_outer, max_loras
)
)
sorted_token_ids, expert_ids, num_tokens_post_padded = _align_block_size(
virtual_topk_ids,
block_size=block_size,
num_experts=virtual_num_experts,
)
# _align_block_size uses a worst-case padded allocation. Trim the routing buffers
# to a tighter upper bound so we keep the real routed work but drop unused padding
num_tokens = topk_ids.numel()
max_nonempty = min(num_tokens, virtual_num_experts)
tight_padded = (
triton.cdiv(num_tokens + max_nonempty * (block_size - 1), block_size)
* block_size
)
sorted_token_ids = sorted_token_ids[:tight_padded]
expert_ids = expert_ids[: tight_padded // block_size]
expert_ids = fused_sanitize_expert_ids(expert_ids, virtual_num_experts)
result = (
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
token_lora_mask,
)
if routing_cache is not None:
routing_cache[cache_key] = result
return result
from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_kernels import (
invoke_fused_moe_kernel,
)
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]
intermediate = torch.zeros(
[token_lora_mapping.shape[0], topk_ids.shape[1], max_lora_rank],
dtype=hidden_states.dtype,
device=hidden_states.device,
)
a_stage_config = _get_stage_config(lora_a_virtual, input_top_k)
(
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
token_lora_mask,
) = _get_routing(
topk_ids,
token_lora_mapping,
num_experts_a,
experts_shared_outer_loras_a,
a_stage_config["BLOCK_SIZE_M"],
)
_invoke_moe_lora_shrink_splitk(
hidden_states,
lora_a_virtual,
intermediate.view(-1, max_lora_rank),
topk_ids,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
input_top_k,
a_stage_config,
)
b_stage_config = _get_stage_config(lora_b_virtual, 1)
(
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
token_lora_mask,
) = _get_routing(
topk_ids,
token_lora_mapping,
num_experts_b,
experts_shared_outer_loras_b,
b_stage_config["BLOCK_SIZE_M"],
)
invoke_fused_moe_kernel(
intermediate.view(-1, max_lora_rank),
lora_b_virtual,
None,
output,
None,
None,
None,
topk_weights,
topk_ids,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
mul_routed_weight,
1,
b_stage_config,
tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16,
False,
False,
False,
False,
False,
None,
fuse_add_to_output=True,
add_output_mask=token_lora_mask,
router_topk=topk_ids.shape[1],
)
def _merged_experts_fused_moe_lora_add_op(
output: torch.Tensor,
hidden_states: torch.Tensor,
lora_a: torch.Tensor,
lora_b: torch.Tensor,
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
token_lora_mapping: torch.Tensor,
mul_routed_weight: bool,
experts_shared_outer_loras_a: bool,
experts_shared_outer_loras_b: bool,
) -> None:
_merged_experts_fused_moe_lora_add_impl(
output,
hidden_states,
lora_a,
lora_b,
topk_ids,
topk_weights,
token_lora_mapping,
mul_routed_weight,
experts_shared_outer_loras_a,
experts_shared_outer_loras_b,
)
from sglang.srt.utils.common import direct_register_custom_op
direct_register_custom_op(
op_name="merged_experts_fused_moe_lora_add",
op_func=_merged_experts_fused_moe_lora_add_op,
mutates_args=["output"],
fake_impl=_merged_experts_fused_moe_lora_add_fake,
)
def merged_experts_fused_moe_lora_add(
output: torch.Tensor,
hidden_states: torch.Tensor,
lora_a: torch.Tensor,
lora_b: torch.Tensor,
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
token_lora_mapping: torch.Tensor,
mul_routed_weight: bool,
experts_shared_outer_loras_a: bool,
experts_shared_outer_loras_b: bool,
routing_cache: dict | None = None,
) -> None:
"""Public API: wraps the registered op with routing_cache support."""
_merged_experts_fused_moe_lora_add_impl(
output,
hidden_states,
lora_a,
lora_b,
topk_ids,
topk_weights,
token_lora_mapping,
mul_routed_weight,
experts_shared_outer_loras_a,
experts_shared_outer_loras_b,
routing_cache,
)

14
python/sglang/srt/server_args.py
View File

@@ -475,6 +475,7 @@ class ServerArgs:
lora_backend: str = "csgmv"
max_lora_chunk_size: Optional[int] = 16
experts_shared_outer_loras: Optional[bool] = None
lora_use_virtual_experts: bool = False
lora_strict_loading: bool = False
# Kernel backend
@@ -4963,6 +4964,12 @@ class ServerArgs:
"(expert_dim=1). Use --no-experts-shared-outer-loras to force disable. "
"By default this is auto-detected from adapter weights.",
)
parser.add_argument(
"--lora-use-virtual-experts",
default=ServerArgs.lora_use_virtual_experts,
action="store_true",
help="Enable virtual expert computation for MoE models. When set, the model will use virtual expert computation.",
)
parser.add_argument(
"--lora-strict-loading",
default=ServerArgs.lora_strict_loading,
@@ -6715,6 +6722,13 @@ class ServerArgs:
and (self.max_lora_chunk_size & (self.max_lora_chunk_size - 1)) == 0
), "--max-lora-chunk-size must be a power of 2 between 16 and 128."
if self.lora_use_virtual_experts:
assert self.lora_backend == "triton", (
"--lora-use-virtual-experts requires --lora-backend triton. "
f"Got: {self.lora_backend}"
)
logger.info("Virtual expert computation enabled.")
def validate_buckets_rule(self, arg_name: str, buckets_rule: List[str]):
if not buckets_rule:
return