Congma/ck tile/preshuffle b

[CongMa13]

This PR will improve the performance of the preshuffle_b

• Introduces a constraint on the per-thread load size along the K dimension from global memory.
• Each thread now loads either:
  • 16 bytes (a single dwordx4 instruction), or
  • Exactly the K required by the MFMA instruction when 16 bytes is inadequate.
• In the 16-byte mode, data from one dwordx4 load can be consumed by one or multiple MFMA instructions.
• In the MFMA-K mode, multiple dwordx4 loads may be consumed by a single MFMA instruction (e.g., f8_16x16x128 on gfx950).
• Both modes are tuned to deliver optimal performance.
• Adds a helper function get_k_warp_tile_for_preshuffle_b to compute the per-lane load size.

[Copilot]

Pull request overview

This PR refactors how K_Warp_Tile is chosen for preshuffle_b to better align per-lane global-memory load sizes with MFMA K requirements, aiming to improve preshuffle_b performance across architectures.

Changes:

• Added get_k_warp_tile_for_preshuffle_b and updated multiple configs (tests/examples) to use it for preshuffle-B kernels.
• Simplified B-shuffle host reference layouts and adjusted warp-lane factoring in tensor_shuffle_utils.hpp.
• Updated WP pipeline policy KB-per-load computation.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
test/ck_tile/grouped_gemm_preshuffle/test_grouped_gemm_preshuffle_util.hpp	Switches grouped preshuffle tests to the new get_k_warp_tile_for_preshuffle_b helper.
test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_util.hpp	Updates test configs to compute K_Warp_Tile via helper logic and adds needed include.
test/ck_tile/gemm_multi_abd/test_gemm_multi_abd_util.hpp	Removes local get_k_warp_tile helper and relies on shared header.
test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp	Removes duplicated K_Warp_Tile derivation helpers and centralizes on shared header.
test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp	Derives K_Warp_Tile dynamically based on whether preshuffle-B is enabled.
include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp	Alters how KB-per-load is computed for the weight preshuffle pipeline policy.
include/ck_tile/ops/gemm/pipeline/tile_gemm_shape.hpp	Introduces get_k_warp_tile_for_preshuffle_b.
include/ck_tile/host/tensor_shuffle_utils.hpp	Updates host reference shuffling for B to use warp-lane factoring (k-lane-per-warp).
example/ck_tile/38_block_scale_gemm/gemm_utils.hpp	Updates example configs to use get_k_warp_tile_for_preshuffle_b.
example/ck_tile/17_grouped_gemm/quant_grouped_gemm_config.hpp	Updates grouped GEMM quant config to use new preshuffle-B K sizing.
example/ck_tile/17_grouped_gemm/grouped_gemm.hpp	Updates grouped GEMM preshuffle configs (incl. WMMA variant) to use new helper.
example/ck_tile/03_gemm/gemm_weight_preshuffle.cpp	Prints CLI help on argument-parse failure.
example/ck_tile/03_gemm/gemm_utils.hpp	Updates GEMM preshuffle configs to use get_k_warp_tile_for_preshuffle_b.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.