ozturkosu · Nov 5, 2026 · Nov 5, 2026 · Nov 5, 2026
@@ -18,6 +18,7 @@ Documentation for Composable Kernel available at [https://rocm.docs.amd.com/proj
 * Added FP6 (E2M3) operands to the CK Tile microscaling (MX) GEMM TDM pipelines on gfx1250, for row-major A, column-major B, a 16x16x128 warp tile, and 32-element E8M0 scale blocks, with fp16 or bf16 output.
 * Added K tails to the CK Tile microscaling (MX) GEMM TDM pipelines on gfx1250: they accept any K that is a whole number of storage elements.
 * Added CK Tile GEMM Stream-K support, including in the CK Tile dispatcher, for the gfx1250 TDM pipelines (comp_tdm and comp_tdm_v2): the CShuffle epilogue for atomic, linear, and tree reduction (the contiguous C dimension must be a multiple of the block tile), the TDM epilogue for linear and tree reduction, and persistent or non-persistent data-parallel tiles.
+* Added CK Tile GEMM Stream-K support for the gfx1250 comp_tdm_v3 pipeline: 4 and 16 waves, bf16 and fp16 inputs (any A/B layout) and fp8 and bf8 inputs (row-major A, column-major B), with the same reductions, epilogues and persistence modes as comp_tdm and comp_tdm_v2.
 
 ### Changed
 

@@ -195,6 +195,46 @@ struct StreamKReductionOps
                       "WaitStorePartialDone not implemented for this architecture.");
     }
 
+    /**
+     * @brief A partial tile in the workspace, returned by LoadPartial and read by AddBlockTile.
+     * @note The partials distribution and the C block distribution assign the same data to each
+     * thread, so a thread buffer read with one is valid under the other.
+     */
+    template <typename OAccTile, typename PartialWindow>
+    struct PartialTileLoader
+    {
+        PartialWindow window;
+    };
+
+    /**
+     * @brief Adds a partial tile from the workspace to an output block tile.
+     * @param in_out_block_tile The output block tile to which values are added.
+     * @param loader The partial tile, as returned by LoadPartial.
+     * @note The adds are fused into the loads, and a scheduling barrier every
+     * partial_adds_per_group elements keeps the compiler from issuing all loads first. Otherwise
+     * the loaded tile and the output tile are live at once, which spills on 16-wave tiles that
+     * already use most of the VGPRs (256x256 on gfx1250).
+     */
+    template <typename OAccTile, typename PartialWindow>
+    CK_TILE_DEVICE void AddBlockTile(OAccTile& in_out_block_tile,
+                                     const PartialTileLoader<OAccTile, PartialWindow>& loader) const
+    {
+        constexpr index_t partial_adds_per_group = 32;
+        using DataType                           = typename remove_cvref_t<OAccTile>::DataType;
+        auto accum_with_partials_dist            = make_static_distributed_tensor<DataType>(
+            MakePartialsDistribution(), in_out_block_tile.get_thread_buffer());
+        index_t elem_cnt = 0;
+        loader.window.load(
+            accum_with_partials_dist, make_tuple(loader.window), [&](auto& dst, const auto& src) {
+                dst += src;
+                if(++elem_cnt % partial_adds_per_group == 0)
+                {
+                    __builtin_amdgcn_sched_barrier(0);
+                }
+            });
+        in_out_block_tile.get_thread_buffer() = accum_with_partials_dist.get_thread_buffer();
+    }
+
     /**
      * @brief Adds the values of a block tile to an output block tile.
      * @param in_out_block_tile The output block tile to which values are added.
@@ -220,15 +260,13 @@ struct StreamKReductionOps
      * @brief Loads a partial block tile from the workspace buffer.
      * @param kargs Kernel arguments, including the workspace pointer.
      * @param cta_idx The index of the thread block (CTA).
-     * @param c_block_tile_dist The tile distribution for the block.
-     * @return The loaded partial block tile.
-     * @note This function calculates the buffer pointer and uses the tile distribution for
-     * loading the partial block tile.
+     * @param c_block_tile_dist The tile distribution of the C block tile (only its type is used).
+     * @return A PartialTileLoader for the partial block tile; AddBlockTile performs the load.
      */
     template <typename DataType, typename OAccTileDist>
     CK_TILE_DEVICE auto LoadPartial(const KernelArgs_& kargs,
                                     index_t cta_idx,
-                                    const OAccTileDist& c_block_tile_dist) const
+                                    [[maybe_unused]] const OAccTileDist& c_block_tile_dist) const
     {
         const auto c_block_tile_buffer_size =
             TilePartitioner::MPerBlock * TilePartitioner::NPerBlock * sizeof(DataType);
@@ -257,15 +295,8 @@ struct StreamKReductionOps
             {0, 0},
             MakePartialsDistribution());
 
