[None][feat] Adding torch ext API for FusedAddRMSNormQuant kernel #9905
Conversation
JintaoPengCS
changed the title
📝 WalkthroughWalkthroughIntroduces a new fused CUDA kernel operation combining Add, RMSNorm/LayerNorm, and FP4 quantization for PyTorch. Changes include C++ implementation with input validation and kernel launcher, Python bindings, a fake implementation, build configuration updates, and compile-time guard optimization in an existing header. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 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: 1
🧹 Nitpick comments (5)
cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh (1)
1-2: Update copyright year to include 2025.The copyright header should be updated to reflect the current year.
/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpp (4)
1-15: Update copyright year to include 2025.The copyright header should reflect the current year per coding guidelines.
/* - * Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2025, NVIDIA CORPORATION. All rights reserved.
74-80: Consider safer access for gamma dimension check.If
gammais a 0-dimensional tensor,gamma.sizes()[0]would cause undefined behavior.- TORCH_CHECK(gamma.sizes()[0] == n, "gamma size must match hidden dimension N."); + TORCH_CHECK(gamma.dim() == 1 && gamma.size(0) == n, "gamma must be 1D tensor with size matching hidden dimension N.");
131-136: Potential narrowing conversion fromint64_ttoint.Lines 131-132 cast
mandnfromint64_ttoint. While the TORCH_CHECK at line 79 ensuresn <= 16384, there's no upper bound check form. For very large token counts, this could silently truncate.+ TORCH_CHECK(m <= std::numeric_limits<int>::max(), "M dimension exceeds int32 limit."); param.m = static_cast<int>(m); param.n = static_cast<int>(n);
133-133: Hardcoded epsilon value.The layernorm epsilon is hardcoded to
1e-5f. Consider whether this should be configurable via a parameter, as different models may use different epsilon values (e.g.,1e-6fis also common).
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh(1 hunks)cpp/tensorrt_llm/thop/CMakeLists.txt(1 hunks)cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpp(1 hunks)tensorrt_llm/_torch/custom_ops/torch_custom_ops.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (3)
**/*.{cpp,h,cu}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.{cpp,h,cu}: Closing braces of namespaces should have a comment saying the namespace it closes (e.g.,} // namespace foo)
Preferconstorconstexprvariables over#definewhenever possible, as the latter are not visible to the compiler
A variable that is not modified after its initialization should be declared asconst
Except0(only used in comparison for checking signness/existence/emptiness) andnullptr,true,false, all other literals should only be used for variable initialization and should be replaced with named constants
Use Allman indentation style for braces in C++
Put the semicolon for an emptyfororwhileloop in a new line
The statement forming the body of aswitch,while,do .. whileorforstatement shall be a compound statement (use brace-delimited statements)
Ifandelseshould always be followed by brace-delimited statements, even if empty or a single statement
C++ filenames should use camel case with first letter lowercase (e.g.,thisIsASubDirandthisIsAFilename.cpp)
All filenames involved in compilation of a compilation target must have case-insensitive unique filenames
All types (including class names) should use camel case with uppercase first letter (e.g.,FooBarClass)
Local variables, methods and namespaces should use camel case with first letter lowercase (e.g.,localFooBar)
Non-magic-number global variables that are non-static and not defined in anonymous namespace should use camel case prefixed by a lower case 'g' (e.g.,gDontUseGlobalFoos)
Non-magic-number global variables that are static or defined in an anonymous namespace should use camel case prefixed by a lower case 's' (e.g.,sMutableStaticGlobal)
Locally visible static variables should use camel case with lowercase prefix 's' as the first letter of the name (e.g.,static std::once_flag sFlag;)
Public, private and protected class member variables should use camel case prefixed with 'm' (e.g.,mNbFooValues), though the 'm' pre...
Files:
cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpp
**/*.{cpp,h,cu,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code files should contain an NVIDIA copyright header that includes the current year at the top
Files:
cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpptensorrt_llm/_torch/custom_ops/torch_custom_ops.py
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces; do not use tabs
Always maintain the namespace when importing in Python, even if only one class or function from a module is used (e.g., usefrom package.subpackage import fooand thenfoo.SomeClass()instead offrom package.subpackage.foo import SomeClass)
Python filenames should use snake_case (e.g.,some_file.py)
Python class names should use PascalCase (e.g.,class SomeClass)
Python function and method names should use snake_case (e.g.,def my_awesome_function():)
Python local variable names should use snake_case, with prefixkfor variable names that start with a number (e.g.,k_99th_percentile = ...)
Python global variables should use upper snake_case with prefixG(e.g.,G_MY_GLOBAL = ...)
Python constants should use upper snake_case (e.g.,MY_CONSTANT = ...)
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Python comments should be reserved for code within a function, or interfaces that are local to a file
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with type and description (e.g.,self.x = 5followed by"""<type>: Description of 'x'""")
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except clause to the smallest set of specific errors possible instead of catching all exceptions
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible and use the else block to implement the logic
Files:
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
🧠 Learnings (14)
📓 Common learnings
Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py:180-182
Timestamp: 2025-10-20T17:09:21.560Z
Learning: In tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py, the _gated_rmsnorm_replacement function does not need to cast the output of torch.ops.auto_deploy.torch_rmsnorm_gated back to the input dtype, even though the custom op returns fp32. The dtype handling is managed elsewhere or the fp32 output is acceptable for downstream consumers.
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/multimem.h:20-30
Timestamp: 2025-09-23T15:13:48.819Z
Learning: TRT-LLM targets modern CUDA toolkits that support FP8 datatypes, so cuda_fp8.h can be included unconditionally without version guards in TRT-LLM code.
📚 Learning: 2025-09-23T15:01:00.070Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:15-17
Timestamp: 2025-09-23T15:01:00.070Z
Learning: In TensorRT-LLM NCCL device kernels, the <sstream> header is not needed as an explicit include in config.cu because it's provided transitively through other headers. Local compilation testing confirms this works without the explicit include.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh
📚 Learning: 2025-09-23T15:13:48.819Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/multimem.h:20-30
Timestamp: 2025-09-23T15:13:48.819Z
Learning: TRT-LLM targets modern CUDA toolkits that support FP8 datatypes, so cuda_fp8.h can be included unconditionally without version guards in TRT-LLM code.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cppcpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh
📚 Learning: 2025-09-23T15:01:00.070Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:15-17
Timestamp: 2025-09-23T15:01:00.070Z
Learning: In TensorRT-LLM NCCL device kernels (cpp/tensorrt_llm/kernels/nccl_device/config.cu), std::ostringstream is used but <sstream> doesn't need to be explicitly included because it's provided transitively through other headers like tensorrt_llm/common/cudaUtils.h or config.h. Local compilation testing confirms this works without the explicit include.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh
📚 Learning: 2025-10-20T17:09:21.560Z
Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py:180-182
Timestamp: 2025-10-20T17:09:21.560Z
Learning: In tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py, the _gated_rmsnorm_replacement function does not need to cast the output of torch.ops.auto_deploy.torch_rmsnorm_gated back to the input dtype, even though the custom op returns fp32. The dtype handling is managed elsewhere or the fp32 output is acceptable for downstream consumers.
Applied to files:
cpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpptensorrt_llm/_torch/custom_ops/torch_custom_ops.py
📚 Learning: 2025-10-20T16:54:09.824Z
Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py:6-6
Timestamp: 2025-10-20T16:54:09.824Z
Learning: In tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py, the import `from ...modules.mamba.layernorm_gated import _layer_norm_fwd` is correct and should not be changed to modules.fla.layernorm_gated. The _layer_norm_fwd function exists in both modules/mamba/layernorm_gated.py and modules/fla/layernorm_gated.py, but the mamba version is the intended implementation for this use case.
Applied to files:
cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cppcpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuhtensorrt_llm/_torch/custom_ops/torch_custom_ops.py
📚 Learning: 2025-09-23T15:12:38.312Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/thop/allreduceOp.cpp:352-446
Timestamp: 2025-09-23T15:12:38.312Z
Learning: In TensorRT-LLM NCCL device allreduce implementation (cpp/tensorrt_llm/thop/allreduceOp.cpp), the goto pattern in runNCCLAllReduceDeviceFusion is intentionally used for future extensibility, allowing multiple switch cases to fallback to the default handler. While not aesthetically ideal, this pattern supports adding more fusion cases later that can reuse the same fallback logic.
Applied to files:
cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpp
📚 Learning: 2025-08-08T22:03:40.707Z
Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1198-1209
Timestamp: 2025-08-08T22:03:40.707Z
Learning: In the CUTLASS MoE kernels (cpp/tensorrt_llm/cutlass_extensions), when `layout_info.fusion` is set to `TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE`, the `router_scales` parameter must be non-null by design. The fused finalize kernel epilogue does not perform nullptr checks and requires valid router scales to function correctly. This is an implicit contract that callers must satisfy when enabling the FINALIZE fusion mode.
Applied to files:
cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cppcpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh
📚 Learning: 2025-08-08T05:10:38.906Z
Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp:0-0
Timestamp: 2025-08-08T05:10:38.906Z
Learning: The ScaledAccPerRowBiasPerColScaleScatter fusion in CUTLASS extensions (cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp) is specifically designed for per-column scaling factors only, so it uses a fixed Stride<_0,_1,int64_t> rather than conditional stride logic.
Applied to files:
cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpp
📚 Learning: 2025-09-23T15:12:38.312Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/thop/allreduceOp.cpp:352-446
Timestamp: 2025-09-23T15:12:38.312Z
Learning: In TensorRT-LLM NCCL device implementation, NCCL version 2.28+ requirements are handled at runtime in the nccl_device/config layer rather than with compile-time guards. This allows the allreduceOp to remain version-agnostic and delegates version compatibility validation to the appropriate lower-level components that can gracefully handle unsupported configurations.
Applied to files:
cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh
📚 Learning: 2025-11-24T17:09:17.870Z
Learnt from: CR
Repo: NVIDIA/TensorRT-LLM PR: 0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-11-24T17:09:17.870Z
Learning: Applies to **/*.h : Use a preprocessor guard in C++ header files with the guard name format `TRTLLM_` followed by the filename in all caps (e.g., `TRTLLM_FOO_BAR_HELLO_H` for file `FooBarHello.h`); do not include directory names in the symbol
Applied to files:
cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh
📚 Learning: 2025-11-14T11:22:03.729Z
Learnt from: nzmora-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 9163
File: tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py:107-113
Timestamp: 2025-11-14T11:22:03.729Z
Learning: In TensorRT-LLM AutoDeploy custom ops, when adding hardware capability checks to select between kernel implementations (e.g., cuBLAS vs. CUDA kernel), use descriptive variable names that identify the specific GPU architectures or families being targeted (e.g., `is_blackwell_geforce_or_ada`) rather than generic names like `enable_cuda_core`. This makes it clear that the code is selecting an implementation path based on hardware capabilities, not enabling/disabling hardware features.
Applied to files:
cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh
📚 Learning: 2025-08-08T05:06:31.596Z
Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp:36-36
Timestamp: 2025-08-08T05:06:31.596Z
Learning: CUTLASS extension files (under cpp/tensorrt_llm/cutlass_extensions/) follow CUTLASS coding style conventions, including using #pragma once instead of TRTLLM_ prefixed header guards, even though they are .hpp files.
Applied to files:
cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh
📚 Learning: 2025-08-19T03:35:20.866Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4616-4626
Timestamp: 2025-08-19T03:35:20.866Z
Learning: In the MOE profiler TMA workspace preparation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu), the overlapping of TMA WS regions for NONE and FINALIZE variants is deliberate design to save memory space, as confirmed by djns99. The comment "reuse the same pointers to save space" reflects this intentional behavior.
Applied to files:
cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh
🧬 Code graph analysis (2)
cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpp (3)
cpp/tensorrt_llm/kernels/quantization.h (1)
computeSwizzledLayoutSFSize(53-58)cpp/include/tensorrt_llm/common/cudaUtils.h (1)
getMultiProcessorCount(407-469)tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)
fused_add_rms_norm_quant(1626-1653)
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)
cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpp (2)
fused_add_rms_norm_quant(49-163)fused_add_rms_norm_quant(49-50)
🪛 Ruff (0.14.8)
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
1659-1659: Unused function argument: residual
(ARG001)
1660-1660: Unused function argument: gamma
(ARG001)
1661-1661: Unused function argument: sf_scale
(ARG001)
1662-1662: Unused function argument: use_rms_norm
(ARG001)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (6)
cpp/tensorrt_llm/thop/CMakeLists.txt (1)
69-69: LGTM!The new source file
fusedAddRMSNormQuant.cppis correctly added to theth_commonshared library target, maintaining alphabetical ordering within the source list.cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh (1)
801-829: LGTM! Simplified preprocessor guard.The consolidation of the two-layer preprocessor guards into a single combined condition
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDACC_VER_MAJOR__ >= 12)is cleaner while preserving the same compile-time protection for SM90+ and CUDA 12+. The runtime branching viaif constexprfor SM9/SM10 is correctly preserved inside.tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (2)
1626-1653: LGTM! Clean wrapper with comprehensive documentation.The wrapper function correctly delegates to the C++ implementation via
torch.ops.trtllm.fused_add_rms_norm_quant. The docstring accurately documents the constraints (SM90+, N ∈ [2048, 16384]) matching the C++ implementation.
1656-1672: Thesf_outdtype correctly matches the C++ implementation.The C++ code defines
SF_DTYPE = torch::ScalarType::Byte(line 62 incpp/tensorrt_llm/thop/thUtils.h), and the Python fake implementation correctly usestorch.uint8, which maps to the same scalar type. The C++fusedAddRMSNormQuantfunction allocatessf_outwithSF_DTYPEat line 101, and the comment at line 45 confirms it stores "scale factors for FP4 (uint8_t)". No changes needed.The note about static analysis warnings for unused arguments is valid—fake implementations must match the full function signature even when not all parameters are used.
cpp/tensorrt_llm/thop/fusedAddRMSNormQuant.cpp (2)
52-66: LGTM! Thorough input validation.The input validation is comprehensive: CUDA device checks, contiguity checks, and architecture verification for SM90+. The error message is informative for debugging.
167-176: LGTM! Torch library bindings are correct.The
TORCH_LIBRARY_FRAGMENTdeclaration correctly defines the operator signature, andTORCH_LIBRARY_IMPLproperly binds the CUDA implementation.
| #define LAUNCH_FUSED_ADD_RMS_NORM_QUANT(T) \ | ||
| do \ | ||
| { \ | ||
| using Param = tensorrt_llm::kernels::GeneralFP4AddBiasResidualPreLayerNormParam<T>; \ | ||
| tensorrt_llm::kernels::WarpSpecializedParam<Param> param; \ | ||
| param.normed_output = reinterpret_cast<uint32_t*>(normed_output.data_ptr()); \ | ||
| param.output = reinterpret_cast<T*>(output.data_ptr()); \ | ||
| param.input = const_cast<T*>(reinterpret_cast<T const*>(input.data_ptr())); \ | ||
| param.sf_scale = sfScalePtr; \ | ||
| param.sf_out = reinterpret_cast<uint32_t*>(sf_out.data_ptr()); \ | ||
| param.residual = reinterpret_cast<T const*>(residual.data_ptr()); \ | ||
| param.bias = nullptr; \ | ||
| param.gamma = reinterpret_cast<T const*>(gamma.data_ptr()); \ | ||
| param.beta = nullptr; \ | ||
| param.m = static_cast<int>(m); \ | ||
| param.n = static_cast<int>(n); \ | ||
| param.layernorm_eps = 1e-5f; \ | ||
| param.stream = stream; \ | ||
| param.counters = counters; \ | ||
| tensorrt_llm::kernels::invokeWSLayerNorm<Param>(param, use_rms_norm, multiProcessorCount); \ | ||
| } while (0) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find kernel files related to fusedLayernormKernels
fd "LayerNorm" --type f -e cu -e h cpp/tensorrt_llm/kernels/
# Search for invokeWSLayerNorm definition
rg "invokeWSLayerNorm" --type cpp --type h -A 10Repository: NVIDIA/TensorRT-LLM
Length of output: 2397
🏁 Script executed:
#!/bin/bash
# Find the implementation file for invokeWSLayerNorm
fd "ws_layernorm" --type f cpp/tensorrt_llm/kernels/
# Also search for GeneralFP4AddBiasResidualPreLayerNormParam definition
rg "GeneralFP4AddBiasResidualPreLayerNormParam" --type cpp --type h -B 2 -A 10Repository: NVIDIA/TensorRT-LLM
Length of output: 5628
🏁 Script executed:
#!/bin/bash
# Check the kernel implementation file
cat -n cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh | head -100
# Search for how input is used in the kernel
rg "\.input" cpp/tensorrt_llm/kernels/fusedLayernormKernels/ --type cpp --type cu -A 2 -B 2Repository: NVIDIA/TensorRT-LLM
Length of output: 4584
🏁 Script executed:
#!/bin/bash
# Search for write operations to input in the kernel
rg "input\s*=" cpp/tensorrt_llm/kernels/fusedLayernormKernels/ --type cpp --type cuh -B 2 -A 2
# Also search for how input is accessed in the kernel
rg "param\.input|\.input\[" cpp/tensorrt_llm/kernels/fusedLayernormKernels/ --type cpp --type cuh -B 1 -A 1Repository: NVIDIA/TensorRT-LLM
Length of output: 118
🏁 Script executed:
#!/bin/bash
# Read the kernel implementation file to check how input is used
cat cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh | grep -n "input" -i -A 1 -B 1 | head -80
# Also check the param structure usage
rg "param\." cpp/tensorrt_llm/kernels/fusedLayernormKernels/ --max-count 50Repository: NVIDIA/TensorRT-LLM
Length of output: 14262
Fix the struct definition instead of using const_cast.
The kernel only reads from input (lines 133, 216 in ws_layernorm.cuh show read operations only), but the struct declares it as T* instead of T const*. Change T* input; to T const* input; in GeneralFP4AddBiasResidualPreLayerNormParam (cpp/tensorrt_llm/kernels/fusedLayernormKernels/layernorm_param.h), then remove the const_cast at line 124.
JintaoPengCS
changed the title
JintaoPengCS
force-pushed
the
main
branch
4 times, most recently
from
46b53cd to
6800838
Compare
|
/bot run --disable-fail-fast |
JintaoPengCS
force-pushed
the
main
branch
2 times, most recently
from
6800838 to
e421f96
Compare
|
/bot run --disable-fail-fast |
2 similar comments
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #28287 [ run ] triggered by Bot. Commit: |
|
PR_Github #28287 [ run ] completed with state |
JintaoPengCS
force-pushed
the
main
branch
2 times, most recently
from
1e4145b to
276d424
Compare
|
/bot run --disable-fail-fast |
JintaoPengCS
force-pushed
the
main
branch
2 times, most recently
from
86e45f8 to
8c5a62c
Compare
JintaoPengCS
force-pushed
the
main
branch
2 times, most recently
from
372c200 to
acff6ae
Compare
update #define for hopper & blackwell Signed-off-by: JtaoPeng <[email protected]>
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #31231 [ run ] triggered by Bot. Commit: |
|
PR_Github #31234 [ run ] triggered by Bot. Commit: |
|
PR_Github #31231 [ run ] completed with state |
|
PR_Github #31234 [ run ] completed with state
|
4 participants
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's 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)
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
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
Details
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.