Merge remote-tracking branch 'origin/develop' into aosewski/ck_tile_gemm_policy

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_bf16_n4096_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t,  1, 2, 1,  256, 8,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t,  1, 4, 1,  256, 4,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t,  1, 2, 1, 1024, 2,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t,  1, 4, 1, 1024, 1,  true,  false, false>>(const S&, A);
-
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_bf16_n4096_tp_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t,  1, 2, 1,  256, 8,  true,  false, true>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t,  1, 4, 1,  256, 4,  true,  false, true>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t,  1, 2, 1, 1024, 2,  true,  false, true>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t,  1, 4, 1, 1024, 1,  true,  false, true>>(const S&, A);
-
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_bf16_n512_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  1,  4,  64, 8,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  2,  4,  64, 4,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  4,  4,  64, 2,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  8,  4,  64, 1,  true , false, false>>(const S&, A);
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_bf16_n64_n128_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  1,  4,  64, 1,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  1,  4,  64, 2,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  2,  4,  64, 1,  true , false, false>>(const S&, A);
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_bf16_n768_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  3,  4,  64, 4,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  6,  4,  64, 2,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::bf16_t, 1, 12,  4,  64, 1,  true , false, false>>(const S&, A);
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_fp16_n1024_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-#if 0
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  2,  4,  64, 8,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  4,  4,  64, 4,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  8,  4,  64, 2,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1, 16,  4,  64, 1,  true , false, false>>(const S&, A);
-
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  1, 256, 4,  true , false, false>>(const S&, A);
-#endif
-
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 1, 2,  128, 8,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 2,  128, 4,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 2,  128, 2,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1,  256, 1,  true,  false, false>>(const S&, A);
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_fp16_n1536_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 4,   64, 8,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 2,  128, 4,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 1,  256, 2,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 6, 1,  256, 1,  true,  false, false>>(const S&, A);
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_fp16_n2048_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 1, 1,  256, 8,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 1,  256, 4,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1,  256, 2,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 8, 1,  256, 1,  true,  false, false>>(const S&, A);
-
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_fp16_n256_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 4,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  2,  4,  64, 2,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  4,  4,  64, 1,  true , false, false>>(const S&, A);
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_fp16_n3072_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 1,  128, 8,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 1,  256, 4,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 6, 1,  256, 2,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 1, 1024, 1,  true,  false, false>>(const S&, A);
-
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_fp16_n4096_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 1,  256, 8,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1,  256, 4,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 1, 1024, 2,  true,  false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1, 1024, 1,  true,  false, false>>(const S&, A);
-
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_fp16_n4096_tp_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 1,  256, 8,  true,  false, true>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1,  256, 4,  true,  false, true>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 1, 1024, 2,  true,  false, true>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1, 1024, 1,  true,  false, true>>(const S&, A);
-
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_fp16_n512_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 8,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  2,  4,  64, 4,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  4,  4,  64, 2,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  8,  4,  64, 1,  true , false, false>>(const S&, A);
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_fp16_n64_n128_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 1,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 2,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  2,  4,  64, 1,  true , false, false>>(const S&, A);
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_fp16_n768_instance.cpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include "rmsnorm2d_fwd_instance_common.hpp"
-
-// clang-format off
-//                                                       rm  rn  tm  tn  vn  pd    rms     2p
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  3,  4,  64, 4,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  6,  4,  64, 2,  true , false, false>>(const S&, A);
-template float rmsnorm2d_fwd_<trait_<ck_tile::fp16_t, 1, 12,  4,  64, 1,  true , false, false>>(const S&, A);
-// clang-format on

example/ck_tile/10_rmsnorm2d/instances/rmsnorm2d_fwd_instance_common.hpp

