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 7, 19:00 UTC
[Misc] Add shard_id validation for MergedColumnLinear (#35055)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Feb 25, 2026 at 14:12 UTC
a6c137521cf7218cc2da5f56aa3e68ad96aa76b1
1 file changed
+67
-7
vllm/model_executor/layers/linear.py
+67
-7
@@ -66,15 +66,23 @@ WEIGHT_LOADER_V2_SUPPORTED = [
66
]
67
68
69
-def adjust_marlin_shard(param, shard_size, shard_offset):
70
- marlin_tile_size = getattr(param, "marlin_tile_size", None)
69
+def adjust_marlin_shard(
70
+ param: Parameter,
71
+ shard_size: int,
72
+ shard_offset: int,
73
+) -> tuple[int, int]:
74
+ marlin_tile_size: int | None = getattr(param, "marlin_tile_size", None)
75
if marlin_tile_size is None:
76
return shard_size, shard_offset
77
78
return shard_size * marlin_tile_size, shard_offset * marlin_tile_size
79
80
77
-def adjust_block_scale_shard(weight_block_size, shard_size, shard_offset):
81
+def adjust_block_scale_shard(
82
+ weight_block_size: tuple[int, ...] | None,
83
+ shard_size: int,
84
+ shard_offset: int,
85
+) -> tuple[int, int]:
86
assert weight_block_size is not None
87
block_n = weight_block_size[0]
88
shard_offset = (shard_offset + block_n - 1) // block_n
@@ -83,7 +91,9 @@ def adjust_block_scale_shard(weight_block_size, shard_size, shard_offset):
91
92
93
def adjust_bitsandbytes_4bit_shard(
86
- param: Parameter, shard_offsets: dict[str, tuple[int, int]], loaded_shard_id: str
94
+ param: Parameter,
95
+ shard_offsets: dict[str, tuple[int, int]],
96
+ loaded_shard_id: str,
97
) -> tuple[int, int]:
98
"""Adjust the quantization offsets and sizes for BitsAndBytes sharding."""
99
@@ -97,7 +107,11 @@ def adjust_bitsandbytes_4bit_shard(
107
return quantized_size, quantized_offset
108
109
100
-def adjust_scalar_to_fused_array(param, loaded_weight, shard_id):
110
+def adjust_scalar_to_fused_array(
111
+ param_data: torch.Tensor,
112
+ loaded_weight: torch.Tensor,
113
+ shard_id: int | str,
114
+) -> tuple[torch.Tensor, torch.Tensor]:
115
"""For fused modules (QKV and MLP) we have an array of length
116
N that holds 1 scale for each "logical" matrix. So the param
117
is an array of length N. The loaded_weight corresponds to
@@ -117,12 +131,14 @@ def adjust_scalar_to_fused_array(param, loaded_weight, shard_id):
131
assert loaded_weight.shape[0] == 1
132
loaded_weight = loaded_weight[0]
133
120
- return param[shard_id], loaded_weight
134
+ return param_data[shard_id], loaded_weight
135
136
137
# TODO(Isotr0py): We might need a more flexible structure to handle
138
# bitsandbytes shard offsets.
125
-def left_shift_bitsandbytes_4bit_shard(bnb_weight_attrs: dict[str, Any]):
139
+def left_shift_bitsandbytes_4bit_shard(
140
+ bnb_weight_attrs: dict[str, Any],
141
+) -> tuple[dict[str, Any], dict[str, Any]]:
142
"""
143
Separate the BitsAndBytes 4-bit shard.
144
@@ -681,12 +697,41 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
697
disable_tp=disable_tp,
698
)
699
700
+ def validate_shard_id(self, loaded_shard_id: int | tuple[int, ...] | None):
701
+ if loaded_shard_id is None:
702
+ return
703
+ if isinstance(loaded_shard_id, tuple):
704
+ for idx in loaded_shard_id:
705
+ if not (0 <= idx < len(self.output_sizes)):
706
+ raise ValueError(
707
+ f"Shard id index {idx} should be between 0 and "
708
+ f"{len(self.output_sizes) - 1}. Got shard id {loaded_shard_id}."
709
+ )
710
+ if len(loaded_shard_id) > 1 and any(
711
+ b - a != 1 for a, b in zip(loaded_shard_id[:-1], loaded_shard_id[1:])
712
+ ):
713
+ raise ValueError(
714
+ "Shard id with multiple indices should be consecutive. "
715
+ f"Got shard id {loaded_shard_id}."
716
+ )
717
+ return
718
+ elif isinstance(loaded_shard_id, int):
719
+ if loaded_shard_id < 0 or loaded_shard_id >= len(self.output_sizes):
720
+ raise ValueError(
721
+ f"Shard id should be between 0 and {len(self.output_sizes) - 1}. "
722
+ f"Got shard id {loaded_shard_id}."
723
+ )
724
+ return
725
+ raise ValueError("This line should not be reached")
726
+
727
def weight_loader(
728
self,
729
param: Parameter,
730
loaded_weight: torch.Tensor,
731
loaded_shard_id: tuple[int, ...] | int | None = None,
732
):
733
+ self.validate_shard_id(loaded_shard_id)
734
+ # FIXME(Isotr0py): Enable tuple shard_id for BNB quantization.
735
if isinstance(loaded_shard_id, tuple):
736
raise NotImplementedError(
737
"Shard id with multiple indices is not supported in weight_loader, "
@@ -874,6 +919,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
919
loaded_weight: torch.Tensor,
920
loaded_shard_id: tuple[int, ...] | int | None = None,
921
):
922
+ self.validate_shard_id(loaded_shard_id)
923
if loaded_shard_id is None or isinstance(loaded_shard_id, tuple):
924
if isinstance(param, PerTensorScaleParameter):
925
param.load_merged_column_weight(loaded_weight=loaded_weight, shard_id=0)
@@ -1005,6 +1051,18 @@ class QKVParallelLinear(ColumnParallelLinear):
1051
disable_tp=disable_tp,
1052
)
1053
1054
+ def validate_shard_id(self, loaded_shard_id: str | None):
1055
+ if loaded_shard_id is None:
1056
+ return
1057
+ if isinstance(loaded_shard_id, str):
1058
+ if loaded_shard_id not in ["q", "k", "v"]:
1059
+ raise ValueError(
1060
+ "Shard id for QKVParallelLinear should be 'q', 'k', or 'v', "
1061
+ f"got shard id {loaded_shard_id}."
1062
+ )
1063
+ return
1064
+ raise ValueError("This line should not be reached")
1065
+
1066
def _get_shard_offset_mapping(self, loaded_shard_id: str):
1067
shard_offset_mapping = {
1068
"q": 0,
@@ -1073,6 +1131,7 @@ class QKVParallelLinear(ColumnParallelLinear):
1131
loaded_weight: torch.Tensor,
1132
loaded_shard_id: str | None = None,
1133
):
1134
+ self.validate_shard_id(loaded_shard_id)
1135
if loaded_shard_id is None: # special case for certain models
1136
if isinstance(param, PerTensorScaleParameter):
1137
param.load_qkv_weight(
@@ -1112,6 +1171,7 @@ class QKVParallelLinear(ColumnParallelLinear):
1171
loaded_weight: torch.Tensor,
1172
loaded_shard_id: str | None = None,
1173
):
1174
+ self.validate_shard_id(loaded_shard_id)
1175
# Special case for GGUF
1176
# initialize GGUF param after we know the quantize type
1177
is_gguf_weight = getattr(param, "is_gguf_weight", False)