diff --git a/python/sglang/srt/layers/moe/moe_runner/base.py b/python/sglang/srt/layers/moe/moe_runner/base.py index 12dd2ba6a..9bfe4cc46 100644 --- a/python/sglang/srt/layers/moe/moe_runner/base.py +++ b/python/sglang/srt/layers/moe/moe_runner/base.py @@ -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 diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 93bb9cbcb..6578fece6 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -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( diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index 7e62a6424..9dcee54da 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -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( diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py index be40253b3..fcae2cadf 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton.py @@ -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], diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py b/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py index b13cd2759..a90add0fa 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py @@ -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, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index 5734b9dc0..f276fca11 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -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, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py index 072fedfae..0ac18784c 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py @@ -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, diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index 4128d0fda..6c79e9063 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -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), diff --git a/python/sglang/srt/lora/lora_moe_runner_marlin.py b/python/sglang/srt/lora/lora_moe_runner_marlin.py new file mode 100644 index 000000000..58ee9ed01 --- /dev/null +++ b/python/sglang/srt/lora/lora_moe_runner_marlin.py @@ -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) diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index 0f626d740..f07d762ee 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/python/sglang/srt/lora/lora_moe_runners.py @@ -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) diff --git a/test/registered/lora/test_lora_moe_runner.py b/test/registered/lora/test_lora_moe_runner.py new file mode 100644 index 000000000..0aad02395 --- /dev/null +++ b/test/registered/lora/test_lora_moe_runner.py @@ -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__]) diff --git a/test/registered/lora/test_marlin_lora_correctness.py b/test/registered/lora/test_marlin_lora_correctness.py new file mode 100644 index 000000000..bc593debc --- /dev/null +++ b/test/registered/lora/test_marlin_lora_correctness.py @@ -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__])