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

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

• [x] 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[bot]]

📝 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

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

[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

[tensorrt-cicd]

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

[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: https://github.com/hyukn/TensorRT-LLM/commit/b5d3b4c52e00c03884f8fc5c202dedf754753261

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]

I benchmarked the best NVFP4 GEMM performance autotuned by each backend on different shapes. After unifying the interface, the globally optimal tactic can always be selected.

I also compared the Qwen2.5-72B NVFP4 performance using trtllm-bench between the default CUTLASS backend on the main branch and the unified interface implemented in this PR, and observed a noticeable end-to-end speedup. (with concurrency ranging from 128 to 2048 in the figure)

Do you think we need any additional performance testing? @kaiyux @hyukn @rosenrodt

[Wong4j]

/bot run

[tensorrt-cicd]

PR_Github #25648 [ run ] triggered by Bot. Commit: 753017e

[tensorrt-cicd]

PR_Github #25648 [ run ] completed with state SUCCESS. Commit: 753017e /LLM/main/L0_MergeRequest_PR pipeline #19436 completed with status: 'FAILURE'

[Wong4j]

/bot run

[tensorrt-cicd]

PR_Github #25697 [ run ] triggered by Bot. Commit: 67e85e1

[tensorrt-cicd]

PR_Github #25697 [ run ] completed with state SUCCESS. Commit: 67e85e1 /LLM/main/L0_MergeRequest_PR pipeline #19479 completed with status: 'SUCCESS' Pipeline passed with automatic retried tests. Check the rerun report for details.

[Wong4j]

/bot run

[Wong4j]

/bot run

[tensorrt-cicd]

PR_Github #25950 [ run ] triggered by Bot. Commit: 5a60909

[tensorrt-cicd]

PR_Github #25950 [ run ] completed with state SUCCESS. Commit: 5a60909 /LLM/main/L0_MergeRequest_PR pipeline #19680 completed with status: 'FAILURE'