Reduce the kernel overhead when num of active loras is smaller than max loras. Multiple cuda graphs are captured for each num of active-loras.

[yugong333]

Purpose

Test Plan

Using llmperf to benchmark concurrency = 1, 2, 4, 8 when max-loras = 4

Test Result

---

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.

---

Note

^{Cursor Bugbot is generating a summary for commit bb040594f60f02a5a5fd8b8bf3e9630da19d47b7. Configure here.}

---

Note

Introduces specialization by active LoRA count and plumbs it end-to-end to reduce kernel/grid overhead when fewer adapters are used.

• Adds LoRAConfig.specialize_active_lora and --specialize-active-lora CLI flag; stores num_active_loras in kernel metadata and BatchDescriptor
• cudagraph_dispatcher and runner capture/dispatch multiple CUDA graphs for LoRA counts (powers of 2 up to max_loras), with fallback to nearest captured count
• Updates Triton ops (lora_shrink, lora_expand, fused_moe_lora) to accept num_active_loras and optionally size the launch grid by it; propagate via Punica wrapper and fake ops
• Adjusts dummy-run LoRA selection to activate an exact number of adapters for capture; API wiring across v1 worker and forward context

^{Written by Cursor Bugbot for commit bb040594f60f02a5a5fd8b8bf3e9630da19d47b7. This will update automatically on new commits. Configure here.}

---

Note

^{Cursor Bugbot is generating a summary for commit cfb1a4e9f4c83820ea67821ade4e5e5dcaee9611. Configure here.}

---

Note

Introduces LoRA-count specialization to reduce kernel/grid overhead when fewer adapters are active.

• Adds LoRAConfig.specialize_active_lora and --specialize-active-lora; stores num_active_loras in BatchDescriptor and LoRA kernel metadata
• CUDAGraph capture/dispatch updated to key on (has_lora, num_active_loras); captures graphs for powers-of-two counts up to max_loras+1 and dispatches to the nearest captured count
• Triton ops (lora_shrink, lora_expand, fused_moe_lora) accept num_active_loras and optionally size the launch grid by it; fake ops and Punica wrapper updated accordingly
• LoRAKernelMeta now tracks num_active_loras and optional captured_lora_counts for rounding during cudagraph use
• Runner/dispatcher changes propagate active LoRA count through dummy-run capture paths and live dispatch

^{Written by Cursor Bugbot for commit cfb1a4e9f4c83820ea67821ade4e5e5dcaee9611. This will update automatically on new commits. Configure here.}

[gemini-code-assist]

Code Review

This pull request introduces an optimization for LoRA serving by specializing CUDA graphs for different numbers of active LoRA adapters, controlled by a new specialize_active_lora flag. The changes are well-structured and span configuration, CUDA graph dispatching logic, LoRA kernel metadata, and Triton kernels to support this specialization. The implementation is mostly correct and consistent. However, I've identified a critical issue in the CUDA graph dispatching logic that prevents the feature from working as intended. The dispatcher doesn't correctly map the runtime number of active LoRAs to the captured graphs. I've provided a detailed comment and a code suggestion to fix this.

[jeejeelee]

@yugong333 Could you please check the bot comments first?

[yugong333]

@yugong333 Could you please check the bot comments first?

Hi @jeejeelee I have fixed the bugs. Thanks for your review!

[jeejeelee]

@ProExpertProg Could you please take a look at this PR? Thank you!

[ProExpertProg]

I am not that familiar with LoRA but the cudagraph changes look good to me

[varun-sundar-rabindranath]

doesn't captured_lora_counts need a reset ?

[dcmaddix]

@yugong333

[yugong333]

Thanks! I have added reset here: https://github.com/vllm-project/vllm/pull/32005/changes#diff-d593fbedbb939694930a23db4958ed49c135b9511d0ed9b054b65e8ba06705c3R94

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @yugong333.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @yugong333, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @yugong333, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint