fix(ck_tile): PermuteN epilogue applies wrong per-token/per-channel scales

projects/composablekernel/include/ck_tile/ops/epilogue/permuten_epilogue.hpp

@@ -59,6 +59,32 @@ struct PermuteNEpilogueProblem
                   "The size of DsDataType and DsLayout should be the same");
 };
 
+// ---------------------------------------------------------------------------
+// Lane -> (m, n) store mapping of this epilogue.
+//
+// The store distribution (IntrThreadShuffleEncode below) decomposes the block
+// tile as
+//   M: (kM0, kM1, kM2) = (MWave, MPerXdl/RowsPerLane, RowsPerLane)
+//   N: (kN0, kN1, kN2) = (NWave, NPerXdl,            NRepeat)
+// with P0 = (m0, n0) = wave id and P1 = (m1, n1) = lane id (lane = m1*NPerXdl
+// + n1), Y0 = m2, Y1 = n2. For a lane `l` of wave (wave_m, wave_n), thread
+// buffer element dst = n2 + m2*NRepeat is stored at the CANONICAL global
+// coordinate
+//   m = block_m + mIter*MWave*MPerXdl + wave_m*MPerXdl
+//       + (l / NPerXdl)*RowsPerLane + m2
+//   n = block_n + wave_n*NPerXdl*NRepeat + (l % NPerXdl)*NRepeat + n2
+// i.e. each lane owns NRepeat *consecutive* N elements. This matches the
+// canonical-N decomposition used by the paired host-side B shuffle
+// (shuffle_b_v1 / shuffle_b_permuteN in host/tensor_shuffle_utils.hpp:249-269:
+// view [n/N_Tile][N_Warp][N_Warp_Tile][NRepeat][k/IPA][IPA], permute
+// {0,3,1,4,2,5}), which hands the weight column n = nBlock*N_Tile +
+// wave_n*NPerXdl*NRepeat + lane_n*NRepeat + n2 to exactly this lane/repeat
+// slot. Stores therefore land at canonical global N coordinates, so any
+// N-indexed side input (per-channel scale / D tensor) gathered through the
+// SAME distribution reads the matching canonical n, and any M-indexed side
+// input reads the matching canonical m PROVIDED its window is advanced by
+// MWave*MPerXdl per mIter like the output window (see the scale_m fix below).
+// ---------------------------------------------------------------------------
 template <typename Problem_, typename Policy_ = void>
 struct PermuteNEpilogue
 {
@@ -344,8 +370,22 @@ struct PermuteNEpilogue
                     }
                     else if constexpr(has_scales && !has_scalar_scales)
                     {
-                        const auto sm = static_cast<float>(sm_tile.get_thread_buffer()[dst]);
-                        const auto sn = static_cast<float>(sn_tile.get_thread_buffer()[dst]);
+                        // Indexing both scale tiles at [dst] applies wrong
+                        // per-token scales: the tiles are
+                        // loaded through the same distribution as the store,
+                        // but stepping the n2 Y-coordinate across scale_m's
+                        // length-1 N dimension carries into its M coordinate.
+                        // The scale values delivered by the distribution are:
+                        //   sm at (y0=m2, y1=n2) = scale_m[m + n2]
+                        //   (wrapped within the current MWave*MPerXdl window),
+                        //   sn delivered at (y0=m2, y1=n2) = scale_n[n].
+                        // scale_m does not depend on n and its n2=0 slot is
+                        // always correct, so gather sm at (m2, n2=0). scale_n
+                        // does not depend on m; gather it at (m2=0, n2) for the
+                        // same robustness.
+                        const auto sm =
+                            static_cast<float>(sm_tile.get_thread_buffer()[m_lane * NRepeat]);
+                        const auto sn = static_cast<float>(sn_tile.get_thread_buffer()[n_idx]);
                         v             = static_cast<AccDataType>(v * sm * sn);
                     }
 
@@ -369,6 +409,16 @@ struct PermuteNEpilogue
             static_for<0, NumDTensor, 1>{}([&](auto idx) {
                 move_tile_window(d_dram_windows[idx], {number<MPerXdl * MWave>{}, number<0>{}});
             });
+            // scale_m is M-indexed (per-token); its window must track the
+            // output rows across MRepeat iterations,
+            // otherwise every iteration re-reads rows [0, MPerXdl*MWave) and
+            // MTile > MWave*MPerXdl silently applies wrong per-token scales.
+            // scale_n's view has M-stride 0 (it is N-indexed), so it does not
+            // need to move.
+            if constexpr(has_scales && !has_scalar_scales)
+            {
+                move_tile_window(scale_m_window, {number<MPerXdl * MWave>{}, number<0>{}});
+            }
         });
     }
 };

projects/composablekernel/test/ck_tile/epilogue/CMakeLists.txt