-        auto partials_tile = load_tile(partial_tile_window);
-
-        // Since the partials distribution is not the same as the C block distribution, we must
-        // describe the contents in the partials tile with the C block distribution.
-        // Note: The data assigned to threads does not change between distributions.
-        auto partials_tile_with_c_distr = make_static_distributed_tensor<DataType>(
-            c_block_tile_dist, partials_tile.get_thread_buffer());
-
-        return partials_tile_with_c_distr;
+        using OAccTile = static_distributed_tensor<DataType, remove_cvref_t<OAccTileDist>>;
+        return PartialTileLoader<OAccTile, decltype(partial_tile_window)>{partial_tile_window};
     }
 
     /**
@@ -290,11 +321,15 @@ struct StreamKReductionOps
     CK_TILE_DEVICE static constexpr auto MakePartialsDistribution()
     {
         // Create the encoding to describe waves within a block
+        using WarpGemmImpl       = typename WarpGemm::WarpGemmAttribute::Impl;
         constexpr index_t m_warp = BlockGemmShape::BlockWarps::at(number<0>{});
         constexpr index_t n_warp = BlockGemmShape::BlockWarps::at(number<1>{});
 
-        constexpr index_t m_iter_per_warp = TilePartitioner::MPerBlock / (m_warp * WarpGemm::kM);
-        constexpr index_t n_iter_per_warp = TilePartitioner::NPerBlock / (n_warp * WarpGemm::kN);
+        constexpr index_t warp_inst_m = WarpGemm::kM / WarpGemmImpl::kAMBlock;
+        constexpr index_t warp_inst_n = WarpGemm::kN / WarpGemmImpl::kBNBlock;
+
+        constexpr index_t m_iter_per_warp = TilePartitioner::MPerBlock / (m_warp * warp_inst_m);
+        constexpr index_t n_iter_per_warp = TilePartitioner::NPerBlock / (n_warp * warp_inst_n);
 
         constexpr auto partials_outer_dstr_encoding = tile_distribution_encoding<
             sequence<>,
@@ -306,8 +341,8 @@ struct StreamKReductionOps
 
         // Create the encoding to describe threads within a wave
         constexpr index_t vector_size         = GetVectorSizePartials();
-        constexpr index_t m_warp_repeat       = WarpGemm::WarpGemmAttribute::Impl::kCM0PerLane;
-        constexpr index_t warp_tile_n_threads = WarpGemm::kN / vector_size;
+        constexpr index_t m_warp_repeat       = WarpGemmImpl::kCM0PerLane;
+        constexpr index_t warp_tile_n_threads = warp_inst_n / vector_size;
         constexpr index_t warp_tile_m_threads = get_warp_size() / warp_tile_n_threads;
 
         // This inner encoding ensures that contiguous threads perform vectorized writes along the

@@ -129,30 +129,29 @@ struct StreamKKernel
 
     /**
      *@brief Stream-K gives every workgroup its own K range, so the TDM pipelines run without
-     *cluster multicast, and the plain-GEMM-only pipelines are not supported.
+     *cluster multicast.
      */
