Skip to content

[TRTLLM-14138][perf] Add fused kernels for Gemma4 serving#16074

Merged
kaiyux merged 4 commits into
NVIDIA:mainfrom
kaiyux:perf/gemma4-fused-kernels-20260707
Jul 11, 2026
Merged

[TRTLLM-14138][perf] Add fused kernels for Gemma4 serving#16074
kaiyux merged 4 commits into
NVIDIA:mainfrom
kaiyux:perf/gemma4-fused-kernels-20260707

Conversation

@kaiyux

@kaiyux kaiyux commented Jul 7, 2026

Copy link
Copy Markdown
Member

Description

Adds four fused kernels for the Gemma4 NVFP4 serving path under tensorrt_llm/_torch/modules/gemma4/, replacing chains of small elementwise/norm/quant ops in the decoder — three bespoke Triton kernels plus one fusion backed by flashinfer:

  • Fused QKV prep (fused_qkv.py, Triton): QK RMSNorm + RoPE + FP8 quantization in one kernel over the fused QKV projection output.
  • Fused GeLU + quant (fused_gelu_quant.py, Triton): GeLU-tanh + mul + NVFP4 quantization, avoiding the intermediate HBM round-trip of flashinfer_gelu_tanh_and_mul followed by a separate quant op.
  • Fused norm + quant (fused_norm_quant.py, flashinfer): pre-feedforward RMSNorm + NVFP4 quantization via flashinfer's CuTe-DSL rmsnorm_fp4quant (SM100+). It emits the same 128x4-swizzled E4M3 block scales with the same e4m3(gs * blockAmax / 6) scale recipe as trtllm::fp4_quantize, so the consumer Linear's static input_scale passes through as the global scale and the outputs are consumed as Fp4QuantizedTensor. Note on numerics: the kernel quantizes the fp32 norm result directly (no intermediate bf16 round), so its outputs are near- rather than byte-identical to the unfused chain, with unchanged quantization error (see Test Coverage and Accuracy below).
  • Fused tail (fused_tail.py, Triton): fused residual-add + RMSNorm tail (with optional scale), optionally also emitting the next layer's input norm.

Each fusion decides lazily from the module/quantization configuration whether it applies (e.g., NVFP4 static-scale gate_up/down projections, plain bf16 norms, no LoRA / MoE block / torch.compile), and the unfused paths remain as the reference fallbacks. The fused QKV prep additionally gates itself on the profiled serving configuration (trtllm-gen backend + FP8 KV cache + standard neox rope layout); the fused norm + quant additionally requires flashinfer's CuTe-DSL kernels to be importable (nvidia-cutlass-dsl) and falls back to the unfused chain otherwise.

Because the fused QKV path feeds attention with Q already quantized to FP8, FlashInferAttention now allocates a BF16 output buffer when Q arrives as FP8 — the trtllm-gen QkvE4m3OBfloat16 cubins emit BF16, so the output must not follow Q's dtype.

Split out of #16033 (3 of 3, together with #16072 and #16073); this PR is independent of the other two.

