chenchenghao/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-26 05: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

6
python/sglang/srt/layers/linear.py
View File

@@ -736,7 +736,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
if isinstance(param, BlockQuantScaleParameter):
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 = (
(sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // block_n
) // self.tp_size
@@ -966,7 +967,8 @@ class QKVParallelLinear(ColumnParallelLinear):
if isinstance(param, BlockQuantScaleParameter):
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_size = (shard_size + block_n - 1) // block_n

6
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -262,7 +262,6 @@ class Fp8LinearMethod(LinearMethodBase):
layer.input_size_per_partition = input_size_per_partition
layer.output_size_per_partition = output_size_per_partition
layer.orig_dtype = params_dtype
layer.executed_weight_requant_ue8m0 = False
# WEIGHT
weight_dtype = (
@@ -300,6 +299,7 @@ class Fp8LinearMethod(LinearMethodBase):
output_dim=0,
weight_loader=weight_loader,
)
scale.format_ue8m0 = False
scale[:] = torch.finfo(torch.float32).min
layer.register_parameter("weight_scale_inv", scale)
else:
@@ -367,14 +367,14 @@ class Fp8LinearMethod(LinearMethodBase):
self.w8a8_block_fp8_linear
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(
layer.weight,
layer.weight_scale_inv,
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
layer.weight.data = weight.data