-    static_assert(!UniversalGemmKernel::kIsPlainGemmOnlyPipeline,
-                  "Stream-K GEMM does not support this pipeline (plain GEMM only)");
     static constexpr bool kClusterLaunch =
         BlockGemmShape::kclusterM * BlockGemmShape::kclusterN * BlockGemmShape::kclusterK > 1;
     static_assert(!UniversalGemmKernel::kIsTdmPipeline || !kClusterLaunch,
                   "TDM Stream-K GEMM: cluster launch unsupported (per-workgroup K ranges)");
 
     /**
-     *@brief True for comp_tdm and comp_tdm_v2: the TDM pipelines other than the wavelet
-     *comp_tdm_pc. They can return with their last, clipped TDM prefetch still writing LDS.
-     *CShuffleEpilogue and TdmEpilogue wait for it; StreamKGemm waits for it after a partial tile,
-     *which runs no epilogue. C must be row-major: CShuffleEpilogue supports nothing else, and
-     *TdmEpilogue's column-major store is not validated with these pipelines.
+     *@brief True for comp_tdm, comp_tdm_v2, and comp_tdm_v3: the TDM pipelines other than the
+     *wavelet comp_tdm_pc. V1 and V2 can return with their last, clipped TDM prefetch still writing
+     *LDS (V3 drains before returning). CShuffleEpilogue and TdmEpilogue wait for it; StreamKGemm
+     *waits for it after a partial tile, which runs no epilogue. C must be row-major:
+     *CShuffleEpilogue supports nothing else, and TdmEpilogue's column-major store is not validated
+     *with these pipelines.
      */
     static constexpr bool kIsCompTdmPipeline =
         UniversalGemmKernel::kIsTdmPipeline && !is_wavelet_pipeline<GemmPipeline>;
     static_assert(!kIsCompTdmPipeline || is_tdm_epilogue_v<EpiloguePipeline> ||
                       is_specialization_of<EpiloguePipeline, CShuffleEpilogue>::value,
-                  "Stream-K with comp_tdm or comp_tdm_v2 requires CShuffleEpilogue or "
-                  "TdmEpilogue (they wait for the pipeline's last TDM prefetch).");
+                  "Stream-K with comp_tdm, comp_tdm_v2, or comp_tdm_v3 requires CShuffleEpilogue "
+                  "or TdmEpilogue (they wait for the pipeline's last TDM prefetch).");
     static_assert(!kIsCompTdmPipeline || std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>,
-                  "Stream-K with comp_tdm or comp_tdm_v2 supports row-major C only.");
+                  "Stream-K with comp_tdm, comp_tdm_v2, or comp_tdm_v3 supports row-major C only.");
 
     struct StreamKKernelArgs : ck_tile::UniversalGemmKernelArgs<>
     {

@@ -58,8 +58,9 @@ struct GemmPipelineAgBgCrCompTDMV3 : public GemmPipelineAgBgCrCompTDMV1<Problem,
                        std::is_same_v<typename Base::BLayout, tensor_layout::gemm::ColumnMajor>),
                   "comp_tdm_v3 with 8-bit A/B supports A row-major, B column-major only");
 
-    // Only tested in the plain GEMM kernel; the multi-D, grouped and batched-contraction kernels
-    // reject a pipeline that sets this (UniversalGemmKernel::kIsPlainGemmOnlyPipeline).
+    // Only tested in the plain and Stream-K GEMM kernels; the multi-D, grouped and
+    // batched-contraction kernels reject a pipeline that sets this
+    // (UniversalGemmKernel::kIsPlainGemmOnlyPipeline).
     static constexpr bool kPlainGemmOnly = true;
 
 #if defined(__gfx125__)
@@ -994,6 +995,28 @@ struct GemmPipelineAgBgCrCompTDMV3 : public GemmPipelineAgBgCrCompTDMV1<Problem,
         return Base::TailHandler(RunPipeline, has_hot_loop, tail_number);
     }
 
