[None][perf] DSv4 prep: attention fusion custom ops

[lfr-0531]

Summary by CodeRabbit

• New Features

  • Added FP8 quantization with UE8M0 scale packing for block-scale GEMM operations
  • Added DeepSeek V4 Q normalization operation
  • Added MLA RoPE inplace operation
  • Added fused inverse RoPE + FP8 quantization for DeepSeek-V4 models

• Tests

  • Added validation tests for DeepSeek V4 Q normalization
  • Added tests for fused inverse RoPE + FP8 quantization operations

Description

This PR is PR-5 from the DeepSeek-V4 split of umbrella PR #14751.

It lands the standalone attention fusion/custom-op optimization primitives without bringing in the full DSv4 sparse MLA backend, cache manager, MoE routing, or DSv4 model/tokenizer files.

Changes included:

• Add DSv4 q-norm CUDA/thop op.
• Add MLA RoPE in-place thop op used by the focused fused quant reference path.
• Add fused inverse-RoPE + FP8 1x128 quantization torch custom op.
• Add packed FP8 1x128 quantization path for SM100.
• Add focused numerical tests for the new custom ops.

Reviewer note: this is intended as a standalone attention optimization PR. Later DSv4 sparse MLA work can consume these ops, but this PR does not functionally depend on the sparse MLA backend PR.

Rebase/conflict note: before push, this branch was rebased from github/main e171875352df96eeb543557fe1c12fe320d37c7e to f49d09f1e0c0df1fa61f0fb89ed416b8e1056326. No conflicts occurred.

Test Coverage

• Built and installed from source after rebasing to latest main:
  python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache --cuda_architectures "90-real;100-real" --configure_cmake
• Installed generated wheel:
  ./.venv-3.12/bin/pip install --force-reinstall --no-deps build/tensorrt_llm-1.3.0rc19-cp312-cp312-linux_x86_64.whl
• Smoke checked that the current worktree source is imported and deepseek_v4_q_norm, fp8_quantize_1x128_packed_ue8m0, and mla_rope_inplace are registered.
• Checked GPU state with nvidia-smi before tests; used idle CUDA_VISIBLE_DEVICES=0 after rebase.
• CUDA_VISIBLE_DEVICES=0 timeout 1800 ./.venv-3.12/bin/python -m pytest -q --tb=short -ra tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py
  • 13 passed
• CUDA_VISIBLE_DEVICES=0 timeout 1800 ./.venv-3.12/bin/python -m pytest -q --tb=short -ra tests/unittest/_torch/custom_ops/test_fused_inv_rope_fp8_quant.py
  • 4 passed
• ./.venv-3.12/bin/pre-commit run --files <PR5 files>
  • Passed
• PR-5 scope-leakage check: no sparse MLA dispatch/backend, cache manager, MoE, compressor/mHC, IndexerTopK/TopK, or DSv4 model/tokenizer files.
• git diff --check HEAD~1..HEAD
  • Passed

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)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds four new custom ops for DeepSeek V4 inference: a CUDA fused FP8 E4M3 quantization + UE8M0 packed scale kernel (1×128 block-scale GEMM), a DeepSeek V4 Q-Norm CUDA kernel, an MLA RoPE inplace Torch op, and a Triton fused inverse-RoPE + FP8 quantization kernel. All ops are wired into PyTorch dispatch with fake implementations for tracing, build system entries, and unit tests.

Changes

DeepSeek V4 custom ops: FP8 packed quantization, QNorm, MLA RoPE inplace, and fused Triton inv-RoPE+FP8

