Plumb FP8+THD

tests/pytorch/attention/test_attention.py

@@ -2096,8 +2096,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.")
@@ -2302,6 +2302,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)
@@ -2340,6 +2344,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)
@@ -2673,6 +2679,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

@@ -402,6 +402,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,
@@ -486,6 +488,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",
@@ -530,6 +533,20 @@ def test_cp_with_fused_attention(
     config.context_parallel = True
     config.cp_comm_type = cp_comm_type
 
+    hopper_fp8_thd_forward = (
+        get_device_compute_capability() == (9, 0)
+        and dtype == "fp8"
+        and model == "cp_1_6"
+        and qkv_format == "thd"
+        and not fp8_bwd
+        and fp8_dpa
+        and not fp8_mha
+        and scaling_mode == "delayed"
+        and not f16_O
+    )
+    if hopper_fp8_thd_forward and get_cudnn_version() < (9, 25, 0):
+        pytest.skip("FP8+THD inference on Hopper requires cuDNN 9.25+.")
+
     if config.head_dim_qk == 256 and config.head_dim_v == 256:
         # D=256 uses this generic CP runner, but only a subset of its axes is supported.
         if get_device_compute_capability() not in ((10, 0), (10, 3)):
@@ -564,8 +581,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!")
 
@@ -607,6 +622,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!")
 
@@ -638,8 +655,9 @@ def test_cp_with_fused_attention(
             MXFP8BlockScaling(fp8_format=Format.E4M3, fp8_dpa=True),
         ]
 
-    # For 111s, dbias calculation is not supported as of cuDNN 9.18, hence, test fwd only for 111s.
-    is_training = False if config.bias_shape == "111s" else True
+    # 111s runs forward-only because its dbias is unsupported. Reuse otherwise-skipped Hopper
+    # FP8+THD nodes for forward-only coverage of the supported cuDNN path.
+    is_training = False if config.bias_shape == "111s" or hopper_fp8_thd_forward else True
     available_backends, _, fused_attn_backends = get_available_attention_backends(
         config,
         qkv_dtype=dtypes[dtype] if dtype != "fp8" else torch.float8_e4m3fn,

tests/pytorch/attention/test_cp_utils.py

@@ -881,6 +881,22 @@ def test_thd_read_half_tensor_reads_each_sequence_half(self):
             torch.equal(kv_second, torch.stack([expected_second, expected_second + 128]))
         )
 
+    def test_thd_grad_correction_copies_byte_half_and_zeros_inactive_half(self):
+        cu_seqlens = torch.tensor([0, 8, 20], dtype=torch.int32, device="cuda")
+        grad_per_step = torch.arange(10 * 2 * 8, dtype=torch.uint8, device="cuda").view(10, 2, 8)
+        first_half_rows = torch.tensor([0, 1, 2, 3, 8, 9, 10, 11, 12, 13], device="cuda")
+        second_half_rows = torch.tensor([4, 5, 6, 7, 14, 15, 16, 17, 18, 19], device="cuda")
+
+        grad = torch.full((20, 2, 8), 255, dtype=torch.uint8, device="cuda")
+        tex.thd_grad_correction(grad, grad_per_step, cu_seqlens, "copy", "zero")
+        self.assertTrue(torch.equal(grad[first_half_rows], grad_per_step))
+        self.assertEqual(torch.count_nonzero(grad[second_half_rows]).item(), 0)
+
+        grad.fill_(255)
+        tex.thd_grad_correction(grad, grad_per_step, cu_seqlens, "zero", "copy")
+        self.assertEqual(torch.count_nonzero(grad[first_half_rows]).item(), 0)
+        self.assertTrue(torch.equal(grad[second_half_rows], grad_per_step))
+
     def test_thd_read_second_half_lse_handles_packed_and_batch_major_lse(self):
         cu_seqlens = torch.tensor([0, 8, 16], dtype=torch.int32, device="cuda")
         lse = torch.arange(2 * 2 * 8, dtype=torch.float32, device="cuda").view(2, 2, 8)

transformer_engine/common/fused_attn/context_parallel.cu

@@ -42,6 +42,12 @@ struct CopyFunctor {
   }
 };
 
