Add block_size=32 support and fix int32 overflow in unified attention

- Add BlockSize_ template parameter to all kernel trait structs,
  enabling block_size=32 instances for d64/GQA-8 (medium/small/tiny tiers).
- 12 new instance files for block_size=32 (bf16/fp16 x mask/nmask x 3 tiers).
- Dispatch branches on page_blk_size to select bs32 vs bs64 instances.
- Fix potential int32 overflow in KV cache tensor view by using long_index_t
  for sizes and strides (overflow threshold was ~66K blocks for d64/GQA-8).

Made-with: Cursor

example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_kernel_traits<unified_attention_args::data_type_enum::bf16, true, 64, 128, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode_s.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_small_kernel_traits<unified_attention_args::data_type_enum::bf16, true, 64, 64, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode_t.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_tiny_kernel_traits<unified_attention_args::data_type_enum::bf16, true, 64, 16, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_kernel_traits<unified_attention_args::data_type_enum::bf16, false, 64, 128, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode_s.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_small_kernel_traits<unified_attention_args::data_type_enum::bf16, false, 64, 64, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode_t.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_tiny_kernel_traits<unified_attention_args::data_type_enum::bf16, false, 64, 16, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_kernel_traits<unified_attention_args::data_type_enum::fp16, true, 64, 128, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode_s.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_small_kernel_traits<unified_attention_args::data_type_enum::fp16, true, 64, 64, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode_t.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_tiny_kernel_traits<unified_attention_args::data_type_enum::fp16, true, 64, 16, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_kernel_traits<unified_attention_args::data_type_enum::fp16, false, 64, 128, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode_s.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_small_kernel_traits<unified_attention_args::data_type_enum::fp16, false, 64, 64, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode_t.cpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "unified_attention.hpp"
+#include "unified_attention_impl.hpp"
+
+namespace ck_tile {
+
+using kernel_traits =
+    unified_attention_decode_tiny_kernel_traits<unified_attention_args::data_type_enum::fp16, false, 64, 16, 8, 32>;
+
+INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits)
+
+} // namespace ck_tile

example/ck_tile/42_unified_attention/unified_attention.cpp

@@ -26,7 +26,7 @@ std::ostream& operator<<(std::ostream& stream,
         return unified_attention_kernel_dispatch<kernel_traits>(args, config); \
     }
 