+    // Stream-K call form (see GemmPipelineAgBgCrCompTDMV1). V3's operator() overloads hide V1's,
+    // so V3 restates it to bind its own PipelineImpl.
+    template <typename ADramBlockWindowTmp,
+              typename BDramBlockWindowTmp,
+              typename std::enable_if_t<!is_detected<is_tuple, ADramBlockWindowTmp>::value &&
+                                            !is_detected<is_tuple, BDramBlockWindowTmp>::value,
+                                        bool>* = nullptr>
+    CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
+                                   const BDramBlockWindowTmp& b_dram_block_window_tmp,
+                                   const index_t num_loop,
+                                   const bool /*has_hot_loop*/,
+                                   const TailNumber /*tail_number*/,
+                                   void* __restrict__ p_smem) const
+    {
+        return operator()(make_tuple(a_dram_block_window_tmp),
+                          element_wise::PassThrough{},
+                          make_tuple(b_dram_block_window_tmp),
+                          element_wise::PassThrough{},
+                          num_loop,
+                          p_smem);
+    }
+
     [[nodiscard]] CK_TILE_HOST static const std::string GetName()
     {
         // clang-format off

@@ -1,7 +1,8 @@
 # Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 # SPDX-License-Identifier: MIT
 
-# Resource check for the comp_tdm_v3 kernels of test_ck_tile_gemm_pipeline_comp_tdm_v3_wmma.
+# Resource check for the comp_tdm_v3 kernels of a test binary: the comp_tdm_v3 GEMM pipeline tests
+# here and the gfx1250 Stream-K tests in ../gemm_streamk.
 # comp_tdm_v3 is laid out to fit 256 VGPRs at 16 waves without spilling; a compiler change that
 # makes it spill would still pass the numerical test, so fail here instead.
 #

@@ -164,16 +164,27 @@ if(GPU_TARGETS MATCHES "gfx1250")
     # The gfx1250 emulator runs one workgroup at a time, so a Linear or Tree tile owner waits
     # forever there for later workgroups. Both full suites are REGRESSION_TESTS; the smoke
     # suite keeps a subset of the atomic rows in the emulator stage. The TDM pipeline suite
-    # (comp_tdm, comp_tdm_v2) follows the same split. It is not built with the legacy buffer
-    # addressing, which has no LDS transpose load for the TDM pipelines.
+    # (comp_tdm, comp_tdm_v2, comp_tdm_v3) follows the same split. It is not built with the
+    # legacy buffer addressing, which has no LDS transpose load for the TDM pipelines.
+    # Re-enable, for the comp_tdm_v3 translation unit only, the post-RA machine scheduler that
+    # the top-level CMakeLists turns off (-enable-post-misched=0; the later option wins).
+    set_source_files_properties(test_gemm_streamk_tdm_v3_gfx1250_wmma.cpp
+        PROPERTIES
+        COMPILE_OPTIONS "-mllvm;-enable-post-misched=1"
+    )
+    set(COMP_TDM_V3_STREAMK_TARGETS)
     foreach(variant IN ITEMS builtin legacy smoke)
         set(test_target test_ck_tile_streamk_gfx1250_wmma)
         set(test_sources test_gemm_streamk_gfx1250_wmma.cpp test_gemm_streamk_util.cpp)
         if(NOT variant STREQUAL "builtin")
             string(APPEND test_target "_${variant}")
         endif()
         if(NOT variant STREQUAL "legacy")
-            list(APPEND test_sources test_gemm_streamk_tdm_gfx1250_wmma.cpp)
+            list(APPEND test_sources
+                test_gemm_streamk_tdm_gfx1250_wmma.cpp
+                test_gemm_streamk_tdm_v3_gfx1250_wmma.cpp
+            )
+            list(APPEND COMP_TDM_V3_STREAMK_TARGETS ${test_target})
         endif()
         add_gtest_executable(${test_target} ${test_sources})
         set_property(TARGET ${test_target}
@@ -185,6 +196,30 @@ if(GPU_TARGETS MATCHES "gfx1250")
         set_property(TEST ${test_target} APPEND PROPERTY LABELS CK_TILE_GEMM_STREAMK_TESTS)
         list(APPEND CK_TILE_GFX1250_STREAMK_TEST_TARGETS ${test_target})
     endforeach()
+
+    # Fails if a compiler change makes the comp_tdm_v3 Stream-K kernels use scratch.
+    get_filename_component(_llvm_bin_dir "${CMAKE_CXX_COMPILER}" DIRECTORY)
+    find_program(CK_LLVM_OBJDUMP llvm-objdump HINTS "${_llvm_bin_dir}" NO_DEFAULT_PATH)
+    find_program(CK_LLVM_READELF llvm-readelf HINTS "${_llvm_bin_dir}" NO_DEFAULT_PATH)
+    if(CK_LLVM_OBJDUMP AND CK_LLVM_READELF)
+        foreach(_sk_target IN LISTS COMP_TDM_V3_STREAMK_TARGETS)
+            set(_sk_check "${_sk_target}_comp_tdm_v3_resources")
+            if("${_sk_target}" IN_LIST REGRESSION_TESTS)
+                set(_sk_stage_label "REGRESSION_TEST")
+            else()
+                set(_sk_stage_label "SMOKE_TEST")
+            endif()
+            add_test(NAME ${_sk_check}
+                COMMAND ${CMAKE_COMMAND}
+                    -DTEST_BINARY=$<TARGET_FILE:${_sk_target}>
+                    -DLLVM_OBJDUMP=${CK_LLVM_OBJDUMP}
+                    -DLLVM_READELF=${CK_LLVM_READELF}
+                    -DWORK_DIR=${CMAKE_CURRENT_BINARY_DIR}/${_sk_target}_comp_tdm_v3_resources
+                    -P ${CMAKE_CURRENT_SOURCE_DIR}/../gemm/check_comp_tdm_v3_resources.cmake)
+            set_tests_properties(${_sk_check}
+                PROPERTIES LABELS "${_sk_stage_label};CK_TILE_GEMM_STREAMK_TESTS")
+        endforeach()
+    endif()
     add_custom_target(ck_tile_gemm_streamk_tests
         COMMAND ${CMAKE_CTEST_COMMAND} --output-on-failure -C ${CMAKE_CFG_INTDIR} -L "CK_TILE_GEMM_STREAMK_TESTS"
         DEPENDS ${CK_TILE_GFX1250_STREAMK_TEST_TARGETS}