feat: nvfp4 attn compile init commit

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:

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,

tensorrt_llm/_torch/compilation/piecewise_optimizer.py

@@ -231,15 +231,15 @@ 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,
         ])):
             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

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"})

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,29 +330,37 @@ 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:
             out_scale = self.o_proj.inv_input_scale
         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)