Improvement in ROCM fmha-backward

.gitignore

@@ -67,6 +67,7 @@ xformers/csrc/attention/hip_fmha/*.hip
 xformers/csrc/attention/hip_fmha/*_hip.h
 xformers/csrc/attention/hip_fmha/instances/*.cu
 xformers/csrc/attention/hip_fmha/instances/*.hip
-xformers/csrc/attention/hip_fmha/instances_tiled/*.cu
-xformers/csrc/attention/hip_fmha/instances_tiled/*.hip
+xformers/csrc/attention/hip_fmha/instances/*.cu
+xformers/csrc/attention/hip_fmha/instances/*.hip
+xformers/csrc/attention/hip_fmha/instances/*_hip.h
 

setup.py

@@ -418,6 +418,14 @@ def get_extensions():
             "--ptxas-options=-allow-expensive-optimizations=true",
         ]
     elif torch.cuda.is_available() and torch.version.hip:
+        disable_hd256_hip_fmha = os.getenv("DISABLE_HD256_HIP_FMHA", "0")
+        if disable_hd256_hip_fmha == "1":
+            source_hip_maxk_256 = []
+            for ff in source_hip:
+                if ff.endswith("maxk_256.cpp"):
+                    source_hip_maxk_256 += [ff]
+            source_hip = list(set(source_hip) - set(source_hip_maxk_256))
+
         rename_cpp_cu(source_hip)
         rocm_home = os.getenv("ROCM_PATH")
         hip_version = get_hip_version(rocm_home)
@@ -436,9 +444,16 @@ def get_extensions():
             Path(this_dir) / "third_party" / "composable_kernel_tiled" / "include"
         ]
 
+        use_rtn_bf16_convert = os.getenv("ENABLE_HIP_FMHA_RTN_BF16_CONVERT", "0")
+
         generator_flag = []
+        if disable_hd256_hip_fmha == "1":
+            generator_flag += ["-DFMHA_SUPPORT_MAX_HEADDIM_128=1"]
 
         cc_flag = ["-DBUILD_PYTHON_PACKAGE"]
+        if use_rtn_bf16_convert == "1":
+            cc_flag += ["-DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=0"]
+
         extra_compile_args = {
             "cxx": ["-O3", "-std=c++17"] + generator_flag,
             "nvcc": [
@@ -447,10 +462,12 @@ def get_extensions():
                 f"--offload-arch={os.getenv('HIP_ARCHITECTURES', 'native')}",
                 "-U__CUDA_NO_HALF_OPERATORS__",
                 "-U__CUDA_NO_HALF_CONVERSIONS__",
-                "-DCK_FMHA_FWD_FAST_EXP2=1",
+                "-DCK_TILE_FMHA_FWD_FAST_EXP2=1",
                 "-fgpu-flush-denormals-to-zero",
                 "-Werror",
                 "-Woverloaded-virtual",
+                "-mllvm",
+                "-enable-post-misched=0",
             ]
             + generator_flag
             + cc_flag,

tests/test_mem_eff_attention.py

@@ -38,13 +38,16 @@
 if torch.cuda.is_available():
     compute_capability = torch.cuda.get_device_capability("cuda")
 sm70_or_better_only = pytest.mark.skipif(
-    compute_capability < (7, 0), reason="requires sm70+"
+    torch.version.cuda is not None and compute_capability < (7, 0),
+    reason="requires sm70+",
 )
 sm75_or_better_only = pytest.mark.skipif(
-    compute_capability < (7, 5), reason="requires sm75+"
+    torch.version.cuda is not None and compute_capability < (7, 5),
+    reason="requires sm75+",
 )
 sm80_or_better_only = pytest.mark.skipif(
-    compute_capability < (8, 0), reason="requires sm80+"
+    torch.version.cuda is not None and compute_capability < (8, 0),
+    reason="requires sm80+",
 )
 skip_if_rocm = pytest.mark.skipif(
     torch.version.hip is not None, reason="not supported on ROCm"
@@ -667,16 +670,8 @@ def test_backward(
 
     if op_bw == fmha.ck.BwOp:
         op_fw = fmha.ck.FwOp
-        if dtype == torch.bfloat16:
-            pytest.skip(
-                "CK Fmha backward for bfloat16 currently is not very accurate for some cases!"
-            )
         if grad_out_contiguous is False:
             pytest.skip("CK Fmha does not support contiguous layout for grad_out!")
-        if k % 2 != 0:
-            pytest.skip(
-                "CK Fmha currently requires the headdim size of query input be an even value!"
-            )
 
     qkv = None
 
@@ -1003,6 +998,38 @@ def test_dropout_backward_ck(dt, q_len, kv_len, batch_size, k, p):
     )
 
 
+@cuda_only
+@disable_tf32
+@disable_on_rocm
+@pytest.mark.parametrize("k_len", [32])
+@pytest.mark.parametrize("batch_size", [1])
+@pytest.mark.parametrize("kv_len", [3 * 32])
+@pytest.mark.parametrize("q_len", [3 * 32])
+def test_memory_efficient_attention_full_block_masked(q_len, kv_len, batch_size, k_len):
+    device = "cuda"
+    op_fw = fmha.small_k.FwOp
+    op_bw = fmha.small_k.BwOp
+
+    scale = 3
+    query = torch.randn((batch_size, q_len, k_len), device=device) * scale
+    key = torch.randn((batch_size, kv_len, k_len), device=device) * scale
+    value = torch.randn((batch_size, kv_len, k_len), device=device) * scale
+
+    # in this case, most of the blocks in a row get masked
+    attn_bias = torch.full((3, 32), float("-inf"), device=device)
+    attn_bias[:2, :4] = 0
+    attn_bias = attn_bias.flatten()[None, None, :].expand(1, q_len, -1)
+
+    out = xformers.ops.memory_efficient_attention(
+        query, key, value, attn_bias, op=(op_fw, op_bw)
+    )
+    ref = ref_attention_for_test(query, key, value, attn_bias)
+
+    assert_allclose(
+        out, ref, atol=op_fw.ERROR_ATOL[query.dtype], rtol=op_fw.ERROR_RTOL[query.dtype]
+    )
+
+
 @pytest.mark.parametrize("fmt", ["BMK", "BMHK"])
 @parametrize_opBW_device_dtype_biasT_B_Mq_Mkv_H_K_Kv__xs
 def test_lowlevel_api_shapes(opBW_device_dtype_biasT_B_Mq_Mkv_H_K_Kv, fmt):
@@ -1581,7 +1608,7 @@ def test_decoder(
     # kv_heads = 1: multiquery
     # kv_heads = None: neither MQA nor GQA
     # kv_heads > 1: BMGHK
-    if dtype == "bf16" and compute_capability < (8, 0):
+    if dtype == "bf16" and torch.version.cuda and compute_capability < (8, 0):
         raise pytest.skip("BF16 is only supported on SM80+")
     import triton
 

xformers/csrc/attention/hip_fmha/attention_backward_generic_ck_tiled.cpp

@@ -122,10 +122,6 @@ efficient_attention_backward_ck(
   int64_t K = query.size(3);
   int64_t Kv = value.size(3);
 
-  if (K % 2 != 0)
-    throw std::runtime_error(
-        "Currently CK Fmha requires the headdim of query/key be an even value!");
-
   auto opts = query.options();
 
   at::Tensor grad_q, grad_k, grad_v, grad_bias;
@@ -143,7 +139,6 @@ efficient_attention_backward_ck(
     grad_q = chunk.select(2, 0);
     grad_k = chunk.select(2, 1);
     grad_v = chunk.select(2, 2);
-    grad_q.fill_(0);
   } else if (
       key.size(3) == value.size(3) &&
       key.storage().is_alias_of(value.storage())) {
@@ -157,14 +152,24 @@ efficient_attention_backward_ck(
     grad_v = chunk.select(2, 1);
 
     grad_q = at::empty_strided(query.sizes(), query.strides(), query.options());
-    grad_q.fill_(0);
   } else {
     grad_q = at::empty_strided(query.sizes(), query.strides(), query.options());
     grad_k = at::empty_strided(key.sizes(), key.strides(), key.options());
     grad_v = at::empty_strided(value.sizes(), value.strides(), value.options());
-    grad_q.fill_(0);
   }
 
+  at::Tensor grad_q_f32;
+  const bool use_grad_q_f32 =
+      (query.scalar_type() == at::ScalarType::BFloat16 ||
+       query.scalar_type() == at::ScalarType::Half);
+
+  if (use_grad_q_f32) {
+    grad_q_f32 = at::empty(grad_q.sizes(), opts.dtype(at::kFloat));
+    grad_q_f32.fill_(0);
+  } else {
+    grad_q.fill_(0);
+  };
+
   // CK-FlashAttn requires q/k/v to have same shapes with dQ/dK/dV respectively
   TORCH_CHECK(query.sizes() == grad_q.sizes());
   TORCH_CHECK(query.strides() == grad_q.strides());
@@ -211,7 +216,7 @@ efficient_attention_backward_ck(
 
     TORCH_CHECK(p.B == logsumexp.size(0));
     TORCH_CHECK(p.Hq == logsumexp.size(1));
-    TORCH_CHECK(p.M <= logsumexp.size(2));
+    TORCH_CHECK(p.M == logsumexp.size(2));
 
     if (scale.has_value()) {
       p.scale = float(*scale);
@@ -229,6 +234,11 @@ efficient_attention_backward_ck(
     p.grad_k_ptr = is_mqa_gqa ? tmp_grad_k.data_ptr() : grad_k.data_ptr();
     p.grad_v_ptr = is_mqa_gqa ? tmp_grad_v.data_ptr() : grad_v.data_ptr();
 
+    if (use_grad_q_f32)
+      p.grad_q_f32_ptr = grad_q_f32.data_ptr();
+    else
+      p.grad_q_f32_ptr = nullptr;
+
     p.q_strides = {
         static_cast<int>(query.stride(0)),
         static_cast<int>(query.stride(1)),
@@ -260,6 +270,14 @@ efficient_attention_backward_ck(
         static_cast<int>(logsumexp.stride(1)),
         static_cast<int>(logsumexp.stride(2))};
 
+    if (use_grad_q_f32) {
+      p.grad_q_f32_strides = {
+          static_cast<int>(grad_q_f32.stride(0)),
+          static_cast<int>(grad_q_f32.stride(1)),
+          static_cast<int>(grad_q_f32.stride(2)),
+          static_cast<int>(grad_q_f32.stride(3))};
+    }
+
     if (is_mqa_gqa) {
       p.grad_k_strides = {
           static_cast<int>(tmp_grad_k.stride(0)),
@@ -335,9 +353,9 @@ efficient_attention_backward_ck(
     p.max_seqlen_q = *max_seqlen_q_;
     p.max_seqlen_k = *max_seqlen_k_;
 
-    TORCH_CHECK(p.num_batches == logsumexp.size(0));
+    // unpadded lse layout required
     TORCH_CHECK(p.Hq == logsumexp.size(1));
-    TORCH_CHECK(p.max_seqlen_q <= logsumexp.size(2));
+    TORCH_CHECK(p.M == logsumexp.size(2));
 
     if (scale.has_value())
       p.scale = float(*scale);
@@ -366,10 +384,16 @@ efficient_attention_backward_ck(
         static_cast<int>(grad_out.stride(3))};
 
     p.lsed_strides = {
-        static_cast<int>(logsumexp.stride(0)),
         static_cast<int>(logsumexp.stride(1)),
         static_cast<int>(logsumexp.stride(2))};
 
+    if (use_grad_q_f32) {
+      p.grad_q_f32_strides = {
+          static_cast<int>(grad_q_f32.stride(1)),
+          static_cast<int>(grad_q_f32.stride(2)),
+          static_cast<int>(grad_q_f32.stride(3))};
+    }
+
     if (is_mqa_gqa) {
       p.grad_k_strides = {
           static_cast<int>(tmp_grad_k.stride(1)),
@@ -480,6 +504,11 @@ efficient_attention_backward_ck(
     p.grad_k_ptr = is_mqa_gqa ? tmp_grad_k.data_ptr() : grad_k.data_ptr();
     p.grad_v_ptr = is_mqa_gqa ? tmp_grad_v.data_ptr() : grad_v.data_ptr();
     p.grad_bias_ptr = bias_requires_grad ? grad_bias.data_ptr() : nullptr;
+
+    if (use_grad_q_f32)
+      p.grad_q_f32_ptr = grad_q_f32.data_ptr();
+    else
+      p.grad_q_f32_ptr = nullptr;
   };
 
   auto inDataType = query.scalar_type();
@@ -515,6 +544,14 @@ efficient_attention_backward_ck(
     grad_v = tmp_grad_v_view.sum(3);
   }
 
+  /*
+    if (inDataType == at::ScalarType::Half)
+      grad_q = grad_q_f32.to(torch::kFloat16);
+
+    if (inDataType == at::ScalarType::BFloat16)
+      grad_q = grad_q_f32.to(torch::kBFloat16);
+  */
+
   return std::make_tuple(grad_q, grad_k, grad_v, grad_bias);
 }
 

