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 4, 21:15 UTC
Add new tp plan styles to the Transformers modelling backend (#40467)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Apr 21, 2026 at 16:51 UTC
6ee081d1d07bbcf7472b7bf76a34ae23310a36af
1 file changed
+15
-3
vllm/model_executor/models/transformers/utils.py
+15
-3
@@ -94,7 +94,15 @@ def init_on_device_without_buffers(device: torch.device):
94
setattr(torch, torch_function_name, old_torch_function)
95
96
97
-Style = Literal["colwise", "colwise_rep", "rowwise", "rowwise_rep", "replicate"]
97
+Style = Literal[
98
+ "colwise",
99
+ "rowwise",
100
+ "replicate",
101
+ "colwise_gather_output",
102
+ "rowwise_split_input",
103
+ "colwise_rep",
104
+ "rowwise_rep",
105
+]
106
107
108
def replace_linear_class(
@@ -120,10 +128,14 @@ def replace_linear_class(
128
129
vllm_linear_cls, vllm_linear_kwargs = {
130
"colwise": (ColumnParallelLinear, {}),
123
- "colwise_rep": (ColumnParallelLinear, {"gather_output": True}),
131
"rowwise": (RowParallelLinear, {}),
125
- "rowwise_rep": (RowParallelLinear, {"input_is_parallel": False}),
132
"replicate": (ReplicatedLinear, {}),
133
+ # Transformers v5
134
+ "colwise_gather_output": (ColumnParallelLinear, {"gather_output": True}),
135
+ "rowwise_split_input": (RowParallelLinear, {"input_is_parallel": False}),
136
+ # Transformers v4
137
+ "colwise_rep": (ColumnParallelLinear, {"gather_output": True}),
138
+ "rowwise_rep": (RowParallelLinear, {"input_is_parallel": False}),
139
}.get(style, (ReplicatedLinear, {}))
140
141
return vllm_linear_cls(