ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 19:18:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

CK-UA: enable FP8 (e4m3) for prefill/m128 and the 32x32x16 small-tile decode variants

Browse Source 
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<AttrNumAccess>`
  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
  <fp8_t, fp8_t, float, 32, 32, 16, true, false, false, EDouble>
  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 <cursoragent@cursor.com>
This commit is contained in:
juuso-oskari
2026-05-15 17:34:50 +00:00
parent c0e985d075
commit 63c75277a0
20 changed files with 560 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask_decode.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_mask_decode_s.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask_decode.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d128_fp8_nmask_decode_s.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask_decode.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_mask_decode_s.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask_decode.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp8_nmask_decode_s.cpp Normal file
View File

@@ -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

33
example/ck_tile/42_unified_attention/unified_attention.cpp
View File

@@ -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<bool, float> dispatch_variant(const unified_attention_args& args,
return unified_attention_kernel_dispatch<
unified_attention_kernel_traits<V, DT::bf16, false>>(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<V, DT::fp8, true>>(args, config);
return unified_attention_kernel_dispatch<
unified_attention_kernel_traits<V, DT::fp8, false>>(args, config);
}
return {false, -1.f};
}
return {false, -1.f};
}

25
example/ck_tile/42_unified_attention/unified_attention.hpp
View File

@@ -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;

19
example/ck_tile/42_unified_attention/unified_attention_impl.hpp
View File

@@ -6,6 +6,7 @@
#include <utility>
#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<unified_attention_args::data_type_enum::
using lse_dtype = float;
};
// FP8 path: Q, K, V are stored / consumed by MFMA as FP8 (e4m3*), but the
// softmax probability P that feeds the second GEMM is *also* FP8 so we can
// reuse the FP8-FP8-F32 MFMA family for PV. The o_dtype stays bf16 (Triton
// reference's output dtype is bf16). Accumulators and LSE remain fp32. The
// type alias resolves to e4m3fnuz on gfx942 and e4m3fn on gfx950 — see
// ck_tile/core/numeric/float8.hpp for the CK_TILE_USE_OCP_FP8 selector.
template <>
struct unified_attention_problem_traits<unified_attention_args::data_type_enum::fp8>
{
using qkvp_dtype = ck_tile::fp8_t;
using acc_dtype = float;
using o_dtype = ck_tile::bf16_t;
using lse_dtype = float;
};
// =============================================================================
// variant_config<V>
//
@@ -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,

12
include/ck_tile/ops/gemm/warp/warp_gemm.hpp
View File

@@ -489,6 +489,18 @@ using WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed =
WarpGemmImpl<WarpGemmAttributeMfmaTransposedCDistribution<
WarpGemmAttributeMfmaImpl_f32_32x32x16_fp8_fp8<WGAttrCtlEnum::Default_>>>;
// 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 <WGAttrNumAccessEnum AttrNumAccess>
using WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed_T =
WarpGemmImpl<WarpGemmAttributeMfmaTransposedCDistribution<
WarpGemmAttributeMfmaImpl_f32_32x32x16_fp8_fp8<WGAttrCtlEnum::Default_>,
AttrNumAccess>>;
using WarpGemmMfma_f32_32x32x16_fp8_bf8_CTransposed =
WarpGemmImpl<WarpGemmAttributeMfmaTransposedCDistribution<
WarpGemmAttributeMfmaImpl_f32_32x32x16_fp8_bf8<WGAttrCtlEnum::Default_>>>;

1
include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp
View File

@@ -124,6 +124,7 @@ template<> struct Dispatcher<bf16_t, bf16_t, float, 32, 32, 16, true, true> { u
template<> struct Dispatcher<fp8_t, fp8_t, float, 32, 32, 16, false> { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8; };
template<> struct Dispatcher<fp8_t, fp8_t, float, 16, 16, 32, false> { using Type = WarpGemmMfma_f32_16x16x32_fp8_fp8; };
template<> struct Dispatcher<fp8_t, fp8_t, float, 32, 32, 16, true> { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed; };
template<> struct Dispatcher<fp8_t, fp8_t, float, 32, 32, 16, true, false, false, EDouble> { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed_T<EDouble>; };
template<> struct Dispatcher<fp8_t, fp8_t, float, 16, 16, 32, true> { using Type = WarpGemmMfma_f32_16x16x32_fp8_fp8_CTransposed; };
template<> struct Dispatcher<fp8_t, bf8_t, float, 32, 32, 16, false> { using Type = WarpGemmMfma_f32_32x32x16_fp8_bf8; };
template<> struct Dispatcher<fp8_t, bf8_t, float, 32, 32, 16, true> { using Type = WarpGemmMfma_f32_32x32x16_fp8_bf8_CTransposed; };

43
include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp
View File

@@ -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<float>(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<float>(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<long_index_t>(kargs.stride_k_cache_1),
static_cast<long_index_t>(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>{}];

252
include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
View File

@@ -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<Problem>() +
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<PDataType, fp16_t>)
if constexpr(std::is_same_v<PDataType, fp8_t>)
{
// 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<fp8_t>(static_cast<fp8_raw_t>((packed >> 0) & 0xFFu));
sp(sp_reg_idx).p.thread_buf_[idx + 1] =
bit_cast<fp8_t>(static_cast<fp8_raw_t>((packed >> 8) & 0xFFu));
sp(sp_reg_idx).p.thread_buf_[idx + 2] =
bit_cast<fp8_t>(static_cast<fp8_raw_t>((packed >> 16) & 0xFFu));
sp(sp_reg_idx).p.thread_buf_[idx + 3] =
bit_cast<fp8_t>(static_cast<fp8_raw_t>((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<uint32_t>(bit_cast<fp8_raw_t>(a)) << 0) |
(static_cast<uint32_t>(bit_cast<fp8_raw_t>(b)) << 8) |
(static_cast<uint32_t>(bit_cast<fp8_raw_t>(c)) << 16) |
(static_cast<uint32_t>(bit_cast<fp8_raw_t>(d)) << 24);
};
auto unpack4 = [](uint32_t v, fp8_t& a, fp8_t& b, fp8_t& c, fp8_t& d) {
a = bit_cast<fp8_t>(static_cast<fp8_raw_t>((v >> 0) & 0xFFu));
b = bit_cast<fp8_t>(static_cast<fp8_raw_t>((v >> 8) & 0xFFu));
c = bit_cast<fp8_t>(static_cast<fp8_raw_t>((v >> 16) & 0xFFu));
d = bit_cast<fp8_t>(static_cast<fp8_raw_t>((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<int>(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<PDataType, fp16_t>)
{
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);
}
};

78
include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline_default_policy.hpp
View File

@@ -39,7 +39,28 @@ struct UnifiedAttentionPipelineDefaultPolicy
using namespace ck_tile;
using KDataType = remove_cvref_t<typename Problem::KDataType>;
#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<KDataType, ck_tile::fp8_t> ||
std::is_same_v<KDataType, ck_tile::bf8_t>)
? 4
: 16;
#else
constexpr index_t MaxReadSizeInBytes = 4;
#endif
@@ -52,7 +73,13 @@ struct UnifiedAttentionPipelineDefaultPolicy
using namespace ck_tile;
using VDataType = remove_cvref_t<typename Problem::VDataType>;
#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<VDataType, ck_tile::fp8_t> ||
std::is_same_v<VDataType, ck_tile::bf8_t>)
? 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<remove_cvref_t<typename Problem::VDataType>, ck_tile::fp8_t> ||
std::is_same_v<remove_cvref_t<typename Problem::VDataType>, ck_tile::bf8_t>
? WGAttrNumAccessEnum::Single
: WGAttrNumAccessEnum::Double;
using WarpGemm =
WarpGemmDispatcher<typename Problem::PDataType,
typename Problem::VDataType,
@@ -291,7 +330,7 @@ struct UnifiedAttentionPipelineDefaultPolicy
true,
false,
false,
WGAttrNumAccessEnum::Double>;
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<Problem>();
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 <typename Problem, ck_tile::index_t IBuf = 0>