From 1a18c1edf3f1016e9f6c1e145abc14760557d888 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 6 Jul 2026 21:28:47 -0500 Subject: [PATCH] [CI][LoRA] Preserve QKV fuser split sizes through LoRA wrappers Signed-off-by: Andreas Karatzas --- tests/lora/test_layers.py | 6 ++++++ vllm/lora/layers/column_parallel_linear.py | 11 +++++++++++ 2 files changed, 17 insertions(+) diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index c366b2cf2976..00465adbcd58 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -797,12 +797,18 @@ def create_column_parallel_packed_layer(idx: int = 0): params_dtype=torch.float16, prefix=f"layer_{idx}", ) + linear.split_sizes = [ + linear.num_heads * linear.head_size, + linear.num_kv_heads * linear.head_size, + linear.num_kv_heads * linear.v_head_size, + ] linear.weight.data = torch.rand_like(linear.weight.data) lora_linear = ( MergedQKVParallelLinearWithLoRA(linear) if not fully_shard else MergedQKVParallelLinearWithShardedLoRA(linear) ) + assert lora_linear.split_sizes == linear.split_sizes else: linear = QKVParallelLinear( 4096, diff --git a/vllm/lora/layers/column_parallel_linear.py b/vllm/lora/layers/column_parallel_linear.py index 12151699ac48..5dfb6f8575c3 100644 --- a/vllm/lora/layers/column_parallel_linear.py +++ b/vllm/lora/layers/column_parallel_linear.py @@ -467,6 +467,17 @@ def __init__(self, base_layer: QKVParallelLinear) -> None: self.kv_shard_id, self.kv_shard_id, ) + self.split_sizes = list( + getattr( + self.base_layer, + "split_sizes", + [ + self.q_proj_shard_size, + self.kv_proj_shard_size, + self.base_layer.num_kv_heads * self.base_layer.v_head_size, + ], + ) + ) def create_lora_weights( self,