[LoRA] Lora shrink swizzle

[li2haipeng]

Purpose

Apply CTA swizzle in lora_shrink Triton kernel to increase L2 cache utilization. It's mostly helpful for TTFT.

We also test the similar approach on lora_expand kernels and only achieved marginal improvement even after tuning. Thus, we only modify lora_shrink kernel in this PR.

Test Result

Server:

export LORA_PATH=Romoamigo/Qwen3-32B-LoRA-adapters
vllm serve Qwen/Qwen3-32B \
    --trust-remote-code \
    -tp 4 \
    --port 8000 \
    --compilation_config='{"cudagraph_capture_sizes": [1,4,8,16], "compile_sizes": [1,4,8,16]}' \
    --enable-lora \
    --max-loras 1 \
    --max-lora-rank 16 \
    --lora-modules  \
        adapter0=$LORA_PATH \
        adapter1=$LORA_PATH \
        adapter2=$LORA_PATH \
        adapter3=$LORA_PATH \
        adapter4=$LORA_PATH \
        adapter5=$LORA_PATH \
        adapter6=$LORA_PATH \
        adapter7=$LORA_PATH \
        adapter8=$LORA_PATH \
        adapter9=$LORA_PATH \
        adapter10=$LORA_PATH \
        adapter11=$LORA_PATH \
        adapter12=$LORA_PATH \
        adapter13=$LORA_PATH \
        adapter14=$LORA_PATH \
        adapter15=$LORA_PATH \

benchmark:

vllm bench serve \
  --model /home/ubuntu/models/Qwen3-32B \
  --lora-modules adapter0 adapter1 adapter2 adapter3 adapter4 adapter5 adapter6 adapter7 adapter8 \
                adapter9 adapter10 adapter11 adapter12 adapter13 adapter14 adapter15 \
  --endpoint /v1/completions \
  --dataset-name sharegpt \
  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
  --max-concurrency 8 \
  --num-prompts 80

Without this PR

============ Serving Benchmark Result ============
Successful requests:                     80        
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  171.52    
Total input tokens:                      19956     
Total generated tokens:                  16766     
Request throughput (req/s):              0.47      
Output token throughput (tok/s):         97.75     
Peak output token throughput (tok/s):    279.00    
Peak concurrent requests:                10.00     
Total Token throughput (tok/s):          214.10    
---------------Time to First Token----------------
Mean TTFT (ms):                          14034.73  
Median TTFT (ms):                        6173.46   
P99 TTFT (ms):                           73440.25  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          11.23     
Median TPOT (ms):                        10.72     
P99 TPOT (ms):                           22.80     
---------------Inter-token Latency----------------
Mean ITL (ms):                           10.77     
Median ITL (ms):                         10.73     
P99 ITL (ms):                            11.10     
==================================================

With this PR (default kernel config with GROUP_SIZE_M=8)

============ Serving Benchmark Result ============
Successful requests:                     80        
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  166.05    
Total input tokens:                      19956     
Total generated tokens:                  16773     
Request throughput (req/s):              0.48      
Output token throughput (tok/s):         101.01    
Peak output token throughput (tok/s):    279.00    
Peak concurrent requests:                11.00     
Total Token throughput (tok/s):          221.19    
---------------Time to First Token----------------
Mean TTFT (ms):                          13536.22  
Median TTFT (ms):                        5056.81   
P99 TTFT (ms):                           66658.13  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          10.62     
Median TPOT (ms):                        10.60     
P99 TPOT (ms):                           11.04     
---------------Inter-token Latency----------------
Mean ITL (ms):                           10.58     
Median ITL (ms):                         10.57     
P99 ITL (ms):                            10.92     
==================================================

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Guard missing group_size_m in tuned configs

The new swizzle passes GROUP_SIZE_M from kernel_config['group_size_m'], but get_lora_op_configs returns the raw dictionary from any tuned JSON without augmenting it. Existing tuned config files shipped before this change don’t contain group_size_m, so the lookup raises a KeyError as soon as someone sets VLLM_TUNED_CONFIG_FOLDER for LoRA shrink kernels. Consider using kernel_config.get('group_size_m', 8) or filling the field when loading tuned configs so old configuration files remain usable.

Useful? React with 👍 / 👎.

[jeejeelee]

This comment looks reasonable

[gemini-code-assist]

Code Review

This pull request introduces a CTA swizzling optimization to the lora_shrink Triton kernel. The goal is to improve L2 cache utilization and thereby reduce Time to First Token (TTFT), which is supported by the benchmark results provided. The implementation changes the CTA scheduling from a simple column-major ordering to a grouped column-major ordering. My analysis of the swizzling logic confirms that it correctly remaps the thread block identifiers to achieve the intended access pattern and is robust against various input dimensions, including edge cases. The changes are well-contained, and the new configuration parameter group_size_m is added with a reasonable default value. Overall, this is a solid performance enhancement.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Handle missing group_size_m in tuned LoRA configs

_lora_shrink now unconditionally reads kernel_config["group_size_m"], but get_lora_op_configs returns raw entries from any tuned JSON config and older tuned files do not define this new key. When such a config is present (a supported workflow before this change) this access raises a KeyError, aborting LoRA serving before the Triton kernel launches. The code should provide a default (e.g. 8) or gracefully fall back when the key is absent so existing tuned configs continue to work.

Useful? React with 👍 / 👎.

[jeejeelee]

LGTM after respondsing to the above comment and fixing the pre-commit failures