Merge recent changes from ROCm xformers

.github/workflows/rocm_build.yml

@@ -24,7 +24,7 @@ jobs:
         python: ['3.11']
         torch_version: ['2.5.1']
         toolkit_type: ['rocm']
-        toolkit_short_version: ['6.1', '6.2']
+        toolkit_short_version: ['6.1', '6.2', '6.3']
 
     uses: ./.github/workflows/wheels_build.yml
     if: github.repository == 'rocm/xformers'

.github/workflows/wheels.yml

@@ -32,7 +32,7 @@ jobs:
         # NOTE: Don't forget to update `upload_pt`'s matrix
         # when changing the CUDA/ROCM versions below!
         CU_VERSIONS = ['118', '121', '124']
-        ROCM_VERSIONS = ["6.1"] # <- 6.0 broken in `manylinux_2_28`
+        ROCM_VERSIONS = ['6.1', '6.2', '6.3'] # <- 6.0 broken in `manylinux_2_28`
         PY_CU = list(itertools.product(PY_VERSIONS, CU_VERSIONS))
         PY_ROCM = list(itertools.product(PY_VERSIONS, ROCM_VERSIONS))
         print("Full matrix PY_CU", PY_CU)
@@ -111,11 +111,12 @@ jobs:
           - cu121
           - cu124
           - rocm6.1
+          - rocm6.2
+          - rocm6.3
     uses: ./.github/workflows/wheels_upload_s3.yml
     with:
       aws_role: "arn:aws:iam::749337293305:role/pytorch_bot_uploader_role"
       s3_path: s3://pytorch/whl/${{ matrix.suffix }}/
       aws_s3_cp_extra_args: --acl public-read
       filter: "*torch2.5.1+${{ matrix.suffix }}*"
       execute: ${{ github.repository == 'facebookresearch/xformers' && github.ref_type == 'tag' }}
-

setup.py

@@ -522,14 +522,6 @@ def get_extensions():
     elif torch.version.hip and (
         torch.cuda.is_available() or os.getenv("HIP_ARCHITECTURES", "") != ""
     ):
-        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)
         hip_version = get_hip_version(ROCM_HOME)
 
@@ -549,8 +541,6 @@ def get_extensions():
         ]
 
         generator_flag = []
-        if disable_hd256_hip_fmha == "1":
-            generator_flag += ["-DFMHA_SUPPORT_MAX_HEADDIM_128=1"]
 
         cc_flag = ["-DBUILD_PYTHON_PACKAGE"]
         use_rtn_bf16_convert = os.getenv("ENABLE_HIP_FMHA_RTN_BF16_CONVERT", "0")
@@ -575,6 +565,7 @@ def get_extensions():
                 "-DCK_TILE_FMHA_FWD_FAST_EXP2=1",
                 "-fgpu-flush-denormals-to-zero",
                 "-Werror",
+                "-Wc++11-narrowing",
                 "-Woverloaded-virtual",
                 "-mllvm",
                 "-enable-post-misched=0",

tests/test_mem_eff_attention.py

@@ -457,6 +457,16 @@ def test_forward(opFW_device_dtype_biasT_B_Mq_Mkv_H_K_Kv, packed, fmt, **kwargs)
     if fmt == "BMK" and not fmha.common._is_bias_type_supported_in_BMK(bias_type):
         pytest.skip("BMK incompatible with this bias")
 
+    if op is fmha.ck.FwOp:
+        if (k > 256 or kv > 256) and issubclass(
+            bias_type,
+            (
+                fmha.attn_bias.PagedBlockDiagonalPaddedKeysMask,
+                fmha.attn_bias.PagedBlockDiagonalGappyKeysMask,
+            ),
+        ):
+            pytest.skip("ck.FwOp hdim-512 is not supported when Paged-KVCache is used!")
+
     query, key, value, attn_bias = create_tensors(
         *opFW_device_dtype_biasT_B_Mq_Mkv_H_K_Kv,
         fmt="BMHK" if packed else fmt,
@@ -545,6 +555,18 @@ def test_logsumexp(opFW_device_dtype_biasT_B_Mq_Mkv_H_K_Kv):
         kv,
     ) = opFW_device_dtype_biasT_B_Mq_Mkv_H_K_Kv
 
