Fix lora buffer sizing for tp>num_kv cases

[opherlieber]

Motivation

• LoRA loading failed (or silently dropped V) at TP sizes that require KV-head replication. Fixes both the buffer allocation and the slice-time offset arithmetic so adapters load correctly at any TP.
• Concrete trigger: Qwen3.5-35B-A3B (num_kv_heads=2) at TP=8 errored with LoRA buffer shape [1152, 32] does not match weight shape [1280, 32]; at TP=4 it loaded but the V slice was empty (0 rows), silently producing wrong logprobs.

Modifications

• QKVParallelLinearWithLoRA.slice_lora_b_weights: compute q_size / k_size from the pre-replication head counts (base_layer.total_num_heads * head_size, base_layer.total_num_kv_heads * head_size) instead of output_sizes.
• New helper lora.utils.get_qkv_lora_kv_total(num_key_value_heads): returns the kv-head count the qkv LoRA buffer must reserve for after replication. Asserts attn_tp_size == tp_size since the rest of the LoRA path (esp. mem_pool.py) sizes by global tp_size — DP-attention / context-parallel are not yet supported with LoRA.
• Use get_qkv_lora_kv_total(...) in the qkv_proj branch of get_hidden_dim in lora/utils.py, models/nemotron_h.py, and models/qwen3_5.py.

Accuracy Tests

Speed Tests and Profiling

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 issues with LoRA weight slicing and buffer sizing when tensor parallelism size exceeds the number of KV heads. It introduces a utility function get_qkv_lora_kv_total to correctly calculate the replicated KV head count and updates the get_hidden_dim logic across several model implementations and utility layers to ensure consistent memory allocation and weight slicing. I have no feedback to provide.

[jybsuper]

Fixed in #24420