ROCm · illsilin · May 19, 2026 · May 2, 2026 · May 5, 2026 · May 5, 2026
diff --git a/projects/composablekernel/include/ck_tile/core/arch/amd_buffer_addressing.hpp b/projects/composablekernel/include/ck_tile/core/arch/amd_buffer_addressing.hpp
@@ -1383,6 +1383,46 @@ CK_TILE_DEVICE void async_buffer_load_dwordxn_v(void* smem,
                                                 index_t /*flag*/       = 0,
                                                 bool_constant<pre_nop> = {})
 {
+#if defined(__gfx103__) || defined(__gfx11__)
+    static_assert(num_dwords == 1 || num_dwords == 3 || num_dwords == 4,
+                  "wrong! not implemented data width");
+
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wold-style-cast"
+    // RDNA uses a synchronous fallback through a VGPR because the direct global-to-LDS
+    // buffer_load ... lds forms below are not available there. The later vmcnt fence
+    // still works, but the load/compute overlap from the async path is lost.
+    if constexpr(pre_nop)
+    {
+        asm volatile("s_nop 4" : : : "memory");
+    }
+
+    // The asm path uses ioffset as the 12-bit immediate offset with soffset = 0.
+    // The raw_buffer_load intrinsics have no immediate field, so pass ioffset as
+    // soffset; the final address is still base + voffset + ioffset.
+    as3_uint32_ptr lds_ptr = (as3_uint32_ptr)(smem);
+    if constexpr(num_dwords == 1)
+    {
+        *lds_ptr = static_cast<uint32_t>(
+            llvm_amdgcn_raw_buffer_load_i32(rsrc, voffset, ioffset, 0));
+    }
+    else if constexpr(num_dwords == 3)
+    {
+        int32x3_t tmp = llvm_amdgcn_raw_buffer_load_i32x3(rsrc, voffset, ioffset, 0);
+        lds_ptr[0]    = static_cast<uint32_t>(tmp[0]);
+        lds_ptr[1]    = static_cast<uint32_t>(tmp[1]);
+        lds_ptr[2]    = static_cast<uint32_t>(tmp[2]);
+    }
+    else if constexpr(num_dwords == 4)
+    {
+        int32x4_t tmp = llvm_amdgcn_raw_buffer_load_i32x4(rsrc, voffset, ioffset, 0);
+        lds_ptr[0]    = static_cast<uint32_t>(tmp[0]);
+        lds_ptr[1]    = static_cast<uint32_t>(tmp[1]);
+        lds_ptr[2]    = static_cast<uint32_t>(tmp[2]);
+        lds_ptr[3]    = static_cast<uint32_t>(tmp[3]);
+    }
+#pragma clang diagnostic pop
+#else
 #define CK_TILE_ASYNC_LOAD_WITH_INSTR(instr)                            \
     if constexpr(pre_nop)                                               \
         asm volatile("s_nop 4\n" instr " %1, %2, 0 offen offset:%3 lds" \
@@ -1414,6 +1454,7 @@ CK_TILE_DEVICE void async_buffer_load_dwordxn_v(void* smem,
         static_assert(false, "wrong! not implemented data width");
     }
 #undef CK_TILE_ASYNC_LOAD_WITH_INSTR
+#endif
 }
 
 CK_TILE_DEVICE void async_buffer_load_fence(index_t cnt = 0)

diff --git a/projects/composablekernel/include/ck_tile/core/tensor/load_tile.hpp b/projects/composablekernel/include/ck_tile/core/tensor/load_tile.hpp
@@ -15,6 +15,7 @@
 #include "ck_tile/core/tensor/tile_window_linear.hpp"
 #include "ck_tile/core/tensor/null_tile_window.hpp"
 #include "ck_tile/core/tensor/null_tensor.hpp"
+#include "ck_tile/core/tensor/store_tile.hpp"
 
 namespace ck_tile {
 // Per-lane read-offset tweaks allow swizzling patterns not representable by tile_distribution.
@@ -159,7 +160,18 @@ CK_TILE_DEVICE void async_load_tile_with_offset(LdsTileWindow_&& lds_tile,
                                                 bool_constant<oob_conditional_check> occ = {},
                                                 bool_constant<static_move_ys> smy        = {})
 {
+#if defined(__gfx12__)
+    // gfx12 does not support the direct global-to-LDS async buffer load used by
+    // tile_window::async_load*. Keep the architecture choice at the tile-load layer by using a
+    // regular buffer load followed by an explicit LDS store instead of hiding it in the primitive.
+    static_cast<void>(smy);
+    static_assert(!static_move_ys,
+                  "gfx12 synchronous async_load_tile fallback does not support static_move_ys");
+    auto tile = load_tile_with_offset(tile_window, offset, number<i_access>{}, occ);
+    store_tile(lds_tile, tile);
+#else
     tile_window.async_load_with_offset(offset, lds_tile, number<i_access>{}, occ, smy);
+#endif
 }
 
 template <typename LdsTileWindow_,
@@ -187,10 +199,16 @@ CK_TILE_DEVICE void async_load_tile_raw(LdsTileWindow_&& lds_tile,
                                         bool_constant<oob_conditional_check> = {},
                                         bool_constant<pre_nop>               = {})
 {
+#if defined(__gfx12__)
+    // See async_load_tile_with_offset: gfx12 uses regular buffer loads plus explicit LDS stores.
+    auto tile = load_tile(tile_window, number<i_access>{}, bool_constant<oob_conditional_check>{});
+    store_tile(lds_tile, tile);
+#else
     tile_window.async_load_raw(lds_tile,
                                number<i_access>{},
                                bool_constant<oob_conditional_check>{},
                                bool_constant<pre_nop>{});
+#endif
 }
 
 CK_TILE_DEVICE void async_load_fence(index_t cnt = 0)

diff --git a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp
@@ -495,12 +495,15 @@ struct tile_window_with_static_distribution
     {
         using LdsTileWindow = remove_cvref_t<LdsTileWindow_>;
         // using LdsTensorView = typename LdsTileWindow::BottomTensorView;
+#if !defined(__gfx12__)
         using LdsDataType = typename LdsTileWindow::DataType;
+#endif
         // using LdsDescriptor = typename LdsTileWindow::BottomTensorDesc;
 
         // issues * warps * lanes
         static_assert(LdsTileWindow::get_num_of_dimension() == 3); // TODO: hard coded
 
+#if !defined(__gfx12__)
         const index_t size_per_buf =
             lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset(
                 make_tuple(number<0>{}, number<0>{}, number<0>{})) *
@@ -523,13 +526,24 @@ struct tile_window_with_static_distribution
             size_per_buf + size_per_wave * get_warp_id(/*ReturnSgpr=*/bool_constant<false>{});
         m0_set_with_memory(
             amd_wave_read_first_lane(m0_init_value)); // This should be wave independent
+#endif
 
         using Traits = typename Base::Traits;
 
         using vector_t = typename Traits::vector_t;
         using SFC_Ys   = typename Traits::SFC_Ys;
 
+#if defined(__gfx12__)
+        // gfx12 does not support the direct global-to-LDS async buffer load. Preserve the
+        // raw LDS issue/warp/lane layout by loading through VGPRs and explicitly storing
+        // each vector to the LDS coordinate the async instruction would have targeted.
+        auto lds_bottom_tensor_view     = lds_tile.get_bottom_tensor_view();
+        const auto& lds_tensor_desc     = lds_bottom_tensor_view.get_tensor_descriptor();
+        const auto lds_lane_id          = get_lane_id();
+        const auto lds_warp_id          = get_warp_id(/*ReturnSgpr=*/bool_constant<false>{});
+#else
         LdsDataType* smem = lds_tile.get_bottom_tensor_view().get_buffer_view().p_data_;
+#endif
 
         // loop over thread tensor space [y0, y1, ...]
         static_for<0, NumCoord, 1>{}([&](auto iCoord) {
@@ -546,9 +560,31 @@ struct tile_window_with_static_distribution
                         return bool_constant<false>{};
                 }();
 
+#if defined(__gfx12__)
+                vector_t vec_value;
+                // async_get_vectorized_elements_raw ignores tensor-coordinate validity; it relies
+                // on raw buffer OOB behavior instead. Keep the synchronous gfx12 fallback aligned
+                // with that raw-load contract.
+                this->get_bottom_tensor_view()
+                    .template get_vectorized_elements_raw<vector_t, false>(
+                    vec_value,
+                    bottom_tensor_thread_coord,
+                    0,
+                    bool_constant<false>{},
+                    pre_nop_);
+
+                const typename LdsTileWindow::BottomTensorIndex lds_thread_idx{
+                    static_cast<index_t>(iCoord * NumAccessPerCoord + iCoordAccess),
+                    lds_warp_id,
+                    lds_lane_id};
+                const auto lds_coord = make_tensor_coordinate(lds_tensor_desc, lds_thread_idx);
+                lds_bottom_tensor_view.template set_vectorized_elements<vector_t, false>(
+                    lds_coord, 0, true, vec_value);
+#else
                 // read from bottom tensor
                 this->get_bottom_tensor_view().template async_get_vectorized_elements_raw<vector_t>(
                     smem, bottom_tensor_thread_coord, 0, pre_nop_);
+#endif
 
                 // move thread coordinate
                 if constexpr(iCoordAccess != (NumAccessPerCoord - 1))
@@ -562,7 +598,9 @@ struct tile_window_with_static_distribution
                     Base::move_window_adaptor_and_bottom_tensor_thread_coordinate(
                         window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys);
 
+#if !defined(__gfx12__)
                     m0_inc_with_memory(size_per_issue);
+#endif
                 }
             });
         });