-
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#include <ck_tile/core.hpp>
-#include "rmsnorm2d_fwd.hpp"
-#include <iostream>
-
-#pragma once
-
-using S = ck_tile::stream_config;
-using A = rmsnorm2d_fwd_args;
-
-template <typename DataType_,
-          ck_tile::index_t Repeat_M_,         // each thread repeat along M
-          ck_tile::index_t Repeat_N_,         // each thread repeat along N
-          ck_tile::index_t ThreadPerBlock_M_, // num threads along M
-          ck_tile::index_t ThreadPerBlock_N_, // num threads along N
-          ck_tile::index_t Vector_N_,         // vector size along N
-          bool kPadN_,
-          bool kSaveInvRms_,
-          bool kTwoPass_>
-using trait_ = rmsnorm2d_fwd_traits_<DataType_,
-                                     Repeat_M_,
-                                     Repeat_N_,
-                                     ThreadPerBlock_M_,
-                                     ThreadPerBlock_N_,
-                                     Vector_N_,
-                                     kPadN_,
-                                     kSaveInvRms_,
-                                     kTwoPass_>;
-
-template <typename Traits_>
-float rmsnorm2d_fwd_(const S& s, A a)
-{
-    using DataType = typename Traits_::DataType;
-
-    using PipelineProblem =
-        ck_tile::Rmsnorm2dFwdPipelineProblem<typename RmsnormTypeConfig<DataType>::XDataType,
-                                             typename RmsnormTypeConfig<DataType>::GammaDataType,
-                                             typename RmsnormTypeConfig<DataType>::ComputeDataType,
-                                             typename RmsnormTypeConfig<DataType>::YDataType,
-                                             typename RmsnormTypeConfig<DataType>::InvRmsDataType,
-                                             typename Traits_::Shape,
-                                             Traits_::kPadN,
-                                             Traits_::kSaveInvRms,
-                                             Traits_::kTwoPass>;
-
-    using OnePassPipeline = ck_tile::Rmsnorm2dFwdPipelineOnePass<PipelineProblem>;
-    using TwoPassPipeline = ck_tile::Rmsnorm2dFwdPipelineTwoPass<PipelineProblem>;
-    using Pipeline        = std::conditional_t<Traits_::kTwoPass, TwoPassPipeline, OnePassPipeline>;
-
-    using Kernel = ck_tile::Rmsnorm2dFwd<Pipeline>;
-
-    const dim3 grids                       = Kernel::GridSize(a);
-    constexpr dim3 blocks                  = Kernel::BlockSize();
-    constexpr ck_tile::index_t kBlockPerCu = 1;
-
-    auto kargs = Kernel::MakeKargs(a);
-    if(s.log_level_ > 0)
-        std::cout << ", " << Kernel::GetName() << std::flush;
-
-    return ck_tile::launch_kernel(
-        s, ck_tile::make_kernel<blocks.x, kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
-}

example/ck_tile/10_rmsnorm2d/rmsnorm2d_fwd.hpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -8,27 +8,34 @@
 #include "ck_tile/ops/rmsnorm2d.hpp"
 #include <string>
 
-template <typename DataType>
+template <typename InType,
+          typename OutType,
+          typename SmoothScaleDataType_,
+          typename YScaleDataType_>
 struct RmsnormTypeConfig;
 
-template <>
-struct RmsnormTypeConfig<ck_tile::half_t>
+template <typename OutType, typename SmoothScaleDataType_, typename YScaleDataType_>
+struct RmsnormTypeConfig<ck_tile::half_t, OutType, SmoothScaleDataType_, YScaleDataType_>
 {
     using XDataType           = ck_tile::half_t;
-    using YDataType       = ck_tile::half_t;
+    using YDataType           = OutType;
     using GammaDataType       = ck_tile::half_t;
     using InvRmsDataType      = ck_tile::half_t;
     using ComputeDataType     = float;
+    using SmoothScaleDataType = SmoothScaleDataType_;
+    using YScaleDataType      = YScaleDataType_;
 };
 
-template <>
-struct RmsnormTypeConfig<ck_tile::bf16_t>
+template <typename OutType, typename SmoothScaleDataType_, typename YScaleDataType_>
+struct RmsnormTypeConfig<ck_tile::bf16_t, OutType, SmoothScaleDataType_, YScaleDataType_>
 {
     using XDataType           = ck_tile::bf16_t;
-    using YDataType       = ck_tile::bf16_t;
+    using YDataType           = OutType;
     using GammaDataType       = ck_tile::bf16_t;
     using InvRmsDataType      = ck_tile::bf16_t;
     using ComputeDataType     = float;
+    using SmoothScaleDataType = SmoothScaleDataType_;
+    using YScaleDataType      = YScaleDataType_;
 };
 
 // runtime args
@@ -36,82 +43,24 @@ struct rmsnorm2d_fwd_args : public ck_tile::Rmsnorm2dFwdHostArgs
 {
 };
 
-// this is used to pattern-match internl kernel implementation, not to instantiate kernel
-template <typename DataType_,
-          ck_tile::index_t Repeat_M_,         // each thread repeat along M
-          ck_tile::index_t Repeat_N_,         // each thread repeat along N
-          ck_tile::index_t ThreadPerBlock_M_, // num threads along M
-          ck_tile::index_t ThreadPerBlock_N_, // num threads along N
-          ck_tile::index_t Vector_N_,         // vector size along N
-          bool kPadN_,
-          bool kSaveInvRms_,
-          bool kTwoPass_>
-struct rmsnorm2d_fwd_traits_
-{
-    using DataType = ck_tile::remove_cvref_t<DataType_>;
-
-    static constexpr bool is_warp_per_row = ThreadPerBlock_N_ <= warpSize;
-    static_assert((ThreadPerBlock_M_ * ThreadPerBlock_N_) % warpSize == 0);
-    static constexpr ck_tile::index_t total_warps =
-        (ThreadPerBlock_M_ * ThreadPerBlock_N_) / warpSize;
-
-    // num of warps along m
-    static constexpr ck_tile::index_t BlockWarps_M = []() {
-        if constexpr(is_warp_per_row)
-        {
-            static_assert(warpSize % ThreadPerBlock_N_ == 0);
-            return total_warps * (warpSize / ThreadPerBlock_N_);
-        }
-        else
-        {
-            // static_assert(warpSize % ThreadPerBlock_M_ == 0);
-            return total_warps / (ThreadPerBlock_N_ / warpSize);
-        }
-    }();
-
-    // num of warps along n
-    static constexpr ck_tile::index_t BlockWarps_N = []() {
-        if constexpr(is_warp_per_row)
-        {
-            static_assert(warpSize % ThreadPerBlock_N_ == 0);
-            return 1;
-        }
-        else
-        {
-            static_assert(ThreadPerBlock_N_ % warpSize == 0);
-            return ThreadPerBlock_N_ / warpSize;
-        }
-    }();
-
-    static constexpr ck_tile::index_t Repeat_M = Repeat_M_;
-    static constexpr ck_tile::index_t Repeat_N = Repeat_N_;
-
-    static constexpr ck_tile::index_t Block_M = Repeat_M_ * ThreadPerBlock_M_;
-    static constexpr ck_tile::index_t Block_N = Repeat_N_ * ThreadPerBlock_N_ * Vector_N_;
-
-    static constexpr ck_tile::index_t Warp_M = ThreadPerBlock_M_ / BlockWarps_M;
-    static constexpr ck_tile::index_t Warp_N = ThreadPerBlock_N_ / BlockWarps_N * Vector_N_;
-
-    using BlockTile  = ck_tile::sequence<Block_M, Block_N>;
-    using BlockWarps = ck_tile::sequence<BlockWarps_M, BlockWarps_N>;
-    using WarpTile   = ck_tile::sequence<Warp_M, Warp_N>;
-    using Vector     = ck_tile::sequence<1, Vector_N_>;
-
-    using Shape = ck_tile::Generic2dBlockShape<BlockTile, BlockWarps, WarpTile, Vector>;
-
-    static constexpr bool kPadN       = kPadN_;
-    static constexpr bool kSaveInvRms = kSaveInvRms_;
-    static constexpr bool kTwoPass    = kTwoPass_;
-};
-
 template <typename Traits_>
 float rmsnorm2d_fwd_(const ck_tile::stream_config& s, rmsnorm2d_fwd_args a);
 
 // This is the public API, will be generated by script
 struct rmsnorm2d_fwd_traits
 {
-    std::string data_type;
+    std::string prec_i; // input precision
+    std::string prec_o; // output precision
+
+    // if fused_quant == 1, need set prec_sm/prec_sy to proper string, otherwise can set
+    // arbitrary(will skip check) if fused_quant == 2, need set prec_sy to proper string, otherwise
+    // can set arbitrary(will skip check)
+    std::string prec_sm; // x-scale, used for [1*N] input smooth quant
+    std::string prec_sy; // y-scale, used for [M*1] output for next layer
+
     bool save_rms;
+    int fused_add;   // 0:no-add, 1:pre-add-store, 2:pre-add
+    int fused_quant; // 0:no-sweep, 1:smooth-dynamic-quant, 2:dynamic-quant
 };
 
 float rmsnorm2d_fwd(rmsnorm2d_fwd_traits, rmsnorm2d_fwd_args, const ck_tile::stream_config&);