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

[lora][moe] Decoupled LoRA MoE backend with Marlin support (#21858)

Browse Source
This commit is contained in:
Kurt Shuster
2026-04-11 14:59:27 -07:00
committed by GitHub
parent 78043d4448
commit 8da1cfb30d
12 changed files with 1629 additions and 543 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/base.py
+6 -2
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable, Optional, Tuple, TypeGuard
from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, TypeGuard
import torch
@@ -82,7 +82,11 @@ class MoeRunnerCore(ABC):
@abstractmethod
def run(
self, runner_input: RunnerInput, quant_info: MoeQuantInfo, running_state: dict
self,
runner_input: RunnerInput,
quant_info: MoeQuantInfo,
running_state: dict,
hooks: Optional[Any] = None,
) -> RunnerOutput:
pass
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
+2 -1
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, List, Optional
from typing import TYPE_CHECKING, Any, List, Optional
import torch
@@ -120,6 +120,7 @@ class DeepGemmRunnerCore(MoeRunnerCore):
runner_input: DeepGemmRunnerInput,
quant_info: DeepGemmMoeQuantInfo,
running_state: dict,
hooks: Optional[Any] = None,
) -> DeepGemmRunnerOutput:
if not runner_input.use_masked_gemm:
hidden_states = self._run_contiguous_gemm(
python/sglang/srt/layers/moe/moe_runner/runner.py
+49 -17
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
import logging
import os
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING, Any, Optional
from sglang.srt.layers.moe.moe_runner.base import (
FusedOpPool,
@@ -19,6 +19,8 @@ if TYPE_CHECKING:
from sglang.srt.layers.moe.moe_runner.base import MoeQuantInfo
from sglang.srt.layers.moe.token_dispatcher.base import CombineInput, DispatchOutput
from sglang.srt.layers.moe.utils import MoeRunnerBackend
from sglang.srt.lora.lora_moe_runners import LoRAHooks
logger = logging.getLogger(__name__)
@@ -37,18 +39,18 @@ class MoeRunner:
self.fused_func = None
if runner_backend.is_triton():
if lora_enabled:
from sglang.srt.lora.lora_moe_runners import TritonRunnerCoreWithLoRA
self.runner_core = TritonRunnerCoreWithLoRA(config)
else:
self.runner_core = TritonRunnerCore(config)
self.runner_core = TritonRunnerCore(config)
elif runner_backend.is_triton_kernels():
self.runner_core = TritonKernelsRunnerCore(config)
elif runner_backend.is_deep_gemm():
self.runner_core = DeepGemmRunnerCore(config)
elif runner_backend.is_marlin():
self.runner_core = None # Marlin only supports fused path
if lora_enabled:
from sglang.srt.lora.lora_moe_runner_marlin import MarlinLoraRunnerCore
self.runner_core = MarlinLoraRunnerCore(config)
else:
self.runner_core = None # Marlin only supports fused path
elif (
runner_backend.is_flashinfer_trtllm()
or runner_backend.is_flashinfer_trtllm_routed()
@@ -94,6 +96,41 @@ class MoeRunner:
return self.fused_func(dispatch_output, quant_info, self.config)
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
if isinstance(_runner_input, DispatchOutput):
hidden_states, topk_ids = (
_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,
)
# Runners that handle dispatch_output directly (e.g., MarlinRunnerCore)
# bypass the pre-permute step and do their own alignment internally.
if hasattr(self.runner_core, "run_from_dispatch"):
hooks = _maybe_build_lora_hooks(dispatch_output)
return self.runner_core.run_from_dispatch(
dispatch_output, quant_info, self.config, hooks=hooks
)
dispatch_format = dispatch_output.format.value
runner_format = self.runner_core.runner_backend.value
self.pre_permute_func = PermuteMethodPool.get_pre_permute(
@@ -110,16 +147,11 @@ class MoeRunner:
dispatch_output, quant_info, self.config, running_state
)
# Pass lora_info to runner_core if LoRA is enabled
if self.lora_enabled:
runner_output = self.runner_core.run(
runner_input, quant_info, running_state, lora_info
)
else:
runner_output = self.runner_core.run(
runner_input, quant_info, running_state
)
hooks = _maybe_build_lora_hooks(runner_input)
runner_output = self.runner_core.run(
runner_input, quant_info, running_state, hooks=hooks
)
runner_format = self.runner_core.runner_backend.value
combine_format = dispatch_output.format.value
self.post_permute_func = PermuteMethodPool.get_post_permute(
python/sglang/srt/layers/moe/moe_runner/triton.py
+22 -4
View File
@@ -3,7 +3,7 @@ from __future__ import annotations
import functools
import os
from dataclasses import dataclass
from typing import TYPE_CHECKING, List, Optional
from typing import TYPE_CHECKING, Any, List, Optional
import torch
import triton.language as tl
@@ -124,6 +124,7 @@ class TritonRunnerCore(MoeRunnerCore):
runner_input: TritonRunnerInput,
quant_info: TritonMoeQuantInfo,
running_state: dict,
hooks: Optional[Any] = None,
) -> TritonRunnerOutput:
# TODO: move these functions to the triton runner
@@ -206,6 +207,11 @@ class TritonRunnerCore(MoeRunnerCore):
block_shape=block_shape,
)
if hooks and hooks.after_gate_up:
hooks.after_gate_up(
hidden_states, intermediate_cache1, topk_weights, topk_ids
)
intermediate_cache2 = torch.empty(
(M * topk_ids.shape[1], N // 2),
device=hidden_states.device,
@@ -258,13 +264,16 @@ class TritonRunnerCore(MoeRunnerCore):
else:
out_hidden_states = torch.empty_like(hidden_states)
# When LoRA hooks are present, always write to intermediate_cache3
# so the hook can modify it before reduction.
_use_intermediate = not no_combine and (topk_ids.shape[1] != 1 or hooks)
invoke_fused_moe_kernel(
intermediate_cache2,
w2,
b2,
(
intermediate_cache3
if not no_combine and topk_ids.shape[1] != 1
if _use_intermediate
else out_hidden_states.unsqueeze(0)
),
a2_scale,
@@ -287,14 +296,23 @@ class TritonRunnerCore(MoeRunnerCore):
block_shape=block_shape,
)
if hooks and hooks.after_down:
hooks.after_down(
intermediate_cache2, intermediate_cache3, topk_weights, topk_ids
)
if routed_scaling_factor is None:
routed_scaling_factor = 1.0
if no_combine:
pass
elif _is_cuda:
if topk_ids.shape[1] == 1 and routed_scaling_factor == 1.0:
pass # we write directly into out_hidden_states
if (
topk_ids.shape[1] == 1
and routed_scaling_factor == 1.0
and not _use_intermediate
):
pass # we wrote directly into out_hidden_states
elif topk_ids.shape[1] == 2 and routed_scaling_factor == 1.0:
torch.add(
intermediate_cache3[:, 0],
python/sglang/srt/layers/moe/moe_runner/triton_kernels.py
+2 -1
View File
@@ -3,7 +3,7 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING, Any, Optional
import torch
@@ -84,6 +84,7 @@ class TritonKernelsRunnerCore(MoeRunnerCore):
runner_input: TritonKernelsRunnerInput,
quant_info: TritonKernelsQuantInfo,
running_state: dict,
hooks: Optional[Any] = None,
) -> TritonKernelsRunnerOutput:
from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import (
triton_kernel_fused_experts,
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
+7 -7
View File
@@ -682,16 +682,13 @@ class CompressedTensorsConfig(QuantizationConfig):
logger.info_once("Using CompressedTensorsWNA16TritonMoE (ROCm)")
return CompressedTensorsWNA16TritonMoE(self)
else:
from sglang.srt.server_args import get_global_server_args
server_args = get_global_server_args()
if server_args and server_args.enable_lora:
moe_backend = get_moe_runner_backend()
if moe_backend.is_triton():
logger.info_once(
"Using CompressedTensorsWNA16TritonMoEMethod "
"(LoRA requires triton-compatible MoE weights)"
"Using CompressedTensorsWNA16TritonMoE "
"(moe_runner_backend=triton)"
)
return CompressedTensorsWNA16TritonMoE(self)
logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
return CompressedTensorsWNA16MoE(self)
else:
@@ -1010,6 +1007,9 @@ class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase):
def get_triton_quant_info(self, layer: torch.nn.Module):
return layer.scheme.get_triton_quant_info(layer)
def get_marlin_quant_info(self, layer: torch.nn.Module):
return layer.scheme.get_marlin_quant_info(layer)
def apply(
self,
layer: torch.nn.Module,
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py
+17
View File
@@ -354,6 +354,23 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme):
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
):
self.moe_runner_config = moe_runner_config
self.runner = MoeRunner(MoeRunnerBackend.MARLIN, moe_runner_config)
def get_marlin_quant_info(self, layer):
from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
return MarlinMoeQuantInfo(
w13_qweight=layer.w13_weight_packed,
w2_qweight=layer.w2_weight_packed,
w13_scales=layer.w13_weight_scale,
w2_scales=layer.w2_weight_scale,
w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None),
w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None),
weight_bits=self.num_bits,
w13_g_idx=getattr(layer, "w13_weight_g_idx", None),
w2_g_idx=getattr(layer, "w2_weight_g_idx", None),
is_k_full=self.is_k_full,
)
def apply_weights(
self,
python/sglang/srt/lora/layers.py
+30 -7
View File
@@ -808,13 +808,20 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
getattr(base_layer.moe_runner_config, "gemm1_alpha", None) is not None
)
# initialize triton_lora moe runner for batches with lora enabled
# Initialize triton_lora moe runner for batches with lora enabled
from sglang.srt.layers.moe import MoeRunnerBackend
from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
from sglang.srt.layers.moe.utils import get_moe_runner_backend
qm = base_layer.quant_method
if hasattr(qm, "runner") and qm.runner is not None:
runner_backend = qm.runner.runner_backend
# Determine runner backend: prefer server arg, fall back to quant method's runner
global_backend = get_moe_runner_backend()
if not global_backend.is_auto():
runner_backend = global_backend
elif (
hasattr(base_layer.quant_method, "runner")
and base_layer.quant_method.runner is not None
):
runner_backend = base_layer.quant_method.runner.runner_backend
else:
runner_backend = MoeRunnerBackend.TRITON
@@ -824,8 +831,25 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
lora_enabled=True,
)
# Pre-compute quant info for efficiency (weights don't change during inference)
self._quant_info = base_layer.quant_method.get_triton_quant_info(base_layer)
if runner_backend.is_marlin():
from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
CompressedTensorsFusedMoEMethod,
)
assert isinstance(
base_layer.quant_method, CompressedTensorsFusedMoEMethod
), (
f"Marlin MoE backend requires CompressedTensorsFusedMoEMethod, "
f"got {type(base_layer.quant_method).__name__}"
)
self._quant_info = base_layer.quant_method.get_marlin_quant_info(base_layer)
elif runner_backend.is_triton():
assert base_layer.quant_method is not None, "Quant method must be set"
self._quant_info = base_layer.quant_method.get_triton_quant_info(base_layer)
else:
raise NotImplementedError(
f"LoRA MoE not supported for backend {runner_backend}"
)
def set_lora_info(
self,
@@ -876,7 +900,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
num_experts=self.base_layer.num_experts,
experts_shared_outer_loras=self.experts_shared_outer_loras,
cg_buffers=cg_buffers,
has_active_lora=batch_info.has_active_lora,
tp_size=self.tp_size,
tp_rank=self.tp_rank,
hidden_size=getattr(self.base_layer, "hidden_size", 0),
python/sglang/srt/lora/lora_moe_runner_marlin.py
+206
View File
@@ -0,0 +1,206 @@
"""Marlin MoE runner core with hook support for LoRA injection.
Uses Marlin int4/int8 kernels for the base MoE projections.
LoRA deltas are injected via hooks.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
from sglang.srt.utils import is_cuda
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import (
StandardCombineInput,
StandardDispatchOutput,
)
_is_cuda = is_cuda()
if _is_cuda:
from sgl_kernel import silu_and_mul
from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm
from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import (
get_scalar_type,
)
from sglang.srt.layers.moe.fused_moe_triton.fused_moe import (
moe_align_block_size,
)
from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_kernels import (
moe_sum_reduce_triton,
)
from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace
_MARLIN_WORKSPACE: Optional[torch.Tensor] = None
class MarlinLoraRunnerCore:
"""
MoE runner using Marlin kernels for base projections, with hooks for LoRA.
Pipeline:
1. moe_wna16_marlin_gemm (gate_up)
1.5. hooks.after_gate_up
2. silu_and_mul
3. moe_wna16_marlin_gemm (down)
3.5. hooks.after_down
4. moe_sum_reduce
"""
def __init__(self, config: MoeRunnerConfig):
self.config = config
def run_from_dispatch(
self,
dispatch_output: StandardDispatchOutput,
quant_info: MarlinMoeQuantInfo,
runner_config: MoeRunnerConfig,
hooks=None,
) -> StandardCombineInput:
global _MARLIN_WORKSPACE
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
assert hooks is not None, "hooks must be provided for MarlinLoraRunnerCore"
hidden_states = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
topk_weights = topk_output.topk_weights
topk_ids = topk_output.topk_ids
assert runner_config.activation == "silu", "Only SiLU activation is supported."
assert (
torch.cuda.get_device_capability(hidden_states.device)[0] >= 9
), "MarlinLoraRunnerCore requires CUDA compute capability >= 9"
inplace = runner_config.inplace
routed_scaling_factor = runner_config.routed_scaling_factor
M, K = hidden_states.shape
E = quant_info.w13_qweight.shape[0]
N = quant_info.w2_qweight.shape[1] * 16
topk = topk_ids.shape[1]
num_bits = quant_info.weight_bits
for block_size_m in [8, 16, 32, 48, 64]:
if M * topk / E / block_size_m < 0.9:
break
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
topk_ids, block_size_m, E
)
if (
_MARLIN_WORKSPACE is None
or _MARLIN_WORKSPACE.device != hidden_states.device
):
_MARLIN_WORKSPACE = marlin_make_workspace(
hidden_states.device, max_blocks_per_sm=4
)
workspace = _MARLIN_WORKSPACE
scalar_type1 = get_scalar_type(num_bits, quant_info.w13_qzeros is not None)
scalar_type2 = get_scalar_type(num_bits, quant_info.w2_qzeros is not None)
# Stage 1: Gate/Up (Marlin)
intermediate_cache1 = torch.empty(
(M * topk, 2 * N), device=hidden_states.device, dtype=hidden_states.dtype
)
intermediate_cache1 = moe_wna16_marlin_gemm(
hidden_states,
intermediate_cache1,
quant_info.w13_qweight,
None,
quant_info.w13_scales,
None,
quant_info.w13_qzeros,
quant_info.w13_g_idx,
quant_info.w13_g_idx_sort_indices,
workspace,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
topk_weights,
moe_block_size=block_size_m,
top_k=topk,
mul_topk_weights=False,
is_ep=quant_info.expert_map is not None,
b_q_type=scalar_type1,
size_m=M,
size_n=2 * N,
size_k=K,
is_k_full=quant_info.is_k_full,
use_atomic_add=True,
use_fp32_reduce=True,
is_zp_float=False,
)
# Hook: after gate_up
if hooks.after_gate_up:
intermediate_cache1_3d = intermediate_cache1.view(M, topk, 2 * N)
hooks.after_gate_up(
hidden_states, intermediate_cache1_3d, topk_weights, topk_ids
)
# Stage 2: Activation
intermediate_cache2 = torch.empty(
(M * topk, N), device=hidden_states.device, dtype=hidden_states.dtype
)
silu_and_mul(intermediate_cache1.view(-1, 2 * N), intermediate_cache2)
# Stage 3: Down (Marlin)
intermediate_cache3 = torch.empty(
(M * topk, K), device=hidden_states.device, dtype=hidden_states.dtype
)
if quant_info.expert_map is not None:
intermediate_cache3.zero_()
intermediate_cache3 = moe_wna16_marlin_gemm(
intermediate_cache2,
intermediate_cache3,
quant_info.w2_qweight,
None,
quant_info.w2_scales,
None,
quant_info.w2_qzeros,
quant_info.w2_g_idx,
quant_info.w2_g_idx_sort_indices,
workspace,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
topk_weights,
moe_block_size=block_size_m,
top_k=1,
mul_topk_weights=True,
is_ep=quant_info.expert_map is not None,
b_q_type=scalar_type2,
size_m=M * topk,
size_n=K,
size_k=N,
is_k_full=quant_info.is_k_full,
use_atomic_add=True,
use_fp32_reduce=True,
is_zp_float=False,
)
intermediate_cache3 = intermediate_cache3.view(M, topk, K)
# Hook: after down
if hooks.after_down:
hooks.after_down(
intermediate_cache2, intermediate_cache3, topk_weights, topk_ids
)
# Stage 4: Reduction
output = hidden_states if inplace else torch.empty_like(hidden_states)
if routed_scaling_factor is None:
routed_scaling_factor = 1.0
# NOTE: fusion opportunity here
moe_sum_reduce_triton(intermediate_cache3, output, routed_scaling_factor)
return StandardCombineInput(hidden_states=output)
python/sglang/srt/lora/lora_moe_runners.py
+365 -504
View File
@@ -12,62 +12,140 @@
# limitations under the License.
# ==============================================================================
"""LoRA-aware MoE runners that integrate LoRA deltas into the MoE computation.
"""LoRA hooks for MoE runners.
The key insight is that LoRA deltas must be added at specific points:
1. After gate_up projection, BEFORE activation (halfway through)
2. After down projection, BEFORE final reduction (at the end)
LoRA deltas are injected at two points in the MoE pipeline:
1. After gate_up projection, BEFORE activation
2. After down projection, BEFORE final reduction
This differs from computing LoRA independently and adding at the very end.
This module provides hook closures that any MoE backend can call at those points,
without needing a per-backend LoRA runner subclass.
"""
from __future__ import annotations
import os
from dataclasses import dataclass
from typing import Optional
from typing import Callable
import torch
import triton.language as tl
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.triton import (
TritonMoeQuantInfo,
TritonRunnerCore,
TritonRunnerInput,
TritonRunnerOutput,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.utils import cpu_has_amx_support, is_cpu, is_cuda, is_hip, is_xpu
from sglang.srt.utils import is_cuda, is_hip, is_xpu, next_power_of_2
_is_hip = is_hip()
_is_cuda = is_cuda()
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu()
_use_aiter = bool(int(os.getenv("SGLANG_USE_AITER", "0")))
_is_hip = is_hip()
_is_xpu = is_xpu()
_MOE_PADDING_SIZE = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0
if _is_cuda or _is_hip:
from sgl_kernel import gelu_and_mul, silu_and_mul
if _is_hip:
from vllm import _custom_ops as vllm_ops # moe_sum
elif _is_cpu and _is_cpu_amx_available:
pass
elif _is_xpu:
from sgl_kernel import silu_and_mul
if _is_cuda or _is_hip or _is_xpu:
from sgl_kernel import ( # noqa: F401
moe_align_block_size as sgl_moe_align_block_size,
)
from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size
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,
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),
)
@dataclass
class LoRAInfo:
"""LoRA weights and dispatch info for MoE computation."""
@@ -102,8 +180,7 @@ class LoRAInfo:
num_experts: int
experts_shared_outer_loras: bool = False
cg_buffers: Optional[dict] = None
has_active_lora: bool = False
cg_buffers: dict | None = None
fully_sharded: bool = False
tp_size: int = 1
@@ -111,168 +188,65 @@ class LoRAInfo:
hidden_size: int = 0
class TritonRunnerCoreWithLoRA(TritonRunnerCore):
@dataclass
class LoRAHooks:
"""Hook callbacks for injecting LoRA deltas into the MoE pipeline."""
after_gate_up: (
Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None] | None
) = None
after_down: (
Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None] | None
) = None
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.
Returns: (sorted_token_ids_reshaped, expert_ids_reshaped, num_tokens_post_padded_lora, lora_ids)
"""
LoRA-aware wrapper around TritonRunnerCore.
cg = lora_info.cg_buffers if get_is_capture_mode() else None
shrink_config = {"BLOCK_SIZE_M": 64}
M = topk_ids.shape[0]
block_size_m = shrink_config["BLOCK_SIZE_M"]
max_loras = len(lora_info.lora_ranks)
Integrates LoRA deltas at the correct points in the MoE forward pass:
1. Base gate_up projection + LoRA gate_up delta -> activation
2. Base down projection + LoRA down delta -> final reduction
max_num_tokens_padded = topk_ids.numel() + lora_info.num_experts * (
block_size_m - 1
)
max_num_tokens_padded = (
(max_num_tokens_padded + block_size_m - 1) // block_size_m
) * block_size_m
max_num_m_blocks = (max_num_tokens_padded + block_size_m - 1) // block_size_m
This follows the vLLM/HF approach where LoRA is fused into the computation
rather than computed independently.
"""
device = topk_ids.device
def __init__(self, config: MoeRunnerConfig):
super().__init__(config)
use_naive = (
cg is None
and M * topk_ids.shape[1] * _SPARSITY_FACTOR
<= lora_info.num_experts * max_loras
)
def run(
self,
runner_input: TritonRunnerInput,
quant_info: TritonMoeQuantInfo,
running_state: dict,
lora_info: Optional[LoRAInfo] = None,
) -> TritonRunnerOutput:
"""
Run MoE with integrated LoRA computation.
This method extends TritonRunnerCore.run() by inserting LoRA delta
computations at the correct points in the MoE forward pass.
Args:
runner_input: Standard Triton runner input
quant_info: Quantization info for base weights
running_state: Running state dict
lora_info: Optional LoRA weights and dispatch info
Returns:
TritonRunnerOutput with combined base + LoRA output
"""
if lora_info is None:
return super().run(runner_input, quant_info, running_state)
if get_is_capture_mode():
# During CUDA graph capture, always enter the LoRA path so that
# the LoRA kernels are recorded in the graph. adapter_enabled is
# all-zeros during capture, so the Triton kernel early-exits per
# program (zero overhead). During replay the tensor is updated
# in-place with the real adapter mask before graph.replay().
has_active_lora = True
else:
has_active_lora = lora_info.has_active_lora
if not has_active_lora:
return super().run(runner_input, quant_info, running_state)
# Extract common variables
hidden_states = runner_input.hidden_states
topk_weights = runner_input.topk_weights
topk_ids = runner_input.topk_ids
sorted_token_ids = runner_input.sorted_token_ids
expert_ids = runner_input.expert_ids
num_tokens_post_padded = runner_input.num_tokens_post_padded
w13 = quant_info.w13_weight
w2 = quant_info.w2_weight
b13 = quant_info.b13
b2 = quant_info.b2
a13_scale = quant_info.a13_scale
a2_scale = quant_info.a2_scale
w13_scale = quant_info.w13_scale
w2_scale = quant_info.w2_scale
w13_zp = quant_info.w13_zp
w2_zp = quant_info.w2_zp
block_shape = quant_info.block_shape
per_channel_quant = quant_info.per_channel_quant
use_fp8_w8a8 = quant_info.use_fp8_w8a8
use_int8_w8a8 = quant_info.use_int8_w8a8
use_int8_w8a16 = quant_info.use_int8_w8a16
use_int4_w4a16 = quant_info.use_int4_w4a16
activation = self.config.activation
no_combine = self.config.no_combine
inplace = self.config.inplace
gemm1_alpha = self.config.gemm1_alpha
gemm1_limit = self.config.gemm1_clamp_limit
routed_scaling_factor = self.config.routed_scaling_factor
apply_router_weight_on_input = self.config.apply_router_weight_on_input
assert self.config.is_gated, "Only gated MoEs are supported for Triton runner"
M = hidden_states.shape[0]
E, N, _ = w13.shape
compute_type = (
tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16
)
# TODO: move these functions to the triton runner
from sglang.srt.layers.moe.fused_moe_triton.fused_moe import (
_swiglu_gpt_oss_sigmoid_alpha,
_swiglu_silu_clamp_mul,
invoke_fused_moe_kernel,
moe_sum_reduce_torch_compile,
moe_sum_reduce_triton,
)
cg = lora_info.cg_buffers if get_is_capture_mode() else None
# ============================================================
# Stage 1: Gate/Up projection (base)
# ============================================================
if cg is not None:
intermediate_cache1 = cg["intermediate_cache1"][:M, : topk_ids.shape[1], :N]
else:
intermediate_cache1 = torch.empty(
(M, topk_ids.shape[1], N),
device=hidden_states.device,
dtype=hidden_states.dtype,
if use_naive:
sorted_token_ids_lora, expert_ids_lora, num_tokens_post_padded_lora = (
_naive_moe_lora_align_block_size(
topk_ids,
lora_info.seg_indptr,
lora_info.req_to_lora,
int(lora_info.num_experts),
int(block_size_m),
int(max_loras),
int(max_num_tokens_padded),
int(max_num_m_blocks),
lora_info.adapter_enabled,
device,
)
invoke_fused_moe_kernel(
hidden_states,
w13,
b13,
intermediate_cache1,
a13_scale,
w13_scale,
w13_zp,
topk_weights,
topk_ids,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
apply_router_weight_on_input,
topk_ids.shape[1],
running_state["config"],
compute_type=compute_type,
use_fp8_w8a8=use_fp8_w8a8,
use_int8_w8a8=use_int8_w8a8,
use_int8_w8a16=use_int8_w8a16,
use_int4_w4a16=use_int4_w4a16,
per_channel_quant=per_channel_quant,
block_shape=block_shape,
)
# ==============================
# Perform LoRA alignment for both gate up and gate down operations
# Define shrink_config for LoRA alignment
# TODO: Add autotuning for block sizes across different GPU architectures and problem sizes
shrink_config = {"BLOCK_SIZE_M": 64}
# Prepare inputs for the kernel
block_size_m = shrink_config["BLOCK_SIZE_M"]
max_loras = len(lora_info.lora_ranks)
# Calculate max_num_tokens_padded
max_num_tokens_padded = topk_ids.numel() + lora_info.num_experts * (
block_size_m - 1
)
max_num_tokens_padded = (
(max_num_tokens_padded + block_size_m - 1) // block_size_m
) * block_size_m
max_num_m_blocks = (max_num_tokens_padded + block_size_m - 1) // block_size_m
device = topk_ids.device
lora_ids = torch.arange(max_loras, dtype=torch.int32, device=device)
else:
if cg is not None:
sorted_token_ids_lora = cg["sorted_token_ids_lora"][
: max_loras * max_num_tokens_padded
@@ -313,328 +287,215 @@ class TritonRunnerCoreWithLoRA(TritonRunnerCore):
num_tokens_post_padded_lora,
lora_info.adapter_enabled,
lora_ids,
None, # expert_map
cumsum_buffer=cg["cumsum_buffer"] if cg is not None else None,
token_mask=(
cg["token_mask"][: max_loras * topk_ids.shape[0]]
if cg is not None
else None
),
cumsum_buffer=cg.get("cumsum_buffer") if cg is not None else None,
token_mask=cg.get("token_mask") if cg is not None else None,
)
# Reshape the sorted tensors for fused_moe_lora (expects 2D: max_loras x max_num_tokens_padded)
sorted_token_ids_reshaped = sorted_token_ids_lora.view(max_loras, -1)
expert_ids_reshaped = expert_ids_lora.view(max_loras, -1)
return (
sorted_token_ids_lora.view(max_loras, -1),
expert_ids_lora.view(max_loras, -1),
num_tokens_post_padded_lora,
lora_ids,
)
# ============================================================
# Stage 1.5: Add LoRA gate_up delta BEFORE activation
# ============================================================
self._add_lora_gate_up_delta(
hidden_states=hidden_states,
intermediate_cache=intermediate_cache1,
topk_weights=topk_weights,
lora_info=lora_info,
sorted_token_ids_reshaped=sorted_token_ids_reshaped,
expert_ids_reshaped=expert_ids_reshaped,
num_tokens_post_padded_lora=num_tokens_post_padded_lora,
lora_ids=lora_ids,
def _add_lora_gate_up_delta(
hidden_states: torch.Tensor,
intermediate_cache: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
lora_info: LoRAInfo,
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,
) -> None:
"""Add LoRA gate_up delta to intermediate_cache in-place."""
from sglang.srt.lora.triton_ops import fused_moe_lora
if get_is_capture_mode():
# During CUDA graph capture, always enter the LoRA path so that
# the LoRA kernels are recorded in the graph. adapter_enabled is
# all-zeros during capture, so the Triton kernel early-exits per
# program (zero overhead). During replay the tensor is updated
# in-place with the real adapter mask before graph.replay().
has_active_lora = True
else:
num_loras = len(lora_info.lora_ranks)
has_active_lora = (
(
lora_info.adapter_enabled[:num_loras]
* (lora_info.lora_ranks > 0).to(lora_info.adapter_enabled.dtype)
)
.any()
.item()
)
if not has_active_lora or lora_info is None or lora_info.max_lora_rank == 0:
return
# ============================================================
# Stage 2: Activation (SiLU or GELU)
# ============================================================
if cg is not None:
intermediate_cache2 = cg["intermediate_cache2"][
: M * topk_ids.shape[1], : N // 2
]
else:
intermediate_cache2 = torch.empty(
(M * topk_ids.shape[1], N // 2),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
if activation == "silu":
if gemm1_alpha is not None:
assert gemm1_limit is not None
intermediate_cache2 = _swiglu_gpt_oss_sigmoid_alpha(
intermediate_cache1.view(-1, N), gemm1_alpha, gemm1_limit
)
elif gemm1_limit is not None:
intermediate_cache2 = _swiglu_silu_clamp_mul(
intermediate_cache1.view(-1, N), gemm1_limit
)
elif _is_cuda or _is_hip or _is_xpu:
silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
else:
vllm_ops.silu_and_mul(
intermediate_cache2, intermediate_cache1.view(-1, N)
)
elif activation == "gelu":
assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu"
assert gemm1_limit is None, "gemm1_limit is not supported for gelu"
if _is_cuda or _is_hip:
gelu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
else:
vllm_ops.gelu_and_mul(
intermediate_cache2, intermediate_cache1.view(-1, N)
)
else:
raise ValueError(f"Unsupported activation: {activation=}")
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
# ============================================================
# Stage 3: Down projection (base)
# ============================================================
if cg is not None:
intermediate_cache3 = cg["intermediate_cache3"][
:M, : topk_ids.shape[1], : w2.shape[1]
]
else:
intermediate_cache3 = torch.empty(
(M, topk_ids.shape[1], w2.shape[1]),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
if lora_info.experts_shared_outer_loras:
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:, :]]
if no_combine:
assert not inplace
out_hidden_states = torch.empty(
(M, topk_ids.shape[1], w2.shape[1]),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
elif inplace:
out_hidden_states = hidden_states
elif cg is not None:
out_hidden_states = cg["out_hidden_states"][:M, : hidden_states.shape[1]]
else:
out_hidden_states = torch.empty_like(hidden_states)
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,
)
invoke_fused_moe_kernel(
intermediate_cache2,
w2,
b2,
intermediate_cache3,
a2_scale,
w2_scale,
w2_zp,
def _add_lora_down_delta(
intermediate_input: torch.Tensor,
intermediate_cache: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
lora_info: LoRAInfo,
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,
) -> None:
"""Add LoRA down delta to intermediate_cache in-place."""
from sglang.srt.lora.triton_ops import fused_moe_lora
if lora_info.max_lora_rank == 0:
return
M, top_k, hidden_dim = intermediate_cache.shape
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:
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
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(
hidden_states: torch.Tensor,
lora_info: LoRAInfo,
topk_ids: torch.Tensor,
) -> 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.
"""
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)
def after_gate_up(
hidden_states: torch.Tensor,
intermediate_cache1: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
) -> None:
_add_lora_gate_up_delta(
hidden_states,
intermediate_cache1,
topk_weights,
topk_ids,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
not apply_router_weight_on_input,
1,
running_state["config"],
compute_type=compute_type,
use_fp8_w8a8=use_fp8_w8a8,
use_int8_w8a8=use_int8_w8a8,
use_int8_w8a16=use_int8_w8a16,
use_int4_w4a16=use_int4_w4a16,
per_channel_quant=per_channel_quant,
block_shape=block_shape,
lora_info,
sorted_token_ids_reshaped,
expert_ids_reshaped,
num_tokens_post_padded_lora,
lora_ids,
)
# ============================================================
# Stage 3.5: Add LoRA down delta BEFORE final reduction
# ============================================================
self._add_lora_down_delta(
intermediate_input=intermediate_cache2,
intermediate_cache=intermediate_cache3,
topk_weights=topk_weights,
lora_info=lora_info,
sorted_token_ids_reshaped=sorted_token_ids_reshaped,
expert_ids_reshaped=expert_ids_reshaped,
num_tokens_post_padded_lora=num_tokens_post_padded_lora,
lora_ids=lora_ids,
)
# ============================================================
# Stage 4: Final reduction (sum across top_k)
# ============================================================
if routed_scaling_factor is None:
routed_scaling_factor = 1.0
if no_combine:
pass
elif _is_cuda:
if topk_ids.shape[1] == 1 and routed_scaling_factor == 1.0:
out_hidden_states[:] = intermediate_cache3.squeeze(1)
elif topk_ids.shape[1] == 2 and routed_scaling_factor == 1.0:
torch.add(
intermediate_cache3[:, 0],
intermediate_cache3[:, 1],
out=out_hidden_states,
).squeeze(dim=1)
else:
if M <= 32:
moe_sum_reduce_torch_compile(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
routed_scaling_factor,
)
else:
moe_sum_reduce_triton(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
routed_scaling_factor,
)
elif _is_hip:
from vllm import _custom_ops as vllm_ops
vllm_ops.moe_sum(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
)
else:
from vllm import _custom_ops as vllm_ops
vllm_ops.moe_sum(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
)
return TritonRunnerOutput(
hidden_states=out_hidden_states,
)
def _add_lora_gate_up_delta(
self,
hidden_states: torch.Tensor, # [M, hidden_dim]
intermediate_cache: torch.Tensor, # [M, top_k, gate_up_dim]
topk_weights: torch.Tensor, # [M, top_k]
lora_info: LoRAInfo,
sorted_token_ids_reshaped: torch.Tensor,
expert_ids_reshaped: torch.Tensor,
num_tokens_post_padded_lora: torch.Tensor,
lora_ids: torch.Tensor,
def after_down(
intermediate_input: torch.Tensor,
intermediate_cache3: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
) -> None:
"""
Add LoRA gate_up delta to intermediate_cache in-place.
For each (token, expert) pair, computes:
delta = scaling * B @ (A @ hidden_states[token])
and adds it to intermediate_cache[token, k] where k is the top_k index.
"""
from sglang.srt.lora.triton_ops import fused_moe_lora
M, top_k, gate_up_dim = intermediate_cache.shape
# Skip LoRA computation if no LoRA adapters have non-zero rank
if lora_info.max_lora_rank == 0:
return
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
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:, :],
]
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,
# TODO: Replace hardcoded block sizes with autotuned configs
shrink_block_size_m=64,
shrink_block_size_n=64,
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=64,
expand_group_size_m=8,
expand_num_warps=4,
expand_num_stages=2,
expand_split_k=1,
fully_sharded=lora_info.fully_sharded,
_add_lora_down_delta(
intermediate_input,
intermediate_cache3,
topk_weights,
topk_ids,
lora_info,
sorted_token_ids_reshaped,
expert_ids_reshaped,
num_tokens_post_padded_lora,
lora_ids,
)
def _add_lora_down_delta(
self,
intermediate_input: torch.Tensor, # [M * top_k, intermediate_dim]
intermediate_cache: torch.Tensor, # [M, top_k, hidden_dim]
topk_weights: torch.Tensor, # [M, top_k]
lora_info: LoRAInfo,
sorted_token_ids_reshaped: torch.Tensor,
expert_ids_reshaped: torch.Tensor,
num_tokens_post_padded_lora: torch.Tensor,
lora_ids: torch.Tensor,
) -> None:
"""
Add LoRA down delta to intermediate_cache in-place.
For each (token, expert) pair, computes:
delta = scaling * B @ (A @ intermediate_input[dispatched_idx])
and adds it to intermediate_cache[token, k].
"""
from sglang.srt.lora.triton_ops import fused_moe_lora
M, top_k, hidden_dim = intermediate_cache.shape
# Skip LoRA computation if no LoRA adapters have non-zero rank
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 = [down_lora_b]
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
fused_moe_lora(
output=intermediate_cache,
qcurr_hidden_states=intermediate_input,
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=lora_info.max_lora_rank,
top_k_num=top_k,
lora_ids=lora_ids,
adapter_enabled=lora_info.adapter_enabled,
# TODO: Replace hardcoded block sizes with autotuned configs
shrink_block_size_m=64,
shrink_block_size_n=64,
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=64,
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,
)
return LoRAHooks(after_gate_up=after_gate_up, after_down=after_down)