[rocm-libraries] ROCm/rocm-libraries#5082 (commit 9313659)

ck_tile: add gtest unit tests for MX flatmm (gfx950) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Add correctness unit tests for the MX-format flatmm kernel (`example/ck_tile/18_flatmm/mxgemm`) under `test/ck_tile/flatmm/` - Tests cover all five dtype combinations: FP4×FP4, FP8×FP8, FP6×FP6, FP8×FP4, FP4×FP8 - Tests cover all four kernel dispatch paths (the `has_hot_loop` × `tail_num` product): - `has_hot_loop=false, tail=ODD` (K=256, num_loop=1) - `has_hot_loop=false, tail=EVEN` (K=512, num_loop=2) - `has_hot_loop=true, tail=ODD` (K=768, num_loop=3) - `has_hot_loop=true, tail=EVEN` (K=1024, num_loop=4) - Remove unsupported `-split_k` CLI option from `tile_example_mx_flatmm`; the pre-shuffled B layout is incompatible with K-splitting and the option silently produced wrong results ## Changes **New files (`test/ck_tile/flatmm/`):** - `CMakeLists.txt` — builds 40 kernel instances as a shared OBJECT library, links into 5 per-dtype test executables; forwards `-DCK_TILE_USE_OCP_FP8` when `CK_USE_OCP_FP8` is ON - `test_mx_flatmm_base.hpp` — base test fixture with `run_test_with_validation(M, N, K, kbatch=1)` - `test_mx_flatmm_fixtures.hpp` — concrete `TestMXFlatmm` typed test class and type aliases - `test_mx_flatmm_fp{4fp4,8fp8,6fp6,8fp4,4fp8}.cpp` — per-dtype `TYPED_TEST_SUITE` files **Modified files:** - `example/ck_tile/18_flatmm/mxgemm/mx_flatmm_arch_traits.hpp` — moved `preShuffleWeight` here (was in `mx_flatmm.cpp`) so it is includeable by both the example and the tests - `example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp` / `run_mx_flatmm.inc` — removed `-split_k` CLI arg, hardcoded `k_batch=1`, fixed `k_split` formula, updated call sites after `preShuffleWeight` move - `test/ck_tile/CMakeLists.txt` — added `add_subdirectory(flatmm)`
2026-04-20 06:49:15 +00:00 · 2026-03-11 22:47:59 +00:00
parent 2169367735
commit 1a4aa7fd89
14 changed files with 627 additions and 73 deletions
--- a/test/ck_tile/flatmm/CMakeLists.txt
+++ b/test/ck_tile/flatmm/CMakeLists.txt
@@ -0,0 +1,79 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+set(TEST_FLATMM_COMPILE_OPTIONS)
+list(APPEND TEST_FLATMM_COMPILE_OPTIONS -mllvm -enable-noalias-to-md-conversion=0)
+
+if(CK_USE_OCP_FP8)
+    list(APPEND TEST_FLATMM_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8)
+endif()
+
+if(GPU_TARGETS MATCHES "gfx95")
+    set(MXGEMM_EXAMPLE_DIR ${CMAKE_SOURCE_DIR}/example/ck_tile/18_flatmm/mxgemm)
+
+    # Generate the 40 kernel instance .cpp files.
+    # We inline the generation here (rather than calling mx_flatmm_instance_generate)
+    # so that configure_file paths resolve correctly from this directory.
+    set(C_DATA_TYPE FP16)
+    set(A_LAYOUT ROW)
+    set(B_LAYOUT COL)
+    set(C_LAYOUT ROW)
+
+    set(FLATMM_INSTANCE_FILES)
+    foreach(PERSISTENT false)
+        foreach(DATA_TYPE FP4xFP4 FP8xFP8 FP6xFP6 FP8xFP4 FP4xFP8)
+            string(REPLACE "x" ";" DATA_TYPE_AB ${DATA_TYPE})
+            list(GET DATA_TYPE_AB 0 A_DATA_TYPE)
+            list(GET DATA_TYPE_AB 1 B_DATA_TYPE)
+            set(ARCH MXFlatmm_GFX950_)
+            set(MXFLATMM_ARCH_TRAITS "${ARCH}${A_DATA_TYPE}${B_DATA_TYPE}_Traits")
+            foreach(SPLIT_K false)
+                foreach(HAS_HOT_LOOP false true)
+                    foreach(TAIL_NUMBER ODD EVEN)
+                        set(KERNEL_FILE instance_${ARCH}${DATA_TYPE}_${PERSISTENT}_${SPLIT_K}_${HAS_HOT_LOOP}_${TAIL_NUMBER}.cpp)
+                        string(TOLOWER ${KERNEL_FILE} KERNEL_FILE)
+                        configure_file(
+                            ${MXGEMM_EXAMPLE_DIR}/mx_flatmm_instance.cpp.in
+                            ${CMAKE_CURRENT_BINARY_DIR}/${KERNEL_FILE}
+                            @ONLY)
+                        list(APPEND FLATMM_INSTANCE_FILES ${CMAKE_CURRENT_BINARY_DIR}/${KERNEL_FILE})
+                    endforeach()
+                endforeach()
+            endforeach()
+        endforeach()
+    endforeach()
+
+    # Compile the 20 kernel instances once into an object library,
+    # shared across all 5 test executables to avoid redundant GPU compilation.
+    # SPLIT_K=true instances are omitted: split-K is confirmed broken at the
+    # kernel level for all dtype combinations and is not tested.
+    add_library(mx_flatmm_test_instances OBJECT ${FLATMM_INSTANCE_FILES})
+    target_include_directories(mx_flatmm_test_instances PRIVATE
+        ${MXGEMM_EXAMPLE_DIR}
+    )
+    target_compile_options(mx_flatmm_test_instances PRIVATE ${TEST_FLATMM_COMPILE_OPTIONS})
+
+    foreach(DTYPE fp4fp4 fp8fp8 fp6fp6 fp8fp4 fp4fp8)
+        add_gtest_executable(test_tile_mx_flatmm_${DTYPE}
+            test_mx_flatmm_${DTYPE}.cpp
+        )
+        target_include_directories(test_tile_mx_flatmm_${DTYPE} PRIVATE
+            ${CMAKE_CURRENT_SOURCE_DIR}
+            ${MXGEMM_EXAMPLE_DIR}
+        )
+        target_compile_options(test_tile_mx_flatmm_${DTYPE} PRIVATE ${TEST_FLATMM_COMPILE_OPTIONS})
+        target_link_libraries(test_tile_mx_flatmm_${DTYPE} PRIVATE mx_flatmm_test_instances)
+    endforeach()
+
+    # Umbrella target to build all flatmm tests at once
+    add_custom_target(test_tile_mx_flatmm_all)
+    add_dependencies(test_tile_mx_flatmm_all
+        test_tile_mx_flatmm_fp4fp4
+        test_tile_mx_flatmm_fp8fp8
+        test_tile_mx_flatmm_fp6fp6
+        test_tile_mx_flatmm_fp8fp4
+        test_tile_mx_flatmm_fp4fp8
+    )
+else()
+    message(DEBUG "Skipping ck_tile flatmm tests for current target")
+endif()
--- a/test/ck_tile/flatmm/test_mx_flatmm_base.hpp
+++ b/test/ck_tile/flatmm/test_mx_flatmm_base.hpp
@@ -0,0 +1,251 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <gtest/gtest.h>
+#include <cstring>
+#include <optional>
+#include <random>
+#include <stdexcept>
+#include <type_traits>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/host.hpp"
+#include "ck_tile/host/check_err.hpp"
+#include "ck_tile/host/reference/reference_gemm.hpp"
+#include "ck_tile/ops/flatmm.hpp"
+#include "ck_tile/ops/gemm.hpp"
+
+#include "mx_flatmm.hpp"
+
+// Base class for MX Flatmm unit tests.
+//
+// Tuple layout: <ADataType, BDataType, CDataType, MXFlatmmArchTraits>
+template <typename Tuple>
+class TestMXFlatmmBase : public ::testing::Test
+{
+    protected:
+    using ADataType          = std::tuple_element_t<0, Tuple>;
+    using BDataType          = std::tuple_element_t<1, Tuple>;
+    using CDataType          = std::tuple_element_t<2, Tuple>;
+    using MXFlatmmArchTraits = std::tuple_element_t<3, Tuple>;
+
+    using FlatmmConfig = typename MXFlatmmArchTraits::Config;
+    using AccDataType  = float;
+    using ScaleType    = ck_tile::e8m0_t;
+
+    using ALayout = ck_tile::tensor_layout::gemm::RowMajor;
+    using BLayout = ck_tile::tensor_layout::gemm::ColumnMajor;
+    using CLayout = ck_tile::tensor_layout::gemm::RowMajor;
+
+    static constexpr int ScaleGranularityM = 1;
+    static constexpr int ScaleGranularityN = 1;
+    static constexpr int ScaleGranularityK = 32;
+
+    using ScaleA = ck_tile::FlatmmScalePointer<ScaleGranularityM, ScaleGranularityK, ScaleType>;
+    using ScaleB = ck_tile::FlatmmScalePointer<ScaleGranularityN, ScaleGranularityK, ScaleType>;
+
+    void
+    run_test_with_validation(ck_tile::index_t M,
+                             ck_tile::index_t N,
+                             ck_tile::index_t K,
+                             ck_tile::index_t kbatch                              = 1,
+                             std::optional<bool> expected_has_hot_loop            = std::nullopt,
+                             std::optional<ck_tile::TailNumber> expected_tail_num = std::nullopt)
+    {
+        constexpr int APackedSize = ck_tile::numeric_traits<ADataType>::PackedSize;
+        constexpr int BPackedSize = ck_tile::numeric_traits<BDataType>::PackedSize;
+
+        ASSERT_EQ(K % ScaleGranularityK, 0) << "K must be a multiple of ScaleGranularityK (32)";
+        ASSERT_EQ(K % APackedSize, 0) << "K must be a multiple of A PackedSize";
+        ASSERT_EQ(K % BPackedSize, 0) << "K must be a multiple of B PackedSize";
+
+        constexpr bool a_row_major = true;
+        constexpr bool b_row_major = false;
+        constexpr bool c_row_major = true;
+
+        const ck_tile::index_t stride_A =
+            ck_tile::get_default_stride(M, K, 0, ck_tile::bool_constant<a_row_major>{});
+        const ck_tile::index_t stride_B =
+            ck_tile::get_default_stride(K, N, 0, ck_tile::bool_constant<b_row_major>{});
+        const ck_tile::index_t stride_C =
+            ck_tile::get_default_stride(M, N, 0, ck_tile::bool_constant<c_row_major>{});
+
+        const auto scale_stride_A = ck_tile::get_default_stride(
+            M / ScaleGranularityM, K / ScaleGranularityK, 0, ck_tile::bool_constant<a_row_major>{});
+        const auto scale_stride_B = ck_tile::get_default_stride(
+            K / ScaleGranularityK, N / ScaleGranularityN, 0, ck_tile::bool_constant<b_row_major>{});
+
+        // Host tensors
+        ck_tile::HostTensor<ADataType> a_host(
+            ck_tile::host_tensor_descriptor(M, K, stride_A, ck_tile::bool_constant<a_row_major>{}));
+        ck_tile::HostTensor<BDataType> b_origin_host(
+            ck_tile::host_tensor_descriptor(K, N, stride_B, ck_tile::bool_constant<b_row_major>{}));
+        ck_tile::HostTensor<CDataType> c_rslt_host(
+            ck_tile::host_tensor_descriptor(M, N, stride_C, ck_tile::bool_constant<c_row_major>{}));
+
+        ck_tile::HostTensor<ScaleType> scale_a(
+            ck_tile::host_tensor_descriptor(M / ScaleGranularityM,
+                                            K / ScaleGranularityK,
+                                            scale_stride_A,
+                                            ck_tile::bool_constant<a_row_major>{}));
+        ck_tile::HostTensor<ScaleType> scale_b(
+            ck_tile::host_tensor_descriptor(K / ScaleGranularityK,
+                                            N / ScaleGranularityN,
+                                            scale_stride_B,
+                                            ck_tile::bool_constant<b_row_major>{}));
+
+        // Initialize data
+        if constexpr(std::is_same_v<ADataType, ck_tile::pk_fp6x16_t>)
+        {
+            // FP6: fill raw bytes with values 1..4 (avoids denormals)
+            auto a_bytes = a_host.get_element_space_size_in_bytes();
+            auto b_bytes = b_origin_host.get_element_space_size_in_bytes();
+            std::vector<int8_t> buf_a(a_bytes), buf_b(b_bytes);
+            std::mt19937 gen(42);
+            std::uniform_int_distribution<int> dis(1, 4);
+            for(auto& v : buf_a)
+                v = static_cast<int8_t>(dis(gen));
+            for(auto& v : buf_b)
+                v = static_cast<int8_t>(dis(gen));
+            memcpy(a_host.data(), buf_a.data(), a_bytes);
+            memcpy(b_origin_host.data(), buf_b.data(), b_bytes);
+            ck_tile::FillUniformDistribution<>{-1.f, 1.f}(scale_a);
+            ck_tile::FillUniformDistribution<>{-1.f, 1.f}(scale_b);
+        }
+        else
+        {
+            ck_tile::FillUniformDistribution<>{0.0f, 1.0f}(a_host);
+            ck_tile::FillUniformDistribution<>{-.5f, .5f}(b_origin_host);
+            ck_tile::FillUniformDistribution<>{-2.f, 2.f}(scale_a);
+            ck_tile::FillUniformDistribution<>{-2.f, 2.f}(scale_b);
+        }
+
+        // Preshuffle B and scales
+        const auto b_shuffled_host  = MXFlatmmArchTraits::preShuffleWeight(b_origin_host);
+        const auto scale_a_shuffled = MXFlatmmArchTraits::template preShuffleScale<true>(scale_a);
+        const auto scale_b_shuffled = MXFlatmmArchTraits::template preShuffleScale<false>(scale_b);
+
+        // Device buffers
+        ck_tile::DeviceMem a_dev_buf(a_host.get_element_space_size_in_bytes());
+        ck_tile::DeviceMem b_shuffled_dev_buf(b_shuffled_host.get_element_space_size_in_bytes());
+        ck_tile::DeviceMem c_dev_buf(c_rslt_host.get_element_space_size_in_bytes());
+        ck_tile::DeviceMem scale_a_dev_buf(scale_a_shuffled.get_element_space_size_in_bytes());
+        ck_tile::DeviceMem scale_b_dev_buf(scale_b_shuffled.get_element_space_size_in_bytes());
+
+        a_dev_buf.ToDevice(a_host.data());
+        b_shuffled_dev_buf.ToDevice(b_shuffled_host.data());
+        c_rslt_host.SetZero();
+        c_dev_buf.ToDevice(c_rslt_host.data());
+        scale_a_dev_buf.ToDevice(scale_a_shuffled.data());
+        scale_b_dev_buf.ToDevice(scale_b_shuffled.data());
+
+        auto scale_a_dev_ptr = ScaleA{static_cast<ScaleType*>(scale_a_dev_buf.GetDeviceBuffer()),
+                                      M / ScaleGranularityM};
+        auto scale_b_dev_ptr = ScaleB{static_cast<ScaleType*>(scale_b_dev_buf.GetDeviceBuffer()),
+                                      N / ScaleGranularityN};
+
+        // Build args
+        ck_tile::ScaleFlatmmHostArgs<ScaleA, ScaleB> args{a_dev_buf.GetDeviceBuffer(),
+                                                          b_shuffled_dev_buf.GetDeviceBuffer(),
+                                                          {},
+                                                          c_dev_buf.GetDeviceBuffer(),
+                                                          kbatch,
+                                                          M,
+                                                          N,
+                                                          K,
+                                                          stride_A,
+                                                          stride_B,
+                                                          {},
+                                                          stride_C,
+                                                          scale_a_dev_ptr,
+                                                          scale_b_dev_ptr};
+
+        // Compute hot_loop / tail_num
+        using FlatmmShape = ck_tile::TileGemmShape<
+            ck_tile::sequence<FlatmmConfig::M_Tile, FlatmmConfig::N_Tile, FlatmmConfig::K_Tile>,
+            ck_tile::sequence<FlatmmConfig::M_Warp, FlatmmConfig::N_Warp, FlatmmConfig::K_Warp>,
+            ck_tile::sequence<FlatmmConfig::M_Warp_Tile,
+                              FlatmmConfig::N_Warp_Tile,
+                              FlatmmConfig::K_Warp_Tile>>;
+
+        using TilePartitioner =
+            ck_tile::GemmSpatiallyLocalTilePartitioner<FlatmmShape,
+                                                       FlatmmConfig::TileParitionerGroupNum,
+                                                       FlatmmConfig::TileParitionerM01>;
+
+        using GemmTraits          = ck_tile::TileGemmTraits<FlatmmConfig::kPadM,
+                                                            FlatmmConfig::kPadN,
+                                                            FlatmmConfig::kPadK,
+                                                            ALayout,
+                                                            BLayout,
+                                                            CLayout,
+                                                            FlatmmConfig::NumWaveGroups>;
+        using GemmPipelineProblem = ck_tile::
+            GemmPipelineProblem<ADataType, BDataType, AccDataType, FlatmmShape, GemmTraits>;
+        using BaseFlatmmPipeline = ck_tile::BaseFlatmmPipelineAGmemBGmemCRegV1<GemmPipelineProblem>;
+
+        const ck_tile::index_t k_grain     = args.k_batch * FlatmmConfig::K_Tile;
+        const ck_tile::index_t k_split     = (K + k_grain - 1) / k_grain * k_grain;
+        const ck_tile::index_t num_loop    = TilePartitioner::GetLoopNum(k_split);
+        const bool has_hot_loop            = BaseFlatmmPipeline::BlockHasHotloop(num_loop);
+        const ck_tile::TailNumber tail_num = BaseFlatmmPipeline::GetBlockLoopTailNum(num_loop);
+
+        if(expected_has_hot_loop.has_value())
+            ASSERT_EQ(has_hot_loop, *expected_has_hot_loop)
+                << "has_hot_loop mismatch for (M=" << M << ", N=" << N << ", K=" << K << ")";
+        if(expected_tail_num.has_value())
+            ASSERT_EQ(tail_num, *expected_tail_num)
+                << "tail_num mismatch for (M=" << M << ", N=" << N << ", K=" << K << ")";
+
+            // Launch kernel (warmup=0, repeat=1 for correctness testing)
+            // mx_flatmm_calc is explicitly instantiated in the linked object library;
+            // suppress the -Wundefined-func-template warning that fires when the
+            // compiler sees only the forward declaration in mx_flatmm.hpp.
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wundefined-func-template"
+        BaseFlatmmPipeline::template TailHandler<true>(
+            [&](auto has_hot_loop_, auto tail_num_) {
+                constexpr auto has_hot_loop_v = has_hot_loop_.value;
+                constexpr auto tail_num_v     = tail_num_.value;
+                // SplitK (kbatch>1) is excluded: confirmed broken at the kernel level.
+                // Always dispatch the kbatch=1 (SPLIT_K=false) path.
+                mx_flatmm_calc<MXFlatmmArchTraits,
+                               ADataType,
+                               BDataType,
+                               ck_tile::tuple<>,
+                               AccDataType,
+                               CDataType,
+                               ALayout,
+                               BLayout,
+                               ck_tile::tuple<>,
+                               CLayout,
+                               ScaleA,
+                               ScaleB,
+                               /*persistent=*/false,
+                               ck_tile::element_wise::PassThrough,
+                               /*split_k=*/false,
+                               has_hot_loop_v,
+                               tail_num_v>(args, ck_tile::stream_config{nullptr, false, 0, 0, 1});
+            },
+            has_hot_loop,
+            tail_num);
+#pragma clang diagnostic pop
+
+        c_dev_buf.FromDevice(c_rslt_host.data());
+
+        // CPU reference
+        ck_tile::HostTensor<CDataType> c_ref(
+            ck_tile::host_tensor_descriptor(M, N, stride_C, ck_tile::bool_constant<c_row_major>{}));
+        c_ref.SetZero();
+
+        ck_tile::reference_mx_gemm<ADataType, BDataType, ScaleType, AccDataType, CDataType>(
+            a_host, b_origin_host, c_ref, scale_a, scale_b);
+
+        const float rtol = 1e-2f;
+        const float atol = 1e-2f;
+        EXPECT_TRUE(
+            ck_tile::check_err(c_rslt_host, c_ref, "MX Flatmm result mismatch", rtol, atol));
+    }
+};
--- a/test/ck_tile/flatmm/test_mx_flatmm_fixtures.hpp
+++ b/test/ck_tile/flatmm/test_mx_flatmm_fixtures.hpp
@@ -0,0 +1,20 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include "test_mx_flatmm_base.hpp"
+#include "mx_flatmm_arch_traits.hpp"
+
+// Convenience type aliases for use in test .cpp files
+using FP4  = ck_tile::pk_fp4_t;
+using FP6  = ck_tile::pk_fp6x16_t;
+using FP8  = ck_tile::fp8_t;
+using FP16 = ck_tile::fp16_t;
+
+// Concrete test fixture — inherits all logic from TestMXFlatmmBase.
+// Tuple layout: <ADataType, BDataType, CDataType, MXFlatmmArchTraits>
+template <typename Tuple>
+class TestMXFlatmm : public TestMXFlatmmBase<Tuple>
+{
+};
--- a/test/ck_tile/flatmm/test_mx_flatmm_fp4fp4.cpp
+++ b/test/ck_tile/flatmm/test_mx_flatmm_fp4fp4.cpp
@@ -0,0 +1,41 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include "ck_tile/host.hpp"
+#include <gtest/gtest.h>
+#include "test_mx_flatmm_fixtures.hpp"
+
+// FP4 x FP4 -> FP16
+// N_Tile = 512 (MXfp4_FlatmmConfig16), so N must be a multiple of 512.
+// K must be a multiple of 32 (ScaleGranularityK) and 8 (FP4 PackedSize) -> multiple of 32.
+// clang-format off
+using FP4FP4Types = ::testing::Types<
+    std::tuple<FP4, FP4, FP16, MXFlatmm_GFX950_FP4FP4_Traits>
+>;
+// clang-format on
+
+TYPED_TEST_SUITE(TestMXFlatmm, FP4FP4Types);
+
+// K=256 -> num_loop=1: has_hot_loop=false, tail=Odd
+TYPED_TEST(TestMXFlatmm, SmallMNK)
+{
+    this->run_test_with_validation(128, 512, 256, 1, false, ck_tile::TailNumber::Odd);
+}
+
+// K=512 -> num_loop=2: has_hot_loop=false, tail=Even
+TYPED_TEST(TestMXFlatmm, MediumMNK)
+{
+    this->run_test_with_validation(256, 1024, 512, 1, false, ck_tile::TailNumber::Even);
+}
+
+// K=768 -> num_loop=3: has_hot_loop=true, tail=Odd
+TYPED_TEST(TestMXFlatmm, LargeK_HotLoopOdd)
+{
+    this->run_test_with_validation(128, 512, 768, 1, true, ck_tile::TailNumber::Odd);
+}
+
+// K=1024 -> num_loop=4: has_hot_loop=true, tail=Even
+TYPED_TEST(TestMXFlatmm, LargeK_HotLoopEven)
+{
+    this->run_test_with_validation(128, 512, 1024, 1, true, ck_tile::TailNumber::Even);
+}
--- a/test/ck_tile/flatmm/test_mx_flatmm_fp4fp8.cpp
+++ b/test/ck_tile/flatmm/test_mx_flatmm_fp4fp8.cpp
@@ -0,0 +1,40 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include "ck_tile/host.hpp"
+#include <gtest/gtest.h>
+#include "test_mx_flatmm_fixtures.hpp"
+
+// FP4 x FP8 -> FP16
+// N_Tile = 256, K must be a multiple of lcm(32, 8) = 32.
+// clang-format off
+using FP4FP8Types = ::testing::Types<
+    std::tuple<FP4, FP8, FP16, MXFlatmm_GFX950_FP4FP8_Traits>
+>;
+// clang-format on
+
+TYPED_TEST_SUITE(TestMXFlatmm, FP4FP8Types);
+
+// K=256 -> num_loop=1: has_hot_loop=false, tail=Odd
+TYPED_TEST(TestMXFlatmm, SmallMNK)
+{
+    this->run_test_with_validation(128, 256, 256, 1, false, ck_tile::TailNumber::Odd);
+}
+
+// K=512 -> num_loop=2: has_hot_loop=false, tail=Even
+TYPED_TEST(TestMXFlatmm, MediumMNK)
+{
+    this->run_test_with_validation(256, 512, 512, 1, false, ck_tile::TailNumber::Even);
+}
+
+// K=768 -> num_loop=3: has_hot_loop=true, tail=Odd
+TYPED_TEST(TestMXFlatmm, LargeK_HotLoopOdd)
+{
+    this->run_test_with_validation(128, 256, 768, 1, true, ck_tile::TailNumber::Odd);
+}
+
+// K=1024 -> num_loop=4: has_hot_loop=true, tail=Even
+TYPED_TEST(TestMXFlatmm, LargeK_HotLoopEven)
+{
+    this->run_test_with_validation(128, 256, 1024, 1, true, ck_tile::TailNumber::Even);
+}
--- a/test/ck_tile/flatmm/test_mx_flatmm_fp6fp6.cpp
+++ b/test/ck_tile/flatmm/test_mx_flatmm_fp6fp6.cpp
@@ -0,0 +1,40 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include "ck_tile/host.hpp"
+#include <gtest/gtest.h>
+#include "test_mx_flatmm_fixtures.hpp"
+
+// FP6 x FP6 -> FP16
+// N_Tile = 256, K must be a multiple of lcm(32, 16) = 32 (FP6 PackedSize=16, lcm(32,16)=32).
+// clang-format off
+using FP6FP6Types = ::testing::Types<
+    std::tuple<FP6, FP6, FP16, MXFlatmm_GFX950_FP6FP6_Traits>
+>;
+// clang-format on
+
+TYPED_TEST_SUITE(TestMXFlatmm, FP6FP6Types);
+
+// K=256 -> num_loop=1: has_hot_loop=false, tail=Odd
+TYPED_TEST(TestMXFlatmm, SmallMNK)
+{
+    this->run_test_with_validation(128, 256, 256, 1, false, ck_tile::TailNumber::Odd);
+}
+
+// K=512 -> num_loop=2: has_hot_loop=false, tail=Even
+TYPED_TEST(TestMXFlatmm, MediumMNK)
+{
+    this->run_test_with_validation(256, 512, 512, 1, false, ck_tile::TailNumber::Even);
+}
+
+// K=768 -> num_loop=3: has_hot_loop=true, tail=Odd
+TYPED_TEST(TestMXFlatmm, LargeK_HotLoopOdd)
+{
+    this->run_test_with_validation(128, 256, 768, 1, true, ck_tile::TailNumber::Odd);
+}
+
+// K=1024 -> num_loop=4: has_hot_loop=true, tail=Even
+TYPED_TEST(TestMXFlatmm, LargeK_HotLoopEven)
+{
+    this->run_test_with_validation(128, 256, 1024, 1, true, ck_tile::TailNumber::Even);
+}
--- a/test/ck_tile/flatmm/test_mx_flatmm_fp8fp4.cpp
+++ b/test/ck_tile/flatmm/test_mx_flatmm_fp8fp4.cpp
@@ -0,0 +1,40 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include "ck_tile/host.hpp"
+#include <gtest/gtest.h>
+#include "test_mx_flatmm_fixtures.hpp"
+
+// FP8 x FP4 -> FP16
+// N_Tile = 256, K must be a multiple of lcm(32, 8) = 32.
+// clang-format off
+using FP8FP4Types = ::testing::Types<
+    std::tuple<FP8, FP4, FP16, MXFlatmm_GFX950_FP8FP4_Traits>
+>;
+// clang-format on
+
+TYPED_TEST_SUITE(TestMXFlatmm, FP8FP4Types);
+
+// K=256 -> num_loop=1: has_hot_loop=false, tail=Odd
+TYPED_TEST(TestMXFlatmm, SmallMNK)
+{
+    this->run_test_with_validation(128, 256, 256, 1, false, ck_tile::TailNumber::Odd);
+}
+
+// K=512 -> num_loop=2: has_hot_loop=false, tail=Even
+TYPED_TEST(TestMXFlatmm, MediumMNK)
+{
+    this->run_test_with_validation(256, 512, 512, 1, false, ck_tile::TailNumber::Even);
+}
+
+// K=768 -> num_loop=3: has_hot_loop=true, tail=Odd
+TYPED_TEST(TestMXFlatmm, LargeK_HotLoopOdd)
+{
+    this->run_test_with_validation(128, 256, 768, 1, true, ck_tile::TailNumber::Odd);
+}
+
+// K=1024 -> num_loop=4: has_hot_loop=true, tail=Even
+TYPED_TEST(TestMXFlatmm, LargeK_HotLoopEven)
+{
+    this->run_test_with_validation(128, 256, 1024, 1, true, ck_tile::TailNumber::Even);
+}
--- a/test/ck_tile/flatmm/test_mx_flatmm_fp8fp8.cpp
+++ b/test/ck_tile/flatmm/test_mx_flatmm_fp8fp8.cpp
@@ -0,0 +1,40 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include "ck_tile/host.hpp"
+#include <gtest/gtest.h>
+#include "test_mx_flatmm_fixtures.hpp"
+
+// FP8 x FP8 -> FP16
+// N_Tile = 256, K must be a multiple of 32.
+// clang-format off
+using FP8FP8Types = ::testing::Types<
+    std::tuple<FP8, FP8, FP16, MXFlatmm_GFX950_FP8FP8_Traits>
+>;
+// clang-format on
+
+TYPED_TEST_SUITE(TestMXFlatmm, FP8FP8Types);
+
+// K=256 -> num_loop=1: has_hot_loop=false, tail=Odd
+TYPED_TEST(TestMXFlatmm, SmallMNK)
+{
+    this->run_test_with_validation(128, 256, 256, 1, false, ck_tile::TailNumber::Odd);
+}
+
+// K=512 -> num_loop=2: has_hot_loop=false, tail=Even
+TYPED_TEST(TestMXFlatmm, MediumMNK)
+{
+    this->run_test_with_validation(256, 512, 512, 1, false, ck_tile::TailNumber::Even);
+}
+
+// K=768 -> num_loop=3: has_hot_loop=true, tail=Odd
+TYPED_TEST(TestMXFlatmm, LargeK_HotLoopOdd)
+{
+    this->run_test_with_validation(128, 256, 768, 1, true, ck_tile::TailNumber::Odd);
+}
+
+// K=1024 -> num_loop=4: has_hot_loop=true, tail=Even
+TYPED_TEST(TestMXFlatmm, LargeK_HotLoopEven)
+{
+    this->run_test_with_validation(128, 256, 1024, 1, true, ck_tile::TailNumber::Even);
+}