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

Fix update weight error for blackwell DeepGEMM (#13910)

Browse Source
This commit is contained in:
fzyzcjy
2025-11-25 13:28:12 -08:00
committed by GitHub
parent 1ab6ce0e62
commit 64a11303ce
2 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/linear.py
+4 -2
View File
@@ -736,7 +736,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
if isinstance(param, BlockQuantScaleParameter): if isinstance(param, BlockQuantScaleParameter):
weight_block_size = self.quant_method.quant_config.weight_block_size weight_block_size = self.quant_method.quant_config.weight_block_size
block_n, _ = weight_block_size[0], weight_block_size[1] raw_block_n, _ = weight_block_size[0], weight_block_size[1]
block_n = 1 if getattr(param, "format_ue8m0", False) else raw_block_n
shard_offset = ( shard_offset = (
(sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // block_n (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // block_n
) // self.tp_size ) // self.tp_size
@@ -966,7 +967,8 @@ class QKVParallelLinear(ColumnParallelLinear):
if isinstance(param, BlockQuantScaleParameter): if isinstance(param, BlockQuantScaleParameter):
weight_block_size = self.quant_method.quant_config.weight_block_size weight_block_size = self.quant_method.quant_config.weight_block_size
block_n, _ = weight_block_size[0], weight_block_size[1] raw_block_n, _ = weight_block_size[0], weight_block_size[1]
block_n = 1 if getattr(param, "format_ue8m0", False) else raw_block_n
shard_offset = (shard_offset + block_n - 1) // block_n shard_offset = (shard_offset + block_n - 1) // block_n
shard_size = (shard_size + block_n - 1) // block_n shard_size = (shard_size + block_n - 1) // block_n
python/sglang/srt/layers/quantization/fp8.py
+3 -3
View File
@@ -262,7 +262,6 @@ class Fp8LinearMethod(LinearMethodBase):
layer.input_size_per_partition = input_size_per_partition layer.input_size_per_partition = input_size_per_partition
layer.output_size_per_partition = output_size_per_partition layer.output_size_per_partition = output_size_per_partition
layer.orig_dtype = params_dtype layer.orig_dtype = params_dtype
layer.executed_weight_requant_ue8m0 = False
# WEIGHT # WEIGHT
weight_dtype = ( weight_dtype = (
@@ -300,6 +299,7 @@ class Fp8LinearMethod(LinearMethodBase):
output_dim=0, output_dim=0,
weight_loader=weight_loader, weight_loader=weight_loader,
) )
scale.format_ue8m0 = False
scale[:] = torch.finfo(torch.float32).min scale[:] = torch.finfo(torch.float32).min
layer.register_parameter("weight_scale_inv", scale) layer.register_parameter("weight_scale_inv", scale)
else: else:
@@ -367,14 +367,14 @@ class Fp8LinearMethod(LinearMethodBase):
self.w8a8_block_fp8_linear self.w8a8_block_fp8_linear
is deepgemm_w8a8_block_fp8_linear_with_fallback is deepgemm_w8a8_block_fp8_linear_with_fallback
) )
and (not layer.executed_weight_requant_ue8m0) and (not layer.weight_scale_inv.format_ue8m0)
): ):
requant_weight_ue8m0_inplace( requant_weight_ue8m0_inplace(
layer.weight, layer.weight,
layer.weight_scale_inv, layer.weight_scale_inv,
self.quant_config.weight_block_size, self.quant_config.weight_block_size,
) )
layer.executed_weight_requant_ue8m0 = True layer.weight_scale_inv.format_ue8m0 = True
weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data
layer.weight.data = weight.data layer.weight.data = weight.data