diff --git a/projects/composablekernel/include/ck_tile/core/tensor/tile_window_linear.hpp b/projects/composablekernel/include/ck_tile/core/tensor/tile_window_linear.hpp
@@ -488,7 +488,9 @@ struct tile_window_linear
                                        bool_constant<pre_nop>               = {}) const
     {
         using LdsTileWindow = remove_cvref_t<LdsTileWindow_>;
+#if !defined(__gfx12__)
         using LdsDataType   = typename LdsTileWindow::DataType;
+#endif
 
         // currently we only support everything is non linear dim
         // actually it's not performant if we have linear dim(e.g. fast changing)
@@ -499,6 +501,7 @@ struct tile_window_linear
         // issues * warps * lanes
         static_assert(LdsTileWindow::get_num_of_dimension() == 3); // TODO: hard coded
 
+#if !defined(__gfx12__)
         const index_t size_per_buf =
             lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset(
                 make_tuple(number<0>{}, number<0>{}, number<0>{})) *
@@ -519,10 +522,21 @@ struct tile_window_linear
         const index_t m0_init_value = size_per_buf + size_per_wave * get_warp_id();
         m0_set_with_memory(
             amd_wave_read_first_lane(m0_init_value)); // This should be wave independent
+#endif
 
         using vector_t = typename Base::Traits::vector_t;
 
+#if defined(__gfx12__)
+        // gfx12 does not support the direct global-to-LDS async buffer load. Preserve the
+        // raw LDS issue/warp/lane layout by loading through VGPRs and explicitly storing
+        // each vector to the LDS coordinate the async instruction would have targeted.
+        auto lds_bottom_tensor_view = lds_tile.get_bottom_tensor_view();
+        const auto& lds_tensor_desc = lds_bottom_tensor_view.get_tensor_descriptor();
+        const auto lds_lane_id      = get_lane_id();
+        const auto lds_warp_id      = get_warp_id();
+#else
         LdsDataType* smem = lds_tile.get_bottom_tensor_view().get_buffer_view().p_data_;
+#endif
 
         // loop over thread tensor space [y0, y1, ...]
         auto issue = [&](auto i_access_) {
@@ -538,14 +552,36 @@ struct tile_window_linear
             auto bottom_tensor_thread_coord = cached_coords_[non_linear_id];
             auto bottom_tensor_flag         = cached_flags_[IAccess]; // get this flag anyway
 
+#if defined(__gfx12__)
+            vector_t vec_value;
+            // async_get_vectorized_elements_raw ignores tensor-coordinate validity; it relies
+            // on raw buffer OOB behavior instead. Keep the synchronous gfx12 fallback aligned
+            // with that raw-load contract.
+            this->get_bottom_tensor_view().template get_vectorized_elements_raw<vector_t, false>(
+                vec_value,
+                bottom_tensor_thread_coord,
+                0,
+                bottom_tensor_flag,
+                bool_constant<false>{},
+                pre_nop_);
+
+            const typename LdsTileWindow::BottomTensorIndex lds_thread_idx{
+                static_cast<index_t>(i_access_), lds_warp_id, lds_lane_id};
+            const auto lds_coord = make_tensor_coordinate(lds_tensor_desc, lds_thread_idx);
+            lds_bottom_tensor_view.template set_vectorized_elements<vector_t, false>(
+                lds_coord, 0, true, vec_value);
+#else
             // read from bottom tensor
             this->get_bottom_tensor_view().template async_get_vectorized_elements_raw<vector_t>(
                 smem, bottom_tensor_thread_coord, 0, bottom_tensor_flag, pre_nop_);
+#endif
 
             // move thread coordinate
             if constexpr(i_access_ != (NumAccess - 1))
             {
+#if !defined(__gfx12__)
                 m0_inc_with_memory(size_per_issue);
+#endif
             }
         };
 

diff --git a/...el/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp b/...el/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp
@@ -1088,7 +1088,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                 tile_elementwise_inout([&k_descale](auto& x) { x *= k_descale; }, s_acc);
             }
 
