diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py index 761b3c592..e2005dd40 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py @@ -81,6 +81,21 @@ def is_symmetric_memory_enabled(): return False +def is_tensor_in_symmetric_mempool(tensor: torch.Tensor) -> bool: + """Check if a tensor's storage is allocated in the NCCL symmetric memory pool.""" + + if _mem_pool is None: + return False # Pool not initialized + + data_ptr = tensor.untyped_storage().data_ptr() + + for segment in _mem_pool.snapshot(): + for block in segment["blocks"]: + if block["address"] == data_ptr: + return True + return False + + def set_graph_pool_id(graph_pool_id): global _graph_pool_id _graph_pool_id = graph_pool_id diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1289c6a47..451f1e588 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -182,6 +182,15 @@ class Envs: SGLANG_SORT_WEIGHT_FILES = EnvBool(False) SGLANG_DISABLED_MODEL_ARCHS = EnvTuple(tuple()) + # NVFP4 / flashinfer (B300 port) + SGLANG_EXPERIMENTAL_LORA_OPTI = EnvBool(False) + SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION = EnvBool(False) + FLASHINFER_NVFP4_4OVER6 = EnvBool(False) + FLASHINFER_NVFP4_4OVER6_E4M3_USE_256 = EnvBool(False) + # B300 port D1: gate modelopt-fp8 flashinfer-bmm OFF by default to keep the + # GLM-5.1-FP8 baseline byte-identical (upstream #28333 defaults it on for sm100). + SGLANG_MODELOPT_FP8_FLASHINFER_BMM = EnvBool(False) + # Logging Options SGLANG_LOG_GC = EnvBool(False) SGLANG_LOG_FORWARD_ITERS = EnvBool(False) diff --git a/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py b/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py index c38126679..cbd43d8b2 100644 --- a/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py @@ -28,6 +28,7 @@ def _fake_fp8_block_scale_moe( enable_pdl: Optional[bool] = None, tune_max_num_tokens: int = 8192, fp8_quantization_type: Optional[int] = None, + activation_type: Optional[int] = None, ) -> torch.Tensor: return torch.empty( hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device @@ -58,6 +59,7 @@ def trtllm_fp8_block_scale_moe_wrapper( enable_pdl: Optional[bool] = None, tune_max_num_tokens: int = 8192, fp8_quantization_type: Optional[int] = None, + activation_type: Optional[int] = None, ) -> torch.Tensor: try: from flashinfer.fused_moe import trtllm_fp8_block_scale_moe @@ -94,6 +96,11 @@ def trtllm_fp8_block_scale_moe_wrapper( kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type) + if activation_type is not None: + from flashinfer.fused_moe.core import ActivationType + + kwargs["activation_type"] = ActivationType(activation_type) + return trtllm_fp8_block_scale_moe(**kwargs) @@ -120,6 +127,7 @@ def _fake_fp8_block_scale_routed_moe( enable_pdl: Optional[bool] = None, tune_max_num_tokens: int = 8192, fp8_quantization_type: Optional[int] = None, + activation_type: Optional[int] = None, ) -> torch.Tensor: return torch.empty( hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device @@ -150,6 +158,7 @@ def trtllm_fp8_block_scale_routed_moe_wrapper( enable_pdl: Optional[bool] = None, tune_max_num_tokens: int = 8192, fp8_quantization_type: Optional[int] = None, + activation_type: Optional[int] = None, ) -> torch.Tensor: try: from flashinfer.fused_moe import trtllm_fp8_block_scale_routed_moe @@ -186,6 +195,11 @@ def trtllm_fp8_block_scale_routed_moe_wrapper( kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type) + if activation_type is not None: + from flashinfer.fused_moe.core import ActivationType + + kwargs["activation_type"] = ActivationType(activation_type) + return trtllm_fp8_block_scale_routed_moe(**kwargs) @@ -210,6 +224,7 @@ def _fake_fp8_per_tensor_scale_moe( routing_method_type: int = 0, enable_pdl: Optional[bool] = None, tune_max_num_tokens: int = 8192, + activation_type: Optional[int] = None, ) -> torch.Tensor: return torch.empty( hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device @@ -238,6 +253,7 @@ def trtllm_fp8_per_tensor_scale_moe_wrapper( routing_method_type: int = 0, enable_pdl: Optional[bool] = None, tune_max_num_tokens: int = 8192, + activation_type: Optional[int] = None, ) -> torch.Tensor: # lazy import try: @@ -271,4 +287,9 @@ def trtllm_fp8_per_tensor_scale_moe_wrapper( "tune_max_num_tokens": tune_max_num_tokens, } + if activation_type is not None: + from flashinfer.fused_moe.core import ActivationType + + kwargs["activation_type"] = ActivationType(activation_type) + return trtllm_fp8_per_tensor_scale_moe(**kwargs) diff --git a/python/sglang/srt/layers/moe/moe_runner/base.py b/python/sglang/srt/layers/moe/moe_runner/base.py index 12dd2ba6a..6363a4de8 100644 --- a/python/sglang/srt/layers/moe/moe_runner/base.py +++ b/python/sglang/srt/layers/moe/moe_runner/base.py @@ -1,8 +1,10 @@ from __future__ import annotations +import contextvars from abc import ABC, abstractmethod +from contextlib import contextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, Optional, Tuple, TypeGuard +from typing import TYPE_CHECKING, Any, Callable, Generator, Optional, Tuple, TypeGuard import torch @@ -26,6 +28,20 @@ if TYPE_CHECKING: ) +_moe_output_buf: contextvars.ContextVar[Optional[torch.Tensor]] = ( + contextvars.ContextVar("moe_output_buf", default=None) +) + + +@contextmanager +def moe_output_buffer_ctx(buf: torch.Tensor) -> Generator[None, None, None]: + token = _moe_output_buf.set(buf) + try: + yield + finally: + _moe_output_buf.reset(token) + + @dataclass class MoeRunnerConfig: # MoE parameters @@ -48,6 +64,7 @@ class MoeRunnerConfig: routed_scaling_factor: Optional[float] = None gemm1_alpha: Optional[float] = None gemm1_clamp_limit: Optional[float] = None + swiglu_limit: Optional[float] = None @dataclass @@ -82,7 +99,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/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 7914265dd..c8a88f8fb 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -1,59 +1,155 @@ from __future__ import annotations +import contextvars +from contextlib import contextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Generator, cast import torch from torch.nn import Module from torch.nn.parameter import Parameter # Import to register custom ops for torch.compile compatibility -import sglang.srt.layers.moe.flashinfer_trtllm_moe # noqa: F401 -from sglang.kernel_api_logging import debug_torch_op from sglang.srt.distributed import get_tp_group from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + is_symmetric_memory_enabled, + is_tensor_in_symmetric_mempool, use_symmetric_memory, ) +from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import is_allocation_symmetric +from sglang.srt.layers.moe.flashinfer_trtllm_moe import ( + trtllm_fp8_block_scale_moe_wrapper, + trtllm_fp8_block_scale_routed_moe_wrapper, + trtllm_fp8_per_tensor_scale_moe_wrapper, +) from sglang.srt.layers.moe.moe_runner.base import ( MoeQuantInfo, MoeRunnerConfig, + _moe_output_buf, register_fused_func, ) from sglang.srt.layers.quantization.fp8_kernel import ( per_token_group_quant_fp8, scaled_fp8_quant, ) +from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import PackTopkIds from sglang.srt.layers.utils import copy_or_rebind_param from sglang.srt.utils.common import ( is_cuda_alike, is_flashinfer_available, - is_sm120_supported, next_power_of_2, ) +_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() + +logger = __import__("logging").getLogger(__name__) + +_deferred_finalize_enabled: contextvars.ContextVar[bool] = contextvars.ContextVar( + "flashinfer_trtllm_deferred_finalize_enabled", default=False +) + + +@dataclass +class FlashInferTrtllmDeferredFinalizeOutput: + gemm2_out: torch.Tensor + expert_weights: torch.Tensor + expanded_idx_to_permuted_idx: torch.Tensor + top_k: int + + +@contextmanager +def flashinfer_trtllm_deferred_finalize_context( + enabled: bool = True, +) -> Generator[None, None, None]: + token = _deferred_finalize_enabled.set(enabled) + try: + yield + finally: + _deferred_finalize_enabled.reset(token) + + +def finalize_flashinfer_trtllm_deferred_output( + deferred_output: FlashInferTrtllmDeferredFinalizeOutput, + shared_output: torch.Tensor, +) -> torch.Tensor: + from sglang.jit_kernel.moe_finalize_fuse_shared import moe_finalize_fuse_shared + from sglang.jit_kernel.utils import is_arch_support_pdl + + return moe_finalize_fuse_shared( + deferred_output.gemm2_out, + deferred_output.expanded_idx_to_permuted_idx, + deferred_output.expert_weights, + shared_output, + deferred_output.top_k, + enable_pdl=is_arch_support_pdl(), + ) + + +def round_up_to_multiple(x: int, m: int) -> int: + """Round up *x* to the nearest multiple of *m*.""" + return (x + m - 1) // m * m + + if TYPE_CHECKING: from sglang.srt.layers.moe.token_dispatcher import ( StandardCombineInput, StandardDispatchOutput, ) -if is_flashinfer_available() and is_sm120_supported(): - from flashinfer import fp4_quantize +if is_flashinfer_available(): + from sglang.srt.layers.quantization.fp4_utils import fp4_quantize elif is_cuda_alike(): from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize else: fp4_quantize = None -_trtllm_fp8_block_scale_routed_moe_wrapper = debug_torch_op( - "trtllm_fp8_block_scale_routed_moe_wrapper" -) -_trtllm_fp8_block_scale_moe_wrapper = debug_torch_op( - "trtllm_fp8_block_scale_moe_wrapper" -) -_trtllm_fp8_per_tensor_scale_moe = debug_torch_op( - "trtllm_fp8_per_tensor_scale_moe_wrapper" -) +_flashinfer_trtllm_shuffle_row_indices_cache_mxfp8: dict[ + tuple, dict[str, torch.Tensor] +] = {} + + +def _is_gated(layer: Module) -> bool: + """Return whether the MoE layer uses a gated activation (default True).""" + is_gated = ( + getattr(layer, "moe_runner_config", None) and layer.moe_runner_config.is_gated + ) + return True if is_gated is None else is_gated + + +def _align_fp8_moe_weights( + w13: torch.Tensor, + w2: torch.Tensor, + is_gated: bool, + min_alignment: int = 16, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """Pad intermediate size so FlashInfer TRTLLM FP8 kernels' alignment holds. + + Returns (w13, w2, padded_intermediate). + """ + num_experts, hidden_size, intermediate = w2.shape + + padded_intermediate = round_up_to_multiple(intermediate, min_alignment) + if padded_intermediate == intermediate: + return w13, w2, intermediate + + logger.info( + "FP8 MoE: padding intermediate size from %d to %d (alignment=%d)", + intermediate, + padded_intermediate, + min_alignment, + ) + + up_mult = 2 if is_gated else 1 + padded_gate_up = up_mult * padded_intermediate + + padded_w13 = w13.new_zeros((num_experts, padded_gate_up, w13.shape[2])) + padded_w13[:, : w13.shape[1], :] = w13 + + padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate)) + padded_w2[:, :, :intermediate] = w2 + + return padded_w13, padded_w2, padded_intermediate def align_fp8_moe_weights_for_flashinfer_trtllm( @@ -67,32 +163,47 @@ def align_fp8_moe_weights_for_flashinfer_trtllm( This is needed for ModelOpt FP8 checkpoints which store weights in [Up, Gate] order, while regular FP8 checkpoints store them in [Gate, Up]. """ - from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a + from flashinfer import shuffle_matrix_a + + is_gated = _is_gated(layer) w13_weight = cast(torch.Tensor, layer.w13_weight) w2_weight = cast(torch.Tensor, layer.w2_weight) - num_experts, two_n, hidden = w13_weight.shape + num_experts, gate_up_dim, hidden = w13_weight.shape - # Optionally swap W13 halves: [Up, Gate] -> [Gate, Up] - if swap_w13_halves: - inter = two_n // 2 + # Optionally swap W13 halves: [Up, Gate] -> [Gate, Up] (only for gated) + if swap_w13_halves and is_gated: + inter = gate_up_dim // 2 w13_weight = ( w13_weight.reshape(num_experts, 2, inter, hidden) .flip(dims=[1]) - .reshape(num_experts, two_n, hidden) + .reshape(num_experts, gate_up_dim, hidden) ) - w13_interleaved_list = [ - reorder_rows_for_gated_act_gemm(w13_weight[i]) for i in range(num_experts) - ] - w13_interleaved: torch.Tensor = torch.stack(w13_interleaved_list).reshape( - num_experts, two_n, hidden + # Pad for kernel alignment (non-gated needs 128, gated needs 16) + min_alignment = 16 if is_gated else 128 + w13_weight, w2_weight, _ = _align_fp8_moe_weights( + w13_weight, w2_weight, is_gated, min_alignment ) + num_experts, gate_up_dim, hidden = w13_weight.shape + + epilogue_tile_m = 128 + + if is_gated: + from flashinfer import reorder_rows_for_gated_act_gemm + + w13_interleaved_list = [ + reorder_rows_for_gated_act_gemm(w13_weight[i]) for i in range(num_experts) + ] + w13_processed: torch.Tensor = torch.stack(w13_interleaved_list).reshape( + num_experts, gate_up_dim, hidden + ) + else: + w13_processed = w13_weight # Shuffle weights for transposed MMA output (both W13, W2) - epilogue_tile_m = 128 w13_shuffled = [ - shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m) + shuffle_matrix_a(w13_processed[i].view(torch.uint8), epilogue_tile_m) for i in range(num_experts) ] w2_shuffled = [ @@ -121,7 +232,16 @@ def align_fp8_moe_weights_for_flashinfer_trtllm( w13_weight_scale = cast(torch.Tensor, layer.w13_weight_scale).to(torch.float32) w2_weight_scale = cast(torch.Tensor, layer.w2_weight_scale).to(torch.float32) - output1_scales_scalar = w13_weight_scale * input_scale * (1.0 / activation_scale) + # For gated (SwiGLU): g1_alphas = w1_scale * a1_scale, g1_scale_c = g1_alphas / a2_scale + # For non-gated (Relu2): g1_scale_c = 1 / a2_scale (no gate dequant contribution) + if is_gated: + output1_scales_scalar = ( + w13_weight_scale * input_scale * (1.0 / activation_scale) + ) + else: + output1_scales_scalar = torch.ones_like(w13_weight_scale) * ( + 1.0 / activation_scale + ) output1_scales_gate_scalar = w13_weight_scale * input_scale output2_scales_scalar = activation_scale * w2_weight_scale @@ -132,13 +252,67 @@ def align_fp8_moe_weights_for_flashinfer_trtllm( layer.output2_scales_scalar = Parameter(output2_scales_scalar, requires_grad=False) +def _align_mxfp8_moe_weights( + w13: torch.Tensor, + w13_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + is_gated: bool, + min_alignment: int = 16, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Pad intermediate size so FlashInfer TRTLLM MXFP8 kernels' alignment holds. + + Returns (w13, w13_scale, w2, w2_scale, padded_intermediate). + """ + num_experts, hidden_size, intermediate = w2.shape + + padded_intermediate = round_up_to_multiple(intermediate, min_alignment) + if padded_intermediate == intermediate: + return w13, w13_scale, w2, w2_scale, intermediate + + logger.info( + "MXFP8 MoE: padding intermediate size from %d to %d (alignment=%d)", + intermediate, + padded_intermediate, + min_alignment, + ) + + up_mult = 2 if is_gated else 1 + padded_gate_up = up_mult * padded_intermediate + + padded_w13 = w13.new_zeros((num_experts, padded_gate_up, w13.shape[2])) + padded_w13[:, : w13.shape[1], :] = w13 + + padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate)) + padded_w2[:, :, :intermediate] = w2 + + padded_w13_scale = w13_scale.new_zeros( + (num_experts, padded_gate_up, w13_scale.shape[2]) + ) + padded_w13_scale[:, : w13_scale.shape[1], :] = w13_scale + + # Scale's last dim tracks intermediate / block_size (MXFP8 block_size = 32) + scale_block_k = intermediate // w2_scale.shape[2] if w2_scale.shape[2] > 0 else 32 + padded_w2_scale = w2_scale.new_zeros( + (num_experts, hidden_size, padded_intermediate // scale_block_k) + ) + padded_w2_scale[:, :, : w2_scale.shape[2]] = w2_scale + + return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate + + def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: """Prepare MXFP8 MoE weights/scales for FlashInfer TRT-LLM kernels.""" - from flashinfer import ( - reorder_rows_for_gated_act_gemm, - shuffle_matrix_a, - shuffle_matrix_sf_a, + from flashinfer import block_scale_interleave + from flashinfer.fused_moe.core import ( + get_reorder_rows_for_gated_act_gemm_row_indices, ) + from flashinfer.utils import ( + get_shuffle_matrix_a_row_indices, + get_shuffle_matrix_sf_a_row_indices, + ) + + is_gated = _is_gated(layer) w13_weight = cast(torch.Tensor, layer.w13_weight).contiguous() w2_weight = cast(torch.Tensor, layer.w2_weight).contiguous() @@ -148,65 +322,186 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: assert w13_scale.dtype == torch.uint8 assert w2_scale.dtype == torch.uint8 - num_experts, two_n, _ = w13_weight.shape + # Pad for kernel alignment (non-gated needs 128, gated needs 16) + min_alignment = 16 if is_gated else 128 + w13_weight, w13_scale, w2_weight, w2_scale, _ = _align_mxfp8_moe_weights( + w13_weight, w13_scale, w2_weight, w2_scale, is_gated, min_alignment + ) + + num_experts, gate_up_dim, _ = w13_weight.shape _, hidden_size, _ = w2_weight.shape epilogue_tile_m = 128 - w13_interleaved = [ - reorder_rows_for_gated_act_gemm(w13_weight[i]) for i in range(num_experts) - ] - w13_scale_interleaved = [ - reorder_rows_for_gated_act_gemm(w13_scale[i]) for i in range(num_experts) - ] + # Reuse precomputed row-index transforms whenever shape/device are unchanged. + w13_weight_u8 = w13_weight.view(torch.uint8) + w2_weight_u8 = w2_weight.view(torch.uint8) + cache_key = ( + gate_up_dim, + hidden_size, + w2_weight.shape[-1], + w13_scale.shape[-1], + w2_scale.shape[-1], + epilogue_tile_m, + (w13_weight.device.type, w13_weight.device.index), + (w2_weight.device.type, w2_weight.device.index), + (w13_scale.device.type, w13_scale.device.index), + (w2_scale.device.type, w2_scale.device.index), + ) + cache = _flashinfer_trtllm_shuffle_row_indices_cache_mxfp8.get(cache_key) + if cache is None: + if is_gated: + reorder_row_indices = get_reorder_rows_for_gated_act_gemm_row_indices( + w13_weight_u8[0] + ).to(w13_weight.device) + else: + reorder_row_indices = torch.arange( + gate_up_dim, device=w13_weight.device, dtype=torch.long + ) + w13_shuffle_row_indices = get_shuffle_matrix_a_row_indices( + w13_weight_u8[0], epilogue_tile_m + ).to(w13_weight.device) + w2_shuffle_row_indices = get_shuffle_matrix_a_row_indices( + w2_weight_u8[0], epilogue_tile_m + ).to(w2_weight.device) + w13_scale_shuffle_row_indices = get_shuffle_matrix_sf_a_row_indices( + w13_scale[0].reshape(gate_up_dim, -1), epilogue_tile_m + ).to(w13_scale.device) + w2_scale_shuffle_row_indices = get_shuffle_matrix_sf_a_row_indices( + w2_scale[0].reshape(hidden_size, -1), epilogue_tile_m + ).to(w2_scale.device) + cache = { + "reorder_row_indices": reorder_row_indices, + "w13_shuffle_row_indices": w13_shuffle_row_indices, + "w2_shuffle_row_indices": w2_shuffle_row_indices, + "w13_scale_shuffle_row_indices": w13_scale_shuffle_row_indices, + "w2_scale_shuffle_row_indices": w2_scale_shuffle_row_indices, + } + _flashinfer_trtllm_shuffle_row_indices_cache_mxfp8[cache_key] = cache - w13_shuffled = [ - shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m) - for i in range(num_experts) - ] - w2_shuffled = [ - shuffle_matrix_a(w2_weight[i].view(torch.uint8), epilogue_tile_m) - for i in range(num_experts) - ] - w13_scale_shuffled = [ - shuffle_matrix_sf_a( - w13_scale_interleaved[i].view(torch.uint8).reshape(two_n, -1), - epilogue_tile_m, + reorder_row_indices = cache["reorder_row_indices"] + w13_shuffle_row_indices = cache["w13_shuffle_row_indices"] + w2_shuffle_row_indices = cache["w2_shuffle_row_indices"] + w13_scale_shuffle_row_indices = cache["w13_scale_shuffle_row_indices"] + w2_scale_shuffle_row_indices = cache["w2_scale_shuffle_row_indices"] + + w13_shuffled_u8 = torch.empty_like(w13_weight_u8) + w2_shuffled_u8 = torch.empty_like(w2_weight_u8) + w13_scale_shuffled = torch.empty_like(w13_scale) + w2_scale_shuffled = torch.empty_like(w2_scale) + + for i in range(num_experts): + w13_interleaved_u8 = w13_weight_u8[i].index_select(0, reorder_row_indices) + w13_scale_interleaved = w13_scale[i].index_select(0, reorder_row_indices) + + w13_shuffled_u8[i].copy_( + w13_interleaved_u8.index_select(0, w13_shuffle_row_indices) ) - for i in range(num_experts) - ] - w2_scale_shuffled = [ - shuffle_matrix_sf_a( - w2_scale[i].view(torch.uint8).reshape(hidden_size, -1), - epilogue_tile_m, + w2_shuffled_u8[i].copy_(w2_weight_u8[i].index_select(0, w2_shuffle_row_indices)) + + w13_scale_linear = w13_scale_interleaved.reshape(gate_up_dim, -1) + w13_scale_shuffled[i].copy_( + block_scale_interleave( + w13_scale_linear.index_select(0, w13_scale_shuffle_row_indices) + ).reshape_as(w13_scale_shuffled[i]) + ) + + w2_scale_linear = w2_scale[i].reshape(hidden_size, -1) + w2_scale_shuffled[i].copy_( + block_scale_interleave( + w2_scale_linear.index_select(0, w2_scale_shuffle_row_indices) + ).reshape_as(w2_scale_shuffled[i]) ) - for i in range(num_experts) - ] # Keep parameter identities stable for CUDA graph capture reuse. - copy_or_rebind_param( - layer, "w13_weight", torch.stack(w13_shuffled).view(torch.float8_e4m3fn) - ) - copy_or_rebind_param( - layer, "w2_weight", torch.stack(w2_shuffled).view(torch.float8_e4m3fn) - ) + copy_or_rebind_param(layer, "w13_weight", w13_shuffled_u8.view(torch.float8_e4m3fn)) + copy_or_rebind_param(layer, "w2_weight", w2_shuffled_u8.view(torch.float8_e4m3fn)) copy_or_rebind_param( layer, "w13_weight_scale_inv", - torch.stack(w13_scale_shuffled).reshape_as(w13_scale).contiguous(), + w13_scale_shuffled.contiguous(), ) copy_or_rebind_param( layer, "w2_weight_scale_inv", - torch.stack(w2_scale_shuffled).reshape_as(w2_scale).contiguous(), + w2_scale_shuffled.contiguous(), ) layer.w13_weight_scale_inv.format_ue8m0 = True layer.w2_weight_scale_inv.format_ue8m0 = True +def _align_fp4_moe_weights( + w13: torch.Tensor, + w13_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + is_gated: bool, + min_alignment: int = 16, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Pad intermediate size so FlashInfer TRTLLM FP4 kernels' alignment holds. + + Returns (w13, w13_scale, w2, w2_scale, padded_intermediate). + """ + num_experts, hidden_size, intermediate_packed = w2.shape + intermediate = intermediate_packed * 2 # FP4 packs 2 values per byte + + padded_intermediate = round_up_to_multiple(intermediate, min_alignment) + if padded_intermediate == intermediate: + return w13, w13_scale, w2, w2_scale, intermediate + + logger.info( + "FP4 MoE: padding intermediate size from %d to %d (alignment=%d)", + intermediate, + padded_intermediate, + min_alignment, + ) + + up_mult = 2 if is_gated else 1 + padded_gate_up = up_mult * padded_intermediate + + padded_w13 = w13.new_zeros((num_experts, padded_gate_up, w13.shape[2])) + padded_w13[:, : w13.shape[1], :] = w13 + + padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate // 2)) + padded_w2[:, :, : w2.shape[2]] = w2 + + padded_w13_scale = w13_scale.new_zeros( + (num_experts, padded_gate_up, w13_scale.shape[2]) + ) + padded_w13_scale[:, : w13_scale.shape[1], :] = w13_scale + + padded_w2_scale = w2_scale.new_zeros( + (num_experts, hidden_size, padded_intermediate // 16) + ) + padded_w2_scale[:, :, : w2_scale.shape[2]] = w2_scale + + return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate + + +def _compute_g1_scale_c( + w2_input_scale_quant: torch.Tensor, + g1_alphas: torch.Tensor, + g1_alphas_up: torch.Tensor, + is_gated: bool, +) -> torch.Tensor: + """TRT-LLM GEMM1-output scale for the up (w3) half. + + TRT-LLM dequantizes the two halves of the fused GEMM1 separately: g1_alphas + covers the gate half, this scalar the up half (hence g1_alphas_up). The + 1/a2_scale factor (w2_input_scale_quant) requantizes GEMM2's input. A shared + scale passes g1_alphas as g1_alphas_up and recovers the single-scale value; + non-gated (Relu2) has no gate half, so it is just 1/a2_scale per expert. + """ + if is_gated: + return (w2_input_scale_quant * g1_alphas_up).to(torch.float32) + num_experts = g1_alphas.shape[0] + return w2_input_scale_quant.to(torch.float32).expand(num_experts).contiguous() + + def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: """Prepare FP4 MoE weights/scales for FlashInfer TRT-LLM kernels. This function handles the weight transformation needed for FP4 TRTLLM MoE: + - Pads intermediate dimension for kernel alignment constraints - Reorders weights for gated activation GEMM - Shuffles weights and scales for transposed MMA output - Computes the output scale factors @@ -220,6 +515,21 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: w13_weight_scale = cast(torch.Tensor, layer.w13_weight_scale) w2_weight_scale = cast(torch.Tensor, layer.w2_weight_scale) + is_gated = layer.moe_runner_config.is_gated + min_alignment = 16 if is_gated else 128 + + # Pad for kernel alignment before shuffle/reorder + w13_weight, w13_weight_scale, w2_weight, w2_weight_scale, intermediate_size = ( + _align_fp4_moe_weights( + w13_weight, + w13_weight_scale, + w2_weight, + w2_weight_scale, + is_gated, + min_alignment, + ) + ) + ( gemm1_weights_fp4_shuffled, gemm1_scales_fp4_shuffled, @@ -231,36 +541,57 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: w13_weight_scale, w2_weight_scale, w2_weight.size(-2), # hidden_size - w13_weight.size(-2) // 2, # intermediate_size + intermediate_size, # padded intermediate_size w13_weight.size(0), # num_experts + is_gated=is_gated, ) - # Set flashinfer parameters + # Set flashinfer parameters in-place + copy_or_rebind_param(layer, "w13_weight", gemm1_weights_fp4_shuffled.contiguous()) + copy_or_rebind_param(layer, "w2_weight", gemm2_weights_fp4_shuffled.contiguous()) copy_or_rebind_param( - layer, "gemm1_weights_fp4_shuffled", gemm1_weights_fp4_shuffled + layer, "w13_weight_scale", gemm1_scales_fp4_shuffled.contiguous() ) copy_or_rebind_param( - layer, "gemm2_weights_fp4_shuffled", gemm2_weights_fp4_shuffled + layer, "w2_weight_scale", gemm2_scales_fp4_shuffled.contiguous() ) - copy_or_rebind_param(layer, "gemm1_scales_fp4_shuffled", gemm1_scales_fp4_shuffled) - copy_or_rebind_param(layer, "gemm2_scales_fp4_shuffled", gemm2_scales_fp4_shuffled) - # Compute additional scaling factor needed for TRT-LLM + # Extra GEMM1-output scalar that TRT-LLM needs (up-half dequant). w2_input_scale_quant = cast(torch.Tensor, layer.w2_input_scale_quant) g1_alphas = cast(torch.Tensor, layer.g1_alphas) - copy_or_rebind_param( - layer, - "g1_scale_c", - (w2_input_scale_quant * g1_alphas).to(torch.float32), + g1_alphas_up = cast(torch.Tensor, getattr(layer, "g1_alphas_up", g1_alphas)) + g1_scale_c = _compute_g1_scale_c( + w2_input_scale_quant, g1_alphas, g1_alphas_up, layer.moe_runner_config.is_gated ) + copy_or_rebind_param(layer, "g1_scale_c", g1_scale_c) - # Clean up weights that won't be used by TRT-LLM - del ( - layer.w2_weight, - layer.w2_weight_scale, - layer.w13_weight, - layer.w13_weight_scale, - ) + # Update intermediate_size_per_partition to reflect any padding applied + layer.intermediate_size_per_partition = intermediate_size + + +def get_activation_type(activation: str, is_gated: bool = True) -> int: + """Map SGLang activation string to FlashInfer ActivationType int value.""" + from flashinfer.fused_moe.core import ActivationType + + if is_gated: + _ACTIVATION_STR_TO_TYPE = { + "silu": ActivationType.Swiglu, + "gelu": ActivationType.Geglu, + } + else: + _ACTIVATION_STR_TO_TYPE = { + "silu": ActivationType.Silu, + "gelu": ActivationType.Gelu, + "relu2": ActivationType.Relu2, + } + act = _ACTIVATION_STR_TO_TYPE.get(activation) + if act is None: + raise ValueError( + f"Unsupported activation '{activation}' for TRTLLM MoE " + f"(is_gated={is_gated}). " + f"Expected one of {list(_ACTIVATION_STR_TO_TYPE.keys())}." + ) + return act.value @dataclass @@ -293,16 +624,8 @@ class FlashInferTrtllmFp8MoeQuantInfo(MoeQuantInfo): output2_scales_scalar: torch.Tensor | None = None use_routing_scales_on_input: bool = False - -def _pack_topk_for_flashinfer_routed( - topk_ids: torch.Tensor, topk_weights: torch.Tensor -) -> torch.Tensor: - """Pack routed top-k tensors into FlashInfer's int32 format.""" - packed_ids = topk_ids.to(torch.int32) - packed_weights = topk_weights.to(torch.bfloat16) - packed = (packed_ids << 16) | packed_weights.view(torch.int16).to(torch.int32) - # SGLang can mark padded tokens with -1 expert ids. - return packed.masked_fill_(packed_ids < 0, 0) + # Activation type (None = kernel default / Swiglu) + activation_type: int | None = None def fused_experts_none_to_flashinfer_trtllm_fp8( @@ -317,7 +640,11 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( from sglang.srt.layers.moe.topk import TopKOutputChecker from sglang.srt.layers.moe.utils import RoutingMethodType - assert runner_config.activation == "silu", "Only silu is supported." + _SUPPORTED_FP8_ACTIVATIONS = {"silu", "relu2"} + assert runner_config.activation in _SUPPORTED_FP8_ACTIVATIONS, ( + f"Only {_SUPPORTED_FP8_ACTIVATIONS} are supported for FP8 MoE, " + f"got '{runner_config.activation}'." + ) assert not runner_config.no_combine, "no_combine is not supported for flashinfer." hidden_states = dispatch_output.hidden_states @@ -325,11 +652,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( if TopKOutputChecker.format_is_bypassed(topk_output): router_logits = topk_output.router_logits topk_config = topk_output.topk_config - correction_bias = ( - None - if topk_config.correction_bias is None - else topk_config.correction_bias.to(hidden_states.dtype) - ) + correction_bias = topk_config.correction_bias else: router_logits = None topk_config = None @@ -358,9 +681,9 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( a_sf_t = a_sf.view(torch.uint8).reshape(hidden_states.shape[0], -1) else: a_q, a_sf = per_token_group_quant_fp8( - hidden_states, quant_info.weight_block_k + hidden_states, quant_info.weight_block_k, column_major_scales=True ) - a_sf_t = a_sf.t().contiguous() + a_sf_t = a_sf.t() # Allocate output inside symmetric memory context with use_symmetric_memory( @@ -369,7 +692,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( symm_output = torch.empty( hidden_states.shape[0], hidden_states.shape[1], - dtype=torch.bfloat16, + dtype=hidden_states.dtype, device=hidden_states.device, ) @@ -381,12 +704,11 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( runner_config.top_k is not None ), "runner_config.top_k is required for flashinfer_trtllm_routed." assert TopKOutputChecker.format_is_standard(topk_output) - packed_topk_ids = _pack_topk_for_flashinfer_routed( - topk_ids=topk_output.topk_ids, - topk_weights=topk_output.topk_weights, + packed_topk_ids = PackTopkIds.execute( + topk_output.topk_ids, topk_output.topk_weights ) - output = _trtllm_fp8_block_scale_routed_moe_wrapper( + output = trtllm_fp8_block_scale_routed_moe_wrapper( topk_ids=packed_topk_ids, routing_bias=None, hidden_states=a_q, @@ -415,16 +737,13 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( use_shuffled_weight=use_shuffled_weight, tune_max_num_tokens=next_power_of_2(a_q.shape[0]), fp8_quantization_type=int(fp8_quantization_type), + activation_type=quant_info.activation_type, ) else: assert TopKOutputChecker.format_is_bypassed(topk_output) - output = _trtllm_fp8_block_scale_moe_wrapper( - routing_logits=( - router_logits.to(torch.float32) - if routing_method_type == RoutingMethodType.DeepSeekV3 - else router_logits - ), + output = trtllm_fp8_block_scale_moe_wrapper( + routing_logits=router_logits, routing_bias=correction_bias, hidden_states=a_q, hidden_states_scale=a_sf_t, @@ -448,10 +767,13 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( use_shuffled_weight=use_shuffled_weight, tune_max_num_tokens=next_power_of_2(a_q.shape[0]), fp8_quantization_type=int(fp8_quantization_type), + activation_type=quant_info.activation_type, ) + # TODO: Once https://github.com/flashinfer-ai/flashinfer/issues/2703 is fixed, pass output to moe kernel and remove this copy. symm_output.copy_(output) output = symm_output else: + assert TopKOutputChecker.format_is_bypassed(topk_output) assert quant_info.w13_input_scale is not None assert quant_info.output1_scales_scalar is not None assert quant_info.output1_scales_gate_scalar is not None @@ -476,8 +798,11 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( # Move kernel call outside context manager to avoid graph breaks # during torch.compile for piecewise cuda graph. # Use custom op wrapper for torch.compile compatibility. - output = _trtllm_fp8_per_tensor_scale_moe( - routing_logits=router_logits.to(torch.bfloat16), + + router_logits = router_logits.to(torch.bfloat16) + + output = trtllm_fp8_per_tensor_scale_moe_wrapper( + routing_logits=router_logits, routing_bias=routing_bias_cast, hidden_states=a_q, gemm1_weights=quant_info.w13_weight, @@ -500,6 +825,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( use_routing_scales_on_input=False, routing_method_type=routing_method_type, tune_max_num_tokens=next_power_of_2(a_q.shape[0]), + activation_type=quant_info.activation_type, ) symm_output.copy_(output) output = symm_output @@ -511,11 +837,10 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( class FlashInferTrtllmFp4MoeQuantInfo(MoeQuantInfo): """Quantization payload consumed by FlashInfer TRT-LLM FP4 MoE kernels.""" - # Shuffled FP4 weights (processed by align_fp4_moe_weights_for_flashinfer_trtllm) - gemm1_weights_fp4_shuffled: torch.Tensor - gemm2_weights_fp4_shuffled: torch.Tensor - gemm1_scales_fp4_shuffled: torch.Tensor - gemm2_scales_fp4_shuffled: torch.Tensor + w13_weight: torch.Tensor + w2_weight: torch.Tensor + w13_weight_scale: torch.Tensor + w2_weight_scale: torch.Tensor # Scaling factors g1_scale_c: torch.Tensor @@ -530,6 +855,7 @@ class FlashInferTrtllmFp4MoeQuantInfo(MoeQuantInfo): intermediate_size_per_partition: int routing_method_type: int + use_per_token_activation: bool = False def quantize_hidden_states_fp4( @@ -567,94 +893,219 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( dispatch_output: StandardDispatchOutput, quant_info: FlashInferTrtllmFp4MoeQuantInfo, runner_config: MoeRunnerConfig, + use_routed_topk: bool = False, ) -> StandardCombineInput: """FlashInfer TRTLLM FP4 MoE forward pass. This function handles the FP4 TRTLLM MoE path that was previously in - FlashInferFP4MoE.forward_impl and ModelOptNvFp4FusedMoEMethod.apply. + ModelOptNvFp4FusedMoEMethod.apply. """ - from flashinfer.fused_moe import trtllm_fp4_block_scale_moe + from flashinfer.fused_moe import ( + trtllm_fp4_block_scale_moe, + trtllm_fp4_block_scale_routed_moe, + ) from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker from sglang.srt.layers.moe.utils import RoutingMethodType - assert runner_config.activation == "silu", "Only silu is supported for FP4 MoE." - assert runner_config.is_gated, "Only gated MoEs are supported for FP4 MoE." + _SUPPORTED_FP4_ACTIVATIONS = {"silu", "relu2", "gelu"} + assert runner_config.activation in _SUPPORTED_FP4_ACTIVATIONS, ( + f"Only {_SUPPORTED_FP4_ACTIVATIONS} are supported for FP4 MoE, " + f"got '{runner_config.activation}'." + ) hidden_states = dispatch_output.hidden_states topk_output = dispatch_output.topk_output - assert TopKOutputChecker.format_is_bypassed(topk_output) - - router_logits = topk_output.router_logits - topk_config = topk_output.topk_config - routing_method_type = quant_info.routing_method_type # Quantize hidden states to FP4 - hs_fp4, hs_scale_linear = quantize_hidden_states_fp4( - hidden_states, quant_info.w13_input_scale_quant + if quant_info.use_per_token_activation: + from flashinfer import SfLayout, nvfp4_quantize + + e4m3_max = 448.0 + if ( + envs.FLASHINFER_NVFP4_4OVER6.get() + and envs.FLASHINFER_NVFP4_4OVER6_E4M3_USE_256.get() + ): + e4m3_max = 256.0 + + hs_fp4_bytes, hs_sf_bytes, per_token_scale = nvfp4_quantize( + hidden_states, + 1.0 / (e4m3_max * 6.0), + sfLayout=SfLayout.layout_linear, + per_token_activation=True, + ) + + seq_len, hidden_size = hidden_states.shape + hs_fp4 = hs_fp4_bytes.reshape(seq_len, hidden_size // 2) + hs_scale_linear = hs_sf_bytes.view(torch.float8_e4m3fn).reshape( + seq_len, hidden_size // 16 + ) + else: + per_token_scale = None + hs_fp4, hs_scale_linear = quantize_hidden_states_fp4( + hidden_states, quant_info.w13_input_scale_quant + ) + hs_scale = hs_scale_linear.view(torch.float8_e4m3fn).reshape( + *hs_scale_linear.shape[:-1], -1 + ) + activation_type = get_activation_type( + runner_config.activation, is_gated=runner_config.is_gated ) - # DeepSeekV3 style routing requires float32 router logits - if routing_method_type == RoutingMethodType.DeepSeekV3: - router_logits = router_logits.to(torch.float32) + # Build per-expert clamp-limit tensor from the per-layer scalar. + _clamp_val = runner_config.gemm1_clamp_limit + if _clamp_val is not None: + gemm1_clamp_limit = torch.full( + (quant_info.local_num_experts,), + _clamp_val, + dtype=torch.float32, + device=hs_fp4.device, + ) + else: + gemm1_clamp_limit = None - correction_bias = ( - None - if topk_config.correction_bias is None - else topk_config.correction_bias.to(hidden_states.dtype) + # Fall back to routed path when topk was already materialized (e.g. sigmoid routing). + if not use_routed_topk and TopKOutputChecker.format_is_standard(topk_output): + use_routed_topk = True + + defer_finalize = ( + _deferred_finalize_enabled.get() + and not use_routed_topk + and TopKOutputChecker.format_is_bypassed(topk_output) ) - with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): + symm_output = None + if not defer_finalize: num_tokens = hs_fp4.shape[0] hidden_size = ( hs_fp4.shape[-1] * 2 if hs_fp4.dtype == torch.uint8 else hs_fp4.shape[-1] ) - symm_output = torch.empty( - num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_fp4.device - ) + _provided = _moe_output_buf.get() + _symm_required = is_allocation_symmetric() + if ( + _provided is not None + and _provided.shape == (num_tokens, hidden_size) + and _provided.dtype == hidden_states.dtype + and _provided.device == hs_fp4.device + and ( + not _symm_required + or not is_symmetric_memory_enabled() + or is_tensor_in_symmetric_mempool(_provided) + ) + ): + symm_output = _provided + else: + with use_symmetric_memory(get_tp_group(), disabled=not _symm_required): + symm_output = torch.empty( + hs_fp4.shape[0], + hidden_size, + dtype=hidden_states.dtype, + device=hs_fp4.device, + ) - result = trtllm_fp4_block_scale_moe( - routing_logits=router_logits, - routing_bias=correction_bias, - hidden_states=hs_fp4, - hidden_states_scale=hs_scale_linear.view(torch.float8_e4m3fn).reshape( - *hs_scale_linear.shape[:-1], -1 - ), - gemm1_weights=quant_info.gemm1_weights_fp4_shuffled, - gemm1_weights_scale=quant_info.gemm1_scales_fp4_shuffled.view( - torch.float8_e4m3fn - ), - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=quant_info.gemm2_weights_fp4_shuffled, - gemm2_weights_scale=quant_info.gemm2_scales_fp4_shuffled.view( - torch.float8_e4m3fn - ), - gemm2_bias=None, - output1_scale_scalar=quant_info.g1_scale_c, - output1_scale_gate_scalar=quant_info.g1_alphas, - output2_scale_scalar=quant_info.g2_alphas, - num_experts=quant_info.global_num_experts, - top_k=topk_config.top_k, - n_group=topk_config.num_expert_group, - topk_group=topk_config.topk_group, - intermediate_size=quant_info.intermediate_size_per_partition, - local_expert_offset=quant_info.local_expert_offset, - local_num_experts=quant_info.local_num_experts, - routed_scaling_factor=runner_config.routed_scaling_factor, - tile_tokens_dim=None, - routing_method_type=( - routing_method_type - if routing_method_type is not None - else RoutingMethodType.Default - ), - do_finalize=True, - tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]), - output=symm_output, - )[0] + if use_routed_topk: + assert TopKOutputChecker.format_is_standard(topk_output) + + packed_topk_ids = PackTopkIds.execute( + topk_output.topk_ids, topk_output.topk_weights + ) + result = trtllm_fp4_block_scale_routed_moe( + topk_ids=packed_topk_ids, + routing_bias=None, + hidden_states=hs_fp4, + hidden_states_scale=hs_scale, + gemm1_weights=quant_info.w13_weight, + gemm1_weights_scale=quant_info.w13_weight_scale.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=gemm1_clamp_limit, + gemm2_weights=quant_info.w2_weight, + gemm2_weights_scale=quant_info.w2_weight_scale.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=quant_info.g1_scale_c, + output1_scale_gate_scalar=quant_info.g1_alphas, + output2_scale_scalar=quant_info.g2_alphas, + per_token_scale=per_token_scale, + num_experts=quant_info.global_num_experts, + top_k=topk_output.topk_ids.shape[1], + n_group=0, + topk_group=0, + intermediate_size=quant_info.intermediate_size_per_partition, + local_expert_offset=quant_info.local_expert_offset, + local_num_experts=quant_info.local_num_experts, + routed_scaling_factor=None, + routing_method_type=1, # Unused, but must be 1 to pass validation. + do_finalize=True, + activation_type=activation_type, + tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]), + output=symm_output, + )[0] + else: + assert TopKOutputChecker.format_is_bypassed(topk_output) + + router_logits = topk_output.router_logits + topk_config = topk_output.topk_config + routing_method_type = quant_info.routing_method_type + + correction_bias = topk_config.correction_bias + moe_kwargs = dict( + routing_logits=router_logits, + routing_bias=correction_bias, + hidden_states=hs_fp4, + hidden_states_scale=hs_scale, + gemm1_weights=quant_info.w13_weight, + gemm1_weights_scale=quant_info.w13_weight_scale.view(torch.float8_e4m3fn), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=gemm1_clamp_limit, + gemm2_weights=quant_info.w2_weight, + gemm2_weights_scale=quant_info.w2_weight_scale.view(torch.float8_e4m3fn), + gemm2_bias=None, + output1_scale_scalar=quant_info.g1_scale_c, + output1_scale_gate_scalar=quant_info.g1_alphas, + output2_scale_scalar=quant_info.g2_alphas, + per_token_scale=per_token_scale, + num_experts=quant_info.global_num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=quant_info.intermediate_size_per_partition, + local_expert_offset=quant_info.local_expert_offset, + local_num_experts=quant_info.local_num_experts, + routed_scaling_factor=runner_config.routed_scaling_factor, + routing_method_type=( + routing_method_type + if routing_method_type is not None + else RoutingMethodType.Default + ), + do_finalize=not defer_finalize, + activation_type=activation_type, + tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]), + ) + if not defer_finalize: + moe_kwargs["output"] = symm_output + + result = trtllm_fp4_block_scale_moe(**moe_kwargs) + if defer_finalize: + gemm2_out, expert_weights, expanded_idx_to_permuted_idx = result[:3] + # FIXME(kpham-sgl): flashinfer sizes this buffer from routing_logits + # dtype (fp32 in DSv3 decode) but always writes bf16 weights into it. + # Reinterpret the live bf16 prefix. Fix upstream alloc to drop this, + # tracking in https://github.com/flashinfer-ai/flashinfer/issues/3595 + if expert_weights.dtype == torch.float32: + n, k = expert_weights.shape + expert_weights = expert_weights.view(torch.bfloat16).view(-1, k)[:n] + result = FlashInferTrtllmDeferredFinalizeOutput( + gemm2_out=gemm2_out, + expert_weights=expert_weights, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + top_k=topk_config.top_k, + ) + else: + result = result[0] return StandardCombineInput(hidden_states=result) @@ -701,9 +1152,11 @@ def fused_experts_none_to_flashinfer_trtllm_bf16( "Please check flashinfer version to use bf16 with flashinfer_trtllm backend." ) from e - assert ( - runner_config.activation == "silu" - ), "Only silu is supported for flashinfer trtllm moe" + _SUPPORTED_BF16_ACTIVATIONS = {"silu", "relu2"} + assert runner_config.activation in _SUPPORTED_BF16_ACTIVATIONS, ( + f"Only {_SUPPORTED_BF16_ACTIVATIONS} are supported for flashinfer trtllm bf16 moe, " + f"got '{runner_config.activation}'." + ) if not use_routed_topk: assert ( dispatch_output.topk_output.topk_config.renormalize @@ -711,9 +1164,9 @@ def fused_experts_none_to_flashinfer_trtllm_bf16( assert ( runner_config.num_fused_shared_experts == 0 ), "Fused shared experts are not supported for flashinfer trtllm moe" - assert ( - runner_config.is_gated - ), "Only gated MoEs are supported for flashinfer trtllm moe" + activation_type = get_activation_type( + runner_config.activation, is_gated=runner_config.is_gated + ) hidden_states = dispatch_output.hidden_states topk_output = dispatch_output.topk_output @@ -730,9 +1183,8 @@ def fused_experts_none_to_flashinfer_trtllm_bf16( elif routing_method_type == RoutingMethodType.DeepSeekV3: routing_method_type = RoutingMethodType.TopK - packed_topk_ids = _pack_topk_for_flashinfer_routed( - topk_ids=topk_output.topk_ids, - topk_weights=topk_output.topk_weights, + packed_topk_ids = PackTopkIds.execute( + topk_output.topk_ids, topk_output.topk_weights ) final_hidden_states = trtllm_bf16_routed_moe( topk_ids=packed_topk_ids, @@ -753,6 +1205,7 @@ def fused_experts_none_to_flashinfer_trtllm_bf16( else 1.0 ), tune_max_num_tokens=next_power_of_2(hidden_states.shape[0]), + activation_type=activation_type, ) else: assert TopKOutputChecker.format_is_bypassed(topk_output) @@ -775,6 +1228,7 @@ def fused_experts_none_to_flashinfer_trtllm_bf16( routing_method_type=runner_config.routing_method_type, routed_scaling_factor=runner_config.routed_scaling_factor, tune_max_num_tokens=next_power_of_2(hidden_states.shape[0]), + activation_type=activation_type, ) return StandardCombineInput(hidden_states=final_hidden_states) @@ -810,6 +1264,13 @@ def fused_experts_none_to_flashinfer_trtllm_routed( quant_info: MoeQuantInfo, runner_config: MoeRunnerConfig, ) -> StandardCombineInput: + if isinstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo): + return fused_experts_none_to_flashinfer_trtllm_fp4( + dispatch_output, + quant_info, + runner_config, + use_routed_topk=True, + ) if isinstance(quant_info, FlashInferTrtllmFp8MoeQuantInfo): return fused_experts_none_to_flashinfer_trtllm_fp8( dispatch_output, @@ -827,3 +1288,9 @@ def fused_experts_none_to_flashinfer_trtllm_routed( raise TypeError( f"Unexpected quant_info type for flashinfer_trtllm_routed: {type(quant_info)}" ) + + +# Register the experimental experimental_sgl_trtllm MoE fused-func (MoeRunner needs it at +# build time even for LoRA); gated by the master switch so the upstream path is untouched. +if _SGLANG_EXPERIMENTAL_LORA_OPTI: + from sglang.srt.lora.trtllm_lora_temp import sgl_backend # noqa: E402,F401 diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 65da83c0e..b5600069f 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -200,6 +200,14 @@ def get_moe_runner_backend() -> MoeRunnerBackend: return MOE_RUNNER_BACKEND +def is_flashinfer_cutedsl_v1_path() -> bool: + """CuteDSL v1 + DeepEP low-latency path (no MoeRunner, no autotune).""" + return ( + get_moe_runner_backend().is_flashinfer_cutedsl() + and get_moe_a2a_backend().is_deepep() + ) + + def get_speculative_moe_runner_backend() -> MoeRunnerBackend: global SPECULATIVE_MOE_RUNNER_BACKEND if SPECULATIVE_MOE_RUNNER_BACKEND is None: diff --git a/python/sglang/srt/layers/quantization/fp4_utils.py b/python/sglang/srt/layers/quantization/fp4_utils.py index 3e913e137..39ffbb935 100644 --- a/python/sglang/srt/layers/quantization/fp4_utils.py +++ b/python/sglang/srt/layers/quantization/fp4_utils.py @@ -2,10 +2,16 @@ from __future__ import annotations import logging from enum import Enum -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional -from sglang.srt.environ import envs -from sglang.srt.utils.common import is_sm120_supported +import torch + +from sglang.srt.utils.common import ( + get_device_capability, + is_cuda, + is_sm100_supported, +) +from sglang.srt.utils.custom_op import register_custom_op_from_extern if TYPE_CHECKING: from sglang.srt.server_args import ServerArgs @@ -13,17 +19,94 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +fp4_quantize = None +try: + from flashinfer import fp4_quantize as _flashinfer_fp4_quantize + + _flashinfer_fp4_quantize_backend = "cute-dsl" if is_sm100_supported() else "cuda" + + def _round_up(x: int, y: int) -> int: + return ((x + y - 1) // y) * y + + def _flashinfer_fp4_quantize_impl( + input: torch.Tensor, + global_scale: Optional[torch.Tensor] = None, + sf_vec_size: int = 16, + sf_use_ue8m0: bool = False, + is_sf_swizzled_layout: bool = True, + is_sf_8x4_layout: bool = False, + enable_pdl: Optional[bool] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return _flashinfer_fp4_quantize( + input=input, + global_scale=global_scale, + sf_vec_size=sf_vec_size, + sf_use_ue8m0=sf_use_ue8m0, + is_sf_swizzled_layout=is_sf_swizzled_layout, + is_sf_8x4_layout=is_sf_8x4_layout, + enable_pdl=enable_pdl, + backend=_flashinfer_fp4_quantize_backend, + ) + + def _flashinfer_fp4_quantize_fake( + input: torch.Tensor, + global_scale: Optional[torch.Tensor] = None, + sf_vec_size: int = 16, + sf_use_ue8m0: bool = False, + is_sf_swizzled_layout: bool = True, + is_sf_8x4_layout: bool = False, + enable_pdl: Optional[bool] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + is_column_major = input.stride(-2) == 1 + if is_column_major: + m = input.shape[-1] + K = input.shape[-2] + else: + m = input.numel() // input.shape[-1] + K = input.shape[-1] + if is_column_major: + x_q = input.new_empty((*input.shape[:-2], K // 2, m), dtype=torch.uint8) + else: + x_q = input.new_empty((*input.shape[:-1], K // 2), dtype=torch.uint8) + if is_sf_swizzled_layout: + row_size = 8 if is_sf_8x4_layout else 128 + sf_rows = _round_up(m, row_size) + sf_cols = _round_up(K // sf_vec_size, 4) + else: + sf_rows = m + sf_cols = K // sf_vec_size + if is_column_major: + sf = input.new_empty((sf_cols, sf_rows), dtype=torch.uint8) + else: + sf = input.new_empty((sf_rows, sf_cols), dtype=torch.uint8) + return x_q, sf + + fp4_quantize = register_custom_op_from_extern( + _flashinfer_fp4_quantize_impl, + op_name="flashinfer_fp4_quantize", + fake_impl=_flashinfer_fp4_quantize_fake, + ) +except ImportError: + fp4_quantize = None + + class Fp4GemmRunnerBackend(Enum): """Enum for FP4 GEMM runner backend selection.""" AUTO = "auto" + CUTLASS = "cutlass" FLASHINFER_CUDNN = "flashinfer_cudnn" + FLASHINFER_CUTEDSL = "flashinfer_cutedsl" FLASHINFER_CUTLASS = "flashinfer_cutlass" FLASHINFER_TRTLLM = "flashinfer_trtllm" + MARLIN = "marlin" def is_auto(self) -> bool: return self == Fp4GemmRunnerBackend.AUTO + def is_cutlass(self) -> bool: + return self == Fp4GemmRunnerBackend.CUTLASS + def is_flashinfer_cudnn(self) -> bool: return self == Fp4GemmRunnerBackend.FLASHINFER_CUDNN @@ -33,6 +116,15 @@ class Fp4GemmRunnerBackend(Enum): def is_flashinfer_trtllm(self) -> bool: return self == Fp4GemmRunnerBackend.FLASHINFER_TRTLLM + def is_flashinfer_cutedsl(self) -> bool: + return self == Fp4GemmRunnerBackend.FLASHINFER_CUTEDSL + + def is_marlin(self) -> bool: + return self == Fp4GemmRunnerBackend.MARLIN + + def is_flashinfer(self) -> bool: + return self.value.startswith("flashinfer_") + def get_flashinfer_backend(self) -> str: """Get the backend string to pass to FlashInfer's mm_fp4 API. @@ -41,7 +133,10 @@ class Fp4GemmRunnerBackend(Enum): 'flashinfer_trtllm' -> 'trtllm' 'flashinfer_cutlass' -> 'cutlass' 'flashinfer_cudnn' -> 'cudnn' + 'flashinfer_cutedsl' -> 'cute-dsl' """ + if self == Fp4GemmRunnerBackend.FLASHINFER_CUTEDSL: + return "cute-dsl" if self.value.startswith("flashinfer_"): return self.value.removeprefix("flashinfer_") else: @@ -56,36 +151,11 @@ def initialize_fp4_gemm_config(server_args: ServerArgs) -> None: global FP4_GEMM_RUNNER_BACKEND backend = server_args.fp4_gemm_runner_backend - - # Handle deprecated env var for backward compatibility - # TODO: Remove this in a future version - if envs.SGLANG_FLASHINFER_FP4_GEMM_BACKEND.is_set(): - env_backend = envs.SGLANG_FLASHINFER_FP4_GEMM_BACKEND.get() - if backend == "auto": - logger.warning( - "SGLANG_FLASHINFER_FP4_GEMM_BACKEND is deprecated. " - f"Please use '--fp4-gemm-backend={env_backend}' instead." - ) - if not env_backend.startswith("flashinfer_"): - env_backend = "flashinfer_" + env_backend - backend = env_backend - else: - logger.warning( - f"FP4 GEMM backend set to '{backend}' via --fp4-gemm-backend overrides " - "environment variable SGLANG_FLASHINFER_FP4_GEMM_BACKEND. " - "Using server argument value." - ) - if backend == "auto": - if is_sm120_supported(): - # flashinfer_cutlass produces NaN in dense MLP layers with - # heterogeneous batches on SM120 (Blackwell). cudnn is stable. - # See: https://github.com/sgl-project/sglang/issues/20043 - backend = "flashinfer_cudnn" - logger.info( - "SM120 (Blackwell) detected: auto-selecting " - "fp4-gemm-backend=flashinfer_cudnn" - ) + if is_sm100_supported(): + backend = "flashinfer_cutedsl" + elif is_cuda() and (10, 0) > get_device_capability() >= (8, 0): + backend = "marlin" else: backend = "flashinfer_cutlass" diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 53e540294..9c1b36b29 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -211,6 +211,7 @@ def _check_cutlass_block_fp8_hardware_support() -> bool: if is_blackwell_supported() and is_flashinfer_available(): + from flashinfer import bmm_fp8 as flashinfer_bmm_fp8 from flashinfer import mm_mxfp8 as _raw_flashinfer_mm_mxfp8 from flashinfer import mxfp8_quantize as _raw_flashinfer_mxfp8_quantize from flashinfer.gemm import gemm_fp8_nt_groupwise as _raw_gemm_fp8_nt_groupwise @@ -1368,6 +1369,28 @@ def _apply_fallback_scaled_mm( return output.to(dtype=input_dtype) +def apply_fp8_linear_bmm_flashinfer( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: torch.Tensor, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Per-tensor static fp8 linear via flashinfer bmm_fp8 (SM10X only). + + B300 port: gated OFF by default (SGLANG_MODELOPT_FP8_FLASHINFER_BMM). `static_quant_fp8` + is a module-level import; `flashinfer_bmm_fp8` is imported under the Blackwell guard above + (this path only runs on sm100, where that guard is taken). + """ + output_shape = [*input.shape[:-1], weight.shape[1]] + input_2d = input.view(-1, input.shape[-1]) + qinput, x_scale = static_quant_fp8(input_2d, input_scale, repeat_scale=False) + output = flashinfer_bmm_fp8(qinput, weight, x_scale, weight_scale, input.dtype) + if bias is not None: + output = output + bias + return output.view(*output_shape) + + def apply_fp8_linear( input: torch.Tensor, weight: torch.Tensor, diff --git a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py new file mode 100644 index 000000000..44cd5181e --- /dev/null +++ b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py @@ -0,0 +1,518 @@ +from __future__ import annotations + +import torch + +from sglang.srt.layers.quantization.marlin_utils import ( + USE_FP32_REDUCE_DEFAULT, + marlin_make_workspace, + marlin_permute_bias, + marlin_permute_scales, + should_use_atomic_add_reduce, +) +from sglang.srt.layers.quantization.utils import get_scalar_types +from sglang.srt.utils import is_cuda +from sglang.srt.utils.custom_op import register_custom_op + +_is_cuda = is_cuda() + +if _is_cuda: + from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm + from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack + +ScalarType, scalar_types = get_scalar_types() + + +def nvfp4_marlin_process_scales(marlin_scales: torch.Tensor) -> torch.Tensor: + if not (marlin_scales >= 0).all(): + # NVFP4 ModelOpt scales are expected to be non-negative. Keep this as + # a warning so unusual checkpoints can still load for diagnosis. + import logging + + logging.getLogger(__name__).warning_once( + "NVFP4 Marlin assumes non-negative scales, but negative scales " + "were found. Accuracy may be degraded." + ) + + marlin_scales = marlin_scales.to(torch.half) + marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view( + marlin_scales.size(0), -1 + ) + marlin_scales = (marlin_scales * (2**7)).view(torch.int16) << 1 + marlin_scales = marlin_scales.view(torch.float8_e4m3fn) + return marlin_scales[:, 1::2].contiguous() + + +def nvfp4_marlin_process_global_scale(global_scale: torch.Tensor) -> torch.Tensor: + assert global_scale.dtype in [torch.half, torch.bfloat16] + global_scale_shape = global_scale.shape + fp4_exponent = 2 + if global_scale.dtype == torch.half: + target_exponent = 5 + elif global_scale.dtype == torch.bfloat16: + target_exponent = 8 + exponent_bias = 2 ** (target_exponent - 1) - 2 ** (fp4_exponent - 1) + global_scale = global_scale * (2.0 ** (exponent_bias - 7)) + if global_scale_shape == torch.Size([]): + global_scale = global_scale.reshape(1) + return global_scale + + +def fake_apply_fp4_marlin_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_global_scale: torch.Tensor, + workspace: torch.Tensor, + size_n: int, + size_k: int, + bias: torch.Tensor | None = None, + use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, +) -> torch.Tensor: + del weight, weight_scale, weight_global_scale, workspace, size_k, bias + out_shape = input.shape[:-1] + (size_n,) + return input.new_empty(out_shape) + + +@register_custom_op(fake_impl=fake_apply_fp4_marlin_linear) +def apply_fp4_marlin_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_global_scale: torch.Tensor, + workspace: torch.Tensor, + size_n: int, + size_k: int, + bias: torch.Tensor | None = None, + use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, +) -> torch.Tensor: + if input.dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError("NVFP4 Marlin requires FP16 or BF16 activations.") + + reshaped_x = input.reshape(-1, input.shape[-1]) + out_shape = input.shape[:-1] + (size_n,) + + use_atomic_add = should_use_atomic_add_reduce( + m=reshaped_x.size(0), + n=size_n, + k=size_k, + device=input.device, + dtype=input.dtype, + ) + + output = gptq_marlin_gemm( + a=reshaped_x, + c=None, + b_q_weight=weight, + b_scales=weight_scale, + global_scale=weight_global_scale, + b_zeros=None, + g_idx=None, + perm=None, + workspace=workspace, + b_q_type=scalar_types.float4_e2m1f, + size_m=reshaped_x.size(0), + size_n=size_n, + size_k=size_k, + is_k_full=True, + use_atomic_add=use_atomic_add, + use_fp32_reduce=use_fp32_reduce, + ) + + if bias is not None: + output.add_(bias) + + return output.reshape(out_shape) + + +def prepare_nvfp4_layer_for_marlin(layer: torch.nn.Module) -> None: + if getattr(layer, "quant_config", None) is not None: + group_size = layer.quant_config.group_size + if group_size != 16: + raise ValueError(f"NVFP4 Marlin requires group_size=16, got {group_size}.") + + part_size_n = layer.output_size_per_partition + part_size_k = layer.input_size_per_partition + param_dtype = getattr(layer, "params_dtype", getattr(layer, "orig_dtype", None)) + if param_dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError("NVFP4 Marlin requires FP16 or BF16 activation dtype.") + + assert layer.weight.shape == (part_size_n, part_size_k // 2) + + if part_size_n % 64 != 0: + raise ValueError( + f"NVFP4 Marlin requires output_size_per_partition to be a multiple of 64, " + f"got {part_size_n}." + ) + + device = layer.weight.device + layer.workspace = marlin_make_workspace(device) + + perm = torch.empty(0, dtype=torch.int, device=device) + qweight = layer.weight.view(torch.int32).T.contiguous() + marlin_qweight = gptq_marlin_repack( + b_q_weight=qweight, + perm=perm, + size_k=part_size_k, + size_n=part_size_n, + num_bits=4, + ) + layer.weight = torch.nn.Parameter(marlin_qweight, requires_grad=False) + + weight_scale = layer.weight_scale.T.contiguous().to(param_dtype) + weight_scale = marlin_permute_scales( + s=weight_scale, + size_k=part_size_k, + size_n=part_size_n, + group_size=16, + ) + weight_scale = nvfp4_marlin_process_scales(weight_scale) + layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + + weight_global_scale = layer.weight_global_scale.to(param_dtype) + weight_global_scale = nvfp4_marlin_process_global_scale(weight_global_scale) + layer.weight_global_scale = torch.nn.Parameter( + weight_global_scale, requires_grad=False + ) + + if hasattr(layer, "bias") and layer.bias is not None: + assert layer.bias.shape == (part_size_n,) + bias = marlin_permute_bias(layer.bias) + layer.bias = torch.nn.Parameter(bias, requires_grad=False) + + +def mxfp4_marlin_process_scales( + marlin_scales: torch.Tensor, + input_dtype: torch.dtype | None = None, +) -> torch.Tensor: + if input_dtype is None or input_dtype.itemsize == 2: + marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view( + marlin_scales.size(0), -1 + ) + marlin_scales = marlin_scales.to(torch.float8_e8m0fnu) + if input_dtype == torch.float8_e4m3fn: + marlin_scales = marlin_scales.view(torch.uint8) + assert marlin_scales.max() <= 249 + # exponent_bias (fp4->fp8) = 2 ** 3 - 2 ** 1 = 6 + marlin_scales = marlin_scales + 6 + marlin_scales = marlin_scales.view(torch.float8_e8m0fnu) + return marlin_scales + + +def _normalize_scale_tensor( + scales: torch.Tensor, target_dtype: torch.dtype +) -> torch.Tensor: + # The kernel consumes E8M0 exponents. Regardless of the placeholder dtype + # the loader used, we want the *numerical* value 2**e in ``target_dtype``. + # float32/bfloat16/float16 containers hold the numerical 2**e directly + # (they were filled via a dtype-promoting copy from uint8/e8m0). + # uint8/int8 containers hold the raw E8M0 byte and must be reinterpreted. + if scales.dtype == torch.float8_e8m0fnu: + return scales.to(target_dtype) + if scales.dtype == torch.uint8: + return scales.view(torch.float8_e8m0fnu).to(target_dtype) + if scales.dtype == torch.int8: + return scales.view(torch.uint8).view(torch.float8_e8m0fnu).to(target_dtype) + if scales.dtype in (torch.float32, torch.bfloat16, torch.float16): + return scales.to(target_dtype) + raise TypeError(f"Unsupported MXFP4 scale dtype for Marlin: {scales.dtype}") + + +def _get_optional_param(layer: torch.nn.Module, *names: str) -> torch.Tensor | None: + for name in names: + value = getattr(layer, name, None) + if value is not None: + return value + return None + + +def deinterleave_moe_mxfp4_w13_for_marlin(layer: torch.nn.Module) -> None: + """Convert GPT-OSS interleaved w13 rows to Marlin's contiguous halves. + + GPT-OSS stores gate/up rows as [gate0, up0, gate1, up1, ...]. The Marlin + fused activation consumes [all_gate_rows, all_up_rows]. + """ + + w13 = layer.w13_weight.data + w13_scale = _get_optional_param(layer, "w13_weight_scale", "w13_weight_scale_inv") + w13_bias = _get_optional_param(layer, "w13_weight_bias", "w13_bias") + + if w13.shape[1] % 2 != 0: + raise ValueError(f"Expected even w13 row dimension, got {w13.shape}.") + + e, n, k = w13.shape + layer.w13_weight.data = ( + w13.view(e, n // 2, 2, k).permute(0, 2, 1, 3).contiguous().view(e, n, k) + ) + + if w13_scale is not None: + scale = w13_scale.data + if scale.shape[1] != n: + raise ValueError( + f"Expected w13 scale row dimension {n}, got {scale.shape}." + ) + w13_scale.data = ( + scale.view(e, n // 2, 2, scale.shape[-1]) + .permute(0, 2, 1, 3) + .contiguous() + .view(e, n, scale.shape[-1]) + ) + + if w13_bias is not None: + bias = w13_bias.data + if bias.shape[1] != n: + raise ValueError(f"Expected w13 bias row dimension {n}, got {bias.shape}.") + w13_bias.data = bias.view(e, n // 2, 2).permute(0, 2, 1).contiguous().view(e, n) + + +def prepare_moe_mxfp4_layer_for_marlin(layer: torch.nn.Module) -> None: + group_size = 32 + w13 = layer.w13_weight.data + w2 = layer.w2_weight.data + w13_scale = _get_optional_param(layer, "w13_weight_scale", "w13_weight_scale_inv") + w2_scale = _get_optional_param(layer, "w2_weight_scale", "w2_weight_scale_inv") + w13_bias = _get_optional_param(layer, "w13_weight_bias", "w13_bias") + w2_bias = _get_optional_param(layer, "w2_weight_bias", "w2_bias") + + if w13_scale is None or w2_scale is None: + raise ValueError("MXFP4 Marlin requires w13/w2 weight scales.") + + w13_scale_data = w13_scale.data if hasattr(w13_scale, "data") else w13_scale + w2_scale_data = w2_scale.data if hasattr(w2_scale, "data") else w2_scale + w13_bias_data = w13_bias.data if hasattr(w13_bias, "data") else w13_bias + w2_bias_data = w2_bias.data if hasattr(w2_bias, "data") else w2_bias + + num_experts = w13.shape[0] + intermediate_size = w13.shape[1] // 2 + hidden_size = w13.shape[2] * 2 + if hidden_size % 128 == 0: + padded_intermediate_size = ((intermediate_size + 63) // 64) * 64 + else: + if hidden_size % 64 != 0: + raise ValueError( + f"MXFP4 Marlin requires hidden_size to be divisible by 64, " + f"got {hidden_size}." + ) + padded_intermediate_size = ((intermediate_size + 127) // 128) * 128 + param_dtype = getattr( + layer, + "orig_dtype", + w13_bias_data.dtype if w13_bias_data is not None else torch.bfloat16, + ) + + device = w13.device + layer.workspace = marlin_make_workspace(device, 4) + perm = torch.empty(0, dtype=torch.int, device=device) + + def _pad_w13(x: torch.Tensor) -> torch.Tensor: + if padded_intermediate_size == intermediate_size: + return x + x = x.view(num_experts, 2, intermediate_size, x.shape[-1]) + x = torch.nn.functional.pad( + x, (0, 0, 0, padded_intermediate_size - intermediate_size) + ) + return x.reshape(num_experts, 2 * padded_intermediate_size, -1) + + def _pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor: + if padded_intermediate_size == intermediate_size: + return x + return torch.nn.functional.pad( + x, (0, (padded_intermediate_size - intermediate_size) // packing) + ) + + w13 = _pad_w13(w13) + w2 = _pad_w2(w2, packing=2) + w13_scale_data = _pad_w13(_normalize_scale_tensor(w13_scale_data, param_dtype)) + w2_scale_data = _pad_w2( + _normalize_scale_tensor(w2_scale_data, param_dtype), + packing=group_size, + ) + if w13_bias_data is not None: + w13_bias_data = _pad_w13(w13_bias_data.unsqueeze(-1)).squeeze(-1) + + def _repack_weight(weight: torch.Tensor, is_w13: bool) -> torch.Tensor: + if is_w13: + size_n, size_k = padded_intermediate_size * 2, hidden_size + else: + size_n, size_k = hidden_size, padded_intermediate_size + assert weight.shape == (num_experts, size_n, size_k // 2) + + tensor_list = [] + for i in range(num_experts): + qweight = weight[i].view(torch.int32).T.contiguous() + marlin_qweight = gptq_marlin_repack( + b_q_weight=qweight, + perm=perm, + size_k=size_k, + size_n=size_n, + num_bits=4, + ) + tensor_list.append(marlin_qweight) + return torch.stack(tensor_list) + + def _permute_scales(scales: torch.Tensor, is_w13: bool) -> torch.Tensor: + if is_w13: + size_n, size_k = padded_intermediate_size * 2, hidden_size + else: + size_n, size_k = hidden_size, padded_intermediate_size + + tensor_list = [] + for i in range(num_experts): + scale = scales[i].T.contiguous() + marlin_scales = marlin_permute_scales( + s=scale, + size_k=size_k, + size_n=size_n, + group_size=group_size, + ) + tensor_list.append( + mxfp4_marlin_process_scales( + marlin_scales, + input_dtype=param_dtype, + ) + ) + return torch.stack(tensor_list) + + def _permute_bias(bias: torch.Tensor | None) -> torch.Tensor | None: + if bias is None: + return None + tensor_list = [] + for i in range(num_experts): + tensor_list.append(marlin_permute_bias(bias[i].to(param_dtype))) + return torch.stack(tensor_list) + + w13_marlin = _repack_weight(w13, True) + w2_marlin = _repack_weight(w2, False) + w13_scale_marlin = _permute_scales(w13_scale_data, True) + w2_scale_marlin = _permute_scales(w2_scale_data, False) + + layer.w13_weight = torch.nn.Parameter(w13_marlin, requires_grad=False) + layer.w2_weight = torch.nn.Parameter(w2_marlin, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter(w13_scale_marlin, requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter(w2_scale_marlin, requires_grad=False) + + if w13_bias_data is not None: + layer.w13_weight_bias = torch.nn.Parameter( + _permute_bias(w13_bias_data), requires_grad=False + ) + if w2_bias_data is not None: + layer.w2_weight_bias = torch.nn.Parameter( + _permute_bias(w2_bias_data), requires_grad=False + ) + + +def prepare_moe_nvfp4_layer_for_marlin(layer: torch.nn.Module) -> None: + if layer.quant_config.group_size != 16: + raise ValueError( + f"NVFP4 Marlin MoE requires group_size=16, got {layer.quant_config.group_size}." + ) + + w13 = layer.w13_weight.data + w2 = layer.w2_weight.data + w13_scale = layer.w13_weight_scale.data + w2_scale = layer.w2_weight_scale.data + w13_global_scale = layer.w13_weight_scale_2.data + w2_global_scale = layer.w2_weight_scale_2.data + w13_bias = getattr(layer, "w13_bias", None) + w2_bias = getattr(layer, "w2_bias", None) + + num_experts = w13.shape[0] + num_shards = 2 if layer.moe_runner_config.is_gated else 1 + intermediate_size = layer.intermediate_size_per_partition + hidden_size = w13.shape[2] * 2 + param_dtype = layer.params_dtype + if param_dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError("NVFP4 Marlin MoE requires FP16 or BF16 activations.") + + device = w13.device + layer.workspace = marlin_make_workspace(device, 4) + perm = torch.empty(0, dtype=torch.int, device=device) + + if not layer.moe_runner_config.is_gated: + padded_intermediate_size = ((intermediate_size + 127) // 128) * 128 + intermediate_size_pad = padded_intermediate_size - intermediate_size + if intermediate_size_pad: + w13 = torch.nn.functional.pad(w13, (0, 0, 0, intermediate_size_pad)) + w13_scale = torch.nn.functional.pad( + w13_scale, (0, 0, 0, intermediate_size_pad) + ) + w2 = torch.nn.functional.pad(w2, (0, intermediate_size_pad // 2, 0, 0)) + w2_scale = torch.nn.functional.pad( + w2_scale, (0, intermediate_size_pad // 16) + ) + if w13_bias is not None: + w13_bias = torch.nn.functional.pad(w13_bias, (0, intermediate_size_pad)) + intermediate_size = padded_intermediate_size + + def _repack_weight(weight: torch.Tensor, is_w13: bool) -> torch.Tensor: + if is_w13: + size_n, size_k = intermediate_size * num_shards, hidden_size + else: + size_n, size_k = hidden_size, intermediate_size + assert weight.shape == (num_experts, size_n, size_k // 2) + + tensor_list = [] + for i in range(num_experts): + qweight = weight[i].view(torch.int32).T.contiguous() + marlin_qweight = gptq_marlin_repack( + b_q_weight=qweight, + perm=perm, + size_k=size_k, + size_n=size_n, + num_bits=4, + ) + tensor_list.append(marlin_qweight) + return torch.stack(tensor_list) + + def _permute_scales(scales: torch.Tensor, is_w13: bool) -> torch.Tensor: + scales = scales.to(param_dtype) + if is_w13: + size_n, size_k = intermediate_size * num_shards, hidden_size + else: + size_n, size_k = hidden_size, intermediate_size + + tensor_list = [] + for i in range(num_experts): + scale = scales[i].T.contiguous() + marlin_scales = marlin_permute_scales( + s=scale, + size_k=size_k, + size_n=size_n, + group_size=16, + ) + tensor_list.append(nvfp4_marlin_process_scales(marlin_scales)) + return torch.stack(tensor_list) + + def _process_global_scale(global_scale: torch.Tensor) -> torch.Tensor: + return nvfp4_marlin_process_global_scale(global_scale.to(param_dtype)) + + def _permute_bias(bias: torch.Tensor | None) -> torch.Tensor | None: + if bias is None: + return None + tensor_list = [] + for i in range(num_experts): + tensor_list.append(marlin_permute_bias(bias[i].to(param_dtype))) + return torch.stack(tensor_list) + + layer.w13_weight = torch.nn.Parameter( + _repack_weight(w13, True), requires_grad=False + ) + layer.w2_weight = torch.nn.Parameter(_repack_weight(w2, False), requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter( + _permute_scales(w13_scale, True), requires_grad=False + ) + layer.w2_weight_scale = torch.nn.Parameter( + _permute_scales(w2_scale, False), requires_grad=False + ) + layer.w13_weight_scale_2 = torch.nn.Parameter( + _process_global_scale(w13_global_scale), requires_grad=False + ) + layer.w2_weight_scale_2 = torch.nn.Parameter( + _process_global_scale(w2_global_scale), requires_grad=False + ) + + if w13_bias is not None: + layer.w13_bias = torch.nn.Parameter( + _permute_bias(w13_bias), requires_grad=False + ) + if w2_bias is not None: + layer.w2_bias = torch.nn.Parameter(_permute_bias(w2_bias), requires_grad=False) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index dc4fa71fc..efa58209b 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py from __future__ import annotations @@ -24,7 +26,10 @@ from sglang.srt.layers.moe import ( ) from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo -from sglang.srt.layers.moe.utils import should_use_flashinfer_cutlass_moe_fp4_allgather +from sglang.srt.layers.moe.utils import ( + is_flashinfer_cutedsl_v1_path, + should_use_flashinfer_cutlass_moe_fp4_allgather, +) from sglang.srt.layers.parameter import ModelWeightParameter, PerTensorScaleParameter from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, @@ -32,14 +37,23 @@ from sglang.srt.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, ) -from sglang.srt.layers.quantization.fp4_utils import get_fp4_gemm_runner_backend +from sglang.srt.layers.quantization.fp4_utils import ( + fp4_quantize, + get_fp4_gemm_runner_backend, +) from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant from sglang.srt.layers.quantization.fp8_utils import ( apply_fp8_linear, + apply_fp8_linear_bmm_flashinfer, cutlass_fp8_supported, is_blackwell_supported, ) from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod +from sglang.srt.layers.quantization.marlin_utils_fp4 import ( + apply_fp4_marlin_linear, + prepare_moe_nvfp4_layer_for_marlin, + prepare_nvfp4_layer_for_marlin, +) from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.quantization.utils import ( convert_to_channelwise, @@ -49,18 +63,20 @@ from sglang.srt.layers.quantization.utils import ( swizzle_blockscale, ) from sglang.srt.layers.radix_attention import RadixAttention -from sglang.srt.layers.utils import copy_or_rebind_param +from sglang.srt.layers.utils import alias_or_bind_derived_param, copy_or_rebind_param from sglang.srt.utils.common import ( - get_bool_env_var, + get_device_capability, is_cuda, + is_flashinfer_available, + is_sm100_supported, is_sm120_supported, next_power_of_2, + round_up, ) from sglang.srt.utils.custom_op import register_custom_op from sglang.srt.utils.patch_torch import register_fake_if_exists if TYPE_CHECKING: - from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.token_dispatcher import ( CombineInput, @@ -68,31 +84,25 @@ if TYPE_CHECKING: ) from sglang.srt.models.utils import WeightsMapper -fp4_quantize = None -try: - if is_sm120_supported(): - try: - from flashinfer import fp4_quantize - except ImportError: - from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize - else: - from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize -except ImportError: - fp4_quantize = None - try: from flashinfer import mm_fp4 as flashinfer_fp4_gemm from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_sf_a enable_flashinfer_fp4_gemm = True except ImportError: - if is_cuda(): - from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm as cutlass_fp4_gemm enable_flashinfer_fp4_gemm = False reorder_rows_for_gated_act_gemm = None shuffle_matrix_a = None shuffle_matrix_sf_a = None +if is_cuda(): + try: + from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm as cutlass_fp4_gemm + except ImportError: + cutlass_fp4_gemm = None +else: + cutlass_fp4_gemm = None + try: from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe from flashinfer.fused_moe.core import ActivationType @@ -102,7 +112,9 @@ except ImportError: # Define a minimal ActivationType enum if flashinfer is not available class ActivationType(IntEnum): Swiglu = 3 + Geglu = 4 Relu2 = 6 + Identity = 7 # Initialize logger for the module @@ -134,7 +146,15 @@ def fp4_gemm( out_features: int, ) -> torch.Tensor: fp4_backend = get_fp4_gemm_runner_backend() - if enable_flashinfer_fp4_gemm: + if fp4_backend.is_cutlass() and cutlass_fp4_gemm is not None: + # flashinfer.fp4_quantize returns scale factors as uint8 (e4m3fn bits + # stored in uint8 memory). The JIT kernel requires float8_e4m3fn dtype. + if input_sf.dtype != torch.float8_e4m3fn: + input_sf = input_sf.view(torch.float8_e4m3fn) + if weight_sf.dtype != torch.float8_e4m3fn: + weight_sf = weight_sf.view(torch.float8_e4m3fn) + return cutlass_fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype) + elif enable_flashinfer_fp4_gemm: # Use the remapping logic to convert SGLang backend names to FlashInfer API names backend = fp4_backend.get_flashinfer_backend() return flashinfer_fp4_gemm( @@ -153,10 +173,6 @@ if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None): return -CUTEDSL_MOE_SCALAR_INPUT_SCALE = get_bool_env_var( - "SGLANG_CUTEDSL_MOE_SCALAR_INPUT_SCALE", "true" -) - # FP4 GEMM alignment constant - CUTLASS/FlashInfer kernels require dimensions divisible by 32 FP4_GEMM_ALIGNMENT = 32 @@ -259,10 +275,8 @@ MOE_NVFP4_DISPATCH = envs.SGLANG_MOE_NVFP4_DISPATCH.get() # Supported activation schemes for the current configuration ACTIVATION_SCHEMES = ["static"] -ACT_STR_TO_TYPE_MAP = { - "silu": ActivationType.Swiglu, # This is the default - "relu2": ActivationType.Relu2, -} + +_SUPPORTED_ACT_STRS = ("silu", "relu2", "gelu") class ModelOptQuantConfig(QuantizationConfig): @@ -276,6 +290,7 @@ class ModelOptQuantConfig(QuantizationConfig): self.packed_modules_mapping = packed_modules_mapping self.exclude_modules = exclude_modules or [] self.kv_cache_quant_algo = kv_cache_quant_algo + self.use_per_token_activation = False def _get_quant_method( self, @@ -313,7 +328,7 @@ class ModelOptQuantConfig(QuantizationConfig): return [] def apply_weight_name_mapper( - self, hf_to_sglang_mapper: "WeightsMapper" + self, hf_to_sglang_mapper: WeightsMapper ): # noqa: B027 # Map excluded module patterns from HF layout to sglang layout. # Ref: HF hf_quant_config.json for nvidia/Kimi-K2.5-NVFP4 @@ -498,6 +513,13 @@ class ModelOptFp8LinearMethod(LinearMethodBase): super().__init__() self.quant_config = quant_config self.cutlass_fp8_supported = cutlass_fp8_supported() + # B300 port D1: gate OFF by default so the GLM-5.1-FP8 baseline (deepep/FlashMLA) + # stays byte-identical; opt-in via SGLANG_MODELOPT_FP8_FLASHINFER_BMM after measuring. + self.enable_flashinfer_bmm = ( + is_sm100_supported() + and is_flashinfer_available() + and envs.SGLANG_MODELOPT_FP8_FLASHINFER_BMM.get() + ) def create_weights( self, @@ -559,8 +581,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase): layer.weight, layer.weight_scale, layer.logical_widths ) layer.weight = Parameter(quantized_weight.t(), requires_grad=False) - # cutlass sgl-kernel only supports per-channel scale - if self.cutlass_fp8_supported: + if self.cutlass_fp8_supported and not self.enable_flashinfer_bmm: max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths) layer.weight_scale = Parameter(max_w_scale, requires_grad=False) layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) @@ -572,6 +593,14 @@ class ModelOptFp8LinearMethod(LinearMethodBase): bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Applies FP8 linear transformation.""" + if self.enable_flashinfer_bmm and layer.input_scale is not None: + return apply_fp8_linear_bmm_flashinfer( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + input_scale=layer.input_scale, + bias=bias, + ) return apply_fp8_linear( input=x, weight=layer.weight, @@ -601,7 +630,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig): packed_modules_mapping: Optional[Dict[str, List[str]]], quantized_layers: Dict[str, Dict[str, Any]], fp8_config: ModelOptFp8Config, - nvfp4_config: "ModelOptFp4Config", + nvfp4_config: ModelOptFp4Config, ) -> None: super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping) self.quantized_layers = quantized_layers @@ -629,7 +658,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig): return ModelOptFp4Config.get_min_capability() @classmethod - def from_config(cls, config: Dict[str, Any]) -> "ModelOptMixedPrecisionConfig": + def from_config(cls, config: Dict[str, Any]) -> ModelOptMixedPrecisionConfig: kv_cache_quant_algo = None exclude_modules = None quantized_layers = {} @@ -700,7 +729,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig): nvfp4_config=nvfp4_config, ) - def apply_weight_name_mapper(self, hf_to_sglang_mapper: "WeightsMapper"): + def apply_weight_name_mapper(self, hf_to_sglang_mapper: WeightsMapper): super().apply_weight_name_mapper(hf_to_sglang_mapper) if self.quantized_layers: self.quantized_layers = hf_to_sglang_mapper.apply_dict( @@ -966,6 +995,43 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): ) layer.fc1_input_dequant = Parameter(input_scale, requires_grad=False) + # flashinfer_cutlass kernel requires intermediate_size to be a + # multiple of 16. Pad weight tensors with zeros after loading. + # For gated activations (swiglu), w13 is [Up, Gate] concatenated + # along dim 1 — we must split, pad each half separately, and + # re-concat so the kernel's half-split stays aligned. + num_shards = 2 if layer.moe_runner_config.is_gated else 1 + isp = layer.w13_weight.shape[1] // num_shards + if isp % 16 != 0: + pad_amount = round_up(isp, 16) - isp + w13_data = layer.w13_weight.data + if num_shards == 2: + up_weight = w13_data[:, :isp, :] + gate_weight = w13_data[:, isp:, :] + layer.w13_weight = Parameter( + torch.cat( + [ + torch.nn.functional.pad( + up_weight, (0, 0, 0, pad_amount) + ), + torch.nn.functional.pad( + gate_weight, (0, 0, 0, pad_amount) + ), + ], + dim=1, + ), + requires_grad=False, + ) + else: + layer.w13_weight = Parameter( + torch.nn.functional.pad(w13_data, (0, 0, 0, pad_amount)), + requires_grad=False, + ) + layer.w2_weight = Parameter( + torch.nn.functional.pad(layer.w2_weight.data, (0, pad_amount)), + requires_grad=False, + ) + def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): @@ -979,9 +1045,10 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): ) -> CombineInput: x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + from sglang.srt.layers.moe.topk import TopKOutputChecker # Fast path: TRT-LLM FP8 per-tensor MoE using BYPASSED TopK routing - from sglang.srt.layers.moe.topk import TopKOutputChecker if ( get_moe_runner_backend().is_flashinfer_trtllm() @@ -995,20 +1062,19 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): topk_config = topk_output.topk_config - # Constraints for ModelOpt FP8 MoE - assert ( - self.moe_runner_config.activation == "silu" - ), "Only silu is supported for flashinfer fp8 moe" + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + get_activation_type, + ) - # Enforce Llama4 routing for ModelOpt FP8 MoE for now. - # TODO(brayden): support other routing methods - assert topk_config.top_k == 1, "ModelOpt FP8 MoE requires top_k==1" - assert ( - not topk_config.num_expert_group - ), "ModelOpt FP8 MoE does not support expert grouping" - assert ( - not topk_config.topk_group - ), "ModelOpt FP8 MoE does not support grouped top-k" + _SUPPORTED_FP8_ACTIVATIONS = {"silu", "relu2"} + assert self.moe_runner_config.activation in _SUPPORTED_FP8_ACTIVATIONS, ( + f"Only {_SUPPORTED_FP8_ACTIVATIONS} are supported for " + f"flashinfer trtllm fp8 moe, got '{self.moe_runner_config.activation}'" + ) + + routing_method_type = getattr( + layer, "routing_method_type", RoutingMethodType.Llama4 + ) quant_info = FlashInferTrtllmFp8MoeQuantInfo( w13_weight=layer.w13_weight, @@ -1017,13 +1083,17 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, local_num_experts=layer.num_local_experts, intermediate_size=layer.w2_weight.shape[2], - routing_method_type=RoutingMethodType.Llama4, + routing_method_type=routing_method_type, block_quant=False, w13_input_scale=layer.w13_input_scale, output1_scales_scalar=layer.output1_scales_scalar, output1_scales_gate_scalar=layer.output1_scales_gate_scalar, output2_scales_scalar=layer.output2_scales_scalar, use_routing_scales_on_input=True, + activation_type=get_activation_type( + self.moe_runner_config.activation, + is_gated=self.moe_runner_config.is_gated, + ), ) return fused_experts_none_to_flashinfer_trtllm_fp8( @@ -1031,15 +1101,33 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): ) if get_moe_runner_backend().is_flashinfer_cutlass(): - activation = ACT_STR_TO_TYPE_MAP[self.moe_runner_config.activation] - assert ( - ( - activation is ActivationType.Relu2 - and not self.moe_runner_config.is_gated + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + get_activation_type, + ) + + activation_str = self.moe_runner_config.activation + assert activation_str in _SUPPORTED_ACT_STRS, ( + f"Activation {activation_str!r} is not supported for " + f"flashinfer cutlass fp8 moe (supported: {_SUPPORTED_ACT_STRS})." + ) + activation = ActivationType( + get_activation_type( + activation_str, is_gated=self.moe_runner_config.is_gated ) - or activation is ActivationType.Swiglu - and self.moe_runner_config.is_gated - ), "Only Relu2 non-gated or Swiglu gated are supported for flashinfer cutlass fp8 moe" + ) + # FlashInfer CUTLASS MoE supports gated Swiglu/Geglu and non-gated + # Relu2/Identity. Non-gated Silu/Gelu are not implemented. + _CUTLASS_SUPPORTED = { + ActivationType.Swiglu, + ActivationType.Geglu, + ActivationType.Relu2, + ActivationType.Identity, + } + assert activation in _CUTLASS_SUPPORTED, ( + f"Activation {activation_str!r} (is_gated=" + f"{self.moe_runner_config.is_gated}) maps to {activation.name}, " + "which is not supported by flashinfer cutlass fp8 moe." + ) topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids x_fp8, _ = scaled_fp8_quant(x, layer.w13_input_scale) output_dtype = x.dtype @@ -1075,7 +1163,11 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): activation_type=activation, )[0] - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + if ( + not layer.should_fuse_routed_scaling_factor_in_topk + and self.moe_runner_config.routed_scaling_factor is not None + ): + output.mul_(self.moe_runner_config.routed_scaling_factor) return StandardCombineInput(hidden_states=output) @@ -1103,6 +1195,7 @@ class ModelOptFp4Config(ModelOptQuantConfig): group_size: int = None, exclude_modules: List[str] = None, packed_modules_mapping: Optional[Dict[str, List[str]]] = None, + use_per_token_activation: Optional[bool] = None, ) -> None: super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping) self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized @@ -1112,6 +1205,10 @@ class ModelOptFp4Config(ModelOptQuantConfig): "format is experimental and subject to change." ) self.group_size = group_size + self.use_per_token_activation = ( + use_per_token_activation + or envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get() + ) @classmethod def override_quantization_method(cls, hf_quant_config, user_quant): @@ -1128,7 +1225,7 @@ class ModelOptFp4Config(ModelOptQuantConfig): @classmethod def get_min_capability(cls) -> int: - return 100 + return 80 @staticmethod def common_group_size(cfg: dict) -> int: @@ -1259,7 +1356,7 @@ class ModelOptFp4Config(ModelOptQuantConfig): layer, prefix, Linear=ModelOptFp4LinearMethod, - Moe=ModelOptNvFp4FusedMoEMethod, # FlashInferFP4MoE needs the same quantization method but with compatible attribute handling + Moe=ModelOptNvFp4FusedMoEMethod, ) @@ -1305,9 +1402,11 @@ class ModelOptFp4LinearMethod(LinearMethodBase): layer.input_size_per_partition = input_size_per_partition layer.output_size_per_partition = output_size_per_partition + layer.params_dtype = params_dtype + layer.quant_config = self.quant_config if input_size_per_partition % 16 != 0: raise ValueError( - "Unsupported model when in features size is " "not multiple of 16" + "Unsupported model when in features size is not multiple of 16" ) weight_dtype = ( @@ -1359,7 +1458,9 @@ class ModelOptFp4LinearMethod(LinearMethodBase): input_scale_2 = layer.input_scale.max().to(torch.float32) weight_scale_2 = layer.weight_scale_2.max().to(torch.float32) - # Keep per-shard scales intact for hot reload; derive scalar params below. + # alpha / input_scale_inv stay as scalar Parameters. Aliasing them into + # the [N_partitions] source slot breaks fused-QKV linears whose + # downstream kernels assume scalar input scale. copy_or_rebind_param( layer, "alpha", (input_scale_2 * weight_scale_2).to(torch.float32) ) @@ -1370,6 +1471,23 @@ class ModelOptFp4LinearMethod(LinearMethodBase): # Store original output size before any padding layer.output_size_per_partition = layer.weight.shape[0] + if get_fp4_gemm_runner_backend().is_marlin(): + if self.quant_config.group_size != 16: + raise ValueError( + f"NVFP4 Marlin requires group_size=16, got {self.quant_config.group_size}." + ) + copy_or_rebind_param(layer, "input_global_scale", input_scale_2) + copy_or_rebind_param(layer, "weight_global_scale", weight_scale_2) + prepare_nvfp4_layer_for_marlin(layer) + layer.weights_padding_cols = 0 + return + + if not is_blackwell_supported(): + raise ValueError( + "ModelOpt NVFP4 native dense GEMM backends require SM100+. " + "Use --fp4-gemm-backend marlin on SM80-SM90." + ) + if get_fp4_gemm_runner_backend().is_flashinfer_trtllm(): # FlashInfer TRTLLM FP4 GEMM requires a different weight layout. # FlashInfer provides nvfp4_quantize to quantize + shuffle the @@ -1416,7 +1534,9 @@ class ModelOptFp4LinearMethod(LinearMethodBase): .view(torch.float8_e4m3fn) ) - copy_or_rebind_param(layer, "weight_scale_interleaved", scale) + alias_or_bind_derived_param( + layer, "weight_scale", "weight_scale_interleaved", scale + ) copy_or_rebind_param(layer, "weight", weight) layer.weights_padding_cols = weights_padding_cols return @@ -1437,6 +1557,15 @@ class ModelOptFp4LinearMethod(LinearMethodBase): K_padded = round_up_to_multiple(K, 4) padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype) padded_scales[:B, :M, :K] = scales + + # Snapshot the raw (pre-swizzle) scale BEFORE alias_or_bind_derived_param + # overwrites layer.weight_scale.data in-place via .copy_() on the broadcast + # path. Without this, the swiglu side-channel below would read the swizzled + # bytes when it later re-reads layer.weight_scale. + raw_scale_snapshot = ( + (scales.squeeze(0) if scale_ndim == 2 else scales).detach().clone() + ) + batches, rows, cols = padded_scales.shape assert rows % 128 == 0 assert cols % 4 == 0 @@ -1448,7 +1577,55 @@ class ModelOptFp4LinearMethod(LinearMethodBase): if scale_ndim == 2 else padded_scales.reshape(B, M_padded, K_padded) ) - copy_or_rebind_param(layer, "weight_scale_interleaved", padded_scales) + alias_or_bind_derived_param( + layer, "weight_scale", "weight_scale_interleaved", padded_scales + ) + + if getattr(layer, "_interleave_for_swiglu_fusion", False): + from sglang.srt.layers.quantization.nvfp4_gemm_swiglu_nvfp4_quant import ( + interleave_linear_and_gate, + swizzle_blockscale_2d, + ) + + w = layer.weight.data + assert weights_padding_cols == 0, ( + "_interleave_for_swiglu_fusion does not support K-padded weights; " + f"got weights_padding_cols={weights_padding_cols}." + ) + assert raw_scale_snapshot.shape[0] == w.shape[0], ( + "_interleave_for_swiglu_fusion requires no N-padding; " + f"raw_scale rows={raw_scale_snapshot.shape[0]} vs weight rows={w.shape[0]}." + ) + assert w.shape[0] % 128 == 0, ( + "_interleave_for_swiglu_fusion requires N % 128 == 0 (group_size=64 " + f"with gate+up halves); got N={w.shape[0]}." + ) + + gate_w, up_w = w.chunk(2, dim=0) + w_swiglu = interleave_linear_and_gate( + torch.cat((up_w, gate_w), dim=0), group_size=64, dim=0 + ) + + gate_s, up_s = raw_scale_snapshot.chunk(2, dim=0) + w_scale_swiglu = swizzle_blockscale_2d( + interleave_linear_and_gate( + torch.cat((up_s, gate_s), dim=0), group_size=64, dim=0 + ) + ) + + layer.weight_swiglu_interleaved = w_swiglu + layer.weight_scale_swiglu_interleaved = w_scale_swiglu + + # Keep the Parameter objects alive so weight reload can refill + # them and re-run this hook; free their storage in the meantime. + layer.weight.data = torch.empty( + 0, dtype=layer.weight.dtype, device=layer.weight.device + ) + layer.weight_scale_interleaved.data = torch.empty( + 0, + dtype=layer.weight_scale_interleaved.dtype, + device=layer.weight_scale_interleaved.device, + ) def apply( self, @@ -1456,17 +1633,33 @@ class ModelOptFp4LinearMethod(LinearMethodBase): x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - output_dtype = x.dtype - x_m, _ = x.shape + if get_fp4_gemm_runner_backend().is_marlin(): + return apply_fp4_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + weight_global_scale=layer.weight_global_scale, + workspace=layer.workspace, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) + + # `_accepts_prequantized_fp4` is the explicit opt-in so an accidental + # tuple from unrelated code can't silently bypass quantization. + if getattr(layer, "_accepts_prequantized_fp4", False) and isinstance(x, tuple): + x_fp4, x_scale_interleaved = x + x_m = x_fp4.shape[0] + output_dtype = layer.params_dtype + else: + x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv) + x_m, _ = x.shape + output_dtype = x.dtype - # Get original output size (before padding) and padded weight size output_size = layer.output_size_per_partition w_n, _ = layer.weight.shape output_shape = [x_m, output_size] - # Quantize BF16 or FP16 to (FP4 and interleaved block scale) - x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv) - assert x_fp4.dtype == torch.uint8 assert layer.weight.dtype == torch.uint8 assert layer.weight_scale_interleaved.dtype == torch.float8_e4m3fn @@ -1478,7 +1671,10 @@ class ModelOptFp4LinearMethod(LinearMethodBase): w = layer.weight w_scale_interleaved = layer.weight_scale_interleaved - if enable_flashinfer_fp4_gemm: + if ( + enable_flashinfer_fp4_gemm + and not get_fp4_gemm_runner_backend().is_cutlass() + ): w = layer.weight.T w_scale_interleaved = layer.weight_scale_interleaved.T @@ -1500,6 +1696,32 @@ class ModelOptFp4LinearMethod(LinearMethodBase): return out.view(*output_shape) +def _compute_gemm1_alphas( + w13_weight_scale_2: torch.Tensor, + w13_input_scale: torch.Tensor, + is_gated: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """GEMM1 weight x input alphas for the gate (w1) and up (w3) halves of w13. + + w13 fuses the gate and up projections, which may carry separate NVFP4 weight + scales stored as [num_experts, 2] (col 0 = gate, col 1 = up). A 1-D (or + [num_experts, 1]) scale, and any non-gated layer, shares one scale across + both halves; the col-1 read is guarded so those cases stay in bounds. + + Returns (g1_alphas, g1_alphas_up), equal for a shared scale. Single-alpha + backends use g1_alphas; the TRT-LLM path also uses g1_alphas_up. + """ + if is_gated and w13_weight_scale_2.dim() == 2 and w13_weight_scale_2.shape[1] >= 2: + gate_scale = w13_weight_scale_2[:, 0] + up_scale = w13_weight_scale_2[:, 1] + else: + gate_scale = w13_weight_scale_2.reshape(w13_weight_scale_2.shape[0]) + up_scale = gate_scale + g1_alphas = (w13_input_scale * gate_scale).to(torch.float32) + g1_alphas_up = (w13_input_scale * up_scale).to(torch.float32) + return g1_alphas, g1_alphas_up + + class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): """ MoE Method for FP4 Quantization with Blockscales and PerTensorScales @@ -1509,14 +1731,21 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): def __init__(self, quant_config: ModelOptFp4Config): self.quant_config = quant_config - if not is_blackwell_supported(): + moe_runner_backend = get_moe_runner_backend() + if moe_runner_backend.is_auto() and is_cuda(): + capability = get_device_capability() + use_marlin_fallback = (8, 0) <= capability < (10, 0) + else: + use_marlin_fallback = moe_runner_backend.is_marlin() + if not is_blackwell_supported() and not use_marlin_fallback: raise ValueError( "Current platform does not support NVFP4" - " quantization. Please use Blackwell and" - " above." + " quantization with the selected MoE backend. Please use " + "Blackwell and above, or use moe_runner_backend=marlin on SM80+." ) self.enable_flashinfer_trtllm_moe = ( get_moe_runner_backend().is_flashinfer_trtllm() + or get_moe_runner_backend().is_flashinfer_trtllm_routed() ) self._cache_permute_indices = {} @@ -1529,11 +1758,34 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): @property def enable_flashinfer_cutedsl_moe(self) -> bool: + """Access the global enable_flashinfer_cutedsl_moe setting.""" from sglang.srt.layers.moe import get_moe_runner_backend - """Access the global enable_flashinfer_cutedsl_moe setting.""" return get_moe_runner_backend().is_flashinfer_cutedsl() + # ----- CuteDSL v1 vs v2 path helpers ----- + # + # "v1": cutedsl + deepep low-latency. + # - MoeRunner fused func calls flashinfer_cutedsl_moe_masked + # (grouped_gemm_nt_masked). + # - Expects W13 in default [Gate, Up] order, NOT interleaved. + # - Uses swizzled blockscales directly (w13_blockscale_swizzled). + # + # "v2" (standard): cutedsl + none/flashinfer a2a. + # - MoeRunner fused func calls CuteDslMoEWrapper kernels. + # - Expects W13 in [Up, Gate] order, interleaved in 64-row chunks. + # - Uses MMA-layout blockscales (w13_blockscale_mma). + + @property + def _is_cutedsl_v1_deepep(self) -> bool: + """CuteDSL v1 + DeepEP low-latency path (masked grouped GEMM).""" + return is_flashinfer_cutedsl_v1_path() + + @property + def _is_cutedsl_v2_standard(self) -> bool: + """CuteDSL v2 standard path (a2a=none or flashinfer, uses CuteDslMoEWrapper).""" + return self.enable_flashinfer_cutedsl_moe and not self._is_cutedsl_v1_deepep + def create_weights( self, layer: torch.nn.Module, @@ -1543,11 +1795,23 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): params_dtype: torch.dtype, **extra_weight_attrs, ): - if not self.quant_config.is_checkpoint_nvfp4_serialized: + is_nvfp4_online = getattr(self.quant_config, "is_nvfp4_online", False) + if not self.quant_config.is_checkpoint_nvfp4_serialized and not is_nvfp4_online: raise ValueError( "NVFP4 quantization was selected, " " dynamic quantization is not supported." ) + # `nvfp4_online` is not a serialized checkpoint format, but after the + # online loader converts each expert it uses the same packed NVFP4 + # weights, block scales, and per-tensor scales as serialized ModelOpt + # NVFP4 checkpoints. Reuse this layout and swap only the weight loader. + if is_nvfp4_online: + if not self.enable_flashinfer_trtllm_moe: + raise ValueError( + "--quantization nvfp4_online supports only " + "--moe-runner-backend flashinfer_trtllm or " + "flashinfer_trtllm_routed." + ) # TODO(ch-wan): check if this is needed layer.intermediate_size_per_partition = intermediate_size_per_partition @@ -1557,6 +1821,8 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): weight_dtype = torch.uint8 weight_scale_dtype = torch.float8_e4m3fn weight_loader = extra_weight_attrs.get("weight_loader") + if is_nvfp4_online: + weight_loader = self.get_online_weight_loader(layer, weight_loader) # GEMM 1 num_shards = 2 if layer.moe_runner_config.is_gated else 1 @@ -1602,10 +1868,15 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): ) layer.register_parameter("w13_weight_scale", w13_weight_scale) - # Only use `swizzle_blockscale` for shapes, not for real content - layer.w13_blockscale_swizzled = Parameter( - swizzle_blockscale(layer.w13_weight_scale), requires_grad=False - ) + # TRTLLM replaces blockscale_swizzled with an alias to weight_scale + # during process_weights_after_loading, so skip the expensive + # swizzle+allocate here to avoid GPU memory fragmentation + if self.enable_flashinfer_trtllm_moe: + layer.w13_blockscale_swizzled = None + else: + layer.w13_blockscale_swizzled = Parameter( + swizzle_blockscale(layer.w13_weight_scale), requires_grad=False + ) w2_weight_scale = ModelWeightParameter( data=torch.empty( @@ -1620,9 +1891,12 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): ) layer.register_parameter("w2_weight_scale", w2_weight_scale) - layer.w2_blockscale_swizzled = Parameter( - swizzle_blockscale(layer.w2_weight_scale), requires_grad=False - ) + if self.enable_flashinfer_trtllm_moe: + layer.w2_blockscale_swizzled = None + else: + layer.w2_blockscale_swizzled = Parameter( + swizzle_blockscale(layer.w2_weight_scale), requires_grad=False + ) from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported @@ -1647,6 +1921,23 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): ) layer.register_parameter("w2_weight_scale_2", w2_weight_scale_2) + if is_nvfp4_online and self.quant_config.is_checkpoint_fp8_serialized: + # FP8 checkpoints usually store expert scales as weight_scale_inv. + # Online NVFP4 consumes them in the loader and writes the generated + # NVFP4 scales into w*_weight_scale / w*_weight_scale_2 instead. + w13_source_weight_scale_inv = PerTensorScaleParameter( + data=torch.empty(0, dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter( + "w13_weight_scale_inv", w13_source_weight_scale_inv + ) + w2_source_weight_scale_inv = PerTensorScaleParameter( + data=torch.empty(0, dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("w2_weight_scale_inv", w2_source_weight_scale_inv) + extra_weight_attrs.update( {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} ) @@ -1671,44 +1962,51 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): Only supports pre-quantized checkpoints with FP8 weights and scales. """ + # GEMM1 scale processing is deferred until the input scale is known; + # see _compute_gemm1_alphas, which splits w13's gate/up weight scales. + moe_runner_backend = getattr( + self, "_moe_runner_backend", get_moe_runner_backend() + ) + if moe_runner_backend.is_marlin(): + # Marlin supports only a single shared w1/w3 weight scale, so collapse + # the gate/up columns to the gate scale here. Other backends keep the + # raw scale and split the halves later (see _compute_gemm1_alphas). + if layer.moe_runner_config.is_gated: + if layer.w13_weight_scale_2.dim() == 1: + # Some checkpoints store a shared scale for w1/w3. + w13_weight_scale_2 = layer.w13_weight_scale_2 + else: + if layer.w13_weight_scale_2.shape[1] >= 2 and not torch.allclose( + layer.w13_weight_scale_2[:, 0], + layer.w13_weight_scale_2[:, 1], + ): + logger.warning_once( + "w1_weight_scale_2 must match w3_weight_scale_2. " + "Accuracy may be affected." + ) - # GEMM 1 scale processing - if layer.moe_runner_config.is_gated: - if layer.w13_weight_scale_2.dim() == 1: - # Some checkpoints store a shared scale for w1/w3. - w13_weight_scale_2 = layer.w13_weight_scale_2 + w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0] else: - if layer.w13_weight_scale_2.shape[1] >= 2 and not torch.allclose( - layer.w13_weight_scale_2[:, 0], - layer.w13_weight_scale_2[:, 1], - ): - logger.warning_once( - "w1_weight_scale_2 must match w3_weight_scale_2. " - "Accuracy may be affected." - ) - - w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0] - else: - w13_weight_scale_2 = layer.w13_weight_scale_2[:] + w13_weight_scale_2 = layer.w13_weight_scale_2[:] + copy_or_rebind_param( + layer, + "w13_weight_scale_2", + w13_weight_scale_2.contiguous(), + ) + prepare_moe_nvfp4_layer_for_marlin(layer) + return # Calculate input scales based on strategy if self.enable_flashinfer_cutlass_moe or self.enable_flashinfer_trtllm_moe: w13_input_scale = layer.w13_input_scale.max().to(torch.float32) w2_input_scale = layer.w2_input_scale.max().to(torch.float32) elif self.enable_flashinfer_cutedsl_moe: - # All-expert-one-input-scale is mathematically different from default per-expert-input-scale - # Thus we allow users to switch the flag to do thorough testing - if CUTEDSL_MOE_SCALAR_INPUT_SCALE: - w13_input_scale = ( - layer.w13_input_scale.max() - .to(torch.float32) - .repeat(layer.w13_input_scale.shape[0]) - ) - else: - w13_input_scale = layer.w13_input_scale.max(dim=1).values.to( - torch.float32 - ) - + # CuteDSL standard path uses a single scalar input scale (all experts). + w13_input_scale = ( + layer.w13_input_scale.max() + .to(torch.float32) + .repeat(layer.w13_input_scale.shape[0]) + ) w2_input_scale = layer.w2_input_scale def _slice_scale(w): @@ -1730,12 +2028,21 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): w13_input_scale = layer.w13_input_scale.max(dim=-1).values.to(torch.float32) w2_input_scale = layer.w2_input_scale - # Create shared parameters - copy_or_rebind_param( - layer, - "g1_alphas", - (w13_input_scale * w13_weight_scale_2).to(torch.float32), + if self.quant_config.use_per_token_activation: + # FlashInfer computes activation scales dynamically per token, so + # the static checkpoint activation scale is intentionally neutral. + w13_input_scale = torch.ones_like(w13_input_scale, dtype=torch.float32) + w2_input_scale = torch.ones_like(w2_input_scale, dtype=torch.float32) + + # Create shared parameters. g1_alphas / g1_alphas_up are the gate (w1) + # and up (w3) GEMM1 scales (equal for shared-scale checkpoints). + g1_alphas, g1_alphas_up = _compute_gemm1_alphas( + layer.w13_weight_scale_2, + w13_input_scale, + layer.moe_runner_config.is_gated, ) + copy_or_rebind_param(layer, "g1_alphas", g1_alphas) + copy_or_rebind_param(layer, "g1_alphas_up", g1_alphas_up) copy_or_rebind_param( layer, "g2_alphas", @@ -1803,14 +2110,43 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): # FlashInfer TRTLLM processing - handles both w13 and w2 align_fp4_moe_weights_for_flashinfer_trtllm(layer) + # TRTLLM doesn't read *_blockscale_swizzled; alias to free the + # placeholders from create_weights. + layer.w13_blockscale_swizzled = layer.w13_weight_scale + layer.w2_blockscale_swizzled = layer.w2_weight_scale else: # CUTLASS processing - handle w13 and w2 separately + if self._is_cutedsl_v2_standard and layer.moe_runner_config.is_gated: + # CuteDSL v2 only: interleave the two logical W13 halves in + # 64-row chunks for the fused SwiGLU GEMM1 layout expected by + # CuteDslMoEWrapper. The v1 (deepep) path uses + # grouped_gemm_nt_masked which expects plain contiguous halves. + from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( + interleave_w13_halves, + ) + + layer.w13_weight = Parameter( + interleave_w13_halves( + layer.w13_weight.view(torch.uint8), group_size=64, dim=1 + ).contiguous(), + requires_grad=False, + ) + layer.w13_weight_scale = Parameter( + interleave_w13_halves( + layer.w13_weight_scale, group_size=64, dim=1 + ).contiguous(), + requires_grad=False, + ) + # Process w13 weights w13_blockscale_swizzled = swizzle_blockscale(layer.w13_weight_scale) - copy_or_rebind_param( - layer, "w13_blockscale_swizzled", w13_blockscale_swizzled + alias_or_bind_derived_param( + layer, + "w13_weight_scale", + "w13_blockscale_swizzled", + w13_blockscale_swizzled, ) w13_weight = layer.w13_weight @@ -1847,10 +2183,54 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): # Process w2 weights w2_blockscale_swizzled = swizzle_blockscale(layer.w2_weight_scale) - copy_or_rebind_param( - layer, "w2_blockscale_swizzled", w2_blockscale_swizzled + alias_or_bind_derived_param( + layer, + "w2_weight_scale", + "w2_blockscale_swizzled", + w2_blockscale_swizzled, ) + if self._is_cutedsl_v2_standard: + # CuteDSL v2 only: convert blockscales to MMA layout for + # CuteDslMoEWrapper. The v1 (deepep) path uses the + # swizzled blockscales directly via flashinfer_cutedsl_moe_masked. + from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout + + from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( + _FP4_SF_VEC_SIZE, + ) + + sf_vec_size = _FP4_SF_VEC_SIZE + num_local_experts = layer.w13_weight.shape[0] + w13_m = layer.w13_weight.shape[1] + w13_k = layer.w13_weight.shape[2] * 2 + w2_m = layer.w2_weight.shape[1] + w2_k = layer.w2_weight.shape[2] * 2 + layer.w13_blockscale_mma = Parameter( + convert_sf_to_mma_layout( + layer.w13_blockscale_swizzled.contiguous() + .view(torch.uint8) + .reshape(-1), + m=w13_m, + k=w13_k, + num_groups=num_local_experts, + sf_vec_size=sf_vec_size, + ), + requires_grad=False, + ) + layer.w2_blockscale_mma = Parameter( + convert_sf_to_mma_layout( + layer.w2_blockscale_swizzled.contiguous() + .view(torch.uint8) + .reshape(-1), + m=w2_m, + k=w2_k, + num_groups=num_local_experts, + sf_vec_size=sf_vec_size, + ), + requires_grad=False, + ) + # Both flashinfer cutlass and regular cutlass use same processing for w2 # Set up CUTLASS MoE parameters (reuse to keep CUDA graph stable) @@ -1876,38 +2256,84 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): @property def load_up_proj_weight_first(self) -> bool: - # FlashInfer CUTLASS kernel assumes [Up, Gate] Proj as W13 - return self.enable_flashinfer_cutlass_moe and self.moe_runner_config.is_gated + # Load W13 as [Up, Gate] for FlashInfer CUTLASS and CuteDSL v2 kernels. + # The CuteDSL v1 (deepep) path uses [Gate, Up] -- do NOT flip. + return self.moe_runner_config.is_gated and ( + self.enable_flashinfer_cutlass_moe or self._is_cutedsl_v2_standard + ) def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): self.moe_runner_config = moe_runner_config - if get_moe_runner_backend().is_flashinfer_trtllm(): - self.runner = MoeRunner( - MoeRunnerBackend.FLASHINFER_TRTLLM, moe_runner_config - ) + moe_runner_backend = get_moe_runner_backend() + + if moe_runner_backend.is_auto(): + if is_cuda() and (8, 0) <= get_device_capability() < (10, 0): + moe_runner_backend = MoeRunnerBackend.MARLIN + else: + # TRTLLM is currently the most performant and tested FP4 MoE + # backend, so use it as the default. + moe_runner_backend = MoeRunnerBackend.FLASHINFER_TRTLLM + + self._moe_runner_backend = moe_runner_backend + + if moe_runner_backend.is_flashinfer_cutedsl(): + import sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl # noqa: F401 – triggers @register_fused_func + + if not moe_runner_backend.is_flashinfer_cutlass(): + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) def apply( self, layer: FusedMoE, dispatch_output: StandardDispatchOutput, ) -> CombineInput: + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - x = dispatch_output.hidden_states - x_sf = dispatch_output.hidden_states_scale - topk_output = dispatch_output.topk_output - + # Note: dispatch_output may be a DeepEPLLDispatchOutput (no topk_output + # attribute -- topk_ids/topk_weights live directly on the dispatch + # tuple). Defer per-attribute access to the branches that actually + # consume them. activation = self.moe_runner_config.activation + moe_runner_backend = getattr( + self, "_moe_runner_backend", get_moe_runner_backend() + ) assert ( - activation in ACT_STR_TO_TYPE_MAP - ), f"{activation=} missing from {ACT_STR_TO_TYPE_MAP.keys()=}" + activation in _SUPPORTED_ACT_STRS + ), f"{activation=} not in supported {_SUPPORTED_ACT_STRS}" moe_runner_config = self.moe_runner_config - # FlashInfer TRTLLM FP4 path - layer has shuffled weights only when - # backend is flashinfer_trtllm - if hasattr(layer, "gemm1_weights_fp4_shuffled"): + if moe_runner_backend.is_marlin(): + from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo + + expert_map = None + global_num_experts = -1 + if hasattr(layer, "dispatcher") and hasattr( + layer.dispatcher, "local_expert_mapping" + ): + expert_map = layer.dispatcher.local_expert_mapping + if expert_map is not None: + global_num_experts = self.moe_runner_config.num_experts + + quant_info = MarlinMoeQuantInfo( + w13_qweight=layer.w13_weight, + w2_qweight=layer.w2_weight, + w13_scales=layer.w13_weight_scale, + w2_scales=layer.w2_weight_scale, + w13_g_idx_sort_indices=None, + w2_g_idx_sort_indices=None, + weight_bits=4, + w13_global_scale=layer.w13_weight_scale_2, + w2_global_scale=layer.w2_weight_scale_2, + expert_map=expert_map, + global_num_experts=global_num_experts, + ) + return self.runner.run(dispatch_output, quant_info) + + # FlashInfer TRTLLM FP4 path + if self.enable_flashinfer_trtllm_moe and hasattr(layer, "g1_scale_c"): from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( FlashInferTrtllmFp4MoeQuantInfo, ) @@ -1919,10 +2345,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): ) quant_info = FlashInferTrtllmFp4MoeQuantInfo( - gemm1_weights_fp4_shuffled=layer.gemm1_weights_fp4_shuffled.data, - gemm2_weights_fp4_shuffled=layer.gemm2_weights_fp4_shuffled.data, - gemm1_scales_fp4_shuffled=layer.gemm1_scales_fp4_shuffled.data, - gemm2_scales_fp4_shuffled=layer.gemm2_scales_fp4_shuffled.data, + w13_weight=layer.w13_weight.data, + w2_weight=layer.w2_weight.data, + w13_weight_scale=layer.w13_weight_scale.data, + w2_weight_scale=layer.w2_weight_scale.data, g1_scale_c=layer.g1_scale_c.data, g1_alphas=layer.g1_alphas.data, g2_alphas=layer.g2_alphas.data, @@ -1932,18 +2358,87 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): local_num_experts=layer.num_local_experts, intermediate_size_per_partition=layer.intermediate_size_per_partition, routing_method_type=routing_method_type, + use_per_token_activation=self.quant_config.use_per_token_activation, ) return self.runner.run(dispatch_output, quant_info) + if self.enable_flashinfer_cutedsl_moe: + from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( + CuteDslFp4MoeQuantInfo, + ensure_cutedsl_wrapper, + ) + + if self._is_cutedsl_v1_deepep: + # v1 path: DeepEP low-latency + flashinfer_cutedsl_moe_masked. + # Weights are [Gate, Up] (non-interleaved) with swizzled blockscales. + quant_info = CuteDslFp4MoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_sf=layer.w13_blockscale_swizzled, + w2_weight_sf=layer.w2_blockscale_swizzled, + w1_alpha=layer.g1_alphas, + w2_alpha=layer.g2_alphas, + a1_scale=layer.w13_input_scale_quant, + a2_scale=layer.w2_input_scale_quant, + use_nvfp4_dispatch=MOE_NVFP4_DISPATCH, + down_gemm_overlap_args=getattr( + self.runner, "down_gemm_overlap_args", None + ), + ) + return self.runner.run(dispatch_output, quant_info) + + # v2 standard path (a2a=none/flashinfer): uses CuteDslMoEWrapper + # with [Up, Gate] interleaved weights and MMA blockscales. + ensure_cutedsl_wrapper(layer) + w1_alpha, fc2_input_scale, w2_alpha = layer._cutedsl_scales + quant_info = CuteDslFp4MoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_sf=getattr( + layer, "w13_blockscale_mma", layer.w13_blockscale_swizzled + ), + w2_weight_sf=getattr( + layer, "w2_blockscale_mma", layer.w2_blockscale_swizzled + ), + w1_alpha=w1_alpha, + w2_alpha=w2_alpha, + a1_scale=layer._cutedsl_input_scale, + a2_scale=fc2_input_scale, + wrapper=layer._cutedsl_wrapper, + ) + return self.runner.run(dispatch_output, quant_info) + if self.enable_flashinfer_cutlass_moe: + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + get_activation_type, + ) from sglang.srt.layers.moe.token_dispatcher import DispatchOutputChecker assert ( not moe_runner_config.apply_router_weight_on_input ), "apply_router_weight_on_input is not supported for Flashinfer" + # Resolve the FlashInfer ActivationType honoring the gated flag, + # then verify the CUTLASS FP4 kernel supports it. + fi_activation = ActivationType( + get_activation_type(activation, is_gated=moe_runner_config.is_gated) + ) + _CUTLASS_FP4_SUPPORTED = { + ActivationType.Swiglu, + ActivationType.Geglu, + ActivationType.Relu2, + ActivationType.Identity, + } + assert fi_activation in _CUTLASS_FP4_SUPPORTED, ( + f"Activation {activation!r} (is_gated={moe_runner_config.is_gated}) " + f"maps to {fi_activation.name}, which is not supported by the " + "flashinfer cutlass fp4 moe kernel." + ) # TRTLLM Cutlass moe takes in activations in BF16/Half/nvfp4 precision # and fp4 quantized weights loaded from the checkpoint + x = dispatch_output.hidden_states + x_sf = dispatch_output.hidden_states_scale + topk_output = dispatch_output.topk_output topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids output_dtype = torch.bfloat16 @@ -1975,7 +2470,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): fc2_expert_weights=layer.w2_weight.view(torch.long), output_dtype=output_dtype, input_sf=x_sf, - # swizzled_input_sf=not get_moe_a2a_backend().is_flashinfer(), + # swizzled_input_sf intentionally omitted; not used for this path. quant_scales=[ layer.w13_input_scale_quant, layer.w13_blockscale_swizzled.view(torch.int32), @@ -1989,16 +2484,16 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): tp_size=layer.moe_tp_size, tp_rank=layer.moe_tp_rank, tune_max_num_tokens=next_power_of_2(x.shape[0]), - activation_type=ACT_STR_TO_TYPE_MAP[activation], + activation_type=fi_activation, enable_alltoall=get_moe_a2a_backend().is_flashinfer(), )[0] - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - return StandardCombineInput(hidden_states=output) from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4 + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids output = cutlass_moe_fp4( a=x, @@ -2014,57 +2509,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): topk_ids=topk_ids, params=layer.cutlass_moe_params, apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input, + no_combine=moe_runner_config.no_combine, ).to(x.dtype) # Scale by routed_scaling_factor is fused into select_experts. - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - return StandardCombineInput(hidden_states=output) - - def apply_without_routing_weights( - self, - layer: FusedMoE, - x: tuple[torch.Tensor, Optional[torch.Tensor]], - masked_m: torch.Tensor, - moe_runner_config: MoeRunnerConfig, - ) -> torch.Tensor: - assert ( - moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." - - assert self.enable_flashinfer_cutedsl_moe, "only support flashinfer cutedsl moe" - assert ( - not moe_runner_config.apply_router_weight_on_input - ), "apply_router_weight_on_input is not supported for Flashinfer" - - from sglang.srt.layers.moe.flashinfer_cutedsl_moe import ( - flashinfer_cutedsl_moe_masked, - ) - - down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = getattr( - layer, "down_gemm_overlap_args", None - ) - - out = flashinfer_cutedsl_moe_masked( - hidden_states=x, - input_global_scale=( - None if MOE_NVFP4_DISPATCH else layer.w13_input_scale_quant - ), - w1=layer.w13_weight, - w1_blockscale=layer.w13_blockscale_swizzled, - w1_alpha=layer.g1_alphas, - w2=layer.w2_weight, - a2_global_scale=layer.w2_input_scale_quant, - w2_blockscale=layer.w2_blockscale_swizzled, - w2_alpha=layer.g2_alphas, - masked_m=masked_m, - **( - dict( - down_sm_count=down_gemm_overlap_args.num_sms, - down_signals=down_gemm_overlap_args.signal, - down_start_event=down_gemm_overlap_args.start_event, - ) - if down_gemm_overlap_args is not None - else {} - ), - ) - return out diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py new file mode 100644 index 000000000..7a2e34d47 --- /dev/null +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -0,0 +1,474 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import torch +import triton +import triton.language as tl +from torch.nn import Module +from torch.nn.parameter import Parameter + +from sglang.srt.distributed import get_tp_group +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.layers.dp_attention import is_allocation_symmetric +from sglang.srt.layers.moe.utils import RoutingMethodType +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import ( + is_flashinfer_available, + log_info_on_rank0, + set_weight_attrs, +) +from sglang.srt.utils.common import is_sm100_supported, next_power_of_2 + +_MXFP8_QUANTIZE_BACKEND = "cute-dsl" if is_sm100_supported() else "cuda" + +if is_flashinfer_available(): + from flashinfer import mxfp8_quantize, shuffle_matrix_a, shuffle_matrix_sf_a + from flashinfer.fp4_quantization import block_scale_interleave + from flashinfer.fused_moe import trtllm_fp4_block_scale_routed_moe + from flashinfer.fused_moe.core import ( + _maybe_get_cached_w3_w1_permute_indices, + get_w2_permute_indices_with_cache, + ) + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput + +from sglang.srt.utils.common import get_bool_env_var + +_USE_OFFICIAL_SHUFFLE = get_bool_env_var( + "SGLANG_MXFP4_USE_OFFICIAL_SHUFFLE", default="true" +) + + +class PackTopkIds: + + @classmethod + def execute( + cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor + ) -> torch.Tensor: + return cls.triton(topk_ids, topk_weights) + + @classmethod + def vanilla( + cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor + ) -> torch.Tensor: + weight_bits = ( + topk_weights.to(torch.bfloat16).view(torch.int16).to(torch.int32) & 0xFFFF + ) + return (topk_ids.to(torch.int32) << 16) | weight_bits + + @classmethod + def triton(cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor) -> torch.Tensor: + assert ( + topk_ids.shape == topk_weights.shape + ), f"shape mismatch: {topk_ids.shape=} vs {topk_weights.shape=}" + assert topk_ids.ndim >= 1, f"expected >=1D, got {topk_ids.shape=}" + + assert ( + topk_ids.dtype == torch.int32 + ), f"topk_ids must be int32, got {topk_ids.dtype}" + assert ( + topk_weights.dtype == torch.float32 + ), f"topk_weights must be float32, got {topk_weights.dtype}" + + assert topk_ids.is_contiguous(), "topk_ids must be contiguous" + assert topk_weights.is_contiguous(), "topk_weights must be contiguous" + + out = torch.empty_like(topk_ids, dtype=torch.int32) + numel = out.numel() + if numel == 0: + return out + + BLOCK_SIZE = 1024 + grid = (triton.cdiv(numel, BLOCK_SIZE),) + _pack_topk_ids_triton_kernel[grid]( + topk_ids, + topk_weights, + out, + numel, + BLOCK_SIZE=BLOCK_SIZE, + ) + return out + + +@triton.jit +def _pack_topk_ids_triton_kernel( + topk_ids_ptr, + topk_weights_ptr, + out_ptr, + numel, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + + ids = tl.load(topk_ids_ptr + offsets, mask=mask, other=0) + w = tl.load(topk_weights_ptr + offsets, mask=mask, other=0.0) + + w_bf16 = w.to(tl.bfloat16) + w_i16 = w_bf16.to(tl.int16, bitcast=True) + w_i32 = w_i16.to(tl.int32) & 0xFFFF + + ids_i32 = ids.to(tl.int32) + packed = (ids_i32 << 16) | w_i32 + + tl.store(out_ptr + offsets, packed, mask=mask) + + +class Mxfp4FlashinferTrtllmMoEMethod: + + def __init__(self, fp8_method, prefix: str): + self._fp8 = fp8_method + self.prefix = prefix + self.flashinfer_mxfp4_moe_precision = ( + get_global_server_args().flashinfer_mxfp4_moe_precision + ) + + def create_moe_runner(self, layer, moe_runner_config): + self.moe_runner_config = moe_runner_config + + swiglu_limit = moe_runner_config.swiglu_limit + assert ( + swiglu_limit is not None + ), f"swiglu_limit must be non-None for DeepSeek V4 (got {swiglu_limit!r})" + self._gemm1_clamp_limit_tensor = ( + torch.full( + (layer.num_local_experts,), + swiglu_limit, + dtype=torch.float32, + device=layer.w13_weight.device, + ) + if swiglu_limit is not None + else None + ) + + def create_weights( + self, + layer, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype, + **extra_weight_attrs, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + fp4_block_k = 32 + + w13_weight = Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // 2, + dtype=torch.int8, + ), + requires_grad=False, + ) + w2_weight = Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // 2, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w13_weight_scale = Parameter( + torch.ones( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // fp4_block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + w2_weight_scale = Parameter( + torch.ones( + num_experts, + hidden_size, + intermediate_size_per_partition // fp4_block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + w13_weight_scale.format_ue8m0 = False + w2_weight_scale.format_ue8m0 = False + scale_attrs = dict(extra_weight_attrs) + scale_attrs["quant_method"] = FusedMoeWeightScaleSupported.BLOCK.value + layer.register_parameter("w13_weight_scale_inv", w13_weight_scale) + set_weight_attrs(w13_weight_scale, scale_attrs) + layer.register_parameter("w2_weight_scale_inv", w2_weight_scale) + set_weight_attrs(w2_weight_scale, scale_attrs) + + def process_weights_after_loading(self, layer: Module) -> None: + from sglang.srt.layers.quantization.utils import reorder_w1w3_to_w3w1 + + self._fp8.process_weights_after_loading(layer) + + if getattr(layer, "_mega_moe_weights_built", False): + return + + w13_w, w13_s = reorder_w1w3_to_w3w1( + layer.w13_weight.data, layer.w13_weight_scale_inv.data + ) + layer.w13_weight = Parameter(w13_w, requires_grad=False) + layer.w13_weight_scale_inv = Parameter(w13_s, requires_grad=False) + + log_info_on_rank0( + logger, + f"Shuffling FP4 expert weights for TRT-LLM MxFP4 kernel " + f"(layer: {self.prefix})...", + ) + + w13 = layer.w13_weight.data + w2 = layer.w2_weight.data + w13_scale = layer.w13_weight_scale_inv.data + w2_scale = layer.w2_weight_scale_inv.data + num_experts = w13.shape[0] + + if w13_scale.dtype == torch.float32: + w13_scale = w13_scale.to(torch.float8_e8m0fnu) + w2_scale = w2_scale.to(torch.float8_e8m0fnu) + + epilogue_tile_m = 128 + g1_w, g1_s, g2_w, g2_s = [], [], [], [] + if _USE_OFFICIAL_SHUFFLE: + cache: dict = {} + for i in range(num_experts): + w13_u8 = w13[i].view(torch.uint8) + w13_s_u8 = w13_scale[i].view(torch.uint8) + w2_u8 = w2[i].view(torch.uint8) + w2_s_u8 = w2_scale[i].view(torch.uint8) + + perm = _maybe_get_cached_w3_w1_permute_indices( + cache, + w13_u8, + epilogue_tile_m, + ) + g1_w.append(w13_u8[perm.to(w13_u8.device)].contiguous()) + perm_sf = _maybe_get_cached_w3_w1_permute_indices( + cache, + w13_s_u8, + epilogue_tile_m, + num_elts_per_sf=16, + ) + g1_s.append( + block_scale_interleave( + w13_s_u8[perm_sf.to(w13_s_u8.device)].contiguous() + ) + ) + + perm = get_w2_permute_indices_with_cache( + cache, + w2_u8, + epilogue_tile_m, + ) + g2_w.append(w2_u8[perm.to(w2_u8.device)].contiguous()) + perm_sf = get_w2_permute_indices_with_cache( + cache, + w2_s_u8, + epilogue_tile_m, + num_elts_per_sf=16, + ) + g2_s.append( + block_scale_interleave( + w2_s_u8[perm_sf.to(w2_s_u8.device)].contiguous() + ) + ) + else: + for i in range(num_experts): + g1_w.append(shuffle_matrix_a(w13[i].view(torch.uint8), epilogue_tile_m)) + g1_s.append( + shuffle_matrix_sf_a(w13_scale[i].view(torch.uint8), epilogue_tile_m) + ) + g2_w.append(shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m)) + g2_s.append( + shuffle_matrix_sf_a(w2_scale[i].view(torch.uint8), epilogue_tile_m) + ) + + layer.w13_weight = Parameter(torch.stack(g1_w), requires_grad=False) + layer.w13_weight_scale_inv = Parameter( + torch.stack(g1_s) + .view(torch.float8_e4m3fn) + .reshape(num_experts, w13.shape[1], -1), + requires_grad=False, + ) + layer.w2_weight = Parameter(torch.stack(g2_w), requires_grad=False) + layer.w2_weight_scale_inv = Parameter( + torch.stack(g2_s) + .view(torch.float8_e4m3fn) + .reshape(num_experts, w2.shape[1], -1), + requires_grad=False, + ) + + self._register_static_scale_ones(layer) + torch.cuda.empty_cache() + + def _register_static_scale_ones(self, layer: Module) -> None: + device = layer.w13_weight.device + for name in ( + "output1_scale_scalar", + "output1_scale_gate_scalar", + "output2_scale_scalar", + ): + layer.register_buffer( + name, + torch.ones(layer.num_local_experts, device=device, dtype=torch.float32), + persistent=False, + ) + + def apply( + self, + layer: Module, + dispatch_output: DispatchOutput, + ) -> CombineInput: + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + from sglang.srt.layers.moe.topk import TopKOutputChecker + + hidden_states = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + w13 = layer.w13_weight + w2 = layer.w2_weight + w13_scale = layer.w13_weight_scale_inv + w2_scale = layer.w2_weight_scale_inv + + intermediate_size = w2.shape[2] * 2 if w2.dtype == torch.uint8 else w2.shape[2] + hidden_size = w13.shape[2] * 2 if w13.dtype == torch.uint8 else w13.shape[2] + + num_local_experts = layer.num_local_experts + if w13_scale.dim() == 2: + w13_scale = w13_scale.reshape(num_local_experts, 2 * intermediate_size, -1) + if w2_scale.dim() == 2: + w2_scale = w2_scale.reshape(num_local_experts, hidden_size, -1) + + if TopKOutputChecker.format_is_standard(topk_output): + topk_ids = topk_output.topk_ids + topk_weights = topk_output.topk_weights + elif TopKOutputChecker.format_is_bypassed(topk_output): + raise NotImplementedError( + "the old code in this branch is WRONG. e.g. it does not consider HashTopK, and may miss args" + ) + else: + raise ValueError(f"Unsupported topk output format: {topk_output.format}") + + packed_topk = PackTopkIds.execute(topk_ids, topk_weights) + + precision = self.flashinfer_mxfp4_moe_precision + if precision == "bf16": + assert hidden_states.dtype == torch.bfloat16 + x_quant = hidden_states + x_scale = None + origin_dim = x_quant.shape[-1] + if hidden_size != origin_dim: + x_quant = torch.nn.functional.pad( + x_quant, + (0, hidden_size - origin_dim), + mode="constant", + value=0.0, + ) + elif precision == "default": + x_quant, x_scale = mxfp8_quantize( + hidden_states, + False, + alignment=hidden_size, + backend=_MXFP8_QUANTIZE_BACKEND, + ) + x_scale = x_scale.view(torch.float8_e4m3fn).reshape( + *hidden_states.shape[:-1], -1 + ) + else: + raise NotImplementedError(f"Unsupported mxfp4 moe precision: {precision}") + + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + num_tokens = x_quant.shape[0] + out_hidden_size = ( + x_quant.shape[-1] * 2 + if x_quant.dtype == torch.uint8 + else x_quant.shape[-1] + ) + symm_output = torch.empty( + num_tokens, out_hidden_size, dtype=torch.bfloat16, device=x_quant.device + ) + + output = trtllm_fp4_block_scale_routed_moe( + topk_ids=packed_topk, + routing_bias=None, + hidden_states=x_quant, + hidden_states_scale=x_scale, + gemm1_weights=w13, + gemm1_weights_scale=w13_scale, + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=self._gemm1_clamp_limit_tensor, + gemm2_weights=w2, + gemm2_weights_scale=w2_scale, + gemm2_bias=None, + output1_scale_scalar=layer.output1_scale_scalar, + output1_scale_gate_scalar=layer.output1_scale_gate_scalar, + output2_scale_scalar=layer.output2_scale_scalar, + num_experts=layer.num_experts, + top_k=packed_topk.shape[1], + n_group=1, + topk_group=1, + intermediate_size=intermediate_size, + local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, + local_num_experts=num_local_experts, + routed_scaling_factor=1.0, + routing_method_type=int(RoutingMethodType.TopK), + do_finalize=True, + tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), + output=symm_output, + )[0] + + return StandardCombineInput(hidden_states=output) + + +def maybe_fuse_routed_scale_and_shared_add( + experts, + routed: torch.Tensor, + shared: torch.Tensor | None, + routed_scaling_factor: float, +) -> torch.Tensor: + # When MxFP4 fusion is on, the upstream `routed *= scale` is skipped and + # the scaling is folded into the shared-add via `shared.add_(routed, + # alpha=scale)`. With no shared output, the missing scale is applied + # in-place. Otherwise `routed` is already scale-final and we just add + # `shared` (or pass through if there is none). + from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import ( + Mxfp4FlashinferCutlassMoEMethod, + ) + from sglang.srt.layers.quantization.mxfp4_marlin_moe import ( + Mxfp4MarlinMoEMethod, + ) + + fused = isinstance( + experts.quant_method, + ( + Mxfp4FlashinferTrtllmMoEMethod, + Mxfp4FlashinferCutlassMoEMethod, + Mxfp4MarlinMoEMethod, + ), + ) + if fused: + if shared is not None: + return shared.add_(routed, alpha=routed_scaling_factor) + return routed.mul_(routed_scaling_factor) + if shared is not None: + routed += shared + return routed diff --git a/python/sglang/srt/layers/quantization/nvfp4_online.py b/python/sglang/srt/layers/quantization/nvfp4_online.py new file mode 100644 index 000000000..89103ab0f --- /dev/null +++ b/python/sglang/srt/layers/quantization/nvfp4_online.py @@ -0,0 +1,583 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import logging +import re +import threading +from typing import Any, Dict, List, Optional + +import torch + +from sglang.srt.environ import envs +from sglang.srt.layers.quantization.fp8_utils import ( + block_quant_dequant, + inverse_transform_scale_ue8m0, +) +from sglang.srt.layers.quantization.modelopt_quant import ( + ModelOptNvFp4FusedMoEMethod, + ModelOptQuantConfig, +) +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.layers.quantization.utils import ( + is_layer_skipped, + per_tensor_dequantize, +) + +logger = logging.getLogger(__name__) + + +class NvFp4OnlineConfig(ModelOptQuantConfig): + """Config for `--quantization nvfp4_online`. + + This mode is a load-time conversion path, not a serialized NVFP4 checkpoint + format. It reuses the ModelOpt NVFP4 MoE parameter layout and fills those + parameters by converting BF16/FP16/FP8 expert tensors as they are loaded. + Dense layers stay in the source checkpoint precision or quantization path. + """ + + # Marker consumed by the ModelOpt FP4 layout and the model loader. Serialized + # NVFP4 checkpoints use ModelOptFp4Config instead. + is_nvfp4_online = True + is_checkpoint_nvfp4_serialized = False + group_size = 16 + + @staticmethod + def _normalize_ignored_layers( + ignored_layers: Optional[List[str]], + ) -> List[str]: + if not ignored_layers: + return [] + normalized_ignored_layers = [] + for layer in ignored_layers: + base = layer.removeprefix("model.") + normalized_ignored_layers.append(base) + normalized_ignored_layers.append(f"model.{base}") + return list(dict.fromkeys(normalized_ignored_layers)) + + def __init__( + self, + exclude_modules: Optional[List[str]] = None, + packed_modules_mapping: Optional[Dict[str, List[str]]] = None, + is_checkpoint_fp8_serialized: bool = False, + activation_scheme: str = "dynamic", + weight_block_size: Optional[List[int]] = None, + use_mxfp8: bool = False, + ) -> None: + source_ignored_layers = self._normalize_ignored_layers(exclude_modules) + fp4_ignored_layers = list(source_ignored_layers) + if ignored_layers_str := envs.SGLANG_FP4_IGNORED_LAYERS.get(): + fp4_ignored_layers.extend( + layer.strip() + for layer in ignored_layers_str.split(",") + if layer.strip() + ) + fp4_ignored_layers = self._normalize_ignored_layers(fp4_ignored_layers) + super().__init__( + kv_cache_quant_algo=None, + exclude_modules=source_ignored_layers, + packed_modules_mapping=packed_modules_mapping or {}, + ) + self.fp4_ignored_layers = fp4_ignored_layers + # Weights use static NVFP4 scales, while FlashInfer computes activation + # FP32 scales dynamically per token at runtime. + self.use_per_token_activation = True + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + self.is_fp4_experts = False + self.activation_scheme = activation_scheme + self.weight_block_size = weight_block_size + self.use_mxfp8 = use_mxfp8 + + @classmethod + def get_name(cls) -> str: + return "nvfp4_online" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 100 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> NvFp4OnlineConfig: + quant_method = str(config.get("quant_method", "")).lower() + use_mxfp8 = "mxfp8" in quant_method + is_checkpoint_fp8_serialized = "fp8" in quant_method or use_mxfp8 + ignored_layers = config.get("ignored_layers") or config.get( + "modules_to_not_convert" + ) + if isinstance(ignored_layers, str): + ignored_layers = [ignored_layers] + return cls( + exclude_modules=ignored_layers, + packed_modules_mapping=config.get("packed_modules_mapping"), + is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, + activation_scheme=config.get("activation_scheme", "dynamic"), + weight_block_size=config.get("weight_block_size"), + use_mxfp8=use_mxfp8, + ) + + def get_quant_method(self, layer: torch.nn.Module, prefix: str): + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod, Fp8MoEMethod + + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix, self.exclude_modules, self.packed_modules_mapping + ) or self.is_layer_excluded(prefix): + return UnquantizedLinearMethod() + if self.is_checkpoint_fp8_serialized: + return Fp8LinearMethod(self) + return UnquantizedLinearMethod() + if isinstance(layer, FusedMoE): + if is_layer_skipped( + prefix, self.exclude_modules, self.packed_modules_mapping + ) or self.is_layer_excluded(prefix): + return None + if is_layer_skipped( + prefix, self.fp4_ignored_layers, self.packed_modules_mapping + ): + if self.is_checkpoint_fp8_serialized: + return Fp8MoEMethod(self) + return None + return ModelOptNvFp4OnlineFusedMoEMethod(self, prefix) + return None + + +class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): + """MoE method that converts source expert weights to NVFP4 during loading.""" + + def __init__(self, quant_config: NvFp4OnlineConfig, layer_prefix: str): + super().__init__(quant_config) + self.layer_prefix = layer_prefix + layer_match = re.search(r"(?:^|\.)layers\.(\d+)(?:\.|$)", layer_prefix) + self.layer_log_name = ( + f"layer {layer_match.group(1)} ({layer_prefix})" + if layer_match is not None + else layer_prefix + ) + if not self.enable_flashinfer_trtllm_moe: + raise ValueError( + "--quantization nvfp4_online supports only " + "--moe-runner-backend flashinfer_trtllm or " + "flashinfer_trtllm_routed." + ) + + @staticmethod + def _quantize_weight_nvfp4( + weight: torch.Tensor, + weight_scale_2: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return packed NVFP4 weight, block scales, and per-tensor decode scale. + + The weight scale is static and per tensor. Callers pass an existing + scale when multiple shards must share one global scale, for example the + gated w1/w3 pair. + """ + from flashinfer import SfLayout, nvfp4_quantize + + if weight.ndim != 2: + raise ValueError( + "--quantization nvfp4_online expects 2D expert weights, " + f"got shape {tuple(weight.shape)}." + ) + if weight.shape[-1] % 16 != 0: + raise ValueError( + "--quantization nvfp4_online requires expert weight K to be " + f"a multiple of 16, got shape {tuple(weight.shape)}." + ) + + if weight_scale_2 is None: + # weight_scale_2 is the NVFP4 decode scale. FlashInfer consumes its + # reciprocal as the global encode scale, matching 448 * 6 / amax. + weight_amax = ( + weight.abs() + .nan_to_num() + .amax() + .to(device=weight.device, dtype=torch.float32) + ) + e4m3_max = ( + 256.0 + if envs.FLASHINFER_NVFP4_4OVER6.get() + and envs.FLASHINFER_NVFP4_4OVER6_E4M3_USE_256.get() + else float(torch.finfo(torch.float8_e4m3fn).max) + ) + fp8_fp4_max = e4m3_max * 6.0 + weight_scale_2 = torch.where( + weight_amax > 0, + weight_amax / fp8_fp4_max, + torch.ones_like(weight_amax), + ) + else: + weight_scale_2 = weight_scale_2.to( + device=weight.device, dtype=torch.float32 + ) + fp4_weight, weight_sf = nvfp4_quantize( + weight.contiguous(), + 1.0 / weight_scale_2, + sfLayout=SfLayout.layout_linear, + backend="cuda", + ) + rows, cols = weight.shape + weight_sf = weight_sf.view(torch.float8_e4m3fn).reshape(rows, cols // 16) + return ( + fp4_weight.reshape(rows, cols // 2), + weight_sf.contiguous(), + weight_scale_2, + ) + + @staticmethod + def _is_fp8_weight(weight: torch.Tensor) -> bool: + fp8_dtypes = { + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None), + ) + if dtype is not None + } + return weight.dtype in fp8_dtypes + + @staticmethod + def _is_fp8_weight_scale_name(weight_name: str) -> bool: + return "weight_scale" in weight_name and "weight_scale_2" not in weight_name + + def _dequantize_fp8_weight( + self, + weight: torch.Tensor, + weight_scale: torch.Tensor, + device: torch.device, + ) -> torch.Tensor: + if self.quant_config.use_mxfp8: + raise ValueError( + "--quantization nvfp4_online does not support online " + "requantization from MXFP8 expert checkpoints." + ) + + weight = weight.to(device).contiguous() + weight_scale = weight_scale.to(device=device).contiguous() + if weight_scale.dtype == torch.int32: + weight_scale = inverse_transform_scale_ue8m0( + weight_scale, mn=weight.shape[-2] + ) + weight_scale = weight_scale.to(dtype=torch.float32).contiguous() + + if weight_scale.numel() == 1 or self.quant_config.weight_block_size is None: + return ( + per_tensor_dequantize(weight, weight_scale) + .to(torch.bfloat16) + .contiguous() + ) + + return block_quant_dequant( + weight, + weight_scale, + self.quant_config.weight_block_size, + torch.bfloat16, + ).contiguous() + + @staticmethod + def _should_skip_loaded_expert( + layer: torch.nn.Module, + param: torch.nn.Parameter, + expert_id: Optional[int], + ) -> bool: + if expert_id is None: + return False + if getattr(param, "_sglang_require_global_experts", False): + return False + # With EPLB or explicit expert placement, logical expert IDs can map to + # one or more physical experts. Let the canonical MoE loader do that + # mapping instead of pre-skipping from the trivial EP layout. + from sglang.srt.eplb.expert_location import get_global_expert_location_metadata + + if get_global_expert_location_metadata() is not None: + return False + return layer._map_global_expert_id_to_local_expert_id(expert_id) == -1 + + @staticmethod + def _scale_weight_name(weight_name: str) -> str: + if "weight" in weight_name: + prefix, suffix = weight_name.rsplit("weight", 1) + return f"{prefix}weight_scale{suffix}" + return f"{weight_name}.weight_scale" + + @staticmethod + def _scale_2_weight_name(weight_name: str) -> str: + if "weight" in weight_name: + prefix, suffix = weight_name.rsplit("weight", 1) + return f"{prefix}weight_scale_2{suffix}" + return f"{weight_name}.weight_scale_2" + + def get_online_weight_loader(self, layer, original_weight_loader): + """Wrap the normal MoE loader with load-time NVFP4 conversion. + + The wrapper quantizes each eligible expert shard as soon as the loader + sees enough source data, which avoids materializing and then converting + the full checkpoint. FP8 checkpoints stream weight and scale tensors + separately, so those pairs are staged until both sides have arrived. + """ + pending_w13_weights = {} + pending_w13_lock = threading.Lock() + pending_fp8_weights = {} + pending_fp8_weight_scales = {} + pending_fp8_lock = threading.Lock() + quantization_log_lock = threading.Lock() + did_log_quantization = False + + def log_quantization_start() -> None: + nonlocal did_log_quantization + if did_log_quantization: + return + with quantization_log_lock: + if did_log_quantization: + return + logger.info( + "Running online NVFP4 quantization for MoE expert weights in %s.", + self.layer_log_name, + ) + did_log_quantization = True + + def store_quantized_weight( + param: torch.nn.Parameter, + fp4_weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_scale_2: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: Optional[int], + ) -> None: + original_weight_loader( + param, + fp4_weight, + weight_name=weight_name, + shard_id=shard_id, + expert_id=expert_id, + ) + + scale_param = ( + layer.w13_weight_scale + if shard_id in ("w1", "w3") + else layer.w2_weight_scale + ) + original_weight_loader( + scale_param, + weight_scale, + weight_name=self._scale_weight_name(weight_name), + shard_id=shard_id, + expert_id=expert_id, + ) + scale_2_param = ( + layer.w13_weight_scale_2 + if shard_id in ("w1", "w3") + else layer.w2_weight_scale_2 + ) + original_weight_loader( + scale_2_param, + weight_scale_2, + weight_name=self._scale_2_weight_name(weight_name), + shard_id=shard_id, + expert_id=expert_id, + ) + + def process_loaded_weight( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: Optional[int], + ) -> None: + log_quantization_start() + if shard_id == "w2": + loaded_weight = loaded_weight.to(param.device) + fp4_weight, weight_scale, weight_scale_2 = self._quantize_weight_nvfp4( + loaded_weight + ) + store_quantized_weight( + param, + fp4_weight, + weight_scale, + weight_scale_2, + weight_name, + shard_id, + expert_id, + ) + return + + pending_key = expert_id + current = ( + param, + loaded_weight, + weight_name, + shard_id, + expert_id, + ) + with pending_w13_lock: + pending = pending_w13_weights.pop(pending_key, None) + if pending is None: + pending_w13_weights[pending_key] = current + return + + ( + pending_param, + pending_weight, + pending_name, + pending_shard_id, + pending_eid, + ) = pending + if pending_shard_id == shard_id: + raise ValueError( + "--quantization nvfp4_online expects paired w1/w3 expert " + f"weights, got two {shard_id} tensors for expert {expert_id}." + ) + pending_weight = pending_weight.to(param.device) + loaded_weight = loaded_weight.to(param.device) + pending_rows = pending_weight.shape[0] + loaded_rows = loaded_weight.shape[0] + # Quantize the gated pair together so w1/w3 share one amax-derived + # per-tensor FP32 scale, matching the serialized NVFP4 convention. + fp4_weight, weight_scale, weight_scale_2 = self._quantize_weight_nvfp4( + torch.cat([pending_weight, loaded_weight], dim=0) + ) + pending_fp4_weight, loaded_fp4_weight = fp4_weight.split( + [pending_rows, loaded_rows], dim=0 + ) + pending_weight_scale, loaded_weight_scale = weight_scale.split( + [pending_rows, loaded_rows], dim=0 + ) + store_quantized_weight( + pending_param, + pending_fp4_weight.contiguous(), + pending_weight_scale.contiguous(), + weight_scale_2, + pending_name, + pending_shard_id, + pending_eid, + ) + store_quantized_weight( + param, + loaded_fp4_weight.contiguous(), + loaded_weight_scale.contiguous(), + weight_scale_2, + weight_name, + shard_id, + expert_id, + ) + + def process_fp8_weight( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: Optional[int], + ) -> None: + if not self._is_fp8_weight(loaded_weight): + process_loaded_weight( + param, loaded_weight, weight_name, shard_id, expert_id + ) + return + if not self.quant_config.is_checkpoint_fp8_serialized: + raise ValueError( + "--quantization nvfp4_online received an FP8 expert " + "weight, but the checkpoint quantization config does not " + "declare serialized FP8 weights." + ) + + key = (expert_id, shard_id) + with pending_fp8_lock: + weight_scale = pending_fp8_weight_scales.pop(key, None) + if weight_scale is None: + pending_fp8_weights[key] = ( + param, + loaded_weight, + weight_name, + shard_id, + expert_id, + ) + return + + log_quantization_start() + loaded_weight = self._dequantize_fp8_weight( + loaded_weight, weight_scale, param.device + ) + process_loaded_weight( + param, loaded_weight, weight_name, shard_id, expert_id + ) + + def process_fp8_weight_scale( + loaded_weight: torch.Tensor, + shard_id: str, + expert_id: Optional[int], + ) -> None: + key = (expert_id, shard_id) + with pending_fp8_lock: + pending = pending_fp8_weights.pop(key, None) + if pending is None: + pending_fp8_weight_scales[key] = loaded_weight + return + + log_quantization_start() + ( + pending_param, + pending_weight, + pending_name, + pending_shard_id, + pending_eid, + ) = pending + loaded_weight = self._dequantize_fp8_weight( + pending_weight, loaded_weight, pending_param.device + ) + process_loaded_weight( + pending_param, + loaded_weight, + pending_name, + pending_shard_id, + pending_eid, + ) + + def nvfp4_online_weight_loader( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: Optional[int], + ): + if shard_id not in ("w1", "w2", "w3"): + original_weight_loader( + param, + loaded_weight, + weight_name=weight_name, + shard_id=shard_id, + expert_id=expert_id, + ) + return + if self._should_skip_loaded_expert(layer, param, expert_id): + return + + if self._is_fp8_weight_scale_name(weight_name): + process_fp8_weight_scale(loaded_weight, shard_id, expert_id) + return + + if "weight" in weight_name: + process_fp8_weight( + param, loaded_weight, weight_name, shard_id, expert_id + ) + return + + original_weight_loader( + param, + loaded_weight, + weight_name=weight_name, + shard_id=shard_id, + expert_id=expert_id, + ) + + return nvfp4_online_weight_loader diff --git a/python/sglang/srt/layers/utils/common.py b/python/sglang/srt/layers/utils/common.py index b982b247e..95fd9e2b8 100644 --- a/python/sglang/srt/layers/utils/common.py +++ b/python/sglang/srt/layers/utils/common.py @@ -54,6 +54,41 @@ def copy_or_rebind_param( setattr(module, name, Parameter(new_value, requires_grad=False)) +def alias_or_bind_derived_param( + module: torch.nn.Module, + source_name: str, + derived_name: str, + derived_value: torch.Tensor, +) -> None: + """Bind a post-processed (derived) tensor to a derived attribute name. + + When `derived_value` is broadcastable to the source Parameter's shape (and + dtype matches), write it broadcast-filled into the source's storage in + place and register `derived_name` as an alias of the source Parameter. The + two attribute names then share one underlying buffer, so: + - apply() can read via `derived_name` + - update_weights_from_disk can keep refilling `source_name` (the loader + re-runs process_weights_after_loading which re-derives in place) + - peak GPU memory is the source size, not source + derived. + + When the shapes are not broadcast-compatible, fall back to allocating a + separate Parameter under `derived_name` via copy_or_rebind_param. + """ + derived_value = derived_value.detach() + source = getattr(module, source_name, None) + if isinstance(source, Parameter) and source.data.dtype == derived_value.dtype: + try: + broadcast = torch.broadcast_to(derived_value, source.data.shape) + except RuntimeError: + broadcast = None + if broadcast is not None: + source.data.copy_(broadcast) + source.requires_grad_(False) + setattr(module, derived_name, source) + return + copy_or_rebind_param(module, derived_name, derived_value) + + class PPMissingLayer(torch.nn.Identity): # Adapted from # https://github.com/vllm-project/vllm/blob/18ed3132d2bfe1df9a74729457b69243955221e8/vllm/model_executor/models/utils.py#L468C1-L486C1