[Bug][XPU] --enable-lora on AutoRound int4 (INC ark_linear) crashes at startup; unconditional all_gather in _mcp_apply breaks non-sharded LoRA at TP>1 (patches included)

[marcorigodanzo]

Summary

Enabling --enable-lora on an AutoRound int4 model (INC ark_linear path) on XPU hits two independent bugs in v0.24.0. With the two small patches below, LoRA serving works end-to-end on 2× Arc Pro B70 (Battlemage), TP=2: clean boot, 12/12 guided-JSON validity, runtime adapter loading via /v1/load_lora_adapter, and a zero-weight (lora_B = 0) test adapter produces outputs bit-identical to the base model at temperature 0. I'm happy to send PR(s) for both.

1. _get_lora_device doesn't know the INC/AutoRound ark_linear layout → --enable-lora crash-loops at startup for any INC WNA16 (AutoRound int4) checkpoint.
2. Unconditional tensor_model_parallel_all_gather in _mcp_apply → after fixing (1), non-sharded (default) LoRA at TP>1 crashes during profiling with a rank-dimension shape mismatch. From code reading this one is not XPU-specific (the gather lives in the platform-agnostic layer code), but I have only verified it on XPU.

Related context: #40601 (INC scheme refactor that owns inc_wna16_linear.py), #39778 (W4A16 / auto_round_kernel support), #41663 (oneCCL stable-fallback env used on this machine).

Environment

Component	Version
GPUs	2× Intel Arc Pro B70 (Battlemage G31), TP=2
CPU	AMD EPYC 7413
OS / kernel	Ubuntu 26.04 LTS, 7.0.0-27-generic, xe driver
PyTorch	2.12.0+xpu (XPU runtime 20250302, oneAPI 2025.3.3, Level Zero driver 26.18.38308.1)
vLLM	0.24.0+xpu, source build from upstream Dockerfile.xpu at v0.24.0 (ee0da84), unmodified
Model	Intel/Qwen3.6-27B-int4-AutoRound (W4A16 group_size=128, hybrid GDN+attention VLM arch, served text-only)
Relevant flags	--tensor-parallel-size 2 --enforce-eager --max-model-len 32768 --enable-lora --max-lora-rank 64 --max-loras 2

oneCCL env from #41663 applied (CCL_ENABLE_SYCL_KERNELS=0, CCL_ATL_TRANSPORT=ofi, CCL_ZE_IPC_EXCHANGE=sockets, CCL_TOPO_FABRIC_VERTEX_CONNECTION_CHECK=0, SYCL_UR_USE_LEVEL_ZERO_V2=0). The vision-tower "no matching PunicaWrapper … will be ignored" warnings are unrelated/benign.

---

Bug 1 — _get_lora_device: unsupported INC/AutoRound ark_linear base layer

Symptom: with --enable-lora, every worker dies while wrapping layers at load_model, container crash-loops:

  File ".../vllm/lora/layers/base_linear.py", line 80, in __init__
    self.device = _get_lora_device(self.base_layer)
  File ".../vllm/lora/layers/utils.py", line 70, in _get_lora_device
    raise ValueError(f"Unsupported base layer: {base_layer}")
ValueError: Unsupported base layer: MergedColumnParallelLinear(
  in_features=5120, output_features=8192, bias=False, tp_size=2, gather_output=False
  (ark_linear): QuantLinear(in_features=5120, out_features=8192, bias=True, w_bit=4, group_size=128)
)

Root cause: INCWNA16LinearScheme.process_weights_after_loading (vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py) repacks the quantized tensors into a layer.ark_linear submodule and then deletes layer.qweight / layer.qzeros / layer.scales from the parallel layer. _get_lora_device probes weight / weight_packed / qweight / w2_* — none of which exist anymore.

Fix (verified): add an ark_linear branch. QuantLinear.post_init() keeps qweight (repacked in place), so the device lookup is stable:

--- a/vllm/lora/layers/utils.py
+++ b/vllm/lora/layers/utils.py
@@ def _get_lora_device(base_layer: nn.Module) -> torch.device:
     # GPTQ/AWQ
     elif hasattr(base_layer, "qweight"):
         return base_layer.qweight.device
+    # AutoRound ark QuantLinear (INC WNA16 scheme): quantized tensors are
+    # repacked into the `ark_linear` submodule and the originals are deleted
+    # from the parallel layer after weight loading.
+    elif hasattr(base_layer, "ark_linear"):
+        return base_layer.ark_linear.qweight.device
     # MoE layer

With this patch, TP=1 boots and serves correctly with --enable-lora.

---

Bug 2 — unconditional all_gather in _mcp_apply breaks non-sharded LoRA at TP>1