-            const auto p = [&]() {
+            const auto p_cast = [&]() {
                 const auto bias_tile = load_tile(bias_dram_window); // load bias tile
 
                 // STAGE 2, scale_s, add bias, mask, softmax
@@ -1423,6 +1423,15 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
                         tile_elementwise_in(p_compute_element_func, p_compute));
 #endif
             }();
+#if defined(__gfx11__)
+            // gfx11 WMMA uses different lane layouts for GEMM C and GEMM A tiles, so remap
+            // softmax P from GEMM0's C layout into GEMM1's A layout before the PV GEMM.
+            auto p = make_static_distributed_tensor<PDataType>(
+                decltype(gemm_1)::template MakeABlockTileDistribution<kM0, kN0>());
+            PermuteWarpGemmCToA(p, p_cast);
+#else
+            const auto p = p_cast;
+#endif
 
             // STAGE 3, KV gemm
             // KV_BLOCKSCALE: accumulate P*V into temporary tile before applying v_descale

diff --git a/...composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp b/...composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp
@@ -762,7 +762,7 @@ struct BlockFmhaPipelineQRKSVSAsync
                     randval_ptr, seq_offset, p_compute, randval_dram_window);
             }
 
-            const auto p = [&]() {
+            const auto p_cast = [&]() {
 #if CK_TILE_FMHA_FLOAT_TO_FLOAT16_RTN
                 // For fp32 to fp16,
                 // impl::cast_tile_pkrtz_fp16_fp32 would cause precision issue,
@@ -777,6 +777,15 @@ struct BlockFmhaPipelineQRKSVSAsync
                         tile_elementwise_in(p_compute_element_func, p_compute));
 #endif
             }();
