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

Support fp4 fp8 non gated moe (#13794)

Browse Source 
Co-authored-by: Roi Koren <roik@nvidia.com>
Co-authored-by: Tomer Natan <tbarnatan@computelab-frontend-8.nvidia.com>
This commit is contained in:
TomerBN-Nvidia
2025-12-02 01:26:28 +02:00
committed by GitHub
parent eb5008846a
commit 02af51e4fc
3 changed files with 199 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View File

@@ -281,7 +281,10 @@ class FusedMoE(torch.nn.Module):
# We have to keep the weight scales of w1 and w3 because
# we need to re-quantize w1/w3 weights after weight loading.
idx = 0 if shard_id == "w1" else 1
param_data[expert_id][idx] = loaded_weight
if self.moe_runner_config.is_gated:
param_data[expert_id][idx] = loaded_weight
else:
param_data[expert_id] = loaded_weight
# If we are in the row parallel case (down_proj)
elif shard_id == "w2":
param_data[expert_id] = loaded_weight
@@ -345,7 +348,6 @@ class FusedMoE(torch.nn.Module):
tp_rank: int,
is_bias: bool = False,
):
# Index the loaded weight for tp sharding.
# gate_up_proj: "MergedColumnParallel", so tp sharding on output_dim
assert shard_id in {"w1", "w3", "w13"}
@@ -481,7 +483,6 @@ class FusedMoE(torch.nn.Module):
loaded_weight: torch.Tensor,
tp_rank: int,
):
if shard_id == "w2":
self._load_w2(
shard_id=shard_id,
@@ -518,7 +519,6 @@ class FusedMoE(torch.nn.Module):
shard_id: str,
expert_id: Optional[int],
) -> None:
# if expert_id is None, then
# all the experts are loaded at the same time
if (
@@ -612,7 +612,6 @@ class FusedMoE(torch.nn.Module):
shard_id: str,
expert_id: int,
) -> None:
tp_rank = self.moe_tp_rank
# compressed-tensors checkpoints with packed weights are stored flipped
@@ -925,7 +924,6 @@ class FusedMoE(torch.nn.Module):
ckpt_up_proj_name: str,
num_experts: int,
) -> List[Tuple[str, str, int, str]]:
return [
# (param_name, weight_name, expert_id, shard_id)
(

202
python/sglang/srt/layers/quantization/modelopt_quant.py
View File

@@ -2,6 +2,7 @@
from __future__ import annotations
import logging
from enum import IntEnum
from typing import TYPE_CHECKING, Any, Dict, List, Optional
import torch
@@ -85,9 +86,16 @@ except ImportError:
try:
from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe
from flashinfer.fused_moe.core import ActivationType
except ImportError:
flashinfer_cutlass_fused_moe = None
# Define a minimal ActivationType enum if flashinfer is not available
class ActivationType(IntEnum):
Swiglu = 3
Relu2 = 6
# Initialize logger for the module
logger = logging.getLogger(__name__)
@@ -145,6 +153,11 @@ FLASHINFER_FP4_GEMM_BACKEND = envs.SGLANG_FLASHINFER_FP4_GEMM_BACKEND.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,
}
class ModelOptQuantConfig(QuantizationConfig):
def __init__(
@@ -443,11 +456,12 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
else params_dtype
)
weight_loader = extra_weight_attrs.get("weight_loader")
num_shards = 2 if layer.moe_runner_config.is_gated else 1
intermediate_size = num_shards * intermediate_size_per_partition
w13_weight = ModelWeightParameter(
data=torch.empty(
num_experts,
2 * intermediate_size_per_partition,
intermediate_size,
hidden_size,
dtype=weight_dtype,
),
@@ -474,9 +488,10 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
# WEIGHT SCALES - Per-tensor scaling for ModelOpts
# Allocate 2 scales for w1 and w3 respectively.
# They will be combined to a single scale after weight loading.
w13_scale_shape = (num_experts, num_shards)
w13_weight_scale = PerTensorScaleParameter(
data=torch.full(
(num_experts, 2),
w13_scale_shape,
torch.finfo(torch.float32).min,
dtype=torch.float32,
),
@@ -528,10 +543,13 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
# Requantize each expert's weights using the combined scale
# w13_weight has shape (num_experts, 2 * intermediate_size_per_partition, hidden_size)
# where the first intermediate_size_per_partition rows are w1, the next are w3
intermediate_size_per_partition = layer.w13_weight.shape[1] // 2
num_shards = 2 if layer.moe_runner_config.is_gated else 1
intermediate_size_per_partition = (
layer.w13_weight.shape[1] // num_shards
)
for expert_id in range(layer.w13_weight.shape[0]):
start = 0
for shard_id in range(2): # w1 and w3
for shard_id in range(num_shards): # (w1 and w3) or w13
# Dequantize using the original scale for this shard
dq_weight = per_tensor_dequantize(
layer.w13_weight[expert_id][
@@ -646,6 +664,34 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
layer.output2_scales_scalar = Parameter(
output2_scales_scalar, requires_grad=False
)
elif get_moe_runner_backend().is_flashinfer_cutlass():
assert (
hasattr(layer, "w13_input_scale") and layer.w13_input_scale is not None
)
assert hasattr(layer, "w2_input_scale") and layer.w2_input_scale is not None
assert (
hasattr(layer, "w13_weight_scale")
and layer.w13_weight_scale is not None
)
assert (
hasattr(layer, "w2_weight_scale") and layer.w2_weight_scale is not None
)
input_scale = layer.w13_input_scale.to(torch.float32)
activation_scale = layer.w2_input_scale.to(torch.float32)
w13_weight_scale = layer.w13_weight_scale.to(torch.float32)
w2_weight_scale = layer.w2_weight_scale.to(torch.float32)
layer.fc1_dequant = Parameter(
w13_weight_scale * input_scale, requires_grad=False
)
layer.fc2_quant = Parameter(
activation_scale.reciprocal(), requires_grad=False
)
layer.fc2_dequant = Parameter(
activation_scale * w2_weight_scale, requires_grad=False
)
layer.fc1_input_dequant = Parameter(input_scale, requires_grad=False)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
@@ -744,6 +790,55 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
return StandardCombineInput(hidden_states=output)
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
)
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"
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
original_col = x.shape[1]
x_sf = None
with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric()
):
symm_output = torch.empty(
x.shape[0], original_col, dtype=output_dtype, device=x.device
)
output = flashinfer_cutlass_fused_moe(
output=symm_output,
input=x_fp8,
token_selected_experts=topk_ids.to(torch.int),
token_final_scales=topk_weights,
fc1_expert_weights=layer.w13_weight,
fc2_expert_weights=layer.w2_weight,
output_dtype=output_dtype,
input_sf=x_sf,
quant_scales=[
layer.fc1_dequant,
layer.fc2_quant,
layer.fc2_dequant,
layer.fc1_input_dequant,
],
ep_size=layer.moe_ep_size,
ep_rank=layer.moe_ep_rank,
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=activation,
)[0]
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
return StandardCombineInput(hidden_states=output)
quant_info = TritonMoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
@@ -1192,10 +1287,12 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
weight_scale_dtype = torch.float8_e4m3fn
weight_loader = extra_weight_attrs.get("weight_loader")
# GEMM 1
num_shards = 2 if layer.moe_runner_config.is_gated else 1
w13_weight = ModelWeightParameter(
data=torch.empty(
layer.num_local_experts,
2 * intermediate_size_per_partition,
num_shards * intermediate_size_per_partition,
# 2 fp4 items are packed in the input dimension
hidden_size // 2,
dtype=weight_dtype,
@@ -1224,7 +1321,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
w13_weight_scale = ModelWeightParameter(
data=torch.empty(
layer.num_local_experts,
2 * intermediate_size_per_partition,
num_shards * intermediate_size_per_partition,
hidden_size // self.quant_config.group_size,
dtype=weight_scale_dtype,
),
@@ -1262,8 +1359,13 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
{"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
)
w13_weight_scale_shape = (
(layer.num_local_experts, 2)
if layer.moe_runner_config.is_gated
else (layer.num_local_experts,)
)
w13_weight_scale_2 = PerTensorScaleParameter(
data=torch.empty(layer.num_local_experts, 2, dtype=torch.float32),
data=torch.empty(w13_weight_scale_shape, dtype=torch.float32),
weight_loader=weight_loader,
)
layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2)
@@ -1278,8 +1380,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
)
w13_input_scale_shape = (layer.num_experts, num_shards)
w13_input_scale = PerTensorScaleParameter(
data=torch.empty(layer.num_experts, 2, dtype=torch.float32),
data=torch.empty(w13_input_scale_shape, dtype=torch.float32),
weight_loader=weight_loader,
)
w13_input_scale._sglang_require_global_experts = True
@@ -1448,15 +1551,18 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
"""
# GEMM 1 scale processing
if 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."
)
if layer.moe_runner_config.is_gated:
if 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]
w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0]
else:
w13_weight_scale_2 = layer.w13_weight_scale_2[:]
layer.w13_weight_scale_2 = Parameter(w13_weight_scale_2, requires_grad=False)
# Calculate input scales based on strategy
@@ -1495,7 +1601,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
assert torch.all(w13_input_scale == w13_input_scale[0])
w13_input_scale = w13_input_scale[0]
else:
w13_input_scale = layer.w13_input_scale.max(dim=1).values.to(torch.float32)
w13_input_scale = layer.w13_input_scale.max(dim=-1).values.to(torch.float32)
w2_input_scale = layer.w2_input_scale
# Create shared parameters
@@ -1521,15 +1627,15 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
)
}
)
# Validate weight scales
assert_dim = 2 if layer.moe_runner_config.is_gated else 1
for name, weight_scale in [
("w13", layer.w13_weight_scale),
("w2", layer.w2_weight_scale),
]:
assert (
weight_scale.shape[2] % 16 == 0
), f"Expected {name}_weight_scale.dim(2) to be divisible by 16"
weight_scale.shape[assert_dim] % 16 == 0
), f"Expected {name}_weight_scale.dim({assert_dim}) to be divisible by 16"
assert (
weight_scale.dtype == torch.float8_e4m3fn
), f"{name} Weight Blockscale must be represented as FP8-E4M3"
@@ -1591,13 +1697,45 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
w13_blockscale_swizzled = self.swizzle_blockscale(layer.w13_weight_scale)
del layer.w13_weight_scale
layer.w13_blockscale_swizzled.data.copy_(w13_blockscale_swizzled)
w13_weight = layer.w13_weight
intermediate_size_pad = w13_blockscale_swizzled.size(1) - w13_weight.size(1)
if intermediate_size_pad:
# padding gated activations will require to split w1 and w3
# and pad them individually
assert not layer.moe_runner_config.is_gated, (
"The intermediate size required padding, "
"but padding is also implemented for gated activations"
)
layer.w13_weight = Parameter(
torch.nn.functional.pad(
w13_weight, (0, 0, 0, intermediate_size_pad)
),
requires_grad=False,
)
layer.w2_weight = Parameter(
torch.nn.functional.pad(
layer.w2_weight, (0, intermediate_size_pad // 2, 0, 0)
),
requires_grad=False,
)
layer.w2_weight_scale = Parameter(
torch.nn.functional.pad(
layer.w2_weight_scale, (0, intermediate_size_pad // 16)
),
requires_grad=False,
)
layer.w2_blockscale_swizzled = Parameter(
self.swizzle_blockscale(layer.w2_weight_scale), requires_grad=False
)
layer.w13_weight = Parameter(layer.w13_weight.data, requires_grad=False)
# Process w2 weights
w2_blockscale_swizzled = self.swizzle_blockscale(layer.w2_weight_scale)
del layer.w2_weight_scale
layer.w2_blockscale_swizzled.data.copy_(w2_blockscale_swizzled)
layer.w2_weight = Parameter(layer.w2_weight.data, requires_grad=False)
# Both flashinfer cutlass and regular cutlass use same processing for w2
@@ -1614,7 +1752,7 @@ 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
return self.enable_flashinfer_cutlass_moe and self.moe_runner_config.is_gated
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
@@ -1631,10 +1769,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
x_sf = dispatch_output.hidden_states_scale
topk_output = dispatch_output.topk_output
assert (
self.moe_runner_config.activation == "silu"
), "Only SiLU activation is supported."
activation = self.moe_runner_config.activation
assert (
activation in ACT_STR_TO_TYPE_MAP
), f"{activation=} missing from {ACT_STR_TO_TYPE_MAP.keys()=}"
moe_runner_config = self.moe_runner_config
# Check if this is a FlashInferFP4MoE layer that should handle its own forward
@@ -1656,13 +1795,17 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
# If x_sf is not None, x is FP4 packed (half size), so we need * 2
# If x_sf is None, x is not packed, so output_col = x.shape[1]
output_col = x.shape[1] * 2 if x_sf is not None else x.shape[1]
output_col = x.shape[1]
if x_sf is not None and layer.moe_runner_config.is_gated:
output_col *= 2
with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric()
):
symm_output = torch.empty(
x.shape[0], output_col, dtype=output_dtype, device=x.device
x.shape[0],
output_col,
dtype=output_dtype,
device=x.device,
)
output = flashinfer_cutlass_fused_moe(
@@ -1687,6 +1830,7 @@ 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],
)[0]
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput

25
python/sglang/srt/server_args.py
View File

@@ -1194,6 +1194,23 @@ class ServerArgs:
f"Disabling Radix Cache for {model_arch} as it is not yet supported."
)
self.disable_radix_cache = True
elif model_arch in ["NemotronHForCausalLM"]:
if self.model_config.quantization in [
"modelopt",
"modelopt_fp8",
"modelopt_fp4",
]:
assert self.model_config.hf_config.mlp_hidden_act == "relu2"
if self.model_config.quantization == "modelopt":
self.quantization = (
"modelopt_fp4"
if self.model_config.hf_config.quantization_config["quant_algo"]
== "NVFP4"
else "modelopt_fp8"
)
else:
self.quantization = self.model_config.quantization
self.moe_runner_backend = "flashinfer_cutlass"
elif model_arch in [
"Qwen3MoeForCausalLM",
"Qwen3VLMoeForConditionalGeneration",
@@ -1491,9 +1508,11 @@ class ServerArgs:
def _handle_moe_kernel_config(self):
if self.moe_runner_backend == "flashinfer_cutlass":
assert (
self.quantization == "modelopt_fp4" or self.quantization is None
), "modelopt_fp4 quantization or bf16 is required for Flashinfer Cutlass MOE"
assert self.quantization in [
"modelopt_fp4",
"modelopt_fp8",
None,
], f"Invalid quantization '{self.quantization}'. \nFlashInfer Cutlass MOE supports only: 'modelopt_fp4', 'modelopt_fp8', or bfloat16 (None)."
assert self.ep_size in [
1,
self.tp_size,