[None][feat] Unify nvfp4 gemm backend

[Wong4j]

Summary by CodeRabbit

• New Features

  • Introduced unified NVFP4 GEMM interface with automatic or manual backend selection (CUTLASS, cuBLASLt, CuteDSL).

• Deprecations

  • Deprecated existing NVFP4 entry points; users should migrate to the new unified interface.

• Breaking Changes

  • Linear module constructor now uses nvfp4_backend parameter instead of individual backend flags.
  • All original nvfp4_gemm in the repo have been renamed to nvfp4_gemm_cutlass. The new nvfp4_gemm now serves as a unified API and accepts a backend (options: auto, cutlass, cublaslt, cutedsl) as an input parameter. (updated on November 26)
  • The input argument alpha in cutedsl has been changed from a host-side float to a device-side tensor to stay consistent with other backends (cutlass, cublaslt) and support the unified interface. (updated on November 26)

• Tests

  • Added comprehensive test coverage for unified backend selection and tactic handling.

Description

This PR introduces a unified NVFP4 GEMM interface that consolidates multiple backend implementations (CUTLASS, cuBLASLt, and CuteDSL) into a single, easy-to-use API with automatic performance optimization.

Introduced torch.ops.trtllm.nvfp4_gemm_unified with a backend parameter supporting:

• "auto" (default): Automatically profiles all available backends and selects the best one
• "cutlass": Force CUTLASS backend
• "cublaslt": Force cuBLASLt backend
• "cutedsl": Force CuteDSL backend

Example:

output = torch.ops.trtllm.nvfp4_gemm_unified(
    act_fp4, weight, act_sf, weight_scale, alpha, 
    output_dtype, backend='auto'
)

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.

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 the stage-list parameter 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.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip 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-pipeline

Reuse 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.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes consolidate multiple NVFP4 GEMM backends (CUTLASS, cuBLASLt, CuteDSL) into a unified entry point nvfp4_gemm_unified with automatic or explicit backend selection. Existing backend-specific functions and Boolean flags are deprecated with warnings, while the Linear module is refactored to replace multiple Boolean parameters with a single string-based nvfp4_backend parameter for runtime backend selection.

Changes

Cohort / File(s)	Summary
Deprecation notices tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py, tensorrt_llm/_torch/custom_ops/torch_custom_ops.py	Added deprecation docstrings and logger.warning_once to cute_dsl_nvfp4_gemm_blackwell, nvfp4_gemm_cublaslt, and nvfp4_gemm functions, directing users to the new nvfp4_gemm_unified entry point.
Unified NVFP4 interface tensorrt_llm/_torch/custom_ops/torch_custom_ops.py	Introduced new public function nvfp4_gemm_unified with auto/explicit backend selection (CUTLASS, cuBLASLt, CuteDSL). Added CuteDSLNVFP4Wrapper class to normalize CuteDSL backend interface. Added conditional imports and capability checks (IS_CUBLASLT_AVAILABLE, IS_CUTLASS_DSL_AVAILABLE).
Linear module refactoring tensorrt_llm/_torch/modules/linear.py	Replaced Boolean backend flags (use_cute_dsl_nvfp4_blockscaling_mm, use_cublaslt_nvfp4_blockscaling_mm) with single string parameter nvfp4_backend (default "auto") in Linear class constructor. Consolidated backend selection branching logic to use unified nvfp4_gemm_unified call. Updated NVFP4LinearMethod to propagate nvfp4_backend parameter.
Test suite expansions tests/unittest/_torch/thop/parallel/test_fp4_linear.py	Updated existing tests to use nvfp4_backend='cutedsl' instead of Boolean flags. Added comprehensive test suite for nvfp4_gemm_unified including auto-backend selection, explicit backend testing (CUTLASS, cuBLASLt, CuteDSL), tactic discovery/replay, and autotuning validation. Included hardware capability and SM version gates for Blackwell-specific tests.

Sequence Diagram(s)

