diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 7369cbd066d5..26fed0bf12bd 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -1980,6 +1980,11 @@ def forward(self, # MLA generation: output has kv_lora_rank per head output = q.new_empty( [q.shape[0], self.num_heads * self.kv_lora_rank]) + elif q.dtype == torch.float8_e4m3fn: + # Q may arrive pre-quantized to FP8 (fused model-side QKV + # prep); the trtllm-gen FP8 cubins emit BF16 + # (QkvE4m3OBfloat16), so the output must not follow q's dtype. + output = q.new_empty(q.shape, dtype=torch.bfloat16) else: output = torch.empty_like(q) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index cc9ae9e71020..bf356d5e3ba1 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -15,7 +15,7 @@ """TensorRT-LLM PyTorch backend implementation for Gemma4 text model.""" import math -from typing import Dict, Optional, Tuple +from typing import Dict, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -29,6 +29,7 @@ from tensorrt_llm._torch.modules.fused_moe.routing import BaseMoeRoutingMethod from tensorrt_llm._torch.modules.qk_norm_attention import QKNormRoPEAttention from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from ..attention_backend import AttentionMetadata, FlashInferAttentionMetadata @@ -43,10 +44,17 @@ from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding +from ..modules.fused_ops.gelu_tanh_mul_fp4_quant import gelu_tanh_mul_fp4_quant +from ..modules.fused_ops.rmsnorm_fp4_quant import rmsnorm_fp4_quant, rmsnorm_fp4_quant_available +from ..modules.fused_ops.rmsnorm_residual_add import ( + rmsnorm_residual_add, + rmsnorm_residual_add_scale, +) from ..modules.gated_mlp import GatedMLP +from ..modules.gemma4.fused_qkv import gemma4_fused_qkv_norm_rope_quant from ..modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig from ..modules.rms_norm import RMSNorm -from ..utils import ActivationType +from ..utils import ActivationType, Fp4QuantizedTensor, is_torch_compiling from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0" @@ -101,6 +109,59 @@ def gelu_tanh(gate_x: torch.Tensor) -> torch.Tensor: return nn.functional.gelu(gate, approximate="tanh") * x +class _Gemma4GeluQuantMLP(GatedMLP): + """GatedMLP that fuses gelu_tanh+mul into the down_proj NVFP4 quantize. + + On the NVFP4 down_proj path the unfused chain writes the bf16 activation + to HBM (flashinfer_gelu_tanh_and_mul) and immediately reads it back to + quantize (fp4_quantize). The fused Triton kernel does both in one pass + and returns an Fp4QuantizedTensor, which Linear's NVFP4 method consumes + directly (using the same static input_scale / alpha the unfused quantize + would use) - so down_proj sees byte-identical inputs. The unfused path + remains only for configurations the kernel does not support. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Decided lazily on first use: down_proj quant attributes are + # finalized after __init__. + self._fused_gelu_quant: Optional[bool] = None + + def _fused_gelu_quant_enabled(self) -> bool: + if self._fused_gelu_quant is None: + dp = self.down_proj + self._fused_gelu_quant = ( + self.activation is gelu_tanh + # Mirror the checks the pre-quantized-input path in + # Linear._input_prepare enforces for Fp4QuantizedTensor. + and getattr(dp, "has_nvfp4", False) + and not getattr(dp, "force_dynamic_quantization", True) + and getattr(dp, "input_scale", None) is not None + and getattr(dp, "pre_quant_scale", None) is None + and getattr(dp, "scaling_vector_size", None) == 16 + ) + logger.info_once( + f"Gemma4 fused gelu_tanh+NVFP4-quant MLP path: " + f"{'enabled' if self._fused_gelu_quant else 'disabled'}", + key="gemma4_fused_gelu_quant", + ) + return self._fused_gelu_quant + + def _apply_activation(self, x, *, has_lora: bool = False): + if ( + not has_lora + and isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.dtype == torch.bfloat16 + and x.shape[-1] % 32 == 0 + and not is_torch_compiling() + and self._fused_gelu_quant_enabled() + ): + fp4, sf = gelu_tanh_mul_fp4_quant(x, self.down_proj.input_scale) + return Fp4QuantizedTensor(fp4, sf) + return super()._apply_activation(x, has_lora=has_lora) + + # --------------------------------------------------------------------------- # Q-only linear for KV shared layers # --------------------------------------------------------------------------- @@ -308,6 +369,41 @@ def __init__( has_weights=False, ) + # Fused QKV prep (norm+rope+FP8 quant in one Triton kernel). Decided + # lazily on first apply_rope because attn.flashinfer_backend and + # has_fp8_kv_cache are finalized after __init__. The unfused path + # below stays as the reference / fallback. + self._fused_qkv_prep: Optional[bool] = None + self._fused_prep_blocked = False + + def _fused_qkv_prep_enabled(self) -> bool: + if self._fused_qkv_prep is None: + rot = self.rotary_emb + self._fused_qkv_prep = ( + not self.is_kv_shared + and not self.fuse_qk_norm_rope + and not self.skip_rope + # The kernel emits KV-cache-dtype FP8 and replicates the + # flashinfer 2-target rope; only the profiled trtllm-gen + + # FP8-KV serving path is routed through it. + and getattr(self.attn, "flashinfer_backend", None) == "trtllm-gen" + and getattr(self.attn, "has_fp8_kv_cache", False) + and rot is not None + and getattr(rot, "is_neox", False) + and not getattr(rot, "inverse", False) + and rot.head_dim == self.head_dim + and rot.rotary_cos_sin.dtype == torch.float32 + and rot.rotary_cos_sin.dim() == 3 + and rot.rotary_cos_sin.shape[1] == 2 + and rot.rotary_cos_sin.shape[2] * 2 == self.head_dim + and rot.rotary_cos_sin.is_contiguous() + and self.q_norm.weight.shape == (self.head_dim,) + and self.k_norm.weight.shape == (self.head_dim,) + and self.q_norm.variance_epsilon == self.k_norm.variance_epsilon + and self.q_norm.variance_epsilon == self.v_norm.variance_epsilon + ) + return self._fused_qkv_prep + def apply_rope( self, q: torch.Tensor, @@ -332,6 +428,33 @@ def apply_rope( # Return None for k/v so FlashInfer skips cache append. return q_raw, None, None + # Fused path: one Triton kernel reads the packed QKV GEMM output + # (strided per-head views), applies q/k/v RMSNorm + RoPE, and + # emits FP8 Q/K/V directly — replacing the reshape copies, three + # norms, the rope launch, and the backend's three + # .to(float8_e4m3fn) casts. + if ( + k is None + and v is None + and position_ids is not None + and q.dtype == torch.bfloat16 + and not self._fused_prep_blocked + and not is_torch_compiling() + and self._fused_qkv_prep_enabled() + ): + return gemma4_fused_qkv_norm_rope_quant( + q, + position_ids, + self.rotary_emb.rotary_cos_sin, + self.q_norm.weight, + self.k_norm.weight, + self.q_norm.variance_epsilon, + self.num_heads, + self.num_key_value_heads, + self.head_dim, + out_fp8=True, + ) + q, k, v = self.split_qkv(q, k, v) # For K=V layers, weight mapper duplicates k_proj weights into # v_proj, so after split_qkv v already equals k_proj(x). @@ -374,6 +497,9 @@ def forward( "Only FlashInfer backend supports custom attention mask currently." ) assert attention_mask == CustomAttentionMask.CUSTOM + # Custom-mask (multimodal) prefill uses the Triton prefill fallback, + # which consumes BF16 q/k/v — keep the unfused prep for those calls. + self._fused_prep_blocked = attention_mask_data is not None return super().forward( position_ids=position_ids, hidden_states=hidden_states, @@ -563,7 +689,7 @@ def __init__( # TP we take the default (full tp_size) and rely on the down_proj # allreduce to produce a full-rank activation. mlp_tp_size = 1 if model_config.mapping.enable_attention_dp else None - self.mlp = GatedMLP( + self.mlp = _Gemma4GeluQuantMLP( hidden_size=config.hidden_size, intermediate_size=intermediate_size, bias=False, @@ -599,6 +725,20 @@ def __init__( # Layer scalar self.register_buffer("layer_scalar", torch.ones(1)) + # Fused layer tail (post_ffn norm + residual add + layer_scalar). + # Decided lazily on first forward: norm/scalar buffers are finalized + # after weight loading. + self._fused_tail: Optional[bool] = None + # Fused post_attention norm + residual add, and fused pre_ffn norm + + # gate_up NVFP4 quantize (both decided lazily like the tail). + self._fused_norm_add: Optional[bool] = None + self._fused_norm_quant: Optional[bool] = None + # The next layer's input_layernorm (wired by Gemma4TextModel); when + # set, the fused tail can emit that norm as a second output so the + # next layer skips its standalone input-norm pass. + self._next_input_layernorm: Optional[RMSNorm] = None + self._fused_tail_norm2: Optional[bool] = None + # MoE block (parallel with dense MLP) self.enable_moe_block = getattr(config, "enable_moe_block", False) if self.enable_moe_block: @@ -635,6 +775,83 @@ def __init__( dtype=config.torch_dtype, ) + def _fused_tail_enabled(self) -> bool: + if self._fused_tail is None: + norm = self.post_feedforward_layernorm + self._fused_tail = ( + not self.enable_moe_block + and not self.hidden_size_per_layer_input + # The kernel replicates the plain (use_gemma=False) + # flashinfer rmsnorm the module dispatches to. + and not getattr(norm, "use_gemma", True) + and isinstance(getattr(norm, "weight", None), torch.Tensor) + and norm.weight.dtype == torch.bfloat16 + # aten promotes the scalar mul to fp32 only for an fp32 + # buffer; the kernel replicates exactly that recipe. + and self.layer_scalar.dtype == torch.float32 + and self.layer_scalar.numel() == 1 + ) + logger.info_once( + f"Gemma4 fused layer-tail (norm+add+scale) path: " + f"{'enabled' if self._fused_tail else 'disabled'}", + key="gemma4_fused_tail", + ) + return self._fused_tail + + def _norm_is_plain_bf16(self, norm: RMSNorm) -> bool: + # The fused kernels replicate the plain (use_gemma=False) flashinfer + # rmsnorm the module dispatches to for bf16 weights. + return ( + not getattr(norm, "use_gemma", True) + and isinstance(getattr(norm, "weight", None), torch.Tensor) + and norm.weight.dtype == torch.bfloat16 + ) + + def _fused_norm_add_enabled(self) -> bool: + if self._fused_norm_add is None: + self._fused_norm_add = self._norm_is_plain_bf16(self.post_attention_layernorm) + logger.info_once( + f"Gemma4 fused post-attention norm+add path: " + f"{'enabled' if self._fused_norm_add else 'disabled'}", + key="gemma4_fused_norm_add", + ) + return self._fused_norm_add + + def _fused_norm_quant_enabled(self) -> bool: + if self._fused_norm_quant is None: + gu = self.mlp.gate_up_proj + self._fused_norm_quant = ( + # flashinfer's CuTe-DSL kernel backs this fusion (SM100+, + # needs nvidia-cutlass-dsl importable). + rmsnorm_fp4_quant_available() + and not self.enable_moe_block + and self._norm_is_plain_bf16(self.pre_feedforward_layernorm) + # Mirror the checks the pre-quantized-input path in + # Linear._input_prepare enforces for Fp4QuantizedTensor. + and getattr(gu, "has_nvfp4", False) + and not getattr(gu, "force_dynamic_quantization", True) + and getattr(gu, "input_scale", None) is not None + and getattr(gu, "pre_quant_scale", None) is None + and getattr(gu, "scaling_vector_size", None) == 16 + ) + logger.info_once( + f"Gemma4 fused pre-ffn norm+NVFP4-quant path: " + f"{'enabled' if self._fused_norm_quant else 'disabled'}", + key="gemma4_fused_norm_quant", + ) + return self._fused_norm_quant + + def _fused_tail_norm2_enabled(self) -> bool: + if self._fused_tail_norm2 is None: + nxt = self._next_input_layernorm + self._fused_tail_norm2 = nxt is not None and self._norm_is_plain_bf16(nxt) + logger.info_once( + f"Gemma4 fused tail next-layer-norm output path: " + f"{'enabled' if self._fused_tail_norm2 else 'disabled'}", + key="gemma4_fused_tail_norm2", + ) + return self._fused_tail_norm2 + @torch.inference_mode() def forward( self, @@ -644,8 +861,9 @@ def forward( residual: Optional[torch.Tensor] = None, attention_mask_data: Optional[torch.Tensor] = None, per_layer_input: Optional[torch.Tensor] = None, + pre_normed: Optional[torch.Tensor] = None, **kwargs, - ) -> torch.Tensor: + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: # lora_params is handled explicitly by the MLP call below; drop it # from kwargs so it is not forwarded into self.self_attn (base # Attention.forward would raise on the extra kwarg). @@ -655,10 +873,17 @@ def forward( target_dtype = self.input_layernorm.weight.dtype if hidden_states.dtype != target_dtype: hidden_states = hidden_states.to(target_dtype) + # pre_normed was computed from the un-cast tensor; recompute. + pre_normed = None # Self-attention residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) + if pre_normed is not None: + # The previous layer's fused tail already emitted this layer's + # input norm as its second output. + hidden_states = pre_normed + else: + hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn( position_ids=position_ids, hidden_states=hidden_states, @@ -669,13 +894,54 @@ def forward( attention_mask_data=attention_mask_data, **kwargs, ) - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = residual + hidden_states + # Fused post_attention RMSNorm + residual add (one kernel instead of + # a norm round-trip plus a separate add). The unfused sequence remains + # for configurations the kernel does not support. + if ( + isinstance(hidden_states, torch.Tensor) + and hidden_states.dim() == 2 + and hidden_states.dtype == torch.bfloat16 + and residual.shape == hidden_states.shape + and not is_torch_compiling() + and self._fused_norm_add_enabled() + ): + hidden_states = rmsnorm_residual_add( + hidden_states, + residual, + self.post_attention_layernorm.weight, + self.post_attention_layernorm.variance_epsilon, + ) + else: + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states # Feed-forward (dense MLP + optional MoE in parallel) residual = hidden_states - hidden_states = self.pre_feedforward_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states, lora_params=lora_params) + # Fused pre_feedforward RMSNorm + gate_up NVFP4 quantize: the normed + # tensor is consumed only by gate_up_proj's input quantize, so emit + # the FP4 payload + swizzled scales directly (same static + # input_scale / alpha the unfused quantize would use; the fused + # kernel skips the intermediate bf16 round, so the GEMM inputs are + # near- rather than byte-identical). + if ( + not lora_params + and isinstance(hidden_states, torch.Tensor) + and hidden_states.dim() == 2 + and hidden_states.dtype == torch.bfloat16 + and hidden_states.shape[-1] % 32 == 0 + and not is_torch_compiling() + and self._fused_norm_quant_enabled() + ): + fp4, sf = rmsnorm_fp4_quant( + hidden_states, + self.pre_feedforward_layernorm.weight, + self.pre_feedforward_layernorm.variance_epsilon, + self.mlp.gate_up_proj.input_scale, + ) + hidden_states = self.mlp(Fp4QuantizedTensor(fp4, sf), lora_params=lora_params) + else: + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states, lora_params=lora_params) if self.enable_moe_block: # MLP path: post-norm the MLP output @@ -697,6 +963,42 @@ def forward( # Combine MLP + MoE hidden_states = hidden_states_mlp + hidden_states_moe + # Fused tail: post_ffn RMSNorm + residual add + fp32 layer_scalar mul + # + bf16 cast in one kernel (the unfused chain materializes a full + # fp32 [M, H] tensor because layer_scalar is an fp32 buffer). The + # unfused sequence below remains for configurations the kernel does + # not support (MoE block, PLE, non-bf16). + if ( + per_layer_input is None + and isinstance(hidden_states, torch.Tensor) + and hidden_states.dim() == 2 + and hidden_states.dtype == torch.bfloat16 + and residual.shape == hidden_states.shape + and not is_torch_compiling() + and self._fused_tail_enabled() + ): + if self._fused_tail_norm2_enabled(): + # Also emit the next layer's input norm as a second output + # (returned as a (hidden, pre_normed) pair the model loop + # hands to the next layer). + nxt = self._next_input_layernorm + return rmsnorm_residual_add_scale( + hidden_states, + residual, + self.post_feedforward_layernorm.weight, + self.layer_scalar, + self.post_feedforward_layernorm.variance_epsilon, + next_norm_weight=nxt.weight, + next_norm_eps=nxt.variance_epsilon, + ) + return rmsnorm_residual_add_scale( + hidden_states, + residual, + self.post_feedforward_layernorm.weight, + self.layer_scalar, + self.post_feedforward_layernorm.variance_epsilon, + ) + hidden_states = self.post_feedforward_layernorm(hidden_states) hidden_states = residual + hidden_states @@ -754,6 +1056,11 @@ def __init__(self, model_config: ModelConfig[Gemma4TextConfig]): for layer_idx in range(pretrained.num_hidden_layers) ] ) + # Wire each layer to its successor's input norm so the fused tail + # can emit that norm as a second output (the last layer keeps the + # plain single-output tail; the model-final self.norm is separate). + for prev_layer, next_layer in zip(self.layers[:-1], self.layers[1:]): + prev_layer._next_input_layernorm = next_layer.input_layernorm self.norm = RMSNorm( hidden_size=pretrained.hidden_size, @@ -884,10 +1191,11 @@ def forward( ple_ids = ple_input_ids if ple_input_ids is not None else input_ids per_layer_inputs = self._compute_per_layer_inputs(ple_ids, hidden_states) + pre_normed = None for i, decoder_layer in enumerate(self.layers): per_layer_input = per_layer_inputs[:, i, :] if per_layer_inputs is not None else None - hidden_states = decoder_layer( + layer_out = decoder_layer( position_ids=position_ids, hidden_states=hidden_states, attn_metadata=attn_metadata, @@ -895,8 +1203,15 @@ def forward( local_attention_mask_data if decoder_layer.is_sliding else None ), per_layer_input=per_layer_input, + pre_normed=pre_normed, **kwargs, ) + # A layer whose fused tail also produced the next layer's input + # norm returns a (hidden, pre_normed) pair. + if isinstance(layer_out, tuple): + hidden_states, pre_normed = layer_out + else: + hidden_states, pre_normed = layer_out, None if hidden_states.dtype != self.dtype: hidden_states = hidden_states.to(self.dtype) diff --git a/tensorrt_llm/_torch/modules/fused_ops/__init__.py b/tensorrt_llm/_torch/modules/fused_ops/__init__.py new file mode 100644 index 000000000000..c060e71c98d2 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_ops/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Standalone fused elementwise/norm/quant operators shared across models. + +Each submodule hosts one fused operator as a plain Python callable (not a +registered custom op). Callers own the enablement checks and keep the +unfused op chains as fallbacks for unsupported configurations; see the +Gemma4 decoder (tensorrt_llm/_torch/models/modeling_gemma4.py) for the +reference usage pattern. +""" diff --git a/tensorrt_llm/_torch/modules/fused_ops/gelu_tanh_mul_fp4_quant.py b/tensorrt_llm/_torch/modules/fused_ops/gelu_tanh_mul_fp4_quant.py new file mode 100644 index 000000000000..8369be52a0d9 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_ops/gelu_tanh_mul_fp4_quant.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Fused gelu_tanh+mul+NVFP4-quantize (single Triton kernel). + +Replaces the unfused pair between a packed [gate | up] GEMM output and an +NVFP4 GEMM consumer (e.g. a gated-MLP down_proj): + + trtllm::flashinfer_gelu_tanh_and_mul (bf16 out, full HBM round-trip) + -> trtllm::fp4_quantize (reads it back, emits FP4 + swizzled scales) + +with one Triton kernel that reads the packed [gate | up] GEMM output once and +emits the packed E2M1 payload plus the 128x4-swizzled E4M3 block scales +directly - eliminating the intermediate activation write+read. + +Numerics replicate the unfused chain byte-for-byte on SM100: +- gelu_tanh uses flashinfer's exact expression order and its `tanh.approx.f32` + hardware instruction (flashinfer JIT builds with -use_fast_math); +- the product is rounded to bf16 (the write the unfused path performs) before + quantization; +- the scale/quant math mirrors cvt_warp_fp16_to_fp4 in + cpp/tensorrt_llm/kernels/quantization.cuh: rcp.approx.ftz.f32 reciprocals, + cvt.rn.satfinite e4m3 scale rounding, and the cvt.rn.satfinite.e2m1x2.f32 + payload conversion, with scales stored at the 128x4 swizzled offsets of + get_sf_out_offset_128x4 (rows padded to 128; padding is zero-filled here, + uninitialized in the unfused op). + +Callers own enablement and keep the unfused pair as the fallback for +configurations this kernel does not support (non-NVFP4 consumer, LoRA, +torch.compile, ...); see the gelu_tanh + down_proj quantize fusion in +modeling_gemma4.py. +""" + +from typing import Tuple + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _rcp_approx(x): + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=f,f", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _gelu_tanh_mul_fp4_kernel( + x_ptr, # [M, 2I] bf16 packed [gate | up], row stride SX + out_ptr, # [M, I//2] uint8, two e2m1 per byte (element 0 = low nibble) + sf_ptr, # [pad128(M) * 4 * NKT] uint8, swizzled 128x4 layout + gs_ptr, # [1] fp32 global scale (down_proj.input_scale) + M, + SX, # x row stride (elements) + NKT, # numKTiles = ceil((IDIM/16) / 4) + IDIM: tl.constexpr, # intermediate size (elements, multiple of 16) + BM: tl.constexpr, + BK: tl.constexpr, # multiple of 16 +): + rows = tl.program_id(0) * BM + tl.arange(0, BM).to(tl.int64) + rmask = rows < M + cols = tl.program_id(1) * BK + tl.arange(0, BK) + mask = rmask[:, None] & (cols[None, :] < IDIM) + + gate = tl.load(x_ptr + rows[:, None] * SX + cols[None, :], mask=mask, other=0.0).to(tl.float32) + up = tl.load(x_ptr + rows[:, None] * SX + IDIM + cols[None, :], mask=mask, other=0.0).to( + tl.float32 + ) + + # gelu_tanh exactly as flashinfer's JIT kernel computes it, then the bf16 + # round the unfused path performs when writing its output tensor. + inner = 0.7978845608028654 * (gate + 0.044715 * gate * gate * gate) + t = tl.inline_asm_elementwise( + "tanh.approx.f32 $0, $1;", "=f,f", [inner], dtype=tl.float32, is_pure=True, pack=1 + ) + act = gate * (0.5 * (1.0 + t)) + v = (act * up).to(tl.bfloat16).to(tl.float32) + + # Per-16-element block amax (exact; max is order-insensitive). + v16 = tl.reshape(v, (BM, BK // 16, 16)) + vmax = tl.max(tl.abs(v16), axis=2) # [BM, BK/16] + + # Scale math replicating cvt_warp_fp16_to_fp4 (e4m3 branch): + # SFValue = gs * (vecMax * rcp.approx(6)); sf8 = e4m3(SFValue) + # outputScale = vecMax != 0 ? rcp.approx(f32(sf8) * rcp.approx(gs)) : 0 + gs = tl.load(gs_ptr) + rcp6 = _rcp_approx(tl.full((BM, BK // 16), 6.0, tl.float32)) + sf8 = (gs * (vmax * rcp6)).to(tl.float8e4nv) + rcpgs = _rcp_approx(tl.full((BM, BK // 16), 0.0, tl.float32) + gs) + oscale = _rcp_approx(sf8.to(tl.float32) * rcpgs) + oscale = tl.where(vmax != 0.0, oscale, 0.0) + + # Scale store at the swizzled 128x4 offsets (get_sf_out_offset_128x4). + kvec = (tl.program_id(1) * (BK // 16) + tl.arange(0, BK // 16)).to(tl.int64) + kvmask = rmask[:, None] & (kvec[None, :] * 16 < IDIM) + m2 = rows[:, None] + k2 = kvec[None, :] + sfoff = ( + (m2 // 128) * (NKT * 512) + + (k2 // 4) * 512 + + (m2 % 32) * 16 + + ((m2 % 128) // 32) * 4 + + (k2 % 4) + ) + tl.store(sf_ptr + sfoff, sf8.to(tl.uint8, bitcast=True), mask=kvmask) + + # E2M1 conversion + pairwise packing via the same PTX instruction the + # CUDA quantize kernel uses (first source operand -> high nibble). + y = tl.reshape(v16 * oscale[:, :, None], (BM, BK)) + lo, hi = tl.split(tl.reshape(y, (BM, BK // 2, 2))) + byte = tl.inline_asm_elementwise( + "{ .reg .b8 t; cvt.rn.satfinite.e2m1x2.f32 t, $2, $1; cvt.u16.u8 $0, t; }", + "=h,f,f", + [lo, hi], + dtype=tl.uint16, + is_pure=True, + pack=1, + ).to(tl.uint8) + + ocols = tl.program_id(1) * (BK // 2) + tl.arange(0, BK // 2) + omask = rmask[:, None] & (ocols[None, :] * 2 < IDIM) + tl.store(out_ptr + rows[:, None] * (IDIM // 2) + ocols[None, :], byte, mask=omask) + + +def sf_swizzled_offsets(m: int, nvec: int, device: torch.device) -> torch.Tensor: + """Flat swizzled offsets of the valid (row, kvec) scale region. + + Mirrors get_sf_out_offset_128x4; used by the parity tests to compare only + the valid region (the unfused op leaves the 128-row padding + uninitialized). + """ + nkt = (nvec + 3) // 4 + mm = torch.arange(m, device=device, dtype=torch.int64)[:, None] + kk = torch.arange(nvec, device=device, dtype=torch.int64)[None, :] + off = ( + (mm // 128) * (nkt * 512) + + (kk // 4) * 512 + + (mm % 32) * 16 + + ((mm % 128) // 32) * 4 + + (kk % 4) + ) + return off.reshape(-1) + + +def gelu_tanh_mul_fp4_quant( + x: torch.Tensor, global_scale: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fused gelu_tanh(gate) * up + NVFP4 block-scale quantize. + + Args: + x: [M, 2*I] bf16, packed [gate | up] (a row stride larger than the + width is allowed; the innermost dim must be contiguous). + global_scale: [1] fp32 tensor - the consumer Linear's static + ``input_scale`` (448*6/amax convention). + + Returns: + (fp4, sf): the packed E2M1 payload [M, I//2] (uint8, element 2j in + the low nibble of byte j) and the E4M3 block scales (uint8, 1D, + swizzled 128x4 layout padded to 128 rows, padding zero-filled) - + byte-compatible with ``trtllm::fp4_quantize``'s outputs and + consumable as ``Fp4QuantizedTensor(fp4, sf)``. + """ + assert x.dim() == 2 and x.stride(-1) == 1 + assert x.dtype == torch.bfloat16 + assert global_scale.dtype == torch.float32 + m, two_i = x.shape + i = two_i // 2 + assert two_i % 32 == 0, "intermediate size must be a multiple of 16" + + nkt = triton.cdiv(i // 16, 4) + out = torch.empty((m, i // 2), dtype=torch.uint8, device=x.device) + sf = torch.zeros((((m + 127) // 128) * 128 * nkt * 4,), dtype=torch.uint8, device=x.device) + if m == 0: + return out, sf + + # B200-tuned; fixed (no runtime autotune) so launches stay deterministic + # under CUDA-graph capture. Measured 234 -> 145 us vs the unfused pair + # at [6455, 2x21504]; 13.6 -> 7.6 us graph-replayed at the 228-token + # decode shape. + bm, bk = 8, 512 + grid = (triton.cdiv(m, bm), triton.cdiv(i, bk)) + _gelu_tanh_mul_fp4_kernel[grid]( + x, out, sf, global_scale, m, x.stride(0), nkt, IDIM=i, BM=bm, BK=bk, num_warps=8 + ) + return out, sf diff --git a/tensorrt_llm/_torch/modules/fused_ops/rmsnorm_fp4_quant.py b/tensorrt_llm/_torch/modules/fused_ops/rmsnorm_fp4_quant.py new file mode 100644 index 000000000000..9b323de3e8c9 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_ops/rmsnorm_fp4_quant.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Fused RMSNorm + NVFP4 quantize (flashinfer CuTe-DSL ``rmsnorm_fp4quant``). + +Replaces the unfused pair in front of an NVFP4 GEMM whose input is the +RMSNorm of a bf16 activation (e.g. a pre-feedforward norm feeding a +gate_up projection): + + trtllm::flashinfer_rmsnorm (bf16 out, full [M, H] HBM round-trip) + -> trtllm::fp4_quantize (reads it back, emits FP4 + swizzled scales) + +with flashinfer's CuTe-DSL ``rmsnorm_fp4quant`` (SM100+), which reads the +input once and emits the packed E2M1 payload plus the 128x4-swizzled E4M3 +block scales directly - eliminating the intermediate bf16 activation +write+read. + +TRT-LLM's static ``input_scale`` (448*6/amax convention) is exactly +flashinfer's ``global_scale``: both the unfused ``trtllm::fp4_quantize`` and +the fused kernel store ``e4m3(gs * blockAmax / 6)`` as the block scale, and +``is_sf_swizzled_layout=True`` emits the same 128x4 layout as +``get_sf_out_offset_128x4`` with the identical padded size +(``ceil(M/128) * 512 * numKTiles``) - so the outputs are consumable as +``Fp4QuantizedTensor(fp4, sf)`` wherever the unfused op's would be. + +Numerics: the fused kernel quantizes the fp32 norm result directly, without +the intermediate bf16 round the unfused chain performs when materializing +the normed tensor. Outputs are therefore near- but not byte-identical to the +unfused pair (~1.4% of payload nibbles differ by one code step at Gemma4 +serving shapes); the quantization error against the fp32 norm is statistically +identical (see tests/unittest/_torch/modules/fused_ops/test_rmsnorm_fp4_quant.py). + +Callers own enablement and keep the unfused pair as the fallback for +configurations this kernel does not support (non-NVFP4 consumer, LoRA, +torch.compile, flashinfer's CuTe-DSL kernels unavailable, ...); see the +pre-feedforward norm + gate_up quantize fusion in modeling_gemma4.py. +""" + +from typing import Tuple + +import torch + +from ...flashinfer_utils import IS_FLASHINFER_AVAILABLE + +if IS_FLASHINFER_AVAILABLE: + try: + # None (rather than an ImportError) when flashinfer's optional + # CuTe-DSL dependency (nvidia-cutlass-dsl) is not importable. + from flashinfer.norm import rmsnorm_fp4quant as _rmsnorm_fp4quant + except ImportError: # flashinfer version without the CuTe-DSL kernels + _rmsnorm_fp4quant = None +else: + _rmsnorm_fp4quant = None + + +def rmsnorm_fp4_quant_available() -> bool: + """Whether the flashinfer CuTe-DSL fused norm+quant kernel is usable.""" + return _rmsnorm_fp4quant is not None + + +def rmsnorm_fp4_quant( + x: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + global_scale: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fused RMSNorm + NVFP4 block-scale quantize. + + Args: + x: [M, H] bf16 input (a row stride larger than the width is + allowed; the innermost dim must be contiguous). + norm_weight: [H] bf16 RMSNorm weight (plain multiplier convention, + ``use_gemma=False``). + eps: RMSNorm epsilon. + global_scale: [1] fp32 tensor - the consumer Linear's static + ``input_scale`` (448*6/amax convention). + + Returns: + (fp4, sf): the packed E2M1 payload [M, H//2] (uint8, element 2j in + the low nibble of byte j) and the E4M3 block scales (uint8, 1D, + swizzled 128x4 layout padded to 128 rows) - layout-compatible with + ``trtllm::fp4_quantize``'s outputs and consumable as + ``Fp4QuantizedTensor(fp4, sf)``. + """ + assert rmsnorm_fp4_quant_available() + assert x.dim() == 2 and x.stride(-1) == 1 + assert x.dtype == torch.bfloat16 + assert global_scale.dtype == torch.float32 + m, h = x.shape + assert h % 16 == 0 and h >= 64, "hidden size must be a multiple of 16 and >= 64" + assert norm_weight.shape == (h,) + + if m == 0: + fp4 = torch.empty((0, h // 2), dtype=torch.uint8, device=x.device) + sf = torch.empty((0,), dtype=torch.uint8, device=x.device) + return fp4, sf + + fp4, sf = _rmsnorm_fp4quant( + x, + norm_weight, + global_scale=global_scale.reshape(1), + eps=eps, + block_size=16, + scale_format="e4m3", + is_sf_swizzled_layout=True, + ) + return fp4.view(torch.uint8), sf.view(torch.uint8) diff --git a/tensorrt_llm/_torch/modules/fused_ops/rmsnorm_residual_add.py b/tensorrt_llm/_torch/modules/fused_ops/rmsnorm_residual_add.py new file mode 100644 index 000000000000..8bf8cbc32954 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_ops/rmsnorm_residual_add.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Fused post-norm residual kernels: ``residual + rmsnorm(x)`` variants. + +One Triton kernel body (specialized via constexpr flags) serves two fusions +on a post-norm residual stream (norm the branch output, then add it to the +residual - the Gemma-style decoder-layer pattern): + +1. ``rmsnorm_residual_add_scale``: RMSNorm + residual add + fp32 scalar mul + + bf16 cast in one kernel, replacing + + flashinfer_rmsnorm(x) (read + write [M, H] bf16) + residual + normed (2 reads + write, bf16) + hidden * scale (fp32[1] buffer promotes the mul: + read bf16 + WRITE FP32 [M, H]) + .to(bf16) at the consumer's entry (read fp32 + write bf16) + + Optionally (``next_norm_weight``) the kernel also emits the RMSNorm of + the result as a second output (e.g. the next decoder layer's input norm), + so the consumer skips its standalone norm pass entirely (one extra [M, H] + write here replaces a read + write there). + +2. ``rmsnorm_residual_add``: the same body without the scalar stage, + replacing + + flashinfer_rmsnorm(x) (read + write [M, H] bf16) + residual + normed (2 reads + write, bf16) + + with one kernel (2 bf16 reads + 1 bf16 write). + +Numerics replicate the unfused chains op-for-op: each norm mirrors +flashinfer's RMSNormKernel (fp32 sum of squares, +`rsqrt.approx.ftz.f32(ssq/d + eps)`, `(x * rcp) * w`, bf16 round), the add +rounds to bf16 exactly like the aten bf16 add (fp32 compute, single round), +the fp32 scalar multiply and final bf16 round match the promoted aten mul + +`.to(bfloat16)`, and the secondary norm reads the bf16-rounded primary +output exactly as a standalone downstream norm would. The only residual +difference is the fp32 reduction order inside the sums of squares (measured +~5e-6 one-step bf16 flips at Gemma4 serving shapes). + +Callers own enablement and keep the unfused chains as fallbacks for +configurations these kernels do not support (non-bf16 norms, torch.compile, +...); see the residual-chain fusions in modeling_gemma4.py. +""" + +from typing import Optional, Tuple, Union + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _rmsnorm_residual_add_kernel( + x_ptr, # [M, H] bf16 kernel input (the branch output being normed) + r_ptr, # [M, H] bf16 residual + w_ptr, # [H] bf16 RMSNorm weight for x + sc_ptr, # [1] fp32 scale (unused when HAS_SCALE == False) + o_ptr, # [M, H] bf16 out + w2_ptr, # [H] bf16 secondary norm weight (unused when HAS_NORM2 == False) + n2_ptr, # [M, H] bf16 secondary norm out (unused when HAS_NORM2 == False) + M, + SX, + SR, + SO, + SN2, + EPS, + EPS2, + H: tl.constexpr, + BH: tl.constexpr, # next_power_of_2(H) + BM: tl.constexpr, + HAS_SCALE: tl.constexpr, + HAS_NORM2: tl.constexpr, +): + rows = tl.program_id(0) * BM + tl.arange(0, BM).to(tl.int64) + rmask = rows < M + cols = tl.arange(0, BH) + cmask = cols < H + mask = rmask[:, None] & cmask[None, :] + + x = tl.load(x_ptr + rows[:, None] * SX + cols[None, :], mask=mask, other=0.0).to(tl.float32) + # flashinfer RMSNormKernel: rms_rcp = rsqrt.approx.ftz(sum_sq/d + eps); + # out = (x * rms_rcp) * (weight_bias=0 + w), bf16 round on store. + ssq = tl.sum(x * x, axis=1) + rcp = tl.inline_asm_elementwise( + "rsqrt.approx.ftz.f32 $0, $1;", + "=f,f", + [ssq / H + EPS], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + w = tl.load(w_ptr + cols, mask=cmask, other=0.0).to(tl.float32) + normed = ((x * rcp[:, None]) * w[None, :]).to(tl.bfloat16) + # aten bf16 add computes in fp32 and rounds once. + r = tl.load(r_ptr + rows[:, None] * SR + cols[None, :], mask=mask, other=0.0).to(tl.float32) + s = (r + normed.to(tl.float32)).to(tl.bfloat16) + if HAS_SCALE: + # aten mul(bf16, fp32[1]) promotes to fp32; the consumer's dtype + # guard rounds back to bf16. + sc = tl.load(sc_ptr) + out = (s.to(tl.float32) * sc).to(tl.bfloat16) + else: + out = s + tl.store(o_ptr + rows[:, None] * SO + cols[None, :], out, mask=mask) + + if HAS_NORM2: + # The secondary RMSNorm, computed on the bf16-rounded primary + # output exactly as a standalone flashinfer_rmsnorm would read it. + # Masked columns hold zeros (x, r and w all load 0 there), so they + # contribute nothing to the sum of squares. + outf = out.to(tl.float32) + ssq2 = tl.sum(outf * outf, axis=1) + rcp2 = tl.inline_asm_elementwise( + "rsqrt.approx.ftz.f32 $0, $1;", + "=f,f", + [ssq2 / H + EPS2], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + w2 = tl.load(w2_ptr + cols, mask=cmask, other=0.0).to(tl.float32) + n2 = ((outf * rcp2[:, None]) * w2[None, :]).to(tl.bfloat16) + tl.store(n2_ptr + rows[:, None] * SN2 + cols[None, :], n2, mask=mask) + + +def _launch_norm_add( + x: torch.Tensor, + residual: torch.Tensor, + norm_weight: torch.Tensor, + scale: Optional[torch.Tensor], + eps: float, + next_norm_weight: Optional[torch.Tensor], + next_norm_eps: float, +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + assert x.dim() == 2 and x.stride(-1) == 1 + assert x.dtype == torch.bfloat16 and residual.dtype == torch.bfloat16 + assert residual.shape == x.shape and residual.stride(-1) == 1 + m, h = x.shape + assert norm_weight.shape == (h,) + if scale is not None: + assert scale.dtype == torch.float32 and scale.numel() == 1 + has_norm2 = next_norm_weight is not None + if has_norm2: + assert next_norm_weight.shape == (h,) + o = torch.empty_like(x) + n2 = torch.empty_like(x) if has_norm2 else None + if m == 0: + return (o, n2) if has_norm2 else o + + # B200-tuned; fixed (no runtime autotune) so launches stay deterministic + # under CUDA-graph capture. Measured 239 -> 39 us vs the 4-op chain at + # [7700, 5376]; 34 -> 16 us at the 228-token decode shape; ~6 TB/s + # effective (memory roofline). A block-size/warp sweep confirmed + # BM=1/num_warps=4 optimal for all three variants. + grid = (triton.cdiv(m, 1),) + _rmsnorm_residual_add_kernel[grid]( + x, + residual, + norm_weight, + scale if scale is not None else norm_weight, + o, + next_norm_weight if has_norm2 else norm_weight, + n2 if has_norm2 else o, + m, + x.stride(0), + residual.stride(0), + o.stride(0), + n2.stride(0) if has_norm2 else 0, + eps, + next_norm_eps, + H=h, + BH=triton.next_power_of_2(h), + BM=1, + HAS_SCALE=scale is not None, + HAS_NORM2=has_norm2, + num_warps=4, + ) + return (o, n2) if has_norm2 else o + + +def rmsnorm_residual_add_scale( + x: torch.Tensor, + residual: torch.Tensor, + norm_weight: torch.Tensor, + scale: torch.Tensor, + eps: float, + next_norm_weight: Optional[torch.Tensor] = None, + next_norm_eps: float = 0.0, +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """Fused ``bf16(fp32(residual + rmsnorm(x)) * scale)``. + + Args: + x: [M, H] bf16 - the branch output (e.g. an MLP down_proj output). + residual: [M, H] bf16 - the residual stream. + norm_weight: [H] bf16 RMSNorm weight for ``x`` (plain multiplier + convention, ``use_gemma=False``). + scale: [1] fp32 tensor (e.g. a checkpoint-loaded per-layer scalar; + any value). + eps: RMSNorm epsilon. + next_norm_weight: optional [H] bf16 - a second RMSNorm weight. When + given, the kernel also returns ``rmsnorm(out, next_norm_weight)`` + computed on the bf16-rounded primary output, replacing the + consumer's standalone norm (e.g. the next decoder layer's input + norm). + next_norm_eps: epsilon for the secondary norm. + + Returns: + [M, H] bf16 output - or ``(out, normed_next)`` when + ``next_norm_weight`` is given. The intermediate fp32 tensor of the + unfused chain is never materialized. + """ + return _launch_norm_add(x, residual, norm_weight, scale, eps, next_norm_weight, next_norm_eps) + + +def rmsnorm_residual_add( + x: torch.Tensor, + residual: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Fused ``bf16(fp32(residual) + fp32(rmsnorm(x)))``. + + The post-norm half of a residual chain: RMSNorm on the branch output + (e.g. an attention o_proj output) followed by the aten bf16 residual + add, in one kernel. + + Args: + x: [M, H] bf16 - the branch output. + residual: [M, H] bf16 - the residual stream. + norm_weight: [H] bf16 RMSNorm weight for ``x`` (plain multiplier + convention, ``use_gemma=False``). + eps: RMSNorm epsilon. + + Returns: + [M, H] bf16 - identical (modulo fp32 reduction order in the sum of + squares) to ``residual + flashinfer_rmsnorm(x, norm_weight, eps)``. + """ + out = _launch_norm_add(x, residual, norm_weight, None, eps, None, 0.0) + assert isinstance(out, torch.Tensor) + return out diff --git a/tensorrt_llm/_torch/modules/gemma4/__init__.py b/tensorrt_llm/_torch/modules/gemma4/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/modules/gemma4/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/modules/gemma4/fused_qkv.py b/tensorrt_llm/_torch/modules/gemma4/fused_qkv.py new file mode 100644 index 000000000000..e2cbbcd17997 --- /dev/null +++ b/tensorrt_llm/_torch/modules/gemma4/fused_qkv.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Fused QKV prep for Gemma4: per-head RMSNorm + RoPE + FP8 quant in one kernel. + +Replaces the unfused chain executed per layer per step on the FlashInfer +FP8-KV path: + + split_qkv (strided views) -> reshape copies (3x direct_copy) + -> q_norm / k_norm / v_norm (3x rmsnorm) + -> flashinfer rope in-place on q, k (1 kernel) + -> q/k/v .to(float8_e4m3fn) in the backend (3x float8_copy) + +with a single Triton kernel that reads the packed QKV GEMM output directly +(per-head strided access, no contiguous intermediate), normalizes each head, +applies neox-style RoPE to Q/K heads from the module's fp32 cos/sin table, +and writes packed FP8 (or BF16) Q, K, V. + +Numerics deliberately replicate the unfused path: fp32 accumulation with a +round to bf16 after the norm and again after RoPE, so the final FP8 values +match the reference chain (see tests/unittest/_torch/modules/ +test_gemma4_fused_qkv_prep.py). + +The kernel processes BLOCK_N tokens per program (2D [BLOCK_N, HALF] tiles) +so each thread issues wide vectorized loads/stores instead of one scalar +element; tile size and warp count are fixed host-side as a pure function of +head_dim (no runtime autotuning, so launches stay deterministic under +CUDA-graph capture). Tuned on B200: 363 -> ~75 us/call at ~6.5k tokens for +the 64-head hd=256 shape (~4.3 TB/s effective), 2.8x on the graph-replayed +decode shape. + +The unfused path is kept only for configurations this kernel does not +support (KV-shared layers, non-FP8 KV cache, custom-mask multimodal +prefill, torch.compile). +""" + +from typing import Tuple + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _gemma4_qkv_norm_rope_quant_kernel( + qkv_ptr, # [N, W] bf16, row stride SW; heads packed [q | k | v], head h at col h*HD + q_out_ptr, # [N, NQ*HD] fp8/bf16 contiguous + k_out_ptr, # [N, NK*HD] + v_out_ptr, # [N, NK*HD] + pos_ptr, # [N] int positions + cos_sin_ptr, # [max_pos, 2, HALF] fp32: [:, 0, :]=cos, [:, 1, :]=sin + qw_ptr, # [HD] bf16 q_norm weight + kw_ptr, # [HD] bf16 k_norm weight + SW, # qkv row stride (elements) + NQ, # num q heads + NK, # num kv heads + EPS, + N, # num tokens (rows) + HD: tl.constexpr, # head dim + HALF: tl.constexpr, # HD // 2 (power of two) + OUT_FP8: tl.constexpr, + BLOCK_N: tl.constexpr, # tokens per program (power of two) +): + h = tl.program_id(0) + rows = tl.program_id(1).to(tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N).to(tl.int64) + rmask = rows < N + offs = tl.arange(0, HALF) + + # Head h occupies columns [h*HD, (h+1)*HD) of the packed qkv rows. + base = qkv_ptr + rows[:, None] * SW + h * HD + offs[None, :] + x1 = tl.load(base, mask=rmask[:, None], other=0.0).to(tl.float32) + x2 = tl.load(base + HALF, mask=rmask[:, None], other=0.0).to(tl.float32) + + ssq = tl.sum(x1 * x1, axis=1) + tl.sum(x2 * x2, axis=1) + rms = tl.math.rsqrt(ssq / HD + EPS) # [BLOCK_N] + + if h < NQ + NK: + # Q/K head: norm (with weight) -> bf16 round -> rope -> bf16 round. + wp = qw_ptr if h < NQ else kw_ptr + w1 = tl.load(wp + offs).to(tl.float32) + w2 = tl.load(wp + HALF + offs).to(tl.float32) + y1 = (x1 * rms[:, None] * w1[None, :]).to(tl.bfloat16).to(tl.float32) + y2 = (x2 * rms[:, None] * w2[None, :]).to(tl.bfloat16).to(tl.float32) + pos = tl.load(pos_ptr + rows, mask=rmask, other=0).to(tl.int64) + cs = cos_sin_ptr + pos[:, None] * (2 * HALF) + offs[None, :] + cos = tl.load(cs, mask=rmask[:, None], other=0.0) + sin = tl.load(cs + HALF, mask=rmask[:, None], other=0.0) + o1 = (y1 * cos - y2 * sin).to(tl.bfloat16) + o2 = (y2 * cos + y1 * sin).to(tl.bfloat16) + if h < NQ: + out = q_out_ptr + rows[:, None] * (NQ * HD) + h * HD + offs[None, :] + else: + out = k_out_ptr + rows[:, None] * (NK * HD) + (h - NQ) * HD + offs[None, :] + if OUT_FP8: + tl.store(out, o1.to(tl.float8e4nv), mask=rmask[:, None]) + tl.store(out + HALF, o2.to(tl.float8e4nv), mask=rmask[:, None]) + else: + tl.store(out, o1, mask=rmask[:, None]) + tl.store(out + HALF, o2, mask=rmask[:, None]) + else: + # V head: weightless norm, no rope (v_norm is applied to the raw v + # slice; reads happen before any write, matching HF's + # v_norm-before-k_norm ordering for K=V layers). + o1 = (x1 * rms[:, None]).to(tl.bfloat16) + o2 = (x2 * rms[:, None]).to(tl.bfloat16) + out = v_out_ptr + rows[:, None] * (NK * HD) + (h - NQ - NK) * HD + offs[None, :] + if OUT_FP8: + tl.store(out, o1.to(tl.float8e4nv), mask=rmask[:, None]) + tl.store(out + HALF, o2.to(tl.float8e4nv), mask=rmask[:, None]) + else: + tl.store(out, o1, mask=rmask[:, None]) + tl.store(out + HALF, o2, mask=rmask[:, None]) + + +def gemma4_fused_qkv_norm_rope_quant( + qkv: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + out_fp8: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the fused per-head norm + rope + quant over a packed QKV tensor. + + Args: + qkv: [N, num_q_heads*head_dim + 2*num_kv_heads*head_dim] bf16, packed + [q | k | v]; may have a row stride larger than its width. + position_ids: flattenable to [N]; integer rope positions per token. + cos_sin: [max_positions, 2, head_dim // 2] fp32 contiguous table + (RotaryEmbedding.rotary_cos_sin layout). + q_weight/k_weight: [head_dim] q_norm/k_norm weights. + eps: shared RMSNorm epsilon. + out_fp8: emit float8_e4m3fn outputs (KV-cache dtype) when True, + bf16 otherwise. + + Returns: + (q, k, v): contiguous [N, q_size], [N, kv_size], [N, kv_size] in the + requested output dtype; q/k are roped, v is norm-only. + """ + assert qkv.dim() == 2 and qkv.stride(-1) == 1 + assert qkv.dtype == torch.bfloat16 + assert cos_sin.is_contiguous() and cos_sin.dtype == torch.float32 + half = head_dim // 2 + assert cos_sin.dim() == 3 and cos_sin.shape[1] == 2 and cos_sin.shape[2] == half + q_size = num_q_heads * head_dim + kv_size = num_kv_heads * head_dim + assert qkv.shape[-1] == q_size + 2 * kv_size + + n = qkv.shape[0] + out_dtype = torch.float8_e4m3fn if out_fp8 else torch.bfloat16 + q_out = torch.empty((n, q_size), dtype=out_dtype, device=qkv.device) + k_out = torch.empty((n, kv_size), dtype=out_dtype, device=qkv.device) + v_out = torch.empty((n, kv_size), dtype=out_dtype, device=qkv.device) + if n == 0: + return q_out, k_out, v_out + + # ~2k-element tiles keep every thread on wide vectorized accesses without + # spilling: 16 tokens/program for hd=256, 8 for hd=512 (B200-tuned, see + # module docstring). head_dim is a power of two on every gated shape, so + # block_n is too (required by tl.arange). + block_n = max(1, min(32, 4096 // head_dim)) + grid = (num_q_heads + 2 * num_kv_heads, triton.cdiv(n, block_n)) + _gemma4_qkv_norm_rope_quant_kernel[grid]( + qkv, + q_out, + k_out, + v_out, + position_ids.view(-1), + cos_sin, + q_weight, + k_weight, + qkv.stride(0), + num_q_heads, + num_kv_heads, + eps, + n, + HD=head_dim, + HALF=half, + OUT_FP8=out_fp8, + BLOCK_N=block_n, + num_warps=8, + ) + return q_out, k_out, v_out diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index f43b08b9fb72..14122d35683a 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -123,6 +123,10 @@ l0_b200: - unittest/_torch/modules/test_mla_helix.py - unittest/_torch/modules/test_fused_add_rms_norm_quant.py - unittest/_torch/modules/test_fused_activation_quant.py + - unittest/_torch/modules/fused_ops/test_rmsnorm_fp4_quant.py + - unittest/_torch/modules/fused_ops/test_gelu_tanh_mul_fp4_quant.py + - unittest/_torch/modules/fused_ops/test_rmsnorm_residual_add.py + - unittest/_torch/modules/test_gemma4_fused_qkv_prep.py - unittest/_torch/modules/test_awq_quantization.py - unittest/_torch/modules/test_triton_linear.py - unittest/_torch/modules/test_group_rmn_norm.py diff --git a/tests/unittest/_torch/modules/fused_ops/test_gelu_tanh_mul_fp4_quant.py b/tests/unittest/_torch/modules/fused_ops/test_gelu_tanh_mul_fp4_quant.py new file mode 100644 index 000000000000..65338f5e05a1 --- /dev/null +++ b/tests/unittest/_torch/modules/fused_ops/test_gelu_tanh_mul_fp4_quant.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Parity test: fused gelu_tanh+mul+NVFP4-quantize (modules/fused_ops/ +gelu_tanh_mul_fp4_quant) vs the unfused reference chain +(trtllm::flashinfer_gelu_tanh_and_mul -> trtllm::fp4_quantize). + +Both the packed E2M1 payload and the swizzled E4M3 scale factors are compared +at the byte level; the scale comparison is restricted to the valid (row, kvec) +region because the unfused op leaves its 128-row padding uninitialized (the +fused op zero-fills it). Byte-exact on SM100 (measured mismatch 0.0); the +tolerance below only guards hypothetical cross-arch fp contraction drift. +""" + +import pytest +import torch +from utils.util import skip_pre_blackwell + +import tensorrt_llm # noqa: F401 (registers trtllm torch ops) +import tensorrt_llm._torch.custom_ops.flashinfer_custom_ops # noqa: F401 +from tensorrt_llm._torch.modules.fused_ops.gelu_tanh_mul_fp4_quant import ( + gelu_tanh_mul_fp4_quant, + sf_swizzled_offsets, +) + +pytestmark = skip_pre_blackwell + + +def _reference(x, gs): + h = torch.ops.trtllm.flashinfer_gelu_tanh_and_mul(x) + return torch.ops.trtllm.fp4_quantize(h, gs, 16, False) + + +def _check(x, gs): + fq, fsf = gelu_tanh_mul_fp4_quant(x, gs) + rq, rsf = _reference(x, gs) + assert fq.shape == rq.shape and fsf.numel() == rsf.numel() + mm_fp4 = (fq != rq).float().mean().item() + valid = sf_swizzled_offsets(x.shape[0], x.shape[1] // 2 // 16, x.device) + mm_sf = (fsf[valid] != rsf[valid]).float().mean().item() + assert mm_fp4 < 1e-4, f"fp4 payload mismatch fraction {mm_fp4}" + assert mm_sf < 1e-4, f"scale-factor mismatch fraction {mm_sf}" + + +# 21504 is the Gemma4-31B intermediate size; 228 the pinned decode batch; +# 6455 the profiled serving prefill size; 333 exercises masked tail rows. +@pytest.mark.parametrize("n_tokens", [1, 7, 228, 333, 6455]) +@pytest.mark.parametrize("intermediate", [21504, 1024]) +def test_gelu_tanh_mul_fp4_quant_parity(n_tokens, intermediate): + torch.manual_seed(1234) + x = ( + torch.randn((n_tokens, 2 * intermediate), dtype=torch.bfloat16, device="cuda") + * torch.rand((n_tokens, 1), dtype=torch.bfloat16, device="cuda") + * 3 + ) + h = torch.ops.trtllm.flashinfer_gelu_tanh_and_mul(x) + gs = (448.0 * 6.0 / h.abs().max().float()).reshape(1) + _check(x, gs) + + +@pytest.mark.parametrize("gs_value", [1e6, 1.0, 1e-6]) +def test_gelu_tanh_mul_fp4_quant_extreme_scales(gs_value): + """Saturating / degenerate global scales must match the unfused op.""" + torch.manual_seed(7) + x = torch.randn((333, 2 * 21504), dtype=torch.bfloat16, device="cuda") + gs = torch.tensor([gs_value], dtype=torch.float32, device="cuda") + _check(x, gs) + + +def test_gelu_tanh_mul_fp4_quant_strided_rows(): + """Row-strided input (a view of a wider buffer).""" + torch.manual_seed(3) + n, i = 65, 21504 + buf = torch.randn((n, 2 * i + 512), dtype=torch.bfloat16, device="cuda") + x = buf[:, : 2 * i] + h = torch.ops.trtllm.flashinfer_gelu_tanh_and_mul(x.contiguous()) + gs = (448.0 * 6.0 / h.abs().max().float()).reshape(1) + fq, fsf = gelu_tanh_mul_fp4_quant(x, gs) + rq, rsf = torch.ops.trtllm.fp4_quantize(h, gs, 16, False) + valid = sf_swizzled_offsets(n, i // 16, x.device) + assert (fq != rq).float().mean().item() < 1e-4 + assert (fsf[valid] != rsf[valid]).float().mean().item() < 1e-4 + + +if __name__ == "__main__": + test_gelu_tanh_mul_fp4_quant_parity(6455, 21504) + test_gelu_tanh_mul_fp4_quant_parity(228, 21504) + test_gelu_tanh_mul_fp4_quant_extreme_scales(1e6) + test_gelu_tanh_mul_fp4_quant_strided_rows() + print("ALL PARITY CHECKS PASSED") diff --git a/tests/unittest/_torch/modules/fused_ops/test_rmsnorm_fp4_quant.py b/tests/unittest/_torch/modules/fused_ops/test_rmsnorm_fp4_quant.py new file mode 100644 index 000000000000..fbe1fa4cac14 --- /dev/null +++ b/tests/unittest/_torch/modules/fused_ops/test_rmsnorm_fp4_quant.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Quality test: fused RMSNorm+NVFP4-quantize (modules/fused_ops/ +rmsnorm_fp4_quant, backed by flashinfer's CuTe-DSL ``rmsnorm_fp4quant``) vs +the unfused reference chain +(trtllm::flashinfer_rmsnorm -> trtllm::fp4_quantize). + +The fused kernel quantizes the fp32 norm result directly, without the +intermediate bf16 round the unfused chain performs when materializing the +normed tensor, so its outputs are near- but not byte-identical (~1.4% of +payload nibbles / ~2% of block scales differ by one code step at serving +shapes). The tests therefore compare dequantized values, not bytes: + +- the fused path's quantization error against the fp32 norm must match the + unfused chain's (both are roundings of the same quantity); +- dequantization reads the scale factors at the swizzled offsets of + get_sf_out_offset_128x4, so any layout or scale-convention bug shows up + as a gross error, not a tolerance miss; +- a small byte-mismatch bound documents the near-parity property itself. +""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers trtllm torch ops) +import tensorrt_llm._torch.custom_ops.flashinfer_custom_ops # noqa: F401 +from tensorrt_llm._torch.modules.fused_ops.gelu_tanh_mul_fp4_quant import sf_swizzled_offsets +from tensorrt_llm._torch.modules.fused_ops.rmsnorm_fp4_quant import ( + rmsnorm_fp4_quant, + rmsnorm_fp4_quant_available, +) + +pytestmark = pytest.mark.skipif( + not rmsnorm_fp4_quant_available(), + reason="flashinfer CuTe-DSL rmsnorm_fp4quant unavailable", +) + +_E2M1_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] +_E2M1_LUT = _E2M1_LUT + [-v for v in _E2M1_LUT] + + +def _reference(x, w, eps, gs): + n = torch.ops.trtllm.flashinfer_rmsnorm(x, w, eps) + q, sf = torch.ops.trtllm.fp4_quantize(n, gs, 16, False) + return q.view(torch.uint8), sf.view(torch.uint8) + + +def _norm_fp32(x, w, eps): + xf = x.float() + return xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * w.float() + + +def _block_scales(sf, m, h): + """Block scales [m, h//16] read at the swizzled offsets (as fp32).""" + valid = sf_swizzled_offsets(m, h // 16, sf.device) + return sf.view(-1)[valid].view(torch.float8_e4m3fn).float().view(m, h // 16) + + +def _dequant(q, sf, m, h, gs): + """Dequantize payload + swizzled scales to fp32 [m, h].""" + lut = torch.tensor(_E2M1_LUT, device=q.device) + vals = torch.empty((m, h), device=q.device, dtype=torch.float32) + vals[:, 0::2] = lut[(q & 0xF).long()] + vals[:, 1::2] = lut[(q >> 4).long()] + return vals * _block_scales(sf, m, h).repeat_interleave(16, dim=1) / gs + + +def _check(x, w, eps, gs): + fq, fsf = rmsnorm_fp4_quant(x, w, eps, gs) + rq, rsf = _reference(x.contiguous(), w, eps, gs) + assert fq.shape == rq.shape and fsf.numel() == rsf.numel() + assert fq.dtype == torch.uint8 and fsf.dtype == torch.uint8 + + m, h = x.shape + n32 = _norm_fp32(x.contiguous(), w, eps) + err_fused = (_dequant(fq, fsf, m, h, gs) - n32).abs().mean().item() + err_ref = (_dequant(rq, rsf, m, h, gs) - n32).abs().mean().item() + # Both outputs are NVFP4 roundings of the same fp32 norm; the fused + # path must not be a worse quantizer than the unfused chain (measured + # equal to ~0.1% at serving shapes; the slack covers rounding-path + # differences on small inputs). + assert err_fused <= err_ref * 1.05 + 1e-6, f"{err_fused=} vs {err_ref=}" + # Near-parity: the two paths differ only where a rounding boundary is + # crossed (~1.4% of bytes at serving shapes). Blocks whose e4m3 scale + # underflows to zero are excluded: their payload never contributes to + # the dequantized value and the implementations fill it differently. + live = (_block_scales(rsf, m, h) != 0) & (_block_scales(fsf, m, h) != 0) + live_bytes = live.repeat_interleave(8, dim=1) # 8 payload bytes / block + if live_bytes.any(): + mm = (fq != rq)[live_bytes].float().mean().item() + assert mm < 5e-2, f"payload byte mismatch fraction {mm}" + + +# 5376 is the Gemma4-31B hidden size; 228 the pinned decode batch; 7700 the +# profiled serving prefill size; 333 exercises padded tail rows. +@pytest.mark.parametrize("n_tokens", [1, 7, 228, 333, 7700]) +@pytest.mark.parametrize("hidden", [5376, 512]) +def test_rmsnorm_fp4_quant_quality(n_tokens, hidden): + torch.manual_seed(1234) + x = ( + torch.randn((n_tokens, hidden), dtype=torch.bfloat16, device="cuda") + * torch.rand((n_tokens, 1), dtype=torch.bfloat16, device="cuda") + * 3 + ) + w = torch.rand((hidden,), dtype=torch.bfloat16, device="cuda") + 0.5 + n = torch.ops.trtllm.flashinfer_rmsnorm(x, w, 1e-6) + gs = (448.0 * 6.0 / n.abs().max().float()).reshape(1) + _check(x, w, 1e-6, gs) + + +@pytest.mark.parametrize("gs_value", [1e6, 1.0, 1e-6]) +def test_rmsnorm_fp4_quant_extreme_scales(gs_value): + """Saturating / degenerate global scales must behave like the unfused op.""" + torch.manual_seed(7) + x = torch.randn((333, 5376), dtype=torch.bfloat16, device="cuda") + w = torch.rand((5376,), dtype=torch.bfloat16, device="cuda") + 0.5 + gs = torch.tensor([gs_value], dtype=torch.float32, device="cuda") + _check(x, w, 1e-6, gs) + + +def test_rmsnorm_fp4_quant_strided_rows(): + """Row-strided input (a view of a wider buffer).""" + torch.manual_seed(3) + n, h = 65, 5376 + buf = torch.randn((n, h + 512), dtype=torch.bfloat16, device="cuda") + x = buf[:, :h] + w = torch.rand((h,), dtype=torch.bfloat16, device="cuda") + 0.5 + normed = torch.ops.trtllm.flashinfer_rmsnorm(x.contiguous(), w, 1e-6) + gs = (448.0 * 6.0 / normed.abs().max().float()).reshape(1) + _check(x, w, 1e-6, gs) + + +def test_rmsnorm_fp4_quant_zero_row(): + """An all-zero row must dequantize to exactly zero (zero block scales).""" + torch.manual_seed(5) + x = torch.randn((9, 5376), dtype=torch.bfloat16, device="cuda") + x[4] = 0 + w = torch.rand((5376,), dtype=torch.bfloat16, device="cuda") + 0.5 + gs = torch.tensor([100.0], dtype=torch.float32, device="cuda") + fq, fsf = rmsnorm_fp4_quant(x, w, 1e-6, gs) + dq = _dequant(fq, fsf, 9, 5376, gs) + assert dq[4].abs().max().item() == 0.0 + _check(x, w, 1e-6, gs) + + +def test_rmsnorm_fp4_quant_cuda_graph(): + """The kernel must be CUDA-graph capturable and replay deterministically + (the serving decode path replays it inside captured graphs).""" + torch.manual_seed(11) + x = torch.randn((228, 5376), dtype=torch.bfloat16, device="cuda") + w = torch.rand((5376,), dtype=torch.bfloat16, device="cuda") + 0.5 + gs = torch.tensor([100.0], dtype=torch.float32, device="cuda") + + eager_q, eager_sf = rmsnorm_fp4_quant(x, w, 1e-6, gs) # also JIT warmup + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + graph_q, graph_sf = rmsnorm_fp4_quant(x, w, 1e-6, gs) + g.replay() + torch.cuda.synchronize() + assert torch.equal(graph_q, eager_q) + valid = sf_swizzled_offsets(228, 5376 // 16, x.device) + assert torch.equal(graph_sf.view(-1)[valid], eager_sf.view(-1)[valid]) + + +if __name__ == "__main__": + test_rmsnorm_fp4_quant_quality(7700, 5376) + test_rmsnorm_fp4_quant_quality(228, 5376) + test_rmsnorm_fp4_quant_extreme_scales(1e6) + test_rmsnorm_fp4_quant_strided_rows() + test_rmsnorm_fp4_quant_zero_row() + test_rmsnorm_fp4_quant_cuda_graph() + print("ALL QUALITY CHECKS PASSED") diff --git a/tests/unittest/_torch/modules/fused_ops/test_rmsnorm_residual_add.py b/tests/unittest/_torch/modules/fused_ops/test_rmsnorm_residual_add.py new file mode 100644 index 000000000000..5c08e1e00be5 --- /dev/null +++ b/tests/unittest/_torch/modules/fused_ops/test_rmsnorm_residual_add.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Parity tests: fused post-norm residual kernels (modules/fused_ops/ +rmsnorm_residual_add) vs the unfused chains. + +Covers the three variants served by the shared kernel body: +- ``rmsnorm_residual_add_scale`` (norm + residual add + fp32 scalar mul + + bf16 cast), +- ``rmsnorm_residual_add`` (no scalar stage), +- the optional second output (the RMSNorm of the bf16-rounded primary + output, e.g. the next decoder layer's input norm). + +Each reference reproduces the exact serving-path op sequence it replaces +(flashinfer_rmsnorm + aten ops). The only tolerated difference is the fp32 +reduction-order inside the sums of squares (one-step bf16 flips, measured +~5e-6 at serving shapes on SM100). +""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers trtllm torch ops) +import tensorrt_llm._torch.custom_ops.flashinfer_custom_ops # noqa: F401 +from tensorrt_llm._torch.modules.fused_ops.rmsnorm_residual_add import ( + rmsnorm_residual_add, + rmsnorm_residual_add_scale, +) + + +def _reference(x, r, w, sc, eps): + n = torch.ops.trtllm.flashinfer_rmsnorm(x, w, eps) + s = r + n + t = s * sc + return t.to(torch.bfloat16) + + +def _mismatch(a, b): + return (a.view(torch.uint16) != b.view(torch.uint16)).float().mean().item() + + +def _check(x, r, w, sc, eps): + fused = rmsnorm_residual_add_scale(x, r, w, sc, eps) + ref = _reference(x, r, w, sc, eps) + mm = _mismatch(fused, ref) + assert mm < 1e-4, f"bf16 mismatch fraction {mm}" + + +# 5376 is the Gemma4-31B hidden size; 228 the pinned decode batch; 7700 the +# round-3 profiled prefill token count; 333 exercises the row tail. +@pytest.mark.parametrize("n_tokens", [1, 7, 228, 333, 7700]) +@pytest.mark.parametrize("hidden", [5376, 512]) +def test_rmsnorm_residual_add_scale_parity(n_tokens, hidden): + torch.manual_seed(1234) + x = torch.randn((n_tokens, hidden), dtype=torch.bfloat16, device="cuda") + r = torch.randn((n_tokens, hidden), dtype=torch.bfloat16, device="cuda") + w = torch.rand((hidden,), dtype=torch.bfloat16, device="cuda") + 0.5 + sc = torch.tensor([1.0], dtype=torch.float32, device="cuda") + _check(x, r, w, sc, 1e-6) + + +@pytest.mark.parametrize("scalar", [0.987654, 2.5, 0.0]) +def test_rmsnorm_residual_add_scale_nontrivial_scalar(scalar): + """layer_scalar loads from the checkpoint - never assume 1.0.""" + torch.manual_seed(7) + x = torch.randn((333, 5376), dtype=torch.bfloat16, device="cuda") + r = torch.randn((333, 5376), dtype=torch.bfloat16, device="cuda") + w = torch.rand((5376,), dtype=torch.bfloat16, device="cuda") + 0.5 + sc = torch.tensor([scalar], dtype=torch.float32, device="cuda") + _check(x, r, w, sc, 1e-6) + + +def test_rmsnorm_residual_add_scale_non_pow2_hidden_and_strided(): + """Masked tail columns (non-power-of-2 H) + row-strided inputs.""" + torch.manual_seed(3) + n, h = 65, 384 + xbuf = torch.randn((n, h + 128), dtype=torch.bfloat16, device="cuda") + rbuf = torch.randn((n, h + 64), dtype=torch.bfloat16, device="cuda") + x, r = xbuf[:, :h], rbuf[:, :h] + w = torch.rand((h,), dtype=torch.bfloat16, device="cuda") + 0.5 + sc = torch.tensor([1.25], dtype=torch.float32, device="cuda") + fused = rmsnorm_residual_add_scale(x, r, w, sc, 1e-6) + ref = _reference(x.contiguous(), r.contiguous(), w, sc, 1e-6) + mm = _mismatch(fused, ref) + assert mm < 1e-4, f"bf16 mismatch fraction {mm}" + + +@pytest.mark.parametrize("n_tokens", [1, 7, 228, 333, 7700]) +@pytest.mark.parametrize("hidden", [5376, 512]) +def test_rmsnorm_residual_add_parity(n_tokens, hidden): + """Post-attention variant: rmsnorm(x) then the aten bf16 residual add.""" + torch.manual_seed(42) + x = torch.randn((n_tokens, hidden), dtype=torch.bfloat16, device="cuda") + r = torch.randn((n_tokens, hidden), dtype=torch.bfloat16, device="cuda") + w = torch.rand((hidden,), dtype=torch.bfloat16, device="cuda") + 0.5 + fused = rmsnorm_residual_add(x, r, w, 1e-6) + ref = r + torch.ops.trtllm.flashinfer_rmsnorm(x, w, 1e-6) + mm = _mismatch(fused, ref) + assert mm < 1e-4, f"bf16 mismatch fraction {mm}" + + +@pytest.mark.parametrize("n_tokens", [1, 228, 333, 7700]) +@pytest.mark.parametrize("hidden", [5376, 512]) +def test_rmsnorm_residual_add_scale_dual_norm_parity(n_tokens, hidden): + """Tail with the secondary next-layer input-norm output. + + The primary output must match the plain tail; the secondary output must + match a standalone flashinfer_rmsnorm applied to the (bf16-rounded) + primary output - that is exactly the tensor the next layer's standalone + input norm would read. + """ + torch.manual_seed(11) + x = torch.randn((n_tokens, hidden), dtype=torch.bfloat16, device="cuda") + r = torch.randn((n_tokens, hidden), dtype=torch.bfloat16, device="cuda") + w = torch.rand((hidden,), dtype=torch.bfloat16, device="cuda") + 0.5 + w2 = torch.rand((hidden,), dtype=torch.bfloat16, device="cuda") + 0.5 + sc = torch.tensor([0.987654], dtype=torch.float32, device="cuda") + out, n2 = rmsnorm_residual_add_scale(x, r, w, sc, 1e-6, next_norm_weight=w2, next_norm_eps=1e-6) + ref = _reference(x, r, w, sc, 1e-6) + mm = _mismatch(out, ref) + assert mm < 1e-4, f"primary-output bf16 mismatch fraction {mm}" + ref_n2 = torch.ops.trtllm.flashinfer_rmsnorm(out, w2, 1e-6) + mm2 = _mismatch(n2, ref_n2) + assert mm2 < 1e-4, f"secondary-norm bf16 mismatch fraction {mm2}" + + +if __name__ == "__main__": + test_rmsnorm_residual_add_scale_parity(7700, 5376) + test_rmsnorm_residual_add_scale_parity(228, 5376) + test_rmsnorm_residual_add_scale_nontrivial_scalar(0.987654) + test_rmsnorm_residual_add_scale_non_pow2_hidden_and_strided() + test_rmsnorm_residual_add_parity(7700, 5376) + test_rmsnorm_residual_add_parity(228, 5376) + test_rmsnorm_residual_add_scale_dual_norm_parity(7700, 5376) + test_rmsnorm_residual_add_scale_dual_norm_parity(228, 5376) + print("ALL PARITY CHECKS PASSED") diff --git a/tests/unittest/_torch/modules/test_gemma4_fused_qkv_prep.py b/tests/unittest/_torch/modules/test_gemma4_fused_qkv_prep.py new file mode 100644 index 000000000000..9ff768736329 --- /dev/null +++ b/tests/unittest/_torch/modules/test_gemma4_fused_qkv_prep.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Parity test: fused Gemma4 QKV prep vs the unfused reference chain. + +The reference reproduces the exact serving-path math it replaces: +split_qkv strided views -> RMSNorm (fp32 accum, bf16 round) -> neox RoPE from +the fp32 cos/sin table (bf16 round) -> .to(float8_e4m3fn). +""" + +import pytest +import torch + +from tensorrt_llm._torch.modules.gemma4.fused_qkv import gemma4_fused_qkv_norm_rope_quant + + +def _make_cos_sin(max_pos: int, head_dim: int, rotary_frac: float, theta: float) -> torch.Tensor: + """Build a [max_pos, 2, head_dim//2] fp32 table like RotaryEmbedding. + + rotary_frac < 1 emulates Gemma4 full-attention layers: only the first + rotary_frac*half frequency pairs are non-trivial, the rest are + zero-frequency (cos=1, sin=0). + """ + half = head_dim // 2 + n_rot = int(half * rotary_frac) + inv_freq = 1.0 / ( + theta + ** ( + torch.arange(0, 2 * n_rot, 2, dtype=torch.float32) + / (2 * n_rot if rotary_frac == 1.0 else head_dim) + ) + ) + inv_freq = torch.cat([inv_freq, torch.zeros(half - n_rot, dtype=torch.float32)]) + pos = torch.arange(max_pos, dtype=torch.float32) + sinusoid = torch.einsum("i,j->ij", pos, inv_freq) + return torch.stack([sinusoid.cos(), sinusoid.sin()], dim=1).cuda().contiguous() + + +def _ref_chain(qkv, position_ids, cos_sin, q_w, k_w, eps, nq, nk, hd, out_fp8): + q_size, kv_size = nq * hd, nk * hd + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) + + def norm(x, w): + xf = x.reshape(-1, hd).float() + r = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) + y = xf * r + if w is not None: + y = y * w.float() + return y.to(torch.bfloat16) + + qn = norm(q, q_w).view(-1, nq, hd) + kn = norm(k, k_w).view(-1, nk, hd) + vn = norm(v, None).view(-1, kv_size) + + half = hd // 2 + cs = cos_sin[position_ids.view(-1).long()] # [N, 2, half] fp32 + cos = cs[:, 0].unsqueeze(1) + sin = cs[:, 1].unsqueeze(1) + + def rope(x): + xf = x.float() + x1, x2 = xf[..., :half], xf[..., half:] + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + return torch.cat([o1, o2], dim=-1).to(torch.bfloat16).reshape(x.shape[0], -1) + + qr, kr = rope(qn), rope(kn) + if out_fp8: + f8 = torch.float8_e4m3fn + return qr.to(f8), kr.to(f8), vn.to(f8) + return qr, kr, vn + + +def _assert_close_fp8(fused: torch.Tensor, ref: torch.Tensor, name: str): + f, r = fused.float(), ref.float() + # Bitwise-equal for the overwhelming majority; reduction-order ULP + # differences may flip an fp8 code by at most one step. + mismatch = (fused.view(torch.uint8) != ref.view(torch.uint8)).float().mean() + assert mismatch < 1e-3, f"{name}: fp8 mismatch fraction {mismatch}" + assert torch.allclose(f, r, atol=0.0625, rtol=0.13), f"{name}: max diff {(f - r).abs().max()}" + + +@pytest.mark.parametrize( + "nq,nk,hd,rotary_frac,theta", + [ + (32, 16, 256, 1.0, 10000.0), # Gemma4-31B sliding layers + (32, 4, 512, 0.25, 1000000.0), # Gemma4-31B full layers (padded freqs) + (8, 2, 128, 1.0, 10000.0), # generic small + ], +) +# 256 is an exact multiple of every BLOCK_N tile; 333/6455 exercise the +# masked tail rows; 6455 is the profiled serving prefill size. +@pytest.mark.parametrize("n_tokens", [1, 7, 256, 333, 6455]) +def test_fused_qkv_prep_parity_fp8(nq, nk, hd, rotary_frac, theta, n_tokens): + torch.manual_seed(1234) + max_pos = 4096 + qkv = torch.randn((n_tokens, (nq + 2 * nk) * hd), dtype=torch.bfloat16, device="cuda") + position_ids = torch.randint(0, max_pos, (n_tokens,), dtype=torch.int32, device="cuda") + cos_sin = _make_cos_sin(max_pos, hd, rotary_frac, theta) + q_w = torch.rand((hd,), dtype=torch.bfloat16, device="cuda") + 0.5 + k_w = torch.rand((hd,), dtype=torch.bfloat16, device="cuda") + 0.5 + eps = 1e-6 + + fq, fk, fv = gemma4_fused_qkv_norm_rope_quant( + qkv, position_ids, cos_sin, q_w, k_w, eps, nq, nk, hd, out_fp8=True + ) + rq, rk, rv = _ref_chain(qkv, position_ids, cos_sin, q_w, k_w, eps, nq, nk, hd, out_fp8=True) + + _assert_close_fp8(fq, rq, "q") + _assert_close_fp8(fk, rk, "k") + _assert_close_fp8(fv, rv, "v") + + +def test_fused_qkv_prep_parity_bf16_and_strided(): + """BF16 output mode + a row-strided qkv input (view of a wider buffer).""" + torch.manual_seed(7) + nq, nk, hd = 32, 16, 256 + n_tokens, max_pos = 65, 2048 + width = (nq + 2 * nk) * hd + buf = torch.randn((n_tokens, width + 512), dtype=torch.bfloat16, device="cuda") + qkv = buf[:, :width] # non-trivial row stride + position_ids = torch.randint(0, max_pos, (n_tokens,), dtype=torch.int32, device="cuda") + cos_sin = _make_cos_sin(max_pos, hd, 1.0, 10000.0) + q_w = torch.rand((hd,), dtype=torch.bfloat16, device="cuda") + 0.5 + k_w = torch.rand((hd,), dtype=torch.bfloat16, device="cuda") + 0.5 + eps = 1e-6 + + fq, fk, fv = gemma4_fused_qkv_norm_rope_quant( + qkv, position_ids, cos_sin, q_w, k_w, eps, nq, nk, hd, out_fp8=False + ) + rq, rk, rv = _ref_chain( + qkv.contiguous(), position_ids, cos_sin, q_w, k_w, eps, nq, nk, hd, out_fp8=False + ) + + for f, r, name in ((fq, rq, "q"), (fk, rk, "k"), (fv, rv, "v")): + assert torch.allclose(f.float(), r.float(), atol=0.02, rtol=0.02), ( + f"{name}: max diff {(f.float() - r.float()).abs().max()}" + ) + + +if __name__ == "__main__": + test_fused_qkv_prep_parity_fp8(32, 16, 256, 1.0, 10000.0, 333) + test_fused_qkv_prep_parity_fp8(32, 16, 256, 1.0, 10000.0, 6455) + test_fused_qkv_prep_parity_fp8(32, 4, 512, 0.25, 1000000.0, 6455) + test_fused_qkv_prep_parity_bf16_and_strided() + print("ALL PARITY CHECKS PASSED")