[LoRA] Support dual CUDA streams-Linear Layer

[jeejeelee]

Purpose

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.

[jeejeelee]

WIP. Triggering CI to test for any uncovered cases.

[gemini-code-assist]

Code Review

This pull request introduces support for dual CUDA streams to enable overlapping base layer and LoRA computations, which is a great performance optimization. The implementation correctly sets up a custom PyTorch operation and an auxiliary stream. However, I've found a critical issue in the asynchronous implementation that serializes the computations, defeating the purpose of using dual streams. My review includes a comment with a suggested fix for this issue. The other changes related to plumbing for this feature seem correct.

[gemini-code-assist]

The comment on line 212 is incorrect, and the wait_stream call on line 213 introduces a serialization point that prevents the intended overlap between the base layer and LoRA computations. The LoRA computation depends on the input x, not the output of the base layer. The wait on the current stream here forces the LoRA computation to wait until the base layer computation is complete, defeating the purpose of using a separate stream. Removing this wait will allow the two computations to run in parallel as intended.

[mergify]

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

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

[mergify]

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

uv pip install pre-commit>=4.5.1
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 failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

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

[mergify]

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

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

[jhaotingc]

Hi @jeejeelee, will you extend this to MoE LoRA layers in the future? Thanks!

[varun-sundar-rabindranath]

Thanks @jeejeelee . left some comments - generally looks good to me. Will take another look tomorrow.

[jeejeelee]

@jhaotingc Yeah, I will

[varun-sundar-rabindranath]

a lot of changes in the file appear to be linting - is it maybe a linter version mismatch ? can you check please. Thanks.

[varun-sundar-rabindranath]

LGTM ! Thanks @jeejeelee .

Left a comment on linting in envs.py. PTAL. Thanks 🙌