sequenceDiagram
    participant App as Application Code
    participant Unified as nvfp4_gemm_unified
    participant Router as Backend Router
    participant CUTLASS as CUTLASS Backend
    participant cuBLASLt as cuBLASLt Backend
    participant CuteDSL as CuteDSL Backend
    participant Wrapper as CuteDSLNVFP4Wrapper

    App->>Unified: nvfp4_gemm_unified(..., backend="auto"|"cutlass"|"cublaslt"|"cutedsl")
    Unified->>Router: Determine backend availability & select runner
    
    alt backend == "auto"
        Router->>Router: Check availability & select default
    else backend == explicit
        Router->>Router: Validate backend availability
    end
    
    alt Selected: CUTLASS
        Router->>CUTLASS: Execute GEMM
        CUTLASS-->>Unified: Result
    else Selected: cuBLASLt
        Router->>cuBLASLt: Execute GEMM
        cuBLASLt-->>Unified: Result
    else Selected: CuteDSL
        Router->>Wrapper: Create/call CuteDSLNVFP4Wrapper
        Wrapper->>CuteDSL: Execute via normalized interface
        CuteDSL-->>Wrapper: Result
        Wrapper-->>Unified: Adapted result
    end
    
    Unified-->>App: Output tensor

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

• New public API surface: nvfp4_gemm_unified function and CuteDSLNVFP4Wrapper class require careful validation of backend selection logic, input validation, and error handling.
• Refactored module interface: Linear class constructor signature changed from multiple Boolean flags to a string parameter; verify all initialization paths, weight loading, and backward compatibility considerations.
• Deprecation propagation: Ensure deprecation warnings are correctly routed and logged without disrupting functionality in existing code paths.
• Multi-backend routing logic: The backend selection and runner initialization in nvfp4_gemm_unified and wrapper class involves conditional imports and runtime capability checks that need verification across different hardware/software configurations.
• Test coverage heterogeneity: New tests span multiple backend implementations, autotuning flows, and hardware gates; each test variant may require separate reasoning.

Pre-merge checks and finishing touches

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.
Description check	❓ Inconclusive	PR description includes title, description explaining the unified interface, and partially completed checklist, but Test Coverage section is empty.	Complete the Test Coverage section by listing the specific test files and test cases that validate the new nvfp4_gemm_unified functionality and backend selection logic.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: unifying the NVFP4 GEMM backend into a single interface, which aligns with the raw summary showing consolidation of CUTLASS, cuBLASLt, and CuteDSL backends.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e822184 and ddf6d3c.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py (1 hunks)
• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (4 hunks)
• tensorrt_llm/_torch/modules/linear.py (3 hunks)
• tests/unittest/_torch/thop/parallel/test_fp4_linear.py (4 hunks)

🧰 Additional context used

📓 Path-based instructions (3)

**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Use only spaces, no tabs; indent with 4 spaces.

Files:

• tensorrt_llm/_torch/modules/linear.py
• tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
• tests/unittest/_torch/thop/parallel/test_fp4_linear.py

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Python code must target Python 3.8+.
Indent Python code with 4 spaces; do not use tabs.
Maintain module namespace when importing; prefer 'from package.subpackage import foo' then 'foo.SomeClass()' instead of importing the class directly.
Python filenames should be snake_case (e.g., some_file.py).
Python classes use PascalCase names.
Functions and methods use snake_case names.
Local variables use snake_case; prefix 'k' for variables that start with a number (e.g., k_99th_percentile).
Global variables use upper SNAKE_CASE prefixed with 'G' (e.g., G_MY_GLOBAL).
Constants use upper SNAKE_CASE (e.g., MY_CONSTANT).
Avoid shadowing variables from an outer scope.
Initialize all externally visible members of a class in the constructor.
Prefer docstrings for interfaces that may be used outside a file; comments for in-function or file-local interfaces.
Use Google-style docstrings for classes and functions (Sphinx-parsable).
Document attributes and variables inline so they render under the class/function docstring.
Avoid reflection when a simpler, explicit approach suffices (e.g., avoid dict(**locals()) patterns).
In try/except, catch the most specific exceptions possible.
For duck-typing try/except, keep the try body minimal and use else for the main logic.

Files:

• tensorrt_llm/_torch/modules/linear.py
• tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
• tests/unittest/_torch/thop/parallel/test_fp4_linear.py

**/*.{cpp,cxx,cc,h,hpp,hh,hxx,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Prepend the NVIDIA Apache-2.0 copyright header with current year to the top of all source files (e.g., .cpp, .h, .cu, .py).

Files:

• tensorrt_llm/_torch/modules/linear.py
• tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
• tests/unittest/_torch/thop/parallel/test_fp4_linear.py

🧠 Learnings (8)

📓 Common learnings

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1475-1480
Timestamp: 2025-08-21T02:39:12.009Z
Learning: The min latency mode functionality in TensorRT-LLM MOE kernels (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu) is deprecated and no longer being maintained/updated, as confirmed by djns99. Bug reports and optimization suggestions for the computeStridesTmaWarpSpecializedLowLatencyKernel and related min latency code paths should be deprioritized.

📚 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:

• tensorrt_llm/_torch/modules/linear.py
• tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

📚 Learning: 2025-08-19T12:45:11.997Z

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:0-0
Timestamp: 2025-08-19T12:45:11.997Z
Learning: In tensorrt_llm/_torch/pyexecutor/model_engine.py, DoRA (Delta Orthogonal Rank Adaptation) functionality was removed from the PyTorch flow to eliminate issues with inverted DoRA detection logic. The original is_dora condition was checking if scaling_vec_pointer == 0, which was potentially incorrect.

Applied to files:

• tensorrt_llm/_torch/modules/linear.py

📚 Learning: 2025-08-21T21:48:35.135Z

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/fusion/sm90_visitor_scatter.hpp:399-417
Timestamp: 2025-08-21T21:48:35.135Z
Learning: CUTLASS extensions in TensorRT-LLM (located under cpp/tensorrt_llm/cutlass_extensions/) are designed to integrate with and extend functionality in the external CUTLASS repository. When analyzing these extensions, their consumers and functionality wiring may exist in the CUTLASS codebase rather than within TensorRT-LLM itself.

Applied to files:

• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

📚 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:

• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

📚 Learning: 2025-08-26T09:49:04.956Z

Learnt from: pengbowang-nv
Repo: NVIDIA/TensorRT-LLM PR: 7192
File: tests/integration/test_lists/test-db/l0_dgx_b200.yml:56-72
Timestamp: 2025-08-26T09:49:04.956Z
Learning: In TensorRT-LLM test configuration files, the test scheduling system handles wildcard matching with special rules that prevent duplicate test execution even when the same tests appear in multiple yaml files with overlapping GPU wildcards (e.g., "*b200*" and "*gb200*").

Applied to files:

• tests/unittest/_torch/thop/parallel/test_fp4_linear.py

📚 Learning: 2025-09-09T09:40:45.658Z

Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7645
File: tests/integration/test_lists/qa/llm_function_core.txt:648-648
Timestamp: 2025-09-09T09:40:45.658Z
Learning: In TensorRT-LLM test lists, it's common and intentional for the same test to appear in multiple test list files when they serve different purposes (e.g., llm_function_core.txt for comprehensive core functionality testing and llm_function_core_sanity.txt for quick sanity checks). This duplication allows tests to be run in different testing contexts.

Applied to files:

• tests/unittest/_torch/thop/parallel/test_fp4_linear.py

📚 Learning: 2025-07-28T17:06:08.621Z

Learnt from: moraxu
Repo: NVIDIA/TensorRT-LLM PR: 6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.

Applied to files:

• tests/unittest/_torch/thop/parallel/test_fp4_linear.py

🧬 Code graph analysis (4)

tensorrt_llm/_torch/modules/linear.py (1)

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)

• nvfp4_gemm_unified (734-881)

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py (1)

tensorrt_llm/logger.py (1)

• warning_once (135-136)

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (4)

tensorrt_llm/quantization/utils/fp4_utils.py (2)

• pad_up (22-23)
• FP4GemmType (26-28)

tensorrt_llm/_torch/autotuner.py (10)

• AutoTuner (514-1177)
• ConstraintSpec (39-49)
• DynamicTensorSpec (23-35)
• OptimizationProfile (127-142)
• TunableRunner (153-209)
• TuningConfig (53-101)
• get_valid_tactics (156-174)
• forward (180-206)
• get (545-548)
• choose_one (623-778)

tensorrt_llm/logger.py (2)

• warning_once (135-136)
• debug (144-145)

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py (4)

• CuteDSLNVFP4BlackwellLinear (31-294)
• get_valid_tactics (61-135)
• forward (146-294)
• _ (337-351)

tests/unittest/_torch/thop/parallel/test_fp4_linear.py (2)

tensorrt_llm/_torch/autotuner.py (2)

• AutoTuner (514-1177)
• autotune (213-245)

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)

• nvfp4_gemm_unified (734-881)

🪛 Gitleaks (8.28.0)

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

[high] 495-495: Detected a Generic API Key, potentially exposing access to various services and sensitive operations.

(generic-api-key)

---

[high] 551-551: Detected a Generic API Key, potentially exposing access to various services and sensitive operations.

(generic-api-key)

🪛 Ruff (0.14.3)

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

697-697: Do not catch blind exception: Exception

(BLE001)

---

705-705: Do not catch blind exception: Exception

(BLE001)

---

716-716: Unpacked variable alpha is never used

Prefix it with an underscore or any other dummy variable pattern

(RUF059)

---

719-721: Avoid specifying long messages outside the exception class

(TRY003)

---

786-787: Avoid specifying long messages outside the exception class

(TRY003)

---

841-843: Avoid specifying long messages outside the exception class

(TRY003)

---

854-856: Avoid specifying long messages outside the exception class

(TRY003)

---

865-865: Do not catch blind exception: Exception

(BLE001)

---

866-868: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling

(B904)

---

866-868: Avoid specifying long messages outside the exception class

(TRY003)

---

871-877: Avoid specifying long messages outside the exception class

(TRY003)

---

888-888: Unused function argument: act_sf

(ARG001)

---

889-889: Unused function argument: weight_scale

(ARG001)

---

890-890: Unused function argument: alpha

(ARG001)

---

892-892: Unused function argument: to_userbuffers

(ARG001)

---

893-893: Unused function argument: backend

(ARG001)

[Wong4j]

/bot run

[tensorrt-cicd]

PR_Github #24362 [ run ] triggered by Bot. Commit: 491d2ea

[Wong4j]

In single tests, multiple backends may be tested, which can lead to the following situation. For example:

The first ut, using “auto”:

runners = [
    FP4GemmRunner(...),           # idx=0  (CUTLASS)
    CublasLtFP4GemmRunner(...),   # idx=1  (cuBLASLt)
    CuteDSLNVFP4Wrapper(...),     # idx=2  (CuteDSL)
]

The second ut, forcing “cublaslt”:

runners = [
    CublasLtFP4GemmRunner(...),   # idx=0  (only this one!)
]

In this case, the cached idx becomes incorrect, leading to an IndexError: list index out of range.
So I modified the autotuner.py code, but I’m not sure if this will cause any side effects.
Could you please take a look? @rosenrodt

[hyukn]

Just put some general ideas about how to define a nested tuning op. This might be clearer and more expandable.

[tensorrt-cicd]

PR_Github #24362 [ run ] completed with state SUCCESS. Commit: 491d2ea
/LLM/main/L0_MergeRequest_PR pipeline #18385 completed with status: 'SUCCESS'

[kaiyux]

This is critical and very helpful change for DS R1 performance - we probably need to verify the performance before merging it to avoid perf regression.

[rosenrodt]

This is critical and very helpful change for DS R1 performance - we probably need to verify the performance before merging it to avoid perf regression.

@kaiyux Doesn't DS-R1 NVFP4 checkpoint actually use very few FP4 GEMMs? I see most of the GEMMs in the up/down projection is still in BF16. And while MoE is indeed NVFP4, this PR touches only the dense GEMMs, not MoE grouped GEMMs

[kaiyux]

This is critical and very helpful change for DS R1 performance - we probably need to verify the performance before merging it to avoid perf regression.

@kaiyux Doesn't DS-R1 NVFP4 checkpoint actually use very few FP4 GEMMs? I see most of the GEMMs in the up/down projection is still in BF16. And while MoE is indeed NVFP4, this PR touches only the dense GEMMs, not MoE grouped GEMMs

We're currently working on moving more dense gemms to nvfp4, it will helpfully be landed soon. (that should not block this PR though)

[hyukn]

To simplify the nested tuning process, we want :

• The inner op is not forced to have forward and get_valid_tactics to be implemented (whether it is a tunable one or not).
• The interface of the inner op is not required to be the same as any other candidate op. (wrapper is not necessary).

This commit might be helpful to illustrate the idea: hyukn@b5d3b4c

I just took minutes to write the draft commit based on @Wong4j 's original changes, but without any local validation. Maybe @Wong4j can try this idea to see if it achieves the same tuning purpose as the original code. Truly appreciate.

[Wong4j]

To simplify the nested tuning process, we want :

• The inner op is not forced to have forward and get_valid_tactics to be implemented (whether it is a tunable one or not).
• The interface of the inner op is not required to be the same as any other candidate op. (wrapper is not necessary).

This commit might be helpful to illustrate the idea: hyukn@b5d3b4c

I just took minutes to write the draft commit based on @Wong4j 's original changes, but without any local validation. Maybe @Wong4j can try this idea to see if it achieves the same tuning purpose as the original code. Truly appreciate.

Sure, I will try it.

[hyukn]

Sure, I will try it.
Thanks a lot for the effort.

I have just pushed another commit to clean the code and make UT work. Because this is the first practical nested tuning process, it is a good opportunity to explore if we can do things in a tidy and efficient way. Some concerns:

• AutoTuner will do redundant profiling generation, which introduces a lot of host overhead even if the inputs are already in the profiling cache. This will destroy the outer tuning. Thus, I did some minor changes to the AutoTuner to eliminate this unacceptable
  overhead.
• When doing nested tuning, capture-replay mechanisms will encounter some issues. I guess it might be the status of the counter that is shared among all the ops, which will be incorrectly updated for the nested tuning process. Therefore, I just disabled that part in the UT for now. Maybe we can do some extra work to make this correct later. cc @rosenrodt
• I suggest @Wong4j observing the final profiling cache status. It should contain all the results for each low-level NVFP4 gemm tuning result, followed by the unified op tuning result.

Hope this will help.

[hyukn]

Hi @Wong4j. Thanks a lot for the effort!
I just moved the common code changes in AutoTuner to a standalone PR #9348 because it might be required by other tunable op as well.

[Wong4j]

@liji-nv

How does the CuteDSL and cuda_core DSL supports to_userbuffers=True? I do not see related change. to_userbuffers=1 means the op must writes to a UB buffer. Or the following allreduce would fail.

cuda_core:

TensorRT-LLM/cpp/tensorrt_llm/thop/cudaNvfp4MM.cpp

Lines 120 to 122 in 9d2df04

	if (to_userbuffers)
	{
	out = torch_ext::create_userbuffers_tensor(output_size, out_dtype_).first;

cubalslt:

TensorRT-LLM/cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp

Lines 189 to 191 in 9d2df04

	if (to_userbuffers)
	{
	out = torch_ext::create_userbuffers_tensor(output_size, mOutputDtype).first;

cutlass:

TensorRT-LLM/cpp/tensorrt_llm/thop/fp4Gemm.cpp

Lines 198 to 200 in 9d2df04

	out = torch_ext::create_userbuffers_tensor(out_shape, out_dtype.value()).first;
	}
	else

For CuteDSL, I add to_userbuffers here:
https://github.com/Wong4j/TensorRT-LLM/blob/205d297bba85f134f896d00a4db8f450dd6dd031/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py#L242

[Wong4j]

/bot run

[tensorrt-cicd]

PR_Github #26410 [ run ] triggered by Bot. Commit: a70c06d

[tensorrt-cicd]

PR_Github #26410 [ run ] completed with state SUCCESS. Commit: a70c06d
/LLM/main/L0_MergeRequest_PR pipeline #20066 completed with status: 'FAILURE'

[tensorrt-cicd]

PR_Github #26426 [ run ] completed with state FAILURE. Commit: d06dfa0
/LLM/main/L0_MergeRequest_PR pipeline #20082 completed with status: 'FAILURE'

[Wong4j]

/bot run

[tensorrt-cicd]

PR_Github #26429 [ run ] triggered by Bot. Commit: 906391b

[tensorrt-cicd]

PR_Github #26429 [ run ] completed with state SUCCESS. Commit: 906391b
/LLM/main/L0_MergeRequest_PR pipeline #20085 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.