[draft][don't review now] Add CuTe DSL nvfp4 linear #7480
Conversation
Signed-off-by: Mindy Li <[email protected]>
📝 WalkthroughWalkthroughAdds a new CuTe-based NVFP4 Blackwell GEMM custom op and fake registration, integrates it into NVFP4LinearMethod.apply with a new scalar_alpha, updates weight-loading to set scalar_alpha, and expands/updates FP4 linear unit tests with new shapes and parameters. Includes auxiliary utilities and debug prints; duplication observed in helper/class definitions. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Torch as NVFP4LinearMethod.apply
participant Op as trtllm::cute_dsl_nvfp4_gemm_blackwell
participant Tuner as AutoTuner
participant Runner as CuteDSLNVFP4BlackwellLinear
participant CuTe as cute.compile/kernel
participant CUDA as CUDA Stream
Torch->>Op: call(input_fp4, weight_fp4, input_scale, weight_scale, scalar_alpha, dtype)
Op->>Tuner: select tactic
Tuner-->>Op: tactic (kernel config)
Op->>Runner: forward(tactic, tensors, alpha)
Runner->>CuTe: compile(config, shapes)
CuTe-->>Runner: kernel handle
Runner->>CUDA: launch kernel(ptrs, aligned shapes)
CUDA-->>Runner: completion
Runner-->>Op: output (bf16)
Op-->>Torch: output
rect rgba(200,230,255,0.25)
Note over Op,Runner: New CuTe NVFP4 Blackwell path
end
sequenceDiagram
autonumber
participant Torch as Fake mode (meta)
participant Op as trtllm::cute_dsl_nvfp4_gemm_blackwell (fake)
Torch->>Op: invoke with meta tensors
Op-->>Torch: return meta bf16 tensor with computed shape
Note over Torch,Op: Fake registration path (no kernel compile/launch)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, Documentation and Community
|
…inear Signed-off-by: Mindy Li <[email protected]>
limin2021
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
🧹 Nitpick comments (7)
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (2)
1163-1173: Missing error handling in custom opThe
cute_dsl_nvfp4_gemm_blackwellfunction doesn't include error handling for invalid inputs or failed kernel execution.Consider adding input validation and error handling:
def cute_dsl_nvfp4_gemm_blackwell( input: torch.Tensor, weight: torch.Tensor, input_scale: torch.Tensor, weight_scale: torch.Tensor, alpha: float, output_dtype: torch.dtype, ) -> torch.Tensor: + # Validate inputs + if not input.is_cuda or not weight.is_cuda: + raise ValueError("Input and weight tensors must be on CUDA device") + if output_dtype != torch.bfloat16: + raise ValueError(f"Currently only bfloat16 output is supported, got {output_dtype}") + tuner = AutoTuner.get()
1045-1046: Replace assertion with explicit dtype check
Instead of relying on anassert, add a runtime check that raises a clear exception whenoutput_dtypeisn’ttorch.bfloat16. For example:- a_tensor, b_tensor, a_sf_tensor, b_sf_tensor, alpha, output_dtype = inputs - assert output_dtype == torch.bfloat16 + a_tensor, b_tensor, a_sf_tensor, b_sf_tensor, alpha, output_dtype = inputs + if output_dtype is not torch.bfloat16: + raise TypeError(f"Unsupported output_dtype {output_dtype!r}; only torch.bfloat16 is supported")This ensures unsupported dtypes produce a descriptive error rather than an assertion failure.
tests/unittest/_torch/thop/test_fp4_linear.py (5)
9-10: Commented imports and decorators should be cleaned upThe test file has commented-out imports and decorators that should either be removed or properly implemented.
Clean up the commented code:
-# from utils.util import skip_pre_blackwell - scaling_vector_size = 16 -# @skip_pre_blackwell @pytest.mark.parametrize("dtype", [torch.bfloat16])Also applies to: 14-15
15-20: Limited test coverage for dtypesThe test only covers
torch.bfloat16dtype. The TODO comment on line 19 questions whether float32 testing is needed, noting thatfp4_quantizesupports fp16, bf16, and fp8_e4m3.Consider expanding test coverage to include other supported dtypes:
-@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
21-23: Commented-out test values should be removedThe test contains commented-out hardcoded values that should be cleaned up.
Remove the commented lines:
- # SEQ_LEN = 10 - # HIDDEN_SIZE = 128 - # OUTPUT_SIZE = 256 torch.manual_seed(0) x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() x_sf_global = (448 * 6) / x.abs().max().float() - # x_sf_global = torch.tensor(1.0).cuda() w = torch.randn((OUTPUT_SIZE, HIDDEN_SIZE), dtype=dtype).cuda() w_sf_global = (448 * 6) / w.abs().max().float() - # w_sf_global = torch.tensor(1.0).cuda()Also applies to: 28-29, 32-33
36-39: Debug print statements in test codeMultiple debug print statements using "limin:" prefix should be converted to proper test logging or removed.
Consider using proper test logging or removing these prints:
- print(f"limin: w_fp4.shape = {w_fp4.shape}") - print(f"limin: w_fp4.dtype = {w_fp4.dtype}") - print(f"limin: w_sf_block.shape = {w_sf_block.shape}") - print(f"limin: w_sf_block.dtype = {w_sf_block.dtype}") + # Use pytest logging if needed for debugging + # pytest.param logging can be controlled via command line flagsAlso applies to: 53-53, 73-73, 91-91
94-95: Script execution block may cause issues in test discoveryThe
if __name__ == "__main__"block could interfere with pytest test discovery and execution.Consider using pytest's command-line interface instead:
-if __name__ == "__main__": - test_fp4_linear(torch.bfloat16, 128, 7680, 1536) +# Run with: pytest tests/unittest/_torch/thop/test_fp4_linear.py -vOr if you need a debug entry point, make it more explicit:
if __name__ == "__main__": # Debug entry point - use pytest for actual testing import sys sys.exit(pytest.main([__file__, "-v"]))
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py(2 hunks)tensorrt_llm/_torch/modules/linear.py(5 hunks)tests/unittest/_torch/thop/test_fp4_linear.py(4 hunks)
🧰 Additional context used
📓 Path-based instructions (4)
**/*
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Filenames compiled into a target must be case-insensitively unique
Files:
tensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/custom_ops/torch_custom_ops.pytests/unittest/_torch/thop/test_fp4_linear.py
**/*.{h,hpp,hh,hxx,cc,cpp,cxx,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Use spaces, not tabs; indent 4 spaces
Files:
tensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/custom_ops/torch_custom_ops.pytests/unittest/_torch/thop/test_fp4_linear.py
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: Code must target Python 3.8+
Indent with 4 spaces; do not use tabs (Python)
Maintain module namespace on import: prefer from package.subpackage import foo; use foo.Symbol()
Python filenames use snake_case
Python class names use PascalCase
Python functions and methods use snake_case
Python local variables use snake_case; if starting with a number concept, prefix with k (e.g., k_99th_percentile)
Python global variables use G_ prefix with UPPER_SNAKE_CASE
Python constants use UPPER_SNAKE_CASE
Avoid shadowing variables from outer scopes
Initialize all externally visible class members in init
For public interfaces, prefer docstrings over comments; comments should be for in-function or file-local interfaces
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes and variables inline with docstrings immediately after assignment
Avoid reflection when a non-reflective approach suffices
Limit except clauses to specific exceptions where possible
When using try/except for duck-typing, keep try body minimal and move logic to else
Files:
tensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/custom_ops/torch_custom_ops.pytests/unittest/_torch/thop/test_fp4_linear.py
**/*.{cpp,cc,cxx,h,hpp,hh,hxx,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/custom_ops/torch_custom_ops.pytests/unittest/_torch/thop/test_fp4_linear.py
🧬 Code graph analysis (3)
tensorrt_llm/_torch/modules/linear.py (1)
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)
cute_dsl_nvfp4_gemm_blackwell(1166-1186)
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (2)
tensorrt_llm/_torch/custom_ops/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent_pf.py (2)
Sm100BlockScaledPersistentDenseGemmKernel(142-2311)Sm100BlockScaledPersistentDenseGemmKernelWrapper(2329-2419)tensorrt_llm/quantization/utils/fp4_utils.py (1)
pad_up(22-23)
tests/unittest/_torch/thop/test_fp4_linear.py (3)
tensorrt_llm/models/modeling_utils.py (2)
QuantConfig(128-268)quant_algo(547-548)tensorrt_llm/_torch/modules/linear.py (2)
Linear(1511-1721)forward(1685-1714)tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (7)
forward(99-125)forward(336-349)forward(421-435)forward(508-527)forward(753-766)forward(832-849)forward(1009-1159)
| import cutlass | ||
| import cutlass.cute as cute | ||
| import torch | ||
| import triton # type: ignore[import] | ||
| from cuda import cuda | ||
| # todo: use optimized make_ptr | ||
| from cutlass.cute.runtime import make_ptr | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Temporary imports and code organization issues
The code has several issues that need to be addressed:
- A TODO comment on line 8 indicates that
make_ptrshould be replaced with an optimized version - The CuTe DSL imports and related code appear to be duplicated (lines 919-1205 seem to be a duplicate of the functionality)
Consider:
- Tracking the
make_ptroptimization in a proper issue tracker - Removing the duplicate code sections to maintain a single implementation
| def pad_up(x: int, y: int) -> int: | ||
| return ((x + y - 1) // y) * y | ||
|
|
There was a problem hiding this comment.
Duplicate pad_up function definition
The pad_up function is already defined in tensorrt_llm.quantization.utils.fp4_utils (as shown in the relevant code snippets) but is being redefined here. This creates unnecessary code duplication.
Remove this duplicate definition and import it from the existing module:
-def pad_up(x: int, y: int) -> int:
- return ((x + y - 1) // y) * y
-
+from tensorrt_llm.quantization.utils.fp4_utils import pad_up📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def pad_up(x: int, y: int) -> int: | |
| return ((x + y - 1) // y) * y | |
| -from tensorrt_llm/_torch/custom_ops/torch_custom_ops.py | |
| -def pad_up(x: int, y: int) -> int: | |
| - return ((x + y - 1) // y) * y | |
| from tensorrt_llm.quantization.utils.fp4_utils import pad_up |
🤖 Prompt for AI Agents
In tensorrt_llm/_torch/custom_ops/torch_custom_ops.py around lines 919 to 921,
there is a duplicate definition of pad_up; remove this local definition and
instead import pad_up from tensorrt_llm.quantization.utils.fp4_utils. Replace
the function block with a single import statement from that module (ensure the
import is placed with other imports and respects existing import ordering), and
run tests/linters to confirm no unused imports or name conflicts.
| class CuteDSLNVFP4BlackwellLinear(TunableRunner): | ||
| kernel_dict = dict() | ||
|
|
||
| class fp8SwapABGemmRunner(TunableRunner): | ||
| # TODO: | ||
| tuning_config = TuningConfig( | ||
| dynamic_tensor_specs=(DynamicTensorSpec( | ||
| 0, 0, fp8_swap_ab_gen_tuning_buckets), ), | ||
| tune_max_num_tokens=4096, | ||
| 0, 0, get_last_power_of_2_num_tokens_buckets, | ||
| last_positive_power_of_2), ), | ||
| constraint_specs=(ConstraintSpec(2, 0, fp4_scale_infer_shape), ), | ||
| ) | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Incomplete tuning configuration and hard-coded tactics
The CuteDSLNVFP4BlackwellLinear class has several issues:
- TODO comment on line 940 indicates incomplete tuning configuration
- Hard-coded tactics in
get_valid_tactics(lines 971-992) with limited options - The fallback tactic in
forward(lines 1034-1043) uses different values than the generated tactics
Consider:
- Completing the tuning configuration implementation
- Making the tactics configurable rather than hard-coded
- Ensuring consistency between generated tactics and fallback values
| print(f"limin: m = {m}, real_k = {real_k}, n = {n}") | ||
| print(f"limin: sf_m = {sf_m}, sf_k = {sf_k}, sf_n = {sf_n}") | ||
| print( | ||
| f"limin: a_tensor.shape = {a_tensor.shape}, b_tensor.shape = {b_tensor.shape}" | ||
| ) | ||
| print( | ||
| f"limin: a_sf_tensor.shape = {a_sf_tensor.shape}, b_sf_tensor.shape = {b_sf_tensor.shape}" | ||
| ) | ||
| print(f"limin: alpha = {alpha}") | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Remove debug print statements before production
Multiple debug print statements are present throughout the code that should be removed or converted to proper logging before merging to production.
Remove or convert these debug prints to proper logging:
- print(f"limin: m = {m}, real_k = {real_k}, n = {n}")
- print(f"limin: sf_m = {sf_m}, sf_k = {sf_k}, sf_n = {sf_n}")
- print(
- f"limin: a_tensor.shape = {a_tensor.shape}, b_tensor.shape = {b_tensor.shape}"
- )
- print(
- f"limin: a_sf_tensor.shape = {a_sf_tensor.shape}, b_sf_tensor.shape = {b_sf_tensor.shape}"
- )
- print(f"limin: alpha = {alpha}")
+ # Use proper logging if debug output is needed
+ # logger.debug(f"Tensor dimensions: m={m}, real_k={real_k}, n={n}")Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In tensorrt_llm/_torch/custom_ops/torch_custom_ops.py around lines 1054 to 1063
there are leftover debug print statements that should not be in production;
replace them with an appropriate logger (e.g., module-level logger =
logging.getLogger(__name__)) and use logger.debug(...) with the same message
content or remove them entirely if not needed, ensuring logging is imported and
configured where appropriate and sensitive data is not logged.
| a_sf_ptr = make_ptr( | ||
| # cutlass.Float8E8M0FNU, | ||
| cutlass.Float8E4M3FN, | ||
| a_sf_tensor.data_ptr(), | ||
| cute.AddressSpace.gmem, | ||
| assumed_align=16, | ||
| ) | ||
| # TODO: add swap_ab=tactic == 0 to detemrmine the swap_ab value | ||
| # Treat the default tactic=-1 as swap_ab=False | ||
| deep_gemm.fp8_gemm_nt( | ||
| (a, a_sf), | ||
| (weight, weight_scale), | ||
| output, | ||
| disable_ue8m0_cast=self.disable_ue8m0_cast, | ||
| b_sf_ptr = make_ptr( | ||
| # cutlass.Float8E8M0FNU, | ||
| cutlass.Float8E4M3FN, | ||
| b_sf_tensor.data_ptr(), | ||
| cute.AddressSpace.gmem, | ||
| assumed_align=16, | ||
| ) |
There was a problem hiding this comment.
Inconsistent data types for scale tensors
The code uses cutlass.Float8E4M3FN for both a_sf_ptr and b_sf_ptr (lines 1080, 1087), but comments suggest it should be cutlass.Float8E8M0FNU (lines 1079, 1086). This inconsistency could lead to incorrect computations.
Verify and fix the correct data type for scale tensors:
a_sf_ptr = make_ptr(
- # cutlass.Float8E8M0FNU,
- cutlass.Float8E4M3FN,
+ cutlass.Float8E8M0FNU, # Use the correct scale type
a_sf_tensor.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=16,
)
b_sf_ptr = make_ptr(
- # cutlass.Float8E8M0FNU,
- cutlass.Float8E4M3FN,
+ cutlass.Float8E8M0FNU, # Use the correct scale type
b_sf_tensor.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=16,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| a_sf_ptr = make_ptr( | |
| # cutlass.Float8E8M0FNU, | |
| cutlass.Float8E4M3FN, | |
| a_sf_tensor.data_ptr(), | |
| cute.AddressSpace.gmem, | |
| assumed_align=16, | |
| ) | |
| # TODO: add swap_ab=tactic == 0 to detemrmine the swap_ab value | |
| # Treat the default tactic=-1 as swap_ab=False | |
| deep_gemm.fp8_gemm_nt( | |
| (a, a_sf), | |
| (weight, weight_scale), | |
| output, | |
| disable_ue8m0_cast=self.disable_ue8m0_cast, | |
| b_sf_ptr = make_ptr( | |
| # cutlass.Float8E8M0FNU, | |
| cutlass.Float8E4M3FN, | |
| b_sf_tensor.data_ptr(), | |
| cute.AddressSpace.gmem, | |
| assumed_align=16, | |
| ) | |
| a_sf_ptr = make_ptr( | |
| cutlass.Float8E8M0FNU, # Use the correct scale type | |
| a_sf_tensor.data_ptr(), | |
| cute.AddressSpace.gmem, | |
| assumed_align=16, | |
| ) | |
| b_sf_ptr = make_ptr( | |
| cutlass.Float8E8M0FNU, # Use the correct scale type | |
| b_sf_tensor.data_ptr(), | |
| cute.AddressSpace.gmem, | |
| assumed_align=16, | |
| ) |
🤖 Prompt for AI Agents
In tensorrt_llm/_torch/custom_ops/torch_custom_ops.py around lines 1078 to 1091,
the scale pointer types for a_sf_ptr and b_sf_ptr are inconsistent (currently
cutlass.Float8E4M3FN) while the preceding comments indicate they should be
cutlass.Float8E8M0FNU; update both make_ptr calls to use cutlass.Float8E8M0FNU
so the scale tensors use the correct Float8 format, keep the same data_ptr(),
AddressSpace.gmem and assumed_align=16 arguments unchanged, and remove or update
the stale comments if needed.
| # output = torch.ops.trtllm.nvfp4_gemm(act_fp4, module.weight, act_sf, | ||
| # module.weight_scale, module.alpha, | ||
| # module.dtype) | ||
| print( | ||
| f"limin: linear layer, act_fp4.shape = {act_fp4.shape}, module.weight.shape = {module.weight.shape}, act_sf.shape = {act_sf.shape}, module.weight_scale.shape = {module.weight_scale.shape}, module.alpha.shape = {module.alpha.shape}" | ||
| ) | ||
| print(f"limin: module.alpha = {module.alpha}") | ||
| print(f"limin: module.alpha.dtype = {module.alpha.dtype}") | ||
| print(f"limin: module.scalar_alpha = {module.scalar_alpha}") | ||
| output = torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell( | ||
| act_fp4, module.weight, act_sf, module.weight_scale, | ||
| module.scalar_alpha, module.dtype) | ||
| if bias is not None: |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Debug prints and commented code should be removed
The production code contains:
- Commented-out code for the previous
nvfp4_gemmimplementation (lines 714-716) - Multiple debug print statements (lines 717-722)
- Debug print for weight loading (line 1720)
These should be cleaned up before merging.
Remove debug code and clean up:
- # output = torch.ops.trtllm.nvfp4_gemm(act_fp4, module.weight, act_sf,
- # module.weight_scale, module.alpha,
- # module.dtype)
- print(
- f"limin: linear layer, act_fp4.shape = {act_fp4.shape}, module.weight.shape = {module.weight.shape}, act_sf.shape = {act_sf.shape}, module.weight_scale.shape = {module.weight_scale.shape}, module.alpha.shape = {module.alpha.shape}"
- )
- print(f"limin: module.alpha = {module.alpha}")
- print(f"limin: module.alpha.dtype = {module.alpha.dtype}")
- print(f"limin: module.scalar_alpha = {module.scalar_alpha}")
output = torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell(
act_fp4, module.weight, act_sf, module.weight_scale,
module.scalar_alpha, module.dtype)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # output = torch.ops.trtllm.nvfp4_gemm(act_fp4, module.weight, act_sf, | |
| # module.weight_scale, module.alpha, | |
| # module.dtype) | |
| print( | |
| f"limin: linear layer, act_fp4.shape = {act_fp4.shape}, module.weight.shape = {module.weight.shape}, act_sf.shape = {act_sf.shape}, module.weight_scale.shape = {module.weight_scale.shape}, module.alpha.shape = {module.alpha.shape}" | |
| ) | |
| print(f"limin: module.alpha = {module.alpha}") | |
| print(f"limin: module.alpha.dtype = {module.alpha.dtype}") | |
| print(f"limin: module.scalar_alpha = {module.scalar_alpha}") | |
| output = torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell( | |
| act_fp4, module.weight, act_sf, module.weight_scale, | |
| module.scalar_alpha, module.dtype) | |
| if bias is not None: | |
| output = torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell( | |
| act_fp4, module.weight, act_sf, module.weight_scale, | |
| module.scalar_alpha, module.dtype) | |
| if bias is not None: |
🤖 Prompt for AI Agents
In tensorrt_llm/_torch/modules/linear.py around lines 714 to 726 (and the debug
print at line 1720), remove the commented-out nvfp4_gemm code and all debug
print statements so the production code is clean; keep the intended call to
torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell and subsequent logic (including
bias handling) unchanged, and delete the extra prints related to shapes,
module.alpha, and module.scalar_alpha as well as the weight-loading debug print
at line 1720.
| output = torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell( | ||
| act_fp4, module.weight, act_sf, module.weight_scale, | ||
| module.scalar_alpha, module.dtype) |
There was a problem hiding this comment.
Type inconsistency in kernel invocation
The new kernel cute_dsl_nvfp4_gemm_blackwell expects alpha as a float parameter (per the signature in torch_custom_ops.py), but module.scalar_alpha might not be initialized if weights haven't been loaded yet.
Add defensive check to ensure scalar_alpha is initialized:
+ if not hasattr(module, 'scalar_alpha'):
+ raise RuntimeError("scalar_alpha not initialized. Ensure weights are loaded before forward pass.")
output = torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell(
act_fp4, module.weight, act_sf, module.weight_scale,
module.scalar_alpha, module.dtype)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| output = torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell( | |
| act_fp4, module.weight, act_sf, module.weight_scale, | |
| module.scalar_alpha, module.dtype) | |
| # defensive: ensure scalar_alpha was set (weights loaded) before calling the kernel | |
| if not hasattr(module, 'scalar_alpha'): | |
| raise RuntimeError("scalar_alpha not initialized. Ensure weights are loaded before forward pass.") | |
| output = torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell( | |
| act_fp4, module.weight, act_sf, module.weight_scale, | |
| module.scalar_alpha, module.dtype) |
🤖 Prompt for AI Agents
In tensorrt_llm/_torch/modules/linear.py around lines 723 to 725, the kernel
call passes module.scalar_alpha which may be uninitialized; before calling
torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell ensure module.scalar_alpha exists
and is a float (e.g. set a default like 1.0 if missing or derive from module
attributes), explicitly cast it to float, and then pass that safe float value
into the kernel; if scalar_alpha being missing is an unexpected state, raise a
clear exception early instead of passing None.
| module.scalar_alpha = alpha.item() | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
New scalar_alpha attribute lacks documentation
A new scalar_alpha attribute is being added alongside the existing alpha tensor parameter, but:
- The purpose and difference between
alphaandscalar_alphais not documented - Both attributes store the same value but in different formats (tensor vs scalar)
Add documentation and consider if both are necessary:
copy_weight(module.alpha, alpha)
+ # Store scalar version for CuTe DSL kernel which requires Python float
module.scalar_alpha = alpha.item()Consider whether maintaining both alpha (tensor) and scalar_alpha (float) is necessary, or if the conversion could be done at the point of use.
Also applies to: 809-810, 831-832
🤖 Prompt for AI Agents
tensorrt_llm/_torch/modules/linear.py lines ~791-792 (also review 809-810 and
831-832): the code adds a new scalar_alpha attribute mirroring the alpha tensor
but lacks documentation and may be redundant; either document the purpose and
difference between alpha (torch.Tensor) and scalar_alpha (float) or remove
scalar_alpha and convert alpha.item() where a plain float is needed; update
docstring/comments near the class/constructor to explain why both
representations are kept (if kept), ensure scalar_alpha is kept in sync with
alpha (or made a property that returns alpha.item()), and remove duplicate
storage if unnecessary by converting at use sites.
| print(f"limin: weight_mode = {weight_mode}") | ||
| self.quant_method.load_weights(self, weights, weight_mode) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Debug print statement in production code
Debug print statement should be removed before merging.
- print(f"limin: weight_mode = {weight_mode}")
self.quant_method.load_weights(self, weights, weight_mode)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| print(f"limin: weight_mode = {weight_mode}") | |
| self.quant_method.load_weights(self, weights, weight_mode) | |
| self.quant_method.load_weights(self, weights, weight_mode) |
🤖 Prompt for AI Agents
In tensorrt_llm/_torch/modules/linear.py around lines 1720-1721 there is a
leftover debug print: remove the print(f"limin: weight_mode = {weight_mode}")
statement and ensure no other stray debug prints remain; simply call
self.quant_method.load_weights(self, weights, weight_mode) as the production
behavior.
| output = l_fp4.forward(x) | ||
|
|
There was a problem hiding this comment.
Duplicate forward pass call
The forward pass is called twice - once within the autotune context (line 76) and once without (line 78). The second call appears to overwrite the first result without any assertion or comparison.
This looks like it might be unintentional. Either remove the duplicate or add a comment explaining why both are needed:
with torch.inference_mode(), autotune():
output = l_fp4.forward(x)
- output = l_fp4.forward(x)
+ # Run again without autotune to get the final output
+ # output = l_fp4.forward(x)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| output = l_fp4.forward(x) | |
| with torch.inference_mode(), autotune(): | |
| output = l_fp4.forward(x) | |
| # Run again without autotune to get the final output | |
| # output = l_fp4.forward(x) |
🤖 Prompt for AI Agents
In tests/unittest/_torch/thop/test_fp4_linear.py around lines 78-79, the test
calls l_fp4.forward(x) a second time (outside the autotune context) which
overwrites the prior result from the autotune invocation; remove the duplicate
call or, if both outputs are intentionally required, add a clarifying comment
and a comparison/assertion between the two outputs to justify the second call
(e.g., assert equality or expected difference) so the test behavior is explicit.
Signed-off-by: Mindy Li <[email protected]>
Signed-off-by: Mindy Li <[email protected]>
|
|
||
|
|
||
| # add nvfp4 cute dsl gemm | ||
| class CuteDSLNVFP4BlackwellLinear(TunableRunner): |
There was a problem hiding this comment.
I suggest adding the cutedsl op to a standalone file for clearer and easier maintenance.
|
|
||
|
|
||
| # https://gitlab-master.nvidia.com/dlarch-fastkernels/dynamic-kernel-generator/-/merge_requests/11520 | ||
| class CuptiProfiler: |
There was a problem hiding this comment.
Is there a simplified way to use it? Or do we need to define this type of "general" profiler tool every time we use the API?
2 participants
Summary by CodeRabbit
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
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.
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.