From d4c4b36f58768b0ab9c4ff5da5b84389f595ecc8 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Thu, 19 Feb 2026 15:57:43 +0000 Subject: [PATCH 1/4] [CK_TILE] Fix FP8 MXGEMM numerical error in async load path The FP8 MXGEMM was producing half the expected result due to an inconsistency in the static_move_ys code path in tile_window.hpp. The issue was that lds_ys_offset was computed using a default-constructed tensor descriptor (decltype(tensor_descriptor){}), while dram_ys_offset used the actual DRAM tensor descriptor. When these descriptors have different strides, the imm_valid calculation becomes incorrect: imm_valid = lds_ys_offset % IMM_RANGE wave_offset = dram_ys_offset - imm_valid This fix sets imm_valid=0 to avoid the inconsistency, putting the full offset in wave_offset instead of splitting it between wave_offset and the IMM parameter. Tested: FP8 and FP4 MXGEMM verification passes. --- .../include/ck_tile/core/tensor/tile_window.hpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp index 2c9b227124d5..d31273155483 100644 --- a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp +++ b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp @@ -627,10 +627,11 @@ struct tile_window_with_static_distribution const auto lds_coord = make_tensor_coordinate(tensor_descriptor, lds_bottom_tensor_thread_idx); - constexpr auto IMM_RANGE = + [[maybe_unused]] constexpr auto IMM_RANGE = (1 << 12) / sizeof(typename Base::DataType) * Traits::PackedSize; - constexpr auto imm_total = lds_ys_offset; - constexpr auto imm_valid = imm_total % IMM_RANGE; + constexpr auto imm_total = lds_ys_offset; + constexpr auto imm_valid = + 0; // FIX: Avoid inconsistency between lds_ys_offset and dram_ys_offset constexpr auto imm_overflow = imm_total - imm_valid; // Calculate SMEM address using base pointer From 8dd4faff97a0d15db45469b1613036b4df62429a Mon Sep 17 00:00:00 2001 From: ThomasNing Date: Thu, 19 Feb 2026 16:04:00 -0600 Subject: [PATCH 2/4] Finish the fix --- .../ck_tile/core/tensor/tile_window.hpp | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp index d31273155483..64cfeae53311 100644 --- a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp +++ b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp @@ -627,17 +627,9 @@ struct tile_window_with_static_distribution const auto lds_coord = make_tensor_coordinate(tensor_descriptor, lds_bottom_tensor_thread_idx); - [[maybe_unused]] constexpr auto IMM_RANGE = - (1 << 12) / sizeof(typename Base::DataType) * Traits::PackedSize; - constexpr auto imm_total = lds_ys_offset; - constexpr auto imm_valid = - 0; // FIX: Avoid inconsistency between lds_ys_offset and dram_ys_offset - constexpr auto imm_overflow = imm_total - imm_valid; - // Calculate SMEM address using base pointer - CK_TILE_LDS_ADDR LdsDataType* smem = lds_base_ptr + - lds_coord.get_offset() / Traits::PackedSize + - imm_overflow / Traits::PackedSize; + CK_TILE_LDS_ADDR LdsDataType* smem = + lds_base_ptr + (lds_coord.get_offset() + lds_ys_offset) / Traits::PackedSize; const auto dram_ys_offset = [&]() { if constexpr(static_move_ys) @@ -657,13 +649,21 @@ struct tile_window_with_static_distribution offset + dram_ys_offset, bool_constant{}); else + { + using Tsrc = + remove_cvref_tget_bottom_tensor_view())::DataType>; + constexpr index_t imm_max_elems = ((1 << 12) - 1) / sizeof(Tsrc); + constexpr index_t align = Traits::PackedSize; + constexpr index_t imm_valid = (dram_ys_offset <= imm_max_elems) + ? (dram_ys_offset / align) * align + : (imm_max_elems / align) * align; this->get_bottom_tensor_view().template async_get_vectorized_elements( smem, bottom_tensor_thread_coord.get_offset() + offset, dram_ys_offset - imm_valid, number{}, bool_constant{}); - + } // Move thread coordinate if not last access if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) { From 3ec0228a8038b738f984586f1bc61e6fce0f5840 Mon Sep 17 00:00:00 2001 From: ThomasNing Date: Fri, 20 Feb 2026 13:53:42 -0600 Subject: [PATCH 3/4] fix the error --- .../include/ck_tile/core/tensor/tile_window.hpp | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp index 64cfeae53311..41f0ab25da45 100644 --- a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp +++ b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp @@ -635,7 +635,8 @@ struct tile_window_with_static_distribution if constexpr(static_move_ys) { const auto coord_ys_offset = make_tensor_coordinate( - this->get_bottom_tensor_view().get_tensor_descriptor(), idx_ys_offset); + this->get_bottom_tensor_view().get_tensor_descriptor(), + idx_ys_offset); return coord_ys_offset.get_offset(); } else @@ -650,18 +651,11 @@ struct tile_window_with_static_distribution bool_constant{}); else { - using Tsrc = - remove_cvref_tget_bottom_tensor_view())::DataType>; - constexpr index_t imm_max_elems = ((1 << 12) - 1) / sizeof(Tsrc); - constexpr index_t align = Traits::PackedSize; - constexpr index_t imm_valid = (dram_ys_offset <= imm_max_elems) - ? (dram_ys_offset / align) * align - : (imm_max_elems / align) * align; this->get_bottom_tensor_view().template async_get_vectorized_elements( smem, bottom_tensor_thread_coord.get_offset() + offset, - dram_ys_offset - imm_valid, - number{}, + dram_ys_offset, + number<0>{}, bool_constant{}); } // Move thread coordinate if not last access From 32b10fb5cc9b2560ca02667a83d21e8d68276417 Mon Sep 17 00:00:00 2001 From: ThomasNing Date: Fri, 20 Feb 2026 14:56:29 -0600 Subject: [PATCH 4/4] clang format --- .../include/ck_tile/core/tensor/tile_window.hpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp index 41f0ab25da45..ba7eeb193647 100644 --- a/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp +++ b/projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp @@ -635,8 +635,7 @@ struct tile_window_with_static_distribution if constexpr(static_move_ys) { const auto coord_ys_offset = make_tensor_coordinate( - this->get_bottom_tensor_view().get_tensor_descriptor(), - idx_ys_offset); + this->get_bottom_tensor_view().get_tensor_descriptor(), idx_ys_offset); return coord_ys_offset.get_offset(); } else