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

[2/n] lora - Shared outer experts and support qwen3_30b_a3b_instruct (#21466)

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Ethan (Yusheng) Su
2026-03-31 14:06:23 -07:00
committed by GitHub
co-authored by Baizhou Zhang
parent f4505e2ee3
commit 3c91ebdf55
8 changed files with 440 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/lora/layers.py
+60 -11
View File
@@ -1,4 +1,4 @@
from typing import Optional
from typing import Dict, Optional, Union
import torch
import torch.nn.functional as F
@@ -711,6 +711,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
# initializes FusedMoE with its own moe_runner for base path
super().__init__(base_layer, lora_backend)
self.experts_shared_outer_loras: bool = False
self.quant_method = base_layer.quant_method
self.tp_size = getattr(base_layer, "moe_tp_size", 1)
self.tp_rank = getattr(base_layer, "moe_tp_rank", 0)
self.intermediate_size_per_partition = getattr(
@@ -782,6 +785,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
adapter_enabled=adapter_enabled,
max_lora_rank=max_lora_rank,
num_experts=self.base_layer.num_experts,
experts_shared_outer_loras=self.experts_shared_outer_loras,
tp_size=self.tp_size,
tp_rank=self.tp_rank,
hidden_size=getattr(self.base_layer, "hidden_size", 0),
@@ -839,34 +843,79 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
return B
def slice_moe_lora_a_weights(
self, A: torch.Tensor, tp_rank: int, target_module: str
) -> torch.Tensor:
self,
A: Union[torch.Tensor, Dict[int, torch.Tensor]],
tp_rank: int,
target_module: str,
):
"""Slice LoRA A weights for MoE with TP.
Accepts:
- 2D tensor [rank, hidden] (single expert)
- 3D tensor [num_experts_or_1, rank, hidden]
- dict {expert_id: 2D tensor}
Per-expert weight shapes:
gate_up_proj_moe A: [rank, hidden_size] — input is full hidden_states, no slice
down_proj_moe A: [rank, intermediate_size] — input is sharded intermediate
"""
if self.tp_size <= 1:
return A
if target_module == "down_proj_moe":
shard_size = self.intermediate_size_per_partition
start = tp_rank * shard_size
end = start + shard_size
return A[:, start:end].contiguous()
return A
if target_module != "down_proj_moe":
return A
if isinstance(A, dict):
return {
eid: self._slice_moe_a(w, tp_rank, target_module)
for eid, w in A.items()
}
return self._slice_moe_a(A, tp_rank, target_module)
def _slice_moe_a(
self, A: torch.Tensor, tp_rank: int, target_module: str
) -> torch.Tensor:
shard_size = self.intermediate_size_per_partition
start = tp_rank * shard_size
end = start + shard_size
return A[..., start:end].contiguous()
def slice_moe_lora_b_weights(
self, B: torch.Tensor, tp_rank: int, target_module: str
) -> torch.Tensor:
self,
B: Union[torch.Tensor, Dict[int, torch.Tensor]],
tp_rank: int,
target_module: str,
):
"""Slice LoRA B weights for MoE with TP.
Accepts:
- 2D tensor [output_dim, rank] (single expert)
- 3D tensor [num_experts_or_1, output_dim, rank]
- dict {expert_id: 2D tensor}
Per-expert weight shapes:
gate_up_proj_moe B: [intermediate_size*2, rank] — output matches sharded base w13
down_proj_moe B: [hidden_size, rank] — output is all-reduced, no slice
"""
if self.tp_size <= 1:
return B
if target_module != "gate_up_proj_moe":
return B
if isinstance(B, dict):
return {
eid: self._slice_moe_b_2d(w, tp_rank, target_module)
for eid, w in B.items()
}
if isinstance(B, torch.Tensor) and B.dim() == 3:
return torch.stack(
[
self._slice_moe_b_2d(B[i], tp_rank, target_module)
for i in range(B.shape[0])
]
)
return self._slice_moe_b_2d(B, tp_rank, target_module)
def _slice_moe_b_2d(
self, B: torch.Tensor, tp_rank: int, target_module: str
) -> torch.Tensor:
if target_module == "gate_up_proj_moe":
shard_size = self.intermediate_size_per_partition
start = tp_rank * shard_size
python/sglang/srt/lora/lora.py
+25 -2
View File
@@ -137,6 +137,8 @@ class LoRAAdapter(nn.Module):
for layer in self.layers:
weight_names = list(layer.weights.keys())
self.normalize_qkv_proj(weight_names, layer.weights)
self._rename_expert_w_to_proj(layer.weights)
weight_names = list(layer.weights.keys())
self.normalize_gate_up_proj(weight_names, layer.weights)
def normalize_qkv_proj(
@@ -192,6 +194,23 @@ class LoRAAdapter(nn.Module):
weights[qkv_name] = weights[qkv_name].repeat(3, 1)
# else: no-op as LoRA B weight is already stacked.
def _rename_expert_w_to_proj(self, weights: Dict[str, torch.Tensor]):
"""Rename w1 -> gate_proj, w3 -> up_proj, w2 -> down_proj so that
normalize_gate_up_proj can stack them into gate_up_proj."""
renames = {}
for name in list(weights.keys()):
new_name = name
if ".w1." in name:
new_name = name.replace(".w1.", ".gate_proj.")
elif ".w3." in name:
new_name = name.replace(".w3.", ".up_proj.")
elif ".w2." in name:
new_name = name.replace(".w2.", ".down_proj.")
if new_name != name:
renames[name] = new_name
for old_name, new_name in renames.items():
weights[new_name] = weights.pop(old_name)
def normalize_gate_up_proj(
self, weight_names: List[str], weights: Dict[str, torch.Tensor]
):
@@ -206,8 +225,9 @@ class LoRAAdapter(nn.Module):
f"Received backend: {self.lora_backend.name}. Please verify your backend configuration "
f"or consider implementing custom initialization logic for other backends."
)
cat_dim = weights[weight_name].dim() - 2
weights[gate_up_name] = torch.cat(
(weights[weight_name], weights[up_name]), 0
(weights[weight_name], weights[up_name]), cat_dim
)
weights.pop(weight_name)
if up_name in weights:
@@ -216,7 +236,10 @@ class LoRAAdapter(nn.Module):
# If gate_up_proj is already stacked, we normalize it following the SGL convention
gate_up_name = weight_name
if "lora_A" in weight_name:
weights[gate_up_name] = weights[gate_up_name].repeat(2, 1)
ndim = weights[gate_up_name].dim()
repeat_dims = [1] * ndim
repeat_dims[ndim - 2] = 2
weights[gate_up_name] = weights[gate_up_name].repeat(*repeat_dims)
# else: no-op as LoRA B weight is already stacked.
def pin_weights_in_cpu(self):
python/sglang/srt/lora/lora_manager.py
+51 -9
View File
@@ -78,8 +78,10 @@ class LoRAManager:
server_args.enable_lora_overlap_loading
)
# Store eviction policy from server args
self.eviction_policy = server_args.lora_eviction_policy
self._experts_shared_outer_override: Optional[bool] = (
server_args.experts_shared_outer_loras
)
# LoRA backend for running sgemm kernels
logger.info(f"Using {lora_backend} as backend of LoRA kernels.")
@@ -303,23 +305,33 @@ class LoRAManager:
if isinstance(module, FusedMoEWithLoRA) and all(
x in self.target_modules for x in ["gate_up_proj", "down_proj"]
):
gate_up_key = (
"gate_up_proj_moe"
if "gate_up_proj_moe" in self.memory_pool.A_buffer
else "gate_up_proj"
)
down_key = (
"down_proj_moe"
if "down_proj_moe" in self.memory_pool.A_buffer
else "down_proj"
)
gate_up_a = self.memory_pool.get_tensor(
target_module="gate_up_proj_moe",
target_module=gate_up_key,
layer_id=layer_id,
lora_type=LoRAType.LORA_A,
)
gate_up_b = self.memory_pool.get_tensor(
target_module="gate_up_proj_moe",
target_module=gate_up_key,
layer_id=layer_id,
lora_type=LoRAType.LORA_B,
)
down_a = self.memory_pool.get_tensor(
target_module="down_proj_moe",
target_module=down_key,
layer_id=layer_id,
lora_type=LoRAType.LORA_A,
)
down_b = self.memory_pool.get_tensor(
target_module="down_proj_moe",
target_module=down_key,
layer_id=layer_id,
lora_type=LoRAType.LORA_B,
)
@@ -387,6 +399,16 @@ class LoRAManager:
target_modules=target_modules,
)
if self._experts_shared_outer_override is not None:
self.experts_shared_outer_loras = self._experts_shared_outer_override
else:
self.experts_shared_outer_loras = self._detect_shared_outer_loras()
if self.experts_shared_outer_loras:
logger.info(
"Shared outer LoRA mode enabled: gate_up lora_A and "
"down lora_B will be shared across experts (expert_dim=1)."
)
self.init_lora_modules()
self.init_memory_pool()
self.update_lora_info()
@@ -412,6 +434,26 @@ class LoRAManager:
f"Failed to load LoRA adapter {lora_ref.lora_name}: {result.error_message}"
)
def _detect_shared_outer_loras(self) -> bool:
"""Auto-detect shared outer LoRA format from loaded adapter weights.
MoE adapters with shared outer experts store 3D tensors where
dim[0]=1 indicates weights shared across all experts, while
dim[0]=num_experts indicates per-expert weights.
Returns True if gate_up lora_A has expert_dim=1 (shared).
"""
for adapter in self.loras.values():
for layer in adapter.layers:
for name, weight in layer.weights.items():
if (
"gate_up_proj" in name
and "lora_A" in name
and weight.dim() == 3
):
return weight.shape[0] == 1
break
return False
def init_lora_shapes(
self,
max_lora_rank: Optional[int] = None,
@@ -589,6 +631,7 @@ class LoRAManager:
base_model=self.base_model,
eviction_policy=self.eviction_policy,
lora_added_tokens_size=self.lora_added_tokens_size,
experts_shared_outer_loras=self.experts_shared_outer_loras,
)
# Initializing memory pool with base model
@@ -683,11 +726,10 @@ class LoRAManager:
)
continue
# Temporarily workaround for FusedMoE layer
if isinstance(module, FusedMoE) and all(
x in self.target_modules for x in ["gate_up_proj", "down_proj"]
):
layer_id = get_layer_id(module_name)
self.lora_modules[layer_id][module_name] = self.set_lora_module(
module_name, module
)
lora_module = self.set_lora_module(module_name, module)
lora_module.experts_shared_outer_loras = self.experts_shared_outer_loras
self.lora_modules[layer_id][module_name] = lora_module
python/sglang/srt/lora/lora_moe_runners.py
+16 -11
View File
@@ -71,17 +71,22 @@ if _is_cuda or _is_hip or _is_xpu:
class LoRAInfo:
"""LoRA weights and dispatch info for MoE computation."""
# LoRA weights: [num_loras, num_experts, dim1, dim2]
# LoRA weights: [num_loras, num_experts_or_1, dim1, dim2]
# When experts_shared_outer_loras=True:
# gate_up_lora_a: [num_loras, 1, max_rank, hidden_dim] (shared)
# down_lora_b: [num_loras, 1, hidden_dim, max_rank] (shared)
gate_up_lora_a_weights: (
torch.Tensor
) # [num_loras, num_experts, max_rank, hidden_dim]
) # [num_loras, num_experts_or_1, max_rank, hidden_dim]
gate_up_lora_b_weights: (
torch.Tensor
) # [num_loras, num_experts, gate_up_dim, max_rank]
down_lora_a_weights: (
torch.Tensor
) # [num_loras, num_experts, max_rank, intermediate_dim]
down_lora_b_weights: torch.Tensor # [num_loras, num_experts, hidden_dim, max_rank]
down_lora_b_weights: (
torch.Tensor
) # [num_loras, num_experts_or_1, hidden_dim, max_rank]
# Indice pointers of each segment in shape (num_segments + 1, )
seg_indptr: torch.Tensor
@@ -95,6 +100,7 @@ class LoRAInfo:
max_lora_rank: int # Maximum LoRA rank across all adapters
num_experts: int
experts_shared_outer_loras: bool = False
fully_sharded: bool = False
tp_size: int = 1
@@ -469,16 +475,11 @@ class TritonRunnerCoreWithLoRA(TritonRunnerCore):
r = lora_info.max_lora_rank
gate_up_a = lora_info.gate_up_lora_a_weights
if lora_info.experts_shared_outer_loras:
gate_up_a = gate_up_a.expand(-1, lora_info.num_experts, -1, -1)
gate_up_b = lora_info.gate_up_lora_b_weights
inter_size = gate_up_b.shape[2] // 2
# Split packed gate_up weights into separate gate and up slices.
# gate_up_lora_a has shape [max_loras, num_experts, 2*r, hidden_dim]
# where the first r rows are gate_lora_a and the next r are up_lora_a.
# gate_up_lora_b has shape [max_loras, num_experts, 2*inter_size, r]
# where the first inter_size rows are gate_lora_b and the rest up_lora_b.
# Using num_slices=2 lets the kernel handle gate and up independently,
# keeping the rank dimension at r so shrink and expand both match.
lora_a_stacked = [gate_up_a[:, :, :r, :], gate_up_a[:, :, r : 2 * r, :]]
lora_b_stacked = [
gate_up_b[:, :, :inter_size, :],
@@ -542,8 +543,12 @@ class TritonRunnerCoreWithLoRA(TritonRunnerCore):
if lora_info.max_lora_rank == 0:
return
down_lora_b = lora_info.down_lora_b_weights
if lora_info.experts_shared_outer_loras:
down_lora_b = down_lora_b.expand(-1, lora_info.num_experts, -1, -1)
lora_a_stacked = [lora_info.down_lora_a_weights]
lora_b_stacked = [lora_info.down_lora_b_weights]
lora_b_stacked = [down_lora_b]
if lora_info.fully_sharded and lora_info.tp_size > 1:
shard_size = lora_info.hidden_size // lora_info.tp_size
python/sglang/srt/lora/mem_pool.py
+124 -54
View File
@@ -60,6 +60,7 @@ class LoRAMemoryPool:
base_model: torch.nn.Module,
eviction_policy: str,
lora_added_tokens_size: int,
experts_shared_outer_loras: bool = False,
):
self.base_hf_config: AutoConfig = base_hf_config
self.num_layer: int = base_hf_config.num_hidden_layers
@@ -70,6 +71,7 @@ class LoRAMemoryPool:
self.lora_added_tokens_size: int = lora_added_tokens_size
self.max_lora_rank: int = max_lora_rank
self.target_modules: Set[str] = target_modules
self.experts_shared_outer_loras: bool = experts_shared_outer_loras
# Initialize eviction policy
self.eviction_policy = get_eviction_policy(eviction_policy)
@@ -140,6 +142,18 @@ class LoRAMemoryPool:
"""Check if module is part of MoE experts."""
return "moe" in module_name
@staticmethod
def _get_num_experts(base_model: torch.nn.Module) -> int:
cfg = base_model.config
if hasattr(cfg, "get_text_config"):
cfg = cfg.get_text_config()
return (
getattr(cfg, "num_experts", None)
or getattr(cfg, "num_local_experts", None)
or getattr(cfg, "n_routed_experts", None)
or 1
)
def _get_standard_shape(
self,
module_name: str,
@@ -178,10 +192,13 @@ class LoRAMemoryPool:
input_dim = divide(input_dim, self.tp_size)
if self.is_moe_module(module_name):
num_experts = base_model.config.num_experts
num_experts = self._get_num_experts(base_model)
expert_dim = num_experts
if self.experts_shared_outer_loras and module_name == "gate_up_proj_moe":
expert_dim = 1
return (
self.max_loras_per_batch,
num_experts,
expert_dim,
max_lora_dim * c,
input_dim,
)
@@ -228,8 +245,11 @@ class LoRAMemoryPool:
# Check if MoE module and return appropriate shape
if self.is_moe_module(module_name):
num_experts = base_model.config.num_experts
return (self.max_loras_per_batch, num_experts, output_dim, max_lora_dim)
num_experts = self._get_num_experts(base_model)
expert_dim = num_experts
if self.experts_shared_outer_loras and module_name == "down_proj_moe":
expert_dim = 1
return (self.max_loras_per_batch, expert_dim, output_dim, max_lora_dim)
else:
return (self.max_loras_per_batch, output_dim, max_lora_dim)
@@ -264,32 +284,33 @@ class LoRAMemoryPool:
target_modules: Set[str],
get_lora_shape_fn: Callable[[str, torch.nn.Module, int, int], Tuple[int]],
):
# Check if model has both shared experts and MoE experts
cfg = base_model.config
if hasattr(cfg, "get_text_config"):
cfg = cfg.get_text_config()
has_shared_experts = (
hasattr(base_model.config, "shared_expert_intermediate_size")
and base_model.config.shared_expert_intermediate_size > 0
)
has_moe = getattr(base_model.config, "num_experts", 1) > 1
hasattr(cfg, "shared_expert_intermediate_size")
and cfg.shared_expert_intermediate_size > 0
) or (getattr(cfg, "n_shared_experts", 0) or 0) > 0
has_moe = self._get_num_experts(base_model) > 1
# Shape functions automatically handle both 3D (standard) and 4D (MoE)
target_modules = target_modules - set(EMBEDDING_NAMES)
for module_name in target_modules:
# Special handling for ambiguous target modules that can be in different contexts
ambiguous_modules = {"gate_up_proj", "down_proj"}
if module_name in ambiguous_modules and has_shared_experts and has_moe:
# Allocate separate buffers for shared and MoE contexts
# Shared expert version (3D)
shared_key = module_name
buffer[shared_key] = [
torch.empty(
get_lora_shape_fn(
module_name, base_model, self.max_lora_rank, idx
),
dtype=self.dtype,
device=device,
)
for idx in range(self.num_layer)
]
if module_name in ambiguous_modules and has_moe:
# Allocate shared expert version (3D) only when model has shared experts
if has_shared_experts:
buffer[module_name] = [
torch.zeros(
get_lora_shape_fn(
module_name, base_model, self.max_lora_rank, idx
),
dtype=self.dtype,
device=device,
)
for idx in range(self.num_layer)
]
# MoE expert version (4D)
moe_key = f"{module_name}_moe"
@@ -521,8 +542,8 @@ class LoRAMemoryPool:
expert_match = re.search(r"experts\.(\d+)\.", name)
if expert_match:
# Per-expert MoE weight — 2D tensors, one per expert
target_module = target_module + "_moe"
# MoE weight - multiple tensors per module (one per expert)
if temp_A_buffer[target_module] is None:
temp_A_buffer[target_module] = {}
temp_B_buffer[target_module] = {}
@@ -532,8 +553,15 @@ class LoRAMemoryPool:
temp_A_buffer[target_module][expert_id] = weights
else:
temp_B_buffer[target_module][expert_id] = weights
elif "experts" in name and weights.dim() == 3:
# Shared outer MoE weight — 3D tensor [expert_dim, rank, hidden]
target_module = target_module + "_moe"
if "lora_A" in name:
temp_A_buffer[target_module] = weights
else:
temp_B_buffer[target_module] = weights
else:
# Standard weight - single tensor per module
# Standard weight — single tensor per module
if "lora_A" in name:
temp_A_buffer[target_module] = weights
else:
@@ -549,20 +577,18 @@ class LoRAMemoryPool:
if isinstance(module, FusedMoEWithLoRA):
moe_target_modules = ["gate_up_proj_moe", "down_proj_moe"]
for target_module in moe_target_modules:
if temp_A_buffer[target_module] is None:
continue
for expert_id in temp_A_buffer[target_module].keys():
temp_A_buffer[target_module][expert_id] = (
if temp_A_buffer.get(target_module) is not None:
temp_A_buffer[target_module] = (
module.slice_moe_lora_a_weights(
temp_A_buffer[target_module][expert_id],
temp_A_buffer[target_module],
self.tp_rank,
target_module,
)
)
temp_B_buffer[target_module][expert_id] = (
if temp_B_buffer.get(target_module) is not None:
temp_B_buffer[target_module] = (
module.slice_moe_lora_b_weights(
temp_B_buffer[target_module][expert_id],
temp_B_buffer[target_module],
self.tp_rank,
target_module,
)
@@ -587,22 +613,42 @@ class LoRAMemoryPool:
temp_B_buffer[target_module], self.tp_rank
)
# Load weights into buffers (handles both 3D standard and 4D MoE)
for name, weights in temp_A_buffer.items():
c = get_stacked_multiply(name)
target_buffer = self.A_buffer[name][layer_id]
if name in ["gate_up_proj_moe", "down_proj_moe"]:
# MoE: multiple tensors per module (one per expert)
for expert_id, expert_weight in weights.items():
# Buffer shape: [num_loras, num_experts, max_rank, hidden_dim]
buffer_view = target_buffer[
buffer_id, expert_id, : lora_rank * c, :
]
load_lora_weight_tensor(buffer_view, expert_weight)
if self.experts_shared_outer_loras and name == "gate_up_proj_moe":
if isinstance(weights, torch.Tensor) and weights.dim() == 3:
buffer_view = target_buffer[
buffer_id, 0, : lora_rank * c, :
]
load_lora_weight_tensor(buffer_view, weights[0])
elif isinstance(weights, dict) and len(weights) > 0:
rep = next(iter(weights.values()))
buffer_view = target_buffer[
buffer_id, 0, : lora_rank * c, :
]
load_lora_weight_tensor(buffer_view, rep)
else:
raise ValueError(
f"Unexpected weight format for shared outer gate_up_proj_moe lora_A: "
f"type={type(weights)}, "
f"shape={weights.shape if isinstance(weights, torch.Tensor) else 'N/A'}"
)
elif isinstance(weights, torch.Tensor) and weights.dim() == 3:
for eid in range(weights.shape[0]):
buffer_view = target_buffer[
buffer_id, eid, : lora_rank * c, :
]
load_lora_weight_tensor(buffer_view, weights[eid])
elif isinstance(weights, dict):
for expert_id, expert_weight in weights.items():
buffer_view = target_buffer[
buffer_id, expert_id, : lora_rank * c, :
]
load_lora_weight_tensor(buffer_view, expert_weight)
else:
# Standard: single tensor per module
c = get_stacked_multiply(name)
buffer_view = target_buffer[buffer_id, : lora_rank * c, :]
load_lora_weight_tensor(buffer_view, weights)
@@ -610,18 +656,42 @@ class LoRAMemoryPool:
target_buffer = self.B_buffer[name][layer_id]
if name in ["gate_up_proj_moe", "down_proj_moe"]:
# MoE: multiple tensors per module (one per expert)
for expert_id, expert_weight in weights.items():
# Buffer shape: [num_loras, num_experts, intermediate_dim, max_rank]
buffer_view = target_buffer[buffer_id, expert_id, :, :lora_rank]
weight_to_load = expert_weight
if weight_to_load is not None:
weight_to_load = weight_to_load * lora_adapter.scaling
load_lora_weight_tensor(buffer_view, weight_to_load)
if self.experts_shared_outer_loras and name == "down_proj_moe":
if isinstance(weights, torch.Tensor) and weights.dim() == 3:
buffer_view = target_buffer[buffer_id, 0, :, :lora_rank]
w = weights[0]
if w is not None:
w = w * lora_adapter.scaling
load_lora_weight_tensor(buffer_view, w)
elif isinstance(weights, dict) and len(weights) > 0:
rep = next(iter(weights.values()))
buffer_view = target_buffer[buffer_id, 0, :, :lora_rank]
if rep is not None:
rep = rep * lora_adapter.scaling
load_lora_weight_tensor(buffer_view, rep)
else:
raise ValueError(
f"Unexpected weight format for shared outer down_proj_moe lora_B: "
f"type={type(weights)}, "
f"shape={weights.shape if isinstance(weights, torch.Tensor) else 'N/A'}"
)
elif isinstance(weights, torch.Tensor) and weights.dim() == 3:
for eid in range(weights.shape[0]):
buffer_view = target_buffer[buffer_id, eid, :, :lora_rank]
w = weights[eid]
if w is not None:
w = w * lora_adapter.scaling
load_lora_weight_tensor(buffer_view, w)
elif isinstance(weights, dict):
for expert_id, expert_weight in weights.items():
buffer_view = target_buffer[
buffer_id, expert_id, :, :lora_rank
]
w = expert_weight
if w is not None:
w = w * lora_adapter.scaling
load_lora_weight_tensor(buffer_view, w)
else:
# Standard: single tensor per module
buffer_view = target_buffer[buffer_id, :, :lora_rank]
load_lora_weight_tensor(buffer_view, weights)
python/sglang/srt/lora/triton_ops/sgemm_lora_b.py
+4 -3
View File
@@ -87,6 +87,7 @@ def _sgemm_lora_b_kernel(
)
# Iterate to compute the block in output matrix
n_mask = n_offset[None, :] < N
partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_K)):
x_tile = tl.load(
@@ -96,7 +97,7 @@ def _sgemm_lora_b_kernel(
)
w_tile = tl.load(
w_ptrs,
mask=(k_offset[:, None] < K - k * BLOCK_K),
mask=(k_offset[:, None] < K - k * BLOCK_K) & n_mask,
other=0.0,
)
partial_sum += tl.dot(x_tile, w_tile)
@@ -110,8 +111,8 @@ def _sgemm_lora_b_kernel(
output_ptr = (output + seg_start * output_stride_0) + (
s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
)
output_mask = s_offset[:, None] < seg_len
partial_sum += tl.load(output_ptr, mask=output_mask)
output_mask = (s_offset[:, None] < seg_len) & n_mask
partial_sum += tl.load(output_ptr, mask=output_mask, other=0.0)
tl.store(output_ptr, partial_sum, mask=output_mask)
python/sglang/srt/server_args.py
+9
View File
@@ -464,6 +464,7 @@ class ServerArgs:
lora_eviction_policy: str = "lru"
lora_backend: str = "csgmv"
max_lora_chunk_size: Optional[int] = 16
experts_shared_outer_loras: Optional[bool] = None
# Kernel backend
attention_backend: Optional[str] = None
@@ -4595,6 +4596,14 @@ class ServerArgs:
choices=[16, 32, 64, 128],
help="Maximum chunk size for the ChunkedSGMV LoRA backend. Only used when --lora-backend is 'csgmv'. Choosing a larger value might improve performance.",
)
parser.add_argument(
"--experts-shared-outer-loras",
default=ServerArgs.experts_shared_outer_loras,
action="store_true",
help="Force shared outer LoRA mode for MoE models. "
"When set, w1/w3 lora_A and w2 lora_B are shared across experts "
"(expert_dim=1). By default this is auto-detected from adapter weights.",
)
# Kernel backend
parser.add_argument(