diff --git a/example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp b/example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp
index 354d774980..a34c3600e0 100644
--- a/example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp
+++ b/example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp
@@ -587,35 +587,32 @@ bool run_no_group_hstu(const ck_tile::ArgParser& arg_parser, bool is_jagged)
     {
         using GemmAccDataType = typename HstuAttentionFwdTypeConfig<InOutDataType>::GemmAccDataType;
 
-        BOOL_SWITCH_3(is_jagged, kIsJagged, use_softmax, kUseSoftmax, use_causal, kUseCausal, [&] {
-            BOOL_SWITCH(is_training, kIsTraining, [&] {
-                constexpr bool kStoreLSE = (kIsTraining && kUseSoftmax);
-                ck_tile::reference_no_group_hstu_attention_fwd<
-                    InOutDataType,
-                    GemmAccDataType,
-                    CompDataType,
-                    kIsJagged,
-                    kUseSoftmax,
-                    kStoreLSE,
-                    kUseCausal>::Run(is_cross_attention,
-                                     q_host,
-                                     k_host,
-                                     v_host,
-                                     o_host_ref,
-                                     lse_host_ref,
-                                     mask_host,
-                                     num_batch,
-                                     scale_s,
-                                     attn_scale,
-                                     max_seqlen_q,
-                                     max_seqlen_kv,
-                                     seq_offsets_q,
-                                     seq_offsets_kv,
-                                     num_targets,
-                                     contextual_seqlen,
-                                     window_size,
-                                     min_full_attn_seqlen);
-            });
+        BOOL_SWITCH_2(is_jagged, kIsJagged, use_causal, kUseCausal, [&] {
+            bool store_lse = (is_training && use_softmax);
+            ck_tile::reference_no_group_hstu_attention_fwd<InOutDataType,
+                                                           GemmAccDataType,
+                                                           CompDataType,
+                                                           kIsJagged,
+                                                           kUseCausal>::Run(is_cross_attention,
+                                                                            use_softmax,
+                                                                            store_lse,
+                                                                            q_host,
+                                                                            k_host,
+                                                                            v_host,
+                                                                            o_host_ref,
+                                                                            lse_host_ref,
+                                                                            mask_host,
+                                                                            num_batch,
+                                                                            scale_s,
+                                                                            attn_scale,
+                                                                            max_seqlen_q,
+                                                                            max_seqlen_kv,
+                                                                            seq_offsets_q,
+                                                                            seq_offsets_kv,
+                                                                            num_targets,
+                                                                            contextual_seqlen,
+                                                                            window_size,
+                                                                            min_full_attn_seqlen);
         });
 
         ck_tile::HostTensor<InOutDataType> o_host(
@@ -1042,36 +1039,34 @@ bool run_group_hstu(const ck_tile::ArgParser& arg_parser, int num_group)
     {
         using GemmAccDataType = typename HstuAttentionFwdTypeConfig<InOutDataType>::GemmAccDataType;
 
-        BOOL_SWITCH_2(use_softmax, kUseSoftmax, use_causal, kUseCausal, [&] {
-            BOOL_SWITCH(is_training, kIsTraining, [&] {
-                constexpr bool kStoreLSE = (kIsTraining && kUseSoftmax);
-                ck_tile::reference_group_hstu_attention_fwd<
-                    InOutDataType,
-                    GemmAccDataType,
-                    CompDataType,
-                    kUseSoftmax,
-                    kStoreLSE,
-                    kUseCausal>::Run(is_cross_attention,
-                                     q_host,
-                                     k_host,
-                                     v_host,
-                                     o_host_ref,
-                                     lse_host_ref,
-                                     mask_host,
-                                     num_batch,
-                                     num_batch / num_group,
-                                     scale_s,
-                                     max_max_seqlen_q,
-                                     max_max_seqlen_kv,
-                                     seq_offsets_q,
-                                     seq_offsets_kv,
-                                     num_targets,
-                                     group_max_seqlens_q,
-                                     group_contextual_seqlens,
-                                     group_window_sizes,
-                                     group_min_full_attn_seqlens,
-                                     group_attn_scales);
-            });
+        BOOL_SWITCH(use_causal, kUseCausal, [&] {
+            bool store_lse = (is_training && use_softmax);
+            ck_tile::reference_group_hstu_attention_fwd<
+                InOutDataType,
+                GemmAccDataType,
+                CompDataType,
+                kUseCausal>::Run(is_cross_attention,
+                                 use_softmax,
+                                 store_lse,
+                                 q_host,
+                                 k_host,
+                                 v_host,
+                                 o_host_ref,
+                                 lse_host_ref,
+                                 mask_host,
+                                 num_batch,
+                                 num_batch / num_group,
+                                 scale_s,
+                                 max_max_seqlen_q,
+                                 max_max_seqlen_kv,
+                                 seq_offsets_q,
+                                 seq_offsets_kv,
+                                 num_targets,
+                                 group_max_seqlens_q,
+                                 group_contextual_seqlens,
+                                 group_window_sizes,
+                                 group_min_full_attn_seqlens,
+                                 group_attn_scales);
         });
 
         ck_tile::HostTensor<InOutDataType> o_host(
diff --git a/example/ck_tile/18_hstu_attention/reference_hstu_attention_fwd.hpp b/example/ck_tile/18_hstu_attention/reference_hstu_attention_fwd.hpp
index 3529d2f64c..2eb885ecf0 100644
--- a/example/ck_tile/18_hstu_attention/reference_hstu_attention_fwd.hpp
+++ b/example/ck_tile/18_hstu_attention/reference_hstu_attention_fwd.hpp
@@ -29,12 +29,12 @@ template <typename InOutDataType,
           typename GemmAccDataType,
           typename CompDataType,
           bool kIsJagged,
-          bool kUseSoftmax,
-          bool kStoreLSE,
           bool kUseCausal>
 struct reference_no_group_hstu_attention_fwd
 {
     static void Run(bool is_cross_attention,
+                    bool use_softmax,
+                    bool store_lse,
                     const HostTensor<InOutDataType>& q_batch_seq_nhead_hdim,
                     const HostTensor<InOutDataType>& k_batch_seq_nhead_hdim,
                     const HostTensor<InOutDataType>& v_batch_seq_nhead_hdim,
@@ -66,7 +66,7 @@ struct reference_no_group_hstu_attention_fwd
             assert(v_batch_seq_nhead_hdim.get_lengths()[0] == 1);
             assert(o_batch_seq_nhead_hdim.get_lengths()[0] == 1);
 
-            if constexpr(kStoreLSE)
+            if(store_lse)
             {
                 assert(lse_batch_seq_nhead.get_lengths()[0] == 1);
             }
@@ -80,7 +80,7 @@ struct reference_no_group_hstu_attention_fwd
             assert(v_batch_seq_nhead_hdim.get_lengths()[0] == num_batch);
             assert(o_batch_seq_nhead_hdim.get_lengths()[0] == num_batch);
 
-            if constexpr(kStoreLSE)
+            if(store_lse)
             {
                 assert(lse_batch_seq_nhead.get_lengths()[0] == num_batch);
             }
@@ -90,7 +90,7 @@ struct reference_no_group_hstu_attention_fwd
         assert(q_batch_seq_nhead_hdim.get_lengths()[1] == k_batch_seq_nhead_hdim.get_lengths()[1]);
         assert(q_batch_seq_nhead_hdim.get_lengths()[1] == v_batch_seq_nhead_hdim.get_lengths()[1]);
         assert(q_batch_seq_nhead_hdim.get_lengths()[1] == o_batch_seq_nhead_hdim.get_lengths()[1]);
-        if constexpr(kStoreLSE)
+        if(store_lse)
         {
             assert(q_batch_seq_nhead_hdim.get_lengths()[1] == lse_batch_seq_nhead.get_lengths()[1]);
         }
@@ -100,7 +100,7 @@ struct reference_no_group_hstu_attention_fwd
         assert(num_head == k_batch_seq_nhead_hdim.get_lengths()[2]);
         assert(num_head == v_batch_seq_nhead_hdim.get_lengths()[2]);
         assert(num_head == o_batch_seq_nhead_hdim.get_lengths()[2]);
-        if constexpr(kStoreLSE)
+        if(store_lse)
         {
             assert(num_head == lse_batch_seq_nhead.get_lengths()[2]);
         }
@@ -258,14 +258,14 @@ struct reference_no_group_hstu_attention_fwd
                         }
                         else
                         {
-                            if constexpr(!kUseSoftmax)
+                            if(!use_softmax)
                                 locals.push_back(ck_tile::type_convert<CompDataType>(0.0f));
                             else
                                 locals.push_back(-ck_tile::numeric<CompDataType>::infinity());
                         };
                     };
 
-                    if constexpr(!kUseSoftmax)
+                    if(!use_softmax)
                     {
                         // SiLu element-wise
                         for(CompDataType& elem : locals)
@@ -292,7 +292,7 @@ struct reference_no_group_hstu_attention_fwd
                                 elem = std::exp(elem - m) / l;
                         }
 
-                        if constexpr(kStoreLSE)
+                        if(store_lse)
                         {
                             lse_batch_seq_nhead(i_batch, sq, i_head) = std::log(l) + m;
                         }
@@ -341,16 +341,13 @@ struct reference_no_group_hstu_attention_fwd
     }
 };
 
-template <typename InOutDataType,
-          typename GemmAccDataType,
-          typename CompDataType,
-          bool kUseSoftmax,
-          bool kStoreLSE,
-          bool kUseCausal>
+template <typename InOutDataType, typename GemmAccDataType, typename CompDataType, bool kUseCausal>
 struct reference_group_hstu_attention_fwd
 {
     static void
     Run(bool is_cross_attention,
+        bool use_softmax,
+        bool store_lse,
         const HostTensor<InOutDataType>& q_batch_seq_nhead_hdim,
         const HostTensor<InOutDataType>& k_batch_seq_nhead_hdim,
         const HostTensor<InOutDataType>& v_batch_seq_nhead_hdim,
@@ -380,7 +377,7 @@ struct reference_group_hstu_attention_fwd
         assert(k_batch_seq_nhead_hdim.get_lengths()[0] == 1);
         assert(v_batch_seq_nhead_hdim.get_lengths()[0] == 1);
         assert(o_batch_seq_nhead_hdim.get_lengths()[0] == 1);
-        if constexpr(kStoreLSE)
+        if(store_lse)
         {
             assert(lse_batch_seq_nhead.get_lengths()[0] == 1);
         }
@@ -389,7 +386,7 @@ struct reference_group_hstu_attention_fwd
         assert(q_batch_seq_nhead_hdim.get_lengths()[1] == k_batch_seq_nhead_hdim.get_lengths()[1]);
         assert(q_batch_seq_nhead_hdim.get_lengths()[1] == v_batch_seq_nhead_hdim.get_lengths()[1]);
         assert(q_batch_seq_nhead_hdim.get_lengths()[1] == o_batch_seq_nhead_hdim.get_lengths()[1]);
-        if constexpr(kStoreLSE)
+        if(store_lse)
         {
             assert(q_batch_seq_nhead_hdim.get_lengths()[1] == lse_batch_seq_nhead.get_lengths()[1]);
         }
@@ -399,7 +396,7 @@ struct reference_group_hstu_attention_fwd
         assert(num_head == k_batch_seq_nhead_hdim.get_lengths()[2]);
         assert(num_head == v_batch_seq_nhead_hdim.get_lengths()[2]);
         assert(num_head == o_batch_seq_nhead_hdim.get_lengths()[2]);
-        if constexpr(kStoreLSE)
+        if(store_lse)
         {
             assert(num_head == lse_batch_seq_nhead.get_lengths()[2]);
         }
@@ -550,14 +547,14 @@ struct reference_group_hstu_attention_fwd
                         }
                         else
                         {
-                            if constexpr(!kUseSoftmax)
+                            if(!use_softmax)
                                 locals.push_back(ck_tile::type_convert<CompDataType>(0.0f));
                             else
                                 locals.push_back(-ck_tile::numeric<CompDataType>::infinity());
                         };
                     };
 
-                    if constexpr(!kUseSoftmax)
+                    if(!use_softmax)
                     {
                         // SiLu element-wise
                         for(CompDataType& elem : locals)
@@ -584,7 +581,7 @@ struct reference_group_hstu_attention_fwd
                                 elem = std::exp(elem - m) / l;
                         }
 
-                        if constexpr(kStoreLSE)
+                        if(store_lse)
                         {
                             lse_batch_seq_nhead(i_batch, sq, i_head) = std::log(l) + m;
                         }