+#if defined(__gfx11__)
+            // gfx11 WMMA uses different lane layouts for GEMM C and GEMM A tiles, so remap
+            // softmax P from GEMM0's C layout into GEMM1's A layout before the PV GEMM.
+            auto p = make_static_distributed_tensor<PDataType>(
+                decltype(gemm_1)::template MakeABlockTileDistribution<kM0, kN0>());
+            PermuteWarpGemmCToA(p, p_cast);
+#else
+            const auto p = p_cast;
+#endif
 
             float v_descale = 1.0f;
             if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE)

@@ -685,7 +685,17 @@ struct BlockFmhaPipelineQRKSVSWholeKPrefetch
 
             __builtin_amdgcn_sched_barrier(0x00000001);
 
-            auto p = cast_tile<PDataType>(tile_elementwise_in(p_compute_element_func, pcomp_tile));
+            auto p_cast =
+                cast_tile<PDataType>(tile_elementwise_in(p_compute_element_func, pcomp_tile));
+#if defined(__gfx11__)
+            // gfx11 WMMA uses different lane layouts for GEMM C and GEMM A tiles, so remap
+            // softmax P from GEMM0's C layout into GEMM1's A layout before the PV GEMM.
+            auto p = make_static_distributed_tensor<PDataType>(
+                decltype(gemm_1)::template MakeABlockTileDistribution<kM0, kN0>());
+            PermuteWarpGemmCToA(p, p_cast);
+#else
+            const auto p = p_cast;
+#endif
 
             __builtin_amdgcn_sched_barrier(0x00000001);
 

