A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 19:00 UTC
[Bugfix] Fix weight_loader v1 block scale (#31103)
Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>
Kyuyeun Kim committed
Jan 1, 2026 at 21:14 UTC
cc410e864480aee7fc01cfc6017abca677b198df
1 file changed
+40
-27
vllm/model_executor/layers/linear.py
+40
-27
@@ -80,6 +80,14 @@ def adjust_marlin_shard(param, shard_size, shard_offset):
80
return shard_size * marlin_tile_size, shard_offset * marlin_tile_size
81
82
83
+def adjust_block_scale_shard(weight_block_size, shard_size, shard_offset):
84
+ assert weight_block_size is not None
85
+ block_n = weight_block_size[0]
86
+ shard_offset = (shard_offset + block_n - 1) // block_n
87
+ shard_size = (shard_size + block_n - 1) // block_n
88
+ return shard_size, shard_offset
89
+
90
+
91
def adjust_bitsandbytes_4bit_shard(
92
param: Parameter, shard_offsets: dict[str, tuple[int, int]], loaded_shard_id: str
93
) -> tuple[int, int]:
@@ -763,8 +771,18 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
771
772
assert loaded_shard_id < len(self.output_sizes)
773
if output_dim is not None:
766
- shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size
767
- shard_size = self.output_sizes[loaded_shard_id] // self.tp_size
774
+ shard_offset = sum(self.output_sizes[:loaded_shard_id])
775
+ shard_size = self.output_sizes[loaded_shard_id]
776
+
777
+ if isinstance(param, BlockQuantScaleParameter):
778
+ weight_block_size = getattr(self, "weight_block_size", None)
779
+ shard_size, shard_offset = adjust_block_scale_shard(
780
+ weight_block_size, shard_size, shard_offset
781
+ )
782
+
783
+ shard_offset //= self.tp_size
784
+ shard_size //= self.tp_size
785
+
786
# Special case for quantization.
787
# If quantized, we need to adjust the offset and size to account
788
# for the packing.
@@ -867,24 +885,17 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
885
886
assert loaded_shard_id < len(self.output_sizes)
887
888
+ shard_offset = sum(self.output_sizes[:loaded_shard_id])
889
+ shard_size = self.output_sizes[loaded_shard_id]
890
+
891
if isinstance(param, BlockQuantScaleParameter):
871
- assert self.quant_method is not None
872
- # Assume the weight block size has been set by quant method
873
- assert hasattr(self, "weight_block_size")
874
- weight_block_size = self.weight_block_size
875
- assert weight_block_size is not None
876
- block_n, _ = weight_block_size[0], weight_block_size[1]
877
- shard_offset = (
878
- (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // block_n
879
- ) // self.tp_size
880
- shard_size = (
881
- (self.output_sizes[loaded_shard_id] + block_n - 1)
882
- // block_n
883
- // self.tp_size
892
+ weight_block_size = getattr(self, "weight_block_size", None)
893
+ shard_size, shard_offset = adjust_block_scale_shard(
894
+ weight_block_size, shard_size, shard_offset
895
)
885
- else:
886
- shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size
887
- shard_size = self.output_sizes[loaded_shard_id] // self.tp_size
896
+
897
+ shard_offset //= self.tp_size
898
+ shard_size //= self.tp_size
899
900
param.load_merged_column_weight(
901
loaded_weight=loaded_weight,
@@ -1066,16 +1077,11 @@ class QKVParallelLinear(ColumnParallelLinear):
1077
shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
1078
shard_size = self._get_shard_size_mapping(loaded_shard_id)
1079
1069
- # Note(simon): This is needed for Qwen3's fp8 quantization.
1080
if isinstance(param, BlockQuantScaleParameter):
1071
- assert self.quant_method is not None
1072
- # Assume the weight block size has been set by quant method
1073
- assert hasattr(self, "weight_block_size")
1074
- weight_block_size = self.weight_block_size
1075
- assert weight_block_size is not None
1076
- block_n, _ = weight_block_size[0], weight_block_size[1]
1077
- shard_offset = (shard_offset + block_n - 1) // block_n
1078
- shard_size = (shard_size + block_n - 1) // block_n
1081
+ weight_block_size = getattr(self, "weight_block_size", None)
1082
+ shard_size, shard_offset = adjust_block_scale_shard(
1083
+ weight_block_size, shard_size, shard_offset
1084
+ )
1085
1086
param.load_qkv_weight(
1087
loaded_weight=loaded_weight,
@@ -1208,6 +1214,13 @@ class QKVParallelLinear(ColumnParallelLinear):
1214
elif loaded_shard_id == "v":
1215
shard_offset = (self.num_heads + self.num_kv_heads) * self.head_size
1216
shard_size = self.num_kv_heads * self.v_head_size
1217
+
1218
+ if isinstance(param, BlockQuantScaleParameter):
1219
+ weight_block_size = getattr(self, "weight_block_size", None)
1220
+ shard_size, shard_offset = adjust_block_scale_shard(
1221
+ weight_block_size, shard_size, shard_offset
1222
+ )
1223
+
1224
# Special case for Quantized Weights.
1225
# If quantized, we need to adjust the offset and size to account
1226
# for the packing.