xformers/csrc/attention/hip_fmha/attention_forward_decoder.cpp

@@ -57,7 +57,7 @@ template <
     int32_t ThreadsPerWavefront,
     int32_t WavefrontsPerBlock,
     int32_t KV_M_MAX = 8192,
-    int32_t K_MAX = 256>
+    int32_t K_MAX = K_MAX>
 at::Tensor& efficient_attention_forward_decoder_ck_out_impl(
     const at::Tensor& XQ, // [B, 1, G, H, D]
     const at::Tensor& cache_K, // [B, KV_M_MAX, G, H or 1, D]

xformers/csrc/attention/hip_fmha/attention_forward_generic_ck_tiled.cpp

@@ -316,18 +316,14 @@ efficient_attention_forward_ck(
       p.dropout_prob = 0.0f;
 
     if (p.compute_logsumexp) {
-      // align the access of logsumexp by each thread-group in cache-line size
-      int aligned_seqlen_q = (p.max_seqlen_q + 15) / 16 * 16;
-      logsumexp = at::empty(
-          {p.num_batches, Hq, aligned_seqlen_q}, opts.dtype(at::kFloat));
+      logsumexp = at::empty({1, Hq, M}, opts.dtype(at::kFloat));
       p.logsumexp_ptr = logsumexp.data_ptr();
       p.lse_strides = {
-          static_cast<int>(logsumexp.stride(0)),
           static_cast<int>(logsumexp.stride(1)),
           static_cast<int>(logsumexp.stride(2))};
     } else {
       p.logsumexp_ptr = nullptr;
-      p.lse_strides = {0, 0, 0};
+      p.lse_strides = {0, 0};
     }
   };