From 4cd3432361791ddd780b083b12f23053b3f9ae20 Mon Sep 17 00:00:00 2001
From: "PoYen, Chen" <PoYen.Chen@amd.com>
Date: Sun, 18 Aug 2024 18:27:37 +0000
Subject: [PATCH] Avoid using too small rotary_cos & rotary_sin

---
 example/ck_tile/01_fmha/fmha_fwd.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/example/ck_tile/01_fmha/fmha_fwd.cpp b/example/ck_tile/01_fmha/fmha_fwd.cpp
index 6e9133806e..2389e763dd 100644
--- a/example/ck_tile/01_fmha/fmha_fwd.cpp
+++ b/example/ck_tile/01_fmha/fmha_fwd.cpp
@@ -603,8 +603,8 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                    : std::array<ck_tile::index_t, 2>{batch, nhead})
             : std::array<ck_tile::index_t, 2>{1, 1});
 
-    auto [rotary_cos_host, rotary_sin_host] =
-        generate_rotary_cos_sin<KDataType>(shape_seqlen_k, rotary_dim, seed);
+    auto [rotary_cos_host, rotary_sin_host] = generate_rotary_cos_sin<KDataType>(
+        std::max(shape_seqlen_q, shape_seqlen_k), rotary_dim, seed);
 
     ck_tile::HostTensor<LSEDataType> lse_acc_host(
         1 < num_splits || use_kvcache