-// Dispatch macros for three tile tiers.
+// Dispatch macros for three tile tiers (default block_size).
 #define DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM(DType, IsMask, HSize, BM, NQPKV) \
     { \
         using kernel_traits = unified_attention_decode_kernel_traits<DType, IsMask, HSize, BM, NQPKV>; \
@@ -45,6 +45,25 @@ std::ostream& operator<<(std::ostream& stream,
         return unified_attention_kernel_dispatch_decode<kernel_traits>(args, config); \
     }
 
+// block_size=32 dispatch macros (6th template arg = 32).
+#define DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM_BS32(DType, IsMask, HSize, BM, NQPKV) \
+    { \
+        using kernel_traits = unified_attention_decode_kernel_traits<DType, IsMask, HSize, BM, NQPKV, 32>; \
+        return unified_attention_kernel_dispatch<kernel_traits>(args, config); \
+    }
+
+#define DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL_BS32(DType, IsMask, HSize, BM, NQPKV) \
+    { \
+        using kernel_traits = unified_attention_decode_small_kernel_traits<DType, IsMask, HSize, BM, NQPKV, 32>; \
+        return unified_attention_kernel_dispatch_decode<kernel_traits>(args, config); \
+    }
+
+#define DISPATCH_UNIFIED_ATTENTION_DECODE_TINY_BS32(DType, IsMask, HSize, BM, NQPKV) \
+    { \
+        using kernel_traits = unified_attention_decode_tiny_kernel_traits<DType, IsMask, HSize, BM, NQPKV, 32>; \
+        return unified_attention_kernel_dispatch_decode<kernel_traits>(args, config); \
+    }
+
 enum class tile_tier { large, medium, small, tiny };
 
 static tile_tier select_tile_tier(const unified_attention_args& args)
@@ -88,51 +107,81 @@ std::pair<bool, float> unified_attention(const unified_attention_args& args,
     // d64, GQA-8 (num_queries_per_kv == 8)
     if(args.hdim == 64 && args.num_queries_per_kv == 8)
     {
+        const bool use_bs32 = (args.page_blk_size < 64);
+
         if(tier == tile_tier::tiny)
         {
-            // Tiny decode: 1 warp, 16x16 MFMA, kBlockM=16 (kBlockQ=2)
             if(args.data_type == unified_attention_args::data_type_enum::fp16)
             {
-                if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::fp16, false, 64, 16, 8)
-                else         DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::fp16, true,  64, 16, 8)
+                if(use_bs32) {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_TINY_BS32(unified_attention_args::data_type_enum::fp16, false, 64, 16, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_TINY_BS32(unified_attention_args::data_type_enum::fp16, true,  64, 16, 8)
+                } else {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::fp16, false, 64, 16, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::fp16, true,  64, 16, 8)
+                }
             }
             else if(args.data_type == unified_attention_args::data_type_enum::bf16)
             {
-                if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::bf16, false, 64, 16, 8)
-                else         DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::bf16, true,  64, 16, 8)
+                if(use_bs32) {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_TINY_BS32(unified_attention_args::data_type_enum::bf16, false, 64, 16, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_TINY_BS32(unified_attention_args::data_type_enum::bf16, true,  64, 16, 8)
+                } else {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::bf16, false, 64, 16, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_TINY(unified_attention_args::data_type_enum::bf16, true,  64, 16, 8)
+                }
             }
         }
         else if(tier == tile_tier::small)
         {
-            // Small decode: 2 warps, kBlockM=64 (kBlockQ=8)
             if(args.data_type == unified_attention_args::data_type_enum::fp16)
             {
-                if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL(unified_attention_args::data_type_enum::fp16, false, 64, 64, 8)
-                else         DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL(unified_attention_args::data_type_enum::fp16, true,  64, 64, 8)
+                if(use_bs32) {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL_BS32(unified_attention_args::data_type_enum::fp16, false, 64, 64, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL_BS32(unified_attention_args::data_type_enum::fp16, true,  64, 64, 8)
+                } else {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL(unified_attention_args::data_type_enum::fp16, false, 64, 64, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL(unified_attention_args::data_type_enum::fp16, true,  64, 64, 8)
+                }
             }
             else if(args.data_type == unified_attention_args::data_type_enum::bf16)
             {
-                if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL(unified_attention_args::data_type_enum::bf16, false, 64, 64, 8)
-                else         DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL(unified_attention_args::data_type_enum::bf16, true,  64, 64, 8)
+                if(use_bs32) {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL_BS32(unified_attention_args::data_type_enum::bf16, false, 64, 64, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL_BS32(unified_attention_args::data_type_enum::bf16, true,  64, 64, 8)
+                } else {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL(unified_attention_args::data_type_enum::bf16, false, 64, 64, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL(unified_attention_args::data_type_enum::bf16, true,  64, 64, 8)
+                }
             }
         }
         else if(tier == tile_tier::medium)
         {
-            // Medium: 4 warps, kBlockM=128 (kBlockQ=16)
             if(args.data_type == unified_attention_args::data_type_enum::fp16)
             {
-                if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM(unified_attention_args::data_type_enum::fp16, false, 64, 128, 8)
-                else         DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM(unified_attention_args::data_type_enum::fp16, true,  64, 128, 8)
+                if(use_bs32) {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM_BS32(unified_attention_args::data_type_enum::fp16, false, 64, 128, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM_BS32(unified_attention_args::data_type_enum::fp16, true,  64, 128, 8)
+                } else {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM(unified_attention_args::data_type_enum::fp16, false, 64, 128, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM(unified_attention_args::data_type_enum::fp16, true,  64, 128, 8)
+                }
             }
             else if(args.data_type == unified_attention_args::data_type_enum::bf16)
             {
-                if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM(unified_attention_args::data_type_enum::bf16, false, 64, 128, 8)
-                else         DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM(unified_attention_args::data_type_enum::bf16, true,  64, 128, 8)
+                if(use_bs32) {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM_BS32(unified_attention_args::data_type_enum::bf16, false, 64, 128, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM_BS32(unified_attention_args::data_type_enum::bf16, true,  64, 128, 8)
+                } else {
+                    if(!is_mask) DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM(unified_attention_args::data_type_enum::bf16, false, 64, 128, 8)
+                    else         DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM(unified_attention_args::data_type_enum::bf16, true,  64, 128, 8)
+                }
             }
         }
         else
         {
             // Large prefill: 8 warps, kBlockM=256 (kBlockQ=32)
+            // No bs32 variant -- NumIssues < 1 for 8-warp tier with block_size=32.
             if(args.data_type == unified_attention_args::data_type_enum::fp16)
             {
                 if(!is_mask) DISPATCH_UNIFIED_ATTENTION(unified_attention_args::data_type_enum::fp16, false, 64, 256, 8)
@@ -152,6 +201,9 @@ std::pair<bool, float> unified_attention(const unified_attention_args& args,
     return std::make_pair(false, -1.f);
 }
 
+#undef DISPATCH_UNIFIED_ATTENTION_DECODE_TINY_BS32
+#undef DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL_BS32
+#undef DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM_BS32
 #undef DISPATCH_UNIFIED_ATTENTION_DECODE_TINY
 #undef DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL
 #undef DISPATCH_UNIFIED_ATTENTION_DECODE_MEDIUM

example/ck_tile/42_unified_attention/unified_attention_impl.hpp

@@ -61,12 +61,13 @@ struct unified_attention_problem_traits<unified_attention_args::data_type_enum::
     using lse_dtype  = float;
 };
 
-// Parameterized kernel traits: DataType, IsMasking, HeadSize, BlockM, NumQueriesPerKV
+// Parameterized kernel traits: DataType, IsMasking, HeadSize, BlockM, NumQueriesPerKV, BlockSize
 template <unified_attention_args::data_type_enum DataType,
           bool IsMasking,
           index_t HeadSize_     = 128,
           index_t BlockM_       = 256,
-          index_t NumQPerKV_    = 1>
+          index_t NumQPerKV_    = 1,
+          index_t BlockSize_    = (HeadSize_ <= 64) ? 64 : 32>
 struct unified_attention_kernel_traits
 {
     static constexpr auto date_type  = DataType;
@@ -74,9 +75,7 @@ struct unified_attention_kernel_traits
 
     static constexpr index_t kBlockM    = BlockM_;
     static constexpr index_t HEAD_SIZE  = HeadSize_;
-    // On gfx950 with 128-bit loads (KVector=8), NumIssues = kPageBlockSize*HeadSize/4096.
-    // For HeadSize<=64 we need kPageBlockSize>=64 to keep NumIssues>=1.
-    static constexpr index_t BLOCK_SIZE = (HEAD_SIZE <= 64) ? 64 : 32;
+    static constexpr index_t BLOCK_SIZE = BlockSize_;
 
     static constexpr index_t num_queries_per_kv = NumQPerKV_;
     static constexpr index_t kBlockQ            = kBlockM / num_queries_per_kv;
@@ -137,7 +136,8 @@ template <unified_attention_args::data_type_enum DataType,
           bool IsMasking,
           index_t HeadSize_  = 128,
           index_t BlockM_    = 128,
-          index_t NumQPerKV_ = 1>
+          index_t NumQPerKV_ = 1,
+          index_t BlockSize_ = (HeadSize_ <= 64) ? 64 : 32>
 struct unified_attention_decode_kernel_traits
 {
     static constexpr auto date_type  = DataType;
@@ -145,7 +145,7 @@ struct unified_attention_decode_kernel_traits
 
     static constexpr index_t kBlockM    = BlockM_;
     static constexpr index_t HEAD_SIZE  = HeadSize_;
-    static constexpr index_t BLOCK_SIZE = (HEAD_SIZE <= 64) ? 64 : 32;
+    static constexpr index_t BLOCK_SIZE = BlockSize_;
 
     static constexpr index_t num_queries_per_kv = NumQPerKV_;
     static constexpr index_t kBlockQ            = kBlockM / num_queries_per_kv;
@@ -197,7 +197,8 @@ template <unified_attention_args::data_type_enum DataType,
           bool IsMasking,
           index_t HeadSize_  = 64,
           index_t BlockM_    = 64,
-          index_t NumQPerKV_ = 8>
+          index_t NumQPerKV_ = 8,
+          index_t BlockSize_ = (HeadSize_ <= 64) ? 64 : 32>
 struct unified_attention_decode_small_kernel_traits
 {
     static constexpr auto date_type  = DataType;
@@ -205,7 +206,7 @@ struct unified_attention_decode_small_kernel_traits
 
     static constexpr index_t kBlockM    = BlockM_;
     static constexpr index_t HEAD_SIZE  = HeadSize_;
-    static constexpr index_t BLOCK_SIZE = (HEAD_SIZE <= 64) ? 64 : 32;
+    static constexpr index_t BLOCK_SIZE = BlockSize_;
 
     static constexpr index_t num_queries_per_kv = NumQPerKV_;
     static constexpr index_t kBlockQ            = kBlockM / num_queries_per_kv;
@@ -259,7 +260,8 @@ template <unified_attention_args::data_type_enum DataType,
           bool IsMasking,
           index_t HeadSize_  = 64,
           index_t BlockM_    = 16,
-          index_t NumQPerKV_ = 8>
+          index_t NumQPerKV_ = 8,
+          index_t BlockSize_ = (HeadSize_ <= 64) ? 64 : 32>
 struct unified_attention_decode_tiny_kernel_traits
 {
     static constexpr auto date_type  = DataType;
@@ -267,7 +269,7 @@ struct unified_attention_decode_tiny_kernel_traits
 
     static constexpr index_t kBlockM    = BlockM_;
     static constexpr index_t HEAD_SIZE  = HeadSize_;
-    static constexpr index_t BLOCK_SIZE = (HEAD_SIZE <= 64) ? 64 : 32;
+    static constexpr index_t BLOCK_SIZE = BlockSize_;
 
     static constexpr index_t num_queries_per_kv = NumQPerKV_;
     static constexpr index_t kBlockQ            = kBlockM / num_queries_per_kv;

include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp

@@ -310,7 +310,7 @@ struct UnifiedAttentionKernel
 
         // TODO sliding window
         const index_t num_blocks_start = 0;
-        index_t kv_head_offset         = kv_head_idx * kargs.stride_k_cache_2;
+        long_index_t kv_head_offset    = static_cast<long_index_t>(kv_head_idx) * kargs.stride_k_cache_2;
 
         // Q/K/V DRAM and DRAM window
         index_t q_ptr_offset_0 = cur_batch_in_all_start_index *
@@ -376,17 +376,19 @@ struct UnifiedAttentionKernel
                              {query_pos * num_queries_per_kv, 0});
 
         const auto k_dram = [&]() {
-            // HEAD dim is skipped as defined in the ptrs
+            // Use long_index_t for size/strides to prevent int32 overflow
+            // when row * stride exceeds 2^31 (happens at ~66K blocks for d64/GQA-8).
             const auto k_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                 k_ptr,
-                make_tuple(kargs.num_blks * kargs.page_size, kHeadDim),
-                make_tuple(kargs.stride_k_cache_1, kargs.stride_k_cache_3),
+                make_tuple(static_cast<long_index_t>(kargs.num_blks) * kargs.page_size,
+                           static_cast<long_index_t>(kHeadDim)),
+                make_tuple(static_cast<long_index_t>(kargs.stride_k_cache_1),
+                           static_cast<long_index_t>(kargs.stride_k_cache_3)),
                 number<UnifiedAttentionPipeline::kAlignmentK>{},
                 number<1>{});
 
             const auto k_dram_pad =
                 pad_tensor_view(k_dram_naive,
-                                // TODO can the kPageBlockSize_RAW needs padding?
                                 make_tuple(kPageBlockSize, kHeadDimPadded),
                                 sequence<false, kPadHeadDimQ>{});
 
@@ -399,8 +401,10 @@ struct UnifiedAttentionKernel
         const auto v_dram = [&]() {
             const auto v_dram_naive = make_naive_tensor_view<address_space_enum::global>(
                 v_ptr,
-                make_tuple(kargs.num_blks * kargs.page_size, kHeadDim),
-                make_tuple(kargs.stride_v_cache_1, kargs.stride_v_cache_3),
+                make_tuple(static_cast<long_index_t>(kargs.num_blks) * kargs.page_size,
+                           static_cast<long_index_t>(kHeadDim)),
+                make_tuple(static_cast<long_index_t>(kargs.stride_v_cache_1),
+                           static_cast<long_index_t>(kargs.stride_v_cache_3)),
                 number<UnifiedAttentionPipeline::kAlignmentV>{},
                 number<1>{});
 

include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp

@@ -353,7 +353,6 @@ struct UnifiedAttentionPipeline
             }
         }
 
-        // TODO check correctness of this
         index_t i_total_loops  = num_blocks_start;
         const index_t PageSize = kv_page_size_in_blocks * kPageBlockSize;
         const ck_tile::index_t* block_tables_ptr_ =
@@ -484,7 +483,6 @@ struct UnifiedAttentionPipeline
 
         auto V_mem_load = [&](auto v_lds_write_idx) {
             async_load_tile_raw(v_lds_window_store(v_lds_write_idx), v_dram_window);
-            // prefetch next V tile (only if not at the end of loop)
             v_block_idx++;
 
             index_t v_page_blk_idx =
@@ -493,7 +491,6 @@ struct UnifiedAttentionPipeline
                 {v_page_blk_idx * PageSize +
                      (v_block_idx % kv_page_size_in_blocks) * kPageBlockSize,
                  0});
-            // we assume that v load is always after k
         };
 
         auto K_lds_load = [&](auto k_lds_read_idx) {