Test Coverage

  • New unit tests for every fused module (78 cases, all passing on B200): tests/unittest/_torch/modules/test_gemma4_fused_qkv_prep.py, test_gemma4_fused_gelu_quant.py, test_gemma4_fused_norm_quant.py, test_gemma4_fused_tail.py. The Triton kernels are byte-parity-checked against the unfused reference paths, including strided inputs, extreme scales, and zero-row cases. The flashinfer-backed norm + quant suite is quality-based instead (dequantized error against the fp32 norm must match the unfused chain's; byte comparison restricted to blocks with nonzero scales) and includes a CUDA-graph capture/replay test.
  • Perf: measured standalone on top of main (045705139d) on 1x B200 (Gemma-4-31B-IT-NVFP4, ISL 1024 / OSL 128, 1024 prompts, concurrency 1024, trtllm-serve + benchmark_serving):
Metric main (045705139d) this PR delta
Output throughput 1331.8 tok/s 1607.9 tok/s +20.7%
Median TTFT 49.45 s 40.98 s -17.1%
Median TPOT 141.0 ms 117.1 ms -17.0%
P99 ITL 409.5 ms 348.9 ms -14.8%
Median E2EL 67.41 s 55.77 s -17.3%

Together with #16072 and #16073, the three PRs compose to the +30.9% measured in #16033. The flashinfer-backed norm + quant implementation was additionally A/B-verified against a same-node, same-tree Triton implementation of the fusion: end-to-end serving throughput is unchanged (1599 vs 1605 tok/s, within noise) with the fused-kernel slot at ~1.2% of GPU time (in-context prefill calls 6-8% faster, decode calls within ~0.6us).

  • Accuracy (GSM8K 5-shot via trtllm-eval, Gemma-4-31B-IT-NVFP4 on 1x B200, same commit, fused norm + quant enabled vs disabled):
Config flexible-extract strict-match average
fused (this PR) 65.96 ± 1.31 65.13 ± 1.31 65.54
unfused norm + quant (control) 65.35 ± 1.31 64.37 ± 1.32 64.86

The difference is within the sampling stderr — no accuracy regression from the fused path.

PR Checklist

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

🤖 Generated with Claude Code

@kaiyux
kaiyux requested review from a team as code owners July 7, 2026 16:28
@coderabbitai

coderabbitai Bot commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

Adds four new Triton-based fused kernel modules for Gemma4 (GELU+mul→NVFP4 quant, QKV RMSNorm+RoPE+quant, RMSNorm→NVFP4 quant, decoder-layer tail residual/scale/next-norm), integrates them into the Gemma4 PyTorch model's attention and decoder layer with enablement gating and cross-layer state propagation, adds corresponding parity tests, and applies an unrelated FP8-output allocation fix in the FlashInfer attention backend.

Changes

FlashInfer FP8 output allocation fix

Layer / File(s) Summary
BF16 output allocation for FP8 Q input
tensorrt_llm/_torch/attention_backend/flashinfer.py
Adds an elif branch allocating a bfloat16 output tensor with q's shape when q is already torch.float8_e4m3fn, instead of inheriting FP8 dtype via empty_like.

Estimated code review effort: 4 (Complex) | ~45 minutes

Gemma4 fused kernel feature

Layer / File(s) Summary
Fused GELU-tanh+mul → NVFP4 quant kernel
tensorrt_llm/_torch/modules/gemma4/fused_gelu_quant.py, tests/unittest/_torch/modules/test_gemma4_fused_gelu_quant.py
New Triton kernel and gemma4_fused_gelu_mul_fp4 wrapper fuse gelu_tanh+mul with NVFP4 E4M3/E2M1 quantization and swizzled scale packing; parity tests validate against unfused reference across sizes, extreme scales, and strided inputs.
Fused QKV RMSNorm + RoPE + quant kernel
tensorrt_llm/_torch/modules/gemma4/fused_qkv.py, tests/unittest/_torch/modules/test_gemma4_fused_qkv_prep.py
New Triton kernel and gemma4_fused_qkv_norm_rope_quant wrapper perform per-head RMSNorm, Neox RoPE on Q/K, and optional FP8 output for packed QKV tensors; parity tests cover fp8 and bf16/strided modes.
Fused RMSNorm → NVFP4 quant kernel
tensorrt_llm/_torch/modules/gemma4/fused_norm_quant.py, tests/unittest/_torch/modules/test_gemma4_fused_norm_quant.py
New Triton kernel and gemma4_fused_norm_fp4 wrapper compute RMSNorm then NVFP4-quantize with swizzled scales; parity tests cover extreme scales, strided rows, and zero rows.
Fused decoder-layer tail (norm + residual + scale + next-norm)
tensorrt_llm/_torch/modules/gemma4/fused_tail.py, tests/unittest/_torch/modules/test_gemma4_fused_tail.py
New Triton kernel and gemma4_fused_norm_add_scale/gemma4_fused_norm_add wrappers fuse post-FFN and post-attention RMSNorm/residual/scale, with optional dual output for the next layer's input norm; parity tests cover scalars, non-power-of-two/strided shapes, and dual-norm output.
Gemma4 model integration of fused kernels
tensorrt_llm/_torch/models/modeling_gemma4.py, tensorrt_llm/_torch/modules/gemma4/__init__.py
Wires the four fused kernels into _Gemma4GeluQuantMLP, Gemma4Attention (fused QKV prep/RoPE), and Gemma4DecoderLayer/Gemma4TextModel (fused norm/add, norm/quant, tail with pre_normed propagation across layers) behind lazy enablement checks; adds SPDX header to the package init.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Gemma4TextModel
  participant Gemma4DecoderLayer
  participant fused_norm_quant
  participant fused_gelu_quant
  participant fused_tail

  Gemma4TextModel->>Gemma4DecoderLayer: forward(hidden_states, pre_normed)
  Gemma4DecoderLayer->>fused_norm_quant: gemma4_fused_norm_fp4(pre-FFN norm)
  fused_norm_quant-->>Gemma4DecoderLayer: Fp4QuantizedTensor
  Gemma4DecoderLayer->>fused_gelu_quant: gemma4_fused_gelu_mul_fp4(MLP down-proj input)
  fused_gelu_quant-->>Gemma4DecoderLayer: Fp4QuantizedTensor
  Gemma4DecoderLayer->>fused_tail: gemma4_fused_norm_add_scale(hidden_states, residual)
  fused_tail-->>Gemma4DecoderLayer: (hidden_states, pre_normed)
  Gemma4DecoderLayer-->>Gemma4TextModel: (hidden_states, pre_normed)
Loading

Suggested reviewers: moraxu, yizhang-nv, yechank-nvidia, yuxianq, brb-nv, hyukn, StanleySun639

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Title check ✅ Passed The title is concise, specific, and accurately summarizes the main change: adding fused Gemma4 serving kernels.
Description check ✅ Passed The description includes the required Description, Test Coverage, and PR Checklist sections and explains the change and validation clearly.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (7)
tensorrt_llm/_torch/modules/gemma4/fused_norm_quant.py (1)

45-45: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Prefer built-in tuple over typing.Tuple.

♻️ Suggested change
-from typing import Tuple
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:

As per coding guidelines: "Prefer built-in types list, dict, tuple over typing.List, typing.Dict, typing.Tuple."

🤖 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 `@tensorrt_llm/_torch/modules/gemma4/fused_norm_quant.py` at line 45, The
import in the fused norm quant module uses typing.Tuple instead of the preferred
built-in tuple type. Update the type annotations in fused_norm_quant.py to use
tuple directly and remove the Tuple import if it is no longer needed, keeping
the change localized to the symbols used in this module.

Source: Coding guidelines

tensorrt_llm/_torch/modules/gemma4/fused_gelu_quant.py (1)

42-42: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Prefer built-in tuple generics over typing.Tuple.

The codebase targets Python 3.10+, so typing.Tuple can be dropped in favor of the built-in.

♻️ Suggested change
-from typing import Tuple
-def gemma4_fused_gelu_mul_fp4(
-    x: torch.Tensor, global_scale: torch.Tensor
-) -> Tuple[torch.Tensor, torch.Tensor]:
+def gemma4_fused_gelu_mul_fp4(
+    x: torch.Tensor, global_scale: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:

As per coding guidelines: "Prefer built-in types list, dict, tuple over typing.List, typing.Dict, typing.Tuple."

🤖 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 `@tensorrt_llm/_torch/modules/gemma4/fused_gelu_quant.py` at line 42, Replace
the `typing.Tuple` import in `fused_gelu_quant.py` with the built-in `tuple`
generic since the codebase targets Python 3.10+ and follows the guideline to
prefer built-in collection types. Check any type annotations in the same module,
especially around `fused_gelu_quant` and related symbols, and update them to use
`tuple[...]` so the import can be removed cleanly.

Source: Coding guidelines

tensorrt_llm/_torch/modules/gemma4/fused_qkv.py (2)

48-48: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Prefer built-in tuple over typing.Tuple.

♻️ Suggested change
-from typing import Tuple
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

As per coding guidelines: "Prefer built-in types list, dict, tuple over typing.List, typing.Dict, typing.Tuple."

🤖 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 `@tensorrt_llm/_torch/modules/gemma4/fused_qkv.py` at line 48, The import in
fused_qkv.py still uses typing.Tuple instead of the built-in tuple type, which
conflicts with the coding guideline. Update the type hints in the affected
module to use tuple directly and remove the Tuple import from typing if it is no
longer needed; check the symbols in fused_qkv.py that reference Tuple so they
all use the built-in type consistently.

Source: Coding guidelines


183-183: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value

position_ids.view(-1) can raise on a non-contiguous input.

The docstring promises inputs "flattenable to [N]", but .view(-1) fails if position_ids is non-contiguous (e.g. a strided slice). .reshape(-1) is the safe flatten. Low risk given the current call-site passes contiguous ids, but cheap to harden.

🛡️ Suggested change
-        position_ids.view(-1),
+        position_ids.reshape(-1),
🤖 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 `@tensorrt_llm/_torch/modules/gemma4/fused_qkv.py` at line 183, The flattening
of position_ids in the fused QKV path can fail for non-contiguous tensors
because the current use of position_ids.view(-1) assumes contiguous storage.
Update the logic in the fused_qkv handling to use a safe flattening approach
such as position_ids.reshape(-1), keeping the same behavior for contiguous
inputs while making the code robust for strided or sliced tensors.
tests/unittest/_torch/modules/test_gemma4_fused_gelu_quant.py (1)

77-89: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Coverage nit: add an all-zero row case.

The kernel's oscale = where(vmax != 0, ...) branch (zero block scales) is exercised by test_gemma4_fused_norm_quant.py::test_fused_norm_quant_zero_row but not here. A zeroed input row (gate=up=0 ⇒ v=0 ⇒ vmax=0) would close the parity gap for this kernel too. Sufficient otherwise; treat as a 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/modules/test_gemma4_fused_gelu_quant.py` around lines
77 - 89, Add a dedicated zero-row coverage case to
test_fused_gelu_quant_strided_rows so the fused GELU quant path exercises the
oscale = where(vmax != 0, ...) branch; create an input where one row is all
zeros (gate/up = 0 so v and vmax are 0), then run gemma4_fused_gelu_mul_fp4 and
compare against torch.ops.trtllm.fp4_quantize using the same validation pattern
as the existing test. Keep the current strided-row parity assertions, and extend
the test to explicitly verify the zero block scale behavior for this kernel.

Source: Path instructions

tensorrt_llm/_torch/modules/gemma4/fused_tail.py (1)

57-57: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Modernize typing: tuple, |, and X | None.

Per the 3.10+ target, Tuple/Union/Optional can be replaced with built-in generics and |.

♻️ Suggested change (signatures)
-from typing import Optional, Tuple, Union
-    scale: Optional[torch.Tensor],
+    scale: torch.Tensor | None,
     eps: float,
-    next_norm_weight: Optional[torch.Tensor],
+    next_norm_weight: torch.Tensor | None,
     next_norm_eps: float,
-) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

Apply the same to gemma4_fused_norm_add_scale's next_norm_weight param and return type.

As per coding guidelines: "Prefer built-in types ... over typing.List, typing.Dict, typing.Tuple; use | syntax instead of typing.Union."

🤖 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 `@tensorrt_llm/_torch/modules/gemma4/fused_tail.py` at line 57, Update the
typing imports and annotations in fused_tail.py to use Python 3.10+ style
built-ins instead of typing.Tuple/Union/Optional. In particular, adjust the
`gemma4_fused_norm_add_scale` signature and any related type hints so
`next_norm_weight` uses `X | None` style and the return type uses `tuple[...]`
and `|` rather than `Tuple`, `Union`, or `Optional`, while keeping the function
behavior unchanged.

Source: Coding guidelines

tensorrt_llm/_torch/models/modeling_gemma4.py (1)

109-159: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Missing return-type/parameter annotations on new methods.

__init__ and _apply_activation are newly added but lack annotations. As per coding guidelines, "Always annotate functions with return types (use None if no return)".

♻️ Proposed annotations
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
-    def _apply_activation(self, x, *, has_lora: bool = False):
+    def _apply_activation(
+        self, x: torch.Tensor | Fp4QuantizedTensor, *, has_lora: bool = False
+    ) -> torch.Tensor | Fp4QuantizedTensor:
🤖 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 `@tensorrt_llm/_torch/models/modeling_gemma4.py` around lines 109 - 159, In
_Gemma4GeluQuantMLP, add the missing type annotations for the new methods so
they follow the project’s typing rules. Update __init__ to explicitly return
None, and annotate _apply_activation’s parameters and return type consistently
with the surrounding class methods. Use the existing symbols
_Gemma4GeluQuantMLP, __init__, and _apply_activation to locate the changes.

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 `@tests/unittest/_torch/modules/test_gemma4_fused_gelu_quant.py`:
- Around line 54-65: The FP4 parity test in test_fused_gelu_quant_parity should
be guarded so it only runs on Blackwell/SM100+, since the NVFP4 path used by
torch.ops.trtllm.flashinfer_gelu_tanh_and_mul depends on that architecture. Add
a pytest.mark.skipif(get_sm_version() < 100, ...) to the test (or its
parametrized wrapper) and apply the same architecture check in
test_gemma4_fused_norm_quant so both FP4 parity suites are consistently skipped
on older GPUs.

---

Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_gemma4.py`:
- Around line 109-159: In _Gemma4GeluQuantMLP, add the missing type annotations
for the new methods so they follow the project’s typing rules. Update __init__
to explicitly return None, and annotate _apply_activation’s parameters and
return type consistently with the surrounding class methods. Use the existing
symbols _Gemma4GeluQuantMLP, __init__, and _apply_activation to locate the
changes.

In `@tensorrt_llm/_torch/modules/gemma4/fused_gelu_quant.py`:
- Line 42: Replace the `typing.Tuple` import in `fused_gelu_quant.py` with the
built-in `tuple` generic since the codebase targets Python 3.10+ and follows the
guideline to prefer built-in collection types. Check any type annotations in the
same module, especially around `fused_gelu_quant` and related symbols, and
update them to use `tuple[...]` so the import can be removed cleanly.

In `@tensorrt_llm/_torch/modules/gemma4/fused_norm_quant.py`:
- Line 45: The import in the fused norm quant module uses typing.Tuple instead
of the preferred built-in tuple type. Update the type annotations in
fused_norm_quant.py to use tuple directly and remove the Tuple import if it is
no longer needed, keeping the change localized to the symbols used in this
module.

In `@tensorrt_llm/_torch/modules/gemma4/fused_qkv.py`:
- Line 48: The import in fused_qkv.py still uses typing.Tuple instead of the
built-in tuple type, which conflicts with the coding guideline. Update the type
hints in the affected module to use tuple directly and remove the Tuple import
from typing if it is no longer needed; check the symbols in fused_qkv.py that
reference Tuple so they all use the built-in type consistently.
- Line 183: The flattening of position_ids in the fused QKV path can fail for
non-contiguous tensors because the current use of position_ids.view(-1) assumes
contiguous storage. Update the logic in the fused_qkv handling to use a safe
flattening approach such as position_ids.reshape(-1), keeping the same behavior
for contiguous inputs while making the code robust for strided or sliced
tensors.

In `@tensorrt_llm/_torch/modules/gemma4/fused_tail.py`:
- Line 57: Update the typing imports and annotations in fused_tail.py to use
Python 3.10+ style built-ins instead of typing.Tuple/Union/Optional. In
particular, adjust the `gemma4_fused_norm_add_scale` signature and any related
type hints so `next_norm_weight` uses `X | None` style and the return type uses
`tuple[...]` and `|` rather than `Tuple`, `Union`, or `Optional`, while keeping
the function behavior unchanged.

In `@tests/unittest/_torch/modules/test_gemma4_fused_gelu_quant.py`:
- Around line 77-89: Add a dedicated zero-row coverage case to
test_fused_gelu_quant_strided_rows so the fused GELU quant path exercises the
oscale = where(vmax != 0, ...) branch; create an input where one row is all
zeros (gate/up = 0 so v and vmax are 0), then run gemma4_fused_gelu_mul_fp4 and
compare against torch.ops.trtllm.fp4_quantize using the same validation pattern
as the existing test. Keep the current strided-row parity assertions, and extend
the test to explicitly verify the zero block scale behavior for this kernel.
🪄 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: cd375ccd-fe71-470b-b8f7-e69024efc527

📥 Commits

Reviewing files that changed from the base of the PR and between 409b696 and 6aa7581.

📒 Files selected for processing (11)
  • tensorrt_llm/_torch/attention_backend/flashinfer.py
  • tensorrt_llm/_torch/models/modeling_gemma4.py
  • tensorrt_llm/_torch/modules/gemma4/__init__.py
  • tensorrt_llm/_torch/modules/gemma4/fused_gelu_quant.py
  • tensorrt_llm/_torch/modules/gemma4/fused_norm_quant.py
  • tensorrt_llm/_torch/modules/gemma4/fused_qkv.py
  • tensorrt_llm/_torch/modules/gemma4/fused_tail.py
  • tests/unittest/_torch/modules/test_gemma4_fused_gelu_quant.py
  • tests/unittest/_torch/modules/test_gemma4_fused_norm_quant.py
  • tests/unittest/_torch/modules/test_gemma4_fused_qkv_prep.py
  • tests/unittest/_torch/modules/test_gemma4_fused_tail.py

@kaiyux

kaiyux commented Jul 8, 2026

Copy link
Copy Markdown
Member Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58126 [ run ] triggered by Bot. Commit: 6aa7581 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58126 [ run ] completed with state SUCCESS. Commit: 6aa7581
/LLM/main/L0_MergeRequest_PR pipeline #46787 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

@kaiyux
kaiyux force-pushed the perf/gemma4-fused-kernels-20260707 branch from 6aa7581 to 3d777d1 Compare July 8, 2026 05:45
@kaiyux

kaiyux commented Jul 8, 2026

Copy link
Copy Markdown
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58166 [ run ] triggered by Bot. Commit: 3d777d1 Link to invocation

@kaiyux kaiyux changed the title [TRTLLM-14138][perf] Add fused Triton kernels for Gemma4 serving [TRTLLM-14138][perf] Add fused kernels for Gemma4 serving Jul 8, 2026
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58166 [ run ] completed with state FAILURE. Commit: 3d777d1
/LLM/main/L0_MergeRequest_PR pipeline #46816 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

@kaiyux

kaiyux commented Jul 8, 2026

Copy link
Copy Markdown
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58263 [ run ] triggered by Bot. Commit: a722a6f Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58263 [ run ] completed with state SUCCESS. Commit: a722a6f
/LLM/main/L0_MergeRequest_PR pipeline #46902 completed with status: 'SUCCESS'

CI Report

Link to invocation

Comment thread tensorrt_llm/_torch/modules/gemma4/__init__.py
@kaiyux
kaiyux force-pushed the perf/gemma4-fused-kernels-20260707 branch from a722a6f to e271dca Compare July 9, 2026 05:38

@QiJune QiJune left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@kaiyux

kaiyux commented Jul 9, 2026

Copy link
Copy Markdown
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58394 [ ] completed with state FAILURE. Commit: ``

Link to invocation

@kaiyux

kaiyux commented Jul 9, 2026

Copy link
Copy Markdown
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58402 [ run ] triggered by Bot. Commit: e271dca Link to invocation

@kaiyux

kaiyux commented Jul 9, 2026

Copy link
Copy Markdown
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58408 [ run ] triggered by Bot. Commit: d947b20 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58402 [ run ] completed with state ABORTED. Commit: e271dca

Link to invocation

@kaiyux
kaiyux enabled auto-merge (squash) July 9, 2026 07:22
@kaiyux

kaiyux commented Jul 10, 2026

Copy link
Copy Markdown
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58559 [ run ] triggered by Bot. Commit: d947b20 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58408 [ run ] completed with state ABORTED. Commit: d947b20
/LLM/main/L0_MergeRequest_PR pipeline #47026 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

kaiyux added 4 commits July 10, 2026 11:33
Add four fused Triton kernels for the Gemma4 decoder under
tensorrt_llm/_torch/modules/gemma4/, replacing chains of small
elementwise/norm/quant ops on the NVFP4 serving path:

- fused_qkv: QK RMSNorm + RoPE + FP8 quantization in one kernel over
  the fused QKV projection output.
- fused_gelu_quant: GeLU-tanh + mul + NVFP4 quantization, avoiding the
  intermediate HBM round-trip of flashinfer_gelu_tanh_and_mul followed
  by a separate quant op.
- fused_norm_quant: RMSNorm + NVFP4 quantization.
- fused_tail: fused residual-add + RMSNorm tail (with optional scale).

Each fusion is enabled by default with a dedicated opt-out env var
(TRTLLM_GEMMA4_DISABLE_FUSED_QKV_PREP, ..._FUSED_GELU_QUANT,
..._FUSED_NORM_QUANT, ..._FUSED_TAIL, ..._FUSED_TAIL_NORM2,
..._FUSED_NORM_ADD), and the unfused paths remain as reference
fallbacks. The fused QKV path feeds attention with pre-quantized FP8
Q, so FlashInferAttention now allocates BF16 output when Q arrives as
FP8 (the trtllm-gen QkvE4m3OBfloat16 cubins emit BF16).

Unit tests cover parity against the unfused reference paths, including
strided inputs, extreme scales, and zero-row cases.

Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
…infer rmsnorm_fp4quant

The pre-feedforward RMSNorm + NVFP4 quantize fusion is now backed by
flashinfer's CuTe-DSL rmsnorm_fp4quant (SM100+) instead of a bespoke
Triton kernel: it uses the same 128x4-swizzled scale layout and the same
e4m3(gs * blockAmax / 6) scale recipe, so the consumer Linear's static
input_scale passes through as the global scale and the outputs remain
consumable as Fp4QuantizedTensor. Kernel-level performance is on par
(prefill calls measured 6-8% faster, decode within ~0.6us); end-to-end
serving throughput is unchanged.

The flashinfer kernel quantizes the fp32 norm result directly, without
the intermediate bf16 round the unfused chain performs, so outputs are
near- rather than byte-identical to the unfused pair while quantization
error against the fp32 norm is unchanged. The parity test is rewritten
accordingly (dequantized-quality based, with the byte comparison
restricted to blocks whose e4m3 scale is nonzero) and a CUDA-graph
capture/replay test is added. The fusion gate now also checks kernel
availability (flashinfer with nvidia-cutlass-dsl importable), falling
back to the unfused path otherwise.

Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Relocate the three model-agnostic fused operators (RMSNorm+NVFP4 quant,
gelu_tanh+mul+NVFP4 quant, RMSNorm residual add) from modules/gemma4/ to
modules/fused_ops/, renaming each module and function after the operator
it implements. modeling_gemma4.py keeps ownership of the enablement
checks and unfused fallback paths; only the model-specific fused QKV
kernel stays under modules/gemma4/. Unit tests move to
tests/unittest/_torch/modules/fused_ops/ to mirror the split.

Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Add the four new fused-op test files to the l0_b200.yml pre_merge/pytorch
list; nothing referenced them, so CI never ran them. Also add the missing
skip_pre_blackwell guard to test_gelu_tanh_mul_fp4_quant.py, which needs
SM100+ (fp4_quantize) but had no capability gate.

Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
@kaiyux
kaiyux force-pushed the perf/gemma4-fused-kernels-20260707 branch from d947b20 to 3d92c59 Compare July 10, 2026 03:33
@kaiyux

kaiyux commented Jul 10, 2026

Copy link
Copy Markdown
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58579 [ run ] triggered by Bot. Commit: 3d92c59 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58559 [ run ] completed with state ABORTED. Commit: d947b20

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58579 [ run ] completed with state SUCCESS. Commit: 3d92c59
/LLM/main/L0_MergeRequest_PR pipeline #47172 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

@kaiyux

kaiyux commented Jul 11, 2026

Copy link
Copy Markdown
Member Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58726 [ run ] triggered by Bot. Commit: 3d92c59 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58726 [ run ] completed with state SUCCESS. Commit: 3d92c59
/LLM/main/L0_MergeRequest_PR pipeline #47309 completed with status: 'SUCCESS'

CI Report

Link to invocation

@kaiyux
kaiyux merged commit cf82c5d into NVIDIA:main Jul 11, 2026
8 checks passed
@kaiyux
kaiyux deleted the perf/gemma4-fused-kernels-20260707 branch July 12, 2026 15:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants