From 5b591af58efa5a8d59942fac8ce0412c9f5caabd Mon Sep 17 00:00:00 2001 From: Amir Ghamarian Date: Mon, 30 Mar 2026 14:16:17 +0000 Subject: [PATCH] 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 --- ...tention_d64_bf16_mask_gqa8_bs32_decode.cpp | 14 ++++ ...ntion_d64_bf16_mask_gqa8_bs32_decode_s.cpp | 14 ++++ ...ntion_d64_bf16_mask_gqa8_bs32_decode_t.cpp | 14 ++++ ...ention_d64_bf16_nmask_gqa8_bs32_decode.cpp | 14 ++++ ...tion_d64_bf16_nmask_gqa8_bs32_decode_s.cpp | 14 ++++ ...tion_d64_bf16_nmask_gqa8_bs32_decode_t.cpp | 14 ++++ ...tention_d64_fp16_mask_gqa8_bs32_decode.cpp | 14 ++++ ...ntion_d64_fp16_mask_gqa8_bs32_decode_s.cpp | 14 ++++ ...ntion_d64_fp16_mask_gqa8_bs32_decode_t.cpp | 14 ++++ ...ention_d64_fp16_nmask_gqa8_bs32_decode.cpp | 14 ++++ ...tion_d64_fp16_nmask_gqa8_bs32_decode_s.cpp | 14 ++++ ...tion_d64_fp16_nmask_gqa8_bs32_decode_t.cpp | 14 ++++ .../unified_attention.cpp | 84 +++++++++++++++---- .../unified_attention_impl.hpp | 24 +++--- .../kernel/unified_attention_kernel.hpp | 18 ++-- .../pipeline/unified_attention_pipeline.hpp | 3 - 16 files changed, 260 insertions(+), 37 deletions(-) create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode_s.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode_t.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode_s.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode_t.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode_s.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode_t.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode_s.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode_t.cpp diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode.cpp new file mode 100644 index 0000000000..112efe1222 --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode_s.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode_s.cpp new file mode 100644 index 0000000000..ef17fc1971 --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode_t.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_gqa8_bs32_decode_t.cpp new file mode 100644 index 0000000000..2c6531c835 --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode.cpp new file mode 100644 index 0000000000..f0c3617a52 --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode_s.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode_s.cpp new file mode 100644 index 0000000000..ead32cf0bf --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode_t.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_nmask_gqa8_bs32_decode_t.cpp new file mode 100644 index 0000000000..cc77cf7726 --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode.cpp new file mode 100644 index 0000000000..b79fe8eeb2 --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode_s.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode_s.cpp new file mode 100644 index 0000000000..272439ecb0 --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode_t.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_gqa8_bs32_decode_t.cpp new file mode 100644 index 0000000000..1420e3fa40 --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode.cpp new file mode 100644 index 0000000000..c883749ac2 --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode_s.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode_s.cpp new file mode 100644 index 0000000000..b76f03fe0c --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode_t.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_nmask_gqa8_bs32_decode_t.cpp new file mode 100644 index 0000000000..134ab386b5 --- /dev/null +++ b/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; + +INST_UNIFIED_ATTENTION_DISPATCH_DECODE(kernel_traits) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/unified_attention.cpp b/example/ck_tile/42_unified_attention/unified_attention.cpp index 9cc1231664..cbf2fd1e15 100644 --- a/example/ck_tile/42_unified_attention/unified_attention.cpp +++ b/example/ck_tile/42_unified_attention/unified_attention.cpp @@ -26,7 +26,7 @@ std::ostream& operator<<(std::ostream& stream, return unified_attention_kernel_dispatch(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; \ @@ -45,6 +45,25 @@ std::ostream& operator<<(std::ostream& stream, return unified_attention_kernel_dispatch_decode(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; \ + return unified_attention_kernel_dispatch(args, config); \ + } + +#define DISPATCH_UNIFIED_ATTENTION_DECODE_SMALL_BS32(DType, IsMask, HSize, BM, NQPKV) \ + { \ + using kernel_traits = unified_attention_decode_small_kernel_traits; \ + return unified_attention_kernel_dispatch_decode(args, config); \ + } + +#define DISPATCH_UNIFIED_ATTENTION_DECODE_TINY_BS32(DType, IsMask, HSize, BM, NQPKV) \ + { \ + using kernel_traits = unified_attention_decode_tiny_kernel_traits; \ + return unified_attention_kernel_dispatch_decode(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 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 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 diff --git a/example/ck_tile/42_unified_attention/unified_attention_impl.hpp b/example/ck_tile/42_unified_attention/unified_attention_impl.hpp index fd2e2d825a..52deb491b7 100644 --- a/example/ck_tile/42_unified_attention/unified_attention_impl.hpp +++ b/example/ck_tile/42_unified_attention/unified_attention_impl.hpp @@ -61,12 +61,13 @@ struct unified_attention_problem_traits + 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 + 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 + 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 + 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; 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 279fa68f85..e970230611 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 @@ -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(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( 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(kargs.num_blks) * kargs.page_size, + static_cast(kHeadDim)), + make_tuple(static_cast(kargs.stride_k_cache_1), + static_cast(kargs.stride_k_cache_3)), number{}, 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{}); @@ -399,8 +401,10 @@ struct UnifiedAttentionKernel const auto v_dram = [&]() { const auto v_dram_naive = make_naive_tensor_view( 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(kargs.num_blks) * kargs.page_size, + static_cast(kHeadDim)), + make_tuple(static_cast(kargs.stride_v_cache_1), + static_cast(kargs.stride_v_cache_3)), number{}, number<1>{}); diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp index 394532b49d..937bfdf5f2 100644 --- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp +++ b/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) {