diff --git a/example/ck_tile/18_hstu_attention/example_hstu_attention.cpp b/example/ck_tile/18_hstu_attention/example_hstu_attention.cpp
index 2052b6df86..da579605c1 100644
--- a/example/ck_tile/18_hstu_attention/example_hstu_attention.cpp
+++ b/example/ck_tile/18_hstu_attention/example_hstu_attention.cpp
@@ -290,7 +290,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
         // only consider num_batch values even if more values are provided by the user
         for(int i = 0; i < num_batch; i++)
         {
-            max_uih_seqlen = max(max_uih_seqlen, seq_lengths_q[i]);
+            max_uih_seqlen = max(max_uih_seqlen, max(seq_lengths_q[i], seq_lengths_kv[i]));
         };
     }
     else