[None][feat] Support dense FP8 LoRA end to end - #16810
Conversation
9db2adf to
c47972f
Compare
|
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 changes add homogeneous LoRA dtype tracking, FP8 LoRA loading and validation, FP8 grouped GEMM and CUDA-graph execution, dtype propagation through Python bindings, shared LoRA result casting, and regression tests. ChangesFP8 LoRA support
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant PEFTCacheManager
participant PyExecutor
participant LoraLayer
participant groupedGemm
PEFTCacheManager->>PyExecutor: expose data_type
PyExecutor->>LoraLayer: pass LoRA parameters and data_type
LoraLayer->>groupedGemm: execute aligned FP8 LoRA GEMM
groupedGemm-->>LoraLayer: return FP8 LoRA result
LoraLayer-->>PyExecutor: cast result to base output dtype
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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: 8
🧹 Nitpick comments (8)
cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu (2)
315-329: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winUnchecked CUDA API return codes.
cudaMemcpyAsync(x4) andcudaGetDevicereturn values are discarded, so a failure here surfaces later as a confusing CUTLASS/launch error instead of at the source. Wrap them in the existingTLLM_CUDA_CHECKhelper.🛡️ Proposed fix
- cudaMemcpyAsync(devPtrA, ptrAGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(devPtrB, ptrBGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(devPtrC, ptrCGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(devPtrD, ptrDGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); + TLLM_CUDA_CHECK(cudaMemcpyAsync(devPtrA, ptrAGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream)); + TLLM_CUDA_CHECK(cudaMemcpyAsync(devPtrB, ptrBGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream)); + TLLM_CUDA_CHECK(cudaMemcpyAsync(devPtrC, ptrCGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream)); + TLLM_CUDA_CHECK(cudaMemcpyAsync(devPtrD, ptrDGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream)); @@ - cudaGetDevice(&hwInfo.device_id); + TLLM_CUDA_CHECK(cudaGetDevice(&hwInfo.device_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 `@cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu` around lines 315 - 329, Wrap all four cudaMemcpyAsync calls and the cudaGetDevice call in the existing TLLM_CUDA_CHECK helper, preserving their current arguments and ordering so CUDA failures are reported at the originating operation.
50-71: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDuplicated FP8 alignment validation.
This is nearly identical to
checkFp8GroupedGemmAlignmentincpp/tensorrt_llm/kernels/groupGemm.cu(lines 51-66), differing only in taking a raw pointer + count instead of a vector. Consider a single shared helper (e.g. declared ingroupGemm.h, which this file already includes) takingGemmCoord const*, int, char const*and having the vector-based caller forwarddata()/size().🤖 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 `@cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu` around lines 50 - 71, Consolidate the duplicated FP8 alignment validation by moving the shared logic from checkFp8CudaGraphAlignment and checkFp8GroupedGemmAlignment into a helper declared in groupGemm.h that accepts GemmCoord const*, an integer count, and kernelName. Update the vector-based caller to forward its data pointer and size, and have both call sites use the shared helper while preserving the existing validation behavior and diagnostics.cpp/tensorrt_llm/kernels/groupGemm.cu (2)
429-437: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
cudaGetDevicereturn value is discarded.
hwInfo.device_idis pre-set to 0 and thecudaGetDeviceresult is ignored, so a failing call silently proceeds with device 0 and an SM count queried for the wrong device.🛡️ Check the call
- cutlass::KernelHardwareInfo hwInfo; - hwInfo.device_id = 0; - cudaGetDevice(&hwInfo.device_id); + cutlass::KernelHardwareInfo hwInfo; + hwInfo.device_id = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&hwInfo.device_id)); hwInfo.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hwInfo.device_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 `@cpp/tensorrt_llm/kernels/groupGemm.cu` around lines 429 - 437, Handle the return status of cudaGetDevice before using hwInfo.device_id in the surrounding GEMM setup. Propagate or return the CUDA error according to the enclosing function’s existing error-handling convention, and only query the multiprocessor count and construct Gemm::Arguments after successfully obtaining the active device; do not fall back to the preset device 0.
94-108: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winUnify FP8 workspace sizing with the launch layout The helper still hardcodes 12-byte problem shapes and 16-byte strides, while the launch path sizes the same buffers from
sizeof(...). If CUTLASS changes either type, workspace allocation will undercount and fail at runtime; move the FP8 layout sizes behind a shared constexpr helper or addstatic_asserts on the CUTLASS types.🤖 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 `@cpp/tensorrt_llm/kernels/groupGemm.cu` around lines 94 - 108, The getFp8GroupedGemmParamsWorkSpaceSize helper hardcodes CUTLASS buffer sizes, risking divergence from the launch layout. Replace the literal shape and stride byte counts with sizes derived from the same CUTLASS types used by the FP8 launch path, preferably through a shared constexpr layout helper; otherwise add static_asserts validating the assumptions while keeping workspace alignment and pointer-array sizing unchanged.cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp (1)
146-168: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTest coverage summary.
Added:
PeftCacheManagerTest.adapterSelectsHomogeneousCacheDataType(FP8-guarded) — covers dtype selection from the first adapter and rejection of a second adapter with a different dtype. No test functions were modified or removed, and C++ gtest cases are registered via CMake rather than thetests/integration/test_lists/files, so no list updates apply here.Verdict: sufficient for the dtype-selection path, with two gaps worth adding when convenient — (1) asserting the dtype actually reached the underlying host/device
LoraCache(LoraCache::getDataType()), and (2) a case pinning behavior for a dtype that is neither the model dtype nor FP8, which is currently accepted (see thepeftCacheManager.cppcomment).🤖 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 `@cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp` around lines 146 - 168, Extend adapterSelectsHomogeneousCacheDataType to verify that the selected dtype is propagated to the underlying host/device LoraCache by asserting LoraCache::getDataType(). Add coverage for a non-model, non-FP8 dtype and pin the currently accepted behavior described in peftCacheManager.cpp.tests/unittest/others/test_lora_manager.py (1)
27-27: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winLGTM on the rest of the new FP8 helper/test additions (rank/dimension alignment coverage, MoE dtype-conversion coverage).
QA test coverage summary: Added
TestLoraManagerFp8(test_fp8_dora_is_rejected,test_fp8_moe_weights_are_converted_to_model_dtype) andTestLoraManagerFp8Alignment(test_misaligned_fp8_adapter_is_rejected_before_cuda_transfer, parametrized over rank/input/output misalignment). These are unittest-style tests undertests/unittest/others/, not the integration suite. Coverage verdict: needs follow-up — I cannot confirm from the provided context whether/howtests/unittest/others/test_lora_manager.pyis registered intests/integration/test_lists/test-db/for CI execution; please confirm this file's list entry still covers the new FP8 test classes.As per path instructions for
tests/**: "Always produce a test coverage summary, even if no issues are found."Also applies to: 53-124, 214-273
🤖 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/others/test_lora_manager.py` at line 27, Confirm that tests/unittest/others/test_lora_manager.py is registered in the relevant tests/integration/test_lists/test-db/ entry so CI executes the new TestLoraManagerFp8 and TestLoraManagerFp8Alignment classes. Add or update the list entry if necessary, preserving coverage for the existing test file.Source: Path instructions
tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py (2)
124-188: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKernel-facing tests validate source-code substrings, not runtime behavior.
These tests (
test_fp8_cuda_graph_grouped_gemm_uses_live_device_problem_metadata,test_fp8_grouped_gemm_dispatch_has_explicit_unsupported_cutlass_guard,test_fp8_grouped_gemm_alignment_checks_require_multiples_of_16,test_fp8_splitk_grouped_gemm_delegates_to_regular_grouped_gemm,test_fp8_tma_alignment_has_one_cpp_definition,test_fp8_cuda_graph_alignment_check_requires_rank_multiple_of_16) read.cu/.hfiles as raw text and assert on specific substrings. They will break on harmless refactors (renames, reformatting, comment changes) and won't catch numerical/behavioral regressions that happen to preserve the checked strings. Consider replacing these with (or supplementing them with) actual GPU-executed regression tests of the grouped-GEMM kernels where feasible.🤖 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_fp8_lora_grouped_gemm_regressions.py` around lines 124 - 188, Replace the raw source-substring assertions in the listed FP8 grouped-GEMM tests with GPU-executed regression coverage that invokes the relevant grouped-GEMM and CUDA-graph kernels and validates numerical results, dispatch behavior, alignment handling, and unsupported CUTLASS behavior. Preserve source-level checks only where runtime execution cannot verify the requirement, especially the single-definition check for kFp8TmaAlignment, and retain coverage for live device metadata and split-K delegation through observable kernel behavior.
1-188: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftPrefer behavior-level checks over raw source-string assertions — the FP8 Python tests are useful, but the
.cu/.hsubstring checks are brittle and can miss real kernel regressions while breaking on harmless refactors.
- Coverage: new tests cover FP8 cache casting, dtype-mismatch rejection, rank/shape alignment, and grouped-GEMM guard paths; no matches were found in
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/.🤖 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_fp8_lora_grouped_gemm_regressions.py` around lines 1 - 188, Replace the raw `.cu`/`.h` substring assertions in the kernel-source tests with behavior-level tests that exercise the FP8 grouped-GEMM and CUDA-graph dispatch paths, including alignment validation, unsupported CUTLASS guards, live metadata usage, split-K delegation, and single-definition expectations. Preserve the existing Python coverage while making tests fail only when observable kernel behavior regresses rather than when source formatting or implementation details change.
🤖 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 `@cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp`:
- Around line 265-272: Update the LoRA request flow around configureDataType and
loraValidateRequestTensors so validation occurs before configuring the cache
dtype, or otherwise enforce the model dtype plus supported FP8 dtype allowlist.
Ensure unsupported request dtypes are rejected against the model configuration
before the first request can pin host and device caches, while preserving
homogeneous dtype checks for subsequent requests.
In `@cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu`:
- Around line 365-387: Add a runtime device-capability check before the FP8
dispatch and unsupported-path handling in cudaGraphGroupedGemm, requiring SM90
or newer. Ensure both ENABLE_FP8 branches use this guard so older GPUs receive a
clear actionable error before invoking fp8CudaGraphGroupedGemm or the existing
unsupported-feature check.
In `@cpp/tensorrt_llm/kernels/groupGemm.cu`:
- Around line 48-68: Move checkFp8GroupedGemmAlignment under the same ENABLE_FP8
and CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED preprocessor guards as its
only call site, ensuring it is not compiled in configurations where it cannot be
referenced.
In `@cpp/tensorrt_llm/thop/loraOp.cpp`:
- Around line 236-238: Fix both dtype validation default cases in
cpp/tensorrt_llm/thop/loraOp.cpp at lines 236-238 and 321-323, including
lora_group_gemm_param_fill_row_reorder_fusion, by removing the printf-style %s
placeholder from the TORCH_CHECK message and leaving a trailing separator before
c10::toString(dtype), so stream concatenation produces a correctly formatted
error.
In `@tensorrt_llm/_torch/peft/lora/layer.py`:
- Around line 507-520: Update the CUDA-graph key construction used by the LoRA
grouped-GEMM path to include the active adapter/rank state, especially the FP8
value derived by _validate_fp8_lora_cuda_graph_alignment from
cuda_graph_params.slot_ranks_host. Ensure graph replay distinguishes changes to
slot_ranks_host/min_kn and cannot reuse a capture with a stale FP8
specialization; keep the existing non-FP8 min_kn calculation unchanged.
In `@tensorrt_llm/lora_manager.py`:
- Around line 1081-1098: Restrict FP8 LoRA handling in
tensorrt_llm/lora_manager.py lines 1081-1098, around is_fp8 and use_fp8_kernel,
to torch.float8_e4m3fn only, and validate or cast t_in to match t_out before
using the FP8 kernel. In tensorrt_llm/_torch/peft/lora/layer.py lines 238-261,
make no direct change; retain its E4M3-only dtype handling as the source of
truth unless E5M2 support is explicitly added instead.
- Around line 1110-1113: Update the FP8 scaling path in the use_fp8_kernel
branch to clamp the BF16-scaled t_out values to the representable finite range
before casting back to t_out.dtype. Reuse the existing FP8 clamp behavior or
bounds established by LoraLayer.forward in layer.py, while preserving the
current BF16 scaling flow.
In `@tests/unittest/others/test_lora_manager.py`:
- Around line 86-89: Update the DoRA test setup in the use_dora branch to size
lora_magnitude_vector with output_size rather than hidden_size, ensuring the
tensor matches the layer’s per-output-channel dimension when output_size
differs.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu`:
- Around line 315-329: Wrap all four cudaMemcpyAsync calls and the cudaGetDevice
call in the existing TLLM_CUDA_CHECK helper, preserving their current arguments
and ordering so CUDA failures are reported at the originating operation.
- Around line 50-71: Consolidate the duplicated FP8 alignment validation by
moving the shared logic from checkFp8CudaGraphAlignment and
checkFp8GroupedGemmAlignment into a helper declared in groupGemm.h that accepts
GemmCoord const*, an integer count, and kernelName. Update the vector-based
caller to forward its data pointer and size, and have both call sites use the
shared helper while preserving the existing validation behavior and diagnostics.
In `@cpp/tensorrt_llm/kernels/groupGemm.cu`:
- Around line 429-437: Handle the return status of cudaGetDevice before using
hwInfo.device_id in the surrounding GEMM setup. Propagate or return the CUDA
error according to the enclosing function’s existing error-handling convention,
and only query the multiprocessor count and construct Gemm::Arguments after
successfully obtaining the active device; do not fall back to the preset device
0.
- Around line 94-108: The getFp8GroupedGemmParamsWorkSpaceSize helper hardcodes
CUTLASS buffer sizes, risking divergence from the launch layout. Replace the
literal shape and stride byte counts with sizes derived from the same CUTLASS
types used by the FP8 launch path, preferably through a shared constexpr layout
helper; otherwise add static_asserts validating the assumptions while keeping
workspace alignment and pointer-array sizing unchanged.
In `@cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp`:
- Around line 146-168: Extend adapterSelectsHomogeneousCacheDataType to verify
that the selected dtype is propagated to the underlying host/device LoraCache by
asserting LoraCache::getDataType(). Add coverage for a non-model, non-FP8 dtype
and pin the currently accepted behavior described in peftCacheManager.cpp.
In `@tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py`:
- Around line 124-188: Replace the raw source-substring assertions in the listed
FP8 grouped-GEMM tests with GPU-executed regression coverage that invokes the
relevant grouped-GEMM and CUDA-graph kernels and validates numerical results,
dispatch behavior, alignment handling, and unsupported CUTLASS behavior.
Preserve source-level checks only where runtime execution cannot verify the
requirement, especially the single-definition check for kFp8TmaAlignment, and
retain coverage for live device metadata and split-K delegation through
observable kernel behavior.
- Around line 1-188: Replace the raw `.cu`/`.h` substring assertions in the
kernel-source tests with behavior-level tests that exercise the FP8 grouped-GEMM
and CUDA-graph dispatch paths, including alignment validation, unsupported
CUTLASS guards, live metadata usage, split-K delegation, and single-definition
expectations. Preserve the existing Python coverage while making tests fail only
when observable kernel behavior regresses rather than when source formatting or
implementation details change.
In `@tests/unittest/others/test_lora_manager.py`:
- Line 27: Confirm that tests/unittest/others/test_lora_manager.py is registered
in the relevant tests/integration/test_lists/test-db/ entry so CI executes the
new TestLoraManagerFp8 and TestLoraManagerFp8Alignment classes. Add or update
the list entry if necessary, preserving coverage for the existing test file.
🪄 Autofix (Beta)
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: 90fecfec-8652-4709-9502-58211eb2f1ef
📒 Files selected for processing (27)
cpp/include/tensorrt_llm/batch_manager/peftCacheManager.hcpp/include/tensorrt_llm/runtime/loraCache.hcpp/tensorrt_llm/batch_manager/peftCacheManager.cppcpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cucpp/tensorrt_llm/kernels/groupGemm.cucpp/tensorrt_llm/kernels/groupGemm.hcpp/tensorrt_llm/kernels/lora/lora.cppcpp/tensorrt_llm/kernels/splitkGroupGemm.cucpp/tensorrt_llm/kernels/splitkGroupGemm.hcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cppcpp/tensorrt_llm/runtime/loraCache.cppcpp/tensorrt_llm/runtime/loraUtils.cppcpp/tensorrt_llm/runtime/loraUtils.hcpp/tensorrt_llm/thop/loraOp.cppcpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpptensorrt_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/_torch/pyexecutor/resource_manager.pytensorrt_llm/lora_manager.pytests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.pytests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.pytests/unittest/others/test_lora_manager.py
c47972f to
be1231e
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cpp/tensorrt_llm/kernels/groupGemm.cu (2)
437-440: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
cudaGetDevicereturn value is discarded.A failure here leaves
hwInfo.device_idat 0 and the SM count query targeting the wrong device. Wrap it in the project's CUDA check macro.♻️ Proposed fix
cutlass::KernelHardwareInfo hwInfo; - hwInfo.device_id = 0; - cudaGetDevice(&hwInfo.device_id); + hwInfo.device_id = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&hwInfo.device_id)); hwInfo.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hwInfo.device_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 `@cpp/tensorrt_llm/kernels/groupGemm.cu` around lines 437 - 440, Update the hardware-info initialization around cudaGetDevice so its return status is checked with the project's CUDA check macro before querying multiprocessor count. Preserve the existing hwInfo.device_id assignment and ensure query_device_multiprocessor_count uses the validated active device.
102-116: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winWorkspace size math is duplicated with hard-coded element sizes.
getFp8GroupedGemmParamsWorkSpaceSizehard-codes 12/8/16-byte element sizes whilefp8GroupedGemm(Lines 375-386) recomputes the same layout withsizeof(...). If any CUTLASS stride/shape type changes size, the two silently diverge and the mismatch only surfaces as a runtimeTLLM_CHECKfailure. Consider exposing the layout computation from one place (e.g. a small helper templated on the kernel types, or at least namedk-prefixed constants shared by both), so the sizing and the consumer stay in lockstep.As per coding guidelines, "Avoid magic literals except
0,nullptr,true, andfalse; initialize named constants instead, usingk-prefixed camelCase names."🤖 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 `@cpp/tensorrt_llm/kernels/groupGemm.cu` around lines 102 - 116, The workspace sizing in getFp8GroupedGemmParamsWorkSpaceSize duplicates the layout sizes used by fp8GroupedGemm. Centralize the element-size and alignment calculations in a shared helper or k-prefixed named constants derived from the actual CUTLASS shape, pointer, and stride types, then have both functions reuse that definition so type-size changes cannot diverge and no magic byte literals remain.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 `@cpp/tensorrt_llm/kernels/groupGemm.cu`:
- Around line 437-440: Update the hardware-info initialization around
cudaGetDevice so its return status is checked with the project's CUDA check
macro before querying multiprocessor count. Preserve the existing
hwInfo.device_id assignment and ensure query_device_multiprocessor_count uses
the validated active device.
- Around line 102-116: The workspace sizing in
getFp8GroupedGemmParamsWorkSpaceSize duplicates the layout sizes used by
fp8GroupedGemm. Centralize the element-size and alignment calculations in a
shared helper or k-prefixed named constants derived from the actual CUTLASS
shape, pointer, and stride types, then have both functions reuse that definition
so type-size changes cannot diverge and no magic byte literals remain.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6148191b-b52e-40de-ac45-1f9b0163edc3
📒 Files selected for processing (11)
cpp/tensorrt_llm/batch_manager/peftCacheManager.cppcpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cucpp/tensorrt_llm/kernels/groupGemm.cucpp/tensorrt_llm/kernels/groupGemm.hcpp/tensorrt_llm/kernels/splitkGroupGemm.cucpp/tensorrt_llm/runtime/loraCache.cppcpp/tensorrt_llm/thop/loraOp.cppcpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpptensorrt_llm/lora_manager.pytests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.pytests/unittest/others/test_lora_manager.py
🚧 Files skipped from review as they are similar to previous changes (8)
- cpp/tensorrt_llm/kernels/groupGemm.h
- cpp/tensorrt_llm/thop/loraOp.cpp
- tensorrt_llm/lora_manager.py
- cpp/tensorrt_llm/kernels/splitkGroupGemm.cu
- cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp
- cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu
- tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py
- cpp/tensorrt_llm/runtime/loraCache.cpp
|
/bot run |
|
PR_Github #62206 [ run ] triggered by Bot. Commit: |
|
PR_Github #62541 [ run ] completed with state
|
brb-nv
left a comment
brb-nv
left a comment
There was a problem hiding this comment.
Sorry, I'm not able to complete the review at once. One thing stood out to me for now.
|
/bot run |
|
PR_Github #62601 [ run ] triggered by Bot. Commit: |
| // ==================================================================== | ||
|
|
||
| template <typename ProblemShape, typename StrideA, typename StrideB, typename StrideC, typename StrideD> | ||
| __global__ void fillFp8CudaGraphGroupedGemmParams(cutlass::gemm::GemmCoord const* problemSizesPtr, int problemCount, |
There was a problem hiding this comment.
Is it possible to fuse this into loraGroupGEMMParamFillRowReorderFusionKernel so the latter directly builds strides?
| auto* devStrideC = static_cast<StrideC*>(devPtr(szStrideC)); | ||
| auto* devStrideD = static_cast<StrideD*>(devPtr(szStrideD)); | ||
|
|
||
| cudaMemcpyAsync(devPtrA, ptrAGpu, problemCount * sizeof(void*), cudaMemcpyDeviceToDevice, stream); |
There was a problem hiding this comment.
Are the D2D copies here necessary?
|
PR_Github #62601 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62638 [ run ] triggered by Bot. Commit: |
|
PR_Github #62638 [ run ] completed with state
|
|
[by Codex] @SimengLiu-nv Friendly review reminder: this PR is awaiting your review. Thanks! |
brb-nv
left a comment
brb-nv
left a comment
There was a problem hiding this comment.
Approving with a few comments. Do you think we should reject MoE fp8 early given you mentioned that would be a separate follow-up?
ed215ab to
cf109ae
Compare
|
/bot run |
|
PR_Github #63304 [ run ] triggered by Bot. Commit: |
|
PR_Github #63304 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63309 [ run ] triggered by Bot. Commit: |
|
PR_Github #63309 [ run ] completed with state
|
Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>
Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>
Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>
Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>
Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>
cf109ae to
ad02cba
Compare
|
/bot run |
|
PR_Github #63507 [ run ] triggered by Bot. Commit: |
|
PR_Github #63507 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63536 [ run ] triggered by Bot. Commit: |
|
PR_Github #63536 [ run ] completed with state
|
5 participants
Dev Engineer Review
Homogeneous PEFT cache dtype enforcement
PeftCacheManagerinternal homogeneous LoRA adapter dtype tracking:std::optional<DataType>state guarded by a mutex plusgetDataType().configureDataType(DataType)validates dtype compatibility (FP8 only whenENABLE_FP8) and rejects mixed-dtype caching across requests; includes fallback to model datatype when unset.runtime::LoraCachewithsetDataType(DataType)/getDataType():setDataTypereinitializes only when cache is empty (prevents post-insert dtype changes).loraValidateRequestTensors(..., std::optional<DataType> loraDataType = std::nullopt)validates weights dtype against the provided cache dtype when specified.PeftCacheManager.data_type."data_type"from the PEFT cache manager into Python.Hopper/newer dense FP8 (E4M3) LoRA grouped-GEMM support + CUDA-graph dispatch
kernels::kFp8TmaAlignment = 16.problemdimensions /minKN).GemmUniversalAdapter.groupGemm.cuadds Hopper FP8 grouped-GEMM implementation and dispatch whendataType == kFP8(with explicit compile-time guard errors when unsupported).cuda_graph_grouped_gemm.cuintercepts FP8 for CUDA-graph grouped paths (cudaGraphGroupedGemmandcudaGraphSplitKGroupedGemm), otherwise failing with explicit guard messages.splitkGroupGemm.cuforwards FP8 to the non-split-K grouped GEMM path.FP8 LoRA execution correctness (dtype conversion/clamping, kernel-path selection, workspace sizing)
kernels/lora/lora.cpp:useUnifiedGemm = falsefor FP8.add_lora_result(output, lora_result)to centralize dtype-correct LoRA accumulation.LoraLayer.forward:Nonewhenlora_paramsis falsy.data_typediffers from input dtype, clamps/casts float16/bfloat16/float32 inputs to FP8 e4m3 range before executing, and restores output dtype back to the original base dtype.min_knwhen activations aretorch.float8_e4m3fn.add_lora_resultusage.ENABLE_FP8, mapstorch.float8_e4m3fntoDataType::kFP8for relevant LoRA entry points and improves error messages to list FP8 as supported.FP8 LoRA adapter loading restrictions/conversions
lora_manager.pyadds Hopper-gated FP8 LoRA support:NotImplementedError).QA Engineer Review
Test code changes (files under
tests/):cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cppunsupportedAdapterDataTypeDoesNotConfigureCacheENABLE_FP8)invalidAdapterDoesNotConfigureCacheDataTypeENABLE_FP8)adapterSelectsHomogeneousCacheDataTypetests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py_validate_fp8_lora_cuda_graph_alignmentpass/fail (including misaligned rank/dims expectations)LoraLayerFP8 cache input clamping/casting and output dtype restorationadd_lora_resultcasting semantics (includingNonepassthrough)tests/unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.pylora_rank=16whendtype=torch.float8_e4m3fn(otherwiselora_rank=8).tests/unittest/others/test_lora_manager.pyTestLoraManagerFp8,TestLoraManagerFp8Alignment) covering DoRA rejection, MoE handling/conversion, dtype mismatch rejection, clamping/scaling behavior, and alignment checks (via patchingtorch.cuda.get_device_capability).Coverage in
tests/integration/test_lists/(CBTS mapping):Description
This change keeps E4M3 LoRA adapter weights in FP8 from loading through execution on Hopper and newer GPUs.
The PEFT cache now selects one homogeneous dtype from the first adapter loaded while the cache is empty and rejects later adapters with a different dtype. That dtype is propagated through the C++ cache, nanobind, Python LoRA layers, and grouped-GEMM dispatch. For an FP8 cache, dense LoRA layers saturate and convert FP16, BF16, or FP32 activations to E4M3 before the grouped GEMM, then convert the LoRA result back to the caller activation dtype before accumulation.
The eager, split-K, and CUDA-graph grouped-GEMM implementations gain Hopper FP8 CUTLASS dispatch. The implementation requires E4M3 adapter weights, SM90 or newer, rank 16 or greater, and rank/input/output dimensions divisible by 16. Routed-expert MoE LoRA remains outside this PR.
No public configuration surface or new dependency is introduced.
Test Coverage
Validated on an H100 PCIe with an SM90 Release build:
PeftCacheManagerTest.adapterSelectsHomogeneousCacheDataType: 1 passed.pytest tests/unittest/_torch/lora/test_fp8_lora_grouped_gemm_regressions.py: 16 passed. This covers cache dtype propagation, eager and CUDA-graph activation conversion, grouped-GEMM constraints, and FP8 delta accumulation into BF16 output.TestQwen3LoRA::test_qwen3_fp8_lora: passed end to end againstQwen3-0.6B-FP8with generated nonzero rank-16 FP8 adapters.TestQwen3LoRA::test_qwen3_bf16_lora: passed as the neighboring BF16 adapter control against the same checkpoint.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using the CodeRabbit summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.