Layer / File(s)	Summary
FP8+UE8M0 packed quantization CUDA kernel and header cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h, ...fp8_blockscale_quant_packed.cu	Introduces the CUDA kernel performing per-warp BF16 loading, warp-level amax reduction, UE8M0 scale derivation, FP8 E4M3 quantization with K-bound masking, and int32 packed-scale output writing; includes the host launcher and declaring header.
FP8 packed quantization PyTorch op wiring and dispatch cpp/tensorrt_llm/thop/fp8Quantize.cpp, tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py, tensorrt_llm/_torch/custom_ops/torch_custom_ops.py	Adds SM100-only fp8_quantize_1x128_packed_ue8m0 Torch op that allocates padded buffers, launches the fused kernel, and returns an FP8 tensor + strided packed-scale tensor; registers Torch schema and CUDA impl; adds fake impl for tracing; gates fused dispatch behind TRTLLM_FUSED_FP8_QUANT_PACK env var.
DeepSeek V4 QNorm CUDA kernel and header cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h, ...deepseekV4QNormKernel.cu	Implements per-row RMS normalization for headDim==512 using Vec2Traits-typed warp reductions over half2/bfloat162 pairs, an rsqrtf scale with epsilon, and the public invokeDeepseekV4QNorm dispatcher.
DeepSeek V4 QNorm PyTorch op, fake impl, build, and tests cpp/tensorrt_llm/thop/deepseekV4QNormOp.cpp, cpp/tensorrt_llm/thop/CMakeLists.txt, tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py, tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py	Wires deepseek_v4_q_norm into the Torch CUDA dispatch with dtype/shape validation, adds fake impl for dynamo, registers source in CMakeLists, and tests correctness + torch.compile fullgraph behavior.
MLA RoPE inplace Torch op binding and build wiring cpp/tensorrt_llm/thop/mlaRopeInplaceOp.cpp, cpp/tensorrt_llm/thop/CMakeLists.txt	Adds trtllm::mla_rope_inplace Torch operator with dtype/shape validation dispatching to invokeMLARoPEInplace for BF16/FP16; registers schema, CUDA impl, and source in CMakeLists.
Fused inverse-RoPE + FP8 Triton kernel, op registration, and tests tensorrt_llm/_torch/custom_ops/triton_fused_inv_rope_fp8_quant.py, tensorrt_llm/_torch/custom_ops/__init__.py, tests/unittest/_torch/custom_ops/test_fused_inv_rope_fp8_quant.py	Introduces a Triton kernel applying inverse RoPE (NEOX/GPT-J) then block-scaled FP8 quantization, a Python CUDA wrapper allocating output buffers and driving the kernel, a fake impl for graph capture, eager op registration at import time, and SM100f-gated numerical equivalence tests against the mla_rope_inplace + fp8_batched_quantize_1x128_permute102 reference pipeline.

Sequence Diagrams

sequenceDiagram
  participant Caller as PyTorch Caller
  participant trtllm_Op as trtllm::fused_inv_rope_fp8_quant_vllm_port
  participant CudaWrapper as fused_inv_rope_fp8_quant_cuda
  participant TritonKernel as Triton Kernel
  participant Outputs as fp8_buf / scale_buf

  rect rgba(70, 130, 180, 0.5)
    note over Caller,trtllm_Op: Op dispatch
    Caller->>trtllm_Op: (o, positions, cos_sin_cache, n_groups, heads_per_group, nope_dim, rope_dim, quant_group_size, is_neox)
  end

  rect rgba(120, 80, 160, 0.5)
    note over CudaWrapper: Validation and allocation
    trtllm_Op->>CudaWrapper: validate shapes, flatten cos_sin_cache, coerce positions dtype
    CudaWrapper->>CudaWrapper: allocate fp8_buf [n_groups, num_tokens_aligned, d]
    CudaWrapper->>CudaWrapper: allocate scale_buf [n_groups, num_scale_blocks, tma_aligned_T]
  end

  rect rgba(60, 160, 80, 0.5)
    note over TritonKernel: Fused compute
    CudaWrapper->>TritonKernel: launch with strides/constexprs (IS_NEOX)
    TritonKernel->>TritonKernel: apply inverse RoPE in-register
    TritonKernel->>TritonKernel: per-block FP32 absmax → scale → clamp → FP8 E4M3
    TritonKernel->>Outputs: write FP8 values and FP32 scales
  end

  Caller-->>Caller: receive (fp8_buf, scale_buf)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Suggested reviewers

• chang-l
• hyukn
• byshiue
• xxi-nv
• schetlur-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 19.15% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title '[None][perf] DSv4 prep: attention fusion custom ops' accurately describes the main contribution: adding attention fusion and custom operation optimization primitives for DeepSeek-V4.
Description check	✅ Passed	The PR description is comprehensive and well-structured, covering objectives, implementation details, test coverage with specific test results, and a completed PR checklist.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 6

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

cpp/tensorrt_llm/thop/fp8Quantize.cpp (1)

2-2: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Update the copyright year range in this modified source file.

Line 2 still says 2020-2023, but this file is materially modified in this PR.

Suggested fix

- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.  All rights reserved.

