From 4db7bcdc5cb4d08e23b052c2a028551c686785f6 Mon Sep 17 00:00:00 2001 From: siliangchen-amd Date: Fri, 25 Sep 2026 12:04:15 +0000 Subject: [PATCH] [CK-Tile] blockscale split-K: pick the pipeline from the per-split loop count gemm_a8w8_blockscale_cktile chose has_hot_loop/tail_num from the full K, but each split only runs K / (k_batch * K_Tile) loops. With two loops per split the pipeline did not match and split-K outputs were wrong (69%-97% of elements with the default gfx942 kernel). Compute the loop count per split, as the other CK-Tile launchers in aiter do, and make test_splitk_correctness assert and cover short splits. Co-authored-by: Cursor --- .../include/gemm_a8w8_blockscale_cktile_common.cuh | 8 +++----- op_tests/test_gemm_a8w8_blockscale.py | 7 +++++++ 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/csrc/ck_gemm_a8w8_blockscale/include/gemm_a8w8_blockscale_cktile_common.cuh b/csrc/ck_gemm_a8w8_blockscale/include/gemm_a8w8_blockscale_cktile_common.cuh index 595a8100566..0d090d14d90 100644 --- a/csrc/ck_gemm_a8w8_blockscale/include/gemm_a8w8_blockscale_cktile_common.cuh +++ b/csrc/ck_gemm_a8w8_blockscale/include/gemm_a8w8_blockscale_cktile_common.cuh @@ -169,11 +169,9 @@ void TileGemmComputeImpl(ck_tile::QuantGemmHostArgs& args) ck_tile::BaseGemmPipelineAgBgCrCompV3, ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2>; - // const ck_tile::index_t K_split = - // (args.K + GemmConfig::K_Tile_v - 1) / GemmConfig::K_Tile_v * GemmConfig::K_Tile_v; - // const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); - // const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); - const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, GemmConfig::K_Tile_v); + // Each split runs its own K range, so pick the pipeline from the per-split loop count. + const ck_tile::index_t k_grain = args.k_batch * GemmConfig::K_Tile_v; + const ck_tile::index_t K_split = (args.K + k_grain - 1) / k_grain * GemmConfig::K_Tile_v; const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); diff --git a/op_tests/test_gemm_a8w8_blockscale.py b/op_tests/test_gemm_a8w8_blockscale.py index e4a6a41dcbf..98e420c873c 100644 --- a/op_tests/test_gemm_a8w8_blockscale.py +++ b/op_tests/test_gemm_a8w8_blockscale.py @@ -295,6 +295,10 @@ def test_splitk_correctness(m=4, n=2112, k=7168, dtype=dtypes.bf16, splitK=1): f"test_splitk_correctness(m={m}, n={n}, k={k}, splitK={splitK}): " f"ck_err={ck_err:.4g}, cktile_err={cktile_err:.4g}" ) + assert ck_err < 0.05 and cktile_err < 0.05, ( + f"split-K mismatch (m={m}, n={n}, k={k}, splitK={splitK}): " + f"ck_err={ck_err:.4g}, cktile_err={cktile_err:.4g}" + ) parser = argparse.ArgumentParser( @@ -520,6 +524,9 @@ def test_splitk_correctness(m=4, n=2112, k=7168, dtype=dtypes.bf16, splitK=1): print("\nRunning split-K correctness checks ...") for splitK in [1, 2]: test_splitk_correctness(m=4, n=512, k=16384, splitK=splitK) +# Two K loops per split for the default CK-Tile tile, the shortest split it runs. +for m, n, k, splitK in [(8, 256, 512, 1), (8, 256, 1024, 2), (32, 512, 2048, 3)]: + test_splitk_correctness(m=m, n=n, k=k, splitK=splitK) # Save results from benchmarks if args.output: