Grid construction based on num_active_lora and support CUDA graph capture across various num_active_lora

[yugong333]

Purpose

Reduce overhead of idle kernel launch due to max-loras in grid construction

Test Plan

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.

[gemini-code-assist]

Code Review

This pull request introduces an optimization for LoRA by making the CUDA grid size dependent on the number of active LoRAs rather than the maximum possible number. This is achieved by passing num_active_loras through various layers down to the Triton kernel. The changes are logical and well-implemented. I've identified one area for improvement concerning code duplication, which could impact maintainability and correctness.

[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]

Fused MoE grid skips LoRAs when batch mixes LoRA and no-LoRA tokens

The fused MoE shrink kernel now launches axis=2 programs based on num_active_loras, which is computed in lora_kernel_metadata.prepare_tensors as the count of non-negative IDs. When a batch mixes LoRA and non-LoRA tokens, active_lora_ids is sorted as [-1, <ids…>], so num_active_loras excludes the leading -1 but the grid here iterates only the first num_active_loras entries. The first iteration hits -1 and early-exits, and later positive IDs are never processed, so LoRA contributions are silently dropped for mixed batches. This miscomputes outputs for any MoE batch that contains both LoRA and non-LoRA tokens.

Useful? React with 👍 / 👎.

[jeejeelee]

Could you please provide the performance comparison?

[mergify]

Documentation preview: https://vllm--30984.org.readthedocs.build/en/30984/

[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]

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

[DarkLight1337]

Looks like you messed up the rebase and tagged everyone, can you create a new PR?