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 5ff4722ab..589f0c0ed 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 1a546dddb..d23bd075c 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -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: diff --git a/python/sglang/srt/layers/utils/common.py b/python/sglang/srt/layers/utils/common.py index e88f3a938..b982b247e 100644 --- a/python/sglang/srt/layers/utils/common.py +++ b/python/sglang/srt/layers/utils/common.py @@ -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