diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 531dada135d6..e65c1191382b 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -560,6 +560,9 @@ def support_mla(cls) -> bool: def support_nvfp4_output(cls) -> bool: return False + def create_output(self, *args, **kwargs): + raise NotImplementedError + @dataclass(kw_only=True, unsafe_hash=True) class MLAParams: diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 6c326b2c0a2c..f01bbddcc7df 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -915,6 +915,7 @@ def __init__( self.kv_scale_orig_quant = 1.0 / self.kv_scale_quant_orig if not skip_create_weights_in_init: self.update_quant_config(self.quant_config) + self.use_nvfp4_output = None def update_quant_config(self, new_quant_config: Optional[QuantConfig]): self.quant_config = new_quant_config @@ -936,6 +937,60 @@ def get_local_layer_idx(self, metadata: TrtllmAttentionMetadata) -> int: else: return metadata.kv_cache_manager.layer_offsets[self.layer_idx] + def get_out_dtype( + self, + out_scale: Optional[torch.Tensor] = None) -> Optional[torch.dtype]: + out_dtype = None + if out_scale is not None: + if self.use_nvfp4_output: + # Use UINT8 as the container dtype for NVFP4. + out_dtype = torch.uint8 + elif (self.has_fp8_qdq or self.has_nvfp4 + or self.has_fp8_block_wise) and self.has_fp8_kv_cache: + # TODO(qijun): revisit fp8_context_fmha logic + out_dtype = torch.float8_e4m3fn + return out_dtype + + def create_output(self, + q: torch.Tensor, + out_scale: Optional[torch.Tensor] = None, + **kwargs): + num_tokens = q.size(0) + out_dtype = self.get_out_dtype(out_scale) + if out_dtype is None: + out_dtype = q.dtype + v_head_size = self.head_dim + if out_dtype == torch.uint8: + num_nvfp4_elements_per_container = 2 + scaling_vector_size = 16 + size_per_token = self.num_heads * v_head_size + output = q.new_empty( + (num_tokens, + size_per_token // num_nvfp4_elements_per_container), + dtype=torch.uint8) + # Create a sf (scaling factors) tensor for NVFP4 (use INT8 as the container dtype). + output_sf = q.new_empty(compute_swizzled_sf_shape( + num_tokens, size_per_token // scaling_vector_size), + dtype=torch.uint8) + else: + output = q.new_empty((num_tokens, self.num_heads * v_head_size), + dtype=out_dtype) + # Seems torch does not support None for mutable tensors, so we need to create an empty tensor. + output_sf = q.new_empty(()) + return output, output_sf + + def use_paged_context_fmha(self, metadata: TrtllmAttentionMetadata) -> bool: + use_paged_context_fmha = ( + metadata.runtime_features.chunked_prefill + or metadata.runtime_features.cache_reuse + or metadata.runtime_features.has_speculative_draft_tokens + ) if metadata.runtime_features else False + + if self.is_mla_enable: + use_paged_context_fmha = use_paged_context_fmha and self.has_cached_kv_for_mla_context( + metadata) + return use_paged_context_fmha + def forward( self, q: torch.Tensor, @@ -964,27 +1019,6 @@ def forward( ) assert not metadata.is_cross, "TRT-LLM Attention does not support cross attention yet." - use_paged_context_fmha = ( - metadata.runtime_features.chunked_prefill - or metadata.runtime_features.cache_reuse - or metadata.runtime_features.has_speculative_draft_tokens - ) if metadata.runtime_features else False - - if self.is_mla_enable: - # for MLA, we only use paged_context_fmha when there is cached kv - use_paged_context_fmha = use_paged_context_fmha and self.has_cached_kv_for_mla_context( - metadata) - - use_nvfp4_output = False - if self.has_nvfp4 and self.support_nvfp4_output(): - # Runtime check whether the NVFP4 output kernel is available. - use_nvfp4_output = self.wrapper.is_nvfp4_output_kernel_available( - tokens_per_block=metadata.tokens_per_block, - attention_mask=attention_mask, - use_paged_context_fmha=use_paged_context_fmha, - is_mla_enable=self.is_mla_enable, - ) - self.wrapper.plan( layer_idx=self.get_local_layer_idx(metadata), tokens_per_block=metadata.tokens_per_block, @@ -1011,8 +1045,8 @@ def forward( kv_scale_quant_orig=self.kv_scale_quant_orig, out_scale=out_scale, out_scale_sf=out_scale_sf, - use_nvfp4_output=use_nvfp4_output, - use_paged_context_fmha=use_paged_context_fmha, + use_nvfp4_output=self.use_nvfp4_output, + use_paged_context_fmha=self.use_paged_context_fmha(metadata), attention_input_type=attention_input_type, latent_cache=latent_cache, q_pe=q_pe, @@ -1022,15 +1056,7 @@ def forward( mla_context_kv_cache_block_offsets, softmax_stats_tensor=softmax_stats_tensor, ) - out_dtype = None - if out_scale is not None: - if use_nvfp4_output: - # Use UINT8 as the container dtype for NVFP4. - out_dtype = torch.uint8 - elif (self.has_fp8_qdq or self.has_nvfp4 - or self.has_fp8_block_wise) and self.has_fp8_kv_cache: - # TODO(qijun): revisit fp8_context_fmha logic - out_dtype = torch.float8_e4m3fn + out_dtype = self.get_out_dtype(out_scale) output, output_sf = self.wrapper.run( q, diff --git a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py index 75a9aeff8e5c..9560a8994d89 100644 --- a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py +++ b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py @@ -231,7 +231,7 @@ def piecewise_optimizer( if node.op in ("output", "placeholder"): continue if (not stop_partition and is_call_function(node, [ - torch.ops.trtllm.attention_inplace.default, + torch.ops.trtllm.attn_custom_op_inplace.default, torch.ops.trtllm.mla_custom_op_inplace.default, torch.ops.aten.index.Tensor, torch.ops.aten.cumsum.default, @@ -239,7 +239,7 @@ def piecewise_optimizer( idx += 1 node_to_graph_id[node] = idx exclude_modules_id.append(idx) - if node.target != torch.ops.trtllm.attention_inplace.default and node.target != torch.ops.trtllm.mla_custom_op_inplace.default: + if node.target != torch.ops.trtllm.attn_custom_op_inplace.default and node.target != torch.ops.trtllm.mla_custom_op_inplace.default: # We only know it is safe to continue splitting after attention # since attention_inplace will not produce any new tensor stop_partition = True diff --git a/tensorrt_llm/_torch/compilation/remove_copy_pass.py b/tensorrt_llm/_torch/compilation/remove_copy_pass.py index fe968f020be0..c5278c5ffc43 100644 --- a/tensorrt_llm/_torch/compilation/remove_copy_pass.py +++ b/tensorrt_llm/_torch/compilation/remove_copy_pass.py @@ -55,7 +55,7 @@ def remove_functionalize_inner(node: Node, mutates_args: dict, is_v2=False): }, is_v2=node.target == auto_functionalized_v2, ) - if inplace_func == torch.ops.trtllm.attention_inplace.default: + if inplace_func == torch.ops.trtllm.attn_custom_op_inplace.default: remove_functionalize_inner(node, {1: "output", 2: "output_sf"}) if inplace_func == torch.ops.trtllm.mla_custom_op_inplace.default: remove_functionalize_inner(node, {1: "output"}) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 09671696aa9e..685130fd1fa8 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -24,6 +24,33 @@ from .rotary_embedding import RotaryEmbedding +@torch.library.custom_op("trtllm::attn_custom_op_inplace", + mutates_args=("output", "output_sf")) +def attn_custom_op_inplace( + qkv: torch.Tensor, + position_ids: Optional[torch.Tensor], + layer_idx: str, + attention_mask: str, + out_scale: Optional[torch.Tensor] = None, + out_scale_sf: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + output_sf: Optional[torch.Tensor] = None, + attention_window_size: Optional[int] = None, +) -> None: + metadata, mrope_config, attn_layer = extract_extra_attrs( + layer_idx, "attention") + attn_layer.forward_impl(qkv, + position_ids, + metadata, + PredefinedAttentionMask(attention_mask), + mrope_config, + attention_window_size, + out_scale, + out_scale_sf, + output=output, + output_sf=output_sf) + + class Attention(nn.Module): def __init__( @@ -63,6 +90,7 @@ def __init__( """ super().__init__() self.layer_idx = layer_idx + self.layer_idx_str = str(layer_idx) config = config or ModelConfig() self.hidden_size = hidden_size @@ -92,6 +120,14 @@ def __init__( # 0 0 0 1 1 1 self.attention_chunk_size = attention_chunk_size + self.register_to_config = False + if config is not None: + if "attn_layers" not in config.extra_attrs: + config.extra_attrs["attn_layers"] = {} + config.extra_attrs["attn_layers"][self.layer_idx_str] = weakref.ref( + self) + self.register_to_config = True + if dense_bias is None: self.dense_bias = bias @@ -218,6 +254,40 @@ def convert_qkv(self, q, k, v): q, k, v = qkv, None, None return q, k, v + def forward_impl( + self, + q: torch.Tensor, + position_ids: Optional[torch.IntTensor], + attn_metadata: AttentionMetadata, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask. + CAUSAL, + mrope_config: Optional[dict] = None, + attention_window_size: Optional[int] = None, + out_scale: Optional[torch.Tensor] = None, + out_scale_sf: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + output_sf: Optional[torch.Tensor] = None, + **kwargs): + + q, k, v = q, None, None + + q, k, v = self.apply_rope(q, k, v, position_ids) + + q, k, v = self.convert_qkv(q, k, v) + self.attn.forward( + q, + k, + v, + attn_metadata, + out_scale=out_scale, + out_scale_sf=out_scale_sf, + attention_mask=attention_mask, + mrope_config=mrope_config, + attention_window_size=attention_window_size, + output=output, + output_sf=output_sf, + ) + def forward( self, position_ids: Optional[torch.IntTensor], @@ -260,10 +330,6 @@ def forward( if qkv_lora is not None: qkv = qkv + qkv_lora - q, k, v = qkv, None, None - - q, k, v = self.apply_rope(q, k, v, position_ids) - out_scale = None out_scale_sf = None if self.o_proj.has_fp8_qdq or self.o_proj.has_nvfp4 or self.o_proj.has_fp8_block_scales: @@ -271,18 +337,30 @@ def forward( if self.o_proj.has_nvfp4 and self.support_nvfp4_output: out_scale_sf = self.o_proj.input_scale - q, k, v = self.convert_qkv(q, k, v) - attn_output = self.attn.forward( - q, - k, - v, - attn_metadata, - out_scale=out_scale, - out_scale_sf=out_scale_sf, - attention_mask=attention_mask, - mrope_config=mrope_config, - attention_window_size=attention_window_size) - hidden_states = attn_output + output, output_sf = self.attn.create_output(q=qkv, out_scale=out_scale) + + if self.register_to_config: + torch.ops.trtllm.attn_custom_op_inplace( + qkv, + position_ids, + self.layer_idx_str, + attention_mask.value, + out_scale, + out_scale_sf, + output, + output_sf, + attention_window_size, + ) + else: + self.forward_impl(qkv, position_ids, attn_metadata, attention_mask, + mrope_config, attention_window_size, out_scale, + out_scale_sf) + + if output_sf is not None: + attn_output = Fp4QuantizedTensor(output, output_sf) + else: + attn_output = output + attn_output = self.o_proj(attn_output, all_reduce_params=all_reduce_params, lora_params=lora_params, @@ -311,28 +389,43 @@ def apply_rope(self, q: torch.Tensor, k: Optional[torch.Tensor], return q, k, v -def extract_extra_attrs(layer_idx: str): +def extract_extra_attrs(layer_idx: str, type: str): extra_attrs = get_model_extra_attrs() assert extra_attrs is not None, "Model extra attrs is not set" metadata_ref = extra_attrs.get("attention_metadata", None) assert metadata_ref is not None, "Attention metadata is not set" metadata = metadata_ref() + mrope_config = extra_attrs.get("mrope_config", None) assert isinstance( metadata, TrtllmAttentionMetadata, ) - mla_layers = extra_attrs.get("mla_layers", None) - assert mla_layers is not None, "MLA layers is not registered" - mla_layer_ref = mla_layers.get(layer_idx, None) - assert mla_layer_ref is not None, f"Cannot find MLA layer for layer {layer_idx}" - mla_layer = mla_layer_ref() - assert isinstance( - mla_layer, - MLA), "MLA layer must be a subclass of MLA or an instance of MLA" + layer = None + + if type == "attention": + attn_layers = extra_attrs.get("attn_layers", None) + assert attn_layers is not None, "Attention layers is not set" + attn_layer_ref = attn_layers.get(layer_idx, None) + assert attn_layer_ref is not None, f"Cannot find Attention layer for layer {layer_idx}" + layer = attn_layer_ref() + assert isinstance( + layer, Attention + ), "Attention layer must be a subclass of Attention or an instance of Attention" + elif type == "mla": + mla_layers = extra_attrs.get("mla_layers", None) + assert mla_layers is not None, "MLA layers is not set" + mla_layer_ref = mla_layers.get(layer_idx, None) + assert mla_layer_ref is not None, f"Cannot find MLA layer for layer {layer_idx}" + layer = mla_layer_ref() + assert isinstance( + layer, + MLA), "MLA layer must be a subclass of MLA or an instance of MLA" + else: + raise ValueError(f"Invalid type: {type}") - return metadata, mla_layer + return metadata, mrope_config, layer @torch.library.custom_op("trtllm::mla_custom_op_inplace", @@ -343,7 +436,7 @@ def mla_custom_op_inplace( layer_idx: str, output: torch.Tensor, ) -> None: - metadata, mla_layer = extract_extra_attrs(layer_idx) + metadata, mla_layer = extract_extra_attrs(layer_idx, "mla") mla_layer.forward_impl(position_ids, hidden_states, metadata, output=output) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index aa0484867c1d..2e9d1d18538b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -34,8 +34,9 @@ from tensorrt_llm.quantization.utils.fp4_utils import float4_e2m1x2 from ..attention_backend.interface import (AttentionMetadata, - AttentionRuntimeFeatures) -from ..attention_backend.trtllm import TrtllmAttentionMetadata + AttentionRuntimeFeatures, + PredefinedAttentionMask) +from ..attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata from ..attention_backend.utils import get_attention_backend from ..attention_backend.vanilla import VanillaAttentionMetadata from ..autotuner import AutoTuner, autotune @@ -2099,6 +2100,34 @@ def model_forward(self, **kwargs): attrs = get_model_extra_attrs() assert attrs is not None, "Model extra attrs is not set" attrs["attention_metadata"] = weakref.ref(kwargs['attn_metadata']) + if 'mrope_config' in kwargs: + attrs["mrope_config"] = kwargs['mrope_config'] + if self.pytorch_backend_config.attn_backend == "TRTLLM": + # Ugly hack to support trtllm attention backend nvfp4 output + for attn_layer in self.model.model_config.extra_attrs[ + 'attn_layers'].values(): + if isinstance(attn_layer, weakref.ReferenceType): + attn_layer = attn_layer() + attn_backend = attn_layer.attn + assert isinstance( + attn_backend, + TrtllmAttention), "attn_backend is not a TrtllmAttention" + metadata = kwargs['attn_metadata'] + assert isinstance(metadata, TrtllmAttentionMetadata + ), "metadata is not a TrtllmAttentionMetadata" + + attn_backend.use_nvfp4_output = False + if attn_backend.has_nvfp4 and attn_backend.support_nvfp4_output( + ): + attn_backend.use_nvfp4_output = attn_backend.wrapper.is_nvfp4_output_kernel_available( + tokens_per_block=metadata.tokens_per_block, + attention_mask=kwargs.get( + 'attention_mask', PredefinedAttentionMask.CAUSAL), + use_paged_context_fmha=attn_backend. + use_paged_context_fmha(metadata), + is_mla_enable=attn_backend.is_mla_enable, + ) + attrs.update(self.model.model_config.extra_attrs) if is_trace_enabled("TLLM_TRACE_MODEL_FORWARD"): diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 9c3197e7c93c..1f3d48a56904 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -389,6 +389,7 @@ def torch_dtype_to_binding(dtype): torch.int64: "