As per coding guidelines, “NVIDIA copyright header on ALL new files (update year on modified files).”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/tensorrt_llm/thop/fp8Quantize.cpp` at line 2, Update the copyright year
range in the header comment of fp8Quantize.cpp from the current "2020-2023" to
include the current year since this file has been materially modified in this
PR. Change line 2 to extend the year range to reflect that the file was modified
after 2023, following NVIDIA's coding guidelines for updating years on modified
files.

Source: Coding guidelines

🧹 Nitpick comments (2)

tests/unittest/_torch/custom_ops/test_fused_inv_rope_fp8_quant.py (1)

58-91: ⚡ Quick win

Coverage is currently insufficient: add one is_neox=False parity case.

This test module validates only the NEOX path, but the new kernel has a second interleaved (GPT-J) branch. Please parametrize is_neox (or add one dedicated case) so both paths are numerically checked in tests/unittest/_torch/custom_ops/test_fused_inv_rope_fp8_quant.py.

As per coding guidelines for tests/**, coverage feedback should be actionable and explicitly call out sufficiency; this branch coverage needs follow-up in this file.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/custom_ops/test_fused_inv_rope_fp8_quant.py` around
lines 58 - 91, Add test coverage for the interleaved (GPT-J) path in the kernel.
The current test_fused_inv_rope_fp8_quant_neox function only validates the NEOX
path with is_neox=True, but the kernel also has an is_neox=False branch. Either
parametrize the test to run with both is_neox=True and is_neox=False values, or
add a new dedicated test function (such as
test_fused_inv_rope_fp8_quant_interleaved) that calls _ref_path and _fused_path
with is_neox=False to ensure numerical correctness of the interleaved path.

Source: Coding guidelines

tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py (1)

20-63: ⚡ Quick win

Validation-path coverage is currently insufficient for this custom op contract.

Coverage status:

• tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py: sufficient for numerical parity + torch.compile happy path.
• cpp/tensorrt_llm/thop/deepseekV4QNormOp.cpp guards: insufficiently tested.

Please add negative tests in tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py for:

• non-CUDA input
• non-contiguous input
• unsupported dtype
• non-2D input
• q.shape[1] != num_heads * head_dim
• head_dim != 512

