From fee025432c29f38658e12f32fcf3507105ac41de Mon Sep 17 00:00:00 2001 From: wejoncy <9417365+wejoncy@users.noreply.github.com> Date: Wed, 5 Aug 2026 06:15:57 +0000 Subject: [PATCH] Add nanojet as an auto_deploy attention and fusion backend (prefill only) Five custom ops and five transforms route a prequantized FP8 model onto nanojet's kernels. All ship disabled, and the two existing files touched are additive only, so no other backend changes behaviour. --- .../_torch/auto_deploy/config/default.yaml | 19 + .../custom_ops/attention/nanojet_attention.py | 187 +++++++++ .../nanojet_fused_qkv_gemm_norm_rope.py | 105 +++++ .../linear/nanojet_gemm_fp8_add_inplace.py | 58 +++ .../linear/nanojet_swiglu_gemm_fp8.py | 72 ++++ .../normalization/nanojet_rmsnorm_fp8.py | 53 +++ tensorrt_llm/_torch/auto_deploy/llm_args.py | 31 ++ .../library/fuse_nanojet_attn_quant_fp8.py | 87 +++++ .../fuse_nanojet_fused_qkv_gemm_norm_rope.py | 368 ++++++++++++++++++ .../library/fuse_nanojet_gemm_fp8_add.py | 160 ++++++++ .../library/fuse_nanojet_rmsnorm_fp8.py | 160 ++++++++ .../library/fuse_nanojet_swiglu_gemm_fp8.py | 150 +++++++ .../_torch/auto_deploy/utils/nanojet_graph.py | 153 ++++++++ tensorrt_llm/_torch/nanojet_utils.py | 94 +++++ .../custom_ops/test_nanojet_rejections.py | 133 +++++++ ...nanojet_fused_qkv_gemm_norm_rope_guards.py | 59 +++ ...t_fused_qkv_gemm_norm_rope_shape_guards.py | 119 ++++++ .../library/test_nanojet_residual_fold.py | 138 +++++++ 18 files changed, 2146 insertions(+) create mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/attention/nanojet_attention.py create mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_fused_qkv_gemm_norm_rope.py create mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_gemm_fp8_add_inplace.py create mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_swiglu_gemm_fp8.py create mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/nanojet_rmsnorm_fp8.py create mode 100644 tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_attn_quant_fp8.py create mode 100644 tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_fused_qkv_gemm_norm_rope.py create mode 100644 tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_gemm_fp8_add.py create mode 100644 tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_rmsnorm_fp8.py create mode 100644 tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_swiglu_gemm_fp8.py create mode 100644 tensorrt_llm/_torch/auto_deploy/utils/nanojet_graph.py create mode 100644 tensorrt_llm/_torch/nanojet_utils.py create mode 100644 tests/unittest/auto_deploy/singlegpu/custom_ops/test_nanojet_rejections.py create mode 100644 tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_fused_qkv_gemm_norm_rope_guards.py create mode 100644 tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_fused_qkv_gemm_norm_rope_shape_guards.py create mode 100644 tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_residual_fold.py diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index d6538bf7ba29..1657a92992d1 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -184,10 +184,29 @@ transforms: fuse_fp8_gemms: stage: post_load_fusion enabled: false # TODO: https://github.com/NVIDIA/TensorRT-LLM/issues/4674 this is causing OOMs + fuse_nanojet_fused_qkv_gemm_norm_rope: + stage: post_load_fusion + enabled: false fuse_fp8_linear: stage: post_load_fusion backend: trtllm requires_shape_prop: true + fuse_nanojet_rmsnorm_fp8: + stage: post_load_fusion + enabled: false + requires_shape_prop: true + fuse_nanojet_swiglu_gemm_fp8: + stage: post_load_fusion + enabled: false + requires_shape_prop: true + fuse_nanojet_attn_quant_fp8: + stage: post_load_fusion + enabled: false + requires_shape_prop: true + fuse_nanojet_gemm_fp8_add: + stage: post_load_fusion + enabled: false + requires_shape_prop: true fuse_trtllm_attn_quant_fp8: stage: post_load_fusion enabled: false diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/nanojet_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/nanojet_attention.py new file mode 100644 index 000000000000..85c409ecfc30 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/nanojet_attention.py @@ -0,0 +1,187 @@ +"""nanojet FlashAttention-3 as an auto-deploy attention backend. Prefill only. + +Quantizes to e4m3 in the kernel epilogue when handed an ``output_scale``, so ``o_proj`` can +read the result directly. +""" + +import math +from typing import List, Optional + +import torch +from torch.fx import Node + +from ....nanojet_utils import is_nanojet_available +from ...utils.nanojet_graph import per_tensor_scale +from ..attention_interface import AttentionRegistry, BatchInfo, Constant, MHACallable +from .torch_backend_attention import TorchBackendAttention + +_REGISTERED = False + +# Where ``fuse_nanojet_attn_quant_fp8`` records the scale the reader quantizes by. +NANOJET_ATTENTION_INPUT_SCALE = "nanojet_attention_input_scale" + + +def register() -> bool: + """Define the ops, importing nanojet only now. Idempotent; returns availability.""" + global _REGISTERED + if _REGISTERED: + return True + if not is_nanojet_available(): + return False + _REGISTERED = True + + from nanojet_kernels import ops + + @torch.library.custom_op("auto_deploy::nanojet_attention", mutates_args=()) + def nanojet_mha_with_cache( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + # STANDARD METADATA + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + cu_seqlen: torch.Tensor, + # EXTRA METADATA + # CONSTANTS + scale: Optional[float], + sinks: Optional[torch.Tensor] = None, + sliding_window_size: Optional[int] = None, + logit_cap: Optional[float] = None, + read_cache_only: bool = False, + output_scale: Optional[float] = None, + custom_attn_mask: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Varlen FlashAttention over the Q/K/V in hand. Prefill only.""" + batch_info = BatchInfo(batch_info_host) + num_seq, num_extend, num_decode = batch_info.get_num_sequences() + num_total_tokens, _, _ = batch_info.get_num_tokens() + if num_decode: + raise RuntimeError( + f"nanojet attention backend is prefill-only, but this batch has {num_decode} " + "decode sequence(s). It keeps no KV cache. Use attn_backend='trtllm' or " + "'flashinfer' for generation." + ) + if read_cache_only: + raise RuntimeError( + "nanojet attention backend was asked to read from a KV cache (shared-KV " + "layer), which it does not keep. Use attn_backend='trtllm' or 'flashinfer'." + ) + if num_extend: + raise RuntimeError( + f"nanojet attention backend received {num_extend} continuation chunk(s) " + "(extend request: cached context plus new tokens). The keys and values for the " + "earlier part of the sequence live in a KV cache it does not keep, so attending " + "over only this chunk would be silently wrong. Disable chunked prefill, or use " + "attn_backend='trtllm' or 'flashinfer'." + ) + + batch, seq = q.shape[:2] + qk_head_dim = q.shape[-1] + v_head_dim = v.shape[-1] + num_heads = q.shape[2] if q.ndim == 4 else q.shape[2] // qk_head_dim + num_kv_heads = k.shape[2] if k.ndim == 4 else k.shape[2] // qk_head_dim + output_shape = ( + (batch, seq, num_heads * v_head_dim) + if q.ndim == 3 + else (batch, seq, num_heads, v_head_dim) + ) + bs_view = (batch, seq) if seq == 1 else (batch * seq,) + + q = q.reshape(*bs_view, num_heads, qk_head_dim) + k = k.reshape(*bs_view, num_kv_heads, qk_head_dim) + v = v.reshape(*bs_view, num_kv_heads, v_head_dim) + if scale is None: + scale = 1.0 / math.sqrt(qk_head_dim) + + result_dtype = torch.float8_e4m3fn if output_scale is not None else q.dtype + result = torch.empty(*bs_view, num_heads, v_head_dim, dtype=result_dtype, device=q.device) + cumulative = cu_seqlen[: num_seq + 1].to(torch.int32) + max_seqlen = int(seq) + ops.flash_attention( + q, + k, + v, + out_tensor=result, + cu_seqlens_q=cumulative, + cu_seqlens_k=cumulative, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + softmax_scale=scale, + causal=True, + window_size_left=-1 if sliding_window_size is None else sliding_window_size, + output_scale=1.0 if output_scale is None else output_scale, + ) + if out is not None: + out_flat = out.view(*bs_view, num_heads, v_head_dim) + out_flat[:num_total_tokens].copy_(result[:num_total_tokens]) + if num_total_tokens < out_flat.shape[0]: + out_flat[num_total_tokens:].zero_() + return out.new_empty(0) + if num_total_tokens < result.shape[0]: + result[num_total_tokens:].zero_() + return result.view(*output_shape) + + @nanojet_mha_with_cache.register_fake + def _( + q, + k, + v, + batch_info_host, + seq_len, + input_pos, + slot_idx, + cu_seqlen, + scale, + sinks=None, + sliding_window_size=None, + logit_cap=None, + read_cache_only=False, + output_scale=None, + custom_attn_mask=None, + out=None, + ): + v_head_dim = v.shape[-1] + dtype = torch.float8_e4m3fn if output_scale is not None else q.dtype + if out is not None: + return out.new_empty(0) + if q.ndim == 3: + return torch.empty(*q.shape[:3], dtype=dtype, device=q.device) + return torch.empty(*q.shape[:3], v_head_dim, dtype=dtype, device=q.device) + + return True + + +@AttentionRegistry.register("nanojet") +class NanojetAttention(TorchBackendAttention): + """Varlen FlashAttention-3 over an unpaged cache. + + Everything except the kernel itself — layout, cache shape, metadata, constants — + matches the torch backend, so those are inherited rather than restated. + """ + + @classmethod + def get_cached_attention_op(cls) -> MHACallable: + if not register(): + raise RuntimeError( + "attention backend 'nanojet' was selected but nanojet is not importable" + ) + return torch.ops.auto_deploy.nanojet_attention.default + + @classmethod + def get_constants(cls, source_attn_node: Node) -> List[Constant]: + """Quantize in the epilogue by the scale ``fuse_nanojet_attn_quant_fp8`` recorded.""" + scale = source_attn_node.meta.get(NANOJET_ATTENTION_INPUT_SCALE) + output_scale = ( + None + if scale is None + else 1.0 / per_tensor_scale(source_attn_node.graph.owning_module, scale) + ) + return list(super().get_constants(source_attn_node)) + [output_scale] + + @classmethod + def get_cache_initializers(cls, source_attn_node, cache_config): + """No caches. Prefill has every key and value already.""" + return {} diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_fused_qkv_gemm_norm_rope.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_fused_qkv_gemm_norm_rope.py new file mode 100644 index 000000000000..745f9d60cb6c --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_fused_qkv_gemm_norm_rope.py @@ -0,0 +1,105 @@ +"""nanojet fused QKV projection with Q/K norm and RoPE, in one kernel.""" + +import torch + +from ....nanojet_utils import is_nanojet_available + +_REGISTERED = False + + +def register() -> bool: + """Define the ops, importing nanojet only now. Idempotent; returns availability.""" + global _REGISTERED + if _REGISTERED: + return True + if not is_nanojet_available(): + return False + _REGISTERED = True + + from nanojet_kernels import ops + + @torch.library.custom_op("auto_deploy::nanojet_fused_qkv_gemm_norm_rope", mutates_args=()) + def nanojet_fused_qkv_gemm_norm_rope( + hidden_states: torch.Tensor, + qkv_weight: torch.Tensor, + query_norm_weight: torch.Tensor, + key_norm_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + position_ids: torch.Tensor, + input_scale: torch.Tensor, + eps: float, + query_scale: float, + key_scale: float, + value_scale: float, + query_size: int, + key_value_size: int, + ) -> torch.Tensor: + """Project, normalize and rotate in one launch. + + ``qkv_weight`` is the three projections stacked ``[q + 2kv, hidden]``; Q, K and V + come back stacked on the last dim, for the graph to slice apart as views. + """ + batch, seq, hidden_size = hidden_states.shape + num_tokens = batch * seq + flattened = hidden_states.reshape(num_tokens, hidden_size) + positions = position_ids.reshape(-1) + if positions.numel() != num_tokens: + raise RuntimeError( + f"nanojet_fused_qkv_gemm_norm_rope got {positions.numel()} position ids for " + f"{num_tokens} tokens ({batch}x{seq}). Every token needs its own position." + ) + if flattened.dtype == torch.float8_e4m3fn: + quantized = flattened + else: + quantized, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + flattened, input_scale + ) + + packed = torch.empty( + num_tokens, + query_size + 2 * key_value_size, + dtype=torch.bfloat16, + device=hidden_states.device, + ) + ops.fused_qkv_gemm_norm_rope( + packed, + quantized, + qkv_weight, + query_scale, + key_scale, + value_scale, + query_size, + key_value_size, + query_norm_weight, + key_norm_weight, + eps, + cos_sin_cache, + positions, + ) + return packed.view(batch, seq, -1) + + @nanojet_fused_qkv_gemm_norm_rope.register_fake + def _( + hidden_states: torch.Tensor, + qkv_weight: torch.Tensor, + query_norm_weight: torch.Tensor, + key_norm_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + position_ids: torch.Tensor, + input_scale: torch.Tensor, + eps: float, + query_scale: float, + key_scale: float, + value_scale: float, + query_size: int, + key_value_size: int, + ) -> torch.Tensor: + return torch.empty( + hidden_states.shape[0], + hidden_states.shape[1], + query_size + 2 * key_value_size, + dtype=torch.bfloat16, + device=hidden_states.device, + ) + + return True diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_gemm_fp8_add_inplace.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_gemm_fp8_add_inplace.py new file mode 100644 index 000000000000..d43ac5f00cf2 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_gemm_fp8_add_inplace.py @@ -0,0 +1,58 @@ +"""nanojet FP8 GEMM for projections whose activation is already e4m3.""" + +import torch + +from ....nanojet_utils import is_nanojet_available + +_REGISTERED = False + + +def register() -> bool: + """Define the ops, importing nanojet only now. Idempotent; returns availability.""" + global _REGISTERED + if _REGISTERED: + return True + if not is_nanojet_available(): + return False + _REGISTERED = True + + from nanojet_kernels import ops + + def _check_activation(name: str, activation: torch.Tensor, weight: torch.Tensor) -> None: + """Reject what the kernel cannot take, while there is still a Python frame to see.""" + if activation.dtype != torch.float8_e4m3fn: + raise RuntimeError( + f"{name} expects an e4m3 activation, got {activation.dtype}. " + "The producer was expected to quantize in its epilogue." + ) + if activation.shape[-1] != weight.shape[1]: + raise RuntimeError( + f"{name} shape mismatch: activation {tuple(activation.shape)} against weight " + f"{tuple(weight.shape)} — the activation was probably taken before a reshape " + "that flattens the head dimension." + ) + + @torch.library.custom_op("auto_deploy::nanojet_gemm_fp8_add_inplace", mutates_args={"residual"}) + def nanojet_gemm_fp8_add_inplace( + hidden_states: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + input_scale: float, + weight_scale: float, + ) -> None: + """``residual += (x @ weight^T) * input_scale * weight_scale``, in the epilogue.""" + activation = hidden_states.reshape(-1, hidden_states.shape[-1]) + accumulator = residual.view(-1, residual.shape[-1]) + _check_activation("nanojet_gemm_fp8_add", activation, weight) + if accumulator.shape[0] != activation.shape[0] or accumulator.shape[-1] != weight.shape[0]: + raise RuntimeError( + f"nanojet_gemm_fp8_add accumulator {tuple(accumulator.shape)} does not match " + f"the projection output [{activation.shape[0]}, {weight.shape[0]}]." + ) + ops.gemm_fp8_add(accumulator, activation, weight, input_scale, weight_scale) + + @nanojet_gemm_fp8_add_inplace.register_fake + def _(hidden_states, weight, residual, input_scale, weight_scale): + return + + return True diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_swiglu_gemm_fp8.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_swiglu_gemm_fp8.py new file mode 100644 index 000000000000..09a49f165cd3 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/nanojet_swiglu_gemm_fp8.py @@ -0,0 +1,72 @@ +"""nanojet gated SwiGLU GEMM. + +Both projections, the activation and the multiply in one CUTLASS gated GEMM, emitting e4m3 +so ``down_proj`` needs no quantize either — four kernels collapse to one. This is the shape +nanojet runs natively; TensorRT LLM has no equivalent, because its own SwiGLU fusion needs +gate and up to come from one already-fused GEMM and ``fuse_gemms`` is disabled upstream. +""" + +import torch + +from ....nanojet_utils import is_nanojet_available + +_REGISTERED = False + + +def register() -> bool: + """Define the ops, importing nanojet only now. Idempotent; returns availability.""" + global _REGISTERED + if _REGISTERED: + return True + if not is_nanojet_available(): + return False + _REGISTERED = True + + from nanojet_kernels import ops + + @torch.library.custom_op("auto_deploy::nanojet_swiglu_gemm_fp8", mutates_args=()) + def nanojet_swiglu_gemm_fp8( + hidden_states: torch.Tensor, + input_scale_tensor: torch.Tensor, + gate_up_weight: torch.Tensor, + input_scale: float, + up_weight_scale: float, + gate_weight_scale: float, + output_scale: float, + ) -> torch.Tensor: + """``e4m3(silu(x @ gate^T) * (x @ up^T))`` in one launch. + + ``gate_up_weight`` is ``[up; gate]`` stacked — up first, which is the order nanojet's + kernel indexes. ``hidden_states`` is already e4m3, quantized by the RMSNorm epilogue + that produced it. All scales are host constants, folded at graph-build time. + """ + shape = hidden_states.shape + flattened = hidden_states.reshape(-1, shape[-1]) + if flattened.dtype != torch.float8_e4m3fn: + flattened, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + flattened, input_scale_tensor + ) + output = ops.swiglu( + flattened, gate_up_weight, input_scale, up_weight_scale, gate_weight_scale, output_scale + ) + return output.view(*shape[:-1], output.shape[-1]) + + @nanojet_swiglu_gemm_fp8.register_fake + def _nanojet_swiglu_gemm_fp8_fake( + hidden_states: torch.Tensor, + input_scale_tensor: torch.Tensor, + gate_up_weight: torch.Tensor, + input_scale: float, + up_weight_scale: float, + gate_weight_scale: float, + output_scale: float, + ) -> torch.Tensor: + intermediate = gate_up_weight.shape[0] // 2 + return torch.empty( + *hidden_states.shape[:-1], + intermediate, + dtype=torch.float8_e4m3fn, + device=hidden_states.device, + ) + + return True diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/nanojet_rmsnorm_fp8.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/nanojet_rmsnorm_fp8.py new file mode 100644 index 000000000000..cb37a191593e --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/nanojet_rmsnorm_fp8.py @@ -0,0 +1,53 @@ +"""nanojet RMSNorm with the quantize in its own epilogue. + +Every FP8 consumer downstream of a norm otherwise pays a separate pass over the activations +just to convert them. ``unified_rmsnorm`` already writes e4m3 directly when asked, which is +how nanojet drives this natively — no standalone quantize kernel in the layer at all. +""" + +import torch + +from ....nanojet_utils import is_nanojet_available + +_REGISTERED = False + + +def register() -> bool: + """Define the ops, importing nanojet only now. Idempotent; returns availability.""" + global _REGISTERED + if _REGISTERED: + return True + if not is_nanojet_available(): + return False + _REGISTERED = True + + from nanojet_kernels import ops + + @torch.library.custom_op("auto_deploy::nanojet_rmsnorm_fp8", mutates_args=()) + def nanojet_rmsnorm_fp8( + hidden_states: torch.Tensor, weight: torch.Tensor, eps: float, quantize_scale: float + ) -> torch.Tensor: + """``e4m3(rmsnorm(x))`` in one launch. + + ``quantize_scale`` is the reciprocal of the consumers' dequant scale — nanojet's + ``fp8_scale`` multiplies — folded on the host at graph-build time so nothing syncs + to the device per call. + """ + shape = hidden_states.shape + hidden = shape[-1] + output = ops.unified_rmsnorm( + hidden_states.reshape(-1, hidden), + weight, + eps=eps, + out_dtype=torch.float8_e4m3fn, + fp8_scale=quantize_scale, + ) + return output.view(shape) + + @nanojet_rmsnorm_fp8.register_fake + def _nanojet_rmsnorm_fp8_fake( + hidden_states: torch.Tensor, weight: torch.Tensor, eps: float, quantize_scale: float + ) -> torch.Tensor: + return torch.empty_like(hidden_states, dtype=torch.float8_e4m3fn) + + return True diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index f0c2af09cf6b..bbc01f97a927 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -416,6 +416,37 @@ def cap_max_batch_size_to_max_num_tokens(self): self.max_batch_size = self.max_num_tokens return self + @model_validator(mode="after") + def reject_unsupported_nanojet_attention(self): + """Refuse configurations the nanojet attention backend cannot serve, at config time. + + It keeps no KV cache — every answer comes from the Q/K/V of the current call — so + anything that needs history is out of reach. Catching it here rather than at the + first forward means the user is told before loading weights and building the graph, + and is never handed a silently wrong result. + """ + if self.attn_backend != "nanojet": + return self + + unsupported = [] + if self.enable_chunked_prefill: + unsupported.append( + "chunked prefill (a continuation chunk's earlier keys live in a cache)" + ) + if self.speculative_config is not None: + unsupported.append("speculative decoding (needs cached history)") + if self.cuda_graph_config is not None: + unsupported.append( + "CUDA graph capture (the shapes it captures are decode shapes)" + ) + if unsupported: + raise ValueError( + "attn_backend='nanojet' is prefill-only and keeps no KV cache, so it cannot " + "support: " + "; ".join(unsupported) + ". Use attn_backend='trtllm' or " + "'flashinfer', or disable the feature." + ) + return self + @model_validator(mode="after") def reject_cudagraph_for_speculative_flashinfer(self): if ( diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_attn_quant_fp8.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_attn_quant_fp8.py new file mode 100644 index 000000000000..d18e7f3a0208 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_attn_quant_fp8.py @@ -0,0 +1,87 @@ +"""Let the nanojet attention quantize in its epilogue when its reader wants e4m3. + +Mirrors ``fuse_trtllm_attn_quant_fp8`` for the nanojet backend: the scale is recorded on the +source attention node, which ``NanojetAttention.get_constants`` reads when the cached node is +built. Recording it here, while shapes are still propagated, is also what lets the following +fusions see an e4m3 activation. +""" + +from typing import Tuple, Type + +import torch +from torch.fx import GraphModule + +from ...custom_ops.attention.nanojet_attention import NANOJET_ATTENTION_INPUT_SCALE +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ...utils.logger import ad_logger +from ...utils.nanojet_graph import per_tensor_scale, set_val_meta +from ...utils.node_utils import ( + collect_terminal_users_through_passthrough, + get_shared_input_scale_for_fp8_linears, + is_op, +) +from ..interface import ( + BaseTransform, + SharedConfig, + TransformConfig, + TransformInfo, + TransformRegistry, +) + + +class FuseNanojetAttnQuantFP8Config(TransformConfig): + """Configuration for quantizing the nanojet attention output in its epilogue.""" + + +@TransformRegistry.register("fuse_nanojet_attn_quant_fp8") +class FuseNanojetAttnQuantFP8(BaseTransform): + """Fold the quantize before ``o_proj`` into the attention that feeds it.""" + + config: FuseNanojetAttnQuantFP8Config + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return FuseNanojetAttnQuantFP8Config + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + if not self.config.enabled: + return gm, TransformInfo(skipped=True, num_matches=0) + + num_matches = 0 + for attn_node in list(gm.graph.nodes): + if not is_op(attn_node, torch.ops.auto_deploy.torch_attention.default): + continue + attn_node.meta.pop(NANOJET_ATTENTION_INPUT_SCALE, None) + + readers, traversal_ok = collect_terminal_users_through_passthrough(attn_node) + fp8_readers, scale = get_shared_input_scale_for_fp8_linears(readers) + if not (traversal_ok and fp8_readers and len(fp8_readers) == len(readers)): + continue + if per_tensor_scale(gm, scale) is None: + continue + + attn_node.meta[NANOJET_ATTENTION_INPUT_SCALE] = scale + # The epilogue changes what this node produces, so say so: the fusions after this + # one decide by the dtype of the activation they read. + set_val_meta(attn_node, attn_node, dtype=torch.float8_e4m3fn) + num_matches += 1 + + if num_matches: + ad_logger.info( + f"fuse_nanojet_attn_quant_fp8: {num_matches} attention outputs quantized " + "in the epilogue" + ) + + return gm, TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=True, + has_valid_shapes=True, + ) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_fused_qkv_gemm_norm_rope.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_fused_qkv_gemm_norm_rope.py new file mode 100644 index 000000000000..ae70b74fc392 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_fused_qkv_gemm_norm_rope.py @@ -0,0 +1,368 @@ +"""Fuse three FP8 projections, the Q/K RMSNorms and RoPE into one nanojet kernel.""" + +from collections import Counter +from typing import List, Optional, Tuple, Type + +import torch +from torch.fx import GraphModule, Node + +from ....nanojet_utils import ensure_tune_configs, nanojet_supports +from ...custom_ops.linear.nanojet_fused_qkv_gemm_norm_rope import register +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ...utils.logger import ad_logger +from ...utils.nanojet_graph import ( + Fp8Projection, + get_attr_tensor, + match_fp8_projection, + set_val_meta, +) +from ...utils.node_utils import ( + extract_op_args, + extract_output_tuple, + is_op, + unwrap_input_through_passthrough, +) +from ..interface import ( + BaseTransform, + SharedConfig, + TransformConfig, + TransformInfo, + TransformRegistry, +) + +_HEAD_MAJOR_UNSQUEEZE_DIM = 2 + + +def _match_norm_over_projection(gm: GraphModule, node: Node): + """Resolve ``node`` to RMSNorm-over-FP8-projection, or a reason string on rejection.""" + source, _ = unwrap_input_through_passthrough(node) + if source is None or not is_op(source, torch.ops.auto_deploy.torch_rmsnorm): + return f"not-rmsnorm({source.target if isinstance(source, Node) else source})" + eps = extract_op_args(source, "eps")[0] + weight = extract_op_args(source, "weight")[0] + if not isinstance(eps, float) or not isinstance(weight, Node): + return "rmsnorm-args" + if len(source.users) != 1: + return f"norm-users={len(source.users)}" + projection = match_fp8_projection(gm, source.args[0]) + if projection is None: + return "not-fp8-projection" + return projection, weight, eps + + +def _match_rope_table(node: Node) -> Optional[Tuple[Node, Node]]: + """Resolve a per-token cos/sin to the (table, position index) it was gathered from.""" + source, _ = unwrap_input_through_passthrough(node) + if source is None or not is_op(source, torch.ops.aten.index.Tensor): + return None + if len(source.args) < 2 or not isinstance(source.args[0], Node): + return None + indices = source.args[1] + if not isinstance(indices, (list, tuple)) or len(indices) != 1: + return None + return source.args[0], indices[0] + + +class FuseNanojetFusedQKVGemmNormRopeConfig(TransformConfig): + """Configuration for the nanojet FP8 QKV + norm + RoPE fusion.""" + + +@TransformRegistry.register("fuse_nanojet_fused_qkv_gemm_norm_rope") +class FuseNanojetFusedQKVGemmNormRope(BaseTransform): + """Collapse three FP8 projections, two RMSNorms and RoPE into one nanojet node.""" + + config: FuseNanojetFusedQKVGemmNormRopeConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return FuseNanojetFusedQKVGemmNormRopeConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + if not self.config.enabled or not register(): + return gm, TransformInfo(skipped=True, num_matches=0) + + ensure_tune_configs(factory) + + # tp_mode says colwise even unsharded, so world size is what decides. + if shared_config.world_size > 1: + ad_logger.info("fuse_nanojet_fused_qkv_gemm_norm_rope: skipped, not supported under sharding") + return gm, TransformInfo(skipped=True, num_matches=0) + + graph = gm.graph + num_matches = 0 + cache_nodes: dict = {} + # Only original nodes are looked up, so one snapshot stays valid. + order = {node: index for index, node in enumerate(graph.nodes)} + rejected: Counter = Counter() + + for node in list(graph.nodes): + if not is_op(node, torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin): + continue + match = self._try_fuse(gm, node) + if isinstance(match, str): + rejected[match] += 1 + continue + num_matches += self._rewrite(gm, node, match, cache_nodes, order) + + if num_matches: + graph.eliminate_dead_code() + gm.recompile() + ad_logger.info(f"fuse_nanojet_fused_qkv_gemm_norm_rope: {num_matches} FP8 QKV+norm+RoPE fused") + else: + found = Counter( + str(node.target).rsplit(".", 1)[0] + for node in graph.nodes + if node.op == "call_function" + and ("rope" in str(node.target) or "rotary" in str(node.target)) + ) + + def summarize(counter): + return ", ".join(f"{name} x{count}" for name, count in counter.items()) or "none" + + inventory = Counter( + str(node.target).rsplit(".", 1)[0] + for node in graph.nodes + if node.op == "call_function" and str(node.target).startswith("auto_deploy.") + ) + ad_logger.warning( + f"fuse_nanojet_fused_qkv_gemm_norm_rope matched nothing; rope ops present: {summarize(found)}" + f"; rejected by: {summarize(rejected)}" + f"; auto_deploy ops in graph: {summarize(inventory)}" + ) + + return gm, TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=num_matches == 0, + ) + + @staticmethod + def _try_fuse(gm: GraphModule, rope_node: Node): + """Validate the six-node subgraph, returning its parts or a reason string on rejection. + + Returning why it failed rather than a bare ``None`` is what makes a silent + zero-match run diagnosable: the caller tallies the reasons into its warning. + """ + if extract_op_args(rope_node, "unsqueeze_dim")[0] != _HEAD_MAJOR_UNSQUEEZE_DIM: + return "rope-layout" + query, key = extract_op_args(rope_node, "q")[0], extract_op_args(rope_node, "k")[0] + cos, sin = extract_op_args(rope_node, "cos")[0], extract_op_args(rope_node, "sin")[0] + if not all(isinstance(arg, Node) for arg in (query, key, cos, sin)): + return "rope-args-not-nodes" + + # Head counts cannot tell Q from K when num_kv_heads == num_heads. + query_match = _match_norm_over_projection(gm, query) + key_match = _match_norm_over_projection(gm, key) + if isinstance(query_match, str) or isinstance(key_match, str): + reason = query_match if isinstance(query_match, str) else key_match + return f"norm-over-projection[{reason}]" + if query_match[2] != key_match[2]: + return "eps-mismatch" + + query_projection, query_norm_weight, eps = query_match + key_projection, key_norm_weight, _ = key_match + + value_projection = FuseNanojetFusedQKVGemmNormRope._find_value_projection( + gm, query_projection, key_projection + ) + if value_projection is None: + return "no-value-projection" + + # n_q/n_kv only describes the stacked weight if K and V match. + if key_projection.weight.shape[0] != value_projection.weight.shape[0]: + return "kv-row-mismatch" + if query_projection.weight.shape[1] != key_projection.weight.shape[1]: + return "hidden-size-mismatch" + + scales = {p.input_scale for p in (query_projection, key_projection, value_projection)} + if len(scales) != 1: + return "input-scale-mismatch" + + cos_table = _match_rope_table(cos) + sin_table = _match_rope_table(sin) + if cos_table is None or sin_table is None or cos_table[1] is not sin_table[1]: + return "rope-table" + + norm_value = query_norm_weight.meta.get("val") + if norm_value is None: + norm_value = get_attr_tensor(gm, query_norm_weight) + if norm_value is None: + return "norm-weight" + head_dim = int(norm_value.shape[-1]) + + # The rewrite keeps each table's first half, which needs a head_dim-wide table. + cos_value = cos_table[0].meta.get("val") + if cos_value is None: + cos_value = get_attr_tensor(gm, cos_table[0]) + if cos_value is None or int(cos_value.shape[-1]) != head_dim: + return "rope-table-width" + + if not nanojet_supports( + "fused_qkv_gemm_norm_rope", + input_dtype="float8_e4m3fn", + weight_dtype="float8_e4m3fn", + head_dim=head_dim, + ): + return "nanojet-declined" + + return ( + query_projection, + key_projection, + value_projection, + query_norm_weight, + key_norm_weight, + eps, + cos_table, + sin_table, + head_dim, + ) + + @staticmethod + def _find_value_projection( + gm: GraphModule, query: Fp8Projection, key: Fp8Projection + ) -> Optional[Fp8Projection]: + """The remaining FP8 projection reading the same activations as Q and K.""" + activations = query.node.args[0] + if not isinstance(activations, Node): + return None + candidates: List[Fp8Projection] = [] + for user in activations.users: + if user is query.node or user is key.node: + continue + projection = match_fp8_projection(gm, user) + if projection is not None: + candidates.append(projection) + return candidates[0] if len(candidates) == 1 else None + + def _rewrite( + self, gm: GraphModule, rope_node: Node, match, cache_nodes: dict, order: dict + ) -> int: + ( + query_projection, + key_projection, + value_projection, + query_norm_weight, + key_norm_weight, + eps, + cos_table, + sin_table, + head_dim, + ) = match + graph = gm.graph + anchor = max( + (query_projection.node, key_projection.node, value_projection.node), + key=lambda candidate: order.get(candidate, 0), + ).next + + stacked = torch.cat( + [query_projection.weight, key_projection.weight, value_projection.weight], dim=0 + ).contiguous() + # node.name, not id(): a reused address would make two layers share a weight. + weight_name = f"nanojet_qkv_weight_{rope_node.name}" + gm.register_buffer(weight_name, stacked) + + cache_key = (cos_table[0], sin_table[0]) + if cache_key not in cache_nodes: + half_dim = head_dim // 2 + table = cos_table[0].meta.get("val") + with graph.inserting_before(anchor): + cos_half = graph.call_function( + torch.ops.aten.slice.Tensor, args=(cos_table[0], -1, 0, half_dim) + ) + sin_half = graph.call_function( + torch.ops.aten.slice.Tensor, args=(sin_table[0], -1, 0, half_dim) + ) + stitched = graph.call_function( + torch.ops.aten.cat.default, args=([cos_half, sin_half], -1) + ) + contiguous = graph.call_function( + torch.ops.aten.contiguous.default, args=(stitched,) + ) + if table is not None: + half_shape = (*table.shape[:-1], half_dim) + set_val_meta(cos_half, table, half_shape) + set_val_meta(sin_half, table, half_shape) + set_val_meta(stitched, table, (*table.shape[:-1], 2 * half_dim)) + set_val_meta(contiguous, table, (*table.shape[:-1], 2 * half_dim)) + cache_nodes[cache_key] = contiguous + cache_node = cache_nodes[cache_key] + + with graph.inserting_before(anchor): + weight_node = graph.get_attr(weight_name) + set_val_meta(weight_node, stacked) + position_key = ("positions", cos_table[1]) + if position_key not in cache_nodes: + positions = graph.call_function( + torch.ops.aten.to.dtype, args=(cos_table[1], torch.int32) + ) + index_value = cos_table[1].meta.get("val") + if index_value is not None: + positions.meta["val"] = index_value.new_empty( + index_value.shape, dtype=torch.int32 + ) + positions.meta.pop("tensor_meta", None) + cache_nodes[position_key] = positions + fused = graph.call_function( + torch.ops.auto_deploy.nanojet_fused_qkv_gemm_norm_rope.default, + args=( + query_projection.node.args[0], + weight_node, + query_norm_weight, + key_norm_weight, + cache_node, + cache_nodes[position_key], + query_projection.input_scale_node, + eps, + query_projection.weight_scale * query_projection.input_scale, + key_projection.weight_scale * query_projection.input_scale, + value_projection.weight_scale * query_projection.input_scale, + int(query_projection.weight.shape[0]), + int(key_projection.weight.shape[0]), + ), + ) + query_size = int(query_projection.weight.shape[0]) + kv_size = int(key_projection.weight.shape[0]) + # The rotation returns q and k, so its meta is a tuple; either carries the batch + # and sequence dims and the dtype the fused node keeps. + rope_value = rope_node.meta.get("val") + template = rope_value[0] if isinstance(rope_value, (tuple, list)) else rope_value + leading = tuple(template.shape[:-2]) if template is not None else None + + def slice_out(start: int, stop: int) -> Node: + node = graph.call_function( + torch.ops.aten.slice.Tensor, args=(fused, -1, start, stop) + ) + if leading is not None: + set_val_meta(node, template, (*leading, stop - start)) + return node + + def split_heads(node: Node, width: int) -> Node: + split = graph.call_function( + torch.ops.aten.unflatten.int, args=(node, -1, [width // head_dim, head_dim]) + ) + if leading is not None: + set_val_meta(split, template, (*leading, width // head_dim, head_dim)) + return split + + if leading is not None: + set_val_meta(fused, template, (*leading, query_size + 2 * kv_size)) + + query_out = split_heads(slice_out(0, query_size), query_size) + key_out = split_heads(slice_out(query_size, query_size + kv_size), kv_size) + value_out = slice_out(query_size + kv_size, query_size + 2 * kv_size) + + # The rotation's schema names its arguments q and k, so result 0 is Q's. + query_getitem, key_getitem = extract_output_tuple(rope_node, 2) + for getitem, replacement in ((query_getitem, query_out), (key_getitem, key_out)): + if getitem is not None: + getitem.replace_all_uses_with(replacement) + value_projection.node.replace_all_uses_with(value_out) + return 1 diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_gemm_fp8_add.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_gemm_fp8_add.py new file mode 100644 index 000000000000..544a4d3e65ae --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_gemm_fp8_add.py @@ -0,0 +1,160 @@ +"""Fuse an e4m3-fed FP8 projection and the residual add reading it into one nanojet GEMM.""" + +from collections import Counter +from typing import Tuple, Type + +import torch +from torch.fx import GraphModule, Node + +from ....nanojet_utils import ensure_tune_configs +from ...custom_ops.linear.nanojet_gemm_fp8_add_inplace import register +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ...utils.logger import ad_logger +from ...utils.nanojet_graph import get_attr_tensor, is_fp8_linear, per_tensor_scale, set_val_meta +from ...utils.node_utils import extract_op_args, is_op, unwrap_input_through_passthrough +from ..interface import ( + BaseTransform, + SharedConfig, + TransformConfig, + TransformInfo, + TransformRegistry, +) + + +def _sole_residual_add(node: Node, order: dict): + """The ``add`` this projection feeds and the accumulator it may safely write into. + + Safe only when nothing reads the accumulator after this add, since the op writes into it. + """ + if len(node.users) != 1: + return None + add_node = next(iter(node.users)) + if not is_op(add_node, torch.ops.aten.add.Tensor) or len(add_node.args) != 2: + return None + left, right = add_node.args + if left is node: + other = right + elif right is node: + other = left + else: + return None + if not isinstance(other, Node): + return None + value = other.meta.get("val") + if value is None or value.dtype != torch.bfloat16: + return None + # Writing into a graph input would corrupt a tensor this pass has no claim on. + if other.op == "placeholder": + return None + # Accumulating into its own activation reads and overwrites one buffer. + if other is node.args[0]: + return None + # The GEMM writes into the accumulator, so the shapes must already match. + projection = node.meta.get("val") + if projection is None or tuple(projection.shape) != tuple(value.shape): + return None + add_position = order.get(add_node, -1) + if any(order.get(user, -1) > add_position for user in other.users if user is not add_node): + return None + return add_node, other + + +def _is_fp8(node) -> bool: + """Whether the value reaching this projection is e4m3.""" + source, _ = unwrap_input_through_passthrough(node) + value = source.meta.get("val") if isinstance(source, Node) else None + return value is not None and value.dtype == torch.float8_e4m3fn + + +class FuseNanojetGemmFP8AddConfig(TransformConfig): + """Configuration for folding an FP8 projection into its residual add.""" + + +@TransformRegistry.register("fuse_nanojet_gemm_fp8_add") +class FuseNanojetGemmFP8Add(BaseTransform): + """Fuse an e4m3-fed FP8 linear and the residual add reading it into one nanojet GEMM.""" + + config: FuseNanojetGemmFP8AddConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return FuseNanojetGemmFP8AddConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + if not self.config.enabled or not register(): + return gm, TransformInfo(skipped=True, num_matches=0) + + ensure_tune_configs(factory) + + graph = gm.graph + num_matches = 0 + rejected: Counter = Counter() + + order = {n: i for i, n in enumerate(graph.nodes)} + for node in reversed(list(graph.nodes)): + if not is_fp8_linear(node): + continue + if extract_op_args(node, "bias")[0] is not None: # bias needs folding + rejected["bias"] += 1 + continue + activation = node.args[0] if node.args else None + if not isinstance(activation, Node): + rejected["activation-not-node"] += 1 + continue + if not _is_fp8(activation): + rejected["activation-dtype"] += 1 + continue + weight_node = extract_op_args(node, "weight_fp8")[0] + weight = get_attr_tensor(gm, weight_node) + if weight is None or weight.dtype != torch.float8_e4m3fn: + rejected["weight-dtype"] += 1 + continue + input_scale = per_tensor_scale(gm, extract_op_args(node, "input_scale")[0]) + weight_scale = per_tensor_scale(gm, extract_op_args(node, "weight_scale")[0]) + if input_scale is None or weight_scale is None: + rejected["scales"] += 1 + continue + + residual = _sole_residual_add(node, order) + if residual is None: + rejected["no-residual-add"] += 1 + continue + add_node, accumulator = residual + # Node order is what tells FX "accumulated, then read". + with graph.inserting_before(add_node): + graph.call_function( + torch.ops.auto_deploy.nanojet_gemm_fp8_add_inplace.default, + args=(activation, weight_node, accumulator, input_scale, weight_scale), + ) + set_val_meta(accumulator, add_node) + add_node.replace_all_uses_with(accumulator) + graph.erase_node(add_node) + order = {n: i for i, n in enumerate(graph.nodes)} + num_matches += 1 + + if num_matches: + graph.eliminate_dead_code() + gm.recompile() + ad_logger.info( + f"fuse_nanojet_gemm_fp8_add: {num_matches} FP8 projections folded into " + "their residual add" + ) + if rejected: + ad_logger.info( + "fuse_nanojet_gemm_fp8_add rejected: " + + ", ".join(f"{reason} x{count}" for reason, count in rejected.items()) + ) + + return gm, TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=num_matches == 0, + ) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_rmsnorm_fp8.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_rmsnorm_fp8.py new file mode 100644 index 000000000000..1f6d7656f9c6 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_rmsnorm_fp8.py @@ -0,0 +1,160 @@ +"""Move the FP8 quantize into the RMSNorm that feeds it. + +Where every reader of a norm takes an already-quantized activation, the conversion belongs +in the norm's epilogue rather than in a pass of its own. That is what nanojet runs natively, +and it is why the native layer has no standalone quantize kernel. + +Runs after ``fuse_fp8_linear`` so the readers are the TensorRT LLM FP8 op, and before the +gated SwiGLU GEMM, which needs an FP8 activation to apply at all. +""" + +from typing import Optional, Tuple, Type + +import torch +from torch.fx import GraphModule, Node + +from ....nanojet_utils import ensure_tune_configs, nanojet_supports +from ...custom_ops.normalization.nanojet_rmsnorm_fp8 import register +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ...utils.logger import ad_logger +from ...utils.nanojet_graph import accepts_out_dtype, fp8_linear_ops, per_tensor_scale, set_val_meta +from ...utils.node_utils import ( + collect_terminal_users_through_passthrough, + extract_op_args, + is_op, + set_op_args, +) +from ..interface import ( + BaseTransform, + SharedConfig, + TransformConfig, + TransformInfo, + TransformRegistry, +) + + +# Consumers that read an already-quantized activation, and where each keeps the scale it +# dequantizes by. Resolved lazily: the nanojet ops only exist once nanojet is installed, and +# this module is imported unconditionally by the library's package scan. +def _fp8_consumer_ops(): + """The ops that can take this norm's output already quantized. + + They all spell the argument ``input_scale``, so only the op identities are listed here — + the position is looked up from each op's schema rather than tabulated by hand. + """ + consumers = list(fp8_linear_ops()) + # Registered only if the QKV fusion is also enabled; absent simply means this norm has no + # such consumer to consider. + fused_qkv = getattr(torch.ops.auto_deploy, "nanojet_fused_qkv_gemm_norm_rope", None) + if fused_qkv is not None: + consumers.append(fused_qkv) + return consumers + + +def _shared_quantize_scale(gm: GraphModule, normed: Node) -> Optional[float]: + """The one dequant scale every reader of ``normed`` uses, if they all take FP8. + + Quantizing in the epilogue only works when nothing downstream wants the BF16 value and + everyone agrees on the scale; otherwise the saved kernel reappears elsewhere, or a + reader silently sees a differently-scaled tensor. + + ``get_shared_input_scale_for_fp8_linears`` answers the same question for linears alone; + the readers here also include nanojet's fused QKV, which it does not know about. + """ + readers, traversal_ok = collect_terminal_users_through_passthrough(normed) + if not traversal_ok or not readers: + return None + consumer_ops = _fp8_consumer_ops() + scales = set() + for reader in readers: + if not any(is_op(reader, op) for op in consumer_ops): + return None + scale = per_tensor_scale(gm, extract_op_args(reader, "input_scale")[0]) + if scale is None: + return None + scales.add(scale) + return scales.pop() if len(scales) == 1 else None + + +class FuseNanojetRMSNormFP8Config(TransformConfig): + """Configuration for folding the FP8 quantize into nanojet's RMSNorm.""" + + +@TransformRegistry.register("fuse_nanojet_rmsnorm_fp8") +class FuseNanojetRMSNormFP8(BaseTransform): + """Replace ``rmsnorm`` + downstream quantize with one nanojet norm emitting e4m3.""" + + config: FuseNanojetRMSNormFP8Config + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return FuseNanojetRMSNormFP8Config + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + if not self.config.enabled or not register(): + return gm, TransformInfo(skipped=True, num_matches=0) + + ensure_tune_configs(factory) + + graph = gm.graph + num_matches = 0 + + for node in list(graph.nodes): + if not is_op(node, torch.ops.auto_deploy.torch_rmsnorm) or len(node.args) < 3: + continue + hidden_states, weight, eps = node.args[0], node.args[1], node.args[2] + if not isinstance(hidden_states, Node) or not isinstance(weight, Node): + continue + if not isinstance(eps, float): + continue + value = node.meta.get("val") + if value is None or value.dtype != torch.bfloat16: + continue + if not nanojet_supports( + "unified_rmsnorm", + hidden_size=int(value.shape[-1]), + zero_centered_weight=False, + multiply_in_fp32=False, + input_dtype=value.dtype, + ): + continue + + quantize_scale = _shared_quantize_scale(gm, node) + if quantize_scale is None: + continue + + output_dtype = str(value.dtype).rsplit(".", 1)[-1] + with graph.inserting_before(node): + fused = graph.call_function( + torch.ops.auto_deploy.nanojet_rmsnorm_fp8.default, + args=(hidden_states, weight, eps, 1.0 / quantize_scale), + ) + # An FP8 activation carries no hint of what the linear should emit. + readers, _ = collect_terminal_users_through_passthrough(node) + for reader in readers: + if accepts_out_dtype(reader): + set_op_args(reader, out_dtype=output_dtype) + set_val_meta(fused, node, dtype=torch.float8_e4m3fn) + node.replace_all_uses_with(fused) + num_matches += 1 + + if num_matches: + graph.eliminate_dead_code() + gm.recompile() + ad_logger.info( + f"fuse_nanojet_rmsnorm_fp8: {num_matches} norms quantizing in the epilogue" + ) + + return gm, TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=True, + ) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_swiglu_gemm_fp8.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_swiglu_gemm_fp8.py new file mode 100644 index 000000000000..552cfc4dae48 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_nanojet_swiglu_gemm_fp8.py @@ -0,0 +1,150 @@ +"""Collapse the whole SwiGLU MLP front half into one nanojet gated GEMM. + +The graph arrives as two FP8 projections off one activation, then silu, then a multiply:: + + gate_projection = fp8_linear(x, gate_proj.weight, ..., [in_scale], [gate_scale]) + up_projection = fp8_linear(x, up_proj.weight, ..., [in_scale], [up_scale]) + h = mul(silu(gate_projection), up_projection) + +Four kernels once the quantize before ``down_proj`` is counted. nanojet's ``swiglu`` is a +CUTLASS gated GEMM that does all of it and writes e4m3, which is what the native path runs. +Ordered before ``fuse_nanojet_act_and_mul`` so this takes the MLPs it can and that one keeps +whatever is left. +""" + +from typing import Tuple, Type + +import torch +from torch.fx import GraphModule, Node + +from ....nanojet_utils import ensure_tune_configs +from ...custom_ops.linear.nanojet_swiglu_gemm_fp8 import register +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ...utils.logger import ad_logger +from ...utils.nanojet_graph import ( + accepts_out_dtype, + match_fp8_projection, + per_tensor_scale, + set_val_meta, +) +from ...utils.node_utils import ( + collect_terminal_users_through_passthrough, + get_shared_input_scale_for_fp8_linears, + is_op, + set_op_args, +) +from ..interface import ( + BaseTransform, + SharedConfig, + TransformConfig, + TransformInfo, + TransformRegistry, +) + + +class FuseNanojetSwiGLUGemmFP8Config(TransformConfig): + """Configuration for the nanojet gated SwiGLU GEMM fusion.""" + + +@TransformRegistry.register("fuse_nanojet_swiglu_gemm_fp8") +class FuseNanojetSwiGLUGemmFP8(BaseTransform): + """Replace two FP8 projections + silu + mul with one nanojet gated GEMM.""" + + config: FuseNanojetSwiGLUGemmFP8Config + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return FuseNanojetSwiGLUGemmFP8Config + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + if not self.config.enabled or not register(): + return gm, TransformInfo(skipped=True, num_matches=0) + + ensure_tune_configs(factory) + + graph = gm.graph + num_matches = 0 + + for node in list(graph.nodes): + if not is_op(node, torch.ops.aten.mul.Tensor) or len(node.args) != 2: + continue + first, second = node.args + # The product is symmetric; the activation marks which side is the gate_projection. + if is_op(first, torch.ops.aten.silu): + silu_node, up_node = first, second + elif is_op(second, torch.ops.aten.silu): + silu_node, up_node = second, first + else: + continue + if not isinstance(up_node, Node) or len(silu_node.users) != 1: + continue + + gate = match_fp8_projection(gm, silu_node.args[0]) + up = match_fp8_projection(gm, up_node) + if gate is None or up is None: + continue + # One GEMM means one activation and one activation scale for both halves. + if gate.activation is not up.activation or gate.input_scale != up.input_scale: + continue + if gate.weight.shape != up.weight.shape: + continue + readers, traversal_ok = collect_terminal_users_through_passthrough(node) + fp8_readers, scale = get_shared_input_scale_for_fp8_linears(readers) + if not (traversal_ok and fp8_readers and len(fp8_readers) == len(readers)): + continue + output_scale = per_tensor_scale(gm, scale) + if output_scale is None: + continue + + # ``[up, gate]`` — up first, which is the order nanojet's kernel indexes. + stacked = torch.cat([up.weight, gate.weight], dim=0).contiguous() + # Graph-unique node name, not ``id()``: an address is neither stable across runs + # nor unique over time, since a collected node's address can be reused. + weight_name = f"nanojet_gate_up_weight_{node.name}" + gm.register_buffer(weight_name, stacked) + + with graph.inserting_before(node): + weight_node = graph.get_attr(weight_name) + set_val_meta(weight_node, stacked) + # Fresh get_attr: the projections' scale nodes may sit later in the graph + # than the node being inserted here. + gate_scale_node = graph.get_attr(gate.input_scale_node.target) + fused = graph.call_function( + torch.ops.auto_deploy.nanojet_swiglu_gemm_fp8.default, + args=( + gate.activation, + gate_scale_node, + weight_node, + gate.input_scale, + up.weight_scale, + gate.weight_scale, + 1.0 / output_scale, + ), + ) + # An e4m3 activation carries no hint of the output dtype, but only the TensorRT LLM + # linear takes the hint; the others fix their return dtype in the implementation. + for consumer in fp8_readers: + if accepts_out_dtype(consumer): + set_op_args(consumer, out_dtype=str(node.meta["val"].dtype).rsplit(".", 1)[-1]) + set_val_meta(fused, node, dtype=torch.float8_e4m3fn) + node.replace_all_uses_with(fused) + num_matches += 1 + + if num_matches: + graph.eliminate_dead_code() + gm.recompile() + ad_logger.info(f"fuse_nanojet_swiglu_gemm_fp8: {num_matches} gated SwiGLU GEMMs fused") + + return gm, TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=True, + ) diff --git a/tensorrt_llm/_torch/auto_deploy/utils/nanojet_graph.py b/tensorrt_llm/_torch/auto_deploy/utils/nanojet_graph.py new file mode 100644 index 000000000000..efe3070be4c9 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/utils/nanojet_graph.py @@ -0,0 +1,153 @@ +"""Shared graph queries for the nanojet fusions and the nanojet attention backend.""" + +import math +from typing import NamedTuple, Optional + +import torch +from torch.fx import GraphModule, Node + +from .node_utils import ( + extract_op_args, + get_param_or_buffer, + is_op, + unwrap_input_through_passthrough, +) + + +def get_attr_tensor(gm: GraphModule, node) -> Optional[torch.Tensor]: + """The tensor behind a ``get_attr`` node, or None if it is not one.""" + if not isinstance(node, Node) or node.op != "get_attr": + return None + try: + return get_param_or_buffer(node.target, gm) + except KeyError: + return None + + +def per_tensor_scale(gm: GraphModule, node) -> Optional[float]: + """A quantization scale as a plain float, or None if it is not one this can use. + + Accepts both spellings a quantized linear uses — a bare tensor, or the single-element + list ``torch_fake_quant_fp8_linear`` wraps it in. Rejects anything that is not a usable + divisor: every caller divides by this, so a zero, negative or non-finite value would + become an infinity that spreads through the model with nothing failing. + """ + if isinstance(node, (list, tuple)): + if len(node) != 1: + return None + node = node[0] + tensor = get_attr_tensor(gm, node) + if tensor is None or tensor.numel() != 1: + return None + scale = float(tensor.reshape(-1)[0]) + if not math.isfinite(scale) or scale <= 0.0: + return None + return scale + + +def scale_node(node): + """The node holding a quantization scale, unwrapping the single-element list spelling. + + Separate from :func:`per_tensor_scale` on purpose: one answers "what is the value", the + other "which node carries it". Returning both from one function as a tuple is what let a + caller keep indexing it after the value-only version replaced it. + """ + if isinstance(node, (list, tuple)): + if len(node) != 1: + return None + node = node[0] + return node if isinstance(node, Node) and node.op == "get_attr" else None + + +_FP8_LINEAR_NAMES = ( + "trtllm_quant_fp8_linear", + "torch_quant_fp8_linear", + "torch_fake_quant_fp8_linear", +) + + +def fp8_linear_ops(): + """auto_deploy's FP8 linears, resolved on use rather than at import. + + This module is imported during the custom-op package scan, before those ops exist. + """ + namespace = torch.ops.auto_deploy + return tuple( + getattr(namespace, name) for name in _FP8_LINEAR_NAMES if hasattr(namespace, name) + ) + + +def is_fp8_linear(node: Node) -> bool: + """Any of auto_deploy's FP8 linears. + + They agree on the leading argument names and on scales being amax/448, which is all the + nanojet fusions read; which one the graph holds depends on ``fuse_fp8_linear.backend``. + """ + return any(is_op(node, op) for op in fp8_linear_ops()) + + +def accepts_out_dtype(node: Node) -> bool: + """Whether this linear takes an out_dtype hint; only the TensorRT LLM one does.""" + return any(argument.name == "out_dtype" for argument in node.target._schema.arguments) + + +class Fp8Projection(NamedTuple): + """One matched FP8 projection: the node, its weight and its scales.""" + + node: Node + activation: Node + weight: torch.Tensor + input_scale: float + weight_scale: float + input_scale_node: Optional[Node] + + +def match_fp8_projection(gm: GraphModule, node: Node) -> Optional[Fp8Projection]: + """Resolve ``node``, through any views, to the FP8 linear producing it. + + Only bias-free projections match: a bias would have to be folded into the epilogue, and + nanojet's kernels do not take one. + """ + source, _ = unwrap_input_through_passthrough(node) + if not isinstance(source, Node) or not is_fp8_linear(source) or not source.args: + return None + if extract_op_args(source, "bias")[0] is not None: + return None + # The FP8 linears name their weight differently; everything else they spell the same. + weight_node = extract_op_args(source, "weight_fp8")[0] + if weight_node is None: + weight_node = extract_op_args(source, "weight_quantized")[0] + weight = get_attr_tensor(gm, weight_node) + if weight is None or weight.dtype != torch.float8_e4m3fn: + return None + input_scale_arg = extract_op_args(source, "input_scale")[0] + input_scale = per_tensor_scale(gm, input_scale_arg) + weight_scale = per_tensor_scale(gm, extract_op_args(source, "weight_scale")[0]) + if input_scale is None or weight_scale is None: + return None + return Fp8Projection( + node=source, + activation=source.args[0], + weight=weight, + input_scale=input_scale, + weight_scale=weight_scale, + input_scale_node=scale_node(input_scale_arg), + ) + + +def set_val_meta(node: Node, source, shape=None, dtype=None) -> None: + """Record what ``node`` produces, taking shape and dtype from ``source`` by default. + + ``source`` is a node, a meta value or a real tensor. ``tensor_meta`` is a second record + of the same facts and must not outlive them. + """ + value = source.meta.get("val") if isinstance(source, Node) else source + if value is None: + return + target_shape = value.shape if shape is None else tuple(shape) + target_dtype = value.dtype if dtype is None else dtype + if hasattr(value, "new_empty") and value.device.type == "meta": + node.meta["val"] = value.new_empty(target_shape, dtype=target_dtype) + else: + node.meta["val"] = torch.empty(target_shape, dtype=target_dtype, device="meta") + node.meta.pop("tensor_meta", None) diff --git a/tensorrt_llm/_torch/nanojet_utils.py b/tensorrt_llm/_torch/nanojet_utils.py new file mode 100644 index 000000000000..c45cf7fb43a7 --- /dev/null +++ b/tensorrt_llm/_torch/nanojet_utils.py @@ -0,0 +1,94 @@ +"""Lazy access to nanojet. + +Nothing here imports nanojet at module load. TensorRT LLM scans its custom-op and transform +packages unconditionally, so an eager import would mean every install pays for nanojet — +attempting it, failing, and logging — whether or not the user asked for a nanojet pass. The +import happens on the first call, which only comes from a nanojet transform or backend that +was explicitly enabled. +""" + +import platform +from typing import Optional + +from ..logger import logger + +_CONTRACT = None +_AVAILABLE: Optional[bool] = None + + +def is_nanojet_available() -> bool: + """Whether nanojet and its integration contract can be loaded. Cached after the first call.""" + global _AVAILABLE, _CONTRACT + if _AVAILABLE is not None: + return _AVAILABLE + + _AVAILABLE = False + if platform.system() == "Windows": + return _AVAILABLE + try: + import nanojet_kernels + from nanojet_kernels.interface_contract import trtllm as contract + + _CONTRACT = contract + _AVAILABLE = True + logger.info(f"nanojet is available: {nanojet_kernels.__version__}") + except ImportError: + logger.warning("nanojet requested but not importable; nanojet passes will not apply") + except AttributeError: + # Installed but without the integration contract: too old to drive from here. + logger.warning( + "nanojet is installed but exposes no interface_contract.trtllm; skipping nanojet ops" + ) + return _AVAILABLE + + +_MODELS_WITH_TUNE_CONFIGS_APPLIED: set = set() + +# nanojet files its tuned tiles under these names; TensorRT LLM spells the algorithm its own +# way in ``quant_algo``. +_QUANT_ALGO_TO_NANOJET = {"FP8": "fp8", "FP8_BLOCK_SCALES": "blockwise_fp8"} + + +def ensure_tune_configs(factory) -> None: + """Load nanojet's tuned CUTLASS tiles for the model this factory built. Idempotent. + + Called from every nanojet transform so no combination of enabled passes can miss it. + A missing config or unrecognized quantization leaves the kernels on default tiles. + """ + if not is_nanojet_available(): + return + try: + model_config, _ = factory._get_model_config() + quantization = _QUANT_ALGO_TO_NANOJET.get(factory.get_quant_config().get("quant_algo")) + if quantization is None: + return + shape = dict( + hidden_size=model_config.hidden_size, + intermediate_size=model_config.intermediate_size, + head_dim=getattr( + model_config, + "head_dim", + model_config.hidden_size // model_config.num_attention_heads, + ), + num_attention_heads=model_config.num_attention_heads, + num_key_value_heads=model_config.num_key_value_heads, + ) + model_identity = (model_config.model_type, quantization, tuple(sorted(shape.items()))) + if model_identity in _MODELS_WITH_TUNE_CONFIGS_APPLIED: + return + _MODELS_WITH_TUNE_CONFIGS_APPLIED.add(model_identity) + _CONTRACT.apply_tune_configs(model_config.model_type, quantization, shape) + except Exception as error: + logger.warning(f"could not apply nanojet tune configs: {type(error).__name__}: {error}") + + +def nanojet_supports(op: str, **constraints) -> bool: + """Ask nanojet whether it accepts a concrete configuration of ``op``. + + Graph transforms call this instead of restating nanojet's dispatch tables, so a nanojet + release that widens (or narrows) the shapes it handles takes effect here without any + change on our side. + """ + if not is_nanojet_available(): + return False + return _CONTRACT.supports(op, **constraints) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_nanojet_rejections.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_nanojet_rejections.py new file mode 100644 index 000000000000..4b95fe743540 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_nanojet_rejections.py @@ -0,0 +1,133 @@ +"""The nanojet integration must refuse what it cannot serve, and say so early. + +The backend keeps no KV cache: every answer comes from the Q/K/V of the current call. Any +configuration that needs history is therefore out of reach, and the failure mode if it were +allowed through is the dangerous one — attention over just the tokens present looks perfectly +plausible and is wrong. These tests pin the refusals. +""" + +import pytest +import torch + +MODEL = "/tmp/nanojet_quant_modelopt/Qwen3-Embedding-0.6B" + + +def _batch_info( + num_prefill: int, num_prefill_tokens: int, num_decode: int = 0, num_extend: int = 0 +): + """The 14-slot host tensor the attention ops read their batch composition from.""" + info = torch.zeros(14, dtype=torch.int32) + info[0] = num_prefill + info[1] = num_prefill_tokens + info[2] = num_extend + info[3] = num_extend + info[4] = num_decode + info[5] = num_decode + return info + + +def _call_attention(**overrides): + """Drive the op with a minimal single-sequence prefill, overriding one thing at a time.""" + from tensorrt_llm._torch.auto_deploy.custom_ops.attention import nanojet_attention + + assert nanojet_attention.register(), "nanojet must be installed for this test" + + tokens, heads, kv_heads, head_dim = 8, 4, 2, 128 + kwargs = dict( + q=torch.randn(1, tokens, heads, head_dim, device="cuda", dtype=torch.bfloat16), + k=torch.randn(1, tokens, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16), + v=torch.randn(1, tokens, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16), + batch_info_host=_batch_info(1, tokens), + seq_len=torch.tensor([tokens], dtype=torch.int32, device="cuda"), + input_pos=torch.zeros(1, dtype=torch.int32, device="cuda"), + slot_idx=torch.zeros(1, dtype=torch.int32, device="cuda"), + cu_seqlen=torch.tensor([0, tokens], dtype=torch.int32, device="cuda"), + scale=None, + ) + kwargs.update(overrides) + return torch.ops.auto_deploy.nanojet_attention(**kwargs) + + +# -------------------------------------------------------------------------------------- +# Config time: rejected before weights are loaded or a graph is built. +# -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "unsupported, needle", + [ + ({"enable_chunked_prefill": True}, "chunked prefill"), + ], +) +def test_config_rejects_unsupported_nanojet_attention(unsupported, needle): + from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs + + with pytest.raises(ValueError, match=needle): + LlmArgs(model=MODEL, attn_backend="nanojet", **unsupported) + + +def test_config_accepts_plain_prefill(): + from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs + + args = LlmArgs(model=MODEL, attn_backend="nanojet", nanojet_fusion=True) + assert args.transforms["fuse_nanojet_fused_qkv_gemm_norm_rope"]["enabled"] + + +def test_other_backends_keep_chunked_prefill(): + """The refusal is scoped to this backend and must not restrict the shipped ones.""" + from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs + + assert LlmArgs(model=MODEL, attn_backend="trtllm", enable_chunked_prefill=True) + + +# -------------------------------------------------------------------------------------- +# Import hygiene: nothing about nanojet is touched unless it was asked for. +# -------------------------------------------------------------------------------------- + + +def test_nanojet_not_imported_unless_requested(): + import subprocess + import sys + + probe = ( + "import sys;" + "import tensorrt_llm._torch.auto_deploy.custom_ops;" + "import tensorrt_llm._torch.auto_deploy.transform.library;" + "from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs;" + f"LlmArgs(model='{MODEL}');" + "print([m for m in sys.modules if m.startswith('nanojet_kernels')])" + ) + out = subprocess.run([sys.executable, "-c", probe], capture_output=True, text=True) + assert out.stdout.strip().endswith("[]"), f"nanojet was imported: {out.stdout}" + + +# -------------------------------------------------------------------------------------- +# Runtime: data-dependent cases the config cannot see. +# -------------------------------------------------------------------------------------- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_runtime_rejects_decode(): + with pytest.raises(RuntimeError, match="prefill-only"): + _call_attention(batch_info_host=_batch_info(0, 0, num_decode=1)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_runtime_rejects_continuation_chunk(): + """An extend request carries cached context this backend does not keep.""" + with pytest.raises(RuntimeError, match="continuation chunk"): + _call_attention(batch_info_host=_batch_info(0, 0, num_extend=1)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_runtime_rejects_shared_kv(): + with pytest.raises(RuntimeError, match="KV cache"): + _call_attention(read_cache_only=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_runtime_accepts_plain_prefill(): + """The supported case must still go through, or the guards above prove nothing.""" + out = _call_attention() + assert out.shape == (1, 8, 4, 128) + assert not out.isnan().any() diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_fused_qkv_gemm_norm_rope_guards.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_fused_qkv_gemm_norm_rope_guards.py new file mode 100644 index 000000000000..ee91ba6c4e4b --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_fused_qkv_gemm_norm_rope_guards.py @@ -0,0 +1,59 @@ +"""Regression tests for four silent-miscompute paths in fuse_nanojet_fused_qkv_gemm_norm_rope. + +Each of these produced a plausible, wrong result with nothing to indicate it, and none is +reachable with a GQA checkpoint like Qwen3 (16/8 heads) — which is exactly why the end-to-end +cosine similarity stayed identical while they were broken. They are driven directly instead. + +Covered here: the Q/K ordering contract and the buffer-naming rule. NOT covered: the +kv-row-mismatch and rope-table-width guards, which sit inside _try_fuse and need a full +six-node FP8 subgraph to reach — a traced nn.Module does not produce the quantized ops. +""" + +import pytest + +pytest.importorskip("nanojet_kernels") + +from tensorrt_llm._torch.auto_deploy.transform.library.fuse_nanojet_fused_qkv_gemm_norm_rope import ( # noqa: E402 + _HEAD_MAJOR_UNSQUEEZE_DIM, +) + + +def test_unsqueeze_dim_constant_is_head_major(): + """The fused kernel is written against [batch, seq, heads, head_dim].""" + assert _HEAD_MAJOR_UNSQUEEZE_DIM == 2 + + +def test_query_and_key_come_from_the_op_schema_not_head_counts(): + """Q is args[0] by the rotation's own signature. + + The previous code inferred it from head counts, which is ambiguous the moment + num_kv_heads == num_heads: `first_heads >= second_heads` is then always true, so a graph + handing the rotation K first had Q and K silently swapped. + """ + import torch as _torch + + schema = _torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin.default._schema + names = [argument.name for argument in schema.arguments] + assert names[:2] == ["q", "k"], f"rotation no longer names its inputs q,k: {names}" + assert names[4] == "unsqueeze_dim" + + +def test_buffer_names_are_graph_unique_not_addresses(): + """``id()`` is neither stable across runs nor unique over time. + + A collected node's address can be handed to a later one, which would make two layers + register the same buffer name and silently share one stacked weight. + """ + import inspect + + from tensorrt_llm._torch.auto_deploy.transform.library import ( + fuse_nanojet_fused_qkv_gemm_norm_rope, + fuse_nanojet_swiglu_gemm_fp8, + ) + + for module in (fuse_nanojet_fused_qkv_gemm_norm_rope, fuse_nanojet_swiglu_gemm_fp8): + source = inspect.getsource(module) + assert "id(rope_node)" not in source and "id(node)" not in source, ( + f"{module.__name__} derives a buffer name from an address" + ) + assert ".name}" in source, f"{module.__name__} should key buffers on node.name" diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_fused_qkv_gemm_norm_rope_shape_guards.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_fused_qkv_gemm_norm_rope_shape_guards.py new file mode 100644 index 000000000000..0d277baa84bf --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_fused_qkv_gemm_norm_rope_shape_guards.py @@ -0,0 +1,119 @@ +"""Drive _try_fuse over a hand-built six-node FP8 subgraph. + +symbolic_trace cannot produce this: the quantized linear, the norm and the rotation are +custom ops, so a traced nn.Module gets rejected long before the guards under test. The graph +is built node by node against the real op schemas instead. +""" + +import pytest +import torch + +pytest.importorskip("nanojet_kernels") + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401,E402 +from tensorrt_llm._torch.auto_deploy.custom_ops.linear.nanojet_fused_qkv_gemm_norm_rope import ( # noqa: E402 + register, +) +from tensorrt_llm._torch.auto_deploy.transform.library.fuse_nanojet_fused_qkv_gemm_norm_rope import ( # noqa: E402 + _HEAD_MAJOR_UNSQUEEZE_DIM, + FuseNanojetFusedQKVGemmNormRope, +) + +HIDDEN, HEAD_DIM, NUM_HEADS, NUM_KV_HEADS, EPS = 256, 64, 8, 4, 1e-6 + + +def _build(kv_rows_match: bool = True, table_width: int = HEAD_DIM): + """A q/k/v + norm + rope subgraph, returning (GraphModule, rope node). + + ``kv_rows_match=False`` gives V a different head count from K, which is what the stacked + ``[q + 2kv, hidden]`` weight cannot represent. ``table_width`` controls the rope table. + """ + module = torch.nn.Module() + graph = torch.fx.Graph() + + value_heads = NUM_KV_HEADS if kv_rows_match else NUM_KV_HEADS + 1 + tensors = { + "query_weight": torch.zeros(NUM_HEADS * HEAD_DIM, HIDDEN, dtype=torch.float8_e4m3fn), + "key_weight": torch.zeros(NUM_KV_HEADS * HEAD_DIM, HIDDEN, dtype=torch.float8_e4m3fn), + "value_weight": torch.zeros(value_heads * HEAD_DIM, HIDDEN, dtype=torch.float8_e4m3fn), + "input_scale": torch.tensor([0.5]), + "weight_scale": torch.tensor([0.25]), + "query_norm": torch.ones(HEAD_DIM), + "key_norm": torch.ones(HEAD_DIM), + "cos_table": torch.zeros(128, table_width), + "sin_table": torch.zeros(128, table_width), + } + for name, tensor in tensors.items(): + module.register_buffer(name, tensor) + + attrs = {name: graph.get_attr(name) for name in tensors} + hidden = graph.placeholder("hidden") + positions = graph.placeholder("positions") + + def projection(weight_attr): + return graph.call_function( + torch.ops.auto_deploy.trtllm_quant_fp8_linear.default, + args=(hidden, weight_attr, None, attrs["input_scale"], attrs["weight_scale"]), + ) + + def normed(projection_node, norm_attr): + return graph.call_function( + torch.ops.auto_deploy.torch_rmsnorm.default, + args=(projection_node, norm_attr, EPS), + ) + + query = normed(projection(attrs["query_weight"]), attrs["query_norm"]) + key = normed(projection(attrs["key_weight"]), attrs["key_norm"]) + projection(attrs["value_weight"]) # V has no norm; found by sharing the activation + + def gather(table): + return graph.call_function(torch.ops.aten.index.Tensor, args=(table, [positions])) + + rope = graph.call_function( + torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin.default, + args=( + query, + key, + gather(attrs["cos_table"]), + gather(attrs["sin_table"]), + _HEAD_MAJOR_UNSQUEEZE_DIM, + ), + ) + graph.output(rope) + return torch.fx.GraphModule(module, graph), rope + + +def _reason(**kwargs): + graph_module, rope = _build(**kwargs) + return FuseNanojetFusedQKVGemmNormRope._try_fuse(graph_module, rope) + + +def test_harness_reaches_the_shape_guards(): + """The control: a well-formed subgraph must get past every earlier check. + + Without this, a rejection below could come from the harness rather than the guard, which + is how a test like this quietly stops testing anything. + """ + assert register(), "nanojet must be installed" + result = _reason() + assert result not in ("kv-row-mismatch", "rope-table-width"), result + assert not isinstance(result, str) or "norm-over-projection" not in result, result + + +def test_rejects_value_head_count_differing_from_key(): + """The stacked weight is indexed by n_q/n_kv, which only describes it if K and V match. + + A mismatch produces a correctly-shaped buffer the kernel then reads at wrong offsets — + no error, just wrong numbers. Unreachable with Qwen3, whose K and V are both 8 heads. + """ + assert register(), "nanojet must be installed" + assert _reason(kv_rows_match=False) == "kv-row-mismatch" + + +def test_rejects_rope_table_narrower_than_head_dim(): + """The rewrite keeps each table's first half, which assumes HF duplicated both halves. + + A table already stored at half width would be silently truncated to a quarter. + """ + assert register(), "nanojet must be installed" + assert _reason(table_width=HEAD_DIM // 2) == "rope-table-width" diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_residual_fold.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_residual_fold.py new file mode 100644 index 000000000000..7c772d5ad664 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_nanojet_residual_fold.py @@ -0,0 +1,138 @@ +"""When may the residual add be folded into the projection's GEMM epilogue? + +The fold writes the projection's result *into* the accumulator. That is only sound when the +accumulator is ours to write and nobody reads it afterwards expecting the pre-add value, so +each condition below is a case where folding would be silently wrong — the fusion must +decline, not produce a plausible answer. +""" + +import pytest +import torch +from torch.fx import Graph + +from tensorrt_llm._torch.auto_deploy.transform.library.fuse_nanojet_gemm_fp8_add import ( + _is_fp8, + _sole_residual_add, +) + + +def _fake(shape, dtype=torch.bfloat16): + return torch.empty(shape, dtype=dtype, device="meta") + + +def _build( + *, + accumulator_is_input=False, + extra_reader_after=False, + extra_reader_before=False, + accumulator_dtype=torch.bfloat16, + accumulator_shape=(4, 8), + two_users=False, +): + """A projection feeding an add, with one condition varied at a time.""" + graph = Graph() + activation = graph.placeholder("activation") + activation.meta["val"] = _fake((4, 8)) + + if accumulator_is_input: + accumulator = graph.placeholder("residual") + else: + accumulator = graph.call_function(torch.ops.aten.mul.Tensor, args=(activation, 1.0)) + accumulator.meta["val"] = _fake(accumulator_shape, accumulator_dtype) + + if extra_reader_before: + before = graph.call_function(torch.ops.aten.relu.default, args=(accumulator,)) + before.meta["val"] = _fake(accumulator_shape, accumulator_dtype) + + projection = graph.call_function(torch.ops.aten.mm.default, args=(activation, activation)) + projection.meta["val"] = _fake((4, 8)) + + add = graph.call_function(torch.ops.aten.add.Tensor, args=(accumulator, projection)) + add.meta["val"] = _fake((4, 8)) + + if two_users: + graph.call_function(torch.ops.aten.relu.default, args=(projection,)) + if extra_reader_after: + graph.call_function(torch.ops.aten.neg.default, args=(accumulator,)) + + graph.output(add) + order = {n: i for i, n in enumerate(graph.nodes)} + return projection, order + + +def test_folds_the_ordinary_residual(): + """The supported shape must fold, or every rejection below proves nothing.""" + projection, order = _build() + assert _sole_residual_add(projection, order) is not None + + +def test_declines_when_projection_has_another_reader(): + """Folding erases the add; a second consumer of the projection would lose its producer.""" + projection, order = _build(two_users=True) + assert _sole_residual_add(projection, order) is None + + +def test_declines_when_accumulator_is_a_graph_input(): + """A placeholder belongs to the caller — writing through it corrupts their tensor.""" + projection, order = _build(accumulator_is_input=True) + assert _sole_residual_add(projection, order) is None + + +def test_declines_when_accumulator_read_after_the_add(): + """That reader wants the pre-add value; the in-place write would hand it the sum.""" + projection, order = _build(extra_reader_after=True) + assert _sole_residual_add(projection, order) is None + + +def test_allows_accumulator_read_before_the_add(): + """A reader that already ran is fine — this is the norm in a transformer block.""" + projection, order = _build(extra_reader_before=True) + assert _sole_residual_add(projection, order) is not None + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float8_e4m3fn]) +def test_declines_on_non_bf16_accumulator(dtype): + projection, order = _build(accumulator_dtype=dtype) + assert _sole_residual_add(projection, order) is None + + +def test_declines_on_shape_mismatch(): + """The accumulator is written into, not broadcast against.""" + projection, order = _build(accumulator_shape=(1, 8)) + assert _sole_residual_add(projection, order) is None + + +# -------------------------------------------------------------------------------------- +# The fusion decides by the dtype of the activation it reads, nothing else. +# -------------------------------------------------------------------------------------- + + +def test_declines_when_attention_is_another_backend(): + """Flashinfer's attention emits BF16; treating it as e4m3 would feed the wrong dtype. + + This is the mixed-backend case: nanojet fusions on, attention deliberately left to + another backend. o_proj must simply not fold. + """ + graph = Graph() + hidden = graph.placeholder("hidden") + hidden.meta["val"] = _fake((4, 8)) + other_attention = graph.call_function(torch.ops.aten.mm.default, args=(hidden, hidden)) + other_attention.meta["val"] = _fake((4, 8)) # BF16, as any non-nanojet attention emits + assert not _is_fp8(other_attention) + + +def test_declines_on_bf16_even_behind_a_view(): + """Views are transparent; the dtype behind them still decides.""" + graph = Graph() + hidden = graph.placeholder("hidden") + hidden.meta["val"] = _fake((4, 8)) + view = graph.call_function(torch.ops.aten.reshape.default, args=(hidden, [4, 8])) + view.meta["val"] = _fake((4, 8)) + assert not _is_fp8(view) + + +def test_accepts_fp8_producer(): + graph = Graph() + hidden = graph.placeholder("hidden") + hidden.meta["val"] = _fake((4, 8), torch.float8_e4m3fn) + assert _is_fp8(hidden)