[LoRA] Fix qkv_proj LoRA buffer sizing when tp_size > num_key_value_heads

[gh1595]

Motivation

When QKVParallelLinear is sharded with tp_size > total_num_kv_heads,
KV heads are replicated across tp_size // num_kv_heads ranks rather
than further divided. The current LoRA path doesn't handle that layout
on either side of the per-rank slice, so loading any LoRA adapter that
targets qkv_proj fails with a buffer-shape assert at the first
forward.
Concrete repro on main:

python -m sglang.launch_server \
    --model-path Qwen/Qwen3.5-35B-A3B \
    --tp 4 --ep 4 --enable-lora --max-loras-per-batch 1 --max-lora-rank 32 \
    --lora-target-modules qkv_proj \
    --lora-paths my_lora=<some-qwen3.5-lora-checkpoint>

Error:

AssertionError: LoRA buffer shape torch.Size([576, 16]) does not match
weight shape torch.Size([768, 16])

Modifications

Two cooperating pieces have to agree on the per-rank QKV layout when KV heads are replicated:

LoRAMemoryPool.get_lora_B_shape (mem_pool.py) previously sized the per-rank qkv_proj LoRA B output as divide(total_output_dim, tp_size). That under-counts whenever tp_size > num_kv_heads, because the K/V part of total_output_dim is 2 * num_kv_heads * head_dim, which doesn't divide further by tp_size.

New helper _column_parallel_lora_b_per_rank_dim returns q_per_rank + 2 * head_dim for qkv_proj in the replicated case (each rank owns exactly head_dim of K and head_dim of V — one KV head replica), and falls back to the naive even split for every other column-parallel module and for the no-replication case (tp_size <= num_kv_heads). It also handles attn_output_gate=True (Qwen3.5) implicitly: gate-doubled q heads are encoded in total_output_dim, so the q remainder after subtracting K+V is correct without special-casing.

QKVParallelLinearWithLoRA.slice_lora_b_weights (layers.py) was indexing the PEFT-format (un-replicated) B tensor [q_total, k_total, v_total] with q_size, k_size, _ = base_layer.output_sizes. But base_layer.output_sizes[1] and [2] are the K/V replicated dims (per-rank K/V × tp_size), so they over-count K/V by num_kv_head_replicas. The V slice then falls past the end of B and silently returns 0 rows.

Fix: use output_sizes[0] for the q offset and output_sizes[1] // num_kv_head_replicas for the k offset, so indexing matches B's un-replicated layout.

The two changes together restore the previous buffer-shape contract (allocated dim == sliced dim) for the replicated case and leave the non-replicated case bit-identical.

Accuracy Tests

This is a load-time / shape correctness fix on a path that previously crashed before any forward, so existing accuracy tests for qkv_proj LoRA in the no-replication regime continue to pass unchanged.

The new unit test test/registered/unit/lora/test_qkv_lora_kv_replication.py exercises both pieces hermetically (no CUDA, no distributed init):

_column_parallel_lora_b_per_rank_dim for non-qkv modules, no-replication qkv, replicated qkv with and without attn_output_gate, varying replication factor, and the defensive fallback when num_key_value_heads is missing.
slice_lora_b_weights for the Qwen3.5-shaped (TP=16, 2 KV heads) case: per-rank shape (1024 rows = 512 q + 256 k + 256 v), KV replication groups sharing identical K/V slices, and an unchanged-behavior smoke test for the no-replication regime (tp <= num_kv_heads).
Run:

python -m pytest test/registered/unit/lora/test_qkv_lora_kv_replication.py -v
End-to-end: validated on Qwen3.5-35B-A3B with
--tp 4 --ep 4 --enable-lora --lora-target-modules all, which previously crashed at first request and now serves correctly.

Speed Tests and Profiling

No expected impact: the helper runs once per init_buffers, and slice_lora_b_weights does the same number of slices/concat as before — only the indices change. The non-replication path (the common case) is bit-identical to the previous implementation.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[gemini-code-assist]

Code Review

This pull request addresses an issue with QKV LoRA weight slicing and buffer sizing when tensor parallelism size exceeds the number of KV heads, leading to KV replication. Changes include updating the slicing logic in layers.py to use unreplicated K/V sizes and implementing a specialized per-rank dimension calculation in mem_pool.py. A comprehensive unit test suite is also added. Feedback indicates a potential AttributeError in LoRAMemoryPool due to the use of an uninitialized _text_config attribute; using base_hf_config with proper resolution is suggested instead.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[yushengsu-thu]

you can rerun the last case in stage c

[Fridge003]

We don't need this test

[Fridge003]

We can delete the comments here, since it has been explained in docstring

[Fridge003]

Delete the comments here