Enable Fused QKV ROPE + CS FP8 quantization fusion

megatron/core/transformer/attention.py

@@ -131,8 +131,12 @@
     from transformer_engine.pytorch.attention.rope import apply_fused_qkv_rotary_pos_emb
 
     HAVE_FUSED_QKV_ROPE = True
+    HAVE_FUSED_QKV_ROPE_FP8_ATTENTION = (
+        "defer_for_fp8_attention" in inspect.signature(apply_fused_qkv_rotary_pos_emb).parameters
+    )
 except ImportError:
     HAVE_FUSED_QKV_ROPE = False
+    HAVE_FUSED_QKV_ROPE_FP8_ATTENTION = False
 
 
 class LinearQkvInterface(Protocol):
@@ -1605,9 +1609,30 @@ def forward_pre_attn_and_core_attn(
                         max_seqlen=rope_freqs_max_seqlen,
                     )
             else:
-                query, key, value = apply_fused_qkv_rotary_pos_emb(
-                    mixed_qkv, q_pos_emb, k_pos_emb, qkv_split_arg_list
+                # The FP8 attention fusion hands TE placeholders that its core attention fills
+                # (or fuses with the FP8 quantization), so they must not be offloaded or
+                # recomputed from.
+                defer_for_fp8_attention = (
+                    self.config.fused_qkv_rope_fp8_attention
+                    and not (self.checkpoint_core_attention and self.training)
+                    and not (self.offload_core_attention and self.training)
                 )
+                if defer_for_fp8_attention:
+                    assert HAVE_FUSED_QKV_ROPE_FP8_ATTENTION, (
+                        "fused_qkv_rope_fp8_attention requires a TransformerEngine version whose "
+                        "apply_fused_qkv_rotary_pos_emb supports defer_for_fp8_attention."
+                    )
+                    query, key, value = apply_fused_qkv_rotary_pos_emb(
+                        mixed_qkv,
+                        q_pos_emb,
+                        k_pos_emb,
+                        qkv_split_arg_list,
+                        defer_for_fp8_attention=True,
+                    )
+                else:
+                    query, key, value = apply_fused_qkv_rotary_pos_emb(
+                        mixed_qkv, q_pos_emb, k_pos_emb, qkv_split_arg_list
+                    )
 
             # TODO, can apply positional embedding to value_layer so it has
             # absolute positional embedding.

megatron/core/transformer/transformer_config.py

@@ -591,6 +591,12 @@ class TransformerConfig(ModelParallelConfig):
     fused_single_qkv_rope: bool = False
     """If set, avoid splitting QKV before ROPE forward and avoid concatenating ROPE dgrads."""
 
+    fused_qkv_rope_fp8_attention: bool = False
+    """If set, let TransformerEngine fuse the RoPE of the single QKV tensor with the FP8
+    current-scaling quantization of Q/K/V inside core attention (falls back to the regular fused
+    RoPE whenever attention does not run FP8 current-scaling fused attention). Requires
+    fused_single_qkv_rope and is disabled for checkpointed or offloaded core attention."""
+
     fused_residual_rmsnorm: bool = False
     """If True, fuses residual connection and RMSNorm backward pass when TE is used."""
 
@@ -3104,6 +3110,9 @@ def __post_init__(self):
             if self.attention_output_gate:
                 raise ValueError("fused_single_qkv_rope does not support gated attention for now.")
 
+        if self.fused_qkv_rope_fp8_attention and not self.fused_single_qkv_rope:
+            raise ValueError("fused_qkv_rope_fp8_attention requires fused_single_qkv_rope.")
+
         if self.multi_latent_attention and self.rotary_interleaved:
             raise ValueError("rotary_interleaved does not work with multi_latent_attention.")