[TRTLLM-14867][feat] Auto-tuning for split-K in LoRA grouped GEMM - #17315
[TRTLLM-14867][feat] Auto-tuning for split-K in LoRA grouped GEMM#17315AlessioNetti wants to merge 20 commits into
Conversation
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>
Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
|
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:
WalkthroughThe change adds shared base-and-LoRA forwarding, optional CUDA-stream overlap, and CUDA-graph grouped-GEMM split-K autotuning. Model warmup coordinates autotuning, and tests cover runner behavior, synthetic inputs, and Qwen3 overlap execution. ChangesLoRA execution changes
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant ModelEngine
participant CudaGraphLoraManager
participant ProjectionModules
participant LoraLayer
participant LoraGroupedGemmRunner
ModelEngine->>CudaGraphLoraManager: prepare_cuda_graph_lora_params()
CudaGraphLoraManager-->>ModelEngine: return LoRA parameters and auxiliary stream
ProjectionModules->>LoraLayer: call forward_with_base(...)
LoraLayer->>ProjectionModules: execute base projection
LoraLayer->>LoraGroupedGemmRunner: autotune split-K and run grouped GEMM
LoraGroupedGemmRunner-->>LoraLayer: return LoRA output
LoraLayer-->>ProjectionModules: accumulate base and LoRA outputs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
tests/unittest/_torch/lora/test_lora_autotuner.py (1)
35-193: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd required function annotations.
New test functions and local test doubles omit parameter or return annotations.
tests/unittest/_torch/lora/test_lora_autotuner.py#L35-L193: annotate test functions and local fake methods.tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py#L152-L212: annotate_run_lora_testand the new test functions.As per coding guidelines,
**/*.pyrequires: “Annotate every function.”🤖 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 `@tests/unittest/_torch/lora/test_lora_autotuner.py` around lines 35 - 193, Add type annotations to every function introduced in these LoRA test files: update the test functions in tests/unittest/_torch/lora/test_lora_autotuner.py and the helper/test functions in tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py, including local fake methods like FakeLayerParams.__init__, FakeCudaGraphParams.__init__/get_layer_params/get_problem_count, FakeTuner.choose_one, and fake_forward_impl. Keep the existing test behavior unchanged while adding parameter and return annotations consistently across the affected symbols.Source: Coding guidelines
tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py (1)
32-32: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAlign the constructor default with the config default.
LoraConfig.overlap_lora_and_basedefaults toFalseintensorrt_llm/lora_helper.py, but this parameter defaults toTrue. The only caller in this cohort passes the value explicitly, so behavior is correct today. A future caller that omits the argument would silently enable the auxiliary stream. UseFalsehere so both defaults agree.♻️ Proposed default alignment
- overlap_lora_and_base: bool = True, + overlap_lora_and_base: bool = False,🤖 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/cuda_graph_lora_manager.py` at line 32, Update the constructor parameter overlap_lora_and_base in the LoRA manager to default to False, matching LoraConfig.overlap_lora_and_base and preserving explicitly supplied values.tensorrt_llm/_torch/peft/lora/layer.py (1)
858-891: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winSkip the redundant
copy_lora_paramson the inference path.
_forward_cuda_graph_modealready assignsrunner.lora_params = runner.copy_lora_params(lora_params)(Line 624).forwardthen copies the same structure again at Line 860. The second copy is only required whenlen(inputs) > 1, because that branch rebinds attributes oncuda_graph_paramsandlayer_params.On the inference path this runs once per LoRA layer per decode step and produces two dict allocations, two
CudaGraphLoraParamsshallow copies, and twoLoraLayerParamscopies that are immediately discarded. Copy only when the synthetic inputs are present.♻️ Proposed refactor
del kwargs assert self.lora_params is not None - lora_params = self.copy_lora_params(self.lora_params) - cuda_graph_params = lora_params['cuda_graph_params'] - layer_params = cuda_graph_params.get_layer_params(self.layer_key) - assert layer_params is not None # The list of tensor input arguments is re-packed into # the local lora_params copy such that we can invoke # LoraLayer's _forward_cuda_graph_mode_impl(). x = inputs[0] - if len(inputs) > 1: + if len(inputs) == 1: + # Inference path: reuse the copy made by _forward_cuda_graph_mode. + lora_params = self.lora_params + else: + # Tuning path: the synthetic tensors below overwrite fields on the + # copy, so isolate it from the caller's parameters. + lora_params = self.copy_lora_params(self.lora_params) + cuda_graph_params = lora_params['cuda_graph_params'] + layer_params = cuda_graph_params.get_layer_params(self.layer_key) + assert layer_params is not None ( x, cuda_graph_params.slot_counts,🤖 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 858 - 891, Update the forward path around _forward_cuda_graph_mode_impl to avoid copying self.lora_params when len(inputs) is 1, since the runner has already copied it. Reuse self.lora_params directly for the normal inference path, and only call copy_lora_params when len(inputs) > 1 before rebinding cuda_graph_params and layer_params fields. Preserve the existing parameter setup and output behavior.
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/peft/lora/layer.py`:
- Around line 612-621: Update the _split_k_runner initialization in the relevant
CUDA-graph path to pass x.shape[-1] as input_hidden_size, matching
_forward_cuda_graph_mode_impl and preserving the correct hidden dimension for
both 2-D and 3-D inputs.
In `@tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py`:
- Around line 196-211: Strengthen the overlap coverage in
test_qwen3_bf16_lora_overlap and test_qwen3_fp8_lora_overlap by adding focused
assertions around LoraLayer.forward_with_base that confirm an auxiliary stream
is available and the parallel executor is entered. Ensure the tests fail when
overlap_lora_and_base is ignored or execution falls back to the serial path,
while preserving the existing output checks.
---
Nitpick comments:
In `@tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py`:
- Line 32: Update the constructor parameter overlap_lora_and_base in the LoRA
manager to default to False, matching LoraConfig.overlap_lora_and_base and
preserving explicitly supplied values.
In `@tensorrt_llm/_torch/peft/lora/layer.py`:
- Around line 858-891: Update the forward path around
_forward_cuda_graph_mode_impl to avoid copying self.lora_params when len(inputs)
is 1, since the runner has already copied it. Reuse self.lora_params directly
for the normal inference path, and only call copy_lora_params when len(inputs) >
1 before rebinding cuda_graph_params and layer_params fields. Preserve the
existing parameter setup and output behavior.
In `@tests/unittest/_torch/lora/test_lora_autotuner.py`:
- Around line 35-193: Add type annotations to every function introduced in these
LoRA test files: update the test functions in
tests/unittest/_torch/lora/test_lora_autotuner.py and the helper/test functions
in tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py,
including local fake methods like FakeLayerParams.__init__,
FakeCudaGraphParams.__init__/get_layer_params/get_problem_count,
FakeTuner.choose_one, and fake_forward_impl. Keep the existing test behavior
unchanged while adding parameter and return annotations consistently across the
affected symbols.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fe310c79-065d-4aa2-ab96-4b8bdaea9c45
📒 Files selected for processing (11)
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.jsontests/unittest/_torch/lora/test_lora_autotuner.pytests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py
…alse Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
1bf46c2 to
77550d2
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
♻️ Duplicate comments (1)
tensorrt_llm/_torch/peft/lora/layer.py (1)
612-621: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winUse
x.shape[-1]forinput_hidden_size.Line 616 reads the hidden size as
x.shape[1], but_forward_cuda_graph_mode_implreads it asx.shape[-1]at Line 521. The two values differ for a rank-3 input, which would give wrong synthetic LoRA A shapes and a wrongunique_id.🤖 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 612 - 621, Update the _split_k_runner initialization in _forward_cuda_graph_mode_impl to pass x.shape[-1] as input_hidden_size, matching the existing hidden-size calculation and preserving correct LoRA shapes and unique_id generation for rank-3 inputs.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/peft/lora/layer.py (1)
623-632: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAvoid both LoRA-parameter copies outside tuning. Assign
runner.lora_params = lora_paramsdirectly, and callcopy_lora_paramsin_LoraGroupedGemmRunner.forwardonly whenAutoTuner.get().is_tuning_modeis set. The non-tuning path only reads the parameters and uses[x], so this removes two hierarchy copies from each cached decode call while preserving tuning isolation.🤖 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 623 - 632, Update the split-k execution path around _split_k_runner to assign lora_params directly instead of calling copy_lora_params. In _LoraGroupedGemmRunner.forward, perform copy_lora_params only when AutoTuner.get().is_tuning_mode is enabled, while retaining direct parameter reads and the [x] input in non-tuning execution.
🤖 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.
Duplicate comments:
In `@tensorrt_llm/_torch/peft/lora/layer.py`:
- Around line 612-621: Update the _split_k_runner initialization in
_forward_cuda_graph_mode_impl to pass x.shape[-1] as input_hidden_size, matching
the existing hidden-size calculation and preserving correct LoRA shapes and
unique_id generation for rank-3 inputs.
---
Nitpick comments:
In `@tensorrt_llm/_torch/peft/lora/layer.py`:
- Around line 623-632: Update the split-k execution path around _split_k_runner
to assign lora_params directly instead of calling copy_lora_params. In
_LoraGroupedGemmRunner.forward, perform copy_lora_params only when
AutoTuner.get().is_tuning_mode is enabled, while retaining direct parameter
reads and the [x] input in non-tuning execution.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8aae3487-5ded-4a1d-b685-03c0c38de695
📒 Files selected for processing (11)
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.jsontests/unittest/_torch/lora/test_lora_autotuner.pytests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py
🚧 Files skipped from review as they are similar to previous changes (9)
- tensorrt_llm/lora_helper.py
- tensorrt_llm/_torch/modules/mlp.py
- tensorrt_llm/_torch/modules/linear.py
- tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py
- tensorrt_llm/_torch/modules/attention.py
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tensorrt_llm/_torch/pyexecutor/model_engine.py
- tensorrt_llm/_torch/modules/gated_mlp.py
- tests/unittest/_torch/lora/test_lora_autotuner.py
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>
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>
2c4f9d7 to
3a058df
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (5)
tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py (2)
232-247: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert that overlap preserves the non-overlap result.
Both new tests only call
_assert_lora_changes_output, which compares the LoRA output against the base output. Overlap execution changes only scheduling, so the LoRA output must match the serial LoRA output token for token. The current assertions pass even if overlap produces a different, incorrect result, which is the failure mode a stream-synchronization defect produces.Add an equivalence check between
overlap=Trueandoverlap=Falsefor the same dtype and adapter. Sampling already usestemperature=0.0, so the comparison is deterministic.The four
test_qwen3_*_lora*methods now differ only bydtypeandoverlap. Consider@pytest.mark.parametrize("dtype, overlap", ...)over a single body to remove the duplication.🤖 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 `@tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py` around lines 232 - 247, Update the Qwen3 LoRA tests so overlap=True outputs are compared token-for-token with the corresponding overlap=False result for the same adapter and dtype, while retaining the base-output LoRA assertion. Consolidate the four test_qwen3_*_lora* methods into one parametrized test over dtype and overlap where practical, using the existing deterministic sampling configuration.
178-247: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd output-equivalence and split-K coverage.
- Added:
test_lora_forward_with_base_executes_overlap,TestQwen3LoRA.test_qwen3_bf16_lora_overlap, andTestQwen3LoRA.test_qwen3_fp8_lora_overlap.- Modified:
_run_lora_testgainedoverlap.- Removed: none.
- Existing directory bridges cover this file in
l0_b200.yml,l0_gb300_multi_gpus.yml,l0_b300.yml, andl0_h100.yml. No QA entry is required for this unit-test bridge.- Gaps remain: no comparison proves overlap and non-overlap produce identical output, and no end-to-end test covers split-K autotuning. The autotuner tests use fakes.
- Verdict: needs follow-up because no CBTS coverage artifact is available.
🤖 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 `@tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py` around lines 178 - 247, Extend the LoRA test coverage around _run_lora_test and TestQwen3LoRA to compare overlap and non-overlap outputs for equivalent Qwen3 configurations, including relevant BF16 and FP8 paths, and add an end-to-end split-K autotuning test using the real execution path rather than fakes. Preserve the existing CUDA skip and model-availability setup while ensuring both coverage gaps produce verifiable assertions.Source: Path instructions
tensorrt_llm/_torch/peft/lora/layer.py (2)
623-632: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAvoid the per-forward
copy_lora_paramscall on the inference path.Line 624 copies the LoRA parameter hierarchy on every forward. The copy is only required when the autotuner re-packs synthetic inputs.
forwardalready callscopy_lora_paramsitself at Line 865 whenlen(inputs) > 1, so the inference path never mutates the dict it receives. The outer copy therefore allocates a dict, aCudaGraphLoraParamsshallow copy, and aLoraLayerParamscopy per LoRA layer per forward with no isolation benefit.Assign the parameters directly and let
forwardcopy only when it tunes.♻️ Proposed change
runner = self._split_k_runner - runner.lora_params = runner.copy_lora_params(lora_params) + runner.lora_params = lora_params runner_inputs = [x]Note that
tests/unittest/_torch/lora/test_lora_autotuner.pyLines 124 and 129-132 assert the outer copy happens. Update those assertions if you apply this change.🤖 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 623 - 632, Remove the per-forward runner.copy_lora_params call in the _split_k_runner path and assign lora_params directly to runner.lora_params before autotuning. Keep forward’s existing copy_lora_params behavior for the len(inputs) > 1 tuning case, and update the corresponding test assertions in test_lora_autotuner.py to expect no outer copy.
828-838: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBind the synthetic input layout to a named schema.
_prepare_synthetic_inputsbuilds a nine-element positional list andforwardunpacks the same nine positions at Lines 869-880. The two sites share no constant, so a reordering in one place silently misassigns tensors in the other.tests/unittest/_torch/lora/test_lora_autotuner.pyalso indexes these positions by literal integers at Lines 140-141 and 174-177, which multiplies the coupling.Define the field order once, for example as a module-level tuple of attribute paths, and drive both the pack and the unpack from it.
🤖 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 828 - 838, Define a module-level named schema for the synthetic input fields used by _prepare_synthetic_inputs, forward, and the autotuner tests, then build and consume the positional list through that schema instead of duplicated ordering and literal indexes. Preserve the existing nine-field order and update all references consistently so reordering the schema cannot silently misassign inputs.tests/unittest/_torch/lora/test_lora_autotuner.py (1)
39-49: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest the bucket mapping through
DynamicTensorSpec. Replace the privateAutoTuner._find_nearest_profilecall and cache reset withassert spec.map_to_tuning_buckets(7) == 4. This keeps the test focused on token-bucket behavior without coupling it to AutoTuner internals.🤖 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 `@tests/unittest/_torch/lora/test_lora_autotuner.py` around lines 39 - 49, Replace the private AutoTuner._find_nearest_profile call and its cache_clear setup with an assertion that spec.map_to_tuning_buckets(7) returns 4. Keep the existing spec.input_idx and spec.dim_idx assertions, and remove the profile-based assertion.
🤖 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.
Inline comments:
In `@tests/unittest/_torch/lora/test_lora_autotuner.py`:
- Around line 63-68: Update FakeLayerParams.__init__ so d_b_ptrs and
d_b_prime_ptrs are initialized as 2-D tensors with shape (1, 4), matching
CudaGraphLoraParams._allocate_layer_params and allowing indexed writes in
_prepare_synthetic_inputs.
- Around line 13-32: Add coverage in the existing LoRA autotuner tests around
_make_runner and the relevant runner invocation for tactic=-1, asserting it
resolves to _LORA_DEFAULT_SPLIT_K, and update the CUDA-graph warmup test to
observe the split-K argument passed to
torch.ops.trtllm.lora_grouped_gemm_cuda_graph instead of only patching
_forward_cuda_graph_mode_impl. Preserve the current warmup and active-slot
assertions while validating both operator-boundary behaviors.
---
Nitpick comments:
In `@tensorrt_llm/_torch/peft/lora/layer.py`:
- Around line 623-632: Remove the per-forward runner.copy_lora_params call in
the _split_k_runner path and assign lora_params directly to runner.lora_params
before autotuning. Keep forward’s existing copy_lora_params behavior for the
len(inputs) > 1 tuning case, and update the corresponding test assertions in
test_lora_autotuner.py to expect no outer copy.
- Around line 828-838: Define a module-level named schema for the synthetic
input fields used by _prepare_synthetic_inputs, forward, and the autotuner
tests, then build and consume the positional list through that schema instead of
duplicated ordering and literal indexes. Preserve the existing nine-field order
and update all references consistently so reordering the schema cannot silently
misassign inputs.
In `@tests/unittest/_torch/lora/test_lora_autotuner.py`:
- Around line 39-49: Replace the private AutoTuner._find_nearest_profile call
and its cache_clear setup with an assertion that spec.map_to_tuning_buckets(7)
returns 4. Keep the existing spec.input_idx and spec.dim_idx assertions, and
remove the profile-based assertion.
In `@tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py`:
- Around line 232-247: Update the Qwen3 LoRA tests so overlap=True outputs are
compared token-for-token with the corresponding overlap=False result for the
same adapter and dtype, while retaining the base-output LoRA assertion.
Consolidate the four test_qwen3_*_lora* methods into one parametrized test over
dtype and overlap where practical, using the existing deterministic sampling
configuration.
- Around line 178-247: Extend the LoRA test coverage around _run_lora_test and
TestQwen3LoRA to compare overlap and non-overlap outputs for equivalent Qwen3
configurations, including relevant BF16 and FP8 paths, and add an end-to-end
split-K autotuning test using the real execution path rather than fakes.
Preserve the existing CUDA skip and model-availability setup while ensuring both
coverage gaps produce verifiable assertions.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d6ad3521-1a88-4ca2-982d-6f887f648f54
📒 Files selected for processing (11)
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.jsontests/unittest/_torch/lora/test_lora_autotuner.pytests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py
🚧 Files skipped from review as they are similar to previous changes (8)
- tensorrt_llm/_torch/modules/gated_mlp.py
- tensorrt_llm/lora_helper.py
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tensorrt_llm/_torch/modules/mlp.py
- tensorrt_llm/_torch/modules/attention.py
- tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py
- tensorrt_llm/_torch/pyexecutor/model_engine.py
- tensorrt_llm/_torch/modules/linear.py
… to grouped GEMM properly Signed-off-by: Alessio Netti <26897207+AlessioNetti@users.noreply.github.com>
2 participants
Summary
_LoraGroupedGemmRunnerandModelEngine.maybe_autotune_lora().Dev Engineer Review
layer_idxand autotuner inputs.QA Engineer Review
Added or modified tests for:
No
tests/integration/test_lists/,test-db/, orqa/coverage entries are identified for these tests. CI or manual-QA coverage requires verification.Verdict: needs follow-up
Description
NOTE: this PR is based on and should be merged only after 16951.
As of TRT-LLM
1.3.0rc23, the grouped GEMM kernel used for LoRA (intensorrt_llm/_torch/peft/lora/layer.py) uses a hardcoded split-K value of16, which might not be ideal for all workloads. This PR addresses this by introducing auto-tuning for the split-K value, allowing to find and use optimal split-K values depending on the LoRA module, hidden size as well as GPU type.This was achieved by introducing a new
_LoraGroupedGemmRunner, which encapsulates all logic required in order to auto-tuneLoraLayer._forward_cuda_graph_mode(). Small changes were required inmodel_engine.pyas well, in order to perform one more auto-tuning pass during CUDA graph warmup (as the grouped GEMM kernel in question is only used in the CUDA graph path).Below you can find a latency comparison, benchmarked using a Qwen3 32B FP8 model, 256 ISL/OSL and a rank 32 LoRA adapter (gate_up_proj, down_proj, qkv_proj and o_proj modules), on H100. We are also showing performance achieved as of PR 16951, upon which this PR is based, as a reference point.
Test Coverage
Any test that uses LoRA with CUDA graphs and auto-tuning enabled, such as
tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py. Additional auto-tuner focused tests were added intests/unittest/_torch/lora/test_lora_autotuner.py.