Symptom: with Bug 1 patched, TP=2 with default (non-fully-sharded) LoRA crashes during determine_available_memory → profile_run (dummy LoRA run), first on a merged column layer:

  File ".../vllm/lora/layers/column_parallel_linear.py", line 344, in apply
    return _mcp_apply(x, bias, self)
  File ".../vllm/lora/layers/column_parallel_linear.py", line 66, in _mcp_apply
    lora_output = layer.punica_wrapper.add_expand(...)
  File ".../vllm/lora/punica_wrapper/punica_xpu.py", line 109, in _apply_expand
    bgmv_expand_slice(
  File ".../vllm/lora/ops/xpu_ops/lora_ops.py", line 79, in bgmv_expand_slice
    torch.ops._xpu_C.bgmv_expand_slice(
RuntimeError: inputs.size(1) must match lora_b_weights.size(-1)

Root cause: in _mcp_apply (vllm/lora/layers/column_parallel_linear.py, line 64 at v0.24.0):

local_lora_rank = layer.lora_a_stacked[0].shape[2]
buffer_shape = (layer.n_slices, x.shape[0], local_lora_rank)
...
buffers = tensor_model_parallel_all_gather(buffers)   # gathers along the LAST dim (rank)

The gather is only shape-consistent with --fully-sharded-loras, where lora_a is sharded along the rank dim (r/tp per rank → gather reconstitutes r, matching lora_b_weights.size(-1) == r). In the default non-sharded config, each rank already holds the full rank r, so the gather produces r * tp and the expand asserts (128 != 64 with max_lora_rank=64, TP=2). At TP=1 the gather is a no-op, which is why this only shows at TP>1. It also costs one collective per wrapped layer per step even when it happens to be benign.

Fix (verified):

--- a/vllm/lora/layers/column_parallel_linear.py
+++ b/vllm/lora/layers/column_parallel_linear.py
@@ def _mcp_apply(x, bias, layer):
-    buffers = tensor_model_parallel_all_gather(buffers)
+    # Only the fully-sharded path shards the LoRA rank dim across TP ranks;
+    # gathering in the non-sharded path breaks shapes (rank*tp vs rank) and
+    # costs one collective per wrapped layer per step.
+    if layer.lora_config.fully_sharded_loras:
+        buffers = tensor_model_parallel_all_gather(buffers)

--fully-sharded-loras also works as a workaround without any patch for Bug 2 (it makes the existing code path self-consistent), but see the performance note below.

---

Verification & measurements (TP=2, full production config)

• Boot: clean startup with --enable-lora in both modes (non-sharded + patch, and fully-sharded).
• Output integrity: 12/12 guided-JSON requests (llguidance backend, json_schema mode) valid, no mojibake/garbage tokens.
• Adapter math: rank-8 adapter with lora_B = 0 (exact no-op) on the 16 full-attention layers, loaded at runtime via POST /v1/load_lora_adapter (VLLM_ALLOW_RUNTIME_LORA_UPDATING=True): outputs at temperature 0 are bit-identical to the base model through the sliced TP=2 path.
• Throughput (engine-total, identical workload replayed against each configuration):

Configuration	Engine throughput	vs LoRA off
--enable-lora off (baseline)	~12.0 tok/s	—
non-sharded + Bug-2 patch	~10.5 tok/s	≈ −12%
--fully-sharded-loras (no patch)	~5–6.6 tok/s	≈ 2× slower

The fully-sharded slowdown is expected on this platform: it adds a rank-dim collective per wrapped layer per step, and with the #41663 oneCCL fallback env those collectives run on the host-staging path. Also, enabling LoRA reduced the available KV cache pool from 16.61 GiB to ~14.8 GiB.

Repro

1. Serve Intel/Qwen3.6-27B-int4-AutoRound on XPU with --enable-lora → Bug 1 (any TP).
2. Apply the _get_lora_device patch, serve with TP≥2, default LoRA sharding → Bug 2 during startup profiling (no adapter needs to be loaded; the dummy-LoRA profile run triggers it).

I can open a PR with both fixes (or split them) — let me know which shape you'd prefer.

[marcorigodanzo]

Closing — both bugs are now fixed upstream:

• Bug 2 (_mcp_apply unconditional all_gather, incompatible with non-sharded LoRA at TP>1) was fixed by [LoRA] Gate all_gather on fully_sharded_loras inside _mcp_apply; rewrite regression test #45715 (merged Jun 23) and ships in v0.25.1 (verified via compare API — the merge commit is an ancestor of the tag). Our v0.24.0 build predated the branch cut, which is why we hit it.
• Bug 1 (_get_lora_device failing on the AutoRound/INC ark_linear base layer) was fixed by [Bugfix][LoRA] Support ark_linear base layer in _get_lora_device #47690 (merged Jul 12). It is not in the 0.25 branch (cut before the merge), so it lands in the next minor.

We're retiring our downstream patches accordingly (the _mcp_apply one immediately with v0.25.1; the ark_linear one once the next minor ships). Thanks to the maintainers!