[CI] Fix Transformers modeling backend LoRA test (#47832)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jul 7, 2026 at 12:40 UTC
0ed05b6f824c21dfa93f9a2d4c23ffabe13bdae4
3 files changed
+8
-18
tests/models/transformers/fusers/test_linear.py
+4
-2
@@ -258,7 +258,8 @@ def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser):
258
merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0))
259
if q.bias is not None:
260
merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0))
261
- merged.split_sizes = [q.out_features, k.out_features, v.out_features]
261
+ merged.output_sizes = [q.out_features, k.out_features, v.out_features]
262
+ merged.tp_size = 1
263
setattr(module, fuser.merged_name, merged)
264
for name in (fuser.q_name, fuser.k_name, fuser.v_name):
265
delattr(module, name)
@@ -332,7 +333,8 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
333
# original semantics (branches, kwargs, attribute reads)
334
code = fuser.fused_forward.__code__
335
names = code.co_names
335
- assert "qkv_proj" in names and "split_sizes" in names and "o_proj" in names
336
+ assert "qkv_proj" in names and "output_sizes" in names and "o_proj" in names
337
+ assert "tp_size" in names
338
assert not {"q_proj", "k_proj", "v_proj"} & set(names)
339
if attn_cls is FakeAttention:
340
assert "update" in names # the cache branch survives
vllm/model_executor/layers/linear.py
+1
-5
@@ -1000,16 +1000,12 @@ class QKVParallelLinear(ColumnParallelLinear):
1000
self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
1001
self.num_kv_head_replicas = 1
1002
input_size = self.hidden_size
1003
- output_size = (
1004
- self.num_heads * self.head_size
1005
- + self.num_kv_heads * self.head_size
1006
- + self.num_kv_heads * self.v_head_size
1007
- ) * tp_size
1003
self.output_sizes = [
1004
self.num_heads * self.head_size * tp_size, # q_proj
1005
self.num_kv_heads * self.head_size * tp_size, # k_proj
1006
self.num_kv_heads * self.v_head_size * tp_size, # v_proj
1007
]
1008
+ output_size = sum(self.output_sizes)
1009
1010
super().__init__(
1011
input_size=input_size,
vllm/model_executor/models/transformers/fusers/qkv.py
+3
-11
@@ -127,15 +127,14 @@ class QKVFuser(StackedFuser):
127
if len({id(block) for block, _ in blocks}) != 1:
128
raise ValueError("projection calls are in different blocks")
129
130
- # q(x), k(x), v(x) -> q, k, v = qkv(x).split(self.qkv.split_sizes, -1)
130
+ # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
131
names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
132
temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)]
133
if names & set(temps):
134
raise ValueError("fused temporaries would shadow existing names")
135
merged = f"self.{self.merged_name}"
136
- template = (
137
- f"{', '.join(temps)} = {merged}(__arg__).split({merged}.split_sizes, -1)"
138
- )
136
+ sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
137
+ template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)"
138
assign = ast.parse(template).body[0]
139
arg = next(
140
node
@@ -198,13 +197,6 @@ class QKVFuser(StackedFuser):
197
self.merged_name,
198
merged,
199
)
201
- # The rewritten forward splits the merged projection into this rank's
202
- # shard sizes (see `update_forward`)
203
- merged.split_sizes = [
204
- merged.num_heads * merged.head_size,
205
- merged.num_kv_heads * merged.head_size,
206
- merged.num_kv_heads * merged.v_head_size,
207
- ]
200
setattr(module, self.merged_name, merged)
201
# Drop the consumed submodules so their (meta) params are not expected.
202
for name in (self.q_name, self.k_name, self.v_name):