Plumb FP8+THD

tests/pytorch/attention/test_attention.py

@@ -2043,8 +2043,8 @@ def get_model(dtype, config):
 }
 
 param_types_fp8_vs_f16 = [torch.float16, torch.bfloat16]
-qkv_layout_fp8_vs_f16 = ["sbh3d", "bshd_bshd_bshd", "sbhd_sbhd_sbhd"]
-qkv_format_fp8_vs_f16 = ["bshd", "sbhd"]
+qkv_layout_fp8_vs_f16 = ["sbh3d", "bshd_bshd_bshd", "sbhd_sbhd_sbhd", "thd_thd_thd"]
+qkv_format_fp8_vs_f16 = ["bshd", "sbhd", "thd"]
 
 
 @pytest.mark.skipif(get_cudnn_version() < (9, 2, 1), reason="cuDNN 9.2.1+ is required.")
@@ -2249,6 +2249,10 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker:
         seqlens_kv = torch.full(
             [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda"
         )
+    if qkv_format == "thd":
+        # FP8 Linear flattens THD input to [t, h*d], so align total tokens for cuBLAS.
+        seqlens_q[-1] += -seqlens_q.sum() % 8
+        seqlens_kv[-1] += -seqlens_kv.sum() % 8
     cu_seqlens_q = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda")
     cu_seqlens_kv = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda")
     cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0)
@@ -2287,6 +2291,8 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker:
             rotary_pos_emb=rotary_pos_emb,
             cu_seqlens_q=cu_seqlens_q,
             cu_seqlens_kv=cu_seqlens_kv,
+            # The optimized zero-fill path dereferences device memory on the host.
+            fast_zero_fill=False,
         )
     if is_training:
         out.backward(out_grad)
@@ -2620,6 +2626,9 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker:
             attn_mask_type=config.attn_mask_type,
             checkpoint_core_attention=False,
             core_attention_bias_type=config.attn_bias_type,
+            fp8_output=fp8_dpa,
+            # The optimized zero-fill path dereferences device memory on the host.
+            fast_zero_fill=False,
         )
     if is_training:
         out.backward(out_grad)

tests/pytorch/attention/test_attention_with_cp.py

@@ -407,6 +407,8 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type
         2, 4096, 12, 128, attn_bias_type="post_scale_bias", bias_shape="bhss"
     ),  # MHA
     "cp_1_5": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal", window_size=(512, 512)),  # MHA
+    # Noncausal MHA without bias/max-logit provides an FP8+THD+CP backend-compatible row.
+    "cp_1_6": ModelConfig(2, 4096, 12, 128),
     "cp_2_0": ModelConfig(
         2,
         4096,
@@ -489,6 +491,7 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type
 if test_essential:
     configs = [
         "cp_1_0",
+        "cp_1_6",
         "cp_2_0",
         "cp_2_1",
         "cp_2_2",
@@ -548,8 +551,6 @@ def test_cp_with_fused_attention(
     if dtype != "fp8" and (fp8_mha or fp8_dpa):
         pytest.skip("dtype!=fp8 requires fp8_dpa=False and fp8_mha=False!")
 
-    if dtype == "fp8" and qkv_format == "thd":
-        pytest.skip("No support for FP8 attention with THD format!")
     if dtype == "fp8" and config.attn_bias_type != "no_bias":
         pytest.skip("No support for FP8 attention with bias!")
 
@@ -597,6 +598,8 @@ def test_cp_with_fused_attention(
         pytest.skip("scaling_mode=delayed requires f16_O=False!")
     if scaling_mode == "mxfp8" and not f16_O:
         pytest.skip("scaling_mode=mxfp8 requires f16_O=True!")
+    if scaling_mode == "mxfp8" and qkv_format == "thd":
+        pytest.skip("MXFP8 quantization does not support THD format!")
     if scaling_mode == "mxfp8" and fp8_mha:
         pytest.skip("No support for scaling_mode=mxfp8 with fp8_mha=True!")
 

transformer_engine/common/fused_attn/fused_attn.cpp

@@ -281,13 +281,17 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend(
             attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK))) &&
       // pre-9.21: {bshd, sbhd}, {vanilla}
       // 9.21+: {bshd, sbhd, bhsd}, {vanilla, off-by-one, learnable}
+      // 9.23+ sm100+: {thd} (ragged/packed variable-length)
       ((cudnn_runtime_version < 92100 &&
         (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) &&
         softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) ||
        (cudnn_runtime_version >= 92100 &&
         (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD ||
-         qkv_format == NVTE_QKV_Format::NVTE_BHSD))) &&
-      !requires_64bit_ragged_offset &&
+         qkv_format == NVTE_QKV_Format::NVTE_BHSD)) ||
+       (cudnn_runtime_version >= 92300 && sm_arch_ >= 100 &&
+        qkv_format == NVTE_QKV_Format::NVTE_THD && supported_ragged_offset_size &&
+        (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK ||
+         attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) &&
       // 9.10.0: known bugs with SDPA FP8
       (cudnn_runtime_version != 91000) && !return_max_logit) {
     backend = NVTE_Fused_Attn_Backend::NVTE_FP8;
@@ -627,12 +631,13 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso
         input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded,
         input_page_table_k, input_page_table_v, input_rng_state, wkspace, stream, handle);
   } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) {
-    fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, is_training,
+    fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, is_training,
                        attn_scale, dropout, qkv_layout, o_format, qkv_scale_inv_format, bias_type,
                        attn_mask_type, softmax_type, window_size_left, window_size_right,
                        bottom_right_diagonal, input_Q, input_K, input_V, input_SoftmaxOffset,
                        input_output_S, output_O, Aux_CTX_Tensors, input_cu_seqlens_q,
-                       input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle);
+                       input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded,
+                       input_rng_state, wkspace, stream, handle);
   } else {
     NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n");
   }
@@ -729,14 +734,15 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso
     if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) {
       input_dO_f16 = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]);
     }
-    fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, attn_scale, dropout,
-                       qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format,
+    fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, attn_scale,
+                       dropout, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format,
                        do_scale_inv_format, bias_type, attn_mask_type, softmax_type,
                        window_size_left, window_size_right, bottom_right_diagonal, deterministic,
                        input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, input_S,
                        input_SoftmaxOffset, input_output_dP, output_dQ, output_dK, output_dV,
                        output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv,
-                       input_rng_state, wkspace, stream, handle);
+                       input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state,
+                       wkspace, stream, handle);
   } else {
     NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n");
   }