@@ -145,10 +145,16 @@ struct BlockFmhaPipelineQRKSVSWholeKPrefetchDefaultPolicy
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackV()
     {
+#if defined(__gfx11__)
+        // gfx11 WMMA V loads expect the LDS K-pack to match the warp GEMM K-per-thread;
+        // clamping to 8 under-reserves LDS padding for K-per-thread 16 variants.
+        return GetKVWarpGemmKPerThreadSize<Problem>();
+#else
         if constexpr(GetKVWarpGemmKPerThreadSize<Problem>() >= 8)
             return 8;
         else
             return 4;
+#endif
     }
 
     template <typename Problem>

diff --git a/...nclude/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_whole_k_prefetch_trload.hpp b/...nclude/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_whole_k_prefetch_trload.hpp
@@ -712,7 +712,17 @@ struct BlockFmhaPipelineQRKSVSWholeKPrefetchTrLoad
 
             __builtin_amdgcn_sched_barrier(0x00000001);
 
-            auto p = cast_tile<PDataType>(tile_elementwise_in(p_compute_element_func, pcomp_tile));
+            auto p_cast =
+                cast_tile<PDataType>(tile_elementwise_in(p_compute_element_func, pcomp_tile));
+#if defined(__gfx11__)
+            // gfx11 WMMA uses different lane layouts for GEMM C and GEMM A tiles, so remap
+            // softmax P from GEMM0's C layout into GEMM1's A layout before the PV GEMM.
+            auto p = make_static_distributed_tensor<PDataType>(
+                decltype(gemm_1)::template MakeABlockTileDistribution<kM0, kN0>());
+            PermuteWarpGemmCToA(p, p_cast);
+#else
+            const auto p = p_cast;
+#endif
 
             __builtin_amdgcn_sched_barrier(0x00000001);
 

@@ -512,8 +512,17 @@ struct BlockFmhaPipelineQSKSVS
 
             block_tile_reduce_sync(rowsum_p, f_sum, bool_constant<false>{});
 
+#if defined(__gfx11__)
+            // gfx11 WMMA uses different lane layouts for GEMM C and GEMM A tiles, so remap
+            // softmax P from GEMM0's C layout into GEMM1's A layout before the PV GEMM.
+            auto p = make_static_distributed_tensor<PDataType>(
+                decltype(gemm_1)::template MakeABlockTileDistribution<kM0, kN0>());
+            PermuteWarpGemmCToA(
+                p, cast_tile<PDataType>(tile_elementwise_in(p_compute_element_func, p_compute)));
+#else
             const auto p =
                 cast_tile<PDataType>(tile_elementwise_in(p_compute_element_func, p_compute));
+#endif
 
             __builtin_amdgcn_sched_barrier(0);