FA4 Inference

megatron/core/transformer/attention.py

@@ -81,6 +81,13 @@
     except ImportError as e:
         pass
 
+try:
+    from flash_attn.cute import flash_attn_varlen_func as flash_attn4_varlen_func
+
+    HAVE_FA4 = True
+except ImportError:
+    HAVE_FA4 = False
+
 try:
     from flash_mla import flash_mla_with_kvcache, get_mla_metadata
 
@@ -802,7 +809,20 @@ def flash_decode_and_prefill(
                 softmax_scale = self.softmax_scale
             else:
                 softmax_scale = q.shape[-1] ** -0.5
-            if HAVE_FA3:
+            if HAVE_FA4:
+                output_total, _ = flash_attn4_varlen_func(
+                    q,
+                    k,
+                    v,
+                    cu_seqlens_q=cu_seqlens_q,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_k=max_seqlen_k,
+                    seqused_k=seqlens_k,
+                    page_table=block_table,
+                    softmax_scale=softmax_scale,
+                    causal=True,
+                )
+            elif HAVE_FA3:
                 # TODO(ksanthanam): Replace with call to flash_attn_varlen_func once
                 # it accepts block_table
                 output_total = self._flash_attention_3_forward_wrapper(
@@ -863,22 +883,44 @@ def flash_decode_and_prefill(
                     causal=True,
                 )
             else:
-                flash_attn_args = {
-                    "q": q,
-                    "k_cache": k,
-                    "v_cache": v,
-                    "cache_seqlens": seqlens_k,
-                    "causal": True,
-                    "page_table" if HAVE_FA3 else "block_table": block_table,
-                    "num_splits": 0 if not self.batch_invariant_mode else 1,
-                }
-                if HAVE_FA3:
-                    output_total = flash_attn3_with_kvcache(**flash_attn_args)
+                if HAVE_FA4:
+                    if getattr(self, "softmax_scale", None) is not None:
+                        softmax_scale = self.softmax_scale
+                    else:
+                        softmax_scale = q.shape[-1] ** -0.5
+                    # Reshape q from (batch, 1, nheads, headdim) to (total_q, nheads, headdim)
+                    q_varlen = q.squeeze(1)
+                    output_total, _ = flash_attn4_varlen_func(
+                        q_varlen,
+                        k,
+                        v,
+                        cu_seqlens_q=cu_seqlens_q,
+                        max_seqlen_q=1,
+                        max_seqlen_k=max_seqlen_k,
+                        seqused_k=seqlens_k,
+                        page_table=block_table,
+                        softmax_scale=softmax_scale,
+                        causal=True,
+                    )
+                    # Reshape back to (batch, 1, nheads, headdim)
+                    output_total = output_total.unsqueeze(1)
                 else:
-                    assert (
-                        not self.batch_invariant_mode
-                    ), "Batch invariant mode is not supported for flash attention 2"
-                    output_total = flash_attn_with_kvcache(**flash_attn_args)
+                    flash_attn_args = {
+                        "q": q,
+                        "k_cache": k,
+                        "v_cache": v,
+                        "cache_seqlens": seqlens_k,
+                        "causal": True,
+                        "page_table" if HAVE_FA3 else "block_table": block_table,
+                        "num_splits": 0 if not self.batch_invariant_mode else 1,
+                    }
+                    if HAVE_FA3:
+                        output_total = flash_attn3_with_kvcache(**flash_attn_args)
+                    else:
+                        assert (
+                            not self.batch_invariant_mode
+                        ), "Batch invariant mode is not supported for flash attention 2"
+                        output_total = flash_attn_with_kvcache(**flash_attn_args)
         return output_total
 
     def forward(
@@ -932,8 +974,8 @@ def forward(
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
         if inference_context and inference_context.is_dynamic_batching():
-            assert HAVE_FA3 or is_fa_min_version(
-                "2.7.3"
+            assert (
+                HAVE_FA4 or HAVE_FA3 or is_fa_min_version("2.7.3")
             ), "flash attn verion v2.7.3 and above is required for dynamic batching."
 
         # hidden_states: [sq, b, h]

pyproject.toml

@@ -103,6 +103,7 @@ dev = [
     "megatron-energon[av_decode]~=6.0",
     "av",
     "flashinfer-python~=0.5.0",
+    "flash-attn-4",
     "wget",
     "onnxscript",
     "fastapi~=0.50",                                    # Forcing a little bit more recent version of fastapi to be compatible with pydantic 2.0
@@ -206,6 +207,7 @@ transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git",
 nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "17ae86b64d7f75653351664f5d8c9e466faede00" }
 emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.2.0" }
 nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", rev = "63154570cea17f8805a7fd15cc3b8cc2919ba575" }
+flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", tag = "fa4-v4.0.0.beta7", subdirectory = "flash_attn/cute" }
 
 [tool.isort]
 profile = "black"                                                          # black-compatible

tests/unit_tests/inference/engines/test_dynamic_engine.py

@@ -306,7 +306,7 @@ def _build_test_env(cls, test_config):
                 params_dtype=torch.bfloat16,
                 num_layers=4,
                 mtp_num_layers=test_config.num_speculative_tokens,
-                hidden_size=128 if test_config.fp8 else 32,
+                hidden_size=256 if test_config.fp8 else 256,
                 num_attention_heads=4,
                 use_cpu_initialization=True,
                 cuda_graph_impl=(
@@ -602,14 +602,14 @@ def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None
 
         # Validate generated tokens.
         gpt_expected_generated_tokens = [
-            [69, 85, 55, 74, 56, 89, 64, 59, 55, 67, 15, 58, 6, 37, 54, 47],
-            [29, 54, 33, 72, 45, 76, 41, 56, 28, 25, 17, 2, 61, 6, 98, 76],
-            [35, 78, 54, 16, 79, 98, 22, 5, 60, 0, 1, 76, 77, 11, 25, 7],
-            [25, 75, 57, 85, 81, 37, 88, 17, 71, 15, 70, 64, 50, 0, 64, 45],
-            [32, 5, 85, 75, 30, 68, 23, 33, 20, 26, 89, 20, 49, 28, 38, 81],
-            [33, 69, 32, 49, 93, 24, 33, 6, 54, 89, 92, 97, 42, 80, 50, 53],
-            [82, 78, 78, 65, 26, 5, 69, 36, 37, 99],
-            [51, 70, 22, 1, 87, 42, 36, 26, 27, 56, 82, 32, 8, 80, 20, 43],
+            [69, 85, 55, 74, 56, 89, 64, 59, 55, 67, 15, 58, 6, 37, 34, 47],
+            [29, 16, 33, 93, 45, 76, 41, 74, 28, 17, 36, 2, 61, 6, 98, 76],
+            [35, 78, 54, 16, 79, 98, 22, 5, 43, 0, 1, 76, 77, 11, 25, 66],
+            [25, 75, 57, 85, 81, 37, 88, 17, 71, 15, 70, 64, 40, 0, 64, 45],
+            [32, 5, 85, 75, 30, 68, 23, 16, 20, 26, 89, 20, 49, 28, 34, 81],
+            [62, 69, 32, 49, 93, 24, 33, 6, 54, 89, 92, 74, 42, 80, 50, 53],
+            [82, 78, 78, 65, 26, 5, 97, 36, 26, 99],
+            [51, 70, 22, 1, 87, 42, 36, 26, 27, 67, 97, 32, 8, 80, 20, 43],
         ]
 
         mamba_expected_generated_tokens = [
@@ -3039,7 +3039,7 @@ def _create_model(self, model_provider, num_cuda_graphs):
             config = TransformerConfig(
                 params_dtype=torch.bfloat16,
                 num_layers=4,
-                hidden_size=32,
+                hidden_size=256,
                 num_attention_heads=4,
                 use_cpu_initialization=True,
                 cuda_graph_impl=cuda_graph_impl,