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

Fix nvfp4 weight update (#18085)

Browse Source
This commit is contained in:
Ziang Li
2026-02-27 14:55:08 -08:00
committed by GitHub
parent 6ca7da3e7c
commit 9469ad089b
3 changed files with 99 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
View File

@@ -21,6 +21,7 @@ from sglang.srt.layers.quantization.fp8_kernel import (
per_token_group_quant_fp8,
scaled_fp8_quant,
)
from sglang.srt.layers.utils import copy_or_rebind_param
from sglang.srt.utils.common import (
is_cuda_alike,
is_flashinfer_available,
@@ -151,25 +152,22 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
)
# Set flashinfer parameters
layer.gemm1_weights_fp4_shuffled = Parameter(
gemm1_weights_fp4_shuffled, requires_grad=False
copy_or_rebind_param(
layer, "gemm1_weights_fp4_shuffled", gemm1_weights_fp4_shuffled
)
layer.gemm2_weights_fp4_shuffled = Parameter(
gemm2_weights_fp4_shuffled, requires_grad=False
)
layer.gemm1_scales_fp4_shuffled = Parameter(
gemm1_scales_fp4_shuffled, requires_grad=False
)
layer.gemm2_scales_fp4_shuffled = Parameter(
gemm2_scales_fp4_shuffled, requires_grad=False
copy_or_rebind_param(
layer, "gemm2_weights_fp4_shuffled", gemm2_weights_fp4_shuffled
)
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
w2_input_scale_quant = cast(torch.Tensor, layer.w2_input_scale_quant)
g1_alphas = cast(torch.Tensor, layer.g1_alphas)
layer.g1_scale_c = Parameter(
copy_or_rebind_param(
layer,
"g1_scale_c",
(w2_input_scale_quant * g1_alphas).to(torch.float32),
requires_grad=False,
)
# Clean up weights that won't be used by TRT-LLM

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

@@ -49,6 +49,7 @@ 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.utils.common import (
get_bool_env_var,
is_cuda,
@@ -1180,13 +1181,13 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
input_scale_2 = layer.input_scale.max().to(torch.float32)
weight_scale_2 = layer.weight_scale_2.max().to(torch.float32)
layer.input_scale = Parameter(input_scale_2, requires_grad=False)
layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False)
layer.alpha = Parameter(
layer.input_scale * layer.weight_scale_2, requires_grad=False
# Keep per-shard scales intact for hot reload; derive scalar params below.
copy_or_rebind_param(
layer, "alpha", (input_scale_2 * weight_scale_2).to(torch.float32)
)
layer.input_scale_inv = Parameter(
(1 / input_scale_2).to(torch.float32), requires_grad=False
copy_or_rebind_param(
layer, "input_scale_inv", (1 / input_scale_2).to(torch.float32)
)
# Store original output size before any padding
@@ -1238,15 +1239,15 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
.view(torch.float8_e4m3fn)
)
layer.weight_scale_interleaved = Parameter(scale, requires_grad=False)
layer.weight = Parameter(weight, requires_grad=False)
copy_or_rebind_param(layer, "weight_scale_interleaved", scale)
copy_or_rebind_param(layer, "weight", weight)
layer.weights_padding_cols = weights_padding_cols
return
# Pad weights for CUTLASS/FlashInfer kernel alignment (K and N divisible by 32)
weight, weights_padding_cols = pad_nvfp4_weight(layer.weight.data)
layer.weights_padding_cols = weights_padding_cols
layer.weight = Parameter(weight, requires_grad=False)
copy_or_rebind_param(layer, "weight", weight)
# Pad and blockwise interleave weight_scale
scales = layer.weight_scale
@@ -1270,7 +1271,7 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
if scale_ndim == 2
else padded_scales.reshape(B, M_padded, K_padded)
)
layer.weight_scale_interleaved = Parameter(padded_scales, requires_grad=False)
copy_or_rebind_param(layer, "weight_scale_interleaved", padded_scales)
def apply(
self,
@@ -1496,18 +1497,22 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
# GEMM 1 scale processing
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."
)
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."
)
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
if self.enable_flashinfer_cutlass_moe or self.enable_flashinfer_trtllm_moe:
@@ -1549,19 +1554,25 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
w2_input_scale = layer.w2_input_scale
# Create shared parameters
layer.g1_alphas = Parameter(
copy_or_rebind_param(
layer,
"g1_alphas",
(w13_input_scale * w13_weight_scale_2).to(torch.float32),
requires_grad=False,
)
layer.g2_alphas = Parameter(
copy_or_rebind_param(
layer,
"g2_alphas",
(w2_input_scale * layer.w2_weight_scale_2).to(torch.float32),
requires_grad=False,
)
layer.w13_input_scale_quant = Parameter(
(1 / w13_input_scale).to(torch.float32), requires_grad=False
copy_or_rebind_param(
layer,
"w13_input_scale_quant",
(1 / w13_input_scale).to(torch.float32),
)
layer.w2_input_scale_quant = Parameter(
(1 / w2_input_scale).to(torch.float32), requires_grad=False
copy_or_rebind_param(
layer,
"w2_input_scale_quant",
(1 / w2_input_scale).to(torch.float32),
)
# TODO: for flashinfer always do MOE_NVFP4_DISPATCH
@@ -1621,8 +1632,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
# Process w13 weights
w13_blockscale_swizzled = swizzle_blockscale(layer.w13_weight_scale)
del layer.w13_weight_scale
layer.w13_blockscale_swizzled.data.copy_(w13_blockscale_swizzled)
copy_or_rebind_param(
layer, "w13_blockscale_swizzled", w13_blockscale_swizzled
)
w13_weight = layer.w13_weight
intermediate_size_pad = w13_blockscale_swizzled.size(1) - w13_weight.size(1)
@@ -1634,46 +1646,56 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
"but padding is also implemented for gated activations"
)
layer.w13_weight = Parameter(
copy_or_rebind_param(
layer,
"w13_weight",
torch.nn.functional.pad(
w13_weight, (0, 0, 0, intermediate_size_pad)
),
requires_grad=False,
)
layer.w2_weight = Parameter(
copy_or_rebind_param(
layer,
"w2_weight",
torch.nn.functional.pad(
layer.w2_weight, (0, intermediate_size_pad // 2, 0, 0)
),
requires_grad=False,
)
layer.w2_weight_scale = Parameter(
copy_or_rebind_param(
layer,
"w2_weight_scale",
torch.nn.functional.pad(
layer.w2_weight_scale, (0, intermediate_size_pad // 16)
),
requires_grad=False,
)
layer.w2_blockscale_swizzled = Parameter(
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 = swizzle_blockscale(layer.w2_weight_scale)
del layer.w2_weight_scale
layer.w2_blockscale_swizzled.data.copy_(w2_blockscale_swizzled)
copy_or_rebind_param(
layer, "w2_blockscale_swizzled", w2_blockscale_swizzled
)
# Both flashinfer cutlass and regular cutlass use same processing for w2
# Set up CUTLASS MoE parameters
# Set up CUTLASS MoE parameters (reuse to keep CUDA graph stable)
device = layer.w13_weight.device
layer.cutlass_moe_params = CutlassMoEParams(
CutlassMoEType.BlockscaledFP4,
device,
num_experts=layer.num_experts, # global num experts
intermediate_size_per_partition=layer.w2_weight.shape[2] * 2, # n
hidden_size=layer.w13_weight.shape[2] * 2,
) # k
inter_size = layer.w2_weight.shape[2] * 2
hidden_size = layer.w13_weight.shape[2] * 2
existing_params = getattr(layer, "cutlass_moe_params", None)
if (
existing_params is None
or existing_params.cutlass_moe_type != CutlassMoEType.BlockscaledFP4
or existing_params.num_experts != layer.num_experts
or existing_params.intermediate_size_per_partition != inter_size
or existing_params.hidden_size != hidden_size
or existing_params.device != device
):
layer.cutlass_moe_params = CutlassMoEParams(
CutlassMoEType.BlockscaledFP4,
device,
num_experts=layer.num_experts, # global num experts
intermediate_size_per_partition=inter_size, # n
hidden_size=hidden_size,
) # k
@property
def load_up_proj_weight_first(self) -> bool:

17
python/sglang/srt/layers/utils/common.py
View File

@@ -2,6 +2,7 @@ import logging
import re
import torch
from torch.nn.parameter import Parameter
logger = logging.getLogger(__name__)
@@ -37,6 +38,22 @@ def pad_or_narrow_weight(
)
def copy_or_rebind_param(
module: torch.nn.Module, name: str, new_value: torch.Tensor
) -> None:
"""Keep parameter identities stable for CUDA graph reuse and hot reload."""
new_value = new_value.detach()
param = getattr(module, name, None)
if isinstance(param, Parameter):
if param.data.shape == new_value.shape and param.data.dtype == new_value.dtype:
param.data.copy_(new_value)
else:
param.data = new_value
param.requires_grad_(False)
else:
setattr(module, name, Parameter(new_value, requires_grad=False))
class PPMissingLayer(torch.nn.Identity):
# Adapted from
# https://github.com/vllm-project/vllm/blob/18ed3132d2bfe1df9a74729457b69243955221e8/vllm/model_executor/models/utils.py#L468C1-L486C1