[ozturkosu]

Stacked on #15083 (Stream-K on comp_tdm and comp_tdm_v2), so the diff below includes it until #15083 merges. Only the top commits belong to this PR. It builds on #15005 (comp_tdm_v3 bf16/fp16), which is in develop. This is the comp_tdm_v3 item from the follow-up list in #15083.

Motivation

After #15083, gfx1250 Stream-K runs comp_tdm and comp_tdm_v2, but not comp_tdm_v3, the 16-wave TDM pipeline. StreamKKernel rejected it because it sets kPlainGemmOnly, and it had no Stream-K call form. Enabling it exposed two problems in the shared Stream-K reduction code (streamk_common.hpp):

• Wrong partials layout for multi-block warp tiles. MakePartialsDistribution() assumed one WMMA instruction per warp tile. The 8-bit comp_tdm_v3 warp tile is 32x32x128 (kAMBlock = kBNBlock = 2, four 16x16 instructions), so the partials distribution did not match the accumulator's thread buffer.
• VGPR spills in linear and tree reduction on 16-wave tiles. A 256x256 tile on 16 waves holds 128 accumulator VGPRs per lane. LoadPartial loaded a second full tile, and only then did AddBlockTile add it.

Changes

Pipeline (gemm_pipeline_ag_bg_cr_comp_tdm_v3.hpp)

• New Stream-K operator()(a_window, b_window, num_loop, has_hot_loop, tail_number, p_smem) for non-tuple windows. It forwards to the tuple form, the same as comp_tdm and comp_tdm_v2 in feat(ck-tile): enable gfx1250 Stream-K on the comp_tdm and comp_tdm_v2 pipelines #15083. V3 restates it because its operator() overloads hide V1's.
• The kPlainGemmOnly comment now names the Stream-K kernel too. Multi-D, grouped and batched-contraction still reject the pipeline.

Kernel (streamk_gemm_kernel.hpp)

• Removed the kIsPlainGemmOnlyPipeline static_assert. comp_tdm_v3 is the only pipeline that sets the flag.
• kIsCompTdmPipeline already covers comp_tdm_v3 (TDM, not wavelet). So V3 gets the same CShuffle/TDM epilogue and row-major C checks as comp_tdm and comp_tdm_v2. V3 drains its last prefetch before returning, so the TDM wait after a partial tile does nothing for it.
• Comments and messages updated. The reduction loops are unchanged.

Shared Stream-K reduction (streamk_common.hpp)

• MakePartialsDistribution() uses the per-instruction warp size, WarpGemm::kM / kAMBlock and WarpGemm::kN / kBNBlock. For warp tiles with kAMBlock = kBNBlock = 1 (every MFMA and 16-bit WMMA tile that the Stream-K tests and examples use) this is the old formula.
• LoadPartial returns a PartialTileLoader, which holds the partials window. AddBlockTile(tile, loader) adds each element as it is loaded, with a sched_barrier every 32 elements so that the compiler does not issue all loads first. No second full tile is held in VGPRs.

This file is also used by the Stream-K GEMM on gfx90a/gfx942/gfx950/gfx11/gfx12 and by grouped conv backward-weight Stream-K. See "Effect on existing kernels".

Tests and CMake

• test_gemm_streamk_tdm_v3_gfx1250_wmma.cpp goes into test_ck_tile_streamk_gfx1250_wmma (REGRESSION_TESTS) and _smoke, but not _legacy, the same as the comp_tdm suite. It has 40 typed configs (8 in smoke) plus edge cases:
  • 16-bit (bf16/fp16), over all four A/B layouts: 4-wave 64x64x64, 128x128x64 and 128x128x128; 16-wave 256x256x64.
  • 8-bit (fp8/bf8), row-major A and column-major B: 4-wave 128x128x128 and 128x128x256; 16-wave 256x256x128 and 256x256x256.
  • Atomic, linear and tree; persistent and non-persistent; CShuffle and TDM epilogues; padded and unpadded.
  • Edge cases: a 1-CTA full tile; num_loop 1-4 (both tails, with and without the hot loop); padded unaligned M/N/K over all four layouts; and the N checks (CShuffle rejects an unaligned N, TDM accepts any N).