+    if op is fmha.ck.FwOp:
+        if issubclass(
+            bias_type,
+            (
+                fmha.attn_bias.PagedBlockDiagonalPaddedKeysMask,
+                fmha.attn_bias.PagedBlockDiagonalGappyKeysMask,
+            ),
+        ):
+            pytest.skip(
+                "With ck.FwOp Paged-KVCache has some problem with forward training!"
+            )
+
     query, key, value, attn_bias = create_tensors(
         *opFW_device_dtype_biasT_B_Mq_Mkv_H_K_Kv,
         fmt="BMHK",

xformers/benchmarks/benchmark_attn_decoding.py

@@ -168,6 +168,73 @@ class AttentionDecodingCUTLASS(AttentionDecodingBase):
 class AttentionDecodingCK(AttentionDecodingBase):
     OP = xops.fmha.ck.FwOp
 
+    def __init__(
+        self,
+        B: int,
+        Mq: int,
+        Mkv: int,
+        Hq: int,
+        Hkv: int,
+        K: int,
+        bw: bool,
+        attn_bias_type,
+    ) -> None:
+        dtype = torch.float16
+        torch.manual_seed(10)
+        self.sub_label = (
+            f"B={B} Mq={Mq} Mkv={Mkv} Hq={Hq} Hkv={Hkv} K={K} TotalBytes="
+            f"{((B * Mkv * Hkv * K * 2) + (B * Mq * Hq * K) + (B * Mq * Hq * K)) * 2}"
+        )
+        self.label = "attn_decoding"
+        self.shapes = (B, Mq, Mkv, Hq, Hkv, K)
+
+        assert Hkv <= Hq
+        assert Hq % Hkv == 0
+        self.q = torch.randn(
+            [B, Mq, Hkv, Hq // Hkv, K], device="cuda", dtype=dtype, requires_grad=bw
+        )
+        self.k = torch.randn(
+            [B, Mkv, Hkv, 1, K], device="cuda", dtype=dtype, requires_grad=bw
+        ).expand(-1, -1, -1, Hq // Hkv, -1)
+        self.v = torch.randn(
+            [B, Mkv, Hkv, 1, K], device="cuda", dtype=dtype, requires_grad=bw
+        ).expand(-1, -1, -1, Hq // Hkv, -1)
+
+        if Hq == Hkv:
+            self.q = self.q[:, :, :, 0]
+            self.k = self.k[:, :, :, 0]
+            self.v = self.v[:, :, :, 0]
+
+        self.attn_bias = create_attn_bias(
+            attn_bias_type,
+            batch_size=B,
+            num_heads=Hq,
+            num_heads_groups=Hq // Hkv,
+            q_len=Mq,
+            kv_len=Mkv,
+            dtype=dtype,
+            device=device,
+            requires_grad=False,
+            fmt="BMHK",
+            op=self.OP,
+        )
+
+        if isinstance(
+            self.attn_bias,
+            xops.fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask,
+        ):
+            self.q = self.q.view(1, -1, *self.q.shape[2:])
+            self.k = self.k.view(1, -1, *self.k.shape[2:])
+            self.v = self.v.view(1, -1, *self.v.shape[2:])
+
+        if hasattr(self.OP, "not_supported_reasons"):
+            inp = xops.fmha.Inputs(
+                query=self.q, key=self.k, value=self.v, attn_bias=self.attn_bias
+            )
+            not_supported_reasons = self.OP.not_supported_reasons(inp)
+            if not_supported_reasons:
+                raise NotSupportedInputError(not_supported_reasons)
+
 
 class AttentionDecodingCKDecoder(AttentionDecodingBase):
     OP = xops.fmha.ck_decoder.FwOp

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_forward.h

@@ -23,24 +23,31 @@ template <
 void run_batched_forward_mask_bias_dropout_dispatch(
     BatchedForwardParams& param,
     hipStream_t stream) {
-  // currently split-kv implementation does not support dropout
+  // currently split-kv implementation does not support:
+  // (*) dropout
+  // (*) head dimension > 256
   if constexpr (!kHasDropout) {
-    if (param.use_split_kv) {
-      if (use_splitkv_smallq(param.M, std::max(param.K, param.Kv))) {
-        batched_forward_splitkv_smallq_mask_bias_dropout_dispatch<
-            ScalarType,
-            kHasMask,
-            kHasBias,
-            MaxK>::Run(param, stream);
-      } else {
-        FMHA_FWD_SEQLEN_Q_SWITCH(param.M, MaxSeqlenQ, [&] {
-          batched_forward_splitkv_mask_bias_dropout_dispatch<
+    if (param.use_split_kv && MaxK <= 256) {
+      if constexpr (MaxK <= 256) {
+        if (use_splitkv_smallq(param.M, std::max(param.K, param.Kv))) {
+          batched_forward_splitkv_smallq_mask_bias_dropout_dispatch<
               ScalarType,
               kHasMask,
               kHasBias,
-              MaxK,
-              MaxSeqlenQ>::Run(param, stream);
-        });
+              MaxK>::Run(param, stream);
+        } else {
+          FMHA_FWD_SEQLEN_Q_SWITCH(param.M, MaxSeqlenQ, [&] {
+            batched_forward_splitkv_mask_bias_dropout_dispatch<
+                ScalarType,
+                kHasMask,
+                kHasBias,
+                MaxK,
+                MaxSeqlenQ>::Run(param, stream);
+          });
+        }
+      } else {
+        // Unreachable. Do not instantiate split-kv pipelines with head
+        // dimension > 256
       }
     } else {
       if (get_fmha_fwd_mtile(param.B, param.Hq, param.M) == 128)

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_forward_dispatch.h

@@ -47,7 +47,7 @@ struct batched_forward_mask_bias_dropout_dispatch {
 
     using FmhaFwdShape_ = typename FmhaFwdShape<MaxK, MTile>::Type;
     constexpr ck_tile::index_t occupancy =
-        (MaxK == 64) ? 3 : ((MaxK == 256) ? 1 : 2);
+        (MaxK == 64) ? 3 : ((MaxK >= 256) ? 1 : 2);
 
     constexpr auto kBiasEnum = kHasBias
         ? ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS
@@ -89,8 +89,10 @@ struct batched_forward_mask_bias_dropout_dispatch {
           using FmhaPipelineProblem =
               FmhaPipelineProblemTemp<FmhaFwdTraits_, FmhaMask>;
 
-          using FmhaFwdPipeline_ =
-              ck_tile::BlockFmhaPipelineQRKSVS<FmhaPipelineProblem>;
+          using FmhaFwdPipeline_ = std::conditional_t<
+              MaxK <= 256,
+              ck_tile::BlockFmhaPipelineQRKSVS<FmhaPipelineProblem>,
+              ck_tile::BlockFmhaPipelineQSKSVS<FmhaPipelineProblem>>;
 
           using FmhaFwdEpilogue_ =
               ck_tile::Default2DEpilogue<ck_tile::Default2DEpilogueProblem<

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_forward_splitkv_dispatch.h

@@ -100,6 +100,7 @@ struct batched_forward_splitkv_mask_bias_dropout_dispatch {
                 false, // kDoFp8StaticQuant place-holder
                 false, // kIsPagedKV
                 kHasUnevenSplits,
+                false, // kMergeNumHeadGroupsSeqLenQ
                 occupancy>;
 
             if (param.num_kv_splits > 1) {
@@ -305,7 +306,7 @@ struct batched_forward_splitkv_mask_bias_dropout_dispatch {
     }();
 
     dim3 kGridSize = FmhaFwdSplitKVKernel::GridSize(
-        param.B, param.Hq, param.M, param.Kv, param.num_kv_splits);
+        param.B, param.Hq, param.Hkv, param.M, param.Kv, param.num_kv_splits);
     constexpr dim3 kBlockSize = FmhaFwdSplitKVKernel::BlockSize();
     constexpr ck_tile::index_t kBlockPerCu = FmhaFwdSplitKVKernel::kBlockPerCu;
 

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_forward_splitkv_smallq_dispatch.h

@@ -98,6 +98,7 @@ struct batched_forward_splitkv_smallq_mask_bias_dropout_dispatch {
                 false, // kDoFp8StaticQuant place-holder
                 false, // kIsPagedKV
                 kHasUnevenSplits,
+                false, // kMergeNumHeadGroupsSeqLenQ
                 occupancy>;
 
             if (param.num_kv_splits > 1) {
@@ -304,7 +305,7 @@ struct batched_forward_splitkv_smallq_mask_bias_dropout_dispatch {
     }();
 
     dim3 kGridSize = FmhaFwdSplitKVKernel::GridSize(
-        param.B, param.Hq, param.M, param.Kv, param.num_kv_splits);
+        param.B, param.Hq, param.Hkv, param.M, param.Kv, param.num_kv_splits);
     constexpr dim3 kBlockSize = FmhaFwdSplitKVKernel::BlockSize();
     constexpr ck_tile::index_t kBlockPerCu = FmhaFwdSplitKVKernel::kBlockPerCu;
 

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_infer.h

@@ -23,24 +23,31 @@ template <
 void run_batched_infer_mask_bias_dropout_dispatch(
     BatchedForwardParams& param,
     hipStream_t stream) {
-  // currently split-kv implementation does not support dropout
+  // currently split-kv implementation does not support:
+  // (*) dropout
+  // (*) head dimension > 256
   if constexpr (!kHasDropout) {
-    if (param.use_split_kv) {
-      if (use_splitkv_smallq(param.M, std::max(param.K, param.Kv))) {
-        batched_infer_splitkv_smallq_mask_bias_dropout_dispatch<
-            ScalarType,
-            kHasMask,
-            kHasBias,
-            MaxK>::Run(param, stream);
-      } else {
-        FMHA_FWD_SEQLEN_Q_SWITCH(param.M, MaxSeqlenQ, [&] {
-          batched_infer_splitkv_mask_bias_dropout_dispatch<
+    if (param.use_split_kv && MaxK <= 256) {
+      if constexpr (MaxK <= 256) {
+        if (use_splitkv_smallq(param.M, std::max(param.K, param.Kv))) {
+          batched_infer_splitkv_smallq_mask_bias_dropout_dispatch<
               ScalarType,
               kHasMask,
               kHasBias,
-              MaxK,
-              MaxSeqlenQ>::Run(param, stream);
-        });
+              MaxK>::Run(param, stream);
+        } else {
+          FMHA_FWD_SEQLEN_Q_SWITCH(param.M, MaxSeqlenQ, [&] {
+            batched_infer_splitkv_mask_bias_dropout_dispatch<
+                ScalarType,
+                kHasMask,
+                kHasBias,
+                MaxK,
+                MaxSeqlenQ>::Run(param, stream);
+          });
+        }
+      } else {
+        // Unreachable. Do not instantiate split-kv pipelines with head
+        // dimension > 256
       }
     } else {
       if (get_fmha_fwd_mtile(param.B, param.Hq, param.M) == 128)

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_infer_dispatch.h

@@ -48,7 +48,7 @@ struct batched_infer_mask_bias_dropout_dispatch {
 
     using FmhaShape = typename FmhaFwdShape<MaxK, MTile>::Type;
     constexpr ck_tile::index_t occupancy =
-        (MaxK == 64) ? 3 : ((MaxK == 256) ? 1 : 2);
+        (MaxK == 64) ? 3 : ((MaxK >= 256) ? 1 : 2);
 
     constexpr auto kBiasEnum = kHasBias
         ? ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS
@@ -92,8 +92,10 @@ struct batched_infer_mask_bias_dropout_dispatch {
             using FmhaPipelineProblem =
                 FmhaPipelineProblemTemp<FmhaTraits, FmhaMask>;
 
-            using FmhaPipeline =
-                ck_tile::BlockFmhaPipelineQRKSVS<FmhaPipelineProblem>;
+            using FmhaPipeline = std::conditional_t<
+                MaxK <= 256,
+                ck_tile::BlockFmhaPipelineQRKSVS<FmhaPipelineProblem>,
+                ck_tile::BlockFmhaPipelineQSKSVS<FmhaPipelineProblem>>;
 
             using FmhaEpilogue =
                 ck_tile::Default2DEpilogue<ck_tile::Default2DEpilogueProblem<
@@ -124,8 +126,10 @@ struct batched_infer_mask_bias_dropout_dispatch {
         using FmhaPipelineProblem =
             FmhaPipelineProblemTemp<FmhaTraits, FmhaMask>;
 
-        using FmhaPipeline =
-            ck_tile::BlockFmhaPipelineQRKSVSAsync<FmhaPipelineProblem>;
+        using FmhaPipeline = std::conditional_t<
+            MaxK <= 256,
+            ck_tile::BlockFmhaPipelineQRKSVSAsync<FmhaPipelineProblem>,
+            ck_tile::BlockFmhaPipelineQSKSVS<FmhaPipelineProblem>>;
 
         using FmhaEpilogue =
             ck_tile::Default2DEpilogue<ck_tile::Default2DEpilogueProblem<

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_infer_splitkv_dispatch.h

@@ -101,6 +101,7 @@ struct batched_infer_splitkv_mask_bias_dropout_dispatch {
                   false, // kDoFp8StaticQuant place-holder
                   false, // kIsPagedKV
                   kHasUnevenSplits,
+                  false, // kMergeNumHeadGroupsSeqLenQ
                   occupancy>;
 
               using ODataType =
@@ -136,6 +137,7 @@ struct batched_infer_splitkv_mask_bias_dropout_dispatch {
                   false, // kDoFp8StaticQuant place-holder
                   false, // kIsPagedKV
                   kHasUnevenSplits,
+                  false, // kMergeNumHeadGroupsSeqLenQ
                   occupancy>;
 
               using ODataType =
@@ -318,7 +320,7 @@ struct batched_infer_splitkv_mask_bias_dropout_dispatch {
     }();
 
     dim3 kGridSize = FmhaFwdSplitKVKernel::GridSize(
-        param.B, param.Hq, param.M, param.Kv, param.num_kv_splits);
+        param.B, param.Hq, param.Hkv, param.M, param.Kv, param.num_kv_splits);
     constexpr dim3 kBlockSize = FmhaFwdSplitKVKernel::BlockSize();
     constexpr ck_tile::index_t kBlockPerCu = FmhaFwdSplitKVKernel::kBlockPerCu;