From c3d27abfb8dc5ff4c3ad235b882fc1b08536ec74 Mon Sep 17 00:00:00 2001
From: Juuso Korhonen <40278371+juuso-oskari@users.noreply.github.com>
Date: Tue, 14 Oct 2025 09:49:54 +0000
Subject: [PATCH] fix q window

---
 .../ops/unified_attention/kernel/unified_attention_kernel.hpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp b/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp
index 152af051f8..ba2da985f4 100644
--- a/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp
+++ b/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp
@@ -376,8 +376,8 @@ struct FmhaFwdV3Kernel
         // stride for dim 0 (num_queries_per_kv * head_dim, head_dim, 1)
         auto q_dram_window = make_tile_window(
             q_dram,
-            make_tuple(BLOCK_Q, HEAD_SIZE_PADDED),
-            {0, 0}
+            make_tuple(BLOCK_Q * num_queries_per_kv, HEAD_SIZE_PADDED),
+            {query_pos * num_queries_per_kv, 0}
         );
         
         const auto k_dram = [&]() {