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)
test/registered/lora/test_lora_moe_runner.py
+635
View File
@@ -0,0 +1,635 @@
# Copyright 2023-2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import random
from unittest.mock import patch
import pytest
import torch
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
from sglang.srt.layers.moe.moe_runner.triton import (
TritonMoeQuantInfo,
)
from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
from sglang.srt.layers.moe.topk import StandardTopKOutput
from sglang.srt.layers.moe.utils import MoeRunnerBackend
from sglang.srt.lora.lora_moe_runners import LoRAInfo
from sglang.srt.utils import set_random_seed
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=80, suite="stage-b-test-1-gpu-large")
def generate_request_data(
num_tokens: int, num_sequences: int, max_loras: int, device="cuda"
):
"""
Generates segment-based request data instead of token-based data.
"""
assert num_sequences > 0 and max_loras > 0
assert num_tokens >= num_sequences, "num_tokens must be >= num_sequences"
# 1. Generate random segment lengths
remaining = num_tokens
seg_lens = []
for _ in range(num_sequences - 1):
# Ensure at least 1 token per sequence
max_len = remaining - (num_sequences - len(seg_lens)) + 1
length = random.randint(1, min(max_len, num_tokens // num_sequences * 2))
seg_lens.append(length)
remaining -= length
seg_lens.append(remaining) # Last segment gets the rest
# 2. Build seg_indptr [0, len1, len1+len2, ...]
seg_indptr = torch.cumsum(
torch.tensor([0] + seg_lens, dtype=torch.int32, device=device),
dim=0,
dtype=torch.int32,
)
# 3. Assign one LoRA ID per Request
req_to_lora = torch.randint(
0, max_loras, (num_sequences,), dtype=torch.int32, device=device
)
# 4. Create dense mapping for the Naive verification function
# (Expand req_to_lora based on seg_lens)
token_lora_mapping = torch.repeat_interleave(
req_to_lora, torch.tensor(seg_lens, device=device)
)
return seg_indptr, req_to_lora, token_lora_mapping
def assign_experts_to_tokens(
num_tokens: int, num_experts: int, top_k_num: int, dtype=torch.float32
):
assert top_k_num <= num_experts, "top_k_num must be <= num_experts"
expert_indices = torch.empty((num_tokens, top_k_num), dtype=torch.int32)
for i in range(num_tokens):
selected = torch.randperm(num_experts)[:top_k_num]
expert_indices[i] = selected
expert_weights = torch.rand((num_tokens, top_k_num), dtype=dtype)
expert_weights = expert_weights / expert_weights.sum(dim=1, keepdim=True)
return expert_indices, expert_weights
def sample_data(
num_tokens: int,
num_sequences: int,
max_loras: int,
num_experts: int,
top_k_num: int,
dtype=torch.float32,
device="cuda",
):
topk_ids, topk_weights = assign_experts_to_tokens(
num_tokens, num_experts, top_k_num, dtype
)
seg_indptr, req_to_lora, token_lora_mapping = generate_request_data(
num_tokens, num_sequences, max_loras, device
)
return topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping
def create_lora_info(
seg_indptr,
weight_indices,
topk_ids,
max_loras,
num_experts,
max_lora_rank,
hidden_dim,
intermediate_dim,
gate_up_dim,
dtype,
device,
):
# -------------------------------------------------------------------------
# 1. Deterministic LoRA A Initialization
# -------------------------------------------------------------------------
val_gate_up_a = 0.1
gate_up_lora_a_weights = torch.full(
(max_loras, num_experts, max_lora_rank * 2, hidden_dim),
val_gate_up_a,
dtype=dtype,
device=device,
)
val_down_a = 1.0 / intermediate_dim
down_lora_a_weights = torch.full(
(max_loras, num_experts, max_lora_rank, intermediate_dim),
val_down_a,
dtype=dtype,
device=device,
)
# -------------------------------------------------------------------------
# 2. Deterministic LoRA B Initialization
# -------------------------------------------------------------------------
base_target = 0.05
gate_up_lora_b_weights = torch.zeros(
(max_loras, num_experts, gate_up_dim, max_lora_rank),
dtype=dtype,
device=device,
)
down_lora_b_weights = torch.zeros(
(max_loras, num_experts, hidden_dim, max_lora_rank), dtype=dtype, device=device
)
for i in range(num_experts):
expert_multiplier = i + 1
divisor = max(1, max_lora_rank)
fill_val = (base_target * expert_multiplier) / divisor
gate_up_lora_b_weights[:, i, :, :] = fill_val
down_lora_b_weights[:, i, :, :] = fill_val
# -------------------------------------------------------------------------
# 3. Setup Metadata
# -------------------------------------------------------------------------
lora_ranks = torch.full(
(max_loras,), max_lora_rank, dtype=torch.int32, device=device
)
# Enable all adapters referenced in weight_indices
adapter_enabled = torch.zeros(max_loras + 1, dtype=torch.int32, device=device)
adapter_enabled.index_fill_(0, weight_indices.long(), 1)
return LoRAInfo(
gate_up_lora_a_weights=gate_up_lora_a_weights,
gate_up_lora_b_weights=gate_up_lora_b_weights,
down_lora_a_weights=down_lora_a_weights,
down_lora_b_weights=down_lora_b_weights,
# UPDATED FIELDS
seg_indptr=seg_indptr,
req_to_lora=weight_indices,
lora_ranks=lora_ranks,
adapter_enabled=adapter_enabled,
max_lora_rank=max_lora_rank,
num_experts=num_experts,
)
def torch_naive_moe_with_lora(
hidden_states,
w13,
w2,
b13,
b2,
topk_weights,
topk_ids,
lora_info,
token_lora_mapping,
):
"""
Naive implementation. Note: We pass 'token_lora_mapping' explicitly because
lora_info no longer contains it, but the naive token-loop logic needs it.
"""
num_tokens, hidden_dim = hidden_states.shape
top_k = topk_ids.shape[1]
num_experts = w13.shape[0]
# Expand hidden states for top-k routing
hidden_expanded = (
hidden_states.unsqueeze(1).expand(-1, top_k, -1).reshape(-1, hidden_dim)
)
# 1. Gate/Up Projection (Base)
gate_up_out = torch.zeros(
num_tokens * top_k,
w13.shape[1],
dtype=hidden_states.dtype,
device=hidden_states.device,
)
for expert_id in range(num_experts):
mask = (topk_ids == expert_id).flatten()
if mask.any():
expert_result = hidden_expanded[mask] @ w13[expert_id].T
gate_up_out[mask] = expert_result
if b13 is not None:
gate_up_out[mask] += b13[expert_id]
gate_up_out = gate_up_out.view(num_tokens, top_k, -1)
# 1.5. LoRA Gate/Up Delta
# gate_up_lora_a is packed as [gate_a; up_a] along rank dim → [2*r, hidden_dim]
# gate_up_lora_b is packed as [gate_b; up_b] along output dim → [2*inter, r]
# Correct computation splits them: gate uses first r rows of A with first half of B,
# up uses last r rows of A with second half of B.
if lora_info.max_lora_rank > 0:
r = lora_info.max_lora_rank
for i in range(num_tokens):
for k in range(top_k):
expert_id = topk_ids[i, k]
lora_id = token_lora_mapping[i]
if lora_id < len(lora_info.lora_ranks):
lora_a = lora_info.gate_up_lora_a_weights[lora_id, expert_id]
lora_b = lora_info.gate_up_lora_b_weights[lora_id, expert_id]
half = lora_b.shape[0] // 2
lora_a_result = lora_a @ hidden_states[i]
gate_delta = lora_b[:half, :] @ lora_a_result[:r]
up_delta = lora_b[half:, :] @ lora_a_result[r:]
gate_up_out[i, k] += torch.cat([gate_delta, up_delta])
# 2. Activation
gate_up_dim = gate_up_out.shape[-1]
gate_dim = gate_up_dim // 2
gate = gate_up_out[..., :gate_dim]
up = gate_up_out[..., gate_dim:]
silu_gate = torch.nn.functional.silu(gate)
intermediate_out = silu_gate * up
# 3. Down Projection (Base)
down_out = torch.zeros(
num_tokens,
top_k,
hidden_dim,
dtype=hidden_states.dtype,
device=hidden_states.device,
)
for expert_id in range(num_experts):
mask = topk_ids == expert_id
if mask.any():
masked_intermediate = intermediate_out[mask]
expert_down_result = masked_intermediate @ w2[expert_id].T
down_out[mask] = expert_down_result
if b2 is not None:
down_out[mask] += b2[expert_id]
# 3.5. LoRA Down Delta
if lora_info.max_lora_rank > 0:
for i in range(num_tokens):
for k in range(top_k):
expert_id = topk_ids[i, k]
lora_id = token_lora_mapping[i] # Use explicit mapping
if lora_id < len(lora_info.lora_ranks):
lora_a = lora_info.down_lora_a_weights[lora_id, expert_id]
lora_b = lora_info.down_lora_b_weights[lora_id, expert_id]
lora_a_result = lora_a @ intermediate_out[i, k]
lora_b_result = lora_b @ lora_a_result
down_out[i, k] += lora_b_result
# 4. Final Reduction
weighted_out = down_out * topk_weights.unsqueeze(-1)
final_out = weighted_out.sum(dim=1)
return final_out
@pytest.mark.parametrize("num_tokens", [32, 64])
@pytest.mark.parametrize("top_k_num", [1, 2])
@pytest.mark.parametrize("num_experts", [8, 20])
@pytest.mark.parametrize("max_lora_rank", [8, 16])
def test_lora_moe_runner_multi_expert(
num_tokens, top_k_num, num_experts, max_lora_rank
):
# Fixed parameters
max_loras = 2
hidden_dim = 512
intermediate_dim = 1024
dtype = torch.float32
device = "cuda:0"
seed = 42
torch.set_default_device(device)
set_random_seed(seed)
num_sequences = 4
# Generate Data using the new Request-Based generator
topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping = sample_data(
num_tokens, num_sequences, max_loras, num_experts, top_k_num, dtype, device
)
gate_up_dim = intermediate_dim * 2
# Initialize experts
w13 = torch.randn(num_experts, gate_up_dim, hidden_dim, dtype=dtype) * 0.1
w2 = torch.randn(num_experts, hidden_dim, intermediate_dim, dtype=dtype) * 0.1
b13 = torch.randn(num_experts, gate_up_dim, dtype=dtype) * 0.1
b2 = torch.randn(num_experts, hidden_dim, dtype=dtype) * 0.1
hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype)
# Create LoRA Info using the new fields
lora_info_delta = create_lora_info(
seg_indptr=seg_indptr,
weight_indices=req_to_lora,
topk_ids=topk_ids,
max_loras=max_loras,
num_experts=num_experts,
max_lora_rank=max_lora_rank,
hidden_dim=hidden_dim,
intermediate_dim=intermediate_dim,
gate_up_dim=gate_up_dim,
dtype=dtype,
device=device,
)
lora_info_baseline = create_lora_info(
seg_indptr=seg_indptr,
weight_indices=req_to_lora,
topk_ids=topk_ids,
max_loras=max_loras,
num_experts=num_experts,
max_lora_rank=0, # Set rank to 0 for baseline
hidden_dim=hidden_dim,
intermediate_dim=intermediate_dim,
gate_up_dim=gate_up_dim,
dtype=dtype,
device=device,
)
# Sort tokens for the runner
topk_ids_flat = topk_ids.flatten()
sorted_indices = torch.argsort(topk_ids_flat)
sorted_token_ids = sorted_indices // top_k_num
expert_ids = topk_ids_flat[sorted_indices]
num_dispatched = num_tokens * top_k_num
num_tokens_post_padded = torch.tensor(
[num_dispatched], dtype=torch.int32, device=device
)
quant_info = TritonMoeQuantInfo(
w13_weight=w13,
w2_weight=w2,
b13=b13,
b2=b2,
)
config = MoeRunnerConfig(
activation="silu",
is_gated=True,
inplace=False,
no_combine=False,
gemm1_alpha=None,
gemm1_clamp_limit=None,
routed_scaling_factor=1.0,
apply_router_weight_on_input=False,
num_local_experts=num_experts,
)
# Create StandardTopKOutput
router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device)
topk_output = StandardTopKOutput(
topk_weights=topk_weights,
topk_ids=topk_ids,
router_logits=router_logits,
)
# Create StandardDispatchOutput
dispatch_output = StandardDispatchOutput(
hidden_states=hidden_states,
hidden_states_scale=None,
topk_output=topk_output,
)
class MockServerArgs:
enable_deterministic_inference = False
with patch(
"sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config.get_global_server_args",
return_value=MockServerArgs(),
):
runner = MoeRunner(MoeRunnerBackend.TRITON, config, lora_enabled=True)
# 3. Get outputs for both scenarios
output_with_lora = runner.run(
dispatch_output, quant_info, lora_info_delta
).hidden_states
output_baseline = runner.run(
dispatch_output, quant_info, lora_info_baseline
).hidden_states
# Run Naive Torch Implementation (Uses dense mapping for verification)
torch_output_lora = torch_naive_moe_with_lora(
hidden_states,
w13,
w2,
b13,
b2,
topk_weights,
topk_ids,
lora_info_delta,
token_lora_mapping,
)
torch_output_base = torch_naive_moe_with_lora(
hidden_states,
w13,
w2,
b13,
b2,
topk_weights,
topk_ids,
lora_info_baseline,
token_lora_mapping,
)
# The actual "Delta" (LoRA effect) for both
sglang_delta = output_with_lora - output_baseline
torch_delta = torch_output_lora - torch_output_base
# Larger expert counts accumulate more numerical drift in Triton kernels on GB300
tol = 0.15 if num_experts >= 20 else 5e-2
torch.testing.assert_close(sglang_delta, torch_delta, atol=tol, rtol=tol)
def _setup_marlin_moe_weights(num_experts, n, k, dtype):
"""Quantize float weights into AWQ Marlin format for testing."""
from sgl_kernel.scalar_type import scalar_types
from sglang.test.test_marlin_utils import awq_marlin_quantize
group_size = 128
quant_type = scalar_types.uint4
w = torch.randn((num_experts, n, k), device="cuda", dtype=dtype) / 20
w_ref_l, qweight_l, scales_l, zeros_l = [], [], [], []
for i in range(num_experts):
w_ref, qweight, scales, zeros = awq_marlin_quantize(
w[i].transpose(1, 0), quant_type, group_size
)
w_ref_l.append(w_ref.T)
qweight_l.append(qweight)
scales_l.append(scales)
zeros_l.append(zeros)
def _stack(tensors):
dev = tensors[0].device
return torch.stack(tensors, dim=0).to(dev)
return (
_stack(w_ref_l),
_stack(qweight_l).contiguous(),
_stack(scales_l),
_stack(zeros_l),
)
@pytest.mark.parametrize("num_tokens", [32, 64])
@pytest.mark.parametrize("top_k_num", [1, 2])
@pytest.mark.parametrize("num_experts", [8])
@pytest.mark.parametrize("max_lora_rank", [8, 16])
def test_lora_moe_runner_marlin(num_tokens, top_k_num, num_experts, max_lora_rank):
from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
max_loras = 2
hidden_dim = 512
intermediate_dim = 1024
gate_up_dim = intermediate_dim * 2
dtype = torch.float16
device = "cuda:0"
seed = 42
torch.set_default_device(device)
set_random_seed(seed)
num_sequences = 4
topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping = sample_data(
num_tokens,
num_sequences,
max_loras,
num_experts,
top_k_num,
dtype,
device,
)
# Quantize base weights to Marlin format
_, w13_qweight, w13_scales, w13_qzeros = _setup_marlin_moe_weights(
num_experts, gate_up_dim, hidden_dim, dtype
)
_, w2_qweight, w2_scales, w2_qzeros = _setup_marlin_moe_weights(
num_experts, hidden_dim, intermediate_dim, dtype
)
hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=device)
lora_info_delta = create_lora_info(
seg_indptr=seg_indptr,
weight_indices=req_to_lora,
topk_ids=topk_ids,
max_loras=max_loras,
num_experts=num_experts,
max_lora_rank=max_lora_rank,
hidden_dim=hidden_dim,
intermediate_dim=intermediate_dim,
gate_up_dim=gate_up_dim,
dtype=dtype,
device=device,
)
lora_info_baseline = create_lora_info(
seg_indptr=seg_indptr,
weight_indices=req_to_lora,
topk_ids=topk_ids,
max_loras=max_loras,
num_experts=num_experts,
max_lora_rank=0,
hidden_dim=hidden_dim,
intermediate_dim=intermediate_dim,
gate_up_dim=gate_up_dim,
dtype=dtype,
device=device,
)
quant_info = MarlinMoeQuantInfo(
w13_qweight=w13_qweight,
w2_qweight=w2_qweight,
w13_scales=w13_scales,
w2_scales=w2_scales,
w13_qzeros=w13_qzeros,
w2_qzeros=w2_qzeros,
w13_g_idx=None,
w2_g_idx=None,
w13_g_idx_sort_indices=None,
w2_g_idx_sort_indices=None,
weight_bits=4,
)
config = MoeRunnerConfig(
activation="silu",
is_gated=True,
inplace=False,
no_combine=False,
gemm1_alpha=None,
gemm1_clamp_limit=None,
routed_scaling_factor=1.0,
apply_router_weight_on_input=False,
num_local_experts=num_experts,
)
router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device)
topk_output = StandardTopKOutput(
topk_weights=topk_weights,
topk_ids=topk_ids,
router_logits=router_logits,
)
dispatch_output = StandardDispatchOutput(
hidden_states=hidden_states,
hidden_states_scale=None,
topk_output=topk_output,
)
class MockServerArgs:
enable_deterministic_inference = False
with patch(
"sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config.get_global_server_args",
return_value=MockServerArgs(),
):
runner = MoeRunner(MoeRunnerBackend.MARLIN, config, lora_enabled=True)
output_with_lora = runner.run(
dispatch_output, quant_info, lora_info_delta
).hidden_states
output_baseline = runner.run(
dispatch_output, quant_info, lora_info_baseline
).hidden_states
marlin_delta = output_with_lora - output_baseline
# Verify the LoRA hooks fired and produced a non-trivial delta
assert marlin_delta.abs().max().item() > 1e-4, (
f"LoRA delta is too small ({marlin_delta.abs().max().item():.6f}), "
"hooks may not be firing"
)
assert torch.isfinite(
output_with_lora
).all(), "Marlin+LoRA output contains non-finite values"
assert torch.isfinite(
output_baseline
).all(), "Marlin baseline output contains non-finite values"
if __name__ == "__main__":
pytest.main([__file__])
test/registered/lora/test_marlin_lora_correctness.py
+288
View File
@@ -0,0 +1,288 @@
# Copyright 2023-2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Correctness test: Marlin (int4 base + LoRA) vs Triton (dequantized base + LoRA).
Fake-quantizes random weights to int4/Marlin format and dequantizes them with the
same path, then runs both backends through MoeRunner and compares LoRA deltas.
"""
from unittest.mock import patch
import pytest
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.layers.moe.moe_runner.runner import MoeRunner
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
from sglang.srt.layers.moe.topk import StandardTopKOutput
from sglang.srt.layers.moe.utils import MoeRunnerBackend
from sglang.srt.lora.lora_moe_runners import LoRAInfo
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=600, suite="stage-b-test-1-gpu-large")
# ---------------------------------------------------------------------------
# Fake quantization helpers (symmetric int4, matching Marlin's dequant path)
# ---------------------------------------------------------------------------
def _quantize_per_expert(w_float: torch.Tensor, K: int, group_size: int):
"""Quantize [N, K] float weight to int4. Returns (q_int [N,K], scales_bf16 [N,groups])."""
N = w_float.shape[0]
num_groups = K // group_size
w_grouped = w_float.reshape(N, num_groups, group_size)
scales_fp32 = w_grouped.abs().amax(dim=-1) / 7.0
scales_fp32 = scales_fp32.clamp(min=1e-6)
scales_bf16 = scales_fp32.to(torch.bfloat16)
scales_for_quant = scales_bf16.float()
q_int = torch.zeros(N, K, dtype=torch.int32, device=w_float.device)
for g in range(num_groups):
s = scales_for_quant[:, g : g + 1]
sl = slice(g * group_size, (g + 1) * group_size)
q_int[:, sl] = torch.round(w_float[:, sl] / s).clamp(-8, 7).to(torch.int32) + 8
return q_int, scales_bf16
def _fake_quantize_to_marlin_int4(weight_bf16: torch.Tensor):
"""Fake-quantize [E, N, K] bf16 weight to Marlin int4 format.
Returns: (qweight, scales, g_idx, g_idx_sort_indices)
"""
from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
from sglang.srt.layers.quantization.marlin_utils import marlin_permute_scales
from sglang.srt.layers.quantization.utils import pack_rows
E, N, K = weight_bf16.shape
num_bits = 4
group_size = 128
device = weight_bf16.device
all_qweight, all_scales = [], []
for e in range(E):
q_int, scales_bf16 = _quantize_per_expert(weight_bf16[e].float(), K, group_size)
w_quant_t = q_int.t().contiguous()
packed = pack_rows(w_quant_t, num_bits, K, N)
perm = torch.arange(K, device=device, dtype=torch.int32)
all_qweight.append(gptq_marlin_repack(packed.to(device), perm, K, N, num_bits))
all_scales.append(
marlin_permute_scales(
scales_bf16.t().contiguous().to(device), K, N, group_size
)
)
g_idx = (
(torch.arange(K, device=device, dtype=torch.int32) // group_size)
.unsqueeze(0)
.expand(E, -1)
.contiguous()
)
sort_indices = (
torch.arange(K, device=device, dtype=torch.int32)
.unsqueeze(0)
.expand(E, -1)
.contiguous()
)
return torch.stack(all_qweight), torch.stack(all_scales), g_idx, sort_indices
def _dequantize_from_marlin_int4(weight_bf16_orig: torch.Tensor, group_size: int = 128):
"""Dequantize using the same path as _fake_quantize, so Triton reference matches Marlin."""
E, N, K = weight_bf16_orig.shape
result = torch.zeros_like(weight_bf16_orig)
for e in range(E):
q_int, scales_bf16 = _quantize_per_expert(
weight_bf16_orig[e].float(), K, group_size
)
num_groups = K // group_size
for g in range(num_groups):
sl = slice(g * group_size, (g + 1) * group_size)
s = scales_bf16[:, g : g + 1]
result[e, :, sl] = (q_int[:, sl] - 8).to(torch.bfloat16) * s
return result
# ---------------------------------------------------------------------------
# Test
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("num_tokens", [1, 8, 32])
@pytest.mark.parametrize("top_k", [2, 8])
def test_marlin_vs_triton_lora_correctness(num_tokens, top_k):
torch.manual_seed(42)
device = "cuda"
dtype = torch.bfloat16
hidden_dim = 7168
intermediate_dim = 2048
gate_up_dim = 2 * intermediate_dim
num_experts = 64
lora_rank = 32
num_loras = 1
hidden = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=device)
topk_weights = torch.randn(
num_tokens, top_k, dtype=torch.float32, device=device
).softmax(dim=-1)
topk_ids = torch.randint(
0, num_experts, (num_tokens, top_k), dtype=torch.int32, device=device
)
# Base weights (random bf16)
w13_bf16 = (
torch.randn(num_experts, gate_up_dim, hidden_dim, dtype=dtype, device=device)
* 0.01
)
w2_bf16 = (
torch.randn(
num_experts, hidden_dim, intermediate_dim, dtype=dtype, device=device
)
* 0.01
)
# LoRA weights (shared across both paths)
gu_lora_a = (
torch.randn(num_loras, 1, lora_rank * 2, hidden_dim, dtype=dtype, device=device)
* 0.01
)
gu_lora_b = (
torch.randn(
num_loras, num_experts, gate_up_dim, lora_rank, dtype=dtype, device=device
)
* 0.01
)
dn_lora_a = (
torch.randn(
num_loras,
num_experts,
lora_rank,
intermediate_dim,
dtype=dtype,
device=device,
)
* 0.01
)
dn_lora_b = (
torch.randn(num_loras, 1, hidden_dim, lora_rank, dtype=dtype, device=device)
* 0.01
)
# Token-to-LoRA mapping: all tokens use adapter 0
seg_indptr = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
req_to_lora = torch.tensor([0], dtype=torch.int32, device=device)
def _make_lora_info(rank):
return LoRAInfo(
gate_up_lora_a_weights=gu_lora_a if rank > 0 else gu_lora_a[:, :, :0, :],
gate_up_lora_b_weights=gu_lora_b if rank > 0 else gu_lora_b[:, :, :, :0],
down_lora_a_weights=dn_lora_a if rank > 0 else dn_lora_a[:, :, :0, :],
down_lora_b_weights=dn_lora_b if rank > 0 else dn_lora_b[:, :, :, :0],
seg_indptr=seg_indptr,
req_to_lora=req_to_lora,
lora_ranks=torch.full((num_loras,), rank, dtype=torch.int32, device=device),
adapter_enabled=torch.ones(num_loras + 1, dtype=torch.int32, device=device),
max_lora_rank=rank,
num_experts=num_experts,
experts_shared_outer_loras=True,
)
lora_info = _make_lora_info(lora_rank)
lora_baseline = _make_lora_info(0)
# Quantize for Marlin, dequantize for Triton reference
w13_qw, w13_sc, w13_gidx, w13_si = _fake_quantize_to_marlin_int4(w13_bf16)
w2_qw, w2_sc, w2_gidx, w2_si = _fake_quantize_to_marlin_int4(w2_bf16)
w13_deq = _dequantize_from_marlin_int4(w13_bf16)
w2_deq = _dequantize_from_marlin_int4(w2_bf16)
marlin_qi = MarlinMoeQuantInfo(
w13_qweight=w13_qw,
w2_qweight=w2_qw,
w13_scales=w13_sc,
w2_scales=w2_sc,
w13_g_idx=w13_gidx,
w2_g_idx=w2_gidx,
w13_g_idx_sort_indices=w13_si,
w2_g_idx_sort_indices=w2_si,
weight_bits=4,
)
triton_qi = TritonMoeQuantInfo(
w13_weight=w13_deq, w2_weight=w2_deq, b13=None, b2=None
)
config = MoeRunnerConfig(
activation="silu",
is_gated=True,
inplace=False,
no_combine=False,
gemm1_alpha=None,
gemm1_clamp_limit=None,
routed_scaling_factor=1.0,
apply_router_weight_on_input=False,
num_local_experts=num_experts,
)
router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device)
topk_output = StandardTopKOutput(
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits
)
dispatch_output = StandardDispatchOutput(
hidden_states=hidden, hidden_states_scale=None, topk_output=topk_output
)
class MockServerArgs:
enable_deterministic_inference = False
with patch(
"sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config.get_global_server_args",
return_value=MockServerArgs(),
):
marlin_runner = MoeRunner(MoeRunnerBackend.MARLIN, config, lora_enabled=True)
triton_runner = MoeRunner(MoeRunnerBackend.TRITON, config, lora_enabled=True)
marlin_out = marlin_runner.run(
dispatch_output, marlin_qi, lora_info
).hidden_states
marlin_base = marlin_runner.run(
dispatch_output, marlin_qi, lora_baseline
).hidden_states
triton_out = triton_runner.run(
dispatch_output, triton_qi, lora_info
).hidden_states
triton_base = triton_runner.run(
dispatch_output, triton_qi, lora_baseline
).hidden_states
marlin_delta = marlin_out - marlin_base
triton_delta = triton_out - triton_base
# Remaining error is from kernel-level accumulation differences
# (Marlin fp32 reduce vs Triton bf16 dot), not from quantization mismatch.
torch.testing.assert_close(
marlin_delta.float(), triton_delta.float(), atol=0.01, rtol=0.05
)
if __name__ == "__main__":
pytest.main([__file__])