[TRTLLM-14692][feat] Overlap LoRA and base model computations - #16951
[TRTLLM-14692][feat] Overlap LoRA and base model computations#16951AlessioNetti wants to merge 8 commits into
Conversation
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/peft/lora/layer.py (2)
177-228: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftAdd overlap parity coverage.
Add CUDA-graph tests that compare sequential and auxiliary-stream outputs, verify
overlap_lora_and_base=False, and cover the supported Attention, MLP, GatedMLP, and Linear paths.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/peft/lora/layer.py` around lines 177 - 228, Add CUDA-graph test coverage for LoraLayer.forward_with_base that compares sequential execution with auxiliary-stream overlap outputs. Parameterize or otherwise cover the supported Attention, MLP, GatedMLP, and Linear paths, and explicitly verify that overlap_lora_and_base=False uses the sequential path while preserving matching outputs.
164-201: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse concrete Python 3.10 types for the new helper.
Any,Dict/List/Optional, and the unannotated nested callable obscure the tensor-onlyadd_()contract. Use concrete tensor return types, built-in generics/|, and a precise LoRA-params type. As per coding guidelines, “Annotate every function, useNonefor non-returning functions, avoidAny… [and] prefer built-in generic types and|.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/peft/lora/layer.py` around lines 164 - 201, Update LoraLayer.forward_with_base and its nested lora_forward helper to use concrete tensor-based annotations instead of Any and an unannotated callable. Replace List/Dict/Optional and tuple typing with Python 3.10 built-in generics and |, and use the precise existing LoRA-params type for lora_params. Annotate every callable, including non-returning paths with None where applicable, while preserving the current tensor/tuple-of-tensors return behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/_torch/peft/lora/layer.py`:
- Around line 177-228: Add CUDA-graph test coverage for
LoraLayer.forward_with_base that compares sequential execution with
auxiliary-stream overlap outputs. Parameterize or otherwise cover the supported
Attention, MLP, GatedMLP, and Linear paths, and explicitly verify that
overlap_lora_and_base=False uses the sequential path while preserving matching
outputs.
- Around line 164-201: Update LoraLayer.forward_with_base and its nested
lora_forward helper to use concrete tensor-based annotations instead of Any and
an unannotated callable. Replace List/Dict/Optional and tuple typing with Python
3.10 built-in generics and |, and use the precise existing LoRA-params type for
lora_params. Annotate every callable, including non-returning paths with None
where applicable, while preserving the current tensor/tuple-of-tensors return
behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 83223af9-06d8-4b0f-b4ed-4b4985b44fe6
📒 Files selected for processing (9)
tensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/gated_mlp.pytensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/modules/mlp.pytensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.pytensorrt_llm/_torch/peft/lora/layer.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/lora_helper.pytensorrt_llm/usage/llm_args_golden_manifest.json
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughLoRA configuration now controls optional overlap between base and LoRA computation. Projection paths use a shared helper for execution, stream scheduling, and output merging. ChangesLoRA overlap execution
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant AttentionOrMLP
participant LoraLayer
participant QuantMethod
participant AuxStream
participant LoRAProjection
AttentionOrMLP->>LoraLayer: call forward_with_base
LoraLayer->>QuantMethod: execute base projection
LoraLayer->>AuxStream: optionally schedule LoRA computation
AuxStream->>LoRAProjection: compute LoRA output
LoraLayer->>LoraLayer: merge LoRA output into base output
LoraLayer-->>AttentionOrMLP: return combined output
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
48861ee to
dea6573
Compare
Funatiq
added
the
api-compatible
|
/bot run |
|
PR_Github #62207 [ run ] triggered by Bot. Commit: |
|
PR_Github #62207 [ run ] completed with state
|
dea6573 to
9b59e76
Compare
|
/bot run |
|
PR_Github #62511 [ run ] triggered by Bot. Commit: |
|
PR_Github #62511 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62533 [ run ] triggered by Bot. Commit: |
|
PR_Github #62533 [ run ] completed with state
|
9b59e76 to
1904bcf
Compare
|
PR_Github #63080 [ run ] triggered by Bot. Commit: |
|
PR_Github #63080 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63389 [ run ] triggered by Bot. Commit: |
|
PR_Github #63389 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63434 [ run ] triggered by Bot. Commit: |
|
PR_Github #63434 [ run ] completed with state |
|
/bot run |
|
PR_Github #63440 [ run ] triggered by Bot. Commit: |
|
PR_Github #63440 [ run ] completed with state
|
aee41f8 to
44578d2
Compare
|
/bot run |
|
PR_Github #63461 [ run ] triggered by Bot. Commit: |
|
PR_Github #63461 [ run ] completed with state
|
Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
…o empty_like() for init Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
44578d2 to
93acdfb
Compare
|
/bot run |
|
PR_Github #63491 [ run ] triggered by Bot. Commit: |
|
PR_Github #63491 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63532 [ run ] triggered by Bot. Commit: |
|
PR_Github #63532 [ run ] completed with state
|
4 participants
Dev Engineer Review
LoraConfig.overlap_lora_and_base: bool = Falseto control whether CUDA-graph LoRA execution overlaps with base projection computation via a secondary CUDA stream.CudaGraphLoraManagerandPyTorchModelEnginewiring to propagateoverlap_lora_and_baseinto CUDA-graph LoRA parameters.LoraLayer.forward_with_base(...)to optionally run base forward and a tuple of LoRA layers in parallel (aux stream + CUDA events) when CUDA-graph-compatible conditions are met (do_multi_stream(), not under torch compile, and LoRA entries exist forlayer_idx); otherwise it falls back to sequential execution.add_on the base output buffer (when LoRA outputs are non-None).torch.empty_like(...)instead oftorch.zeros_like(...)prior toindex_copy_.LoraLayer.forward_with_base(...)for LoRA-enabled paths in:attention(QKV)mlp(up projection)gated_mlp(gate/up projection)linear(GEMM path)TorchLlmArgsgolden manifest schema to includelora_config.overlap_lora_and_base.QA Engineer Review
No test changes.
Description
As of TRT-LLM
1.3.0rc22, LoRA computations are enqueued in the same default CUDA stream as the base model, which results in sequential execution. In an effort to improve performance, this PR changes this by introducing a secondary CUDA stream just for LoRA operations, which results in base model projections happening in parallel with LoRA, hiding most of the latency of the latter.Below you can find a latency comparison for the new changes, benchmarked using a Qwen3 32B FP8 model, 256 ISL/OSL and a rank 32 LoRA adapter (
gate_up_proj,down_proj,qkv_projando_projmodules), on H100. It can be seen that the changes improve latency significantly, especially in low-concurrency regimes (~1s or ~15% latency reduction at concurrency 1).The new behavior can be disabled via the new
overlap_lora_and_baseflag added toLoraConfig- this is disabled by default. Some caveats regarding the changes:add_to sum LoRA outputs to base projection outputs, and switching fromzeros_liketoempty_likewhen allocating the LoRA restored output buffer;Test Coverage
Focused tests were added under
tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py, in theTestQwen3LoRAclass.