• The V3 TU is built with -mllvm -enable-post-misched=1, the same as the comp_tdm_v3 GEMM tests.
• check_comp_tdm_v3_resources.cmake (from feat(ck-tile): comp_tdm_v3 bf16/fp16 GEMM on gfx1250 #15005) is now also registered for both Stream-K binaries. It fails if any CompTDMV3 kernel uses scratch or spills VGPRs.
• CHANGELOG entry.

Validation

Resource checks (no scratch, no VGPR spills):

• test_ck_tile_streamk_gfx1250_wmma_comp_tdm_v3_resources: 50 kernels. The 16-wave 256x256x256 fp8/bf8 linear and tree kernels use 251 VGPRs.
• test_ck_tile_streamk_gfx1250_wmma_smoke_comp_tdm_v3_resources: 10 kernels.
• test_ck_tile_gemm_pipeline_comp_tdm_v3_resources: 9 kernels. _bf16_resources: 14 kernels.
• I also built both Stream-K binaries with a second compiler (ROCm 10.1, clang 24). Both checks pass, and no Stream-K kernel in either binary uses scratch or spills: 31 atomic, 49 linear and 36 tree kernels.

gfx1250 B0 board (asicRevision 1, ROCm 10.2), at the head of this PR, with BUILD_DEV=On and -Werror (0 warnings):

• test_ck_tile_streamk_gfx1250_wmma: 423/423 (379 from feat(ck-tile): enable gfx1250 Stream-K on the comp_tdm and comp_tdm_v2 pipelines #15083 + 44 new).
• _smoke: 43/43. _legacy: 91/91.
• The comp_tdm_v3 Stream-K tests with --gtest_repeat=3: 132/132.
• Both Stream-K resource checks above pass on this build: 50 and 10 kernels.

Before the tree-loop change (see the end), I also ran these on an A0 board, with the same counts. On both boards I ran test_ck_tile_gemm_pipeline_comp_tdm_v3_wmma 66/66, _bf16_wmma 84/84 and _unaligned_mnk_wmma 115/115.

gfx1250 emulator (Jenkins smoke path): ctest -L SMOKE_TEST -R streamk: 2/2 (_smoke 43/43 plus its resource check).

Effect on existing kernels

The shared reduction change reaches other architectures. I checked gfx950 (MI350X, ROCm 7.1), comparing #15083's head against this PR.

Device ISA:

• Identical: the atomic kernels, test_ck_tile_streamk_regression and test_ck_tile_streamk_tile_partitioner.
• Changed: the linear and tree kernels, and grouped conv bwd-weight Stream-K. These are the fused partial add.
• Conv bwd-weight Stream-K: no spills before or after.

Tests, all passing:

• Stream-K: atomic_smoke 60, linear_smoke 63, tree_smoke 63, pipelines_smoke 36, extended 26, regression 2, tile_partitioner 40.
• test_ck_tile_grouped_conv_bwd_weight_streamk: 30.

Performance: 40_streamk_gemm example; fp16 and bf16; persistent and non-persistent; M x N x K = 256x256x65536, 512x512x16384, 1024x1024x8192, 2048x2048x4096, 3840x4096x4096 and 4096x4096x8192. Base and PR are interleaved; each value is the median of 3 rounds of 50 warmup + 100 timed runs.

Reduction	Geomean vs. #15083	Range
Linear	+0.8%	-0.6% to +3.1%
Tree	+34%	-1.0% to +86%

The example only offers atomic and linear, so for tree I built it with StreamKReductionStrategy::Tree. The tree gain is largest on small M x N with long K, where most of the time goes into the reduction. Large tiles (4096x4096x8192) are unchanged.

An earlier revision of this PR also moved the tree epilogue and partial store out of the partner loop. On gfx950 that raised tree-kernel VGPR spills from 252 to 11,138 (tree_smoke) and made tree reduction 30% slower (geomean, down to -43%). Without it, comp_tdm_v3 is still spill-free, so the tree loop is the same as in develop.

[therock-pr-bot]

⏳ Waiting on Required Checks

⏳ Required checks have not all concluded yet. This bot run is exiting without failing so those checks can report independently.

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	⏳ Pending	⏳ Still running…
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—

⏳ Required checks are still running or have not started yet.
This PR is not ready for review until they complete.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

[therock-pr-bot]

⏳ Waiting for required checks to start or finish:

• ⏳ pre-commit