lnigam
7b8443ac78
* ggml-cuda: add flash-attn support for DKQ=320/DV=256 with ncols2=32 (GQA=32) Adds MMA-f16 and tile kernel configs, dispatch logic, template instances, and tile .cu file for Mistral Small 4 (head sizes 320/256), restricting to ncols2=32 to support GQA ratio 32 only. * Adding check to return BEST_FATTN_KERNEL_NONE in case GQA!=32 * Apply suggestions from code review Address review comments Co-authored-by: Johannes Gäßler <johannesg@5d6.de> * Address review comments and making kernel config default to DQK=512, DV=512 instead of DQK=256,DV=256 * Fixed bug with sinks=1, with ncols=32, there are two warp-groups created but sinks index is same(0,...,15) for both the groups hence with sinks=1, output is not matching with CPU output. Added sink_base which will be base index for each warp_group (threadIdx.y / np) * Apply suggestions from code review Co-authored-by: Johannes Gäßler <johannesg@5d6.de> * Update ggml/src/ggml-cuda/template-instances/generate_cu_files.py Co-authored-by: Johannes Gäßler <johannesg@5d6.de> --------- Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
58 lines
2.0 KiB
Plaintext
58 lines
2.0 KiB
Plaintext
#include "common.cuh"
|
|
#include "fattn-tile.cuh"
|
|
#include "fattn-wmma-f16.cuh"
|
|
|
|
void ggml_cuda_flash_attn_ext_tile(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
|
const ggml_tensor * K = dst->src[1];
|
|
const ggml_tensor * V = dst->src[2];
|
|
switch (K->ne[0]) {
|
|
case 40: {
|
|
GGML_ASSERT(V->ne[0] == K->ne[0]);
|
|
ggml_cuda_flash_attn_ext_tile_case< 40, 40>(ctx, dst);
|
|
} break;
|
|
case 64: {
|
|
GGML_ASSERT(V->ne[0] == K->ne[0]);
|
|
ggml_cuda_flash_attn_ext_tile_case< 64, 64>(ctx, dst);
|
|
} break;
|
|
case 72: {
|
|
GGML_ASSERT(V->ne[0] == K->ne[0]);
|
|
ggml_cuda_flash_attn_ext_tile_case< 72, 72>(ctx, dst);
|
|
} break;
|
|
case 80: {
|
|
GGML_ASSERT(V->ne[0] == K->ne[0]);
|
|
ggml_cuda_flash_attn_ext_tile_case< 80, 80>(ctx, dst);
|
|
} break;
|
|
case 96: {
|
|
GGML_ASSERT(V->ne[0] == K->ne[0]);
|
|
ggml_cuda_flash_attn_ext_tile_case< 96, 96>(ctx, dst);
|
|
} break;
|
|
case 112: {
|
|
GGML_ASSERT(V->ne[0] == K->ne[0]);
|
|
ggml_cuda_flash_attn_ext_tile_case<112, 112>(ctx, dst);
|
|
} break;
|
|
case 128: {
|
|
GGML_ASSERT(V->ne[0] == K->ne[0]);
|
|
ggml_cuda_flash_attn_ext_tile_case<128, 128>(ctx, dst);
|
|
} break;
|
|
case 256: {
|
|
GGML_ASSERT(V->ne[0] == K->ne[0]);
|
|
ggml_cuda_flash_attn_ext_tile_case<256, 256>(ctx, dst);
|
|
} break;
|
|
case 320: {
|
|
GGML_ASSERT(V->ne[0] == 256);
|
|
ggml_cuda_flash_attn_ext_tile_case<320, 256>(ctx, dst);
|
|
} break;
|
|
case 512: {
|
|
GGML_ASSERT(V->ne[0] == K->ne[0]);
|
|
ggml_cuda_flash_attn_ext_tile_case<512, 512>(ctx, dst);
|
|
} break;
|
|
case 576: {
|
|
GGML_ASSERT(V->ne[0] == 512);
|
|
ggml_cuda_flash_attn_ext_tile_case<576, 512>(ctx, dst);
|
|
} break;
|
|
default: {
|
|
GGML_ABORT("Unsupported head size");
|
|
} break;
|
|
}
|
|
}
|