+struct ZeroFunctor {
+  __forceinline__ __device__ static void run(void *token, void *token_per_step, int idx) {
+    reinterpret_cast<float4 *>(token)[idx] = make_float4(0.f, 0.f, 0.f, 0.f);
+  }
+};
+
 template <typename dtype>
 struct AddFunctor {
   __forceinline__ __device__ static void run(dtype *token, dtype *token_per_step, int idx) {
@@ -357,24 +363,27 @@ __global__ void thd_grad_correction_kernel(dtype *grad, dtype *grad_per_step, in
   for (int token_id = group_id; token_id < num_total_tokens; token_id += num_groups) {
     int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1);
 
-    int token_offset;
-    bool is_first_half;
     if constexpr (functor_idx < 2) {
-      token_offset = cu_seqlens_s[seq_id + functor_idx];
-      is_first_half = (functor_idx == 0);
+      dtype *first_half_token =
+          &grad[(token_id + cu_seqlens_s[seq_id]) * static_cast<size_t>(hidden_size)];
+      dtype *second_half_token =
+          &grad[(token_id + cu_seqlens_s[seq_id + 1]) * static_cast<size_t>(hidden_size)];
+      dtype *token_per_step = &grad_per_step[token_id * static_cast<size_t>(hidden_size)];
+      for (int idx = lane_id; idx < num_inner_loops; idx += group_size) {
+        Functor_0::run(first_half_token, token_per_step, idx);
+        Functor_1::run(second_half_token, token_per_step, idx);
+      }
     } else {
-      token_offset = 0;
       int len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id];
-      is_first_half = (token_id - cu_seqlens_s[seq_id]) < (len / 2);
-    }
-
-    dtype *token = &grad[(token_id + token_offset) * static_cast<size_t>(hidden_size)];
-    dtype *token_per_step = &grad_per_step[token_id * static_cast<size_t>(hidden_size)];
-    for (int idx = lane_id; idx < num_inner_loops; idx += group_size) {
-      if (is_first_half) {
-        Functor_0::run(token, token_per_step, idx);
-      } else {
-        Functor_1::run(token, token_per_step, idx);
+      bool is_first_half = (token_id - cu_seqlens_s[seq_id]) < (len / 2);
+      dtype *token = &grad[token_id * static_cast<size_t>(hidden_size)];
+      dtype *token_per_step = &grad_per_step[token_id * static_cast<size_t>(hidden_size)];
+      for (int idx = lane_id; idx < num_inner_loops; idx += group_size) {
+        if (is_first_half) {
+          Functor_0::run(token, token_per_step, idx);
+        } else {
+          Functor_1::run(token, token_per_step, idx);
+        }
       }
     }
   }
@@ -707,6 +716,12 @@ static void thd_grad_dispatcher(Tensor grad, const Tensor &grad_per_step, const
   } else if (first_half == "none" && second_half == "copy") {
     thd_grad_correction_helper<dtype, EmptyFunctor, CopyFunctor, 1>(grad, grad_per_step, cu_seqlens,
                                                                     stream);
+  } else if (first_half == "copy" && second_half == "zero") {
+    thd_grad_correction_helper<dtype, CopyFunctor, ZeroFunctor, 0>(grad, grad_per_step, cu_seqlens,
+                                                                   stream);
+  } else if (first_half == "zero" && second_half == "copy") {
+    thd_grad_correction_helper<dtype, ZeroFunctor, CopyFunctor, 1>(grad, grad_per_step, cu_seqlens,
+                                                                   stream);
   } else if (first_half == "add" && second_half == "copy") {
     thd_grad_correction_helper<dtype, AddFunctor<dtype>, CopyFunctor, 2>(grad, grad_per_step,
                                                                          cu_seqlens, stream);
@@ -722,6 +737,18 @@ void thd_grad_correction(Tensor grad, const Tensor &grad_per_step, const Tensor
                          const std::string &first_half, const std::string &second_half,
                          cudaStream_t stream) {
   using namespace transformer_engine;
+  if (grad.dtype() == DType::kByte) {
+    if (first_half == "copy" && second_half == "zero") {
+      thd_grad_correction_helper<byte, CopyFunctor, ZeroFunctor, 0>(grad, grad_per_step, cu_seqlens,
+                                                                    stream);
+    } else if (first_half == "zero" && second_half == "copy") {
+      thd_grad_correction_helper<byte, ZeroFunctor, CopyFunctor, 1>(grad, grad_per_step, cu_seqlens,
+                                                                    stream);
+    } else {
+      NVTE_ERROR("Byte gradients require copy/zero or zero/copy correction\n");
+    }
+    return;
+  }
   TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(
       grad.dtype(), dtype,
       thd_grad_dispatcher<dtype>(grad, grad_per_step, cu_seqlens, first_half, second_half,

transformer_engine/common/fused_attn/fused_attn.cpp

@@ -274,7 +274,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend(
            (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK ||
             attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK ||
             attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK ||
-            attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)) ||
+            attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK ||
+            (sm_arch_ >= 100 &&
+             attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK))) ||
           // 9.21: d_qk=192, d_v=128
           (cudnn_runtime_version >= 92100 && sm_arch_ >= 100 && head_dim_qk <= 192 &&
            head_dim_v <= 128 && head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 &&
@@ -283,13 +285,20 @@ 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}; 9.25+ sm90+: {thd} inference-only
       ((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) ||
+         (cudnn_runtime_version >= 92500 && !is_training)) &&
+        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 ||
+         (sm_arch_ >= 100 &&
+          attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)))) &&
       // 9.10.0: known bugs with SDPA FP8
       (cudnn_runtime_version != 91000) && !return_max_logit) {
     backend = NVTE_Fused_Attn_Backend::NVTE_FP8;
@@ -645,12 +654,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");
   }
@@ -747,14 +757,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");
   }