@@ -11,7 +11,9 @@ endif()
 
 if(GPU_TARGETS MATCHES "gfx950")
     add_gtest_executable(test_ck_tile_cshuffle_epilogue_fp8_gfx950 test_cshuffle_epilogue_fp8_gfx950.cpp)
+    add_gtest_executable(test_ck_tile_permuten_epilogue_scale test_permuten_epilogue_scale.cpp)
     if(CK_USE_OCP_FP8)
         target_compile_options(test_ck_tile_cshuffle_epilogue_fp8_gfx950 PRIVATE -DCK_TILE_USE_OCP_FP8)
+        target_compile_options(test_ck_tile_permuten_epilogue_scale PRIVATE -DCK_TILE_USE_OCP_FP8)
     endif()
 endif()

projects/composablekernel/test/ck_tile/epilogue/test_permuten_epilogue_scale.cpp

@@ -0,0 +1,192 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+
+#include <gtest/gtest.h>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/host.hpp"
+#include "ck_tile/host/hip_check_error.hpp"
+#include "ck_tile/ops/elementwise.hpp"
+#include "ck_tile/ops/epilogue/permuten_epilogue.hpp"
+#include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp"
+
+namespace {
+
+constexpr ck_tile::index_t M = 128;
+constexpr ck_tile::index_t N = 256;
+
+using Problem = ck_tile::PermuteNEpilogueProblem<ck_tile::fp8_t,
+                                                 ck_tile::fp8_t,
+                                                 ck_tile::tuple<>,
+                                                 float,
+                                                 ck_tile::half_t,
+                                                 ck_tile::tuple<>,
+                                                 ck_tile::tensor_layout::gemm::RowMajor,
+                                                 ck_tile::element_wise::PassThrough,
+                                                 M,
+                                                 N,
+                                                 1,
+                                                 4,
+                                                 16,
+                                                 16,
+                                                 128,
+                                                 false>;
+
+template <bool ApplyScales>
+__global__ void permuten_epilogue_scale_kernel(const float* input,
+                                               ck_tile::half_t* output,
+                                               float* scale_m,
+                                               float* scale_n)
+{
+    using Epilogue = ck_tile::PermuteNEpilogue<Problem>;
+    using WG       = ck_tile::WarpGemmDispatcher<typename Epilogue::ATypeToUse,
+                                                 typename Epilogue::BTypeToUse,
+                                                 typename Problem::AccDataType,
+                                                 Problem::MPerXdl,
+                                                 Problem::NPerXdl,
+                                                 Problem::KPerXdl,
+                                                 Problem::isCTransposed>;
+
+    constexpr ck_tile::index_t MIterPerWarp = M / (Problem::MWave * Problem::MPerXdl);
+    constexpr ck_tile::index_t NIterPerWarp = N / (Problem::NWave * Problem::NPerXdl);
+
+    constexpr auto outer_distribution = ck_tile::tile_distribution_encoding<
+        ck_tile::sequence<>,
+        ck_tile::tuple<ck_tile::sequence<MIterPerWarp, Problem::MWave>,
+                       ck_tile::sequence<NIterPerWarp, Problem::NWave>>,
+        ck_tile::tuple<ck_tile::sequence<1, 2>>,
+        ck_tile::tuple<ck_tile::sequence<1, 1>>,
+        ck_tile::sequence<1, 2>,
+        ck_tile::sequence<0, 0>>{};
+
+    constexpr auto acc_distribution_encoding =
+        ck_tile::detail::make_embed_tile_distribution_encoding(outer_distribution,
+                                                               typename WG::CWarpDstrEncoding{});
+    constexpr auto acc_distribution =
+        ck_tile::make_static_tile_distribution(acc_distribution_encoding);
+    auto acc_tile = ck_tile::make_static_distributed_tensor<float>(acc_distribution);
+
+    auto input_view = ck_tile::make_naive_tensor_view<ck_tile::address_space_enum::global>(
+        const_cast<float*>(input),
+        ck_tile::make_tuple(M, N),
+        ck_tile::make_tuple(N, 1),
+        ck_tile::number<1>{},
+        ck_tile::number<1>{});
+    auto input_window =
+        ck_tile::make_tile_window(input_view,
+                                  ck_tile::make_tuple(ck_tile::number<M>{}, ck_tile::number<N>{}),
+                                  {0, 0},
+                                  acc_distribution);
+    ck_tile::load_tile(acc_tile, input_window);
+
+    auto output_view = ck_tile::make_naive_tensor_view<ck_tile::address_space_enum::global>(
+        output,
+        ck_tile::make_tuple(M, N),
+        ck_tile::make_tuple(N, 1),
+        ck_tile::number<Epilogue::GetVectorSizeC()>{},
+        ck_tile::number<1>{});
+    auto output_window = ck_tile::make_tile_window(
+        output_view, ck_tile::make_tuple(ck_tile::number<M>{}, ck_tile::number<N>{}), {0, 0});
+
+    auto empty_ds = ck_tile::make_tuple();
+    if constexpr(ApplyScales)
+    {
+        auto scale_m_view = ck_tile::make_naive_tensor_view<ck_tile::address_space_enum::global>(
+            scale_m,
+            ck_tile::make_tuple(M, N),
+            ck_tile::make_tuple(1, 0),
+            ck_tile::number<1>{},
+            ck_tile::number<1>{});
+        auto scale_n_view = ck_tile::make_naive_tensor_view<ck_tile::address_space_enum::global>(
+            scale_n,
+            ck_tile::make_tuple(M, N),
+            ck_tile::make_tuple(0, 1),
+            ck_tile::number<1>{},
+            ck_tile::number<1>{});
+        auto scale_m_window = ck_tile::make_tile_window(
+            scale_m_view, ck_tile::make_tuple(ck_tile::number<M>{}, ck_tile::number<N>{}), {0, 0});
+        auto scale_n_window = ck_tile::make_tile_window(
+            scale_n_view, ck_tile::make_tuple(ck_tile::number<M>{}, ck_tile::number<N>{}), {0, 0});
+
+        Epilogue{}(output_window, acc_tile, empty_ds, nullptr, scale_m_window, scale_n_window);
+    }
+    else
+    {
+        Epilogue{}(output_window, acc_tile, empty_ds, nullptr);
+    }
+}
+
+TEST(PermuteNEpilogueScale, AppliesPerTokenAndPerChannelScalesAtCanonicalCoordinates)
+{
+    ck_tile::HostTensor<float> input({M, N});
+    ck_tile::HostTensor<float> scale_m({M});
+    ck_tile::HostTensor<float> scale_n({N});
+
+    for(ck_tile::index_t m = 0; m < M; ++m)
+    {
+        scale_m.mData[m] = 0.5F + static_cast<float>(m % 31) * 0.015625F;
+        for(ck_tile::index_t n = 0; n < N; ++n)
+        {
+            input(m, n) = 0.5F + static_cast<float>((m * 5 + n * 3) % 16) * 0.0625F;
+        }
+    }
+    for(ck_tile::index_t n = 0; n < N; ++n)
+    {
+        scale_n.mData[n] = 0.75F + static_cast<float>(n % 29) * 0.015625F;
+    }
+
+    ck_tile::DeviceMem input_device(input.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem scale_m_device(scale_m.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem scale_n_device(scale_n.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem unscaled_device(M * N * sizeof(ck_tile::half_t));
+    ck_tile::DeviceMem scaled_device(M * N * sizeof(ck_tile::half_t));
+    input_device.ToDevice(input.data());
+    scale_m_device.ToDevice(scale_m.data());
+    scale_n_device.ToDevice(scale_n.data());
+    unscaled_device.SetZero();
+    scaled_device.SetZero();
+
+    const dim3 grid(1);
+    const dim3 block(Problem::kBlockSize);
+    permuten_epilogue_scale_kernel<false>
+        <<<grid, block>>>(static_cast<const float*>(input_device.GetDeviceBuffer()),
+                          static_cast<ck_tile::half_t*>(unscaled_device.GetDeviceBuffer()),
+                          nullptr,
+                          nullptr);
+    permuten_epilogue_scale_kernel<true>
+        <<<grid, block>>>(static_cast<const float*>(input_device.GetDeviceBuffer()),
+                          static_cast<ck_tile::half_t*>(scaled_device.GetDeviceBuffer()),
+                          static_cast<float*>(scale_m_device.GetDeviceBuffer()),
+                          static_cast<float*>(scale_n_device.GetDeviceBuffer()));
+    HIP_CHECK_ERROR(hipGetLastError());
+    HIP_CHECK_ERROR(hipDeviceSynchronize());
+
+    ck_tile::HostTensor<ck_tile::half_t> unscaled({M, N});
+    ck_tile::HostTensor<ck_tile::half_t> scaled({M, N});
+    unscaled_device.FromDevice(unscaled.data());
+    scaled_device.FromDevice(scaled.data());
+
+    constexpr float tolerance = 0.004F;
+    float max_error           = 0.0F;
+    std::size_t error_count   = 0;
+    for(ck_tile::index_t m = 0; m < M; ++m)
+    {
+        for(ck_tile::index_t n = 0; n < N; ++n)
+        {
+            const float expected =
+                ck_tile::type_convert<float>(unscaled(m, n)) * scale_m.mData[m] * scale_n.mData[n];
+            const float actual = ck_tile::type_convert<float>(scaled(m, n));
+            const float error  = std::abs(actual - expected);
+            max_error          = std::max(max_error, error);
+            error_count += error > tolerance;
+        }
+    }
+
+    EXPECT_EQ(error_count, 0) << "max error: " << max_error;
+}
+
+} // namespace