[Bug]: LoRA set_lora on merged QKV/gate_up raises IndexError when lora_a is a single tensor

### Your current environment Reproduced on the **official upstream nightly image**, `vllm/vllm-openai:nightly-aarch64`. <details> <summary>Output of python3 -m vllm.collect_env</summary> ``` ============================== System Info ============================== OS : Ubuntu 22.04.5 LTS (aarch64) GCC version : (Ubuntu 11.4.0-1ubuntu1~22.04.3) 11.4.0 Libc version : glibc-2.35 ============================== PyTorch Info ============================== PyTorch version : 2.13.0+cu130 Is debug build : False CUDA used to build PyTorch : 13.0 ============================== Python Environment ============================== Python version : 3.12.13 (main, Mar 4 2026, 09:23:07) [GCC 11.4.0] (64-bit runtime) Python platform : Linux-6.17.0-1014-nvidia-64k-aarch64-with-glibc2.35 ============================== CUDA / GPU Info ============================== Is CUDA available : True CUDA runtime version : 13.0.88 GPU models and configuration : GPU 0: NVIDIA GB300 GPU 1: NVIDIA GB300 GPU 2: NVIDIA GB300 GPU 3: NVIDIA GB300 Nvidia driver version : 580.159.04 ============================== Versions of relevant libraries ============================== [pip3] torch==2.13.0+cu130 [pip3] torchaudio==2.11.0+cu130 [pip3] torchvision==0.28.0+cu130 [pip3] transformers==5.14.1 [pip3] triton==3.7.1 [pip3] nvidia-nccl-cu13==2.30.7 ============================== vLLM Info ============================== vLLM Version : 0.26.1rc1.dev457+gc810e5ee9 (git sha: c810e5ee9) vLLM Build Flags: CUDA Archs: 8.0 8.7 8.9 9.0 10.0 11.0 12.0; ROCm: Disabled; XPU: Disabled ============================== Environment Variables ============================== VLLM_IMAGE_TAG=vllm/vllm-openai:nightly-c810e5ee9976ad86b81d1277b53e76d0ee639414 VLLM_BUILD_COMMIT=c810e5ee9976ad86b81d1277b53e76d0ee639414 VLLM_BUILD_URL=https://buildkite.com/vllm/release-v2/builds/4838 VLLM_USAGE_SOURCE=production-docker-image CUDA_VERSION=13.0.2 TORCH_CUDA_ARCH_LIST=8.0 8.7 8.9 9.0 10.0 11.0 12.0 ``` </details> The affected code is unchanged on `main` as of `448344c0e2`. ### 🐛 Describe the bug `MergedColumnParallelLinearWithLoRA.set_lora()` declares ```python lora_a: torch.Tensor | list[torch.Tensor] ``` but only normalizes `lora_b`. A single 2-D `lora_a` of shape `[rank, input_size]` is passed through unexpanded, so `lora_a[i]` yields a 1-D **row** instead of the i-th slice's matrix. This raises `IndexError` on any packed layer — fused QKV (3 slices) or gate_up (2 slices). `vllm/lora/layers/column_parallel_linear.py`: ```python def set_lora(self, index, lora_a: torch.Tensor | list[torch.Tensor], # a single tensor IS a declared input lora_b: torch.Tensor | list[torch.Tensor]): self.reset_lora(index) if isinstance(lora_b, list) and len(lora_b) != self.n_slices: # only lora_b is normalized lora_a, lora_b = self.expand_packed_lora(lora_a, lora_b) # nothing expands a single-tensor lora_a into a per-slice list if self.tp_size > 1: lora_a = self.slice_lora_a(lora_a) for i in range(self.n_slices): if (lora_a_i := lora_a[i]) is not None: self.lora_a_stacked[i][index, 0, : lora_a_i.shape[0], : lora_a_i.shape[1]] ``` Confirmed directly against the shipped nightly binary: ``` $ python3 -c "import inspect; from vllm.lora.layers.column_parallel_linear import \ MergedColumnParallelLinearWithLoRA as M; \ print('NORMALIZES_LORA_A:', 'isinstance(lora_a, torch.Tensor)' in inspect.getsource(M.set_lora))" NORMALIZES_LORA_A: False ``` It breaks at four distinct sites: | Combination | Function | Error | |---|---|---| | fused QKV, `fully_sharded_loras=True` | `MergedQKVParallelLinearWithShardedLoRA.slice_lora_a` | `too many indices for tensor of dimension 1` | | gate_up, `fully_sharded_loras=True` | `MergedColumnParallelLinearWithShardedLoRA.slice_lora_a` | `too many indices for tensor of dimension 1` | | fused QKV, `fully_sharded_loras=False` | `set_lora` copy loop | `tuple index out of range` | | gate_up, `fully_sharded_loras=False` | `set_lora` copy loop | `tuple index out of range` | The two `slice_lora_a` sites sit behind `if self.tp_size > 1`. The two copy-loop sites do **not** — those reproduce at TP=1 as well. ### Why a single tensor is a legitimate input It is a declared input, not caller misuse. Only a *fused* source produces one A — e.g. Megatron, which stores QKV as a single linear. An ordinary PEFT checkpoint stores q/k/v separately and the loader assembles a list, which is why this is not seen in normal serving. Such weights reach `set_lora` when an RL trainer pushes them **in-memory** (`update_weights_from_ipc -> add_lora -> set_lora`), bypassing the loader entirely. This is how we hit it: RL post-training (verl + vLLM + Megatron) with LoRA and TP>1 raises `IndexError` on the first weight sync, the rollout worker dies, and the training run aborts. vLLM already treats `lora_a` as shared across slices elsewhere — `expand_packed_lora()` does `expanded_a.append(a_i)`, with the comment *"Split b_i into per-slice tensors and replicate a_i for each"*. And `MergedColumnParallelLinearVariableSliceWithLoRA.set_lora()` in the same file already contains exactly the missing normalization, under the comment *"Override to handle single tensor weights"*. Only the bare-tensor entry point on the base class was missed. ### Scope Adapters loaded **from disk** cannot reach this. I tested it: `worker_manager._load_adapter()` builds `expected_lora_modules` by expanding packed modules into their constituents (`q_proj`/`k_proj`/`v_proj`), so a checkpoint targeting the fused name (`qkv_proj`) is rejected by `check_unexpected_modules()` with a clean `ValueError` before reaching `set_lora`. The in-memory weight-push path is the only route. ### Reproduce 2 GPUs, no model and no kernel — finishes in seconds: ```python # builds QKVParallelLinear / MergedColumnParallelLinear + the LoRA wrapper, # then calls set_lora(0, lora_a, lora_b) with lora_a as ONE 2-D tensor torchrun --nproc_per_node=2 repro.py ``` Running a 2x2x2 matrix — {fused QKV, gate_up} x {`fully_sharded_loras` True, False} x {`lora_a` single tensor, per-slice list} — on the nightly image above: ``` vllm 0.26.1rc1.dev457+gc810e5ee9 tp=2 [FAIL] qkv fully_sharded=True lora_a=single-tensor File ".../vllm/lora/layers/column_parallel_linear.py", line 317, in set_lora File ".../vllm/lora/layers/column_parallel_linear.py", line 639, in slice_lora_a IndexError: too many indices for tensor of dimension 1 [FAIL] qkv fully_sharded=False lora_a=single-tensor File ".../vllm/lora/layers/column_parallel_linear.py", line 323, in set_lora IndexError: tuple index out of range [FAIL] gate_up fully_sharded=True lora_a=single-tensor File ".../vllm/lora/layers/column_parallel_linear.py", line 559, in slice_lora_a IndexError: too many indices for tensor of dimension 1 [FAIL] gate_up fully_sharded=False lora_a=single-tensor File ".../vllm/lora/layers/column_parallel_linear.py", line 323, in set_lora IndexError: tuple index out of range [PASS] qkv fully_sharded=True lora_a=list [PASS] qkv fully_sharded=False lora_a=list [PASS] gate_up fully_sharded=True lora_a=list [PASS] gate_up fully_sharded=False lora_a=list 4 combination(s) failed ``` 4/4 single-tensor combinations fail on every run; the 4 per-slice-list controls pass, confirming only the single-tensor entry point is affected. ### No existing test covers it `tests/lora/test_layers.py` is the only file in the whole `tests/` tree with a direct `set_lora()` call, and its call sites miss the combination: - `populate_loras()` exercises the merged classes, but feeds them via `PackedLoRALayerWeights.pack()`, which always builds a **list** `lora_a`. - The one site that passes a bare tensor targets `MergedColumnParallelLinearVariableSliceWithLoRA` — the sibling class that already normalizes. ### Related but distinct PR #37019 generalizes the same `slice_lora_a` methods from hardcoded 2/3 subloras to arbitrary N, driven by Qwen3.5 GDN 4-way fusion. I checked its diff: it does not touch the base class's `set_lora`, so this gap survives it. ### Fix PR to follow — replicate the shared `lora_a` to one entry per group, counted by `len(lora_b)` rather than `n_slices` so `expand_packed_lora()`'s zip against `lora_b`'s groups keeps working. ### Before submitting a new issue... - [x] Make sure you already searched for relevant issues, and asked the core devs or community in the [Discussion Forum](https://discuss.vllm.ai/c/help/) first, before raising a new issue.