mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-03 05:01:25 +00:00
1e77695fe8
* Pass hdim to tile_example_fmha_fwd in fp8 tests
* Add WMMA support to fwd FMHA pipelines
* Tune tile sizes a bit for less spilling
fp16 256 is still quite slow
* Fix Q grad tile distribution for warp size = 32 and hdim >= 256
With AccDataType = float and warp size = 32, K0 becomes 0, K repeat is required to correcty distribute the tile.
* Use code based on BlockDropout in BlockDropoutBwd
* Fix split KV combine kernel for gfx12 (warp size 32) and make it more universal
* Fix LSE LDS tensor descriptors: kMaxSplits and kM0 were swapped, it worked on gfx9
because they both equal to 8 while on gfx12 they are 8 and 4;
* Fix Oacc LDS tensor descriptor: it was transposed even though its shape=[4 * kM0, kN1],
it worked on gfx9 because 4 * kM == kN1 == 32;
* Removing these hidden dependecies allows to support:
* any number of warps (power-of-2), not only 4;
* kN1 = 16, not only 32;
* any number of splits;
* Rename ids like o_acc_4 and Oacc4 to eliminate confusion: kNumWarps doesn't have to be 4 now
* Replace hard-coded kN1 in dispatch code with the requested tile size
* Add gfx12-specific tile sizes for split KV
* Pass GPU architecture to kernel generation scripts
This is still a temporary solution.
* Build and run FMHA CI tests for gfx12
* Fix issue after merging
* Fix bwd tile sizes
The current pipelines always read only one tile K and V tile, this
requires bk0 == bhdq and bk2 == bhdv (kK0 == kQKHeaddim and
kK2 == kVHeaddim).
* Use hardware f32->f8 on gfx12, remove v_perm
__builtin_amdgcn_perm is not needed because
__builtin_amdgcn_cvt_pk_fp8_f32 allows to specify which word (16 bit of
32-bit dword) is used to store results (two f8 values).
* Update changelog
* Add WMMA support to pagedkv
* Fix scripts after rebasing
* Support 16x16 (MFMA, WMMA) and 32x32 (MFMA) tiles in fwd and bwd BlockDropout
Add comments with dropout implementation details
Fix performance regression of fwd+dropout
* Remove some usage of type punning (reinterpret_cast with ref or ptr) in Philox;
* "scalarize" seed and offset, they may come either from kernel args or from device memory
(presumably loaded with vector loads).
These changes help the compiler to procude more optimal code and reduce register spilling.
Use WarpGemmDispatcher instead of explicit WarpGemmMfma... to get CWarpDstrEncoding
Use code based on BlockDropout in BlockDropoutBwd
Refactor BlockDropout (fwd)
Implement BlockDropout (fwd) for WMMA
Originally BlockDropout only supported 32x32 tiles (IsWG32 = true),
this version supports 16x16 tiles.
If MPerBlock > MWarp * 16, it can generate numbers for two 16x16 tiles, similarly
to BlockDropoutBwd.
Implement BlockDropoutBwd for WMMA
Remove MakeRandValLds* functions unused in BlockDropoutBwd
Remove unused Run overload from BlockDropoutBwd
* Fix regression with philox seed and offset when they exceed 32-bit int
__builtin_amdgcn_readfirstlane works with 32-bit values, seed and offset
are 64-bit so they get truncated.
* Fix names after cherry-picking
* Fix selection of a fallback tile based on bm0
The assumption that the largest bm0 == 128 is not always true for
current fp32 tiles.
* Do not use filters related to qr_async_trload
They disable tiles/pipelines which are valid for gfx12.
* Use different dstr encoding when C is transposed
* Do not call GetQKBlockGemm (and hence WarpGemmDispatcher) in host code
Some WarpGemmDispatcher instantiations are defined only
for specific archs and undefined on host.
Calculations related to sched barriers are moved from Pipeline's public
fields into pipeline's operator().
* Fix incorrect name WarpGemmMfmaFp8Fp8F32M32N32K16SwizzleBTransposedCDistribution
Correct name is WarpGemmMfmaFp8Fp8F32M32N32K32SwizzleBTransposedCDistribution
because it's 32x32x16 with IterateK = 2 so K = 32, also all tiles used
in codegen scripts are 32, 32, 32.
* Generalize usages of WarpGemmDispatcher for MFMA and WMMA
WarpGemmMfmaFp8Fp8F32M32N32K32SwizzleBTransposedCDistribution is still
used explicitly becaus of swizzle factor = 4.
* Mark has_load_tr as maybe_unused
There are no transpose loading for RDNA.
* Remove CK_TILE_USE_MFMA/WMMA from fmha-related code
* Detect BlockSize on host based on warp size of the current device
If kBlockSize == kNumWarps * get_warp_size(), the kernel is launched with
kBlockSize / 2 because on host get_warp_size() == 64 always.
* Fix calculation of grid size for combine kernel with warp size = 32
* Add missing includes and header
* Support multiple archs in one binary for fwd
* Support multiple archs in one binary for fwd_splitkv, fwd_appendkv, pagedkv_prefill
* Support multiple archs in one binary for bwd
* trload kernels are compiled only for gfx950;
* instances with padding are checked after instances without padding so
they can be used as fallbacks (similarly to fwd);
* Extract common code from register_traits
* Revert "Fix regression with philox seed and offset when they exceed 32-bit int"
To simplify merging , the proper fix is in develop already.
* Support new numerical d paddings in trait ordering checks
* Build fp32 tests only on gfx9
* Do not use hardcoded M0 = 64 for dot bwd kernel
* Use textwrap.indent from standard library
* Make fp8 pipelines on gfx12 consistent with gfx9
* Update tests for current pipelines
* Make ninja check more responsive in CI
ninja buffers output so this job looks hanging.
* Support fp8fp32 by limiting O vector size
The fp32 output type requires storing 8 * sizeof(float) = 32 bytes,
which is not implemented (here 8 is the number of C values per lane for
v_wmma_f32_16x16x16...).
* Remove unused cmake options
* Unify including amd_buffer_addressing.hpp/_builtins.hpp
* Temporarily use amd_buffer_addressing.hpp on >=gfx10
amd_buffer_addressing_builtins.hpp uses inline asm for loads/stores
which is not compatible with >=gfx10:
* 1 scalar for exec masks instead of 2,
* gfx12 uses different instruction names etc.
* Update asm in bf16 conversions to work with warp 32
* Do not generate splitkv/appendkv with vlayout=col for consistency with fwd
* Add arch tags to kernels/host funcs, compile for each arch separately
* Add kM0 to fmha_bwd_dot_do_o kernel name to match filename
* Add workaround for miscompilation of bwd with padded hdim
SWDEV-559729: v_wmma instructions can be incorrectly placed in divergent
branches used to store padded tensors (when some lanes are inactive due
to padding). Inline asm with dummy dependencies on VGPRs of the tensors
prevents the compiler doing this.
* Fix add_gtest_executable for absolute paths
Some tests (like gemm_tile_engine) pass absolute paths to source files.
In CI the branch name is a part of the root dir, and if the branch name
contains "wmma", "xdl" etc., files can be incorrectly excluded.
* Run only hdim 128 smoke tests for fp8fp32
There are no instances for hdim 64 and 256.
* Format py with ruff to simplify merging develop
* Fix incorrect var name
* Codegen for gfx9,gfx950 when --targets is not specified
Aiter and Pytorch require changes for passing their targets to the codegen scripts.
With this temporary solution the files are generated but not all of them
have to be really built (depending on the used --offload-arch=).
* Combine arch-related values into ArchTrait
This more centralized approach removes duplication of various formatting templates.
* Try a workaround for Jenkins error "groovyjarjarasm.asm.MethodTooLargeException: Method too large"
Some code is extracted into a function.
523 lines
24 KiB
C++
523 lines
24 KiB
C++
// SPDX-License-Identifier: MIT
|
|
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
|
|
|
#pragma once
|
|
|
|
#include "ck_tile/core.hpp"
|
|
#include "ck_tile/host/device_prop.hpp"
|
|
#include "ck_tile/host/kernel_launch.hpp"
|
|
#include "ck_tile/ops/fmha.hpp"
|
|
#include "ck_tile/ops/epilogue.hpp"
|
|
#include "mask.hpp"
|
|
#include "bias.hpp"
|
|
|
|
#include <type_traits>
|
|
#include <utility>
|
|
#include <variant>
|
|
|
|
struct FmhaBwdFp32
|
|
{
|
|
};
|
|
|
|
struct FmhaBwdFp16
|
|
{
|
|
};
|
|
|
|
struct FmhaBwdBf16
|
|
{
|
|
};
|
|
|
|
template <typename DataType>
|
|
struct FmhaBwdTypeConfig;
|
|
|
|
template <>
|
|
struct FmhaBwdTypeConfig<FmhaBwdFp32>
|
|
{
|
|
using QDataType = float;
|
|
using KDataType = float;
|
|
using VDataType = float;
|
|
using GemmDataType = float;
|
|
using BiasDataType = float;
|
|
using LSEDataType = float;
|
|
using AccDataType = float; // data type for gemm accumulation
|
|
using DDataType = float;
|
|
using RandValOutputDataType = uint8_t;
|
|
using ODataType = float;
|
|
using OGradDataType = float;
|
|
using QGradDataType = float;
|
|
using KGradDataType = float;
|
|
using VGradDataType = float;
|
|
using BiasGradDataType = float;
|
|
};
|
|
|
|
template <>
|
|
struct FmhaBwdTypeConfig<FmhaBwdFp16>
|
|
{
|
|
using QDataType = ck_tile::half_t;
|
|
using KDataType = ck_tile::half_t;
|
|
using VDataType = ck_tile::half_t;
|
|
using GemmDataType = ck_tile::half_t;
|
|
using BiasDataType = ck_tile::half_t;
|
|
using LSEDataType = float;
|
|
using AccDataType = float; // data type for gemm accumulation
|
|
using DDataType = float;
|
|
using RandValOutputDataType = uint8_t;
|
|
using ODataType = ck_tile::half_t;
|
|
using OGradDataType = ck_tile::half_t;
|
|
using QGradDataType = ck_tile::half_t;
|
|
using KGradDataType = ck_tile::half_t;
|
|
using VGradDataType = ck_tile::half_t;
|
|
using BiasGradDataType = ck_tile::half_t;
|
|
};
|
|
|
|
template <>
|
|
struct FmhaBwdTypeConfig<FmhaBwdBf16>
|
|
{
|
|
using QDataType = ck_tile::bf16_t;
|
|
using KDataType = ck_tile::bf16_t;
|
|
using VDataType = ck_tile::bf16_t;
|
|
using GemmDataType = ck_tile::bf16_t;
|
|
using BiasDataType = ck_tile::bf16_t;
|
|
using LSEDataType = float;
|
|
using AccDataType = float; // data type for gemm accumulation
|
|
using DDataType = float;
|
|
using RandValOutputDataType = uint8_t;
|
|
using ODataType = ck_tile::bf16_t;
|
|
using OGradDataType = ck_tile::bf16_t;
|
|
using QGradDataType = ck_tile::bf16_t;
|
|
using KGradDataType = ck_tile::bf16_t;
|
|
using VGradDataType = ck_tile::bf16_t;
|
|
using BiasGradDataType = ck_tile::bf16_t;
|
|
};
|
|
|
|
struct FmhaMasks
|
|
{
|
|
using NoMask = ck_tile::GenericAttentionMask<false>;
|
|
using GenericMask = ck_tile::GenericAttentionMask<true, true>;
|
|
using CausalMask = ck_tile::GenericAttentionMask<true, false>;
|
|
};
|
|
|
|
// runtime args, some will passed to karg, some will used to compute grids/blocks
|
|
struct fmha_bwd_args
|
|
{
|
|
const void* q_ptr;
|
|
const void* k_ptr;
|
|
const void* v_ptr;
|
|
const void* bias_ptr; // bias or alibi_slope pointer
|
|
const void* o_ptr;
|
|
const void* lse_ptr;
|
|
const void* do_ptr;
|
|
void* d_ptr;
|
|
void* rand_val_ptr;
|
|
void* dq_ptr;
|
|
void* dk_ptr;
|
|
void* dv_ptr;
|
|
void* dbias_ptr;
|
|
void* dq_acc_ptr;
|
|
|
|
// Usage notes for sequence length pointer parameters:
|
|
//
|
|
// [Note: Define "Group mode" vs "Batch mode" here if possible, e.g., "Group mode handles
|
|
// MQA/GQA..."]
|
|
//
|
|
// With padding:
|
|
// Group mode:
|
|
// - seqstart_q_ptr, seqstart_k_ptr: Record cumulative physical (including padding) sequence
|
|
// lengths. [array size: batch + 1]
|
|
// - seqlen_q_ptr/seqlen_k_ptr: Records logical (excluding padding) length for each
|
|
// sequence. [array size: batch]
|
|
// - cu_seqlen_q_ptr/cu_seqlen_k_ptr: Records cumulative logical (excluding padding)
|
|
// sequence lengths. [array size: batch + 1]
|
|
// - seqlen_q_ptr (per-sequence) and cu_seqlen_q_ptr (cumulative logical) are mutually
|
|
// exclusive. Use one set, not both.
|
|
//
|
|
// Batch mode:
|
|
// - cu_seqlen_q_ptr/cu_seqlen_k_ptr: Records cumulative logical (excluding padding)
|
|
// sequence lengths. [array size: batch + 1]
|
|
// - seqstart_* and seqlen_* pointers must be nullptr.
|
|
//
|
|
// Without padding:
|
|
// (Note: Physical length equals logical length)
|
|
//
|
|
// Group mode:
|
|
// - seqstart_q_ptr, seqstart_k_ptr: Record cumulative physical sequence lengths. [array
|
|
// size: batch + 1]
|
|
// - seqlen_q_ptr/seqlen_k_ptr and cu_seqlen_q_ptr/cu_seqlen_k_ptr must be nullptr.
|
|
//
|
|
// Batch mode:
|
|
// - All sequence length pointers (seqstart_*, seqlen_*, cu_seqlen_*) must be nullptr.
|
|
//
|
|
const void* seqstart_q_ptr =
|
|
nullptr; // Cumulative physical sequence length array [batch + 1]. (Used in Group mode)
|
|
const void* seqstart_k_ptr =
|
|
nullptr; // Cumulative physical sequence length array [batch + 1]. (Used in Group mode)
|
|
const void* seqlen_q_ptr = nullptr; // Per-sequence logical (excluding padding) length array
|
|
// [batch]. (Used in Group mode with padding)
|
|
const void* seqlen_k_ptr = nullptr; // Per-sequence logical (excluding padding) length array
|
|
// [batch]. (Used in Group mode with padding)
|
|
const void* cu_seqlen_q_ptr = nullptr; // Cumulative logical (excluding padding) sequence length
|
|
// array [batch + 1]. (Used with padding)
|
|
const void* cu_seqlen_k_ptr = nullptr; // Cumulative logical (excluding padding) sequence length
|
|
// array [batch + 1]. (Used with padding)
|
|
ck_tile::index_t seqlen_q;
|
|
ck_tile::index_t seqlen_k;
|
|
ck_tile::index_t batch;
|
|
ck_tile::index_t max_seqlen_q;
|
|
ck_tile::index_t max_seqlen_k;
|
|
ck_tile::index_t hdim_q;
|
|
ck_tile::index_t hdim_v;
|
|
ck_tile::index_t nhead_q;
|
|
ck_tile::index_t nhead_k;
|
|
float scale;
|
|
ck_tile::index_t stride_q;
|
|
ck_tile::index_t stride_k;
|
|
ck_tile::index_t stride_v;
|
|
ck_tile::index_t stride_bias; // if alibi, b*h need set this to h, 1*h need set this to 0
|
|
ck_tile::index_t stride_o;
|
|
ck_tile::index_t stride_randval;
|
|
ck_tile::index_t stride_do;
|
|
ck_tile::index_t stride_dq_acc;
|
|
ck_tile::index_t stride_dq;
|
|
ck_tile::index_t stride_dk;
|
|
ck_tile::index_t stride_dv;
|
|
ck_tile::index_t stride_dbias;
|
|
ck_tile::index_t nhead_stride_q;
|
|
ck_tile::index_t nhead_stride_k;
|
|
ck_tile::index_t nhead_stride_v;
|
|
ck_tile::index_t nhead_stride_bias;
|
|
ck_tile::index_t nhead_stride_o;
|
|
ck_tile::index_t nhead_stride_randval;
|
|
ck_tile::index_t nhead_stride_do;
|
|
ck_tile::index_t nhead_stride_lsed;
|
|
ck_tile::index_t nhead_stride_dq_acc;
|
|
ck_tile::index_t nhead_stride_dq;
|
|
ck_tile::index_t nhead_stride_dk;
|
|
ck_tile::index_t nhead_stride_dv;
|
|
ck_tile::index_t nhead_stride_dbias;
|
|
ck_tile::index_t batch_stride_q;
|
|
ck_tile::index_t batch_stride_k;
|
|
ck_tile::index_t batch_stride_v;
|
|
ck_tile::index_t batch_stride_bias;
|
|
ck_tile::index_t batch_stride_o;
|
|
ck_tile::index_t batch_stride_randval;
|
|
ck_tile::index_t batch_stride_do;
|
|
ck_tile::index_t batch_stride_lsed;
|
|
ck_tile::index_t batch_stride_dq_acc;
|
|
ck_tile::index_t batch_stride_dq;
|
|
ck_tile::index_t batch_stride_dk;
|
|
ck_tile::index_t batch_stride_dv;
|
|
ck_tile::index_t batch_stride_dbias;
|
|
ck_tile::index_t split_stride_dq_acc;
|
|
ck_tile::index_t window_size_left;
|
|
ck_tile::index_t window_size_right;
|
|
ck_tile::index_t mask_type;
|
|
float p_drop;
|
|
float p_undrop;
|
|
std::variant<std::pair<uint64_t, uint64_t>, std::pair<const void*, const void*>>
|
|
drop_seed_offset;
|
|
};
|
|
|
|
template <typename FmhaBwdDQDKDVKernel>
|
|
auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args)
|
|
{
|
|
assert(args.nhead_q % args.nhead_k == 0);
|
|
auto kargs = [&] {
|
|
constexpr bool dq_uss_acc = FmhaBwdDQDKDVKernel::kMaxSeqLenQ == 0;
|
|
const auto dq_ptr = dq_uss_acc ? args.dq_acc_ptr : args.dq_ptr;
|
|
const auto stride_dq = dq_uss_acc ? args.stride_dq_acc : args.stride_dq;
|
|
const auto nhead_stride_dq = dq_uss_acc ? args.nhead_stride_dq_acc : args.nhead_stride_dq;
|
|
const auto batch_stride_dq = dq_uss_acc ? args.batch_stride_dq_acc : args.batch_stride_dq;
|
|
|
|
// create group mode kernel arguments
|
|
if constexpr(FmhaBwdDQDKDVKernel::kIsGroupMode)
|
|
{
|
|
return FmhaBwdDQDKDVKernel::MakeKargsImpl(args.q_ptr,
|
|
args.k_ptr,
|
|
args.v_ptr,
|
|
args.bias_ptr,
|
|
args.lse_ptr,
|
|
args.do_ptr,
|
|
args.d_ptr,
|
|
args.rand_val_ptr,
|
|
args.dk_ptr,
|
|
args.dv_ptr,
|
|
args.dbias_ptr,
|
|
dq_ptr,
|
|
args.seqstart_q_ptr,
|
|
args.seqstart_k_ptr,
|
|
args.seqlen_q_ptr,
|
|
args.seqlen_k_ptr,
|
|
args.cu_seqlen_q_ptr,
|
|
args.cu_seqlen_k_ptr,
|
|
args.hdim_q,
|
|
args.hdim_v,
|
|
args.nhead_q,
|
|
args.nhead_q / args.nhead_k,
|
|
args.scale,
|
|
args.stride_q,
|
|
args.stride_k,
|
|
args.stride_v,
|
|
args.stride_bias,
|
|
args.stride_randval,
|
|
args.stride_do,
|
|
stride_dq,
|
|
args.stride_dk,
|
|
args.stride_dv,
|
|
args.stride_dbias,
|
|
args.nhead_stride_q,
|
|
args.nhead_stride_k,
|
|
args.nhead_stride_v,
|
|
args.nhead_stride_bias,
|
|
args.nhead_stride_randval,
|
|
args.nhead_stride_do,
|
|
args.nhead_stride_lsed,
|
|
nhead_stride_dq,
|
|
args.nhead_stride_dk,
|
|
args.nhead_stride_dv,
|
|
args.nhead_stride_dbias,
|
|
args.split_stride_dq_acc,
|
|
args.window_size_left,
|
|
args.window_size_right,
|
|
args.mask_type,
|
|
args.p_drop,
|
|
args.drop_seed_offset);
|
|
}
|
|
else
|
|
{ // create batch mode kernel arguments
|
|
return FmhaBwdDQDKDVKernel::MakeKargsImpl(args.q_ptr,
|
|
args.k_ptr,
|
|
args.v_ptr,
|
|
args.bias_ptr,
|
|
args.lse_ptr,
|
|
args.do_ptr,
|
|
args.d_ptr,
|
|
args.rand_val_ptr,
|
|
args.dk_ptr,
|
|
args.dv_ptr,
|
|
args.dbias_ptr,
|
|
dq_ptr,
|
|
args.seqlen_q,
|
|
args.seqlen_k,
|
|
args.hdim_q,
|
|
args.hdim_v,
|
|
args.nhead_q,
|
|
args.nhead_q / args.nhead_k,
|
|
args.scale,
|
|
args.stride_q,
|
|
args.stride_k,
|
|
args.stride_v,
|
|
args.stride_bias,
|
|
args.stride_randval,
|
|
args.stride_do,
|
|
stride_dq,
|
|
args.stride_dk,
|
|
args.stride_dv,
|
|
args.stride_dbias,
|
|
args.nhead_stride_q,
|
|
args.nhead_stride_k,
|
|
args.nhead_stride_v,
|
|
args.nhead_stride_bias,
|
|
args.nhead_stride_randval,
|
|
args.nhead_stride_do,
|
|
args.nhead_stride_lsed,
|
|
nhead_stride_dq,
|
|
args.nhead_stride_dk,
|
|
args.nhead_stride_dv,
|
|
args.nhead_stride_dbias,
|
|
args.batch_stride_q,
|
|
args.batch_stride_k,
|
|
args.batch_stride_v,
|
|
args.batch_stride_bias,
|
|
args.batch_stride_randval,
|
|
args.batch_stride_do,
|
|
args.batch_stride_lsed,
|
|
batch_stride_dq,
|
|
args.batch_stride_dk,
|
|
args.batch_stride_dv,
|
|
args.batch_stride_dbias,
|
|
args.split_stride_dq_acc,
|
|
args.window_size_left,
|
|
args.window_size_right,
|
|
args.mask_type,
|
|
args.p_drop,
|
|
args.drop_seed_offset);
|
|
}
|
|
}();
|
|
|
|
dim3 grids = FmhaBwdDQDKDVKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_k);
|
|
return ck_tile::make_tuple(kargs, grids);
|
|
}
|
|
|
|
template <typename FmhaBwdOGradDotOKernel>
|
|
auto fmha_bwd_dot_do_o_create_kargs_and_grids(fmha_bwd_args args)
|
|
{
|
|
auto kargs = [&] {
|
|
// create group mode kernel arguments
|
|
if constexpr(FmhaBwdOGradDotOKernel::kIsGroupMode)
|
|
{
|
|
return FmhaBwdOGradDotOKernel::MakeKargs(args.o_ptr,
|
|
args.do_ptr,
|
|
args.d_ptr,
|
|
args.p_undrop,
|
|
args.seqstart_q_ptr,
|
|
args.seqlen_q_ptr,
|
|
args.cu_seqlen_q_ptr,
|
|
args.hdim_v,
|
|
args.stride_do,
|
|
args.stride_o,
|
|
args.nhead_stride_do,
|
|
args.nhead_stride_o,
|
|
args.nhead_stride_lsed);
|
|
}
|
|
else
|
|
{ // create batch mode kernel arguments
|
|
return FmhaBwdOGradDotOKernel::MakeKargs(args.o_ptr,
|
|
args.do_ptr,
|
|
args.d_ptr,
|
|
args.p_undrop,
|
|
args.seqlen_q,
|
|
args.hdim_v,
|
|
args.stride_do,
|
|
args.stride_o,
|
|
args.nhead_stride_do,
|
|
args.nhead_stride_o,
|
|
args.nhead_stride_lsed,
|
|
args.batch_stride_do,
|
|
args.batch_stride_o,
|
|
args.batch_stride_lsed);
|
|
}
|
|
}();
|
|
|
|
dim3 grids = FmhaBwdOGradDotOKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q);
|
|
return ck_tile::make_tuple(kargs, grids);
|
|
}
|
|
|
|
template <typename FmhaBwdConvertQGradKernel>
|
|
auto fmha_bwd_convert_dq_create_kargs_and_grids(fmha_bwd_args args)
|
|
{
|
|
auto kargs = [&] {
|
|
// create group mode kernel arguments
|
|
if constexpr(FmhaBwdConvertQGradKernel::kIsGroupMode)
|
|
{
|
|
return FmhaBwdConvertQGradKernel::MakeKargs(args.dq_acc_ptr,
|
|
args.dq_ptr,
|
|
args.seqstart_q_ptr,
|
|
args.seqstart_k_ptr,
|
|
args.seqlen_q_ptr,
|
|
args.seqlen_k_ptr,
|
|
args.cu_seqlen_q_ptr,
|
|
args.cu_seqlen_k_ptr,
|
|
args.hdim_q,
|
|
args.stride_dq,
|
|
args.stride_dq_acc,
|
|
args.nhead_stride_dq,
|
|
args.nhead_stride_dq_acc,
|
|
args.split_stride_dq_acc);
|
|
}
|
|
else
|
|
{ // create batch mode kernel arguments
|
|
return FmhaBwdConvertQGradKernel::MakeKargs(args.dq_acc_ptr,
|
|
args.dq_ptr,
|
|
args.seqlen_q,
|
|
args.seqlen_k,
|
|
args.hdim_q,
|
|
args.stride_dq,
|
|
args.stride_dq_acc,
|
|
args.nhead_stride_dq,
|
|
args.nhead_stride_dq_acc,
|
|
args.batch_stride_dq,
|
|
args.batch_stride_dq_acc,
|
|
args.split_stride_dq_acc);
|
|
}
|
|
}();
|
|
|
|
dim3 grids = FmhaBwdConvertQGradKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q);
|
|
return ck_tile::make_tuple(kargs, grids);
|
|
}
|
|
|
|
// this is used to pattern-match internl kernel implementation, not to instantiate kernel
|
|
template <ck_tile::index_t HDim_,
|
|
typename DataType_,
|
|
bool kIsGroupMode_,
|
|
typename FmhaMask_,
|
|
typename FmhaDropout_,
|
|
ck_tile::BlockAttentionBiasEnum BiasEnum_,
|
|
bool kHasBiasGrad_,
|
|
ck_tile::index_t kPadD_,
|
|
ck_tile::index_t kPadDv_,
|
|
bool kIsDeterministic_,
|
|
bool kUseTrLoad_,
|
|
ck_tile::index_t MaxSeqLenQ_,
|
|
ck_tile::index_t kN0>
|
|
struct fmha_bwd_dq_dk_dv_traits_
|
|
{
|
|
};
|
|
|
|
template <typename Traits_, typename Arch = void>
|
|
float fmha_bwd_dq_dk_dv_(const ck_tile::stream_config&, fmha_bwd_args);
|
|
|
|
template <typename Traits_, typename Arch = void>
|
|
void fmha_bwd_dq_dk_dv_oneshot_(const ck_tile::stream_config&, fmha_bwd_args);
|
|
|
|
template <typename Traits_, typename Arch = void>
|
|
std::string fmha_bwd_dq_dk_dv_get_name_();
|
|
template <typename Traits_, typename Arch = void>
|
|
int fmha_bwd_dq_dk_dv_maxq_();
|
|
|
|
template <ck_tile::index_t HDim_, typename DataType_, bool kIsGroupMode_, bool kPadS_, bool kPadDv_>
|
|
struct fmha_bwd_dot_do_o_traits_
|
|
{
|
|
static constexpr ck_tile::index_t HDim = HDim_;
|
|
using DataType = ck_tile::remove_cvref_t<DataType_>;
|
|
static constexpr bool kIsGroupMode = kIsGroupMode_;
|
|
static constexpr bool kPadS = kPadS_;
|
|
static constexpr bool kPadDv = kPadDv_;
|
|
};
|
|
|
|
template <typename Traits_, typename Arch = void>
|
|
float fmha_bwd_dot_do_o_(const ck_tile::stream_config&, fmha_bwd_args);
|
|
|
|
template <typename Traits_, typename Arch = void>
|
|
void fmha_bwd_dot_do_o_oneshot_(const ck_tile::stream_config&, fmha_bwd_args);
|
|
|
|
template <typename Traits_, typename Arch = void>
|
|
std::string fmha_bwd_dot_do_o_get_name_();
|
|
|
|
template <ck_tile::index_t HDim_,
|
|
typename DataType_,
|
|
bool kIsGroupMode_,
|
|
bool kPadS_,
|
|
bool kPadD_,
|
|
bool kIsDeterministic_,
|
|
ck_tile::index_t kN0>
|
|
struct fmha_bwd_convert_dq_traits_
|
|
{
|
|
};
|
|
|
|
template <typename Traits_, typename Arch = void>
|
|
float fmha_bwd_convert_dq_(const ck_tile::stream_config&, fmha_bwd_args);
|
|
|
|
template <typename Traits_, typename Arch = void>
|
|
void fmha_bwd_convert_dq_oneshot_(const ck_tile::stream_config&, fmha_bwd_args);
|
|
|
|
template <typename Traits_, typename Arch = void>
|
|
std::string fmha_bwd_convert_dq_get_name_();
|
|
|
|
// This is the public API, will be generated by script
|
|
struct fmha_bwd_traits
|
|
{
|
|
int hdim_q;
|
|
int hdim_v;
|
|
std::string data_type;
|
|
bool is_group_mode;
|
|
mask_enum mask_type;
|
|
bias_enum bias_type; // 0:no bias, 1:elementwise bias, 2:alibi. sync with BlockAttentionBiasEnum
|
|
bool has_dbias;
|
|
bool has_dropout;
|
|
bool is_store_randval;
|
|
bool is_deterministic;
|
|
// TODO: padding check is inside this api
|
|
};
|
|
template <int Version = 2>
|
|
float fmha_bwd(fmha_bwd_traits, fmha_bwd_args, const ck_tile::stream_config&);
|