From 63c75277a0234e696a9e3a68cf2da0eebcd5186e Mon Sep 17 00:00:00 2001 From: juuso-oskari Date: Fri, 15 May 2026 17:34:50 +0000 Subject: [PATCH] CK-UA: enable FP8 (e4m3) for prefill/m128 and the 32x32x16 small-tile decode variants Full pipeline support for FP8 (e4m3fn on gfx950 / e4m3fnuz on gfx942) in the unified-attention kernel, gated to the 32x32x16 MFMA tiers in both d=64 and d=128 ladders: prefill_d{64,128}, decode_d{64,128}_m128, decode_d128_m32, and decode_d64_m64. The 16x16x32 _m16 tiers stay BF16/FP16-only -- the QK-C and PV-A per-thread layouts there differ by an M<->N swap that the current slot-swap fixup cannot express; a full per-thread transpose (most likely via LDS) is needed. Pipeline (unified_attention_pipeline.hpp): * `fmha_alu1` now performs a cross-lane P-tile fixup right after the FP8 packing of softmax(P). It's a `ds_bpermute_b32` between paired lanes `lane ^ 32`, swapping sub=0 slot[k_base+4..k_base+7] with sub=1 slot[k_base..k_base+3] for every 8-fp8 chunk. This realigns the FP8 packed P operand with PV-A's `Single` AttrNumAccess per-thread layout, which is necessary because the QK-C output and PV-A input alias byte-for-byte via the sp_compute/p union -- and for FP8 the two warp-gemm layouts no longer agree (BF16/FP16 keep Double AttrNumAccess in the PV gemm, which matches QK-C natively). Gated on `Gemm1WarpTile == 32x32x16`; FP8-only (BF16/FP16 paths take the existing cvt_pk path unchanged). Default policy (unified_attention_pipeline_default_policy.hpp): * PV warp gemm now selects `WGAttrNumAccessEnum::Single` when V is fp8/bf8 and `Double` otherwise. Forced by load_tile_transpose's SubMinDim = 64-bit / sizeof(V) constraint: for FP8 SubMinDim=8 and kABKPerLane=8 only Single satisfies the validation static_asserts. * GetAlignmentK / GetAlignmentV on gfx950 drop to 4 B/lane for fp8/ bf8. The natural 16 B/lane async-load that BF16/FP16 use leaves NumIssues = 0 for the FP8 tile shapes we compile, and 8 B/lane fails the dword / dwordx3 / dwordx4 constraint in amd_buffer_addressing_builtins. 4 B/lane gives NumIssues >= 1 on every targeted variant and is the same alignment the gfx942 fallback already used. BF16/FP16 keep the full 16 B/lane path so existing perf is unchanged. * GetSmemSizeKV adds a `VLoadDescSize` lower bound. The MakeVLdsLoadBlockDescriptor's element span dominates the banked SingleVSize only for FP8 (small per-lane KVector + fixed kVLdsPadInBytes = 64), so without it FP8 hits the GetSmemSizeKV static_asserts. BF16/FP16 are unaffected. Warp-gemm headers + dispatcher: * New `WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed_T` template alias in warp_gemm.hpp (mirrors the existing BF16 32x32x16 CTransposed template), used by the PV gemm to thread the FP8 Single AttrNumAccess through. * New Dispatcher specialization for in warp_gemm_dispatcher.hpp routing to the new template. ABI / dispatcher (unified_attention.{cpp,hpp}, unified_attention_impl.hpp): * New `fp8` value in `unified_attention_args::data_type_enum` (selects e4m3fn on gfx950 via CK_TILE_USE_OCP_FP8, e4m3fnuz elsewhere). * New `unified_attention_problem_traits<...::fp8>` alias: qkvp_dtype = ck_tile::fp8_t, acc_dtype = float, o_dtype = bf16_t (matches the Triton reference), lse_dtype = float. * Per-tensor `q_descale` / `k_descale` / `v_descale` floats on `unified_attention_args` (default 1.0f so non-FP8 round-trips cleanly). The pipeline folds q_descale*k_descale into the softmax scale and applies v_descale once to o_acc after the 1/l norm -- same semantics as Triton's q_scale/k_scale/v_scale. * `dispatch_variant<>` enables FP8 on prefill_d{64,128}, decode_d{64,128}_m128, decode_d128_m32, decode_d64_m64. The 16x16x32 _m16 tiers return (false, -1.f) for now (see top comment). Instances: * 12 new FP8 .cpp files under example/.../42_unified_attention/ instances/ covering the 6 enabled variants x {mask, nmask}. Validation: 112 / 0 / 128 in the FP8 pytest sweep (passed / failed / m16-skipped); 245 / 245 in the BF16/FP16 sweep (no regression). Functional correctness is within the FP8 quant-noise tolerance the Triton FP8 suite uses (atol/rtol = 1.5e-1). Perf still trails Triton across the enabled tiers (CK FP8 / Triton FP8 = 0.39-0.69x on the shapes we benchmarked); that's a separate workstream. Co-authored-by: Cursor --- .../unified_attention_d128_fp8_mask.cpp | 11 + ...unified_attention_d128_fp8_mask_decode.cpp | 11 + ...ified_attention_d128_fp8_mask_decode_s.cpp | 11 + .../unified_attention_d128_fp8_nmask.cpp | 11 + ...nified_attention_d128_fp8_nmask_decode.cpp | 11 + ...fied_attention_d128_fp8_nmask_decode_s.cpp | 11 + .../unified_attention_d64_fp8_mask.cpp | 11 + .../unified_attention_d64_fp8_mask_decode.cpp | 11 + ...nified_attention_d64_fp8_mask_decode_s.cpp | 11 + .../unified_attention_d64_fp8_nmask.cpp | 11 + ...unified_attention_d64_fp8_nmask_decode.cpp | 11 + ...ified_attention_d64_fp8_nmask_decode_s.cpp | 11 + .../unified_attention.cpp | 33 +++ .../unified_attention.hpp | 25 +- .../unified_attention_impl.hpp | 19 ++ include/ck_tile/ops/gemm/warp/warp_gemm.hpp | 12 + .../ops/gemm/warp/warp_gemm_dispatcher.hpp | 1 + .../kernel/unified_attention_kernel.hpp | 43 ++- .../pipeline/unified_attention_pipeline.hpp | 252 ++++++++++++++++-- ...fied_attention_pipeline_default_policy.hpp | 78 +++++- 20 files changed, 560 insertions(+), 35 deletions(-) create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask_decode.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask_decode_s.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask_decode.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask_decode_s.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask_decode.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask_decode_s.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask_decode.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask_decode_s.cpp diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask.cpp new file mode 100644 index 0000000000..b31a71b6b7 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(prefill_d128, fp8, true) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask_decode.cpp new file mode 100644 index 0000000000..aa64d09c5b --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask_decode.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(decode_d128_m128, fp8, true) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask_decode_s.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask_decode_s.cpp new file mode 100644 index 0000000000..7aad26b73d --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask_decode_s.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(decode_d128_m32, fp8, true) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask.cpp new file mode 100644 index 0000000000..33ad5eea10 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(prefill_d128, fp8, false) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask_decode.cpp new file mode 100644 index 0000000000..177bb9b609 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask_decode.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(decode_d128_m128, fp8, false) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask_decode_s.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask_decode_s.cpp new file mode 100644 index 0000000000..760c45bcd5 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask_decode_s.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(decode_d128_m32, fp8, false) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask.cpp new file mode 100644 index 0000000000..0697fa0393 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(prefill_d64, fp8, true) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask_decode.cpp new file mode 100644 index 0000000000..164d6827d4 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask_decode.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(decode_d64_m128, fp8, true) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask_decode_s.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask_decode_s.cpp new file mode 100644 index 0000000000..2284781ef9 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask_decode_s.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(decode_d64_m64, fp8, true) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask.cpp new file mode 100644 index 0000000000..097ecf3d2e --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(prefill_d64, fp8, false) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask_decode.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask_decode.cpp new file mode 100644 index 0000000000..53fd1493a2 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask_decode.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(decode_d64_m128, fp8, false) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask_decode_s.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask_decode_s.cpp new file mode 100644 index 0000000000..2817ba3038 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask_decode_s.cpp @@ -0,0 +1,11 @@ +// 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 { + +INST_UNIFIED_ATTENTION_DISPATCH(decode_d64_m64, fp8, false) + +} // 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 b563e26643..722bb04b1f 100644 --- a/example/ck_tile/42_unified_attention/unified_attention.cpp +++ b/example/ck_tile/42_unified_attention/unified_attention.cpp @@ -14,6 +14,7 @@ std::ostream& operator<<(std::ostream& stream, { case unified_attention_args::data_type_enum::fp16: return stream << "fp16"; case unified_attention_args::data_type_enum::bf16: return stream << "bf16"; + case unified_attention_args::data_type_enum::fp8: return stream << "fp8"; default: return stream << "unknown"; } } @@ -131,6 +132,38 @@ std::pair dispatch_variant(const unified_attention_args& args, return unified_attention_kernel_dispatch< unified_attention_kernel_traits>(args, config); } + if(args.data_type == DT::fp8) + { + // FP8 is enabled on every variant that uses the 32x32x16 MFMA, + // i.e. all of prefill_d{64,128}, decode_d{64,128}_m128, + // decode_d128_m32, and decode_d64_m64. The `fmha_alu1` + // cross-lane P-tile fixup in unified_attention_pipeline.hpp + // (gated on `WarpGemmShape == 32x32x16`) realigns the FP8 P + // operand for these variants regardless of the single-/two-/ + // multi-warp `BlockWarps` policy. + // + // The 16x16x32 MFMA `_m16` tiers (decode_d128_m16, + // decode_d64_m16) are NOT yet enabled: the QK-C and PV-A + // per-thread layouts differ by an M<->N axis swap there + // (16x16 C: M=4 contiguous per lane, N=1 per lane; vs PV-A + // Single: M=1 per lane, K=8 contiguous per lane), which the + // current slot-swap fixup cannot express -- those will need a + // full per-thread transpose (likely via LDS). + if constexpr(V == KernelVariant::prefill_d128 || + V == KernelVariant::decode_d128_m128 || + V == KernelVariant::decode_d128_m32 || + V == KernelVariant::prefill_d64 || + V == KernelVariant::decode_d64_m128 || + V == KernelVariant::decode_d64_m64) + { + if(is_mask) + return unified_attention_kernel_dispatch< + unified_attention_kernel_traits>(args, config); + return unified_attention_kernel_dispatch< + unified_attention_kernel_traits>(args, config); + } + return {false, -1.f}; + } return {false, -1.f}; } diff --git a/example/ck_tile/42_unified_attention/unified_attention.hpp b/example/ck_tile/42_unified_attention/unified_attention.hpp index e5c1015b50..9f99281ba5 100644 --- a/example/ck_tile/42_unified_attention/unified_attention.hpp +++ b/example/ck_tile/42_unified_attention/unified_attention.hpp @@ -17,7 +17,8 @@ struct unified_attention_args enum class data_type_enum { fp16, - bf16 + bf16, + fp8 // e4m3fn on gfx950 (OCP) / e4m3fnuz on gfx942 (FNUZ); selected by CK_TILE_USE_OCP_FP8 }; data_type_enum data_type; @@ -34,11 +35,23 @@ struct unified_attention_args index_t hdim; // TODO window - float scale_s; - float scale; - float scale_k; - float scale_v; - float scale_out; + float scale_s; // softmax scale (1/sqrt(d) by convention); pre-multiplied with log2(e) + // inside MakeKargs so the device-side softmax can use exp2. + // Per-tensor FP8 descales (a.k.a. "scales" in the Triton kernel naming). All three + // default to 1.0f so non-FP8 dtypes round-trip cleanly. q_descale and k_descale are + // folded into scale_s inside the pipeline (so the softmax sees the combined scalar); + // v_descale is applied once to o_acc after the 1/l normalization, outside the K/V + // loop. Matches Triton unified_attention's q_scale/k_scale/v_scale semantics + // (see aiter/ops/triton/_triton_kernels/attention/unified_attention.py:110-114, 351-358). + float q_descale = 1.0f; + float k_descale = 1.0f; + float v_descale = 1.0f; + // Legacy fields kept for ABI stability with downstream callers (csrc glue, examples). + // The pipeline currently uses q_descale/k_descale/v_descale instead. + float scale = 1.0f; + float scale_k = 1.0f; + float scale_v = 1.0f; + float scale_out = 1.0f; const void* q_ptr; index_t query_stride_0; 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 f4b3b2ab72..f4d911737a 100644 --- a/example/ck_tile/42_unified_attention/unified_attention_impl.hpp +++ b/example/ck_tile/42_unified_attention/unified_attention_impl.hpp @@ -6,6 +6,7 @@ #include #include "ck_tile/core/numeric/bfloat16.hpp" +#include "ck_tile/core/numeric/float8.hpp" #include "ck_tile/core/numeric/half.hpp" #include "ck_tile/core/container/sequence.hpp" #include "ck_tile/host/kernel_launch.hpp" @@ -76,6 +77,21 @@ struct unified_attention_problem_traits +struct unified_attention_problem_traits +{ + using qkvp_dtype = ck_tile::fp8_t; + using acc_dtype = float; + using o_dtype = ck_tile::bf16_t; + using lse_dtype = float; +}; + // ============================================================================= // variant_config // @@ -303,6 +319,9 @@ float unified_attention_kernel_launch(const unified_attention_args& args, args.scale_k, args.scale_v, args.scale_out, + args.q_descale, + args.k_descale, + args.v_descale, args.page_blk_size, total_num_q_blocks, args.query_stride_0, diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp index e37af2ef5f..c72868df8b 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm.hpp @@ -489,6 +489,18 @@ using WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed = WarpGemmImpl>>; +// Templated form used by the FP8 unified-attention PV gemm. The PV block +// gemm in CK UA passes WGAttrNumAccessEnum::Double so the V tiles can be +// consumed via `load_tile_transpose()`; the BF16 32x32x16 CTransposed +// warp gemm already exposes this template parameter (see +// WarpGemmMfmaBf16Bf16F32M32N32K16TransposedCDistribution), and we mirror +// that pattern here. +template +using WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed_T = + WarpGemmImpl, + AttrNumAccess>>; + using WarpGemmMfma_f32_32x32x16_fp8_bf8_CTransposed = WarpGemmImpl>>; diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp index 94e0494aac..43a3dd2924 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp @@ -124,6 +124,7 @@ template<> struct Dispatcher { u template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8; }; template<> struct Dispatcher { using Type = WarpGemmMfma_f32_16x16x32_fp8_fp8; }; template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed; }; +template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed_T; }; template<> struct Dispatcher { using Type = WarpGemmMfma_f32_16x16x32_fp8_fp8_CTransposed; }; template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_bf8; }; template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x16_fp8_bf8_CTransposed; }; 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 92fe209fcd..d27a0f4573 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 @@ -66,11 +66,35 @@ struct UnifiedAttentionKernel // if this param is larger than 1, indicate MQA/GQA case const ck_tile::index_t num_queries_per_kv; // scales + // + // `scale_s` is the softmax scale (1/sqrt(d) by convention) AFTER: + // 1. multiplication by log2(e) so the pipeline's exp2 reproduces + // the natural-exponent softmax (preserved from the pre-FP8 + // design — see MakeKargs below); + // 2. fusion of the FP8 per-tensor Q/K descales (q_descale, + // k_descale) for FP8 problems. This matches Triton's + // unified_attention reference, which computes + // `qk_scale = sm_scale * q_scale * k_scale` and bakes the + // log2(e) factor into the exp2 inside the kernel. For + // non-FP8 dtypes the host passes q_descale = k_descale = 1.0 + // so the value reduces to `sm_scale * log2(e)`. + // + // `v_descale` is the FP8 per-tensor V descale, deferred to the + // post-loop `o_acc *= v_descale / l` step inside the pipeline + // (mathematically exact — V is a linear factor on the + // unnormalised attention output). For non-FP8 dtypes the host + // passes 1.0f so this is a free no-op. + // + // The legacy `scale` / `scale_k` / `scale_v` / `scale_out` fields + // are kept in the kargs struct for ABI continuity with downstream + // code that constructed kargs directly; they are not read by the + // pipeline any more. float scale_s; float scale; float scale_k; float scale_v; float scale_out; + float v_descale; ck_tile::index_t page_size; @@ -129,6 +153,9 @@ struct UnifiedAttentionKernel float scale_k, float scale_v, float scale_out, + float q_descale, + float k_descale, + float v_descale, ck_tile::index_t page_size, ck_tile::index_t total_num_q_blocks, ck_tile::index_t query_stride_0, @@ -157,6 +184,16 @@ struct UnifiedAttentionKernel ck_tile::index_t nhead_stride_o_acc = 0, bool cache_ptr_int32_overflow_possible = false) { + // Fuse the Q/K FP8 descales into `scale_s` so the softmax sees a + // single combined scalar — matches the Triton FP8 reference + // (qk_scale = sm_scale * q_scale * k_scale) and avoids extra + // arithmetic per element inside the kernel. The log2(e) factor + // is included here so the device-side exp2 produces the + // natural-exponent softmax. For non-FP8 dtypes the host passes + // q_descale = k_descale = 1.0, which reduces this to the original + // `scale_s * log2(e)`. + const float scale_s_fused = + static_cast(scale_s * q_descale * k_descale * ck_tile::log2e_v<>); Kargs kargs{{q_ptr, k_ptr, v_ptr, @@ -164,11 +201,12 @@ struct UnifiedAttentionKernel num_blks, num_head_q, num_queries_per_kv, - static_cast(scale_s * ck_tile::log2e_v<>), + scale_s_fused, scale, scale_k, scale_v, scale_out, + v_descale, page_size, total_num_q_blocks, query_stride_0, @@ -525,7 +563,8 @@ struct UnifiedAttentionKernel static_cast(kargs.stride_k_cache_1), static_cast(kargs.stride_v_cache_1), num_queries_per_kv, - kargs.cache_ptr_int32_overflow_possible); + kargs.cache_ptr_int32_overflow_possible, + kargs.v_descale); auto& o_acc_tile = pipeline_result[number<0>{}]; auto& lse_tile = pipeline_result[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 44b6e17a7a..194f62205c 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 @@ -97,8 +97,23 @@ struct UnifiedAttentionPipeline CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { - // create another LDS buffer for p - return ck_tile::max(kBlockM * kHeadDimPadded * sizeof(PDataType), + // Two layouts share the same smem base region during the pipeline: + // - the o_lds tile (kBlockM * kHeadDimPadded * sizeof(PDataType)), + // overlapped with the s_lds tile (kBlockM * kPageBlockSize * + // sizeof(SaccDataType)) — the QK gemm writes its float32 sacc + // into s_lds while the PV output is staged through o_lds; + // - K/V double-buffered storage (Policy::GetSmemSize) plus the + // p_lds tile that lives immediately after it (kBlockM * + // kPageBlockSize * sizeof(PDataType)). + // For BF16/FP16 the o_lds term dominates so historically the + // s_lds bound was implicit; for FP8 the PDataType drops to 1 B + // while SaccDataType stays at 4 B, so the s_lds term becomes the + // tightest bound and we must include it explicitly. The + // static_assert further down (sizeof(SaccDataType) * kPageBlockSize + // * kBlockM <= GetSmemSize()) now passes for every FP8 variant we + // compile. + return ck_tile::max(ck_tile::max(kBlockM * kHeadDimPadded * sizeof(PDataType), + kBlockM * kPageBlockSize * sizeof(SaccDataType)), Policy::template GetSmemSize() + kBlockM * kPageBlockSize * sizeof(PDataType)); } @@ -199,7 +214,16 @@ struct UnifiedAttentionPipeline // size can exceed INT32_MAX. Routes K/V async loads through the // 64-bit-base `global_load_lds` path (correct but lower throughput). // False uses the original shared-SRD `buffer_load_dword_lds` path. - const bool cache_ptr_int32_overflow_possible = false) const + const bool cache_ptr_int32_overflow_possible = false, + // Per-tensor FP8 V descale, applied to o_acc once after the 1/l + // normalisation. For non-FP8 dtypes the host passes 1.0f and this + // becomes a no-op multiply. The mathematical identity is: + // sum_j P[i,j] * V_real[j,:] = v_descale * sum_j P[i,j] * V_fp8[j,:] + // so deferring the v_descale outside the K/V loop is exact (not an + // approximation). For split-KV (num_splits > 1) each partial gets + // v_descale baked in; the combine step's affine weighting passes it + // through unchanged so the final O is correct. + const float v_descale = 1.0f) const { using namespace ck_tile; static_assert( @@ -752,22 +776,199 @@ struct UnifiedAttentionPipeline /// the cast_tile() call as inline assembly, forcing the conversions to be /// emitted at this point. static_assert(sp(sp_reg_idx).p.thread_buf_.size() % 2 == 0); - static_for<0, sp(sp_reg_idx).p.thread_buf_.size(), 2>{}([&](auto idx) { - float x = p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx]); - float y = p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx + 1]); - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) + { + // FP8 P packing for the PV gemm. + // + // The CK reference path for fp32 -> fp8 is `cast_tile_pk_fp8_fp32` + // in tile_elementwise.hpp, which CHAINS two `__builtin_amdgcn_cvt_pk_fp8_f32` + // calls per 4-lane group: the second call uses the first call's + // result as its `old` operand so the final uint32_t holds four + // valid fp8 bytes packed into one register. Doing the conversion + // pair-by-pair (one cvt_pk per 2 lanes, ignoring the upper 16 + // bits of the result) is *not* equivalent in practice -- the + // builtin's `old` argument feeds the upper-bits passthrough and + // when fed an uninitialised int the compiler is free to schedule + // the cvt against junk that overlaps another live register, + // which we observed as occasional whole-row output corruption + // for FP8 prefill workloads. + // + // Use the chained 4-lanes-per-iteration pattern to match + // `cast_tile_pk_fp8_fp32` byte-for-byte, then store as 4 fp8_t + // bytes back into `p.thread_buf_`. We still anchor the work + // inline (no `cast_tile(...)` indirection) so the conversions + // stay at the end of `fmha_alu1` like the FP16/BF16 paths. + static_assert(sp(sp_reg_idx).p.thread_buf_.size() % 4 == 0, + "fp8 P conversion expects packs of 4 fp32 lanes per " + "thread; widen the warp gemm M distribution if this " + "trips."); +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wuninitialized" + int dummy_old; + static_for<0, sp(sp_reg_idx).p.thread_buf_.size(), 4>{}([&](auto idx) { + const float a = p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx + 0]); + const float b = p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx + 1]); + const float c = p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx + 2]); + const float d = p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx + 3]); + + uint32_t lo = + __builtin_amdgcn_cvt_pk_fp8_f32(a, b, dummy_old, /*hi=*/false); + uint32_t packed = + __builtin_amdgcn_cvt_pk_fp8_f32(c, d, lo, /*hi=*/true); + sp(sp_reg_idx).p.thread_buf_[idx + 0] = + bit_cast(static_cast((packed >> 0) & 0xFFu)); + sp(sp_reg_idx).p.thread_buf_[idx + 1] = + bit_cast(static_cast((packed >> 8) & 0xFFu)); + sp(sp_reg_idx).p.thread_buf_[idx + 2] = + bit_cast(static_cast((packed >> 16) & 0xFFu)); + sp(sp_reg_idx).p.thread_buf_[idx + 3] = + bit_cast(static_cast((packed >> 24) & 0xFFu)); + }); +#pragma clang diagnostic pop + + // --------------------------------------------------------- + // FP8 P-tile cross-lane permute (kCMLane <-> kABKLane fixup) + // + // The CK UA pipeline relies on the QK-gemm's C output + // aliasing byte-for-byte with the PV-gemm's A operand + // through the `sp_compute` / `p` union. That alias is + // only valid when the two warp gemms' per-thread element + // layouts match. + // + // For BF16/FP16 the PV gemm uses `WGAttrNumAccess::Double` + // which splits `kABKPerLane=8` as (AttrNumAccess=2, + // kABKLane=2, kABKPerLane/Access=4). That last `4` + // matches the QK-C `kCM1PerLane=4` so each lane's + // sp_compute slots line up 1:1 with its p slots. + // + // FP8 cannot use `Double`: `load_tile_transpose` for the + // V tile requires the last H[1] dim to equal + // `64-bit/sizeof(fp8)=8`, so we are forced into `Single` + // where H[1]=(kABKLane=2, kABKPerLane=8). Now the + // sp_compute "kCM0=4, kCMLane=2, kCM1=4" layout no + // longer matches the p "kABKLane=2, kABKPerLane=8" + // layout. For each 8-fp8 PV K-chunk (one warp-gemm K + // iteration, since FP8 Single packs 8 elements/K-iter + // per thread), the slot decomposition is: + // + // sub=0 | slot in chunk | sp_compute (N) | PV (K) | + // -----+---------------+------------------+---------+ + // | [0..3] | N=0..3 | K=0..3 OK + // | [4..7] | N=8..11 | K=4..7 BAD + // + // sub=1 | slot in chunk | sp_compute (N) | PV (K) | + // -----+---------------+------------------+---------+ + // | [0..3] | N=4..7 | K=8..11 BAD + // | [4..7] | N=12..15 | K=12..15 OK + // + // (The above is for QK NIter=1 / PV KIter=1; the same + // 8-slot mismatch pattern repeats for every PV K-chunk + // -- d=128 has KIter=2 ==> 16 fp8/thread, d=64 has + // KIter=4 ==> 32 fp8/thread.) + // + // Each sub's *bad* 4-fp8 chunk carries exactly the data + // the *paired* sub needs for its *bad* slot (N is + // contracted into K in the PV gemm, so the index space + // is shared). The fix is a cross-lane exchange (lane L + // with lane L^32) of those bad chunks, repeated per PV + // K iteration. + // + // Implemented with `__builtin_amdgcn_ds_bpermute_b32`: + // each lane packs its own bad 4-fp8 chunk into a dword, + // bpermutes from the paired lane's same value (which is + // that lane's bad chunk = our needed data), and unpacks + // it into its own bad slot. The good slots are untouched. + // + // Only required when the QK-C and PV-A warp encodings + // disagree on per-thread element ordering, i.e. when the + // PV gemm runs with `Single` AttrNumAccess. The current + // FP8 PV policy is `Single`; BF16/FP16 use `Double` and + // alias for free. + // + // Cost: one `ds_bpermute_b32` per PV K-iteration per + // `fmha_alu1` call, each on a single dword. Same order + // of magnitude as the existing `permlane32_swap` + // reductions in this file. FP8-only. + // + // Gating: this swap pattern assumes the PV gemm uses the + // 32x32x16 MFMA (lanes split as kABMLane=32 / kABKLane=2, + // i.e. paired lane = `lane ^ 32`). The 16x16x32 MFMA + // used by the d128_m16 tiny-decode tier has a different + // lane partitioning (kABKLane=4) and is not FP8-enabled + // at the dispatcher, so we gate this swap on the + // 32x32x16 PV warp shape and leave the 16x16x32 case + // alone -- it would need a different swap pattern (and + // would have a different alias mismatch to fix). + using PVWarpTile = typename UnifiedAttentionShape::Gemm1WarpTile; + if constexpr(PVWarpTile::at(number<0>{}) == 32 && + PVWarpTile::at(number<1>{}) == 32 && + PVWarpTile::at(number<2>{}) == 16) { - auto casted = detail::cvt_pk_fp16_f32(x, y); - sp(sp_reg_idx).p.thread_buf_[idx] = casted.x; - sp(sp_reg_idx).p.thread_buf_[idx + 1] = casted.y; + static_assert(sp(sp_reg_idx).p.thread_buf_.size() % 8 == 0, + "FP8 32x32x16 + Single cross-lane permute " + "expects PV per-thread buffer in chunks of 8 " + "fp8 (one warp-gemm K iteration)."); + + const int lane_id = ck_tile::get_lane_id(); + const int paired_addr = (lane_id ^ 32) << 2; // bytes + const bool is_sub_0 = (lane_id & 32) == 0; + + auto pack4 = [](fp8_t a, fp8_t b, fp8_t c, fp8_t d) -> uint32_t { + return (static_cast(bit_cast(a)) << 0) | + (static_cast(bit_cast(b)) << 8) | + (static_cast(bit_cast(c)) << 16) | + (static_cast(bit_cast(d)) << 24); + }; + auto unpack4 = [](uint32_t v, fp8_t& a, fp8_t& b, fp8_t& c, fp8_t& d) { + a = bit_cast(static_cast((v >> 0) & 0xFFu)); + b = bit_cast(static_cast((v >> 8) & 0xFFu)); + c = bit_cast(static_cast((v >> 16) & 0xFFu)); + d = bit_cast(static_cast((v >> 24) & 0xFFu)); + }; + + // One swap per PV K-iteration: within each 8-fp8 chunk, + // sub=0's slot [k_base+4..k_base+7] <-> sub=1's slot + // [k_base..k_base+3]. + static_for<0, sp(sp_reg_idx).p.thread_buf_.size(), 8>{}([&](auto k_base) { + auto& p = sp(sp_reg_idx).p; + const uint32_t own_bad = + is_sub_0 + ? pack4(p.thread_buf_[k_base + 4], p.thread_buf_[k_base + 5], + p.thread_buf_[k_base + 6], p.thread_buf_[k_base + 7]) + : pack4(p.thread_buf_[k_base + 0], p.thread_buf_[k_base + 1], + p.thread_buf_[k_base + 2], p.thread_buf_[k_base + 3]); + const uint32_t recv = + __builtin_amdgcn_ds_bpermute(paired_addr, static_cast(own_bad)); + if(is_sub_0) + unpack4(recv, + p.thread_buf_[k_base + 4], p.thread_buf_[k_base + 5], + p.thread_buf_[k_base + 6], p.thread_buf_[k_base + 7]); + else + unpack4(recv, + p.thread_buf_[k_base + 0], p.thread_buf_[k_base + 1], + p.thread_buf_[k_base + 2], p.thread_buf_[k_base + 3]); + }); } - else - { - auto casted = cvt_pk_bf16_f32(x, y); - sp(sp_reg_idx).p.thread_buf_[idx] = casted.x; - sp(sp_reg_idx).p.thread_buf_[idx + 1] = casted.y; - } - }); + } + else + { + static_for<0, sp(sp_reg_idx).p.thread_buf_.size(), 2>{}([&](auto idx) { + float x = p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx]); + float y = p_compute_element_func(sp(sp_reg_idx).sp_compute.thread_buf_[idx + 1]); + if constexpr(std::is_same_v) + { + auto casted = detail::cvt_pk_fp16_f32(x, y); + sp(sp_reg_idx).p.thread_buf_[idx] = casted.x; + sp(sp_reg_idx).p.thread_buf_[idx + 1] = casted.y; + } + else + { + auto casted = cvt_pk_bf16_f32(x, y); + sp(sp_reg_idx).p.thread_buf_[idx] = casted.x; + sp(sp_reg_idx).p.thread_buf_[idx + 1] = casted.y; + } + }); + } /// Note: Place fmha_alu1() at the end of the phase. The surrounding inline assembly /// can interfere with the behavior of sched_group_barrier(), so ending the phase here @@ -1253,13 +1454,18 @@ struct UnifiedAttentionPipeline sweep_tile_span(o_spans[number<0>{}], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); + // Fuse the V FP8 descale into the per-row normalisation so the + // post-loop pass touches o_acc only once. v_descale is host-set + // to 1.0f for non-FP8 dtypes so this stays a free no-op there. + // Masked rows that saw no valid keys keep their zeros (the + // l == 0 short-circuit below). const auto tmp = [&]() { if constexpr(FmhaMask::IsMasking) { - return l[i_idx] == 0.f ? 0.f : 1 / l[i_idx]; + return l[i_idx] == 0.f ? 0.f : v_descale / l[i_idx]; } else - return 1 / l[i_idx]; + return v_descale / l[i_idx]; }(); sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { constexpr auto i_j_idx = make_tuple(idx0, idx1); @@ -1323,7 +1529,10 @@ struct UnifiedAttentionPipeline // runtime kBlockQ. See the documentation on that overload. const index_t num_queries_per_kv = 0, // See the doc on the full-args operator(). - const bool cache_ptr_int32_overflow_possible = false) const + const bool cache_ptr_int32_overflow_possible = false, + // See the doc on the full-args operator(). Defaults to 1.0f so + // non-FP8 callers see no behavior change. + const float v_descale = 1.0f) const { using namespace ck_tile; @@ -1347,7 +1556,8 @@ struct UnifiedAttentionPipeline k_row_stride, v_row_stride, num_queries_per_kv, - cache_ptr_int32_overflow_possible); + cache_ptr_int32_overflow_possible, + v_descale); } }; diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp index c2773f9076..eb861ec5f9 100644 --- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp +++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp @@ -39,7 +39,28 @@ struct UnifiedAttentionPipelineDefaultPolicy using namespace ck_tile; using KDataType = remove_cvref_t; #if defined(__gfx950__) - constexpr index_t MaxReadSizeInBytes = 16; + // FP8 caveat on gfx950: + // - The natural 16 B/lane async load (KVector = 16 elements/lane) + // leaves NumIssues = 0 for the prefill / decode kPageBlockSize x + // NumWarps tuples we compile FP8 instances for (kBlockSize * + // KVector exceeds kNPerBlock * kKPerBlock). + // - Dropping to 8 B/lane brings NumIssues back to >=1, but on + // gfx950 the LDS-targeted `global_load_lds` instruction only + // supports dword / dwordx3 / dwordx4 (4 / 12 / 16 B per lane); + // 8 B fails the static_assert in amd_buffer_addressing_builtins. + // - 4 B / lane (one dword) works on every targeted tile and is + // the same as the gfx942 fallback below. Lower per-instruction + // bytes than BF16/FP16's 16 B path, but the FP8 K/V is half the + // size in bytes so the *element* count moved per cycle stays + // reasonable. Verified compile-time on prefill_d{64,128}, + // decode_d{64,128}_m128. + // BF16 / FP16 keep the full 16 B/lane read so existing perf is + // unchanged. + constexpr index_t MaxReadSizeInBytes = + (std::is_same_v || + std::is_same_v) + ? 4 + : 16; #else constexpr index_t MaxReadSizeInBytes = 4; #endif @@ -52,7 +73,13 @@ struct UnifiedAttentionPipelineDefaultPolicy using namespace ck_tile; using VDataType = remove_cvref_t; #if defined(__gfx950__) - constexpr index_t MaxReadSizeInBytes = 16; + // See the FP8 caveat on GetAlignmentK above — symmetric reasoning + // for V. + constexpr index_t MaxReadSizeInBytes = + (std::is_same_v || + std::is_same_v) + ? 4 + : 16; #else constexpr index_t MaxReadSizeInBytes = 4; #endif @@ -279,8 +306,20 @@ struct UnifiedAttentionPipelineDefaultPolicy Problem::UnifiedAttentionShape::kPageBlockSize>, typename Problem::UnifiedAttentionShape::Gemm1BlockWarps, typename Problem::UnifiedAttentionShape::Gemm1WarpTile>>; - /// NOTICE: in order to use load_tile_transpose() later for V tiles, we have to pass - /// WGAttrNumAccessEnum::Double instead of WGAttrNumAccessEnum::Single + // `load_tile_transpose` is only valid when the tile distribution's inner + // packing matches the transpose engine's SubtileMinorDimension = + // 64 bits / sizeof(VDataType_in_bits). For BF16 / FP16 SubMinDim=4 and the + // PV warp gemm produces kABKPerLane / AttrNumAccess = 8 / 2 = 4 elements + // per lane on the K direction — Double access is needed there. For FP8 + // SubMinDim=8 and kABKPerLane=8, so we must pass `Single` (otherwise + // 8/2=4 mismatches the FP8 SubMinDim=8 and the load_tile_transpose + // validation static_asserts fire — see DefaultTranspose::Quad in + // load_tile_transpose.hpp). The select is purely a compile-time alias. + static constexpr WGAttrNumAccessEnum PVAttrNumAccess = + std::is_same_v, ck_tile::fp8_t> || + std::is_same_v, ck_tile::bf8_t> + ? WGAttrNumAccessEnum::Single + : WGAttrNumAccessEnum::Double; using WarpGemm = WarpGemmDispatcher; + PVAttrNumAccess>; using BlockGemmPolicy = BlockGemmARegBRegCRegV2CustomPolicy< typename Problem::PDataType, @@ -451,7 +490,34 @@ struct UnifiedAttentionPipelineDefaultPolicy return (kKPerBlock / kKPack) * (kNPerBlock / NPerRow) * (PixelsPerRow + kKPack); }(); - return max(SingleKSize, SingleVSize); + // Lower-bound on the actual MakeVLdsLoadBlockDescriptor element + // span: it allocates a (NumIssues, LaneGroups, NumWarps, LanesPerK, + // KVector) buffer with the outermost stride NumWarps * (WarpSize * + // KVector + kPad). For BF16/FP16 the existing banked-layout + // SingleVSize above is always larger; for FP8 the small per-lane + // KVector (4 B = 4 fp8 elements) combined with the byte-fixed + // kVLdsPadInBytes = 64 makes the V descriptor's element span + // dominate, so we must include it here or the static_asserts in + // GetSmemSizeKV fire. + constexpr index_t VLoadDescSize = [&]() { + constexpr index_t kNPerBlock = Problem::UnifiedAttentionShape::kPageBlockSize; + constexpr index_t kKPerBlock = Problem::UnifiedAttentionShape::kHeadDim; + constexpr index_t NumWarps = Problem::UnifiedAttentionShape::NumWarps; + constexpr index_t WarpSize = ck_tile::get_warp_size(); + constexpr index_t KVector = GetAlignmentV(); + constexpr index_t kPad = + kVLdsPadInBytes / sizeof(typename Problem::VDataType); + + static_assert(WarpSize * KVector >= kKPerBlock && + WarpSize * KVector % kKPerBlock == 0); + constexpr index_t LanesPerK = kKPerBlock / KVector; + constexpr index_t LaneGroups = WarpSize / LanesPerK; + constexpr index_t NumIssues = kNPerBlock / (LaneGroups * NumWarps); + + return NumIssues * NumWarps * (WarpSize * KVector + kPad); + }(); + + return max(max(SingleKSize, SingleVSize), VLoadDescSize); } template