As per coding guidelines, tests under tests/** should include actionable coverage assessment (sufficient/insufficient/follow-up), and this file currently needs guardrail-coverage follow-up.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py` around lines 20
- 63, The test file for deepseek_v4_q_norm currently lacks negative test
coverage for validation guards present in the C++ implementation. Add the
following negative test functions to
tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py: test cases that
verify the custom op raises appropriate errors for non-CUDA input,
non-contiguous input, unsupported dtype, non-2D input, shape mismatch where
q.shape[1] does not equal num_heads * head_dim, and invalid head_dim values
where head_dim is not 512. Each test should call
torch.ops.trtllm.deepseek_v4_q_norm with invalid parameters and assert that the
expected exception is raised, ensuring all validation guards are covered.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In
`@cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu`:
- Around line 197-205: Add a guard check before the launchWithPdlWhenEnabled
call in the fp8_quantize_1x128_packed_kernel_impl block to validate that k and m
are both non-zero. When either k or m is 0, the computed grid dimensions
(num_packed_sf_k and m_blocks) can become 0, resulting in invalid kernel launch
parameters. Insert an early return or skip the kernel launch entirely when m ==
0 or k == 0 to prevent launching with zero grid dimensions.

In `@cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu`:
- Around line 127-133: The current code uses assert to check that headDim equals
512, but asserts are disabled in release builds, allowing invalid headDim values
to bypass the check and potentially cause out-of-bounds access or incorrect
indexing in the kernel launch. Replace the assert statement with a proper
runtime check (such as an early return, error log, or exception) that verifies
headDim equals 512 before launching the deepseekV4QNormKernel kernel. This
enforcement must apply wherever the public kernel entrypoint calls
deepseekV4QNormKernel with a hardcoded template parameter of 512 to ensure
consistency across all call sites.

In `@cpp/tensorrt_llm/thop/mlaRopeInplaceOp.cpp`:
- Around line 40-50: The code retrieves a CUDA stream at line 40 before
validating that inputs are CUDA tensors, and the auxiliary tensor validations at
lines 48-50 only check cos_sin_cache scalar type without validating its shape.
Add device validation checks (using TORCH_CHECK to confirm data and position_ids
are CUDA tensors) before the at::cuda::getCurrentCUDAStream call to prevent
accessing the CUDA stream with non-CUDA inputs. Additionally, add shape
validation for cos_sin_cache to ensure its dimensions match expected
requirements (verify it is 2D or 3D as appropriate, and check that its
dimensions align with num_tokens, rope_dim, and num_heads as needed). Apply the
same validations at lines 54-62 where similar operations occur with the same
tensors.

In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`:
- Around line 1710-1712: The gate condition on line 1710 uses get_sm_version()
>= 100, but the underlying C++ op
torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0 is restricted to SM100-family
only. Update the condition to match the C++ implementation constraint (likely
checking for SM100 specifically rather than >= 100) so that unsupported
architectures fall through to the fallback quantization path instead of
encountering a runtime error.

In `@tensorrt_llm/_torch/custom_ops/triton_fused_inv_rope_fp8_quant.py`:
- Around line 282-283: The variable `num_heads` is being unpacked from the
`o.shape` tuple assignment but is never used in the subsequent code, which
triggers a RUF059 linting violation. Rename `num_heads` to `_num_heads` in the
tuple unpacking statement to follow Python convention for intentionally unused
variables and satisfy the linter.
- Around line 176-267: Add hard runtime validation checks before the kernel
launch in the _fused_inv_rope_fp8_quant_per_head call to ensure critical
preconditions are met. Specifically, validate that positions.numel() equals
num_tokens to prevent illegal memory access, and verify that all input tensors
(o, positions, cos_sin_cache) are on the same CUDA device. Replace or supplement
the existing assert statements with explicit runtime checks using if conditions
that raise clear exceptions, since Python asserts can be skipped in optimized
runs. These checks should execute unconditionally before the grid-based kernel
launch to catch configuration errors at the Python level rather than as cryptic
kernel failures.

---

Outside diff comments:
In `@cpp/tensorrt_llm/thop/fp8Quantize.cpp`:
- Line 2: Update the copyright year range in the header comment of
fp8Quantize.cpp from the current "2020-2023" to include the current year since
this file has been materially modified in this PR. Change line 2 to extend the
year range to reflect that the file was modified after 2023, following NVIDIA's
coding guidelines for updating years on modified files.

---

Nitpick comments:
In `@tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py`:
- Around line 20-63: The test file for deepseek_v4_q_norm currently lacks
negative test coverage for validation guards present in the C++ implementation.
Add the following negative test functions to
tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py: test cases that
verify the custom op raises appropriate errors for non-CUDA input,
non-contiguous input, unsupported dtype, non-2D input, shape mismatch where
q.shape[1] does not equal num_heads * head_dim, and invalid head_dim values
where head_dim is not 512. Each test should call
torch.ops.trtllm.deepseek_v4_q_norm with invalid parameters and assert that the
expected exception is raised, ensuring all validation guards are covered.

In `@tests/unittest/_torch/custom_ops/test_fused_inv_rope_fp8_quant.py`:
- Around line 58-91: Add test coverage for the interleaved (GPT-J) path in the
kernel. The current test_fused_inv_rope_fp8_quant_neox function only validates
the NEOX path with is_neox=True, but the kernel also has an is_neox=False
branch. Either parametrize the test to run with both is_neox=True and
is_neox=False values, or add a new dedicated test function (such as
test_fused_inv_rope_fp8_quant_interleaved) that calls _ref_path and _fused_path
with is_neox=False to ensure numerical correctness of the interleaved path.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 50cf71ce-f14d-4185-b603-26a836222b7b

📥 Commits

Reviewing files that changed from the base of the PR and between f49d09f and 38cf87f.

📒 Files selected for processing (14)

• cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu
• cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h
• cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu
• cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h
• cpp/tensorrt_llm/thop/CMakeLists.txt
• cpp/tensorrt_llm/thop/deepseekV4QNormOp.cpp
• cpp/tensorrt_llm/thop/fp8Quantize.cpp
• cpp/tensorrt_llm/thop/mlaRopeInplaceOp.cpp
• tensorrt_llm/_torch/custom_ops/__init__.py
• tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
• tensorrt_llm/_torch/custom_ops/triton_fused_inv_rope_fp8_quant.py
• tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py
• tests/unittest/_torch/custom_ops/test_fused_inv_rope_fp8_quant.py

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54448 [ run ] triggered by Bot. Commit: 38cf87f Link to invocation

[tensorrt-cicd]

PR_Github #54448 [ run ] completed with state FAILURE. Commit: 38cf87f
/LLM/main/L0_MergeRequest_PR pipeline #43511 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54505 [ run ] triggered by Bot. Commit: 6d55c86 Link to invocation

[tensorrt-cicd]

PR_Github #54505 [ run ] completed with state SUCCESS. Commit: 6d55c86
/LLM/main/L0_MergeRequest_PR pipeline #43565 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54695 [ run ] triggered by Bot. Commit: 6d55c86 Link to invocation

[mingyangHao]

The redundant triton kernel can be removed later.

[tensorrt-cicd]

PR_Github #54695 [ run ] completed with state FAILURE. Commit: 6d55c86
/LLM/main/L0_MergeRequest_PR pipeline #43725 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54761 [ run ] triggered by Bot. Commit: 6d55c86 Link to invocation

[tensorrt-cicd]

PR_Github #54761 [ run ] completed with state SUCCESS. Commit: 6d55c86
/LLM/main/L0_MergeRequest_PR pipeline #43781 completed with status: 'SUCCESS'

CI Report

Link to invocation