From f72822d3a0b2f5565d74aa20896137b64944dcf4 Mon Sep 17 00:00:00 2001
From: ltqin <letaoqin@amd.com>
Date: Mon, 3 Nov 2025 12:31:45 +0000
Subject: [PATCH] change squant to quant type

---
 example/ck_tile/01_fmha/example_fmha_fwd.cpp | 20 ++++++--------------
 example/ck_tile/01_fmha/fmha_fwd_runner.hpp  | 10 +++++-----
 2 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/example/ck_tile/01_fmha/example_fmha_fwd.cpp b/example/ck_tile/01_fmha/example_fmha_fwd.cpp
index c27a5ce1ae..d7dc696a29 100644
--- a/example/ck_tile/01_fmha/example_fmha_fwd.cpp
+++ b/example/ck_tile/01_fmha/example_fmha_fwd.cpp
@@ -50,13 +50,10 @@ auto create_args(int argc, char* argv[])
                 "scale factor of S. 0 means equal to 1/sqrt(hdim).\n"
                 "note when squant=1, this value will be modified")
         .insert("logits_soft_cap", "0", "attention logits soft capping value.")
-        .insert("squant",
-                "auto",
-                "if using static quantization fusion or not. auto: fp8 will default use squant, "
-                "other will not\n"
-                "0: no static quant(not implemented) 1: apply scale_p and scale_o with respect to "
-                "P and O.\n"
-                "calculate scale_s, scale_p, scale_o auto")
+        .insert("quant",
+                "0",
+                "if using quantization or not.\n"
+                "0: no quant 1: per-tensor quant 2: block scaling.\n")
         .insert("iperm",
                 "1",
                 "permute input\n"
@@ -162,12 +159,7 @@ auto run(const ck_tile::ArgParser& arg_parser)
     std::string init_method          = arg_parser.get_str("init");
     uint32_t seed                    = arg_parser.get_uint32("seed");
 
-    bool squant = [&]() {
-        if(arg_parser.get_str("squant") == "auto")
-            return std::is_same_v<DataTypeConfig, FmhaFwdFp8>;
-        else
-            return arg_parser.get_bool("squant");
-    }();
+    int quant = arg_parser.get_bool("quant");
 
     ck_tile::stream_config stream_config{nullptr,
                                          true,
@@ -208,7 +200,7 @@ auto run(const ck_tile::ArgParser& arg_parser)
                                         drop_offset,
                                         drop_prefs,
                                         mask_str,
-                                        squant,
+                                        quant,
                                         is_rotary_interleaved,
                                         num_splits,
                                         init_method,
diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp
index 0703af71e3..bc2c6a5ef2 100644
--- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp
+++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp
@@ -178,7 +178,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
                         uint64_t drop_offset,
                         bool drop_prefs,
                         std::string mask_str,
-                        bool squant,
+                        int quant,
                         bool is_rotary_interleaved,
                         ck_tile::index_t num_splits,
                         std::string init_method,
@@ -705,7 +705,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
 
     float scale_p = 1.f;
     float scale_o = 1.f;
-    if(squant)
+    if(quant)
     {
         float q_dtype_max = ck_tile::type_convert<float>(ck_tile::numeric<QDataType>::max());
         float k_dtype_max = ck_tile::type_convert<float>(ck_tile::numeric<KDataType>::max());
@@ -816,7 +816,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
               << (seqlen_kpads[0] < 0 ? ""
                                       : (std::string("(") + std::to_string(seqlen_kpads[0]) + ")"))
               << ", d:" << hdim_q << "/" << hdim_v << ", scale_s:" << scale_s << ", bias:" << bias
-              << ", p_drop:" << p_drop << ", lse:" << lse << ", squant:" << squant
+              << ", p_drop:" << p_drop << ", lse:" << lse << ", quant:" << quant
               << ", mask:" << mask << ", v:" << (is_v_rowmajor ? "r" : "c");
 #if CK_TILE_FMHA_FWD_APPENDKV_API
     if(0 < rotary_dim)
@@ -908,7 +908,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
             traits.mask_type           = mask.type;
             traits.bias_type           = bias.type;
             traits.has_lse             = lse;
-            traits.do_fp8_static_quant = squant;
+            traits.do_fp8_static_quant = quant == 1;
 
             if constexpr(std::is_same_v<fmha_fwd_traits, std::decay_t<decltype(traits)>>)
             {
@@ -1774,7 +1774,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
                                    scale_s,
                                    p_drop,
                                    lse,
-                                   squant,
+                                   quant,
                                    bias.type == bias_enum::elementwise_bias
                                        ? "elementwise_bias"
                                        : (bias.type == bias_enum::alibi ? "alibi" : "no_bias"),