From 8998067ad2dc19adcbaf81ceaca0201de118d7a3 Mon Sep 17 00:00:00 2001 From: zichguan-amd Date: Fri, 16 Oct 2026 17:27:17 -0400 Subject: [PATCH 01/13] ck: sync ck_tile examples with TheRock 10.1.0rc3 Sync Libraries/ComposableKernel with ROCm/rocm-libraries@39d8d60365d2272a5366f76ccb3bfba24b219eaf (projects/composablekernel/example/ck_tile), the rocm-libraries pin of the TheRock 10.1.0rc3 release candidate: https://rc.repo.amd.com/rocm/core/tarball/therock-dist-linux-gfx110X-all-tests-10.1.0rc3.tar.gz Synced files are upstream's, license headers included, formatted in rocm-examples style by this repo's pre-commit hooks (clang-format 18.1.4 without a .clang-format, i.e. LLVM style; whitespace; black). Changed rocm-examples-owned files keep their headers, with the copyright year brought up to date. - updated: 220 files, unchanged: 0 - upstream example dir missing (kept): - basic/copy -> 39_copy - overlays (manual CMake/README adaptations): - applied 0001-ck-therock-10.1-adaptations.patch Validated locally at the branch tip (with any follow-up fixup commits): pre-commit, linting.yml checks, CK build + ctest on gfx1100 and CK build for build_libraries.yml GPU_TARGETS in ghcr.io/rocm/rocm-examples-ubuntu-24.04-multiarch:latest with the RC tarball. Co-Authored-By: Claude --- .../ComposableKernel/attention/CMakeLists.txt | 3 +- .../attention/fmha/CMakeLists.txt | 17 +- .../ComposableKernel/attention/fmha/README.md | 2 +- .../ComposableKernel/attention/fmha/bias.hpp | 32 +- .../attention/fmha/codegen/__init__.py | 2 + .../attention/fmha/codegen/arch.py | 2 +- .../attention/fmha/codegen/cmake_config.py | 3 +- .../attention/fmha/codegen/cpp_symbol_map.py | 45 +- .../attention/fmha/codegen/ops/__init__.py | 2 + .../fmha/codegen/ops/fmha_batch_prefill.py | 437 ++++-- .../attention/fmha/codegen/ops/fmha_bwd.py | 244 ++- .../attention/fmha/codegen/ops/fmha_fwd.py | 1317 +++++++++++++---- .../fmha/codegen/ops/fmha_fwd_appendkv.py | 35 +- .../fmha/codegen/ops/fmha_fwd_splitkv.py | 151 +- .../fmha/codegen/ops/fmha_pagedkv_prefill.py | 90 +- .../attention/fmha/codegen/utils.py | 3 +- .../attention/fmha/fmha_bwd.cpp | 33 +- .../attention/fmha/fmha_bwd.hpp | 494 +++++-- .../attention/fmha/fmha_bwd_runner.hpp | 382 +++-- .../attention/fmha/fmha_fwd.cpp | 79 +- .../attention/fmha/fmha_fwd.hpp | 464 ++++-- .../attention/fmha/fmha_fwd_head_grouping.hpp | 415 ++++++ .../attention/fmha/fmha_fwd_runner.hpp | 1241 ++++++++++++---- .../attention/fmha/generate.py | 22 +- .../ComposableKernel/attention/fmha/mask.hpp | 72 +- .../ComposableKernel/attention/fmha/quant.hpp | 68 + .../attention/fmha/rotary.hpp | 28 +- .../ComposableKernel/attention/fmha/utils.hpp | 23 +- .../attention/fmha_fwd_v3/.gitignore | 1 - .../attention/fmha_fwd_v3/CMakeLists.txt | 92 -- .../attention/fmha_fwd_v3/README.md | 96 -- .../attention/fmha_fwd_v3/bias.hpp | 109 -- .../fmha_fwd_v3/example_fmha_fwd_v3.cpp | 581 -------- .../attention/fmha_fwd_v3/fmha_fwd.hpp | 1080 -------------- .../attention/fmha_fwd_v3/fmha_fwd_v3.cpp | 76 - .../attention/fmha_fwd_v3/fmha_fwd_v3.hpp | 91 -- .../fmha_fwd_v3/fmha_fwd_v3_impl.hpp | 170 --- .../instances/fmha_fwd_v3_d128_bf16_mask.cpp | 34 - .../instances/fmha_fwd_v3_d128_bf16_nmask.cpp | 34 - .../instances/fmha_fwd_v3_d128_fp16_mask.cpp | 34 - .../instances/fmha_fwd_v3_d128_fp16_nmask.cpp | 34 - .../attention/fmha_fwd_v3/mask.hpp | 162 -- .../attention/fmha_fwd_v3/rotary.hpp | 101 -- .../fmha_fwd_v3/test_ck_tile_fmha_fwd_v3.sh | 46 - .../attention/topk_softmax/topk_softmax.cpp | 34 +- .../topk_softmax/topk_softmax_api.cpp | 23 +- .../topk_softmax/topk_softmax_api.hpp | 23 +- .../batched_transpose_api.cpp | 84 +- .../batched_transpose_example.cpp | 33 +- .../batched_transpose_example.hpp | 30 +- .../basic/elementwise/elementwise_common.hpp | 26 +- .../basic/elementwise/elementwise_example.cpp | 28 +- .../elementwise_example_add_4d.cpp | 28 +- .../elementwise_example_transpose.cpp | 28 +- .../elementwise/elementwise_example_unary.cpp | 55 +- .../basic/img2col/image_to_column.cpp | 23 +- .../basic/img2col/image_to_column.hpp | 23 +- .../alternative_impl/matrix_core_swizzle.cpp | 23 +- .../alternative_impl/matrix_core_swizzle.hpp | 24 +- .../matrix_core_swizzle_kernel.hpp | 23 +- .../basic/permute/permute.cpp | 23 +- .../basic/permute/permute.hpp | 23 +- .../ComposableKernel/basic/reduce/README.md | 34 +- .../ComposableKernel/basic/reduce/reduce.cpp | 98 +- .../ComposableKernel/basic/reduce/reduce.hpp | 181 --- .../batched_contraction/CMakeLists.txt | 8 +- .../batched_contraction.cpp | 151 +- .../batched_contraction/contraction_utils.hpp | 112 +- .../run_batched_contraction_example.inc | 255 +++- .../grouped_convolution/CMakeLists.txt | 9 +- .../convolution/grouped_convolution/README.md | 4 - .../grouped_convolution/conv_configs.hpp | 145 +- .../grouped_convolution_backward_data.cpp | 26 +- ...uped_convolution_backward_data_invoker.hpp | 201 +-- .../grouped_convolution_backward_weight.cpp | 28 +- ...ed_convolution_backward_weight_invoker.hpp | 246 ++- ..._convolution_backward_weight_two_stage.cpp | 30 +- ...tion_backward_weight_two_stage_invoker.hpp | 339 ++--- .../grouped_convolution_forward.cpp | 35 +- ...grouped_convolution_forward_bias_clamp.cpp | 26 +- .../grouped_convolution_forward_invoker.hpp | 211 +-- ...ouped_convolution_forward_large_tensor.cpp | 74 - ...nvolution_forward_large_tensor_invoker.hpp | 360 ----- .../grouped_convolution_utils.hpp | 70 +- ...n_grouped_convolution_bwd_data_example.inc | 68 +- ...grouped_convolution_bwd_weight_example.inc | 123 +- ...ped_convolution_fwd_bias_clamp_example.inc | 58 +- .../run_grouped_convolution_fwd_example.inc | 69 +- .../gemm/batched_gemm/README.md | 8 - .../gemm/batched_gemm/batched_gemm.cpp | 164 +- .../gemm/batched_gemm/batched_gemm.hpp | 35 +- .../batched_gemm/run_batched_gemm_example.inc | 23 +- .../gemm/block_scale_gemm/CMakeLists.txt | 21 +- .../gemm/block_scale_gemm/README.md | 94 +- .../gemm_bquant_quantgrouped_bf8.cpp | 63 + .../gemm_bquant_quantgrouped_bf8i4.cpp | 63 + .../gemm_bquant_quantgrouped_fp8.cpp | 63 + .../gemm_bquant_quantgrouped_fp8i4.cpp | 63 + .../gemm/block_scale_gemm/gemm_quant.cpp | 124 ++ .../block_scale_gemm/gemm_quant_basic.cpp | 389 ----- .../gemm/block_scale_gemm/gemm_utils.hpp | 315 ++-- .../run_gemm_quant_example.inc | 733 +++++++-- .../ComposableKernel/gemm/flatmm/README.md | 2 - .../gemm/flatmm/flatmm_basic.cpp | 286 ++-- .../gemm/flatmm/flatmm_basic.hpp | 52 +- .../gemm/flatmm/grouped_flatmm.cpp | 221 ++- .../gemm/flatmm/moe_flatmm.cpp | 84 +- .../gemm/flatmm/moe_flatmm.hpp | 88 +- .../gemm/flatmm/run_flatmm_example.inc | 28 +- .../flatmm/run_grouped_flatmm_example.inc | 25 +- .../gemm/flatmm/run_moe_flatmm_example.inc | 50 +- .../ComposableKernel/gemm/gemm/README.md | 6 - .../ComposableKernel/gemm/gemm/gemm_basic.cpp | 41 +- .../gemm/gemm/gemm_basic_invoker.hpp | 214 ++- .../gemm/gemm/gemm_splitk_two_stage.cpp | 28 +- .../gemm/gemm_splitk_two_stage_invoker.hpp | 344 ++--- .../gemm/gemm_splitk_two_stage_reduce.cpp | 271 ++-- .../ComposableKernel/gemm/gemm/gemm_utils.hpp | 315 ++-- .../gemm/gemm/gemm_weight_preshuffle.cpp | 26 +- .../gemm/gemm_weight_preshuffle_invoker.hpp | 250 ++-- .../gemm/gemm/run_gemm_example.inc | 166 +-- .../gemm/gemm/run_gemm_example_common.hpp | 28 +- .../gemm/gemm/universal_gemm.cpp | 245 ++- .../gemm/gemm/universal_gemm_invoker.hpp | 447 ++++-- .../gemm/gemm_multi_abd/CMakeLists.txt | 8 +- .../gemm_multi_abd/gemm_multi_abd_fp16.cpp | 151 +- .../gemm_multi_abd/gemm_multi_abd_fp16.hpp | 35 +- .../run_gemm_multi_abd_fp16_example.inc | 23 +- .../gemm/gemm_multi_abd/utils.hpp | 23 +- .../gemm/gemm_multi_d/README.md | 8 - .../gemm/gemm_multi_d/gemm_multi_d_fp16.cpp | 153 +- .../gemm/gemm_multi_d/gemm_multi_d_fp16.hpp | 35 +- .../run_gemm_multi_d_fp16_example.inc | 23 +- .../gemm/gemm_multi_d/utils.hpp | 23 +- .../gemm/grouped_gemm/CMakeLists.txt | 17 +- .../gemm/grouped_gemm/README.md | 4 - .../gemm/grouped_gemm/grouped_gemm.cpp | 259 ++-- .../gemm/grouped_gemm/grouped_gemm.hpp | 126 +- .../grouped_gemm/grouped_gemm_multi_d.cpp | 253 ++-- .../grouped_gemm/grouped_gemm_multi_d.hpp | 53 +- .../grouped_gemm/grouped_gemm_preshuffle.cpp | 253 ++-- .../gemm/grouped_gemm/quant_grouped_gemm.cpp | 243 ++- .../gemm/grouped_gemm/quant_grouped_gemm.hpp | 179 --- .../quant_grouped_gemm_bf8_aquant.cpp | 8 + .../quant_grouped_gemm_bf8_bquant.cpp | 8 + .../quant_grouped_gemm_bf8_rowcol.cpp | 8 + .../quant_grouped_gemm_bf8_tensor.cpp | 8 + .../quant_grouped_gemm_config.hpp | 143 ++ .../quant_grouped_gemm_fp8_aquant.cpp | 8 + .../quant_grouped_gemm_fp8_bquant.cpp | 8 + .../quant_grouped_gemm_fp8_rowcol.cpp | 8 + .../quant_grouped_gemm_fp8_tensor.cpp | 8 + .../quant_invoke_grouped_gemm_kernel.hpp | 210 +++ .../quant_run_grouped_gemm_example.hpp | 510 +++++++ .../quant_run_grouped_gemm_example.inc | 606 -------- .../grouped_gemm/run_grouped_gemm_example.inc | 30 +- .../run_grouped_gemm_multi_d_example.inc | 49 +- .../gemm/streamk_gemm/CMakeLists.txt | 8 +- .../gemm/streamk_gemm/gemm_utils.hpp | 124 +- .../gemm/streamk_gemm/run_gemm_example.inc | 434 +++--- .../gemm/streamk_gemm/streamk_gemm_basic.cpp | 308 ++-- .../moe/fused_moe/fused_moe.hpp | 26 +- .../moe/fused_moe/fused_moegemm.hpp | 23 +- .../moe/fused_moe/fused_moesorting.hpp | 26 +- .../moe/fused_moe/instances/fused_moe_api.cpp | 26 +- .../fused_moe/instances/fused_moegemm_api.cpp | 26 +- .../instances/fused_moegemm_api_internal.hpp | 26 +- .../instances/fused_moegemm_api_traits.hpp | 23 +- .../instances/fused_moegemm_bf16_m32.cpp | 29 +- .../instances/fused_moegemm_fp16_m32.cpp | 29 +- .../instances/fused_moesorting_api.cpp | 23 +- .../ComposableKernel/moe/fused_moe/main.cpp | 87 +- .../moe_smoothquant_bf16_n1024_instance.cpp | 31 +- .../moe_smoothquant_bf16_n1536_instance.cpp | 23 +- .../moe_smoothquant_bf16_n2048_instance.cpp | 23 +- .../moe_smoothquant_bf16_n256_instance.cpp | 23 +- .../moe_smoothquant_bf16_n3072_instance.cpp | 23 +- .../moe_smoothquant_bf16_n4096_instance.cpp | 23 +- ...moe_smoothquant_bf16_n4096_tp_instance.cpp | 23 +- .../moe_smoothquant_bf16_n512_instance.cpp | 23 +- ...moe_smoothquant_bf16_n64_n128_instance.cpp | 23 +- .../moe_smoothquant_bf16_n768_instance.cpp | 23 +- .../moe_smoothquant_fp16_n1024_instance.cpp | 31 +- .../moe_smoothquant_fp16_n1536_instance.cpp | 23 +- .../moe_smoothquant_fp16_n2048_instance.cpp | 23 +- .../moe_smoothquant_fp16_n256_instance.cpp | 23 +- .../moe_smoothquant_fp16_n3072_instance.cpp | 23 +- .../moe_smoothquant_fp16_n4096_instance.cpp | 23 +- ...moe_smoothquant_fp16_n4096_tp_instance.cpp | 23 +- .../moe_smoothquant_fp16_n512_instance.cpp | 23 +- ...moe_smoothquant_fp16_n64_n128_instance.cpp | 23 +- .../moe_smoothquant_fp16_n768_instance.cpp | 23 +- .../instances/moe_smoothquant_fwd_api.cpp | 26 +- .../moe_smoothquant_instance_common.hpp | 26 +- .../moe/moe_smoothquant/moe_smoothquant.cpp | 27 +- .../moe/moe_smoothquant/moe_smoothquant.hpp | 23 +- .../moe/moe_sorting/moe_sorting.cpp | 41 +- .../moe/moe_sorting/moe_sorting_api.cpp | 23 +- .../moe/moe_sorting/moe_sorting_api.hpp | 26 +- .../add_rmsnorm2d_rdquant_fwd.cpp | 235 +-- .../add_rmsnorm2d_rdquant_fwd.hpp | 23 +- .../example_add_rmsnorm2d_rdquant_fwd.cpp | 26 +- .../add_rmsnorm2d_rdquant_fwd_api.cpp | 26 +- ...norm2d_rdquant_fwd_bf16_n1024_instance.cpp | 31 +- ...norm2d_rdquant_fwd_bf16_n1536_instance.cpp | 23 +- ...norm2d_rdquant_fwd_bf16_n2048_instance.cpp | 23 +- ...snorm2d_rdquant_fwd_bf16_n256_instance.cpp | 23 +- ...norm2d_rdquant_fwd_bf16_n3072_instance.cpp | 23 +- ...norm2d_rdquant_fwd_bf16_n4096_instance.cpp | 23 +- ...snorm2d_rdquant_fwd_bf16_n512_instance.cpp | 23 +- ...m2d_rdquant_fwd_bf16_n64_n128_instance.cpp | 23 +- ...snorm2d_rdquant_fwd_bf16_n768_instance.cpp | 23 +- ...norm2d_rdquant_fwd_bf16_n8192_instance.cpp | 23 +- ...m2d_rdquant_fwd_bf16_n8192_tp_instance.cpp | 23 +- ...norm2d_rdquant_fwd_fp16_n1024_instance.cpp | 31 +- ...norm2d_rdquant_fwd_fp16_n1536_instance.cpp | 23 +- ...norm2d_rdquant_fwd_fp16_n2048_instance.cpp | 23 +- ...snorm2d_rdquant_fwd_fp16_n256_instance.cpp | 23 +- ...norm2d_rdquant_fwd_fp16_n3072_instance.cpp | 23 +- ...norm2d_rdquant_fwd_fp16_n4096_instance.cpp | 23 +- ...snorm2d_rdquant_fwd_fp16_n512_instance.cpp | 23 +- ...m2d_rdquant_fwd_fp16_n64_n128_instance.cpp | 23 +- ...snorm2d_rdquant_fwd_fp16_n768_instance.cpp | 23 +- ...norm2d_rdquant_fwd_fp16_n8192_instance.cpp | 23 +- ...m2d_rdquant_fwd_fp16_n8192_tp_instance.cpp | 23 +- ..._rmsnorm2d_rdquant_fwd_instance_common.hpp | 26 +- .../normalization/layernorm2d/generate.py | 9 +- .../layernorm2d/layernorm2d_fwd.cpp | 27 +- .../layernorm2d/layernorm2d_fwd.hpp | 23 +- .../rmsnorm2d/example_rmsnorm2d_fwd.cpp | 26 +- .../normalization/rmsnorm2d/generate.py | 9 +- .../normalization/rmsnorm2d/rmsnorm2d_fwd.cpp | 30 +- .../normalization/rmsnorm2d/rmsnorm2d_fwd.hpp | 23 +- .../pooling/pool3d/pool3d.cpp | 38 +- .../smoothquant/example_smoothquant.cpp | 476 +++--- .../smoothquant_bf16_n1024_instance.cpp | 53 +- .../smoothquant_bf16_n1536_instance.cpp | 43 +- .../smoothquant_bf16_n2048_instance.cpp | 45 +- .../smoothquant_bf16_n256_instance.cpp | 41 +- .../smoothquant_bf16_n3072_instance.cpp | 45 +- .../smoothquant_bf16_n4096_instance.cpp | 45 +- .../smoothquant_bf16_n4096_tp_instance.cpp | 45 +- .../smoothquant_bf16_n512_instance.cpp | 43 +- .../smoothquant_bf16_n64_n128_instance.cpp | 41 +- .../smoothquant_bf16_n768_instance.cpp | 41 +- .../smoothquant_fp16_n1024_instance.cpp | 53 +- .../smoothquant_fp16_n1536_instance.cpp | 43 +- .../smoothquant_fp16_n2048_instance.cpp | 45 +- .../smoothquant_fp16_n256_instance.cpp | 41 +- .../smoothquant_fp16_n3072_instance.cpp | 45 +- .../smoothquant_fp16_n4096_instance.cpp | 45 +- .../smoothquant_fp16_n4096_tp_instance.cpp | 45 +- .../smoothquant_fp16_n512_instance.cpp | 43 +- .../smoothquant_fp16_n64_n128_instance.cpp | 41 +- .../smoothquant_fp16_n768_instance.cpp | 41 +- .../instances/smoothquant_fwd_api.cpp | 26 +- .../instances/smoothquant_instance_common.hpp | 26 +- .../quantization/smoothquant/smoothquant.cpp | 27 +- .../quantization/smoothquant/smoothquant.hpp | 23 +- README.md | 1 - 260 files changed, 11343 insertions(+), 13867 deletions(-) create mode 100644 Libraries/ComposableKernel/attention/fmha/fmha_fwd_head_grouping.hpp create mode 100644 Libraries/ComposableKernel/attention/fmha/quant.hpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/.gitignore delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/CMakeLists.txt delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/README.md delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/bias.hpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/example_fmha_fwd_v3.cpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd.hpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3.cpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3.hpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3_impl.hpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_bf16_mask.cpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_bf16_nmask.cpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_fp16_mask.cpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_fp16_nmask.cpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/mask.hpp delete mode 100644 Libraries/ComposableKernel/attention/fmha_fwd_v3/rotary.hpp delete mode 100755 Libraries/ComposableKernel/attention/fmha_fwd_v3/test_ck_tile_fmha_fwd_v3.sh delete mode 100644 Libraries/ComposableKernel/basic/reduce/reduce.hpp delete mode 100644 Libraries/ComposableKernel/convolution/grouped_convolution/grouped_convolution_forward_large_tensor.cpp delete mode 100644 Libraries/ComposableKernel/convolution/grouped_convolution/grouped_convolution_forward_large_tensor_invoker.hpp create mode 100644 Libraries/ComposableKernel/gemm/block_scale_gemm/gemm_bquant_quantgrouped_bf8.cpp create mode 100644 Libraries/ComposableKernel/gemm/block_scale_gemm/gemm_bquant_quantgrouped_bf8i4.cpp create mode 100644 Libraries/ComposableKernel/gemm/block_scale_gemm/gemm_bquant_quantgrouped_fp8.cpp create mode 100644 Libraries/ComposableKernel/gemm/block_scale_gemm/gemm_bquant_quantgrouped_fp8i4.cpp create mode 100644 Libraries/ComposableKernel/gemm/block_scale_gemm/gemm_quant.cpp delete mode 100644 Libraries/ComposableKernel/gemm/block_scale_gemm/gemm_quant_basic.cpp delete mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_grouped_gemm.hpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_grouped_gemm_bf8_aquant.cpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_grouped_gemm_bf8_bquant.cpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_grouped_gemm_bf8_rowcol.cpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_grouped_gemm_bf8_tensor.cpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_grouped_gemm_config.hpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_grouped_gemm_fp8_aquant.cpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_grouped_gemm_fp8_bquant.cpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_grouped_gemm_fp8_rowcol.cpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_grouped_gemm_fp8_tensor.cpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_invoke_grouped_gemm_kernel.hpp create mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_run_grouped_gemm_example.hpp delete mode 100644 Libraries/ComposableKernel/gemm/grouped_gemm/quant_run_grouped_gemm_example.inc diff --git a/Libraries/ComposableKernel/attention/CMakeLists.txt b/Libraries/ComposableKernel/attention/CMakeLists.txt index 8c76047991..b0b12fd7f4 100644 --- a/Libraries/ComposableKernel/attention/CMakeLists.txt +++ b/Libraries/ComposableKernel/attention/CMakeLists.txt @@ -1,6 +1,6 @@ # MIT License # -# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal @@ -27,5 +27,4 @@ file(RELATIVE_PATH folder_bin ${CMAKE_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}) set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin/${folder_bin}) add_subdirectory(fmha) -add_subdirectory(fmha_fwd_v3) add_subdirectory(topk_softmax) diff --git a/Libraries/ComposableKernel/attention/fmha/CMakeLists.txt b/Libraries/ComposableKernel/attention/fmha/CMakeLists.txt index 0fdfb4bb32..059b30f822 100644 --- a/Libraries/ComposableKernel/attention/fmha/CMakeLists.txt +++ b/Libraries/ComposableKernel/attention/fmha/CMakeLists.txt @@ -1,6 +1,6 @@ # MIT License # -# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal @@ -28,9 +28,6 @@ set(example_fwd_script test_ck_tile_fmha_fwd.sh) set(example_bwd ${example_name}_bwd) set(example_bwd_script test_ck_tile_fmha_bwd.sh) -set(example_fwd_v3 ${example_name}_fwd_v3) -set(example_fwd_v3_script test_ck_tile_fmha_fwd_v3.sh) - cmake_minimum_required(VERSION 3.21 FATAL_ERROR) project(${example_name} LANGUAGES CXX) @@ -105,7 +102,15 @@ add_dependencies(${example_fwd} fmha_fwd_gen_srcs) set_source_files_properties(fmha_fwd.cpp ${EXAMPLE_FWD_GEN_SRCS} PROPERTIES LANGUAGE ${ROCM_EXAMPLES_GPU_LANGUAGE}) -target_compile_definitions(${example_fwd} PRIVATE CK_TILE_FMHA_FWD_FAST_EXP2=1) +target_compile_definitions( + ${example_fwd} + PRIVATE + CK_TILE_FMHA_FWD_FAST_EXP2=1 + CK_TILE_FMHA_FWD_SPLITKV_API=0 + CK_TILE_FMHA_FWD_APPENDKV_API=0 + CK_TILE_FMHA_FWD_PAGEDKV_API=0 + CK_TILE_FMHA_FWD_BATCH_PREFILL_API=0 +) target_compile_features(${example_fwd} PRIVATE cxx_std_17 hip_std_17) target_compile_options(${example_fwd} PRIVATE -Wno-float-equal -Wno-undefined-func-template @@ -151,7 +156,7 @@ add_custom_target(fmha_bwd_gen_srcs add_dependencies(${example_bwd} fmha_bwd_gen_srcs) target_compile_definitions(${example_bwd} PRIVATE CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=3) target_compile_features(${example_bwd} PRIVATE cxx_std_17 hip_std_17) -target_compile_options(${example_fwd} PRIVATE -Wno-float-equal +target_compile_options(${example_bwd} PRIVATE -Wno-float-equal -Wno-undefined-func-template -fgpu-flush-denormals-to-zero ) diff --git a/Libraries/ComposableKernel/attention/fmha/README.md b/Libraries/ComposableKernel/attention/fmha/README.md index 7c36a5aa6a..bb62bb4dc3 100644 --- a/Libraries/ComposableKernel/attention/fmha/README.md +++ b/Libraries/ComposableKernel/attention/fmha/README.md @@ -69,7 +69,6 @@ For more information on CK Tile terminology, refer to the * `ck_tile::bf8_t` * `ck_tile::BlockAttentionBiasEnum` * `ck_tile::BlockFmhaPipelineEnum` -* `ck_tile::composes` * `ck_tile::DeviceMem` * `ck_tile::FillNormalDistribution` * `ck_tile::FillNormalDistributionIntegerValue` @@ -98,6 +97,7 @@ For more information on CK Tile terminology, refer to the * `ck_tile::get_alibi_slopes` * `ck_tile::integer_divide_ceil` * `ck_tile::make_alibi_from_lr_mask` +* `ck_tile::make_composes` * `ck_tile::make_generic_attention_mask_coordinates_from_lr_window` * `ck_tile::make_tuple` * `ck_tile::numeric::max` diff --git a/Libraries/ComposableKernel/attention/fmha/bias.hpp b/Libraries/ComposableKernel/attention/fmha/bias.hpp index 0b4c728340..2198f46ca4 100644 --- a/Libraries/ComposableKernel/attention/fmha/bias.hpp +++ b/Libraries/ComposableKernel/attention/fmha/bias.hpp @@ -1,32 +1,15 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once -#include "ck_tile/core.hpp" -#include "ck_tile/ops/fmha.hpp" #include + #include +#include "ck_tile/core.hpp" +#include "ck_tile/ops/fmha.hpp" + // keep sync with BlockAttentionBiasEnum enum class bias_enum { no_bias = 0, @@ -102,7 +85,8 @@ struct bias_info { return info; } - friend std::ostream &operator<<(std::ostream &os, const bias_info &bi) { + friend std::ostream &operator<<([[clang::lifetimebound]] std::ostream &os, + const bias_info &bi) { bi.serialize(os); return os; } diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/__init__.py b/Libraries/ComposableKernel/attention/fmha/codegen/__init__.py index e69de29bb2..1df4857184 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/__init__.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/arch.py b/Libraries/ComposableKernel/attention/fmha/codegen/arch.py index 1bfc78d3cd..aeb9a98bbb 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/arch.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/arch.py @@ -1,5 +1,5 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. from dataclasses import dataclass, field from typing import Any, List, Callable diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/cmake_config.py b/Libraries/ComposableKernel/attention/fmha/codegen/cmake_config.py index 483934b03b..3399f58947 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/cmake_config.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/cmake_config.py @@ -1,5 +1,4 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. # generate kernel instances to speed up compilation - GEN_DIR = "" # in Cmake, have to generate files in same folder diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/cpp_symbol_map.py b/Libraries/ComposableKernel/attention/fmha/codegen/cpp_symbol_map.py index 4098eb67c2..61e4071681 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/cpp_symbol_map.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/cpp_symbol_map.py @@ -1,7 +1,6 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. # generate kernel instances to speed up compilation - FWD_DTYPE_MAP = { "fp32": "FmhaFwdFp32", "fp16": "FmhaFwdFp16", @@ -10,6 +9,8 @@ "fp8fp16": "FmhaFwdFp8Fp16", "fp8bf16": "FmhaFwdFp8Bf16", "fp8fp32": "FmhaFwdFp8Fp32", + "mxfp8": "FmhaFwdMxFp8", + "mxfp4": "FmhaFwdMxFp4", } BWD_DTYPE_MAP = {"fp32": "FmhaBwdFp32", "fp16": "FmhaBwdFp16", "bf16": "FmhaBwdBf16"} @@ -31,16 +32,24 @@ } -def get_mask_map(mask: str): - if mask == "generic": +def get_mask_map(mask_impl: str): + if mask_impl == "generic": return _MASK_MAP - elif mask == "simplified": + elif mask_impl == "simplified": return _MASK_SIMPLIFIED_MAP else: assert False return None +def get_mask_impl(mask: str) -> str: + return "simplified" if mask.startswith("s_") else "generic" + + +def get_mask_cpp_type(mask: str) -> str: + return get_mask_map(get_mask_impl(mask))[mask] + + _MASK_CHECK_MAP = { "no": "t.mask_type == mask_enum::no_mask", "causal": "t.mask_type == mask_enum::mask_top_left || t.mask_type == mask_enum::mask_bottom_right", @@ -63,6 +72,26 @@ def get_mask_check_map(mask: str): return None +def get_mask_cpp_check_expr(mask: str) -> str: + return get_mask_check_map(get_mask_impl(mask))[mask] + + +QSCALE_MAP = { + "no": "ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE", + "pertensor": "ck_tile::BlockAttentionQuantScaleEnum::PERTENSOR", + "blockscale": "ck_tile::BlockAttentionQuantScaleEnum::BLOCKSCALE", + "kv_blockscale": "ck_tile::BlockAttentionQuantScaleEnum::KV_BLOCKSCALE", + "mx": "ck_tile::BlockAttentionQuantScaleEnum::MX", +} + +QSCALE_CHECK_MAP = { + "no": "quant_scale_enum::no_scale", + "pertensor": "quant_scale_enum::pertensor", + "blockscale": "quant_scale_enum::blockscale", + "kv_blockscale": "quant_scale_enum::kv_blockscale", + "mx": "quant_scale_enum::mx", +} + BIAS_MAP = { "no": "ck_tile::BlockAttentionBiasEnum::NO_BIAS", "bias": "ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS", @@ -110,18 +139,24 @@ def get_mask_check_map(mask: str): PIPELINE_MAP = { "qr": "ck_tile::BlockFmhaPipelineQRKSVS", + "qr_hpad": "ck_tile::BlockFmhaPipelineQRKSVSHpad", "qr_async": "ck_tile::BlockFmhaPipelineQRKSVSAsync", "qs": "ck_tile::BlockFmhaPipelineQSKSVS", "qr_async_trload": "ck_tile::BlockFmhaPipelineQRKSVSAsyncTrload", + "qr_async_trload_v3": "ck_tile::BlockFmhaFwdV3Pipeline", + "qr_tdm": "ck_tile::BlockFmhaPipelineQRKSVSTdm", } PIPELINE_ENUM_MAP = { "qr": "ck_tile::BlockFmhaPipelineEnum::QRKSVS", + "qr_hpad": "ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD", "qr_async": "ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC", "qr_nwarp_sshuffle": "ck_tile::BlockFmhaPipelineEnum::QRKSVS", "qs": "ck_tile::BlockFmhaPipelineEnum::QSKSVS", "qr_pagedkv": "ck_tile::BlockFmhaPipelineEnum::QRKSVS", "qr_async_trload": "ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD", + "qr_async_trload_v3": "ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD_V3", + "qr_tdm": "ck_tile::BlockFmhaPipelineEnum::QRKSVS_TDM", } BOOL_MAP = { diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/ops/__init__.py b/Libraries/ComposableKernel/attention/fmha/codegen/ops/__init__.py index e69de29bb2..1df4857184 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/ops/__init__.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/ops/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_batch_prefill.py b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_batch_prefill.py index 74db4e084c..b687cb9761 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_batch_prefill.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_batch_prefill.py @@ -1,7 +1,6 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. # generate kernel instances to speed up compilation - import copy from dataclasses import dataclass, field import fnmatch @@ -20,14 +19,60 @@ FWD_DTYPE_MAP, BOOL_MAP, PIPELINE_ENUM_MAP, + QSCALE_CHECK_MAP, + QSCALE_MAP, ) +from codegen.arch import ArchTrait from codegen.utils import update_file +# Architecture trait for kernels requiring global_load_lds (CDNA3+). +# Only used for GLOBAL_LOAD_LDS variants; all other kernels are arch-agnostic. +CDNA3_PLUS_ARCH = ArchTrait( + "cdna3_plus", + preprocessor_check="defined(__gfx94__) || defined(__gfx950__)", +) + +# Architecture trait for tiles that are only valid / only wanted on gfx942 +# (e.g. MI308). Used to gate a tile's device-side compilation so it is skipped +# on other archs +GFX942_ARCH = ArchTrait( + "gfx942", + preprocessor_check="defined(__gfx942__)", +) + +DTYPE_BITS = { + "fp32": 32, + "fp16": 16, + "bf16": 16, + "fp8": 8, + "fp8bf16": 8, + "fp8fp32": 8, + "bf8": 8, +} -DTYPE_BITS = {"fp32": 32, "fp16": 16, "bf16": 16, "fp8": 8, "bf8": 8} +# Element size in bytes per dtype, used by the auto-generated dispatcher to +# decide kv_load_mode per-arm (total KV cache bytes vs INT32_MAX). +DTYPE_BYTES = {k: v // 8 for k, v in DTYPE_BITS.items()} K0_MAX_SUBMAX_MAP = {32: 32, 64: 64, 96: 128, 128: 128, 256: 256} +SUPPORTED_PAGE_SIZE = [1, 16, 1024] +SUPPORTED_KV_MEMORY_LAYOUT = ["vectorized", "linear"] +SUPPORTED_KV_LOOKUP_TABLE = ["vllm", "sglang"] +KV_MEMORY_LAYOUT_ENUM_MAP = { + "vectorized": "ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT", + "linear": "ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::LINEAR_LAYOUT", +} +KV_LOOKUP_TABLE_ENUM_MAP = { + "vllm": "ck_tile::BlockAttentionKVCacheLookupTableEnum::VLLM_BLOCK_TABLE_2D", + "sglang": "ck_tile::BlockAttentionKVCacheLookupTableEnum::SGLANG_PAGE_TABLE_1D", +} +KV_LOAD_MODE_ENUM_MAP = { + False: "ck_tile::BlockAttentionKVCacheLoadModeEnum::BUFFER_LOAD", + True: "ck_tile::BlockAttentionKVCacheLoadModeEnum::GLOBAL_LOAD_LDS", +} + + FMHA_BATCH_PREFILL_PIPELINE_MAP = { "qr_async": "ck_tile::BlockFmhaBatchPrefillPipelineQRKSVSAsync", } @@ -40,6 +85,8 @@ """ FMHA_FWD_KERNEL_BODY = """ +#if !defined(__HIP_DEVICE_COMPILE__) || ({F_arch_check}) + using fmha_dtype_{F_idx} = {F_dtype}; using fmha_block_tile_{F_idx} = ck_tile::sequence<{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}>; @@ -51,7 +98,7 @@ ck_tile::sequence<{F_wm1}, {F_wn1}, {F_wk1}>, {F_vlayout}>; -using fmha_trait_{F_idx} = ck_tile::TileFmhaTraits<{F_spad}, +using fmha_trait_{F_idx} = ck_tile::TileFmhaBatchPrefillTraits<{F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, @@ -60,14 +107,20 @@ false, {F_lse}, {F_dropout}, - {F_squant}, - {F_occupancy}>; + {F_qscale}, + {F_occupancy}, + false, + {F_sink}, + {F_page_size}, + {F_kv_memory_layout}, + {F_kv_lookup_table}, + {F_kv_load_mode}>; using fmha_variant_{F_idx} = ck_tile::ComposedAttention<{F_logits} * ck_tile::LOGITS_SOFT_CAP, CK_TILE_FMHA_FWD_FAST_EXP2>; using fmha_mask_{F_idx} = {F_mask}; -using fmha_pipeline_problem_{F_idx} = ck_tile::BlockFmhaPipelineProblem< +using fmha_pipeline_problem_{F_idx} = ck_tile::BlockFmhaBatchPrefillPipelineProblem< typename FmhaFwdTypeConfig::QDataType, typename FmhaFwdTypeConfig::KDataType, typename FmhaFwdTypeConfig::VDataType, @@ -84,6 +137,7 @@ fmha_variant_{F_idx}, fmha_mask_{F_idx}, false, + {F_page_size}, fmha_trait_{F_idx}>; using fmha_pipeline_{F_idx} = {F_pipeline}< @@ -97,8 +151,8 @@ using fmha_kernel_{F_idx} = ck_tile::FmhaBatchPrefillWithPagedKVCacheKernel; -using trait_{F_idx} = fmha_fwd_traits_<{F_hdim}, {F_dtype}, {F_mode},{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, - {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_dropout}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, false>; +using trait_{F_idx} = fmha_fwd_batch_prefill_traits_<{F_hdim}, {F_dtype}, {F_mode},{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, + {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_dropout}, {F_qscale}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, false, false, {F_sink}, {F_page_size}, {F_kv_memory_layout}, {F_kv_lookup_table}, {F_kv_load_mode}>; #include @@ -107,17 +161,21 @@ {{ using k_ = fmha_kernel_{F_idx}; if(s.log_level_ > 0) - std::cout << ", " << k_::GetName() << std::flush; + std::cout << ", {F_kname}" << std::flush; auto [kargs, grids] = fmha_batch_prefill_create_kargs_and_grids(a); const dim3 blocks = k_::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)); }} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || ({F_arch_check}) """ FMHA_FWD_API_FILENAME = "fmha_batch_prefill_api.cpp" FMHA_FWD_API = """ +#include #include +#include namespace {{ bool get_num_cus(unsigned& num_cu) {{ @@ -157,6 +215,11 @@ return r; }} + // Host-side arch identity, used to mirror device-only arch gates (e.g. tiles + // restricted via F_arch) in the host dispatch so we never select an + // arm whose device image was elided for this arch. + [[maybe_unused]] const std::string device_name = ck_tile::get_device_name(); + [[maybe_unused]] auto get_num_blocks = [&](unsigned kM0) {{ return get_num_thread_blocks(a.batch, a.nhead_q, a.max_seqlen_q, kM0); }}; @@ -167,6 +230,7 @@ """ FMHA_FWD_API_PER_DTYPE = """ {F_if}(t.data_type.compare(\"{F_dtype}\") == 0){{ + constexpr int kElementBytes = {F_element_bytes}; {F_hdim_case} }} """ @@ -175,9 +239,9 @@ }} """ -FMHA_FWD_API_INNER_DISPATCH = """ {F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_lse == {F_lse}) && (t.has_dropout == {F_dropout}) && (t.do_fp8_static_quant == {F_squant}) && - ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck}) && ({F_constraint})) {{ - using trait_ = fmha_fwd_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, {F_lse}, {F_dropout}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, false>; +FMHA_FWD_API_INNER_DISPATCH = """ {F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_lse == {F_lse}) && (t.has_dropout == {F_dropout}) && (t.qscale_type == {F_qscale_check}) && (t.has_sink == {F_sink}) && + ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck}) && ({F_constraint}) && ({F_arch_host}) && (t.kv_memory_layout == {F_kv_memory_layout}) && (t.kv_lookup_table == {F_kv_lookup_table}) && (t.page_size == {F_page_size}) && (fmha_batch_prefill_select_kv_load_mode(a.page_block_size, {F_bn0}, a.num_total_pages, a.batch_stride_k, a.batch_stride_v, kElementBytes, a.k_ptr, a.v_ptr) == {F_kv_load_mode})) {{ + using trait_ = fmha_fwd_batch_prefill_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, {F_lse}, {F_dropout}, {F_qscale}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, false, false, {F_sink}, {F_page_size}, {F_kv_memory_layout}, {F_kv_lookup_table}, {F_kv_load_mode}>; return fmha_batch_prefill_(s, a); }} """ @@ -216,18 +280,27 @@ class FmhaFwdApiTrait: bias: str # lse: str # dropout: str - squant: str # + qscale: str # spad: str skpad: str dpad: str dvpad: str + sink: str # t/f constraint: CppConstraint + kv_memory_layout: str + kv_lookup_table: str + page_size: int = 1 # page block size + use_global_load: bool = False # use global_load_lds_* for >2GB KV cache + # Host-side arch predicate mirroring a tile's device-only F_arch gate. + # "true" means arch-agnostic (no host restriction). + arch_host: str = "true" @property def name(self) -> str: return ( f"{self.hdim}-{self.dtype}-{self.mode}-{self.bm0}-{self.bn0}-{self.bk0}-{self.bn0}-{self.bk1}-{self.bk0max}-" - + f"{self.vlayout}-{self.logits}-{self.mask}-{self.bias}-{self.lse}-{self.dropout}-{self.squant}-{self.spad}-{self.skpad}-{self.dpad}-{self.dvpad}" + + f"{self.vlayout}-{self.logits}-{self.mask}-{self.bias}-{self.lse}-{self.dropout}-{self.qscale}-{self.spad}-{self.skpad}-{self.dpad}-{self.dvpad}-{self.kv_memory_layout}-{self.kv_lookup_table}-ps{self.page_size}" + + ("-gload" if self.use_global_load else "-bload") ) @property @@ -312,8 +385,11 @@ class FmhaFwdPipeline: F_bias: str # true/false F_lse: str # F_dropout: str # - F_squant: str # + F_qscale: str # no/pertensor F_mask: str # value from MASK_MAP + F_sink: str # t/f (StreamLLM sink tokens) + F_kv_memory_layout: str # + F_kv_lookup_table: str # F_constraint: CppConstraint = field(default_factory=lambda: CppConstraint()) @property @@ -370,10 +446,17 @@ def pad_name() -> str: else: n += "_ndropout" - if self.F_squant == "t": - n += "_squant" + if self.F_qscale != "no": + n += f"_{self.F_qscale}" else: - n += "_nsquant" + n += "_nqscale" + + if self.F_sink == "t": + n += "_sink" + else: + n += "_nsink" + + n += "_" + self.F_kv_memory_layout + "_" + self.F_kv_lookup_table return n @@ -413,12 +496,14 @@ def api(self) -> str: F_bias=BIAS_MAP[trait.bias], F_lse=BOOL_MAP[trait.lse], F_dropout=BOOL_MAP[trait.dropout], - F_squant=BOOL_MAP[trait.squant], + F_qscale_check=QSCALE_CHECK_MAP[trait.qscale], + F_qscale=QSCALE_MAP[trait.qscale], F_scheck=trait.scheck, F_skcheck=trait.skcheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck, F_constraint=trait.constraint, + F_arch_host=trait.arch_host, F_spad=BOOL_MAP[trait.spad], F_skpad=BOOL_MAP[trait.skpad], F_dpad=BOOL_MAP[trait.dpad], @@ -431,6 +516,15 @@ def api(self) -> str: F_bk0max=trait.bk0max, F_hdim=hdim, F_dtype=FWD_DTYPE_MAP[dtype], + F_kv_memory_layout=KV_MEMORY_LAYOUT_ENUM_MAP[ + trait.kv_memory_layout + ], + F_kv_lookup_table=KV_LOOKUP_TABLE_ENUM_MAP[ + trait.kv_lookup_table + ], + F_page_size=trait.page_size, + F_sink=BOOL_MAP[trait.sink], + F_kv_load_mode=KV_LOAD_MODE_ENUM_MAP[trait.use_global_load], ) if_j = "if" if j == 0 else "else if" per_hdim_case = per_hdim_case + FMHA_FWD_API_PER_HDIM_CASE.format( @@ -438,7 +532,10 @@ def api(self) -> str: ) if_i = "if" if i == 0 else "else if" per_dtypes = per_dtypes + FMHA_FWD_API_PER_DTYPE.format( - F_if=if_i, F_dtype=dtype, F_hdim_case=per_hdim_case + F_if=if_i, + F_dtype=dtype, + F_element_bytes=DTYPE_BYTES[dtype], + F_hdim_case=per_hdim_case, ) if not per_dtypes: # empty string we add some ignore to suppress warning in api @@ -468,6 +565,11 @@ class FmhaFwdTileSize: F_wk1: int # gemm1 warp size along k F_occupancy: int # occupancy, -1 will let pipeline decide the occupancy, other value will overwrite occupancy F_constraint: CppConstraint = field(default_factory=lambda: CppConstraint()) + # Optional compile-time arch guard for this tile's device code. When set, + # the kernel body is only instantiated during the device pass for matching + # archs (host pass always keeps it so the host symbol still exists). This is + # independent from F_constraint, which is a runtime dispatch predicate. + F_arch: Optional[ArchTrait] = None @property def name(self) -> str: @@ -488,10 +590,13 @@ class FmhaFwdKernel: F_tile: FmhaFwdTileSize F_pipeline: FmhaFwdPipeline mask_impl: str + F_page_size: int = 1 # page block size + F_use_global_load: bool = False # use global_load_lds_* for >2GB KV cache @property def template(self) -> str: return FMHA_FWD_KERNEL_HEADER + FMHA_FWD_KERNEL_BODY.format( + F_kname=self.name, F_idx=self.F_idx, F_hdim=self.F_hdim, F_dtype=FWD_DTYPE_MAP[self.F_dtype], @@ -522,19 +627,41 @@ def template(self) -> str: F_bias=BIAS_MAP[self.F_pipeline.F_bias], F_lse=BOOL_MAP[self.F_pipeline.F_lse], F_dropout=BOOL_MAP[self.F_pipeline.F_dropout], - F_squant=BOOL_MAP[self.F_pipeline.F_squant], + F_qscale=QSCALE_MAP[self.F_pipeline.F_qscale], F_occupancy=self.F_tile.F_occupancy, + F_kv_memory_layout=KV_MEMORY_LAYOUT_ENUM_MAP[ + self.F_pipeline.F_kv_memory_layout + ], + F_kv_lookup_table=KV_LOOKUP_TABLE_ENUM_MAP[ + self.F_pipeline.F_kv_lookup_table + ], F_pipeline_enum=PIPELINE_ENUM_MAP[self.F_pipeline.tag], F_mask=get_mask_map(self.mask_impl)[self.F_pipeline.F_mask], F_mode=MODE_MAP[self.F_mode], F_pipeline=FMHA_BATCH_PREFILL_PIPELINE_MAP[self.F_pipeline.tag], + F_page_size=self.F_page_size, + F_sink=BOOL_MAP[self.F_pipeline.F_sink], + F_kv_load_mode=KV_LOAD_MODE_ENUM_MAP[self.F_use_global_load], + F_arch_check=self._arch_check(), ) + def _arch_check(self) -> str: + # Combine any arch guards that apply to this kernel. + checks = [] + if self.F_use_global_load: + checks.append(CDNA3_PLUS_ARCH.preprocessor_check) + if self.F_tile.F_arch is not None: + checks.append(self.F_tile.F_arch.preprocessor_check) + if not checks: + return "true" + return " && ".join(f"({c})" for c in checks) + @property def name(self) -> str: # TODO: we don't encode idx here return ( - f"fmha_batch_prefill_d{self.F_hdim}_{self.F_dtype}_{self.F_mode}_" + f"fmha_batch_prefill_d{self.F_hdim}_{self.F_dtype}_{self.F_mode}_ps{self.F_page_size}_" + + ("gload_" if self.F_use_global_load else "bload_") + self.F_tile.name + "_" + self.F_pipeline.name @@ -562,21 +689,43 @@ def api_trait(self) -> FmhaFwdApiTrait: bias=self.F_pipeline.F_bias, lse=self.F_pipeline.F_lse, dropout=self.F_pipeline.F_dropout, - squant=self.F_pipeline.F_squant, + qscale=self.F_pipeline.F_qscale, spad=self.F_pipeline.F_spad, skpad=self.F_pipeline.F_skpad, dpad=self.F_pipeline.F_dpad, dvpad=self.F_pipeline.F_dvpad, + sink=self.F_pipeline.F_sink, constraint=self.F_tile.F_constraint & self.F_pipeline.F_constraint, + kv_memory_layout=self.F_pipeline.F_kv_memory_layout, + kv_lookup_table=self.F_pipeline.F_kv_lookup_table, + page_size=self.F_page_size, + use_global_load=self.F_use_global_load, + arch_host=( + self.F_tile.F_arch.device_name_check + if self.F_tile.F_arch is not None + else "true" + ), ) class KernelComponentFactory: @staticmethod def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: - if dtype == "fp16" or dtype == "bf16": + if dtype in ["fp16", "bf16"]: return { 128 : [FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], + 256 : [ + FmhaFwdTileSize(128, 32, 16, 256, 16, 256, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, 2, CppConstraint("num_cus < 128"), GFX942_ARCH), + FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1), + ], + } # fmt: skip + elif dtype in ["fp8bf16"]: + return { + 128 : [FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], + 256 : [ + FmhaFwdTileSize(128, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, 2, CppConstraint("num_cus < 128"), GFX942_ARCH), + FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1), + ], } # fmt: skip else: return None @@ -587,20 +736,59 @@ def get_pipelines(dtype, hdim, receipt, mask_impl) -> List[FmhaFwdPipeline]: # TODO: the order of List matters! the later in this list will be also be checked later # TODO: currently for qr pipeline, let 't' padding to appear later!! # TODO: how to design this more generic? - squant = "t" if dtype == "fp8" else "f" pipelines = [] if dtype in ["fp16", "bf16"]: - for logits, mask, bias, lse, dropout in itertools.product( + qscale = "no" + for ( + logits, + mask, + bias, + lse, + dropout, + sink, + kv_memory_layout, + kv_lookup_table, + ) in itertools.product( ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], + ["t", "f"], + SUPPORTED_KV_MEMORY_LAYOUT, + SUPPORTED_KV_LOOKUP_TABLE, ): - pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "f", "t", "t", logits, bias, lse, dropout, squant, mask)) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "t", "t", "t", logits, bias, lse, dropout, squant, mask)) # fmt: skip - # pipelines.append(FmhaFwdPipeline("qr_async", "col", "t", "f", "t", "t", logits, bias, lse, dropout, squant, mask)) # fmt: skip - # pipelines.append(FmhaFwdPipeline("qr_async", "col", "t", "t", "t", "t", logits, bias, lse, dropout, squant, mask)) # fmt: skip + # sink tokens are only meaningful when masking is enabled; + # skip the sink="t" + nomask combinations to avoid emitting + # kernels that can never be dispatched. + if sink == "t" and mask in ("no", "s_no"): + continue + pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, sink, kv_memory_layout, kv_lookup_table)) # fmt: skip + elif dtype in ["fp8bf16"]: + # no need lse/dropout kernels (sink is supported via kHasSink) + for ( + logits, + qscale, + mask, + bias, + sink, + kv_memory_layout, + kv_lookup_table, + ) in itertools.product( + ["t", "f"], + ["pertensor", "kv_blockscale"], + get_mask_map(mask_impl).keys(), + ["no"], + ["t", "f"], + SUPPORTED_KV_MEMORY_LAYOUT, + SUPPORTED_KV_LOOKUP_TABLE, + ): + # sink tokens are only meaningful when masking is enabled; + # skip the sink="t" + nomask combinations to avoid emitting + # kernels that can never be dispatched. + if sink == "t" and mask in ("no", "s_no"): + continue + pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "t", "t", "t", logits, bias, "f", "f", qscale, mask, sink, kv_memory_layout, kv_lookup_table)) # fmt: skip else: assert False return pipelines @@ -610,27 +798,44 @@ class CustomFactory(KernelComponentFactory): @staticmethod def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: result = KernelComponentFactory.get_hdim_tile_size_dict(dtype) - if dtype == "fp16" or dtype == "bf16": + if dtype in ["fp16", "bf16"]: if 128 in result.keys(): result[128].insert(0, FmhaFwdTileSize( 64, 128, 64, 128, 64, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1, CppConstraint("get_num_blocks(128) < num_cus * min_cu_util_rate"))) # fmt: skip return result def get_fwd_blobs( - kernel_filter: Optional[str], receipt, optdim_list, mask_impl + kernel_filter: Optional[str], + receipt, + optdim_list, + mask_impl, + targets: Optional[List[str]] = None, ) -> Tuple[FmhaFwdApiPool, List[FmhaFwdKernel]]: + # batch_prefill pipeline uses gfx9-specific async scatter-gather buffer addressing + # (amd_buffer_addressing.hpp raw buffer loads) that is not compatible with + # non-gfx9 architectures (gfx11/gfx12/gfx10 are wave32 and use different + # buffer instruction formats). Skip all batch_prefill kernels for non-gfx9 targets. + has_non_gfx9 = targets is not None and any( + not t.startswith("gfx9") for t in targets + ) # TODO: we don't support tuning yet, so pick up one value for vlayout/pipeline/pad # support this in future gen = list() api_pool = FmhaFwdApiPool(mask_impl) + if has_non_gfx9: + return api_pool, gen + for dtype in FWD_DTYPE_MAP.keys(): d = CustomFactory.get_hdim_tile_size_dict(dtype) if d is None: continue # for hdim_str, mode, mask, bias, lse in itertools.product(d.keys(), MODE_MAP.keys(), MASK_MAP.keys(), ["t", "f"], ["t", "f"]): for (hdim, tiles), mode in itertools.product(d.items(), MODE_MAP.keys()): + # batch_prefill pipeline requires group mode (static_assert in pipeline problem) + if mode != "group": + continue for tile, pipeline in itertools.product( tiles, CustomFactory.get_pipelines(dtype, hdim, receipt, mask_impl) ): @@ -652,70 +857,100 @@ def get_fwd_blobs( or pipeline.F_logits == "f" ): continue - k = FmhaFwdKernel( - F_idx=0, - F_hdim=hdim, - F_dtype=dtype, - F_mode=mode, - F_tile=tile, - F_pipeline=pipeline, - mask_impl=mask_impl, - ) - if kernel_filter != "": - if not fnmatch.fnmatch(k.name, kernel_filter): - continue - if optdim_list != [-1]: - if hdim not in optdim_list: - continue - # 2 - Flash attention integration - if receipt in (2, 3): - cond = dtype in ["fp16", "bf16"] - cond &= pipeline.F_vlayout == "row" - cond &= pipeline.F_bias in ["no", "alibi"] - cond &= pipeline.F_squant == "f" - if not cond: - continue - # PyTorch integration - elif receipt == 4: - cond = dtype in ["fp16", "bf16"] - cond &= pipeline.F_vlayout == "row" - cond &= pipeline.F_bias in ["no", "bias"] - cond &= pipeline.F_squant == "f" - if not cond: - continue - # Aiter(mha_fwd) integration - elif receipt == 100: - cond = dtype in ["fp16", "bf16"] - cond &= mode == "batch" - cond &= pipeline.F_vlayout == "row" - cond &= pipeline.F_squant == "f" - if not cond: - continue - # Aiter(mha_batch_prefill) integration - elif receipt == 200: - cond = dtype in ["fp16", "bf16"] - cond &= mode == "group" - cond &= pipeline.F_vlayout == "row" - cond &= pipeline.F_squant == "f" - if not cond: - continue - # aiter::mha_batch_prefill C++ api integration - elif receipt == 600: - cond = dtype in ["fp16", "bf16"] - cond &= mode == "group" - cond &= pipeline.F_vlayout == "row" - cond &= pipeline.F_squant == "f" - if not cond: - continue - # fp32 only - if receipt == 800 or receipt == 801: - cond = dtype == "fp32" - if not cond: + # Generate kernels for both page_size=16 and page_size=1024 + for page_size in SUPPORTED_PAGE_SIZE: + if page_size == 1 and pipeline.F_kv_memory_layout != "linear": continue - - api_pool.register_traits(k.api_trait()) - gen.append(k) + # kv_blockscale requires page_size >= kN0 (tile.F_bn0) + # This ensures all tokens in a main loop iteration belong to the same page + if pipeline.F_qscale == "kv_blockscale" and page_size < tile.F_bn0: + continue + k = FmhaFwdKernel( + F_idx=0, + F_hdim=hdim, + F_dtype=dtype, + F_mode=mode, + F_tile=tile, + F_pipeline=pipeline, + mask_impl=mask_impl, + F_page_size=page_size, + ) + if kernel_filter != "": + if not fnmatch.fnmatch(k.name, kernel_filter): + continue + if optdim_list != [-1]: + if hdim not in optdim_list: + continue + # 2 - Flash attention integration + if receipt in (2, 3): + cond = dtype in ["fp16", "bf16"] + cond &= pipeline.F_vlayout == "row" + cond &= pipeline.F_bias in ["no", "alibi"] + cond &= pipeline.F_qscale == "no" + if not cond: + continue + # PyTorch integration + elif receipt == 4: + cond = dtype in ["fp16", "bf16"] + cond &= pipeline.F_vlayout == "row" + cond &= pipeline.F_bias in ["no", "bias"] + cond &= pipeline.F_qscale == "no" + if not cond: + continue + # Aiter(mha_fwd) integration + elif receipt == 100: + cond = dtype in ["fp16", "bf16"] + cond &= mode == "batch" + cond &= pipeline.F_vlayout == "row" + cond &= pipeline.F_qscale == "no" + if not cond: + continue + # Aiter(mha_batch_prefill) integration + elif receipt == 200: + cond = dtype in ["fp16", "bf16", "fp8bf16"] + cond &= mode == "group" + cond &= pipeline.F_vlayout == "row" + if not cond: + continue + # aiter::mha_batch_prefill C++ api integration + elif receipt == 600: + cond = dtype in ["fp16", "bf16", "fp8bf16"] + cond &= mode == "group" + cond &= pipeline.F_vlayout == "row" + cond &= pipeline.F_qscale == "no" + if not cond: + continue + elif receipt == 700: + continue # TE does not use this API + + # fp32 only + if receipt == 800 or receipt == 801: + cond = dtype == "fp32" + if not cond: + continue + + api_pool.register_traits(k.api_trait()) + gen.append(k) + + # For page_size < kN0 (tile.F_bn0), also generate a GLOBAL_LOAD_LDS + # variant for >2GB KV cache support. The default (BUFFER_LOAD) uses SRD + # buffer_load (fast, <2GB). GLOBAL_LOAD_LDS uses global_load_lds_* + # (slower, handles >2GB). + if page_size < tile.F_bn0: + k_global_load = FmhaFwdKernel( + F_idx=0, + F_hdim=hdim, + F_dtype=dtype, + F_mode=mode, + F_tile=tile, + F_pipeline=pipeline, + mask_impl=mask_impl, + F_page_size=page_size, + F_use_global_load=True, + ) + api_pool.register_traits(k_global_load.api_trait()) + gen.append(k_global_load) return (api_pool, gen) @@ -736,7 +971,9 @@ def write_blobs( optdim_list, mask_impl, ) -> None: - api_pool, kernels = get_fwd_blobs(kernel_filter, receipt, optdim_list, mask_impl) + api_pool, kernels = get_fwd_blobs( + kernel_filter, receipt, optdim_list, mask_impl, targets + ) for kernel in kernels: write_single_fwd_kernel(kernel, output_dir) write_fwd_api(api_pool, output_dir) @@ -751,7 +988,9 @@ def list_blobs( mask_impl, ) -> None: with file_path.open("a") as f: - _, kernels = get_fwd_blobs(kernel_filter, receipt, optdim_list, mask_impl) + _, kernels = get_fwd_blobs( + kernel_filter, receipt, optdim_list, mask_impl, targets + ) for kernel in kernels: - f.write(str(file_path.parent / GEN_DIR / kernel.filename) + "\n") - f.write(str(file_path.parent / GEN_DIR / FMHA_FWD_API_FILENAME) + "\n") + f.write((file_path.parent / GEN_DIR / kernel.filename).as_posix() + "\n") + f.write((file_path.parent / GEN_DIR / FMHA_FWD_API_FILENAME).as_posix() + "\n") diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_bwd.py b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_bwd.py index 79fb35d7be..b7f91deaeb 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_bwd.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_bwd.py @@ -1,7 +1,6 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. # generate kernel instances to speed up compilation - import copy import fnmatch import itertools @@ -29,7 +28,15 @@ FMHA_BWD_KERNEL_HEADER = """// SPDX-License-Identifier: MIT // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.\n // auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && \\ + (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || \\ + defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || \\ + defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__)) +#undef CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT +#define CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT 1 +#endif #include "fmha_bwd.hpp" + """ FMHA_BWD_DQ_DK_DV_KERNEL_BODY = """ @@ -146,6 +153,16 @@ auto [kargs, grids] = fmha_bwd_dq_dk_dv_create_kargs_and_grids(a); const dim3 blocks = k_::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + if constexpr(k_::kNeedsKernelPrezeroDqAcc) + {{ + using pk_ = typename k_::DqAccPrezeroKernel; + return ck_tile::launch_kernel( + s, + ck_tile::make_kernel( + pk_{{}}, pk_::GridSize(), pk_::BlockSize(), 0, + k_::MakeDqAccPrezeroKargs(a.workspace_ptr, a.batch)), + ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)); + }} return ck_tile::launch_kernel( s, ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)); }} @@ -157,6 +174,14 @@ auto [kargs, grids] = fmha_bwd_dq_dk_dv_create_kargs_and_grids(a); const dim3 blocks = k_::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + if constexpr(k_::kNeedsKernelPrezeroDqAcc) + {{ + using pk_ = typename k_::DqAccPrezeroKernel; + ck_tile::make_kernel( + pk_{{}}, pk_::GridSize(), pk_::BlockSize(), 0, + k_::MakeDqAccPrezeroKargs(a.workspace_ptr, a.batch))( + ck_tile::stream_config{{s.stream_id_}}); + }} ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)( ck_tile::stream_config{{s.stream_id_}}); }} @@ -168,6 +193,41 @@ return k_::kMaxSeqLenQ; }} +template <> +size_t fmha_bwd_dq_dk_dv_dq_ws_host_size_(int batch_size) +{{ + using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx}; + return k_::GetWorkspaceHostSize(batch_size); +}} + +template <> +size_t fmha_bwd_dq_dk_dv_dq_ws_device_upper_bound_( + ck_tile::index_t max_batch, ck_tile::index_t hdim_q, ck_tile::index_t nhead_q, + ck_tile::index_t total_seqlen_q_padded, ck_tile::index_t max_seqlen_k) +{{ + using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx}; + return k_::GetWorkspaceDeviceSizeUpperBound( + max_batch, hdim_q, nhead_q, total_seqlen_q_padded, max_seqlen_k); +}} + +template <> +size_t fmha_bwd_dq_dk_dv_dq_prepare_ws_host_( + void* cpu_ws, ck_tile::index_t batch_size, ck_tile::index_t hdim_q, + ck_tile::index_t nhead_q, ck_tile::index_t seqlen_q, ck_tile::index_t seqlen_k, + const ck_tile::index_t* seqstart_qs, const ck_tile::index_t* seqstart_ks) +{{ + using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx}; + return k_::PrepareWorkspaceHost( + cpu_ws, batch_size, hdim_q, nhead_q, seqlen_q, seqlen_k, seqstart_qs, seqstart_ks); +}} + +template <> +bool fmha_bwd_dq_dk_dv_needs_zero_dq_acc_() +{{ + using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx}; + return k_::NeedsZeroDqAcc(); +}} + template <> std::string fmha_bwd_dq_dk_dv_get_name_() {{ @@ -180,34 +240,15 @@ FMHA_BWD_API_FILENAME = "fmha_bwd_api.cpp" FMHA_BWD_API = """ -#include - -template -float fmha_bwd_(const ck_tile::stream_config& s, fmha_bwd_args a) -{{ - if constexpr (!std::is_same_v) - {{ - if(s.log_level_ > 0) - std::cout << ", " << fmha_bwd_dot_do_o_get_name_() << "@" << fmha_bwd_convert_dq_get_name_() << "@" << fmha_bwd_dq_dk_dv_get_name_() << std::flush; - return ck_tile::launch_kernel(s, - [=](const ck_tile::stream_config& s_){{ fmha_bwd_dot_do_o_oneshot_(s_, a); }}, - [=](const ck_tile::stream_config& s_){{ fmha_bwd_dq_dk_dv_oneshot_(s_, a); }}, - [=](const ck_tile::stream_config& s_){{ fmha_bwd_convert_dq_oneshot_(s_, a); }} - ); - }} - else - {{ - if(s.log_level_ > 0) - std::cout << ", " << fmha_bwd_dot_do_o_get_name_() << "@" << fmha_bwd_dq_dk_dv_get_name_() << std::flush; - return ck_tile::launch_kernel(s, - [=](const ck_tile::stream_config& s_){{ fmha_bwd_dot_do_o_oneshot_(s_, a); }}, - [=](const ck_tile::stream_config& s_){{ fmha_bwd_dq_dk_dv_oneshot_(s_, a); }} - ); - }} +fmha_bwd_launcher::fmha_bwd_launcher(const fmha_bwd_traits& t){{ + [[maybe_unused]] const std::string device_name = ck_tile::get_device_name(); +{F_launcher} }} + +// Prefer to use launcher. Leave fmha_bwd here for backward compatibility. template <> -float fmha_bwd<2>(fmha_bwd_traits t, fmha_bwd_args a, const ck_tile::stream_config& s){{ +float fmha_bwd<2>(const fmha_bwd_traits& t, fmha_bwd_args a, const ck_tile::stream_config& s){{ [[maybe_unused]] const std::string device_name = ck_tile::get_device_name(); float r = -1; {F_dispatch} @@ -226,15 +267,22 @@ def FMHA_BWD_API_COND_STATEMENT(F_cond: str, F_body: str, *, if_i=0) -> str: return "\n".join(lines) + "\n" -FMHA_BWD_API_INNER_DISPATCH = """{F_if}((t.is_group_mode == {F_mode}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_dbias == {F_dbias}) && ({F_dropout_check}) && +FMHA_BWD_API_INNER_DISPATCH_COMMON = """{F_if}((t.is_group_mode == {F_mode}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_dbias == {F_dbias}) && ({F_dropout_check}) && ({F_scheck}) && ({F_dcheck}) && ({F_dvcheck}) && (t.is_deterministic == {F_deterministic}){F_max_seq_q_cond}{F_cond_extra}) {{ using dot_do_o_trait_ = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1d}, ({F_dvpad} > 0)>; using dq_dk_dv_trait_ = fmha_bwd_dq_dk_dv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_mask}, {F_dropout}, {F_bias}, {F_dbias}, {F_dpad}, {F_dvpad}, {F_deterministic}, {F_trload}, {F_maxq}, {F_bn0}>; - using convert_dq_trait_ = fmha_bwd_convert_dq_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1d}, ({F_dpad} > 0), {F_deterministic}, {F_convert_dq_bn0}>; + using convert_dq_trait_ = fmha_bwd_convert_dq_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1d}, ({F_dpad} > 0), {F_deterministic}>; +""" +FMHA_BWD_API_INNER_DISPATCH_RUN = """ r = fmha_bwd_, {F_arch.tag}>(s, a); return r; }} """ +FMHA_BWD_API_INNER_DISPATCH_LAUNCHER = """ + this->init, {F_arch.tag}>(t); + return; +}} +""" # M0 size for 1d kernels (dot/convert) M0_1D = 64 @@ -447,7 +495,7 @@ def get_dq_dk_dv_tiles(dtype: str, tr_load: str) -> List[FmhaBwdDQDKDVTileSize]: results = KernelComponentFactoryGfx9.get_dq_dk_dv_tiles(dtype, tr_load) if dtype in ["fp16", "bf16"] and tr_load == "t": results.extend([ - FmhaBwdDQDKDVTileSize( 32, 128, 64, 32, 64, 32, 32, 64, 64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 32, 1), + FmhaBwdDQDKDVTileSize( 32, 256, 64, 32, 64, 32, 32, 64, 64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 32, 1), FmhaBwdDQDKDVTileSize( 32, 128, 128, 32, 128, 32, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 32, 1), FmhaBwdDQDKDVTileSize( 16, 192, 128, 16, 128, 16, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1), # FmhaBwdDQDKDVTileSize( 32, 32, 64, 32, 64, 32, 32, 64, 64, 1, 1, 1, 1, 1, 1, 1, 1, 1, 16, 16, 32, 16, 16, 32, 1, 32), @@ -458,6 +506,24 @@ def get_dq_dk_dv_tiles(dtype: str, tr_load: str) -> List[FmhaBwdDQDKDVTileSize]: return results +class KernelComponentFactoryGfx11(KernelComponentFactoryBase): + arch = ArchTrait("gfx11") + + @staticmethod + def get_dq_dk_dv_tiles(dtype: str, tr_load: str) -> List[FmhaBwdDQDKDVTileSize]: + if tr_load == "t": + return [] + if dtype in ["fp16", "bf16"]: + return [ + # bm0, bn0, bk0, bk1, bk2, bk3, bk4, bhdq, bhdv, + FmhaBwdDQDKDVTileSize( 32, 64, 32, 32, 32, 32, 64, 32, 32, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaBwdDQDKDVTileSize( 32, 64, 64, 32, 64, 32, 32, 64, 64, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaBwdDQDKDVTileSize( 16, 64, 128, 16, 128, 16, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaBwdDQDKDVTileSize( 16, 64, 256, 16, 256, 16, 32, 256, 256, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + ] # fmt: skip + return [] + + class KernelComponentFactoryGfx12(KernelComponentFactoryBase): arch = ArchTrait("gfx12") @@ -476,6 +542,25 @@ def get_dq_dk_dv_tiles(dtype: str, tr_load: str) -> List[FmhaBwdDQDKDVTileSize]: return [] +class KernelComponentFactoryGfx125(KernelComponentFactoryBase): + arch = ArchTrait("gfx125") + + @staticmethod + def get_dq_dk_dv_tiles(dtype: str, tr_load: str) -> List[FmhaBwdDQDKDVTileSize]: + if tr_load == "t": + return [] + if dtype in ["fp16", "bf16"]: + return [ + # bm0, bn0, bk0, bk1, bk2, bk3, bk4, bhdq, bhdv, + FmhaBwdDQDKDVTileSize( 32, 64, 32, 32, 32, 32, 64, 32, 32, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 32, 16, 16, 32, -1), + FmhaBwdDQDKDVTileSize( 32, 64, 64, 32, 64, 32, 32, 64, 64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 32, -1), + #FmhaBwdDQDKDVTileSize( 32, 64, 64, 32, 64, 32, 64, 64, 64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 32, -1), + FmhaBwdDQDKDVTileSize( 32, 64, 128, 32, 128, 32, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 32, -1), + FmhaBwdDQDKDVTileSize( 32, 64, 256, 32, 256, 32, 32, 256, 256, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 32, -1), + ] # fmt: skip + return [] + + def get_factory(target: str): # Place more specific architectures first @@ -483,7 +568,10 @@ def get_factory(target: str): return KernelComponentFactoryGfx950 if target.startswith("gfx9"): return KernelComponentFactoryGfx9 - + if target.startswith("gfx11"): + return KernelComponentFactoryGfx11 + if target.startswith("gfx125"): + return KernelComponentFactoryGfx125 if target.startswith("gfx12"): return KernelComponentFactoryGfx12 @@ -504,6 +592,7 @@ def get_factory(target: str): typename FmhaBwdTypeConfig::ODataType, typename FmhaBwdTypeConfig::OGradDataType, typename FmhaBwdTypeConfig::DDataType, + typename FmhaBwdTypeConfig::LSEDataType, /* BlockSize = M0 = */ {F_bm0}, {F_hdim}, {F_mode}, @@ -620,7 +709,6 @@ def filename(self) -> str: typename FmhaBwdTypeConfig::QGradDataType, /* BlockSize = */ 256, {F_bm0}, - {F_bn0}, {F_hdim}, {F_mode}, {F_deterministic}, @@ -637,8 +725,7 @@ def filename(self) -> str: {F_mode}, {F_spad}, {F_dpad}, - {F_deterministic}, - {F_bn0}>; + {F_deterministic}>; template <> float fmha_bwd_convert_dq_(const ck_tile::stream_config& s, fmha_bwd_args a) @@ -682,7 +769,6 @@ class FmhaBwdConvertQGradKernel: F_hdim: int # hdim F_dtype: str # data type F_bm0: int # tile size along q seqlen (block size) - F_bn0: int # tile size along k seqlen F_spad: str # true/false F_dpad: str # F_mode: str # value from MODE_MAP @@ -698,7 +784,6 @@ def template(self) -> str: F_hdim=self.F_hdim, F_dtype=BWD_DTYPE_MAP[self.F_dtype], F_bm0=self.F_bm0, - F_bn0=self.F_bn0, F_spad=BOOL_MAP[self.F_spad], F_dpad=BOOL_MAP[self.F_dpad], F_mode=MODE_MAP[self.F_mode], @@ -719,7 +804,7 @@ def pad_name() -> str: return n pn = pad_name() - n = f"fmha_bwd_convert_dq_d{self.F_hdim}_{self.F_dtype}_b{self.F_bm0}x{self.F_bn0}_{self.F_mode}_o{self.F_occupancy}" + n = f"fmha_bwd_convert_dq_d{self.F_hdim}_{self.F_dtype}_b{self.F_bm0}_{self.F_mode}_o{self.F_occupancy}" if pn != "": n += f"_{pn}" else: @@ -776,42 +861,49 @@ def scheck(self) -> str: if self.mode == "group": return "true /*spad1d is always true in group mode*/" elif self.spad1d == "t": - return f"true /*a.seqlen_q % {M0_1D} != 0*/" + return f"true /*t.seqlen_q % {M0_1D} != 0*/" else: # self.spad1d == "f" - return f"a.seqlen_q % {M0_1D} == 0" + return f"t.seqlen_q % {M0_1D} == 0" @property def dcheck(self) -> str: if self.dpad == 0: - return f"a.hdim_q % {self.bhdq} == 0" + return f"t.hdim_q % {self.bhdq} == 0" else: - return f"a.hdim_q % {self.dpad} == 0" + return f"t.hdim_q % {self.dpad} == 0" @property def dvcheck(self) -> str: if self.dvpad == 0: - return f"a.hdim_v % {self.bhdv} == 0" + return f"t.hdim_v % {self.bhdv} == 0" else: - return f"a.hdim_v % {self.dvpad} == 0" + return f"t.hdim_v % {self.dvpad} == 0" @property def max_seq_q_cond(self) -> str: if self.tile.max_seq_q != 0: - return f" && (a.seqlen_q <= {self.tile.max_seq_q})" + if self.mode == "group": + return f" && (t.max_seqlen_q <= {self.tile.max_seq_q})" + else: + return f" && (t.seqlen_q <= {self.tile.max_seq_q})" else: return "" @property def extra_cond(self) -> str: - if self.tr_load == "t" and self.tile.max_seq_q == 0 and self.tile.F_bn0 == 128: - return " && (a.seqlen_k <= 256)" + if ( + self.tr_load == "t" + and self.tile.max_seq_q == 0 + and self.tile.F_bn0 == 128 + and self.tile.F_bhdq == 128 + ): + if self.mode == "group": + return " && (t.max_seqlen_k <= 256)" + else: + return " && (t.seqlen_k <= 256)" else: return "" - @property - def convert_dq_bn0(self) -> int: - return self.tile.F_bn0 if self.deterministic == "t" else 0 - @property def dot_do_o_kernel(self) -> FmhaBwdOGradDotOKernel: # TODO: we don't support tuning yet, so pick up one value for pad/occupancy @@ -866,7 +958,6 @@ def get_occupancy(dtype, hdim): F_hdim=self.hdim, F_dtype=self.dtype, F_bm0=M0_1D, - F_bn0=self.convert_dq_bn0, F_spad=self.spad1d, F_dpad=F_dpad, F_mode=self.mode, @@ -891,12 +982,12 @@ def register_dq_dk_dv_traits(self, trait: FmhaBwdApiTrait) -> None: check_duplicates_and_paddings(ts, trait) ts.append(copy.copy(trait)) - def _api_inners(self, traits: List[FmhaBwdApiTrait]) -> str: + def _api_inners(self, traits: List[FmhaBwdApiTrait]) -> tuple[str, str]: inners = "" + inners_launcher = "" for i_trait, trait in enumerate(traits): - inners += FMHA_BWD_API_INNER_DISPATCH.format( + inners_common = FMHA_BWD_API_INNER_DISPATCH_COMMON.format( F_if=if_(i_trait), - F_arch=trait.arch, F_mode=MODE_MAP[trait.mode], F_mask_check=get_mask_check_map(self.mask_impl)[trait.mask], F_mask=get_mask_map(self.mask_impl)[trait.mask], @@ -916,13 +1007,22 @@ def _api_inners(self, traits: List[FmhaBwdApiTrait]) -> str: F_deterministic=BOOL_MAP[trait.deterministic], F_trload=BOOL_MAP[trait.tr_load], F_maxq=trait.tile.max_seq_q, - F_convert_dq_enabled=BOOL_MAP[not trait.convert_dq_kernel.disabled], F_max_seq_q_cond=trait.max_seq_q_cond, F_cond_extra=trait.extra_cond, F_bn0=trait.tile.F_bn0, - F_convert_dq_bn0=trait.convert_dq_bn0, ) - return inners + inners += inners_common + FMHA_BWD_API_INNER_DISPATCH_RUN.format( + F_arch=trait.arch, + F_convert_dq_enabled=BOOL_MAP[not trait.convert_dq_kernel.disabled], + ) + inners_launcher += ( + inners_common + + FMHA_BWD_API_INNER_DISPATCH_LAUNCHER.format( + F_arch=trait.arch, + F_convert_dq_enabled=BOOL_MAP[not trait.convert_dq_kernel.disabled], + ) + ) + return inners, inners_launcher @staticmethod def max_seq_q_sort_key(trait): @@ -938,8 +1038,7 @@ def dtype_cond(dtype: str) -> str: def hdim_cond(hdim: int) -> str: return f"t.hdim_q <= {hdim} && t.hdim_v <= {hdim}" - @property - def api(self) -> str: + def _api_per_arch(self, variant) -> str: per_arch = "" for i_arch, (arch, pool_by_arch) in enumerate(self.dq_dk_dv_pool.items()): per_dtypes = "" @@ -949,7 +1048,7 @@ def api(self) -> str: traits = sorted(pool_by_hdim, key=self.max_seq_q_sort_key) inners = self._api_inners(traits) per_hdim_case += FMHA_BWD_API_COND_STATEMENT( - if_i=i_hdim, F_cond=self.hdim_cond(hdim), F_body=inners + if_i=i_hdim, F_cond=self.hdim_cond(hdim), F_body=inners[variant] ) per_dtypes += FMHA_BWD_API_COND_STATEMENT( if_i=i_dtype, F_cond=self.dtype_cond(dtype), F_body=per_hdim_case @@ -959,9 +1058,14 @@ def api(self) -> str: ) if not per_arch: # empty string we add some ignore to suppress warning in api - per_arch = "(void)t; (void)s; (void)a;" + per_arch = ("(void)t; (void)s; (void)a;", "(void)t;")[variant] + return per_arch + + @property + def api(self) -> str: result = FMHA_BWD_KERNEL_HEADER + FMHA_BWD_API.format( - F_dispatch=indent(per_arch) + F_dispatch=indent(self._api_per_arch(0)), + F_launcher=indent(self._api_per_arch(1)), ) return result.replace("\n\n", "\n") @@ -1027,7 +1131,7 @@ def get_bwd_blobs( hdim = tile.F_bhdq if (mode == "group") and (spad1d == "f"): continue - if (mode == "group" or ("no" not in mask)) and tile.max_seq_q != 0: + if ("no" not in mask) and tile.max_seq_q != 0: continue if (bias == "no" or bias == "alibi") and dbias == "t": continue @@ -1115,6 +1219,12 @@ def get_bwd_blobs( cond = dtype in ["fp16", "bf16"] if not cond: continue + # TransformerEngine integration + elif receipt == 700: + cond = dtype in ["fp16", "bf16"] + cond &= dropout in ["no", "dropout_wg32", "dropout_wg16"] + if not cond: + continue # fp32 only, all variations if receipt == 800: @@ -1186,9 +1296,9 @@ def list_blobs( ) with file_path.open("a") as f: for k in kernels_dot_do_o: - f.write(str(file_path.parent / GEN_DIR / k.filename) + "\n") + f.write((file_path.parent / GEN_DIR / k.filename).as_posix() + "\n") for k in kernels_dq_dk_dv: - f.write(str(file_path.parent / GEN_DIR / k.filename) + "\n") + f.write((file_path.parent / GEN_DIR / k.filename).as_posix() + "\n") for k in kernels_convert_dq: - f.write(str(file_path.parent / GEN_DIR / k.filename) + "\n") - f.write(str(file_path.parent / GEN_DIR / FMHA_BWD_API_FILENAME) + "\n") + f.write((file_path.parent / GEN_DIR / k.filename).as_posix() + "\n") + f.write((file_path.parent / GEN_DIR / FMHA_BWD_API_FILENAME).as_posix() + "\n") diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd.py b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd.py index 0ffb1d510c..87accb313b 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd.py @@ -1,7 +1,6 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. # generate kernel instances to speed up compilation - import copy import fnmatch import itertools @@ -9,14 +8,13 @@ from collections import OrderedDict from dataclasses import dataclass, field from pathlib import Path -from typing import List, Optional, Tuple +from typing import Callable, ClassVar, Iterable, List, Optional, Tuple from codegen.arch import ArchTrait, get_factories_for_targets from codegen.cmake_config import GEN_DIR from codegen.cpp_symbol_map import ( LAYOUT_MAP, BIAS_CHECK_MAP, - get_mask_check_map, BOOL_MAP, PIPELINE_MAP, PIPELINE_ENUM_MAP, @@ -24,13 +22,36 @@ FWD_DTYPE_MAP, BIAS_MAP, get_mask_map, + get_mask_cpp_type, + get_mask_cpp_check_expr, + QSCALE_CHECK_MAP, + QSCALE_MAP, ) from codegen.utils import check_duplicates_and_paddings, if_, indent, update_file -DTYPE_BITS = {"fp32": 32, "fp16": 16, "bf16": 16, "fp8": 8, "bf8": 8} - -K0_MAX_SUBMAX_MAP = {32: 32, 48: 48, 64: 64, 96: 128, 128: 128, 192: 192, 256: 256} +DTYPE_BITS = { + "fp32": 32, + "fp16": 16, + "bf16": 16, + "fp8": 8, + "fp8bf16": 8, + "fp8fp32": 8, + "bf8": 8, + "mxfp8": 8, + "mxfp4": 4, +} + +K0_MAX_SUBMAX_MAP = { + 32: 32, + 48: 48, + 64: 64, + 80: 96, + 96: 128, + 128: 128, + 192: 192, + 256: 256, +} FMHA_FWD_KERNEL_HEADER = """// SPDX-License-Identifier: MIT // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.\n @@ -39,79 +60,98 @@ #include "fmha_fwd.hpp" """ -FMHA_FWD_KERNEL_BODY = """ +FMHA_FWD_KERNEL_HEADER_QR_HPAD = """// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.\n +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && \ + (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || \ + defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || \ + defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || \ + defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" +""" + +FMHA_FWD_KERNEL_BODY_TEMPLATE = """ #include #if !defined(__HIP_DEVICE_COMPILE__) || ({F_arch.preprocessor_check}) -using fmha_dtype_{F_idx} = {F_dtype}; - -using fmha_block_tile_{F_idx} = ck_tile::sequence<{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}>; - -using fmha_shape_{F_idx} = ck_tile::TileFmhaShape, - ck_tile::sequence<{F_wm0}, {F_wn0}, {F_wk0}>, - ck_tile::sequence<{F_rm1}, {F_rn1}, {F_rk1}>, - ck_tile::sequence<{F_wm1}, {F_wn1}, {F_wk1}>, - {F_vlayout}>; - -using fmha_trait_{F_idx} = ck_tile::TileFmhaTraits<{F_spad}, - {F_skpad}, - {F_dpad}, - {F_dvpad}, - {F_logits}, - {F_bias}, - false, - {F_lse}, - {F_dropout}, - {F_squant}, - {F_occupancy}, - {F_skip}>; - -using fmha_variant_{F_idx} = ck_tile::ComposedAttention<{F_logits} * ck_tile::LOGITS_SOFT_CAP, CK_TILE_FMHA_FWD_FAST_EXP2>; - -using fmha_mask_{F_idx} = {F_mask}; - -using fmha_pipeline_problem_{F_idx} = ck_tile::BlockFmhaPipelineProblem< - typename FmhaFwdTypeConfig::QDataType, - typename FmhaFwdTypeConfig::KDataType, - typename FmhaFwdTypeConfig::VDataType, - typename FmhaFwdTypeConfig::SaccDataType, - typename FmhaFwdTypeConfig::SMPLComputeDataType, - typename FmhaFwdTypeConfig::BiasDataType, - typename FmhaFwdTypeConfig::RandValOutputDataType, - typename FmhaFwdTypeConfig::LSEDataType, - typename FmhaFwdTypeConfig::PDataType, - typename FmhaFwdTypeConfig::OaccDataType, - typename FmhaFwdTypeConfig::ODataType, - fmha_shape_{F_idx}, +using fmha_dtype = {F_dtype}; + +using fmha_block_tile = ck_tile::sequence<{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<{F_wm0}, {F_wn0}, {F_wk0}>, + ck_tile::sequence<{F_rm1}, {F_rn1}, {F_rk1}>, + ck_tile::sequence<{F_wm1}, {F_wn1}, {F_wk1}>, + {F_vlayout}>; + +using fmha_traits = ck_tile::TileFmhaTraits<{F_spad}, + {F_skpad}, + {F_dpad}, + {F_dvpad}, + {F_logits}, + {F_bias}, + false, + {F_lse}, + {F_dropout}, + {F_qscale}, + {F_occupancy}, + {F_skip}, + {F_sink}>; + +using fmha_variant = ck_tile::ComposedAttention<{F_logits} * ck_tile::LOGITS_SOFT_CAP, CK_TILE_FMHA_FWD_FAST_EXP2>; + +using fmha_mask = {F_mask}; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, {F_mode}, - fmha_variant_{F_idx}, - fmha_mask_{F_idx}, + fmha_variant, + fmha_mask, {F_trload}, - fmha_trait_{F_idx}>; + fmha_traits>; -using fmha_pipeline_{F_idx} = {F_pipeline}< - fmha_pipeline_problem_{F_idx}>; +using fmha_pipeline = {F_pipeline}< + fmha_pipeline_problem>; -using fmha_epilogue_{F_idx} = - ck_tile::Default2DEpilogue::OaccDataType, - typename FmhaFwdTypeConfig<{F_dtype}>::ODataType, - {F_spad}, {F_dvpad}>>; +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + {F_spad}, {F_dvpad}>>; -using fmha_kernel_{F_idx} = - ck_tile::FmhaFwdKernel; +using fmha_kernel = {F_kernel}; -using trait_{F_idx} = fmha_fwd_traits_<{F_hdim}, {F_dtype}, {F_mode},{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, - {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_dropout}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_trload}, {F_skip}>; + +using trait = fmha_fwd_traits_<{F_hdim}, {F_dtype}, {F_mode},{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, + {F_pipeline_enum}, {F_logits}, fmha_mask, {F_bias}, {F_lse}, {F_dropout}, {F_qscale}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_trload}, {F_skip}, {F_sink}>; template<> -float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) {{ - using k_ = fmha_kernel_{F_idx}; + using k_ = fmha_kernel; + if(a.selected_kernel_name != nullptr) + *a.selected_kernel_name = "{F_kname}"; if(s.log_level_ > 0) - std::cout << ", " << k_::GetName() << std::flush; - auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + std::cout << ", {F_kname}" << std::flush; + auto [kargs, grids] = {F_kargs_creator}(a); const dim3 blocks = k_::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{{}}, grids, blocks, 0, kargs)); @@ -121,40 +161,49 @@ """ FMHA_FWD_API_FILENAME = "fmha_fwd_api.cpp" -FMHA_FWD_API = """ +FMHA_FWD_API_HEADER = """ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.\n +// auto generated by generate.py #include +#include +#include #include -namespace {{ -bool get_num_cus(unsigned& num_cus) {{ +#include "fmha_fwd.hpp" + +namespace { +bool get_num_cus(unsigned& num_cus) { int device; auto status = hipGetDevice(&device); - if(status != hipSuccess) {{ + if(status != hipSuccess) { fprintf(stderr, "failed to get device"); return false; - }} + } - hipDeviceProp_t props{{}}; + hipDeviceProp_t props{}; status = hipGetDeviceProperties(&props, device); - if(status != hipSuccess) {{ + if(status != hipSuccess) { fprintf(stderr, "failed to get device properties"); return false; - }} + } num_cus = props.multiProcessorCount; return true; -}} +} -unsigned get_num_thread_blocks(unsigned batch, unsigned nheads, unsigned max_seqlen_q, unsigned kM0) {{ +unsigned get_num_thread_blocks(unsigned batch, unsigned nheads, unsigned max_seqlen_q, unsigned kM0) { const unsigned num_m_blocks = (max_seqlen_q + kM0 - 1) / kM0; const unsigned num_n_blocks = 1; // we assume that num_n_blocks is always 1 return batch * nheads * num_m_blocks * num_n_blocks; -}} -}} // namespace - -float fmha_fwd(fmha_fwd_traits t, fmha_fwd_args a, const ck_tile::stream_config& s) {{ +} +} // namespace +""" +FMHA_FWD_API_FUNC_TEMPLATE = """ +namespace {{ +float {F_func_name}([[maybe_unused]] fmha_fwd_traits t, [[maybe_unused]] fmha_fwd_args a, [[maybe_unused]] const ck_tile::stream_config& s) {{ float r = -1; [[maybe_unused]] const float min_cu_util_rate = 0.8; // minimum CU utilization rate @@ -173,6 +222,16 @@ {F_dispatch} return r; }} +}} // namespace +""" +FMHA_FWD_API_FOOTER = """ +float fmha_fwd(fmha_fwd_traits traits, fmha_fwd_args args, const ck_tile::stream_config& config) { + if (args.max_seqlen_q > 1) { + float r = fmha_fwd_v3(traits, args, config); + if (r >= 0) return r; + } + return fmha_fwd_v2(traits, args, config); +} """ FMHA_FWD_API_PER_ARCH = """{F_if}({F_arch.device_name_check}) {{ @@ -190,9 +249,9 @@ }} """ -FMHA_FWD_API_INNER_DISPATCH = """{F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_lse == {F_lse}) && (t.has_dropout == {F_dropout}) && (t.do_fp8_static_quant == {F_squant}) && (t.skip_min_seqlen_q == {F_skip}) && +FMHA_FWD_API_INNER_DISPATCH = """{F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_lse == {F_lse}) && (t.has_dropout == {F_dropout}) && (t.qscale_type == {F_qscale_check}) && (t.skip_min_seqlen_q == {F_skip}) &&(t.has_sink == {F_sink}) && ({F_scheck}) && ({F_seqtune}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck}) && ({F_constraint})) {{ - using trait_ = fmha_fwd_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, {F_lse}, {F_dropout}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_trload}, {F_skip}>; + using trait_ = fmha_fwd_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, {F_lse}, {F_dropout}, {F_qscale}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_trload}, {F_skip}, {F_sink}>; return fmha_fwd_(s, a); }} """ @@ -232,32 +291,38 @@ class FmhaFwdApiTrait: bias: str # lse: str # dropout: str - squant: str # + qscale: str # spad: str skpad: str dpad: str dvpad: str skip: str tr_load: str + sink: str constraint: CppConstraint @property def name(self) -> str: return ( f"{self.hdim}-{self.dtype}-{self.mode}-{self.bm0}-{self.bn0}-{self.bk0}-{self.bn0}-{self.bk1}-{self.bk0max}-" - + f"{self.vlayout}-{self.logits}-{self.mask}-{self.bias}-{self.lse}-{self.dropout}-{self.squant}-{self.spad}-{self.skpad}-{self.dpad}-{self.dvpad}-{self.skip}" + + f"{self.vlayout}-{self.logits}-{self.mask}-{self.bias}-{self.lse}-{self.dropout}-{self.qscale}-{self.spad}-{self.skpad}-{self.dpad}-{self.dvpad}-{self.skip}-{self.sink}" ) @property def scheck(self) -> str: if self.mode == "group": return "true/*group mode spad always true*/" # group mode only generate spad/skpad == true - if self.pipeline_tag in ["qr_async", "qr_async_trload"]: + if self.pipeline_tag in [ + "qr_async", + "qr_async_trload", + "qr_async_trload_v3", + "qr_tdm", + ]: if self.spad == "t": return "true" # always support else: return "true" - elif self.pipeline_tag in ["qr", "qs"]: + elif self.pipeline_tag in ["qr", "qr_hpad", "qs"]: if self.spad == "t": return f"true /*a.seqlen_q % {self.bm0} != 0*/" # TODO: order of get_pipelines() matters! (ugly) else: @@ -266,9 +331,11 @@ def scheck(self) -> str: assert False def seqtune(self, max_bm0: int) -> str: - if self.bm0 == max_bm0: + if self.bm0 == max_bm0 or self.bm0 == 64: return "true/*fall back to largest tile*/" else: + if self.mode == "group": + return f"a.max_seqlen_q <= {self.bm0}" return f"a.seqlen_q <= {self.bm0}" @property @@ -280,12 +347,12 @@ def skcheck(self) -> str: return f"(a.cu_seqlen_k_ptr != nullptr) || (a.seqlen_k == 0 || a.seqlen_k % {self.bn0} != 0)" else: return f"(a.cu_seqlen_k_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0)" - elif self.pipeline_tag in ["qr", "qs"]: + elif self.pipeline_tag in ["qr", "qr_hpad", "qs"]: if self.skpad == "t": return f"true /*a.seqlen_k % {self.bn0} != 0*/" # TODO: order of get_pipelines() matters! (ugly) else: return f"(a.cu_seqlen_k_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % {self.bn0} == 0)" - elif self.pipeline_tag == "qr_async_trload": + elif self.pipeline_tag in ["qr_async_trload", "qr_async_trload_v3", "qr_tdm"]: if self.skpad == "t": return "true" else: @@ -295,13 +362,19 @@ def skcheck(self) -> str: @property def dcheck(self) -> str: - if self.pipeline_tag == "qr_async": - vec = int((32 * 4) / DTYPE_BITS[self.dtype]) + if self.pipeline_tag == "qr_hpad": if self.dpad == "t": - return f"a.hdim_q % {vec} == 0" + return "a.hdim_q % 8 == 0" else: assert False - elif self.pipeline_tag in ["qr", "qs", "qr_async_trload"]: + elif self.pipeline_tag in [ + "qr", + "qs", + "qr_async", + "qr_async_trload", + "qr_async_trload_v3", + "qr_tdm", + ]: bk0submax = K0_MAX_SUBMAX_MAP[self.bk0max] if self.dpad == "t": return f"true /*a.hdim_q % {bk0submax} != 0*/" # TODO: order of get_pipelines() matters! (ugly) @@ -312,13 +385,19 @@ def dcheck(self) -> str: @property def dvcheck(self) -> str: - if self.pipeline_tag == "qr_async": - vec = int((32 * 4) / DTYPE_BITS[self.dtype]) + if self.pipeline_tag == "qr_hpad": if self.dvpad == "t": - return f"a.hdim_v % {vec} == 0" + return "a.hdim_v % 8 == 0" else: assert False - elif self.pipeline_tag in ["qr", "qs", "qr_async_trload"]: + elif self.pipeline_tag in [ + "qr", + "qs", + "qr_async", + "qr_async_trload", + "qr_async_trload_v3", + "qr_tdm", + ]: bk0submax = K0_MAX_SUBMAX_MAP[self.bk0max] if self.dvpad == "t": return f"true /*a.hdim_v % {bk0submax} != 0*/" # TODO: order of get_pipelines() matters! (ugly) @@ -341,10 +420,11 @@ class FmhaFwdPipeline: F_bias: str # true/false F_lse: str # F_dropout: str # - F_squant: str # + F_qscale: str # no/pertensor F_mask: str # value from MASK_MAP F_skip: str # true/false F_trload: str # true/false + F_sink: str # true/false F_constraint: CppConstraint = field(default_factory=lambda: CppConstraint()) @property @@ -406,23 +486,26 @@ def pad_name() -> str: else: n += "_nskip" - if self.F_squant == "t": - n += "_squant" + if self.F_qscale != "no": + n += f"_{self.F_qscale}" else: - n += "_nsquant" + n += "_nqscale" if self.F_trload == "t": n += "_trload" else: n += "_ntrload" + if self.F_sink == "t": + n += "_sink" + else: + n += "_nsink" return n class FmhaFwdApiPool: - def __init__(self, mask_impl): + def __init__(self): self.pool = OrderedDict() - self.mask_impl = mask_impl def register_traits(self, trait: FmhaFwdApiTrait) -> None: hdim = trait.hdim, trait.bn1 @@ -434,19 +517,60 @@ def register_traits(self, trait: FmhaFwdApiTrait) -> None: check_duplicates_and_paddings(ts, trait) ts.append(copy.copy(trait)) - @property - def api(self) -> str: + def get_num_traits( + self, filter_fn: Optional[Callable[[FmhaFwdApiTrait], bool]] = None + ) -> int: + if filter_fn is None: + + def accept_all(trait: FmhaFwdApiTrait) -> bool: + return True + + filter_fn = accept_all + + return sum( + sum(1 for trait in pool_by_hdim if filter_fn(trait)) + for pool_by_arch in self.pool.values() + for pool_by_dtype in pool_by_arch.values() + for pool_by_hdim in pool_by_dtype.values() + ) + + def render( + self, func_name, filter_fn: Optional[Callable[[FmhaFwdApiTrait], bool]] = None + ) -> str: + if filter_fn is None: + + def accept_all(trait: FmhaFwdApiTrait) -> bool: + return True + + filter_fn = accept_all + + def has_traits(node) -> bool: + """Recursively traverse nested OrderedDicts and lists to determine if any FmhaFwdApiTrait satisfies filter_fn().""" + if isinstance(node, list): + return any(filter_fn(elem) for elem in node) + elif isinstance(node, OrderedDict): + return any(has_traits(val) for val in node.values()) + return False + per_arch = str() - for i_arch, (arch, pool_by_arch) in enumerate(self.pool.items()): + for i_arch, (arch, pool_by_arch) in enumerate( + item for item in self.pool.items() if has_traits(item[1]) + ): per_dtypes = str() - for i_dtype, (dtype, pool_by_dtype) in enumerate(pool_by_arch.items()): + for i_dtype, (dtype, pool_by_dtype) in enumerate( + item for item in pool_by_arch.items() if has_traits(item[1]) + ): per_hdim_case = str() for i_hdim, ((hdim, hdim_v), pool_by_hdim) in enumerate( - pool_by_dtype.items() + item for item in pool_by_dtype.items() if has_traits(item[1]) ): - max_bm0 = max((t.bm0 for t in pool_by_hdim), default=0) + max_bm0 = max( + (t.bm0 for t in pool_by_hdim if filter_fn(t)), default=0 + ) inners = str() - for i_trait, trait in enumerate(pool_by_hdim): + for i_trait, trait in enumerate( + [trait for trait in pool_by_hdim if filter_fn(trait)] + ): inners += FMHA_FWD_API_INNER_DISPATCH.format( F_if=if_(i_trait), F_arch=arch, @@ -454,15 +578,17 @@ def api(self) -> str: F_vlayout=LAYOUT_MAP[trait.vlayout], F_pipeline_enum=PIPELINE_ENUM_MAP[trait.pipeline_tag], F_logits=BOOL_MAP[trait.logits], - F_mask=get_mask_map(self.mask_impl)[trait.mask], - F_mask_check=get_mask_check_map(self.mask_impl)[trait.mask], + F_mask=get_mask_cpp_type(trait.mask), + F_mask_check=get_mask_cpp_check_expr(trait.mask), F_bias_check=BIAS_CHECK_MAP[trait.bias], F_bias=BIAS_MAP[trait.bias], F_lse=BOOL_MAP[trait.lse], F_dropout=BOOL_MAP[trait.dropout], F_skip=BOOL_MAP[trait.skip], F_trload=BOOL_MAP[trait.tr_load], - F_squant=BOOL_MAP[trait.squant], + F_qscale_check=QSCALE_CHECK_MAP[trait.qscale], + F_qscale=QSCALE_MAP[trait.qscale], + F_sink=BOOL_MAP[trait.sink], F_scheck=trait.scheck, F_seqtune=trait.seqtune(max_bm0), F_skcheck=trait.skcheck, @@ -496,10 +622,9 @@ def api(self) -> str: F_arch=arch, F_dtype_case=indent(per_dtypes), ) - if not per_arch: - # empty string we add some ignore to suppress warning in api - per_arch = "(void)t; (void)s; (void)a;" - return FMHA_FWD_KERNEL_HEADER + FMHA_FWD_API.format(F_dispatch=indent(per_arch)) + return FMHA_FWD_API_FUNC_TEMPLATE.format( + F_func_name=func_name, F_dispatch=indent(per_arch) + ) @dataclass @@ -538,18 +663,41 @@ def name(self) -> str: @dataclass class FmhaFwdKernel: F_arch: ArchTrait - F_idx: int # this is not a tunable, but a counter to differentiate symbol F_hdim: int # hdim F_dtype: str # data type F_mode: str # value from MODE_MAP F_tile: FmhaFwdTileSize F_pipeline: FmhaFwdPipeline - mask_impl: str - @property - def template(self) -> str: - return FMHA_FWD_KERNEL_HEADER + FMHA_FWD_KERNEL_BODY.format( - F_idx=self.F_idx, + _KERNEL_HEADER: ClassVar[str] = FMHA_FWD_KERNEL_HEADER + _KERNEL_HEADER_QR_HPAD: ClassVar[str] = FMHA_FWD_KERNEL_HEADER_QR_HPAD + _KERNEL_BODY_TEMPLATE: ClassVar[str] = FMHA_FWD_KERNEL_BODY_TEMPLATE + + @classmethod + def _get_cpp_kernel_class_name(cls, pipeline_tag): + if pipeline_tag == "qr_async_trload_v3": + return "ck_tile::FmhaFwdV3Kernel" + else: + return "ck_tile::FmhaFwdKernel" + + @classmethod + def _get_kernel_header(cls, pipeline_tag): + if pipeline_tag == "qr_hpad": + return cls._KERNEL_HEADER_QR_HPAD + return cls._KERNEL_HEADER + + @classmethod + def _get_cpp_kargs_creator_func_name(cls, pipeline_tag): + if pipeline_tag == "qr_async_trload_v3": + return "fmha_fwd_v3_create_kargs_and_grids" + else: + return "fmha_fwd_create_kargs_and_grids" + + def render(self) -> str: + return type(self)._get_kernel_header(self.F_pipeline.tag) + type( + self + )._KERNEL_BODY_TEMPLATE.format( + F_kname=self.name, F_arch=self.F_arch, F_hdim=self.F_hdim, F_dtype=FWD_DTYPE_MAP[self.F_dtype], @@ -580,14 +728,17 @@ def template(self) -> str: F_bias=BIAS_MAP[self.F_pipeline.F_bias], F_lse=BOOL_MAP[self.F_pipeline.F_lse], F_dropout=BOOL_MAP[self.F_pipeline.F_dropout], - F_squant=BOOL_MAP[self.F_pipeline.F_squant], + F_qscale=QSCALE_MAP[self.F_pipeline.F_qscale], F_skip=BOOL_MAP[self.F_pipeline.F_skip], F_occupancy=self.F_tile.F_occupancy, F_pipeline_enum=PIPELINE_ENUM_MAP[self.F_pipeline.tag], - F_mask=get_mask_map(self.mask_impl)[self.F_pipeline.F_mask], + F_mask=get_mask_cpp_type(self.F_pipeline.F_mask), F_mode=MODE_MAP[self.F_mode], - F_pipeline=PIPELINE_MAP[self.F_pipeline.tag], F_trload=BOOL_MAP[self.F_pipeline.F_trload], + F_pipeline=PIPELINE_MAP[self.F_pipeline.tag], + F_kernel=self._get_cpp_kernel_class_name(self.F_pipeline.tag), + F_kargs_creator=self._get_cpp_kargs_creator_func_name(self.F_pipeline.tag), + F_sink=BOOL_MAP[self.F_pipeline.F_sink], ) @property @@ -623,27 +774,198 @@ def api_trait(self) -> FmhaFwdApiTrait: bias=self.F_pipeline.F_bias, lse=self.F_pipeline.F_lse, dropout=self.F_pipeline.F_dropout, - squant=self.F_pipeline.F_squant, + qscale=self.F_pipeline.F_qscale, spad=self.F_pipeline.F_spad, skpad=self.F_pipeline.F_skpad, dpad=self.F_pipeline.F_dpad, dvpad=self.F_pipeline.F_dvpad, skip=self.F_pipeline.F_skip, tr_load=self.F_pipeline.F_trload, + sink=self.F_pipeline.F_sink, constraint=self.F_tile.F_constraint & self.F_pipeline.F_constraint, ) -class KernelComponentFactoryGfx9: +@dataclass +class ProblemContext: + dtype: str + mode: str + hdim: int + hdim_v: int + + +@dataclass +class KernelContext: + tile: FmhaFwdTileSize + pipeline: FmhaFwdPipeline + mask_impl: str + + +CompatibilityRule = Callable[[ProblemContext, KernelContext], bool] + + +def is_compatible( + problem_ctx: ProblemContext, + kernel_ctx: KernelContext, + rules: Iterable[CompatibilityRule], +) -> bool: + return all(rule(problem_ctx, kernel_ctx) for rule in rules) + + +def create_kernel( + arch: ArchTrait, problem_ctx: ProblemContext, kernel_ctx: KernelContext +) -> FmhaFwdKernel: + return FmhaFwdKernel( + F_arch=arch, + F_dtype=problem_ctx.dtype, + F_mode=problem_ctx.mode, + F_hdim=problem_ctx.hdim, + F_tile=kernel_ctx.tile, + F_pipeline=kernel_ctx.pipeline, + ) + + +class CompatibilityRuleFactory: + @staticmethod + def get_rules() -> List[CompatibilityRule]: + # in group mode, spad/skpad must be true, since we can't predict if seqlen of current batch need pad or not + def check_mode(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + if problem_ctx.mode == "group": + if ( + kernel_ctx.pipeline.F_spad != "t" + or kernel_ctx.pipeline.F_skpad != "t" + ): + return False + return True + + def check_hdim(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + # NOTE: this is used to speedup deepseek prefill case, we don't gen training + if (problem_ctx.hdim, problem_ctx.hdim_v) == (192, 128): + if ( + kernel_ctx.pipeline.F_bias != "no" + or kernel_ctx.pipeline.F_dropout == "t" + ): + return False + return True + + def check_feature( + problem_ctx: ProblemContext, kernel_ctx: KernelContext + ) -> bool: + # logits_soft_cap is only allowed if no bias + if not ( + ( + kernel_ctx.pipeline.F_logits == "t" + and kernel_ctx.pipeline.F_bias == "no" + ) + or kernel_ctx.pipeline.F_logits == "f" + ): + return False + return True + + return [check_mode, check_hdim, check_feature] + + +class CompatibilityRuleFactoryGfx9(CompatibilityRuleFactory): + _AVAILABLE_PIPELINES = frozenset({"qr", "qr_async", "qs"}) + + @classmethod + def get_rules(cls) -> List[CompatibilityRule]: + rules = CompatibilityRuleFactory.get_rules() + + def check_hdim_tile( + problem_ctx: ProblemContext, kernel_ctx: KernelContext + ) -> bool: + # FIX: too confusing that it has to know about mx types + if problem_ctx.dtype not in ("fp32", "mxfp8", "mxfp4"): + # TODO: update if >=gfx11 archs get qr_async and qr_async_trload support + if kernel_ctx.pipeline.tag in cls._AVAILABLE_PIPELINES and ( + ( + (problem_ctx.hdim, problem_ctx.hdim_v) == (128, 128) + and kernel_ctx.tile.F_bn0 != 128 + ) + or ( + (problem_ctx.hdim, problem_ctx.hdim_v) != (128, 128) + and kernel_ctx.tile.F_bm0 != 128 + ) + or ( + (problem_ctx.hdim, problem_ctx.hdim_v) == (128, 128) + and kernel_ctx.pipeline.tag != "qr_async" + and kernel_ctx.tile.F_bk0 == 64 + ) + ): + # non qr_async_trload only support km0=128 tile size when hdim is not 128 + # non qr_async only support kn0=128 tile size when hdim is 128 + return False + return True + + rules.append(check_hdim_tile) + return rules + + +class CompatibilityRuleFactoryGfx950(CompatibilityRuleFactoryGfx9): + _AVAILABLE_PIPELINES = ( + CompatibilityRuleFactoryGfx9._AVAILABLE_PIPELINES + | frozenset({"qr_async_trload", "qr_async_trload_v3"}) + ) + + @classmethod + def get_rules(cls) -> List[CompatibilityRule]: + rules = CompatibilityRuleFactoryGfx9.get_rules() + + def check_tile_pipeline( + problem_ctx: ProblemContext, kernel_ctx: KernelContext + ) -> bool: + if kernel_ctx.pipeline.tag == "qr_async_trload" and ( + ( + (problem_ctx.hdim, problem_ctx.hdim_v) == (128, 128) + and kernel_ctx.tile.F_bn0 == 128 + ) + or ( + (problem_ctx.hdim, problem_ctx.hdim_v) + not in [(64, 64), (128, 128), (256, 256)] + ) + ): + return False + + # only qr_async_trload_v3 use km0=256 & 8-warps + is_v3_dedicated_tile = ( + kernel_ctx.tile.F_bm0 == 256 + and (kernel_ctx.tile.F_rm0 * kernel_ctx.tile.F_rn0 * kernel_ctx.tile.F_rk0) == 8 + and (kernel_ctx.tile.F_rm1 * kernel_ctx.tile.F_rn1 * kernel_ctx.tile.F_rk1) == 8 + ) # fmt: skip + is_v3_pipeline = kernel_ctx.pipeline.tag == "qr_async_trload_v3" + return is_v3_dedicated_tile == is_v3_pipeline + + rules.extend([check_tile_pipeline]) + return rules + + +class KernelComponentFactoryGfx9(CompatibilityRuleFactoryGfx9): arch = ArchTrait( "gfx9", preprocessor_check="defined(__gfx9__) && !defined(__gfx950__)" ) + _DT_FP32 = ("fp32",) + _DT_FP16_BF16 = ("fp16", "bf16") + _DT_FP8 = ("fp8",) + _DT_FP8BF16 = ("fp8bf16",) + _DT_FP8FP32 = ("fp8fp32",) + + @classmethod + def supported_dtypes(cls) -> Tuple[str]: + return ( + cls._DT_FP32 + + cls._DT_FP16_BF16 + + cls._DT_FP8 + + cls._DT_FP8BF16 + + cls._DT_FP8FP32 + ) + # TODO: design a more practical way to do it # this is current supported tile size per hdim - @staticmethod - def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: - if dtype in ["fp32"]: + @classmethod + def get_hdim_tile_size_dict(cls, dtype: str) -> Optional[dict]: + if dtype in cls._DT_FP32: return { # bm0, bn0, bk0, bn1, bk1, ( 32, 32) : [FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], @@ -656,15 +978,17 @@ def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: (192, 192) : [FmhaFwdTileSize( 64, 64, 32, 192, 32, 192, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], (256, 256) : [FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], } # fmt: skip - elif dtype in ["fp16", "bf16"]: + elif dtype in cls._DT_FP16_BF16: return { ( 32, 32) : [FmhaFwdTileSize(128, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], ( 64, 64) : [FmhaFwdTileSize( 16, 32, 64, 64, 32, 64, 1, 1, 1, 1, 1, 1, 16, 16, 32, 16, 16, 32, -1), FmhaFwdTileSize( 32, 32, 64, 64, 32, 64, 1, 1, 1, 1, 1, 1, 32, 32, 16, 32, 32, 16, -1), FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], + ( 80, 96) : [FmhaFwdTileSize(128, 128, 16, 96, 32, 80, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], ( 96, 128) : [FmhaFwdTileSize(128, 128, 32, 128, 32, 96, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], (128, 128) : [FmhaFwdTileSize( 16, 32, 64, 128, 32, 128, 1, 1, 1, 1, 1, 1, 16, 16, 32, 16, 16, 32, -1), FmhaFwdTileSize( 32, 32, 128, 128, 32, 128, 1, 1, 1, 1, 1, 1, 32, 32, 16, 32, 32, 16, -1), + FmhaFwdTileSize( 64, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 16, -1, CppConstraint('get_num_blocks(64) <= num_cus')), FmhaFwdTileSize(128, 64, 32, 128, 16, 128, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1), FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], # (160, 160) : [FmhaFwdTileSize(128, 128 , 32, 160, 32, 160, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, 1)], @@ -672,172 +996,450 @@ def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: (192, 192) : [FmhaFwdTileSize(128, 128, 32, 192, 32, 192, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, 1)], (256, 256) : [FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], } # fmt: skip - elif dtype in ["fp8", "fp8bf16"]: + elif dtype in cls._DT_FP8 or dtype in cls._DT_FP8BF16: return { ( 64, 64) : [FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 2, 1, 1, 2, 1, 1, 32, 32, 32, 32, 32, 32, -1)], (128, 128) : [FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], + (192, 128) : [FmhaFwdTileSize(128, 128, 32, 128, 32, 192, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], (256, 256) : [FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], } # fmt: skip - elif dtype in ["fp8fp32"]: + elif dtype in cls._DT_FP8FP32: return { (128, 128) : [FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], } # fmt: skip - else: - return None # TODO: we don't support tuning yet, so pick up one value for vlayout/pipeline/pad # support this in future - @staticmethod - def get_pipelines(dtype, hdim, hdim_v, receipt, mask_impl) -> List[FmhaFwdPipeline]: + @classmethod + def get_pipelines( + cls, dtype, hdim, hdim_v, receipt, mask_impl + ) -> List[FmhaFwdPipeline]: # this function will populate a list possible pipelines # TODO: the order of List matters! the later in this list will be also be checked later # TODO: currently for qr pipeline, let "t" padding to appear later!! # TODO: how to design this more generic? pipelines = [] - if dtype in ["fp32"]: - squant = "f" - for logits, mask, bias, lse, dropout, skip in itertools.product( + if dtype in cls._DT_FP32: + qscale = "no" + for logits, mask, bias, lse, dropout, skip, sink in itertools.product( ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"], + ["t", "f"], ): - pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr", "row", "f", "t", "f", "f", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip - elif dtype in ["fp16", "bf16"]: - squant = "f" - for logits, mask, bias, lse, dropout, skip in itertools.product( + pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "f", "t", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + elif dtype in cls._DT_FP16_BF16: + qscale = "no" + for logits, mask, bias, lse, dropout, skip, sink in itertools.product( ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"], + ["t", "f"], ): if hdim == 256 and hdim_v == 256: - pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip # the below two is used for hdim vectorize load - pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "f", "f", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip else: if bias == "bias": # TODO: rocm 6.2 compiler problem if using qr_async for bias case - pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip else: - pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "f", "t", "t", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "t", "t", "t", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip + # skpad=f + dpad=t/dvpad=t variant covers padded head dims (e.g. d=16, d=160) + # when seqlen_k is divisible by bn0. Restored after #6526 changed it to + # "f","f","f","f", which left those cases with no dispatchable qr_async kernel. + pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "f", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip if receipt == 1 and bias != "bias": - pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip # TODO: cover arbitraty hdim# fmt: skip - elif dtype in ["fp8", "fp8bf16", "fp8fp32"]: + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip # TODO: cover arbitraty hdim# fmt: skip + elif dtype in cls._DT_FP8BF16 or dtype in cls._DT_FP8FP32: # no need lse/dropout kernels - for logits, squant, mask, bias in itertools.product( - ["f"], ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys() + for logits, qscale, mask, bias, sink in itertools.product( + ["t", "f"], + ["no", "pertensor", "blockscale"], + get_mask_map(mask_impl).keys(), + ["no"], + ["f", "t"], ): - pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, "f", "f", squant, mask, "f", "f")) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "f", "f", logits, bias, "f", "f", squant, mask, "f", "f")) # fmt: skip - elif dtype in ["fp8fp16", "bf8"]: - # TODO - None - else: - assert False + if hdim == 64: + pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip + else: + # skpad=f + dpad=t/dvpad=t variant covers padded head dims when seqlen_k is + # divisible by bn0. Restored after #6526 changed it to "f","f","f","f". + pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "f", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "t", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip return pipelines -class KernelComponentFactoryGfx950(KernelComponentFactoryGfx9): +class KernelComponentFactoryGfx950( + KernelComponentFactoryGfx9, CompatibilityRuleFactoryGfx950 +): arch = ArchTrait("gfx950") - @staticmethod - def get_pipelines(dtype, hdim, hdim_v, receipt, mask_impl) -> List[FmhaFwdPipeline]: + _DT_MXFP8 = ("mxfp8",) + _DT_MXFP4 = ("mxfp4",) + + @classmethod + def supported_dtypes(cls) -> Tuple[str]: + return ( + KernelComponentFactoryGfx9.supported_dtypes() + + cls._DT_MXFP8 + + cls._DT_MXFP4 + ) + + @classmethod + def get_hdim_tile_size_dict(cls, dtype: str) -> Optional[dict]: + result = KernelComponentFactoryGfx9.get_hdim_tile_size_dict(dtype) + if dtype in cls._DT_FP16_BF16: + # # add tile for qr_async_trload_v3 (bf16/fp16 V3 not ready) + # if (128, 128) in result.keys(): + # result[(128, 128)].append( + # FmhaFwdTileSize(256, 32, 128, 128, 32, 128, 8, 1, 1, 8, 1, 1, 32, 32, 16, 32, 32, 16, -1)) # fmt: skip + if (256, 256) in result.keys(): + result[(256, 256)] = [FmhaFwdTileSize(128, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)] # fmt: skip + elif dtype in cls._DT_MXFP8: + return { + # bm0, bn0, bk0, bn1, bk1, + (128, 128) : [FmhaFwdTileSize(128, 128, 64, 128, 64, 128, 4, 1, 1, 4, 1, 1, 32, 32, 64, 32, 32, 64, -1)], + (256, 256) : [FmhaFwdTileSize(128, 128, 128, 256, 128, 256, 4, 1, 1, 4, 1, 1, 16, 16, 128, 16, 16, 128, -1)], + } # fmt: skip + elif dtype in cls._DT_MXFP4: + return { + # bm0, bn0, bk0, bn1, bk1, + (128, 128) : [FmhaFwdTileSize(128, 128, 64, 128, 64, 128, 4, 1, 1, 4, 1, 1, 32, 32, 64, 32, 32, 64, -1)], + (256, 256) : [FmhaFwdTileSize(128, 128, 128, 256, 128, 256, 4, 1, 1, 4, 1, 1, 16, 16, 128, 16, 16, 128, -1)], + } # fmt: skip + elif dtype in cls._DT_FP8BF16: + if (128, 128) in result.keys(): + result[(128, 128)].append( + FmhaFwdTileSize(256, 64, 128, 128, 64, 128, 8, 1, 1, 8, 1, 1, 32, 32, 32, 32, 32, 32, -1)) # fmt: skip + return result + + @classmethod + def get_pipelines( + cls, dtype, hdim, hdim_v, receipt, mask_impl + ) -> List[FmhaFwdPipeline]: pipelines = KernelComponentFactoryGfx9.get_pipelines( dtype, hdim, hdim_v, receipt, mask_impl ) - if dtype in ["fp16", "bf16"]: - squant = "f" - for logits, mask, bias, lse, dropout, skip in itertools.product( + if dtype in cls._DT_FP16_BF16: + qscale = "no" + for logits, mask, bias, lse, dropout, skip, sink in itertools.product( ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"], + ["t", "f"], ): if ( - (hdim, hdim_v) in [(64, 64), (128, 128)] + (hdim, hdim_v) in [(64, 64), (128, 128), (256, 256)] and logits == "f" and bias == "no" and dropout == "f" and skip == "f" ): - pipelines.append(FmhaFwdPipeline("qr_async_trload", "row", "f", "f", "f", "f", logits, bias, lse, dropout, squant, mask, skip, "t")) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr_async_trload", "row", "f", "f", "t", "t", logits, bias, lse, dropout, squant, mask, skip, "t")) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_async_trload", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "t", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_async_trload", "row", "f", "f", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "t", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_async_trload", "row", "t", "t", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "t", sink)) # fmt: skip # group mode spad + + # # qr_async_trload_v3 bf16/fp16 not ready + # if (hdim, hdim_v) == (128, 128): + # for logits, mask in itertools.product(["t", "f"], ["no", "causal"]): + # pipelines.append(FmhaFwdPipeline("qr_async_trload_v3", "row", "t", "t", "f", "f", + # F_logits=logits, F_bias="no", F_lse="f", F_dropout="f", F_qscale=qscale, F_mask=mask, F_skip="f", F_trload="t", F_sink="f")) # fmt: skip + elif dtype in cls._DT_FP8BF16: + # qr_async_trload_v3 only supports (generic) causal mask + for logits, qscale, mask, lse in itertools.product( + ["t", "f"], + ["no", "pertensor"], + ["no", "causal"], + ["t", "f"], + ): + pipelines.append(FmhaFwdPipeline("qr_async_trload_v3", "row", "t", "t", "f", "f", F_logits=logits, F_bias="no", F_lse=lse, F_dropout="f", F_qscale=qscale, F_mask=mask, F_skip="f", F_trload="t", F_sink="f")) # fmt: skip + + elif dtype in cls._DT_MXFP8 or dtype in cls._DT_MXFP4: + # no need dropout kernels + lse = "t" + dropout = "f" + for logits, qscale, mask, bias, sink in itertools.product( + ["f"], + ["mx"], + get_mask_map(mask_impl).keys(), + ["no"], + ["f", "t"], + ): + pipelines.append(FmhaFwdPipeline("qr", "col", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "col", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip + if hdim > 64 and dtype in cls._DT_MXFP8: + pipelines.append(FmhaFwdPipeline("qr_async", "col", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_async", "col", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip + return pipelines + + +class KernelComponentFactoryGfx11(CompatibilityRuleFactory): + arch = ArchTrait( + "gfx11", + preprocessor_check="defined(__gfx11__) && !defined(__gfx115__)", + ) + + _DT_FP16_BF16 = ("fp16", "bf16") + + @classmethod + def supported_dtypes(cls) -> Tuple[str]: + return cls._DT_FP16_BF16 + + @classmethod + def get_rules(cls) -> List[CompatibilityRule]: + rules = super().get_rules() + + def check_d128_tile_pipeline( + problem_ctx: ProblemContext, kernel_ctx: KernelContext + ) -> bool: + if problem_ctx.dtype not in cls._DT_FP16_BF16: + return True + + if (problem_ctx.hdim, problem_ctx.hdim_v) != (128, 128): + return True + + # For (128, 128) head dims, partial-fragment support in qr_hpad removes the need + # for the previous qr_hpad-specific handling that was added to avoid register spill. + # qr_hpad now reuses the regular 128x64 tile choice. + # The 64x64 tile remains disabled for qr_hpad because it is consistently slower + # in our measurements. + if kernel_ctx.tile.F_bm0 == 64 and kernel_ctx.tile.F_bn0 == 64: + return kernel_ctx.pipeline.tag != "qr_hpad" + + return True + + rules.append(check_d128_tile_pipeline) + return rules + + @classmethod + def get_hdim_tile_size_dict(cls, dtype: str) -> Optional[dict]: + if dtype in cls._DT_FP16_BF16: + return { + # bm0, bn0, bk0, bn1, bk1, + ( 32, 32) : [FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + ( 64, 64) : [FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1, CppConstraint("a.max_seqlen_q < 4096")), + FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 8, 1, 1, 8, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + (128, 128) : [FmhaFwdTileSize(128, 64, 32, 128, 32, 128, 8, 1, 1, 8, 1, 1, 16, 16, 16, 16, 16, 16, 6)], + (192, 128) : [FmhaFwdTileSize( 64, 64, 32, 128, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + (256, 256) : [FmhaFwdTileSize(128, 64, 32, 256, 32, 256, 8, 1, 1, 8, 1, 1, 16, 16, 16, 16, 16, 16, 6)] + } # fmt: skip + else: + raise ValueError(f"unsupported dtype={dtype}") + + @classmethod + def get_pipelines( + cls, dtype, hdim, hdim_v, receipt, mask_impl + ) -> List[FmhaFwdPipeline]: + pipelines = [] + if dtype in cls._DT_FP16_BF16: + qscale = "no" + for logits, mask, bias, lse, dropout, skip, sink in itertools.product( + ["t", "f"], + get_mask_map(mask_impl).keys(), + BIAS_MAP.keys(), + ["t", "f"], + ["t", "f"], + ["t", "f"], + ["t", "f"], + ): + # Keep only ttff/tttt for gfx11: ffff path is often similar or worse + # pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_hpad", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + if receipt == 1: + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip return pipelines -class KernelComponentFactoryGfx12: +class KernelComponentFactoryGfx115(KernelComponentFactoryGfx11): + arch = ArchTrait("gfx115") + + @classmethod + def get_hdim_tile_size_dict(cls, dtype: str) -> Optional[dict]: + result = super().get_hdim_tile_size_dict(dtype) + if dtype in cls._DT_FP16_BF16: + result[(64, 64)] = [FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)] # fmt: skip + result[(256, 256)] = [FmhaFwdTileSize(128, 64, 32, 256, 32, 256, 8, 1, 1, 8, 1, 1, 16, 16, 16, 16, 16, 16, -1)] # fmt: skip + return result + + +class KernelComponentFactoryGfx12(CompatibilityRuleFactory): arch = ArchTrait("gfx12") - @staticmethod - def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: - if dtype in ["fp16", "bf16"]: + _DT_FP16_BF16 = ("fp16", "bf16") + _DT_FP8_FP8BF16 = ("fp8", "fp8bf16") + _DT_FP8FP32 = ("fp8fp32",) + + @classmethod + def supported_dtypes(cls) -> Tuple[str]: + return cls._DT_FP16_BF16 + cls._DT_FP8_FP8BF16 + cls._DT_FP8FP32 + + @classmethod + def get_hdim_tile_size_dict(cls, dtype: str) -> Optional[dict]: + if dtype in cls._DT_FP16_BF16: return { # bm0, bn0, bk0, bn1, bk1, ( 32, 32) : [FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], ( 64, 64) : [FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], - (128, 128) : [FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + (128, 128) : [FmhaFwdTileSize(128, 64, 32, 128, 32, 128, 8, 1, 1, 8, 1, 1, 16, 16, 16, 16, 16, 16, 6)], (192, 128) : [FmhaFwdTileSize( 64, 64, 32, 128, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], (256, 256) : [FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], } # fmt: skip - elif dtype in ["fp8", "fp8bf16"]: + elif dtype in cls._DT_FP8_FP8BF16: return { # bm0, bn0, bk0, bn1, bk1, ( 64, 64) : [FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], (128, 128) : [FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], (256, 256) : [FmhaFwdTileSize( 64, 32, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], } # fmt: skip - elif dtype in ["fp8fp32"]: + elif dtype in cls._DT_FP8FP32: return { # bm0, bn0, bk0, bn1, bk1, (128, 128) : [FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], } # fmt: skip else: - return None + raise ValueError(f"unsupported dtype={dtype}") - @staticmethod - def get_pipelines(dtype, hdim, hdim_v, receipt, mask_impl) -> List[FmhaFwdPipeline]: + @classmethod + def get_pipelines( + cls, dtype, hdim, hdim_v, receipt, mask_impl + ) -> List[FmhaFwdPipeline]: pipelines = [] - if dtype in ["fp16", "bf16"]: - squant = "f" - for logits, mask, bias, lse, dropout, skip in itertools.product( + if dtype in cls._DT_FP16_BF16: + qscale = "no" + for logits, mask, bias, lse, dropout, skip, sink in itertools.product( ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"], + ["t", "f"], ): - pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, squant, mask, skip, "f")) # fmt: skip - elif dtype in ["fp8", "fp8bf16", "fp8fp32"]: + # pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_hpad", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + if receipt == 1: + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + elif dtype in cls._DT_FP8_FP8BF16 or dtype in cls._DT_FP8FP32: # no need lse/dropout kernels - for logits, squant, mask, bias in itertools.product( - ["f"], ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys() + for logits, qscale, mask, bias in itertools.product( + ["f"], + ["no", "pertensor", "blockscale"], + get_mask_map(mask_impl).keys(), + ["no"], ): - pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, "f", "f", squant, mask, "f", "f")) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "f", "f", logits, bias, "f", "f", squant, mask, "f", "f")) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, "f", "f", qscale, mask, "f", "f", "f")) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "f", "f", logits, bias, "f", "f", qscale, mask, "f", "f", "f")) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", "f")) # fmt: skip + return pipelines + + +class KernelComponentFactoryGfx125(CompatibilityRuleFactory): + arch = ArchTrait("gfx125") + + _DT_FP16_BF16 = ("fp16", "bf16") + _DT_FP8_FP8BF16 = ("fp8", "fp8bf16") + _DT_FP8FP32 = ("fp8fp32",) + + @classmethod + def supported_dtypes(cls) -> Tuple[str]: + return cls._DT_FP16_BF16 + cls._DT_FP8_FP8BF16 + cls._DT_FP8FP32 + + @classmethod + def get_hdim_tile_size_dict(cls, dtype: str) -> Optional[dict]: + if dtype in cls._DT_FP16_BF16: + return { + # bm0, bn0, bk0, bn1, bk1, + ( 32, 32) : [FmhaFwdTileSize( 64, 64, 32, 32, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1)], + ( 64, 64) : [FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1)], + (128, 128) : [FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1, CppConstraint("a.max_seqlen_q < 2048")), + FmhaFwdTileSize(128, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1)], + (192, 128) : [FmhaFwdTileSize( 64, 64, 32, 128, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1)], + (256, 256) : [FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1)], + } # fmt: skip + elif dtype in cls._DT_FP8_FP8BF16: + return { + # bm0, bn0, bk0, bn1, bk1, + ( 64, 64) : [FmhaFwdTileSize(128, 64, 64, 64, 64, 128, 4, 1, 1, 4, 1, 1, 16, 16, 64, 16, 16, 64, -1)], + (128, 128) : [FmhaFwdTileSize( 64, 64, 64, 128, 64, 128, 4, 1, 1, 4, 1, 1, 16, 16, 64, 16, 16, 64, -1)], + #(256, 256) : [FmhaFwdTileSize( 64, 32, 64, 256, 64, 256, 4, 1, 1, 4, 1, 1, 16, 16, 64, 16, 16, 64, -1)], + } # fmt: skip + elif dtype in cls._DT_FP8FP32: + return { + # bm0, bn0, bk0, bn1, bk1, + (128, 128) : [FmhaFwdTileSize( 64, 64, 64, 128, 64, 128, 4, 1, 1, 4, 1, 1, 16, 16, 64, 16, 16, 64, -1)], + } # fmt: skip else: - assert False + raise ValueError(f"unsupported dtype={dtype}") + + @classmethod + def get_pipelines( + cls, dtype, hdim, hdim_v, receipt, mask_impl + ) -> List[FmhaFwdPipeline]: + pipelines = [] + if dtype in cls._DT_FP16_BF16: + qscale = "no" + # qr_tdm: gfx1250 TDM pipeline, preferred for d=128. + # Emitted first so runtime dispatcher selects qr_tdm over qr + # when both match (dispatch order = list order in generated code). + # NOTE: dropout is not yet implemented in qr_tdm — only emit + # dropout="f" so dropout workloads fall through to qr. + if hdim == 128 and hdim_v == 128: + for logits, mask, bias, lse, sink in itertools.product( + ["t", "f"], + get_mask_map(mask_impl).keys(), + BIAS_MAP.keys(), + ["t", "f"], + ["t", "f"], + ): + pipelines.append(FmhaFwdPipeline("qr_tdm", "row", "f", "f", "f", "f", logits, bias, lse, "f", qscale, mask, "f", "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_tdm", "row", "f", "f", "t", "t", logits, bias, lse, "f", qscale, mask, "f", "f", sink)) # fmt: skip + + # qr: generic pipeline fallback for trait combos not covered by + # qr_tdm (e.g., bias, dropout, skip, d!=128). + for logits, mask, bias, lse, dropout, skip, sink in itertools.product( + ["t", "f"], + get_mask_map(mask_impl).keys(), + BIAS_MAP.keys(), + ["t", "f"], + ["t", "f"], + ["t", "f"], + ["t", "f"], + ): + pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip + elif dtype in cls._DT_FP8_FP8BF16 or dtype in cls._DT_FP8FP32: + # no need lse/dropout kernels + for logits, qscale, mask, bias in itertools.product( + ["f"], ["no", "pertensor"], get_mask_map(mask_impl).keys(), ["no"] + ): + pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, "f", "f", qscale, mask, "f", "f", "f")) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", "f")) # fmt: skip return pipelines -class CustomFactory(KernelComponentFactoryGfx9): - @staticmethod - def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: +class CustomFactory(KernelComponentFactoryGfx9, CompatibilityRuleFactoryGfx9): + @classmethod + def get_hdim_tile_size_dict(cls, dtype: str) -> Optional[dict]: result = KernelComponentFactoryGfx9.get_hdim_tile_size_dict(dtype) - if dtype == "fp16" or dtype == "bf16": + if dtype in cls._DT_FP16_BF16: if (128, 128) in result.keys(): result[(128, 128)].insert(0, FmhaFwdTileSize( 64, 128, 64, 128, 64, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1, CppConstraint("get_num_blocks(128) < num_cus * min_cu_util_rate"))) # fmt: skip return result @@ -853,157 +1455,203 @@ def get_factory(target: str): return KernelComponentFactoryGfx950 if target.startswith("gfx9"): return KernelComponentFactoryGfx9 - + if target.startswith("gfx115"): + return KernelComponentFactoryGfx115 + if target.startswith("gfx11"): + return KernelComponentFactoryGfx11 + if target.startswith("gfx125"): + return KernelComponentFactoryGfx125 if target.startswith("gfx12"): return KernelComponentFactoryGfx12 raise Exception(f"Unsupported device target {target}") +@dataclass(frozen=True) +class Product: + name: str + rule: CompatibilityRule + + def __call__(self, problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + return self.rule(problem_ctx, kernel_ctx) + + +def get_product(receipt: int) -> Product: + # Flash attention integration + if receipt == 2: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + cond = problem_ctx.dtype in ["fp16", "bf16"] + cond &= kernel_ctx.pipeline.F_vlayout == "row" + cond &= kernel_ctx.pipeline.F_bias in ["no", "alibi"] + cond &= kernel_ctx.pipeline.F_qscale == "no" + cond &= kernel_ctx.pipeline.F_skip == "f" + cond &= kernel_ctx.pipeline.F_sink == "f" + # FlashAttention direct fwd wrappers always use softcap disabled and LSE enabled. + cond &= kernel_ctx.pipeline.F_logits == "f" + cond &= kernel_ctx.pipeline.F_lse == "t" + return cond + + return Product(name="Flash attention integration", rule=fit) + # Receipt 3 forward coverage used by CK library / smoke tests + elif receipt == 3: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + cond = problem_ctx.dtype in ["fp16", "bf16"] + cond &= kernel_ctx.pipeline.F_vlayout == "row" + cond &= kernel_ctx.pipeline.F_bias in ["no", "alibi", "bias"] + cond &= kernel_ctx.pipeline.F_qscale == "no" + cond &= kernel_ctx.pipeline.F_skip == "f" + return cond + + return Product(name="Flash attention integration", rule=fit) + # PyTorch integration + elif receipt == 4: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + cond = problem_ctx.dtype in ["fp16", "bf16"] + cond &= kernel_ctx.pipeline.F_vlayout == "row" + cond &= kernel_ctx.pipeline.F_bias in ["no", "bias"] + cond &= kernel_ctx.pipeline.F_qscale == "no" + cond &= kernel_ctx.pipeline.F_skip == "f" + cond &= kernel_ctx.pipeline.F_logits == "f" + return cond + + return Product(name="PyTorch integration", rule=fit) + # Aiter(mha_fwd) integration + elif receipt == 100: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + cond = problem_ctx.dtype in ["fp16", "bf16", "fp8bf16"] + cond &= problem_ctx.mode == "batch" + cond &= kernel_ctx.pipeline.F_vlayout == "row" + if problem_ctx.dtype == "fp8bf16": + cond &= problem_ctx.hdim == 128 or problem_ctx.hdim == 192 + return cond + + return Product(name="Aiter(mha_fwd) integration", rule=fit) + # Aiter(mha_varlen_fwd) integration + elif receipt == 200: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + cond = problem_ctx.dtype in ["fp16", "bf16", "fp8bf16"] + cond &= problem_ctx.mode == "group" + cond &= kernel_ctx.pipeline.F_vlayout == "row" + if problem_ctx.dtype == "fp8bf16": + cond &= problem_ctx.hdim == 128 or problem_ctx.hdim == 192 + return cond + + return Product(name="Aiter(mha_varlen_fwd) integration", rule=fit) + # aiter::mha_fwd C++ api integration + elif receipt == 600: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + cond = problem_ctx.dtype in ["fp16", "bf16", "fp8bf16"] + cond &= kernel_ctx.pipeline.F_vlayout == "row" + if problem_ctx.dtype == "fp8bf16": + cond &= problem_ctx.hdim == 128 or problem_ctx.hdim == 192 + return cond + + return Product(name="aiter::mha_fwd C++ api integration", rule=fit) + # TransformerEngine integration + elif receipt == 700: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + cond = problem_ctx.dtype in ["fp16", "bf16"] + cond &= kernel_ctx.pipeline.F_vlayout == "row" + cond &= kernel_ctx.pipeline.F_qscale == "no" + cond &= kernel_ctx.pipeline.F_lse == "t" + cond &= kernel_ctx.pipeline.F_skip == "f" + cond &= kernel_ctx.pipeline.F_sink == "f" + cond &= kernel_ctx.pipeline.F_logits == "f" + return cond + + return Product(name="TransformerEngine integration", rule=fit) + elif receipt == 888: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + cond = problem_ctx.dtype in ["fp8bf16", "fp8fp32"] + cond &= kernel_ctx.pipeline.F_vlayout == "row" + cond &= problem_ctx.hdim == 128 or problem_ctx.hdim == 192 + return cond + + return Product(name="receipt = 888", rule=fit) + # fp32 only, all variations + elif receipt == 800: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + cond = problem_ctx.dtype == "fp32" + cond &= kernel_ctx.pipeline.F_skip == "f" + cond &= kernel_ctx.pipeline.F_logits == "f" + return cond + + return Product(name="fp32 only, all variations", rule=fit) + # fp32 only, minimal set of parameters + elif receipt == 801: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + cond = problem_ctx.dtype == "fp32" + cond &= problem_ctx.hdim in [48, 128] + cond &= problem_ctx.mode == "batch" + cond &= kernel_ctx.pipeline.F_bias == "no" + cond &= kernel_ctx.pipeline.F_lse == "f" + cond &= kernel_ctx.pipeline.F_dropout == "f" + cond &= kernel_ctx.pipeline.F_skip == "f" + cond &= kernel_ctx.pipeline.F_logits == "f" + cond &= kernel_ctx.pipeline.F_mask == "s_no" + return cond + + return Product(name="fp32 only, minimal set of parameters", rule=fit) + # Don't build fp32 by default + else: + + def fit(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool: + return problem_ctx.dtype != "fp32" + + return Product(name="Default", rule=fit) + + def get_fwd_blobs( targets: List[str], kernel_filter: Optional[str], receipt, optdim_list, mask_impl ) -> Tuple[FmhaFwdApiPool, List[FmhaFwdKernel]]: gen = list() - api_pool = FmhaFwdApiPool(mask_impl) + api_pool = FmhaFwdApiPool() factories = get_factories_for_targets(targets, get_factory) - for factory, dtype in itertools.product(factories, FWD_DTYPE_MAP.keys()): + for factory, dtype in ((f, t) for f in factories for t in f.supported_dtypes()): d = factory.get_hdim_tile_size_dict(dtype) - if d is None: - continue # for hdim_str, mode, mask, bias, lse in itertools.product(d.keys(), MODE_MAP.keys(), MASK_MAP.keys(), ["t", "f"], ["t", "f"]): for ((hdim, hdim_v), tiles), mode in itertools.product( d.items(), MODE_MAP.keys() ): + if optdim_list != [-1]: + if hdim not in optdim_list: + continue for tile, next_tile in zip(tiles, tiles[1:]): assert ( next_tile.F_bm0 >= tile.F_bm0 ), "Tiles must be ordered by increasing bm0" + for tile, pipeline in itertools.product( tiles, factory.get_pipelines(dtype, hdim, hdim_v, receipt, mask_impl) ): - if mode == "group": - if pipeline.F_spad != "t" or pipeline.F_skpad != "t": - # in group mode, spad/skpad must be true, since we can't predict if seqlen of current batch need pad or not - continue - if (hdim, hdim_v) == (192, 128): - # NOTE: this is used to speedup deepseek prefill case, we don't gen training - if pipeline.F_bias != "no" or pipeline.F_dropout == "t": - continue - if factory.arch.name.startswith("gfx9") and dtype != "fp32": - # TODO: update if >=gfx11 archs get qr_async and qr_async_trload support - if pipeline.tag != "qr_async_trload" and ( - ((hdim, hdim_v) == (128, 128) and tile.F_bn0 != 128) - or ((hdim, hdim_v) != (128, 128) and tile.F_bm0 != 128) - ): - # non qr_async_trload only support km0=128 tile size when hdim is not 128 - # non qr_async only support kn0=128 tile size when hdim is 128 - continue - if pipeline.tag == "qr_async_trload" and ( - ((hdim, hdim_v) == (128, 128) and tile.F_bn0 == 128) - or ((hdim, hdim_v) not in [(64, 64), (128, 128)]) - ): - continue - # logits_soft_cap is only allowed if no bias - if not ( - (pipeline.F_logits == "t" and pipeline.F_bias == "no") - or pipeline.F_logits == "f" - ): - continue - k = FmhaFwdKernel( - F_arch=factory.arch, - F_idx=0, - F_hdim=hdim, - F_dtype=dtype, - F_mode=mode, - F_tile=tile, - F_pipeline=pipeline, - mask_impl=mask_impl, + problem_ctx = ProblemContext( + dtype=dtype, mode=mode, hdim=hdim, hdim_v=hdim_v + ) + kernel_ctx = KernelContext( + tile=tile, pipeline=pipeline, mask_impl=mask_impl ) + rules = factory.get_rules() + product = get_product(receipt) + + if not is_compatible(problem_ctx, kernel_ctx, [*rules, product]): + continue + + k = create_kernel(factory.arch, problem_ctx, kernel_ctx) if kernel_filter != "": if not fnmatch.fnmatch(k.name, kernel_filter): continue - if optdim_list != [-1]: - if hdim not in optdim_list: - continue - # 2 - Flash attention integration - if receipt in (2, 3): - cond = dtype in ["fp16", "bf16"] - cond &= pipeline.F_vlayout == "row" - cond &= pipeline.F_bias in ["no", "alibi"] - cond &= pipeline.F_squant == "f" - cond &= pipeline.F_skip == "f" - if not cond: - continue - # PyTorch integration - elif receipt == 4: - cond = dtype in ["fp16", "bf16"] - cond &= pipeline.F_vlayout == "row" - cond &= pipeline.F_bias in ["no", "bias"] - cond &= pipeline.F_squant == "f" - cond &= mode == "batch" - cond &= pipeline.F_skip == "f" - cond &= pipeline.F_logits == "f" - if not cond: - continue - # Aiter(mha_fwd) integration - elif receipt == 100: - cond = dtype in ["fp16", "bf16", "fp8bf16"] - cond &= mode == "batch" - cond &= pipeline.F_vlayout == "row" - if dtype == "fp8bf16": - cond &= hdim == 128 - if not cond: - continue - # Aiter(mha_varlen_fwd) integration - elif receipt == 200: - cond = dtype in ["fp16", "bf16", "fp8bf16"] - cond &= mode == "group" - cond &= pipeline.F_vlayout == "row" - if dtype == "fp8bf16": - cond &= hdim == 128 - if not cond: - continue - # aiter::mha_fwd C++ api integration - elif receipt == 600: - cond = dtype in ["fp16", "bf16", "fp8bf16"] - cond &= pipeline.F_vlayout == "row" - if dtype == "fp8bf16": - cond &= hdim == 128 - if not cond: - continue - elif receipt == 888: - cond = dtype in ["fp8", "fp8bf16", "fp8fp32"] - cond &= pipeline.F_vlayout == "row" - cond &= hdim == 128 - if not cond: - continue - - # fp32 only, all variations - if receipt == 800: - cond = dtype == "fp32" - cond &= pipeline.F_skip == "f" - cond &= pipeline.F_logits == "f" - if not cond: - continue - # fp32 only, minimal set of parameters - elif receipt == 801: - cond = dtype == "fp32" - cond &= hdim in [48, 128] - cond &= mode == "batch" - cond &= pipeline.F_bias == "no" - cond &= pipeline.F_lse == "f" - cond &= pipeline.F_dropout == "f" - cond &= pipeline.F_skip == "f" - cond &= pipeline.F_logits == "f" - cond &= pipeline.F_mask == "s_no" - if not cond: - continue - else: - # Don't build fp32 by default - if dtype == "fp32": - continue api_pool.register_traits(k.api_trait()) gen.append(k) @@ -1012,11 +1660,28 @@ def get_fwd_blobs( def write_single_fwd_kernel(kernel: FmhaFwdKernel, autogen_dir: Path) -> None: - update_file(autogen_dir / kernel.filename, kernel.template) + update_file(autogen_dir / kernel.filename, kernel.render()) -def write_fwd_api(api_pool: FmhaFwdApiPool, autogen_dir: Path) -> None: - update_file(autogen_dir / FMHA_FWD_API_FILENAME, api_pool.api) +def write_fwd_api( + api_pool: FmhaFwdApiPool, + autogen_dir: Path, +) -> None: + def accept_only_v3(trait: FmhaFwdApiTrait) -> bool: + return trait.pipeline_tag == "qr_async_trload_v3" + + def accept_only_v2(trait: FmhaFwdApiTrait) -> bool: + return not accept_only_v3(trait) + + content = "".join( + [ + FMHA_FWD_API_HEADER, + api_pool.render("fmha_fwd_v2", filter_fn=accept_only_v2), + api_pool.render("fmha_fwd_v3", filter_fn=accept_only_v3), + FMHA_FWD_API_FOOTER, + ] + ) + update_file(autogen_dir / FMHA_FWD_API_FILENAME, content) def write_blobs( @@ -1048,5 +1713,5 @@ def list_blobs( targets, kernel_filter, receipt, optdim_list, mask_impl ) for kernel in kernels: - f.write(str(file_path.parent / GEN_DIR / kernel.filename) + "\n") - f.write(str(file_path.parent / GEN_DIR / FMHA_FWD_API_FILENAME) + "\n") + f.write((file_path.parent / GEN_DIR / kernel.filename).as_posix() + "\n") + f.write((file_path.parent / GEN_DIR / FMHA_FWD_API_FILENAME).as_posix() + "\n") diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd_appendkv.py b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd_appendkv.py index ee24949cca..f4807fb062 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd_appendkv.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd_appendkv.py @@ -1,7 +1,6 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. # generate kernel instances to speed up compilation - import copy import fnmatch import itertools @@ -389,16 +388,39 @@ class KernelComponentFactoryGfx9(KernelComponentFactoryBase): arch = ArchTrait("gfx9") +class KernelComponentFactoryGfx11(KernelComponentFactoryBase): + arch = ArchTrait("gfx11") + + @staticmethod + def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: + if dtype in ["fp16", "bf16"]: + return KernelComponentFactoryBase.get_hdim_tile_size_dict(dtype) + return None + + @staticmethod + def get_pipelines(dtype, hdim) -> List[FmhaFwdAppendKVPipeline]: + if dtype in ["fp16", "bf16"]: + return KernelComponentFactoryBase.get_pipelines(dtype, hdim) + return [] + + class KernelComponentFactoryGfx12(KernelComponentFactoryBase): arch = ArchTrait("gfx12") +class KernelComponentFactoryGfx125(KernelComponentFactoryBase): + arch = ArchTrait("gfx125") + + def get_factory(target: str): # Place more specific architectures first if target.startswith("gfx9"): return KernelComponentFactoryGfx9 - + if target.startswith("gfx11"): + return KernelComponentFactoryGfx11 + if target.startswith("gfx125"): + return KernelComponentFactoryGfx125 if target.startswith("gfx12"): return KernelComponentFactoryGfx12 @@ -498,5 +520,8 @@ def list_blobs( targets, kernel_filter, receipt, mask_impl, optdim_list ) for kernel in kernels: - f.write(str(file_path.parent / GEN_DIR / kernel.filename) + "\n") - f.write(str(file_path.parent / GEN_DIR / FMHA_FWD_APPENDKV_API_FILENAME) + "\n") + f.write((file_path.parent / GEN_DIR / kernel.filename).as_posix() + "\n") + f.write( + (file_path.parent / GEN_DIR / FMHA_FWD_APPENDKV_API_FILENAME).as_posix() + + "\n" + ) diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd_splitkv.py b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd_splitkv.py index db533583b1..8b43571fbc 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd_splitkv.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_fwd_splitkv.py @@ -1,7 +1,6 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. # generate kernel instances to speed up compilation - import copy import fnmatch import itertools @@ -74,7 +73,8 @@ {F_pagedkv}, kHasUnevenSplits, kMergeNumHeadGroupsSeqLenQ, - {F_occupancy}>; + {F_occupancy}, + {F_sink}>; using fmha_pipeline_problem = ck_tile::BlockFmhaFwdSplitKVPipelineProblem< typename FmhaFwdTypeConfig::QDataType, @@ -118,16 +118,17 @@ }} // anonymous namespace using trait_{F_idx} = fmha_fwd_splitkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, - {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_squant}, {F_pagedkv}, {F_spad}, {F_skpad}, {F_dpad}, + {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_squant}, {F_pagedkv}, {F_sink}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>; +#ifdef __clang__ #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wtautological-compare" - +#endif namespace {{ template void run_instance(const ck_tile::stream_config& s, fmha_fwd_splitkv_args a) {{ - if constexpr ({F_hdim} == 128 && {F_bias} == ck_tile::BlockAttentionBiasEnum::NO_BIAS + if constexpr ({F_bias} == ck_tile::BlockAttentionBiasEnum::NO_BIAS && (std::is_same_v<{F_mask}, ck_tile::SimplifiedGenericAttentionMask> || std::is_same_v<{F_mask}, FmhaMasks::NoMask>)) {{ if (a.max_seqlen_q == 1 && a.nhead_k < a.nhead_q) {{ @@ -141,7 +142,9 @@ }} }} // anonymous namespace +#ifdef __clang__ #pragma clang diagnostic pop +#endif template<> void fmha_fwd_splitkv_oneshot_(const ck_tile::stream_config& s, fmha_fwd_splitkv_args a) @@ -280,8 +283,8 @@ """ FMHA_FWD_SPLITKV_API_INNER_DISPATCH = """{F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.do_fp8_static_quant == {F_squant}) && - ((a.block_table_ptr != nullptr) == {F_pagedkv}) && ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{ - using traits_ = fmha_fwd_splitkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, true, {F_squant}, {F_pagedkv}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>; + ((a.block_table_ptr != nullptr) == {F_pagedkv}) && (t.has_sink == {F_sink}) && ({F_scheck}) && ({F_seqtune}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{ + using traits_ = fmha_fwd_splitkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, true, {F_squant}, {F_pagedkv},{F_sink}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>; // get combine kernel tile sizes using OaccDataType = typename FmhaFwdTypeConfig<{F_dtype}>::OaccDataType; @@ -333,6 +336,7 @@ class FmhaFwdSplitKVApiTrait: dpad: str dvpad: str pagedkv: str + sink: str # sink or not bn1comb: int # tile size along v head_dim of combine kernel @property @@ -340,7 +344,7 @@ def name(self) -> str: return ( f"{self.hdim}-{self.dtype}-{self.mode}-{self.bm0}-{self.bn0}-{self.bk0}-{self.bn0}-{self.bk1}-{self.bk0max}-" + f"{self.vlayout}-{self.logits}-{self.mask}-{self.bias}-{self.lse}-{self.squant}-{self.spad}-{self.skpad}-{self.dpad}-" - + f"{self.dvpad}-{self.pagedkv}" + + f"{self.dvpad}-{self.pagedkv}-{self.sink}" ) @property @@ -360,6 +364,14 @@ def scheck(self) -> str: else: assert False + def seqtune(self, max_bm0: int) -> str: + if self.bm0 == max_bm0: + return "true/*fall back to largest tile*/" + else: + if self.mode == "group": + return f"a.max_seqlen_q <= {self.bm0}" + return f"a.seqlen_q <= {self.bm0}" + @property def skcheck(self) -> str: if self.mode == "group": @@ -426,6 +438,7 @@ class FmhaFwdSplitKVPipeline: F_lse: str # F_squant: str # F_pagedkv: str # t/f + F_sink: str # t/f F_mask: str # value from MASK_MAP @property @@ -486,6 +499,10 @@ def pad_name() -> str: n += "_pagedkv" else: n += "_npagedkv" + if self.F_sink == "t": + n += "_sink" + else: + n += "_nsink" return n @@ -552,6 +569,7 @@ def api(self) -> str: for i_dtype, (dtype, pool_by_dtype) in enumerate(pool_by_arch.items()): per_hdim_case = str() for i_hdim, (hdim, pool_by_hdim) in enumerate(pool_by_dtype.items()): + max_bm0 = max((t.bm0 for t in pool_by_hdim), default=0) inners = str() for i_trait, trait in enumerate(pool_by_hdim): inners += FMHA_FWD_SPLITKV_API_INNER_DISPATCH.format( @@ -568,7 +586,9 @@ def api(self) -> str: F_lse=BOOL_MAP[trait.lse], F_squant=BOOL_MAP[trait.squant], F_pagedkv=BOOL_MAP[trait.pagedkv], + F_sink=BOOL_MAP[trait.sink], F_scheck=trait.scheck, + F_seqtune=trait.seqtune(max_bm0), F_skcheck=trait.skcheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck, @@ -668,6 +688,7 @@ def template(self) -> str: F_squant=BOOL_MAP[self.F_pipeline.F_squant], F_pagedkv=BOOL_MAP[self.F_pipeline.F_pagedkv], F_occupancy=self.F_tile.F_occupancy, + F_sink=BOOL_MAP[self.F_pipeline.F_sink], F_pipeline_enum=PIPELINE_ENUM_MAP[self.F_pipeline.tag], F_mask=get_mask_map(self.mask_impl)[self.F_pipeline.F_mask], F_mode=MODE_MAP[self.F_mode], @@ -741,19 +762,24 @@ def get_pipelines(dtype, hdim, mask_impl) -> List[FmhaFwdSplitKVPipeline]: squant = "t" if dtype == "fp8" else "f" pipelines = [] if dtype in ["fp16", "bf16"]: - for logits, mask, bias, pagedkv in itertools.product( - ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"] + for logits, mask, bias, pagedkv, sink in itertools.product( + ["t", "f"], + get_mask_map(mask_impl).keys(), + BIAS_MAP.keys(), + ["t", "f"], + ["t", "f"], ): - pipelines.append(Pipeline("qr", "row", "f", "t", "f", "f", logits, bias, "t", squant, pagedkv, mask)) # fmt: skip - pipelines.append(Pipeline("qr", "row", "t", "f", "f", "f", logits, bias, "t", squant, pagedkv, mask)) # fmt: skip - pipelines.append(Pipeline("qr", "row", "t", "t", "f", "f", logits, bias, "t", squant, pagedkv, mask)) # fmt: skip - pipelines.append(Pipeline("qr", "row", "t", "t", "t", "t", logits, bias, "t", squant, pagedkv, mask)) # fmt: skip + pipelines.append(Pipeline("qr", "row", "f", "t", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip + pipelines.append(Pipeline("qr", "row", "t", "f", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip + pipelines.append(Pipeline("qr", "row", "t", "t", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip + pipelines.append(Pipeline("qr", "row", "t", "t", "t", "t", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip + pipelines.append(Pipeline("qr_nwarp_sshuffle", "row", "t", "t", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip elif dtype in ["fp8", "bf8"]: for logits, mask, bias in itertools.product( ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys() ): - pipelines.append(Pipeline("qr", "row", "f", "f", "f", "f", logits, bias, "t", squant, "f", mask)) # fmt: skip - pipelines.append(Pipeline("qr", "row", "t", "t", "f", "f", logits, bias, "t", squant, "f", mask)) # fmt: skip + pipelines.append(Pipeline("qr", "row", "f", "f", "f", "f", logits, bias, "t", squant, "f", "f", mask)) # fmt: skip + pipelines.append(Pipeline("qr", "row", "t", "t", "f", "f", logits, bias, "t", squant, "f", "f", mask)) # fmt: skip elif dtype in ["fp8fp16", "fp8bf16"]: # TODO None @@ -824,6 +850,27 @@ def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: return None +class KernelComponentFactoryGfx11(KernelComponentFactoryBase): + arch = ArchTrait("gfx11") + + @staticmethod + def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: + if dtype in ["fp16", "bf16"]: + return { + # bm0, bn0, bk0, bn1, bk1, + "32" : [FmhaFwdTileSize( 16, 64, 16, 32, 32, 32, 1, 2, 1, 1, 2, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "64" : [FmhaFwdTileSize( 16, 64, 32, 64, 32, 64, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "128": [FmhaFwdTileSize( 16, 64, 32, 128, 32, 128, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "256": [FmhaFwdTileSize( 16, 64, 32, 256, 32, 256, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + } # fmt: skip + else: + return None + + class KernelComponentFactoryGfx12(KernelComponentFactoryBase): arch = ArchTrait("gfx12") @@ -831,11 +878,15 @@ class KernelComponentFactoryGfx12(KernelComponentFactoryBase): def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: if dtype in ["fp16", "bf16"]: return { - # bm0, bn0, bk0, bn1, bk1, - "32" : FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "64" : FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "128": FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "256": FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + # bm0, bn0, bk0, bn1, bk1, + "32" : [FmhaFwdTileSize( 16, 64, 16, 32, 32, 32, 1, 2, 1, 1, 2, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "64" : [FmhaFwdTileSize( 16, 64, 32, 64, 32, 64, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "128": [FmhaFwdTileSize( 16, 128, 32, 128, 32, 128, 1, 8, 1, 1, 8, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "256": [FmhaFwdTileSize( 16, 128, 32, 256, 32, 256, 1, 8, 1, 1, 8, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], } # fmt: skip elif dtype in ["fp8", "bf8"]: return { @@ -847,12 +898,38 @@ def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: return None +class KernelComponentFactoryGfx125(KernelComponentFactoryBase): + arch = ArchTrait("gfx125") + + @staticmethod + def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: + if dtype in ["fp16", "bf16"]: + return { + # bm0, bn0, bk0, bn1, bk1, + "32" : FmhaFwdTileSize( 64, 64, 32, 32, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1), + "64" : FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1), + "128": FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1), + "256": FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1), + } # fmt: skip + elif dtype in ["fp8", "bf8"]: + return { + # bm0, bn0, bk0, bn1, bk1, + "64" : FmhaFwdTileSize(128, 64, 64, 64, 64, 128, 4, 1, 1, 4, 1, 1, 16, 16, 64, 16, 16, 64, -1), + "128": FmhaFwdTileSize( 64, 64, 64, 128, 64, 128, 4, 1, 1, 4, 1, 1, 16, 16, 64, 16, 16, 64, -1), + } # fmt: skip + else: + return None + + def get_factory(target: str): # Place more specific architectures first if target.startswith("gfx9"): return KernelComponentFactoryGfx9 - + if target.startswith("gfx11"): + return KernelComponentFactoryGfx11 + if target.startswith("gfx125"): + return KernelComponentFactoryGfx125 if target.startswith("gfx12"): return KernelComponentFactoryGfx12 @@ -872,11 +949,17 @@ def get_fwd_splitkv_blobs( d = factory.get_hdim_tile_size_dict(dtype) if d is None: continue - # for hdim_str, mode, mask, bias, lse in itertools.product(d.keys(), MODE_MAP.keys(), MASK_MAP.keys(), ["t", "f"], ["t", "f"]): for hdim_str, mode in itertools.product(d.keys(), MODE_MAP.keys()): - tile = d[hdim_str] + tiles = d[hdim_str] + if not isinstance(tiles, list): + tiles = [tiles] hdim = int(hdim_str) - for pipeline in factory.get_pipelines(dtype, hdim, mask_impl): + for tile, pipeline in itertools.product( + tiles, factory.get_pipelines(dtype, hdim, mask_impl) + ): + # Use qr_nwarp_sshuffle with multiple N warps and qr otherwise + if (tile.F_rn0 != 1) != (pipeline.tag == "qr_nwarp_sshuffle"): + continue if mode == "group": if pipeline.F_spad != "t" or pipeline.F_skpad != "t": # in group mode, spad/skpad must be true, since we can't predict if seqlen of current batch need pad or not @@ -908,7 +991,10 @@ def get_fwd_splitkv_blobs( cond = dtype in ["fp16", "bf16"] cond &= pipeline.F_vlayout == "row" cond &= pipeline.F_bias in ["no", "alibi"] + # FlashAttention splitkv paths use softcap-disabled kernels only. + cond &= pipeline.F_logits == "f" cond &= pipeline.F_squant == "f" + cond &= pipeline.F_sink == "f" if not cond: continue # PyTorch integration @@ -918,6 +1004,7 @@ def get_fwd_splitkv_blobs( cond &= pipeline.F_bias in ["no", "bias"] cond &= pipeline.F_squant == "f" cond &= mode == "batch" + cond &= pipeline.F_sink == "f" if not cond: continue # Aiter(mha_varlen_fwd) integration @@ -935,6 +1022,8 @@ def get_fwd_splitkv_blobs( cond &= pipeline.F_squant == "f" if not cond: continue + elif receipt == 700: + continue # TE does not use this API # fp32 only if receipt == 800 or receipt == 801: @@ -1076,6 +1165,7 @@ def write_blobs( lse=kernel.F_pipeline.F_lse, squant=kernel.F_pipeline.F_squant, pagedkv=kernel.F_pipeline.F_pagedkv, + sink=kernel.F_pipeline.F_sink, spad=kernel.F_pipeline.F_spad, skpad=kernel.F_pipeline.F_skpad, dpad=kernel.F_pipeline.F_dpad, @@ -1102,10 +1192,13 @@ def list_blobs( targets, filter_list[0], receipt, optdim_list ) for kernel in kernels: - f.write(str(file_path.parent / GEN_DIR / kernel.filename) + "\n") + f.write((file_path.parent / GEN_DIR / kernel.filename).as_posix() + "\n") kernels = get_fwd_splitkv_blobs( targets, filter_list[1], receipt, mask_impl, optdim_list ) for kernel in kernels: - f.write(str(file_path.parent / GEN_DIR / kernel.filename) + "\n") - f.write(str(file_path.parent / GEN_DIR / FMHA_FWD_SPLITKV_API_FILENAME) + "\n") + f.write((file_path.parent / GEN_DIR / kernel.filename).as_posix() + "\n") + f.write( + (file_path.parent / GEN_DIR / FMHA_FWD_SPLITKV_API_FILENAME).as_posix() + + "\n" + ) diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_pagedkv_prefill.py b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_pagedkv_prefill.py index 17ac129e64..f90cd0ef8e 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_pagedkv_prefill.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/ops/fmha_pagedkv_prefill.py @@ -1,7 +1,6 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. # generate kernel instances to speed up compilation - import copy import fnmatch import itertools @@ -66,7 +65,8 @@ {F_pagedkv}, //pagedkv {F_squant}, {F_occupancy}, - {F_skip}>; + {F_skip}, + {F_sink}>; using fmha_variant_{F_idx} = ck_tile::ComposedAttention<{F_logits} * ck_tile::LOGITS_SOFT_CAP, CK_TILE_FMHA_FWD_FAST_EXP2>; @@ -101,7 +101,7 @@ ck_tile::FmhaFwdPagedKVKernel; using trait_{F_idx} = fmha_fwd_pagedkv_traits_<{F_hdim}, {F_dtype}, {F_mode},{F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, - {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_pagedkv}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_skip}>; + {F_pipeline_enum}, {F_logits}, fmha_mask_{F_idx}, {F_bias}, {F_lse}, {F_pagedkv}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_skip}, {F_sink}>; template<> float fmha_fwd_pagedkv_(const ck_tile::stream_config& s, fmha_fwd_pagedkv_args a) @@ -130,9 +130,9 @@ }} """ -FMHA_FWD_API_INNER_DISPATCH = """{F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_lse == {F_lse}) && (t.use_pagedkv == {F_pagedkv}) && (t.do_fp8_static_quant == {F_squant}) && (t.skip_min_seqlen_q == {F_skip}) && +FMHA_FWD_API_INNER_DISPATCH = """{F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_lse == {F_lse}) && (t.use_pagedkv == {F_pagedkv}) && (t.do_fp8_static_quant == {F_squant}) && (t.skip_min_seqlen_q == {F_skip}) && (t.has_sink == {F_sink}) && ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{ - using trait_ = fmha_fwd_pagedkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, {F_lse}, {F_pagedkv}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_skip}>; + using trait_ = fmha_fwd_pagedkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, {F_lse}, {F_pagedkv}, {F_squant}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_skip},{F_sink}>; return fmha_fwd_pagedkv_(s, a); }} """ @@ -164,12 +164,13 @@ class FmhaFwdApiTrait: dpad: str dvpad: str skip: str + sink: str @property def name(self) -> str: return ( f"{self.hdim}-{self.dtype}-{self.mode}-{self.bm0}-{self.bn0}-{self.bk0}-{self.bn0}-{self.bk1}-{self.bk0max}-" - + f"{self.vlayout}-{self.logits}-{self.mask}-{self.bias}-{self.lse}-{self.pagedkv}-{self.squant}-{self.spad}-{self.skpad}-{self.dpad}-{self.dvpad}-{self.skip}" + + f"{self.vlayout}-{self.logits}-{self.mask}-{self.bias}-{self.lse}-{self.pagedkv}-{self.squant}-{self.spad}-{self.skpad}-{self.dpad}-{self.dvpad}-{self.skip}-{self.sink}" ) @property @@ -257,6 +258,7 @@ class FmhaFwdPipeline: F_squant: str # F_mask: str # value from MASK_MAP F_skip: str # true/false + F_sink: str # true/false @property def name(self) -> str: @@ -321,6 +323,10 @@ def pad_name() -> str: n += "_pagedkv" else: n += "_npagedkv" + if self.F_sink == "t": + n += "_sink" + else: + n += "_nsink" return n @@ -364,6 +370,7 @@ def api(self) -> str: F_lse=BOOL_MAP[trait.lse], F_pagedkv=BOOL_MAP[trait.pagedkv], F_skip=BOOL_MAP[trait.skip], + F_sink=BOOL_MAP[trait.sink], F_squant=BOOL_MAP[trait.squant], F_scheck=trait.scheck, F_skcheck=trait.skcheck, @@ -481,6 +488,7 @@ def template(self) -> str: F_pagedkv=BOOL_MAP[self.F_pipeline.F_pagedkv], F_squant=BOOL_MAP[self.F_pipeline.F_squant], F_skip=BOOL_MAP[self.F_pipeline.F_skip], + F_sink=BOOL_MAP[self.F_pipeline.F_sink], F_occupancy=self.F_tile.F_occupancy, F_pipeline_enum=PIPELINE_ENUM_MAP[self.F_pipeline.tag], F_mask=get_mask_map(self.mask_impl)[self.F_pipeline.F_mask], @@ -527,6 +535,7 @@ def api_trait(self) -> FmhaFwdApiTrait: dpad=self.F_pipeline.F_dpad, dvpad=self.F_pipeline.F_dvpad, skip=self.F_pipeline.F_skip, + sink=self.F_pipeline.F_sink, ) @@ -540,22 +549,23 @@ def get_pipelines(dtype, hdim, mask_impl) -> List[FmhaFwdPipeline]: squant = "t" if dtype == "fp8" else "f" pipelines = [] if dtype in ["fp16", "bf16"]: - for logits, mask, bias, pagedkv, skip in itertools.product( + for logits, mask, bias, pagedkv, skip, sink in itertools.product( ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t"], ["f"], + ["t", "f"], ): - pipelines.append(FmhaFwdPipeline("qr_pagedkv", "row", "t", "f", "f", "f", logits, bias, "f", pagedkv, squant, mask, skip)) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr_pagedkv", "row", "t", "t", "f", "f", logits, bias, "f", pagedkv, squant, mask, skip)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_pagedkv", "row", "t", "f", "f", "f", logits, bias, "f", pagedkv, squant, mask, skip, sink)) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_pagedkv", "row", "t", "t", "f", "f", logits, bias, "f", pagedkv, squant, mask, skip, sink)) # fmt: skip elif dtype in ["fp8", "bf8"]: # no need lse/dropout kernels for logits, mask, bias in itertools.product( ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys() ): - pipelines.append(FmhaFwdPipeline("qr_pagedkv", "row", "f", "f", "f", "f", logits, bias, "f", "t", squant, mask, "f")) # fmt: skip - pipelines.append(FmhaFwdPipeline("qr_pagedkv", "row", "t", "t", "f", "f", logits, bias, "f", "t", squant, mask, "f")) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_pagedkv", "row", "f", "f", "f", "f", logits, bias, "f", "t", squant, mask, "f", "f")) # fmt: skip + pipelines.append(FmhaFwdPipeline("qr_pagedkv", "row", "t", "t", "f", "f", logits, bias, "f", "t", squant, mask, "f", "f")) # fmt: skip elif dtype in ["fp8fp16", "fp8bf16"]: pass # TODO else: @@ -587,6 +597,24 @@ def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: return None +class KernelComponentFactoryGfx11(KernelComponentFactoryBase): + arch = ArchTrait("gfx11") + + @staticmethod + def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: + if dtype in ["fp16", "bf16"]: + return { + # bm0, bn0, bk0, bn1, bk1, + # "32": FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + # "64": FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + "128": FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + # "192": FmhaFwdTileSize( 64, 64, 32, 128, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + # "256": FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + } # fmt: skip + else: + return None + + class KernelComponentFactoryGfx12(KernelComponentFactoryBase): arch = ArchTrait("gfx12") @@ -612,12 +640,40 @@ def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: return None +class KernelComponentFactoryGfx125(KernelComponentFactoryBase): + arch = ArchTrait("gfx125") + + @staticmethod + def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: + if dtype in ["fp16", "bf16"]: + return { + # bm0, bn0, bk0, bn1, bk1, + # "32": FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 32, -1), + # "64": FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 32, -1), + "128": FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1), + # "192": FmhaFwdTileSize( 64, 64, 32, 128, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 32, -1), + # "256": FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 32, -1), + } # fmt: skip + elif dtype in ["fp8", "bf8"]: + return { + # bm0, bn0, bk0, bn1, bk1, + "64": FmhaFwdTileSize(128, 64, 64, 64, 64, 128, 4, 1, 1, 4, 1, 1, 16, 16, 64, 16, 16, 64, -1), + "128": FmhaFwdTileSize( 64, 64, 64, 128, 64, 128, 4, 1, 1, 4, 1, 1, 16, 16, 64, 16, 16, 64, -1), + "256": FmhaFwdTileSize( 64, 64, 64, 256, 64, 256, 4, 1, 1, 4, 1, 1, 16, 16, 64, 16, 16, 64, -1), + } # fmt: skip + else: + return None + + def get_factory(target: str): # Place more specific architectures first if target.startswith("gfx9"): return KernelComponentFactoryGfx9 - + if target.startswith("gfx11"): + return KernelComponentFactoryGfx11 + if target.startswith("gfx125"): + return KernelComponentFactoryGfx125 if target.startswith("gfx12"): return KernelComponentFactoryGfx12 @@ -679,6 +735,7 @@ def get_fwd_blobs( cond &= pipeline.F_bias in ["no", "alibi"] cond &= pipeline.F_squant == "f" cond &= pipeline.F_skip == "f" + cond &= pipeline.F_sink == "f" if not cond: continue # PyTorch integration @@ -688,6 +745,7 @@ def get_fwd_blobs( cond &= pipeline.F_bias in ["no", "bias"] cond &= pipeline.F_squant == "f" cond &= pipeline.F_skip == "f" + cond &= pipeline.F_sink == "f" if not cond: continue # Aiter(mha_fwd) integration @@ -713,6 +771,8 @@ def get_fwd_blobs( cond &= pipeline.F_squant == "f" if not cond: continue + elif receipt == 700: + continue # TE does not use this API # fp32 only if receipt == 800 or receipt == 801: @@ -763,5 +823,5 @@ def list_blobs( targets, kernel_filter, receipt, optdim_list, mask_impl ) for kernel in kernels: - f.write(str(file_path.parent / GEN_DIR / kernel.filename) + "\n") - f.write(str(file_path.parent / GEN_DIR / FMHA_FWD_API_FILENAME) + "\n") + f.write((file_path.parent / GEN_DIR / kernel.filename).as_posix() + "\n") + f.write((file_path.parent / GEN_DIR / FMHA_FWD_API_FILENAME).as_posix() + "\n") diff --git a/Libraries/ComposableKernel/attention/fmha/codegen/utils.py b/Libraries/ComposableKernel/attention/fmha/codegen/utils.py index 7afa4f6dd8..3fefe73ad9 100644 --- a/Libraries/ComposableKernel/attention/fmha/codegen/utils.py +++ b/Libraries/ComposableKernel/attention/fmha/codegen/utils.py @@ -1,7 +1,6 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. # generate kernel instances to speed up compilation - import dataclasses import os.path as path import textwrap diff --git a/Libraries/ComposableKernel/attention/fmha/fmha_bwd.cpp b/Libraries/ComposableKernel/attention/fmha/fmha_bwd.cpp index e5dbacf211..55d22f9168 100644 --- a/Libraries/ComposableKernel/attention/fmha/fmha_bwd.cpp +++ b/Libraries/ComposableKernel/attention/fmha/fmha_bwd.cpp @@ -1,27 +1,9 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT -#include "fmha_bwd.hpp" #include "ck_tile/host.hpp" + +#include "fmha_bwd.hpp" #include "fmha_bwd_runner.hpp" #include @@ -105,6 +87,8 @@ auto create_args(int argc, char *argv[]) { "if set to 1 will use multi-buffer reduction strategy for dq, " "atomic operation " "will not be used") + .insert("sink_grad", "0", + "if set to 1, compute and validate sink token gradient") .insert("json", "0", "0: No Json, 1: Dump Results in Json format") .insert("jsonfile", "fmha_bwd.json", "json file name to dump results"); @@ -139,6 +123,7 @@ auto run(const ck_tile::ArgParser &arg_parser) { bool deterministic = arg_parser.get_bool("deterministic"); std::string init_method = arg_parser.get_str("init"); uint32_t seed = arg_parser.get_uint32("seed"); + bool sink_grad = arg_parser.get_bool("sink_grad"); ck_tile::stream_config stream_config{ nullptr, @@ -155,8 +140,8 @@ auto run(const ck_tile::ArgParser &arg_parser) { return fmha_bwd_run( mode, batch, nhead, nhead_k, seqlen_qs, seqlen_ks, seqlen_qpads, seqlen_kpads, hdim_q, hdim_v, i_perm, o_perm, scale, bias_str, use_dbias, - p_drop, drop_seed, drop_offset, drop_prefs, mask_str, deterministic, - init_method, seed, do_validation, stream_config, json); + p_drop, drop_seed, drop_offset, drop_prefs, mask_str, sink_grad, + deterministic, init_method, seed, do_validation, stream_config, json); } int main(int argc, char *argv[]) { diff --git a/Libraries/ComposableKernel/attention/fmha/fmha_bwd.hpp b/Libraries/ComposableKernel/attention/fmha/fmha_bwd.hpp index 9e4a734262..5c9374a319 100644 --- a/Libraries/ComposableKernel/attention/fmha/fmha_bwd.hpp +++ b/Libraries/ComposableKernel/attention/fmha/fmha_bwd.hpp @@ -1,39 +1,42 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once -#include "bias.hpp" #include "ck_tile/core.hpp" + #include "ck_tile/host/device_prop.hpp" + #include "ck_tile/host/kernel_launch.hpp" -#include "ck_tile/ops/epilogue.hpp" + #include "ck_tile/ops/fmha.hpp" + +#include "ck_tile/ops/epilogue.hpp" + #include "mask.hpp" +#include "bias.hpp" + +#include +#include +#include #include #include #include +// Function pointer type for workspace packing (no captures, points to code +// segment) +using PrepareWorkspaceHostFunc = size_t (*)( + void *, // host_ws + ck_tile::index_t, // batch + ck_tile::index_t, // hdim_q + ck_tile::index_t, // nhead_q + ck_tile::index_t, // seqlen_q + ck_tile::index_t, // seqlen_k + const ck_tile::index_t *, // seqstart_q + const ck_tile::index_t * // seqstart_k +); + struct FmhaBwdFp32 {}; struct FmhaBwdFp16 {}; @@ -118,7 +121,11 @@ struct fmha_bwd_args { void *dk_ptr; void *dv_ptr; void *dbias_ptr; - void *dq_acc_ptr; + void *workspace_ptr; + const void *sink_ptr; // sink scores [nhead] in log-space (LSEDataType); + // nullptr disables sink + void *d_sink_ptr; // sink gradient output [nhead] (LSEDataType); nullptr + // disables sink gradient // Usage notes for sequence length pointer parameters: // @@ -128,13 +135,17 @@ struct fmha_bwd_args { // With padding: // Group mode: // - seqstart_q_ptr, seqstart_k_ptr: Record cumulative physical (including - // padding) sequence lengths. [array size: batch + 1] + // padding) sequence + // lengths. [array size: batch + 1] // - seqlen_q_ptr/seqlen_k_ptr: Records logical (excluding padding) length - // for each sequence. [array size: batch] + // 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] + // (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. + // are mutually + // exclusive. Use one set, not both. // // Batch mode: // - cu_seqlen_q_ptr/cu_seqlen_k_ptr: Records cumulative logical @@ -190,7 +201,6 @@ struct fmha_bwd_args { 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; @@ -203,7 +213,6 @@ struct fmha_bwd_args { 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; @@ -216,12 +225,10 @@ struct fmha_bwd_args { 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; @@ -236,52 +243,43 @@ template 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.do_ptr, args.d_ptr, args.rand_val_ptr, args.dq_ptr, args.dk_ptr, + args.dv_ptr, args.dbias_ptr, args.workspace_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.batch, 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_randval, args.stride_do, args.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.nhead_stride_lsed, args.nhead_stride_dq, args.nhead_stride_dk, + args.nhead_stride_dv, args.nhead_stride_dbias, 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.dq_ptr, args.dk_ptr, + args.dv_ptr, args.dbias_ptr, args.workspace_ptr, args.seqlen_q, + args.seqlen_k, args.batch, 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, + args.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, args.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, args.batch_stride_dq, + args.batch_stride_dk, args.batch_stride_dv, args.batch_stride_dbias, + args.window_size_left, args.window_size_right, args.mask_type, + args.p_drop, args.drop_seed_offset); } }(); @@ -296,14 +294,16 @@ auto fmha_bwd_dot_do_o_create_kargs_and_grids(fmha_bwd_args args) { // 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.o_ptr, args.do_ptr, args.d_ptr, args.lse_ptr, args.sink_ptr, + args.d_sink_ptr, args.p_undrop, args.seqstart_q_ptr, + args.seqlen_q_ptr, args.cu_seqlen_q_ptr, args.hdim_v, args.nhead_q, + 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.o_ptr, args.do_ptr, args.d_ptr, args.lse_ptr, args.sink_ptr, + args.d_sink_ptr, args.p_undrop, args.seqlen_q, args.hdim_v, + args.nhead_q, 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); } @@ -320,17 +320,15 @@ auto fmha_bwd_convert_dq_create_kargs_and_grids(fmha_bwd_args args) { // 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); + args.workspace_ptr, args.dq_ptr, args.batch, args.nhead_q, + 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.nhead_stride_dq); } 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); + args.workspace_ptr, args.dq_ptr, args.batch, args.nhead_q, + args.seqlen_q, args.seqlen_k, args.hdim_q, args.stride_dq, + args.nhead_stride_dq, args.batch_stride_dq); } }(); @@ -358,6 +356,26 @@ void fmha_bwd_dq_dk_dv_oneshot_(const ck_tile::stream_config &, fmha_bwd_args); template std::string fmha_bwd_dq_dk_dv_get_name_(); template int fmha_bwd_dq_dk_dv_maxq_(); +struct fmha_bwd_traits; +template +size_t fmha_bwd_dq_dk_dv_dq_ws_host_size_(int batch_size); +// `total_seqlen_q_padded` is total q tokens across all batches (incl. per-batch +// padding): +// - batch mode: max_batch * seqlen_q +// - group mode: seqstart_q[batch] (== varlen q tensor's first dim) +template +size_t fmha_bwd_dq_dk_dv_dq_ws_device_upper_bound_( + ck_tile::index_t max_batch, ck_tile::index_t hdim_q, + ck_tile::index_t nhead_q, ck_tile::index_t total_seqlen_q_padded, + ck_tile::index_t max_seqlen_k); +template +size_t fmha_bwd_dq_dk_dv_dq_prepare_ws_host_( + void *cpu_ws, ck_tile::index_t batch_size, ck_tile::index_t hdim_q, + ck_tile::index_t nhead_q, ck_tile::index_t seqlen_q, + ck_tile::index_t seqlen_k, const ck_tile::index_t *seqstart_qs, + const ck_tile::index_t *seqstart_ks); +template +bool fmha_bwd_dq_dk_dv_needs_zero_dq_acc_(); template @@ -379,8 +397,7 @@ template std::string fmha_bwd_dot_do_o_get_name_(); template + bool kPadS_, bool kPadD_, bool kIsDeterministic_> struct fmha_bwd_convert_dq_traits_ {}; template @@ -393,10 +410,18 @@ void fmha_bwd_convert_dq_oneshot_(const ck_tile::stream_config &, template std::string fmha_bwd_convert_dq_get_name_(); -// This is the public API, will be generated by script +// Traits that are used to dispatch different kernel implementations for fmha +// backward struct fmha_bwd_traits { + int seqlen_q; + int seqlen_k; + int batch; + int max_seqlen_q; + int max_seqlen_k; int hdim_q; int hdim_v; + int nhead_q; + int nhead_k; std::string data_type; bool is_group_mode; mask_enum mask_type; @@ -408,5 +433,304 @@ struct fmha_bwd_traits { bool is_deterministic; // TODO: padding check is inside this api }; + +template +float fmha_bwd_(const ck_tile::stream_config &s, fmha_bwd_args a) { + if constexpr (!std::is_same_v) { + if (s.log_level_ > 0) + std::cout << ", " << fmha_bwd_dot_do_o_get_name_() << "@" + << fmha_bwd_convert_dq_get_name_() << "@" + << fmha_bwd_dq_dk_dv_get_name_() << std::flush; + return ck_tile::launch_kernel( + s, + [=](const ck_tile::stream_config &s_) { + fmha_bwd_dot_do_o_oneshot_(s_, a); + }, + [=](const ck_tile::stream_config &s_) { + fmha_bwd_dq_dk_dv_oneshot_(s_, a); + }, + [=](const ck_tile::stream_config &s_) { + fmha_bwd_convert_dq_oneshot_(s_, a); + }); + } else { + if (s.log_level_ > 0) + std::cout << ", " << fmha_bwd_dot_do_o_get_name_() << "@" + << fmha_bwd_dq_dk_dv_get_name_() << std::flush; + return ck_tile::launch_kernel( + s, + [=](const ck_tile::stream_config &s_) { + fmha_bwd_dot_do_o_oneshot_(s_, a); + }, + [=](const ck_tile::stream_config &s_) { + fmha_bwd_dq_dk_dv_oneshot_(s_, a); + }); + } +} + template -float fmha_bwd(fmha_bwd_traits, fmha_bwd_args, const ck_tile::stream_config &); +float fmha_bwd(const fmha_bwd_traits &, fmha_bwd_args, + const ck_tile::stream_config &); + +struct fmha_bwd_launcher { + // POD closure for graph capture (trivially destructible, no heap allocation) + struct GraphClosure { + PrepareWorkspaceHostFunc func_ptr; + void *pin_w_ptr; + const ck_tile::index_t *seqstart_q_ptr; + const ck_tile::index_t *seqstart_k_ptr; + ck_tile::index_t batch; + ck_tile::index_t hdim_q; + ck_tile::index_t nhead_q; + ck_tile::index_t seqlen_q; + ck_tile::index_t seqlen_k; + + static void invoke(void *ud) { + auto *closure = static_cast(ud); + if (closure->func_ptr) { + // Callback runs on the HIP driver helper thread across a C ABI + // boundary; any exception escaping it would call std::terminate. + try { + closure->func_ptr(closure->pin_w_ptr, closure->batch, closure->hdim_q, + closure->nhead_q, closure->seqlen_q, + closure->seqlen_k, closure->seqstart_q_ptr, + closure->seqstart_k_ptr); + } catch (const std::exception &e) { + // The H2D queued after this callback will copy indeterminate + // metadata to device and the kernel will produce wrong results; + // unlikely in practice since pack_workspace_host only throws on + // precondition violations. + std::cerr << "fmha_bwd_launcher: pack_workspace_host threw: " + << e.what() << '\n'; + } catch (...) { + std::cerr << "fmha_bwd_launcher: pack_workspace_host threw unknown\n"; + } + } + } + }; + static_assert(std::is_trivially_destructible_v, + "GraphClosure must be trivially destructible for placement-new " + "without dtor"); + + std::function run{ + [](fmha_bwd_args, const ck_tile::stream_config &) { + std::cerr + << "fmha_bwd: no kernel found for given traits, skipping run\n"; + return -1.0f; + }}; + // Layout: [host_ws_size_ bytes (host-prepared metadata)][dq_acc region] + size_t workspace_size = 0; + + fmha_bwd_launcher(const fmha_bwd_traits &); + fmha_bwd_launcher(fmha_bwd_launcher &&) = delete; + fmha_bwd_launcher &operator=(fmha_bwd_launcher &&) = delete; + + ~fmha_bwd_launcher() noexcept { schedule_pin_staging_release(); } + + // Stream-async: zero dq_acc, D2H seqstart, host-pack metadata, H2D into + // device_ws. + // + // `pinned_host_alloc` returns a shared_ptr to a pinned host buffer. + // + // **Deleter behavior differs by mode**: + // - Normal mode: deleter is invoked on the stream tail after H2D completes. + // - Graph capture mode: deleter is NOT invoked; caller must keep buffer alive + // until hipGraphDestroy (see precondition #2 below). + // + // REQUIRED PRECONDITIONS for `pinned_host_alloc`: + // + // 1. **Capture-safe allocation**: The allocator must NOT call synchronizing + // APIs + // (e.g., bare hipHostMalloc) during active stream capture, as these + // invalidate the capture and cause hipStreamEndCapture to fail. Use a + // caching allocator that serves from cache during capture (e.g., PyTorch + // CachingHostAllocator). + // + // 2. **Buffer lifetime in graph mode**: In graph capture mode (detected via + // hipStreamIsCapturing), the returned buffer is NOT freed automatically + // after the H2D completes. The caller MUST keep the buffer alive for the + // graph's entire lifetime (until hipGraphDestroy), otherwise graph replay + // will read from freed memory. In normal (non-graph) mode, the buffer is + // automatically freed on the stream tail after the H2D. + // + void prepare_workspace_async( // + void *device_ws_ptr, const int *seqstart_q_dev, const int *seqstart_k_dev, + const ck_tile::stream_config &s, + const std::function(size_t)> &pinned_host_alloc) { + hipStream_t stream = s.stream_id_; + + // Fast path: no host-side metadata to stage; just zero dq_acc if needed. + if (host_ws_size_ == 0) { + if (needs_zero_dq_acc_ && workspace_size > 0) + HIP_CHECK_ERROR( + hipMemsetAsync(device_ws_ptr, 0, workspace_size, stream)); + return; + } + + if (!pinned_host_alloc) + throw std::runtime_error("fmha_bwd_launcher::prepare_workspace_async: " + "pinned_host_alloc is required"); + + // Allocate pinned host staging first: if it throws we haven't issued any + // stream work yet, leaving the workspace cleanly un-prepared. + // 16-align each section: pin_w stores alignas(16) + // FmhaBwdGroupPersistentCuState written via x86 SIMD; misaligned + // destinations fault. + const size_t seqstart_bytes = + traits_.is_group_mode ? sizeof(int) * (traits_.batch + 1) : 0; + const size_t seqstart_stride = ck_tile::integer_least_multiple( + seqstart_bytes, static_cast(16)); + const size_t pin_w_offset = 2 * seqstart_stride; + const size_t data_size = pin_w_offset + host_ws_size_; + + // Check if we're in graph capture mode + hipStreamCaptureStatus capture_status; + HIP_CHECK_ERROR(hipStreamIsCapturing(stream, &capture_status)); + const bool is_graph_capture = + (capture_status == hipStreamCaptureStatusActive); + + // Allocate pinned buffer with extra aligned space for closure + // Both modes use placement new for closure (POD, trivially destructible) + constexpr size_t closure_align = alignof(GraphClosure); + constexpr size_t closure_size = sizeof(GraphClosure); + const size_t aligned_data = + (data_size + closure_align - 1) & ~(closure_align - 1); + const size_t total_bytes = aligned_data + closure_size; + auto pin_base = pinned_host_alloc(total_bytes); + + if (needs_zero_dq_acc_ && workspace_size > host_ws_size_) + HIP_CHECK_ERROR( + hipMemsetAsync(static_cast(device_ws_ptr) + host_ws_size_, 0, + workspace_size - host_ws_size_, stream)); + + char *base = static_cast(pin_base.get()); + int *pin_q = reinterpret_cast(base); + int *pin_k = reinterpret_cast(base + seqstart_stride); + void *pin_w = base + pin_w_offset; + const ck_tile::index_t *seqstart_q_pinned = + traits_.is_group_mode + ? reinterpret_cast(pin_q) + : nullptr; + const ck_tile::index_t *seqstart_k_pinned = + traits_.is_group_mode + ? reinterpret_cast(pin_k) + : nullptr; + + if (traits_.is_group_mode) { + if (!seqstart_q_dev || !seqstart_k_dev) + throw std::runtime_error( + "fmha_bwd_launcher::prepare_workspace_async: " + "seqstart_q_dev and seqstart_k_dev are required in " + "group mode"); + HIP_CHECK_ERROR(hipMemcpyAsync(pin_q, seqstart_q_dev, seqstart_bytes, + hipMemcpyDeviceToHost, stream)); + HIP_CHECK_ERROR(hipMemcpyAsync(pin_k, seqstart_k_dev, seqstart_bytes, + hipMemcpyDeviceToHost, stream)); + } + + // === UNIFIED PATH: Both modes use placement new for POD closure === + // Construct closure in aligned location within pinned buffer (all POD, + // trivially destructible) + void *closure_addr = base + aligned_data; + auto *closure = new (closure_addr) GraphClosure{ + prepare_ws_func_, // Function pointer (points to code segment, always + // valid) + pin_w, seqstart_q_pinned, seqstart_k_pinned, + batch_, // Copy captured data (independent of launcher lifetime) + hdim_q_, nhead_q_, seqlen_q_, seqlen_k_}; + + HIP_CHECK_ERROR(hipLaunchHostFunc(stream, GraphClosure::invoke, closure)); + + HIP_CHECK_ERROR(hipMemcpyAsync(device_ws_ptr, pin_w, host_ws_size_, + hipMemcpyHostToDevice, stream)); + + if (!is_graph_capture) { + // Normal mode: transfer ownership of pinned buffer and schedule release. + // The H2D memcpy above copies host-prepared metadata to device workspace. + // Pinned buffer must stay alive until H2D completes, so we schedule its + // release on the stream tail (via hipLaunchHostFunc callback). + schedule_pin_staging_release(); // Release any previous in-flight buffer + // first + pin_staging_ = std::move(pin_base); + release_stream_ = stream; + } + // Graph mode: pin_base stays owned by caller (via shared_ptr ref-count). + // Caller must keep it alive for the graph's lifetime. Closure and data + // pointers within pin_base remain valid as long as caller holds pin_base. + } + +private: + fmha_bwd_traits traits_{}; + size_t host_ws_size_ = 0; + bool needs_zero_dq_acc_ = false; + + // Function pointer (points to code segment, survives launcher destruction) + PrepareWorkspaceHostFunc prepare_ws_func_ = nullptr; + + // Captured data (copied to closure in graph mode) + ck_tile::index_t batch_ = 0; + ck_tile::index_t hdim_q_ = 0; + ck_tile::index_t nhead_q_ = 0; + ck_tile::index_t seqlen_q_ = 0; + ck_tile::index_t seqlen_k_ = 0; + + std::shared_ptr pin_staging_; + hipStream_t release_stream_ = nullptr; + + // The pin_staging_ deleter MUST NOT call any HIP API: it fires from the + // hipLaunchHostFunc callback on the driver helper thread, which holds + // runtime locks (would deadlock against main-thread hipFree). PyTorch's + // CachingHostAllocator is safe; bare hipHostMalloc users should defer + // hipHostFree via ck_tile::pinned_host_releaser. + void schedule_pin_staging_release() noexcept { + if (!pin_staging_) + return; + auto *heap_ref = new std::shared_ptr(std::move(pin_staging_)); + const hipError_t err = hipLaunchHostFunc( + release_stream_, + [](void *ud) { delete static_cast *>(ud); }, + heap_ref); + if (err != hipSuccess) { + std::cerr << "fmha_bwd_launcher: hipLaunchHostFunc failed: " + << hipGetErrorString(err) << "; releasing eagerly\n"; + delete heap_ref; + } + } + + template + void init(const fmha_bwd_traits &t) { + traits_ = t; + run = [](fmha_bwd_args a, const ck_tile::stream_config &s) { + return fmha_bwd_(s, a); + }; + host_ws_size_ = fmha_bwd_dq_dk_dv_dq_ws_host_size_(t.batch); + size_t device_ws_size = 0; + if (host_ws_size_ > 0) { + // In group mode t.seqlen_q is already the padded total (== + // seqstart_q[batch]); in batch mode it's per-batch and the total is batch + // * seqlen_q. + const ck_tile::index_t total_seqlen_q_padded = + t.is_group_mode ? t.seqlen_q : t.batch * t.seqlen_q; + device_ws_size = fmha_bwd_dq_dk_dv_dq_ws_device_upper_bound_( + t.batch, t.hdim_q, t.nhead_q, total_seqlen_q_padded, t.max_seqlen_k); + + // Store function pointer (directly assign template-instantiated function) + prepare_ws_func_ = &fmha_bwd_dq_dk_dv_dq_prepare_ws_host_; + + // Store captured data as member variables + batch_ = t.batch; + hdim_q_ = t.hdim_q; + nhead_q_ = t.nhead_q; + seqlen_q_ = t.seqlen_q; + seqlen_k_ = t.seqlen_k; + } + workspace_size = host_ws_size_ + device_ws_size; + needs_zero_dq_acc_ = fmha_bwd_dq_dk_dv_needs_zero_dq_acc_(); + } + +public: + template float operator()(Args &&...args) const { + return run(std::forward(args)...); + } +}; diff --git a/Libraries/ComposableKernel/attention/fmha/fmha_bwd_runner.hpp b/Libraries/ComposableKernel/attention/fmha/fmha_bwd_runner.hpp index 0f56e5d5cd..7528d2b10f 100644 --- a/Libraries/ComposableKernel/attention/fmha/fmha_bwd_runner.hpp +++ b/Libraries/ComposableKernel/attention/fmha/fmha_bwd_runner.hpp @@ -1,35 +1,23 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once #include "ck_tile/host.hpp" -#include "ck_tile/utility/json_dump.hpp" #include "fmha_bwd.hpp" + #include "utils.hpp" +#include "ck_tile/utility/json_dump.hpp" + +#include "ck_tile/host/pinned_host_releaser.hpp" + #include +#include #include #include +#include +#include #include #include #include @@ -72,8 +60,23 @@ auto get_elimit(ck_tile::index_t hdim_q, ck_tile::index_t hdim_v) { return ck_tile::make_tuple(rtol, atol); } -extern template float fmha_bwd<2>(fmha_bwd_traits, fmha_bwd_args, - const ck_tile::stream_config &); +// Which sink values the runner hands to the kernel; only meaningful when +// sink_grad is on. +enum class sink_regime { + // Small values centered near zero. The default, and the only one that keeps + // d_sink dependent on which sink element was read: a sink far above rowmax + // saturates lse_new = log(exp(lse_old) + exp(sink)) and collapses d_sink to + // -sum_q D[q]. + sensitive, + // -inf on every other head, which both kernel and reference treat as "no + // sink". Pins the sink_ptr [nhead] contract, since only -inf yields a d_sink + // of exactly 0. + neutral_heads, + // Deliberately saturated. Numerically blind for the reason above, so it is + // not a correctness check; it keeps the P_sink -> 1 and lse-saturation + // arithmetic exercised, which the sensitive regime no longer reaches. + saturated, +}; template bwd_result @@ -85,10 +88,12 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::index_t hdim_q, ck_tile::index_t hdim_v, bool i_perm, bool o_perm, float scale, std::string bias_str, bool use_dbias, float p_drop, uint64_t drop_seed, uint64_t drop_offset, - bool drop_prefs, std::string mask_str, bool deterministic, - std::string init_method, uint32_t seed, int do_validation, - const ck_tile::stream_config &stream_config, - std::optional json = std::nullopt) { + bool drop_prefs, std::string mask_str, + bool sink_grad, // if true, compute and validate sink gradient + bool deterministic, std::string init_method, uint32_t seed, + int do_validation, const ck_tile::stream_config &stream_config, + std::optional json = std::nullopt, + sink_regime sink_values = sink_regime::sensitive) { const std::string data_type = []() { if constexpr (std::is_same_v) return "fp32"; @@ -249,12 +254,6 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, (mode == mode_enum::batch ? seqlen_qs[0] : seqstart_q_host.back()); const ck_tile::index_t shape_seqlen_k = (mode == mode_enum::batch ? seqlen_ks[0] : seqstart_k_host.back()); - // Keep it equal to or smaller than minimal bn0 of all tiles in fmha_bwd.py - // TODO: add API for requesting kN0/nsplits/workspace_size? It is not safe to - // rely on internal implementation details in client code. - const ck_tile::index_t kN0 = 16; - const ck_tile::index_t nsplits = - deterministic ? ck_tile::integer_divide_ceil(max_seqlen_k, kN0) : 1; ck_tile::HostTensor q_host( get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, hdim_q)); @@ -277,6 +276,25 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, get_lengths(o_perm, shape_batch, nhead, shape_seqlen_q, hdim_v)); ck_tile::HostTensor lse_host( std::array{shape_batch, nhead, shape_seqlen_q}); + ck_tile::HostTensor sink_host( + sink_grad ? std::array{nhead} + : std::array{ + 1} /* dummy when sink is disabled */); + if (sink_grad) { + // See the sink_regime comments: the sensitive range is deliberately small, + // because a sink far above rowmax collapses d_sink to -sum_q D[q] + // regardless of which sink element was read, which hid a [batch, nhead] vs + // [nhead] indexing mismatch in 12 of 15 SinkGrad cases, and also drives the + // token weights to 0, weakening the dK/dV checks. + const float lo = (sink_values == sink_regime::saturated) ? 30.0f : -1.0f; + const float hi = (sink_values == sink_regime::saturated) ? 60.0f : 1.0f; + std::uniform_real_distribution sink_dist(lo, hi); + sink_host.ForEach([&](auto &self, auto i) { + self(i) = (sink_values == sink_regime::neutral_heads && i[0] % 2 == 0) + ? -std::numeric_limits::infinity() + : static_cast(sink_dist(random_engine)); + }); + } ck_tile::HostTensor d_host( std::array{shape_batch, nhead, shape_seqlen_q}); ck_tile::HostTensor randval_host( @@ -296,11 +314,12 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, max_seqlen_k) : std::array{ 1, 1, 1, 1} /* dummy shape for simplifying code */); - ck_tile::HostTensor dq_acc_host( - i_perm ? std::array{nsplits, shape_batch, nhead, - shape_seqlen_q, hdim_q} - : std::array{nsplits, shape_batch, - shape_seqlen_q, nhead, hdim_q}); + ck_tile::HostTensor d_sink_host( + sink_grad ? std::array{nhead} + : std::array{0}); + if (sink_grad) { + d_sink_host.ForEach([&](auto &self, auto i) { self(i) = 0; }); + } if (init_method == "ui" || init_method == "0") { ck_tile::FillUniformDistributionIntegerValue{ @@ -354,12 +373,16 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::DeviceMem bias_buf(bias_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem o_buf(o_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem lse_buf(lse_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem sink_buf( + sink_grad ? sink_host.get_element_space_size_in_bytes() : 0); ck_tile::DeviceMem d_buf(d_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem randval_buf( randval_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem dq_buf(dq_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem dk_buf(dk_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem dv_buf(dv_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem d_sink_buf( + sink_grad ? d_sink_host.get_element_space_size_in_bytes() : 0); ck_tile::DeviceMem do_buf(do_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem dbias_buf(dbias_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem seqstart_q(seqstart_q_host.size() * sizeof(int32_t)); @@ -372,15 +395,70 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::DeviceMem drop_offset_buf(drop_prefs ? sizeof(uint64_t) : 0); ck_tile::DeviceMem alibi_slope_buf( alibi_slope_host.get_element_space_size_in_bytes()); - ck_tile::DeviceMem dq_acc_buf(dq_acc_host.get_element_space_size_in_bytes()); + const auto t0_launcher = std::chrono::high_resolution_clock::now(); + fmha_bwd_launcher launcher(fmha_bwd_traits{ + shape_seqlen_q, + shape_seqlen_k, + batch, + max_seqlen_q, + max_seqlen_k, + hdim_q, + hdim_v, + nhead, + nhead_k, + data_type, + mode == mode_enum::group, + mask.type, + bias.type, + use_dbias, + p_drop > 0.0f, + s_randval, + deterministic, + }); + const auto t1_launcher = std::chrono::high_resolution_clock::now(); + const double launcher_ctor_ms = + std::chrono::duration(t1_launcher - t0_launcher) + .count(); + const size_t ws_size = launcher.workspace_size; + ck_tile::DeviceMem ws_buf(ws_size); + + // Stage seqstart to device before prepare_workspace_async (which D2Hs it + // back). + seqstart_q.ToDevice(seqstart_q_host.data()); + seqstart_k.ToDevice(seqstart_k_host.data()); + + // Pinned host allocator for the launcher's async prepare pipeline. The + // shared_ptr deleter MUST NOT call any HIP API: it runs from the launcher's + // tail hipLaunchHostFunc on the driver helper thread, which holds HIP + // runtime locks. Deleter enqueues to a worker thread that hipHostFrees off + // the callback path. + auto pinned_host_alloc = [](size_t bytes) -> std::shared_ptr { + void *p = nullptr; + HIP_CHECK_ERROR(hipHostMalloc(&p, bytes, hipHostMallocDefault)); + return std::shared_ptr(p, [](void *q) { + ck_tile::pinned_host_releaser::instance().enqueue(q); + }); + }; + + ck_tile::gpu_timer prepare_ws_timer; + prepare_ws_timer.start(stream_config.stream_id_); + launcher.prepare_workspace_async( + ws_buf.GetDeviceBuffer(), + (mode == mode_enum::group) + ? static_cast(seqstart_q.GetDeviceBuffer()) + : nullptr, + (mode == mode_enum::group) + ? static_cast(seqstart_k.GetDeviceBuffer()) + : nullptr, + stream_config, pinned_host_alloc); + prepare_ws_timer.stop(stream_config.stream_id_); q_buf.ToDevice(q_host.data()); k_buf.ToDevice(k_host.data()); v_buf.ToDevice(v_host.data()); bias_buf.ToDevice(bias_host.data()); do_buf.ToDevice(do_host.data()); - seqstart_q.ToDevice(seqstart_q_host.data()); - seqstart_k.ToDevice(seqstart_k_host.data()); + // seqstart_q/k were already ToDevice'd above before prepare_workspace_async. if (mode == mode_enum::group) { std::vector seqlen_q_host(seqlen_qs.begin(), seqlen_qs.end()); seqlen_q_dev.ToDevice(seqlen_q_host.data()); @@ -390,6 +468,10 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, drop_seed_buf.ToDevice(drop_prefs ? &drop_seed : nullptr); drop_offset_buf.ToDevice(drop_prefs ? &drop_offset : nullptr); alibi_slope_buf.ToDevice(alibi_slope_host.data()); + if (sink_grad) { + sink_buf.ToDevice(sink_host.data()); + d_sink_buf.ToDevice(d_sink_host.data()); + } // clang-format off auto layout_str = [&](bool permute){ @@ -402,26 +484,28 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, }; // clang-format on - const std::size_t workspace_size_in_megabytes = ck_tile::integer_divide_ceil( - dq_acc_host.get_element_space_size_in_bytes(), 1024 * 1024); + const std::size_t workspace_size_in_megabytes = + ck_tile::integer_divide_ceil(ws_size, 1024 * 1024); std::cout << "[" << data_type << "|" << mode << "|" << io_layout(i_perm, o_perm) << "] b:" << batch << ", h:" << nhead << "/" << nhead_k << ", s:" << seqlen_qs[0] << "/" << seqlen_ks[0] << ", d:" << hdim_q << "/" << hdim_v << ", scale:" << scale << ", bias:" << bias << ", dbias:" << use_dbias - << ", p_drop:" << p_drop << ", s_randval:" << s_randval + << ", p_drop:" << p_drop + << (sink_grad ? (sink_values == sink_regime::neutral_heads + ? ", sink:(-inf every other head, grad)" + : (sink_values == sink_regime::saturated + ? ", sink:(rand[30,60] saturated, grad)" + : ", sink:(rand[-1,1], grad)")) + : "") + << ", s_randval:" << s_randval << ", deterministic:" << deterministic - << (deterministic - ? std::string(", workspace:") + - std::to_string(workspace_size_in_megabytes) + "MiB" - : "") - << ", mask:" << mask << std::flush; - - auto fmha_traits = fmha_bwd_traits{ - hdim_q, hdim_v, data_type, mode == mode_enum::group, - mask.type, bias.type, use_dbias, p_drop > 0.0f, - s_randval, deterministic}; + << ", workspace:" << std::to_string(workspace_size_in_megabytes) + << "MiB" << ", mask:" << mask << ", init:" << launcher_ctor_ms + << "ms" << ", prws:" << prepare_ws_timer.duration() << "ms" + << std::flush; + auto fmha_args = [&]() { /// NOTE: we broadcast bias from [1, 1, seqlen_q, seqlen_k] to [batch, /// nhead, seqlen_q, @@ -470,8 +554,8 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, const ck_tile::index_t batch_stride_dv = (nhead * shape_seqlen_k * hdim_v); const ck_tile::index_t batch_stride_dbias = (nhead * shape_seqlen_q * max_seqlen_k); - const ck_tile::index_t split_stride_dq_acc = - (shape_batch * nhead * shape_seqlen_q * hdim_q); + + void *ws_ptr = ws_size > 0 ? ws_buf.GetDeviceBuffer() : nullptr; const auto drop_seed_offset = [&]() -> decltype(fmha_bwd_args::drop_seed_offset) { @@ -487,7 +571,6 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, use_qpadding ? seqlen_q_dev.GetDeviceBuffer() : nullptr; const void *seqlen_k_ptr_dev = use_kpadding ? seqlen_k_dev.GetDeviceBuffer() : nullptr; - return fmha_bwd_args{ q_buf.GetDeviceBuffer(), k_buf.GetDeviceBuffer(), @@ -503,7 +586,9 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, dk_buf.GetDeviceBuffer(), dv_buf.GetDeviceBuffer(), dbias_buf.GetDeviceBuffer(), - dq_acc_buf.GetDeviceBuffer(), + ws_ptr, + sink_buf.GetDeviceBuffer(), + d_sink_buf.GetDeviceBuffer(), seqstart_q.GetDeviceBuffer(), seqstart_k.GetDeviceBuffer(), seqlen_q_ptr_dev, @@ -528,8 +613,7 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, stride_o, stride_randval, stride_do, - stride_q, // stride_dq_acc - stride_q, // stride_dq + stride_q, // stride_dq (same layout as q for dq output) stride_dk, stride_dv, stride_dbias, @@ -541,7 +625,6 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, nhead_stride_randval, nhead_stride_do, nhead_stride_lsed, - nhead_stride_q, // nhead_stride_dq_acc nhead_stride_q, // nhead_stride_dq nhead_stride_k, // nhead_stride_dk nhead_stride_v, // nhead_stride_dv @@ -554,12 +637,10 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, batch_stride_randval, batch_stride_do, batch_stride_lsed, - batch_stride_q, // batch_stride_dq_acc batch_stride_q, // batch_stride_dq batch_stride_dk, batch_stride_dv, batch_stride_dbias, - split_stride_dq_acc, mask.left, mask.right, static_cast(mask.type), @@ -568,7 +649,7 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, drop_seed_offset}; }(); - const float ave_time = fmha_bwd(fmha_traits, fmha_args, stream_config); + const float ave_time = launcher(fmha_args, stream_config); if (ave_time < 0) { std::cout << ", not supported yet" << std::flush << std::endl; return bwd_result::no_instance; @@ -593,6 +674,7 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, std::vector> randval_host_refs; std::vector> p_hp_host_refs; std::vector> p_lp_host_refs; + std::vector> p_sink_host_refs; randval_buf.FromDevice(randval_host.data()); @@ -638,8 +720,13 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::HostTensor p_dropped_hp_host_ref( {nhead, real_seqlen_q, real_seqlen_k}); // p_dropped_hp_g_m_n high precision + + // p_lp_g_m_n low precision used for fwd (with rp_undrop) + ck_tile::HostTensor p_fwd_host_ref( + {nhead, real_seqlen_q, real_seqlen_k}); + // p_lp_g_m_n low precision used for bwd (no rp_undrop) ck_tile::HostTensor p_lp_host_ref( - {nhead, real_seqlen_q, real_seqlen_k}); // p_lp_g_m_n low precision + {nhead, real_seqlen_q, real_seqlen_k}); ck_tile::index_t nr = nhead / nhead_k; @@ -748,15 +835,65 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::reference_batched_softmax( s_host_ref, p_hp_host_ref, ck_tile::identity{}, lse_host_ref); + // Incorporate sink token into the softmax distribution (reference + // computation). The sink acts as an extra key whose score is + // sink_host(i_h) (in log-space), which is a per-head random value in + // [-1., 1.]. + // lse_new = log(exp(lse_old) + exp(sink)) + // P_new = P_old * exp(lse_old - lse_new) (rescaled token attention) + // P_sink = exp(sink - lse_new) (sink attention weight) + ck_tile::HostTensor p_sink_host_ref( + sink_grad ? std::array{nhead, real_seqlen_q} + : std::array{0, 0}); + if (sink_grad) { + for (int i_h = 0; i_h < nhead; ++i_h) { + AccDataType sink_val = sink_host(i_h); + for (int i_q = 0; i_q < real_seqlen_q; ++i_q) { + // Use numerically stable log-sum-exp: lse_new = + // log(exp(lse_old)+exp(sink)) + // = max(lse_old, sink) + log(1 + exp(min - max)) + // This handles lse_old = -inf (fully-masked rows) without producing + // NaN: + // if lse_old=-inf: max=sink, min=-inf, exp(-inf-sink)=0, + // lse_new=sink + // It also avoids exp(lse_old) overflow when lse_old is large. + // p_scale = exp(lse_old - lse_new) [fraction kept by regular + // tokens] p_sink = exp(sink - lse_new) [sink attention + // weight] + AccDataType lse_old = lse_host_ref(i_h, i_q); + AccDataType hi = lse_old > sink_val ? lse_old : sink_val; + AccDataType lo = lse_old > sink_val ? sink_val : lse_old; + // Both -inf (a fully masked row on a head whose sink is disabled) + // would make lo - hi a NaN, so short-circuit it: no sink and no + // token means lse stays -inf. + AccDataType lse_new = + (std::isinf(hi) && hi < AccDataType(0)) + ? hi + : hi + ck_tile::log(AccDataType(1) + ck_tile::exp(lo - hi)); + AccDataType p_scale = ck_tile::exp(lse_old - lse_new); + + lse_host_ref(i_h, i_q) = lse_new; + + for (int i_k = 0; i_k < real_seqlen_k; ++i_k) + p_hp_host_ref(i_h, i_q, i_k) *= p_scale; + + p_sink_host_ref(i_h, i_q) = ck_tile::exp(sink_val - lse_new); + } + } + } + if (p_drop > 0) { p_dropped_hp_host_ref = p_hp_host_ref; ck_tile::reference_batched_dropout_randval(randval_host_ref, wb, drop_seed, drop_offset); - ck_tile::reference_batched_dropout(p_dropped_hp_host_ref, - randval_host_ref, - p_undrop_in_uint8_t, rp_undrop); + ck_tile::reference_batched_dropout( + p_dropped_hp_host_ref, randval_host_ref, p_undrop_in_uint8_t, 1.f); p_lp_host_ref = p_dropped_hp_host_ref.template CopyAsType(); + p_dropped_hp_host_ref.ForEach( + [&](auto &self, const auto &idx) { self(idx) *= rp_undrop; }); + p_fwd_host_ref = + p_dropped_hp_host_ref.template CopyAsType(); ck_tile::HostTensor randval_host_result( {nhead, real_seqlen_q, real_seqlen_k}); @@ -779,12 +916,13 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, } } else { p_lp_host_ref = p_hp_host_ref.template CopyAsType(); + p_fwd_host_ref = p_lp_host_ref; } // O = P * V ck_tile::reference_batched_gemm( - p_lp_host_ref, v_host_ref, + p_fwd_host_ref, v_host_ref, o_host_ref); // o_g_m_o = p_lp_g_m_n@v_g_o_n // clang-format off @@ -801,6 +939,7 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, o_host_refs.push_back(o_host_ref); p_hp_host_refs.push_back(p_hp_host_ref); p_lp_host_refs.push_back(p_lp_host_ref); + p_sink_host_refs.push_back(p_sink_host_ref); if (p_drop > 0) { randval_host_refs.push_back(randval_host_ref); } @@ -813,35 +952,47 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::numeric::infinity()}(dk_host); ck_tile::FillConstant{ ck_tile::numeric::infinity()}(dv_host); - ck_tile::FillConstant{ - ck_tile::numeric::infinity()}(dq_acc_host); dq_buf.ToDevice(dq_host.data()); dk_buf.ToDevice(dk_host.data()); dv_buf.ToDevice(dv_host.data()); - dq_acc_buf.ToDevice(dq_acc_host.data()); o_buf.ToDevice(o_host.data()); lse_buf.ToDevice(lse_host.data()); dbias_buf.SetZero(); - - // non-deterministic kernels use atomic add to write dq - // Some block may be skipped with causal mask and dq are not set to zeros - // In these cases thus we need to zero out it first - if (!deterministic || mask.type != mask_enum::no_mask) - dq_acc_buf.SetZero(); + if (sink_grad) + d_sink_buf.SetZero(); ck_tile::stream_config stream_config_v{nullptr, true, 0, 0, 1}; - fmha_bwd(fmha_traits, fmha_args, stream_config_v); + // re-initialize workspace for validation run + launcher.prepare_workspace_async( + ws_buf.GetDeviceBuffer(), + (mode == mode_enum::group) + ? static_cast(seqstart_q.GetDeviceBuffer()) + : nullptr, + (mode == mode_enum::group) + ? static_cast(seqstart_k.GetDeviceBuffer()) + : nullptr, + stream_config_v, pinned_host_alloc); + launcher(fmha_args, stream_config_v); dq_buf.FromDevice(dq_host.data()); dk_buf.FromDevice(dk_host.data()); dv_buf.FromDevice(dv_host.data()); dbias_buf.FromDevice(dbias_host.data()); + if (sink_grad) + d_sink_buf.FromDevice(d_sink_host.data()); // Track the index into reference vectors (may differ from wb if batches // were skipped) ck_tile::index_t ref_idx = 0; + // validation sink accumulator: global over batch, shape [nhead] + ck_tile::HostTensor d_sink_host_ref( + sink_grad ? std::array{nhead} + : std::array{0}); + if (sink_grad) + d_sink_host_ref.ForEach([&](auto &self, auto i) { self(i) = 0; }); + for (ck_tile::index_t wb = 0; wb < batch; ++wb) { // When padding is enabled, use logical lengths instead of computing from // padded prefix-sum @@ -898,7 +1049,7 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, if (p_drop > 0) { ck_tile::reference_batched_dropout(dp_hp_host_ref, randval_host_refs[ref_idx], - p_undrop_in_uint8_t, rp_undrop); + p_undrop_in_uint8_t, 1.f); } // dS_i_j = P_i_j .* (dP_i_j - dO_i dot O_i) @@ -909,7 +1060,8 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, do_dot_o += ck_tile::type_convert(do_host_ref(i0, i1, o)) * ck_tile::type_convert( - o_host_refs[ref_idx](i0, i1, o)); + o_host_refs[ref_idx](i0, i1, o)) * + p_undrop; } ds_hp_host_ref(i0, i1, i2) = ck_tile::type_convert( p_hp_host_refs[ref_idx](i0, i1, i2) * @@ -920,6 +1072,26 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ds_hp_host_ref.mDesc.get_lengths()[2])( std::thread::hardware_concurrency()); + if (sink_grad) { + // Reference: dSink[h] = -sum_q( P_sink[h,q] * D[h,q] ) + // where D[h,q] = sum_j(dO[h,q,j] * O[h,q,j]) * p_undrop + for (int i_h = 0; i_h < nhead; ++i_h) { + AccDataType d_sink_head_acc = 0; + for (int i_q = 0; i_q < real_seqlen_q; ++i_q) { + AccDataType do_dot_o = 0; + for (int o = 0; o < hdim_v; o++) { + do_dot_o += + ck_tile::type_convert(do_host_ref(i_h, i_q, o)) * + ck_tile::type_convert( + o_host_refs[ref_idx](i_h, i_q, o)) * + p_undrop; + } + d_sink_head_acc += -p_sink_host_refs[ref_idx](i_h, i_q) * do_dot_o; + } + d_sink_host_ref(i_h) += d_sink_head_acc; + } + } + if (use_dbias) { dbias_host_ref = ds_hp_host_ref.template CopyAsType(); } @@ -934,8 +1106,9 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, do_host_ref.transpose({0, 2, 1}); // do_g_m_o -> do_g_o_m ck_tile::reference_batched_gemm( - p_t_lp_host_ref, do_t_host_ref, - dv_host_ref); // dv_g_n_o = p_lp_g_n_m@do_g_o_m + p_t_lp_host_ref, do_t_host_ref, dv_host_ref, ck_tile::identity{}, + ck_tile::identity{}, + ck_tile::scales(rp_undrop)); // dv_g_n_o = p_lp_g_n_m@do_g_o_m // dQ = scale * dS@K^T auto k_t_host_ref = @@ -944,7 +1117,7 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, QGradDataType>( ds_lp_host_ref, k_t_host_ref, dq_host_ref, ck_tile::identity{}, ck_tile::identity{}, - ck_tile::scales(scale)); // dq_g_m_k = ds_g_m_n@k_g_k_n + ck_tile::scales(scale * rp_undrop)); // dq_g_m_k = ds_g_m_n@k_g_k_n // dK = scale * dS^T@Q^T auto ds_t_lp_host_ref = @@ -955,7 +1128,7 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, KGradDataType>( ds_t_lp_host_ref, q_t_host_ref, dk_host_ref, ck_tile::identity{}, ck_tile::identity{}, - ck_tile::scales(scale)); // dk_g_n_k = ds_g_n_m@q_g_k_m + ck_tile::scales(scale * rp_undrop)); // dk_g_n_k = ds_g_n_m@q_g_k_m ck_tile::HostTensor dq_host_result( {nhead, real_seqlen_q, hdim_q}); // dq_g_m_k @@ -1016,6 +1189,43 @@ fmha_bwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ref_idx++; } + if (pass && sink_grad) { + auto [rtol, atol] = get_elimit(hdim_q, hdim_v); + bool dsink_pass = ck_tile::check_err( + d_sink_host, d_sink_host_ref, + std::string("Error: SinkGrad Incorrect results!"), rtol, atol); + pass &= dsink_pass; + } + + if (sink_grad && sink_values == sink_regime::neutral_heads) { + // sink_ptr is [nhead], shared by every batch. Even heads were given -inf, + // so they must contribute nothing and land on exactly 0, while odd heads + // must be non-zero. Only -inf yields exactly 0, so a shifted index - or a + // read past the end of the buffer - breaks the pattern instead of merely + // perturbing a number, which is what makes this independent of both + // tolerances and the init range. + bool identity_pass = true; + for (ck_tile::index_t h = 0; h < nhead; ++h) { + const bool neutral = (h % 2 == 0); + const double gpu = ck_tile::type_convert(d_sink_host(h)); + + if (neutral && std::abs(gpu) > 0) { + identity_pass = false; + std::cerr << "Error: head " << h + << " has sink -inf so d_sink must be exactly " << "0, got " + << gpu << ". The kernel read some other sink element\n"; + } + if (!neutral && !(std::abs(gpu) > 0)) { + identity_pass = false; + std::cerr + << "Error: head " << h << " has a finite sink so d_sink must be " + << "non-zero, got 0. The kernel read a disabled head's sink\n"; + } + } + + pass &= identity_pass; + } + std::cout << ", valid:" << (pass ? "y" : "n") << std::flush << std::endl; } diff --git a/Libraries/ComposableKernel/attention/fmha/fmha_fwd.cpp b/Libraries/ComposableKernel/attention/fmha/fmha_fwd.cpp index 8fe2a6efe2..1cc1a643e0 100644 --- a/Libraries/ComposableKernel/attention/fmha/fmha_fwd.cpp +++ b/Libraries/ComposableKernel/attention/fmha/fmha_fwd.cpp @@ -1,27 +1,9 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT -#include "fmha_fwd.hpp" #include "ck_tile/host.hpp" + +#include "fmha_fwd.hpp" #include "fmha_fwd_runner.hpp" #include @@ -67,17 +49,16 @@ auto create_args(int argc, char *argv[]) { .insert("d", "128", "head dim for q, k") .insert("d_v", "-1", "head dim for v, -1 means equal to d") .insert("scale_s", "0", - "scale factor of S. 0 means equal to 1/sqrt(hdim).\n" - "note when squant=1, this value will be modified") + "scale factor of S. 0 means equal to 1/sqrt(hdim)") + .insert("qscale", "n", + "quant scale:\n" + " n or 0, no scale\n" + " pt or 1, per-tensor scale\n" + " bs or 2, block scale\n" + " kvbs or 3, Q per-tensor, K/V per-page block scale\n" + " mx or 4, microscaling (exclusively for data types like mxfp8 " + "and mxfp4)") .insert("logits_soft_cap", "0", "attention logits soft capping value.") - .insert("squant", "auto", - "if using static quantization fusion or not. auto: fp8 will " - "default use squant, " - "other will not\n" - "0: no static quant(not implemented) 1: apply scale_p and " - "scale_o with respect to " - "P and O.\n" - "calculate scale_s, scale_p, scale_o auto") .insert("iperm", "1", "permute input\n" "if true, will be b*h*s*d, else b*s*h*d") @@ -87,7 +68,8 @@ auto create_args(int argc, char *argv[]) { "e(lementwise) or 1, elementwise bias with 1*1*s*s. e:1, " "1*h*s*s. e:2, b*h*s*s\n" "a(libi) or 2, alibi with 1*h. a:1, b*h") - .insert("prec", "fp16", "data type. fp32/fp16/bf16/fp8/bf8") + .insert("prec", "fp16", + "data type: fp32/fp16/bf16/fp8/fp8bf16/fp8fp32/mxfp8/mxfp4") .insert( "mask", "0", "0: no mask, 1: top-left(same as 't'), 2:bottom-right(same as 'b')\n" @@ -114,7 +96,9 @@ auto create_args(int argc, char *argv[]) { "init method:\n ui or 0 - uniform random int\n ni - normalized " "random int" "\n uf or 1 - uniform random float\n nf - normalized random float" - "\n tf or 2 - trig float\n") + "\n tf or 2 - trig float" + "\n tf or 3 - uniform random float, min max is the max of the " + "type\n") .insert("seed", "11939", "random seed used for initializing input tensors. 0 for " "non-deterministic seed") @@ -146,7 +130,12 @@ auto create_args(int argc, char *argv[]) { .insert( "kv_eff_lens", "", "Batch-mode only: per-batch effective seqlen for KV (exclude PAD).\n" - "Comma-separated list of length 'b'. If empty, no override."); + "Comma-separated list of length 'b'. If empty, no override.") + .insert("init_sink", "0", + "value to init the output tensor sink value for validation") + .insert( + "pack_gqa", "1", + "1: enable Pack-GQA (fold GQA Q heads into seqlen for non-causal)"); bool result = arg_parser.parse(argc, argv); return std::make_tuple(result, arg_parser); @@ -178,6 +167,7 @@ auto run(const ck_tile::ArgParser &arg_parser) { ck_tile::index_t page_block_size = arg_parser.get_int("page_block_size"); bool use_cache_batch_idx = arg_parser.get_bool("cache_batch_idx"); std::string bias_str = arg_parser.get_str("bias"); + std::string qscale_str = arg_parser.get_str("qscale"); float p_drop = arg_parser.get_float("p_drop"); uint64_t drop_seed = arg_parser.get_uint64("drop_seed"); uint64_t drop_offset = arg_parser.get_uint64("drop_offset"); @@ -187,13 +177,8 @@ auto run(const ck_tile::ArgParser &arg_parser) { ck_tile::index_t num_splits = arg_parser.get_int("num_splits"); std::string init_method = arg_parser.get_str("init"); uint32_t seed = arg_parser.get_uint32("seed"); - - bool squant = [&]() { - if (arg_parser.get_str("squant") == "auto") - return std::is_same_v; - else - return arg_parser.get_bool("squant"); - }(); + int init_sink_value = arg_parser.get_int("init_sink"); + int pack_gqa = arg_parser.get_int("pack_gqa"); ck_tile::stream_config stream_config{ nullptr, @@ -212,9 +197,9 @@ auto run(const ck_tile::ArgParser &arg_parser) { seqlen_knew, seqlen_qpads, seqlen_kpads, q_eff_lens_per_batch, kv_eff_lens_per_batch, rotary_dim, i_perm, o_perm, scale_s, logits_soft_cap, is_v_rowmajor, lse, page_block_size, use_cache_batch_idx, - bias_str, p_drop, drop_seed, drop_offset, drop_prefs, mask_str, squant, - is_rotary_interleaved, num_splits, init_method, seed, do_validation, - stream_config, json); + bias_str, p_drop, drop_seed, drop_offset, drop_prefs, mask_str, + qscale_str, is_rotary_interleaved, num_splits, init_method, seed, + do_validation, init_sink_value, pack_gqa, stream_config, json); } int main(int argc, char *argv[]) { @@ -236,6 +221,10 @@ int main(int argc, char *argv[]) { return run(arg_parser) == fwd_result::success ? 0 : -2; } else if (data_type == "fp8fp32") { return run(arg_parser) == fwd_result::success ? 0 : -2; + } else if (data_type == "mxfp8") { + return run(arg_parser) == fwd_result::success ? 0 : -2; + } else if (data_type == "mxfp4") { + return run(arg_parser) == fwd_result::success ? 0 : -2; } std::cerr << "Unsupported precision: " << data_type << std::endl; return -1; diff --git a/Libraries/ComposableKernel/attention/fmha/fmha_fwd.hpp b/Libraries/ComposableKernel/attention/fmha/fmha_fwd.hpp index daae03c5d8..8cdca66bee 100644 --- a/Libraries/ComposableKernel/attention/fmha/fmha_fwd.hpp +++ b/Libraries/ComposableKernel/attention/fmha/fmha_fwd.hpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once @@ -30,6 +11,7 @@ #include "bias.hpp" #include "mask.hpp" +#include "quant.hpp" #include "rotary.hpp" #include @@ -52,6 +34,10 @@ struct FmhaFwdFp8Bf16 {}; struct FmhaFwdFp8Fp32 {}; +struct FmhaFwdMxFp8 {}; + +struct FmhaFwdMxFp4 {}; + template struct FmhaFwdTypeConfig; template <> struct FmhaFwdTypeConfig { @@ -159,6 +145,55 @@ template <> struct FmhaFwdTypeConfig { using ODataType = float; }; +template <> struct FmhaFwdTypeConfig { + using QDataType = ck_tile::fp8_t; + using KDataType = ck_tile::fp8_t; + using VDataType = ck_tile::fp8_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = + float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = float; + + using QScaleDataType = ck_tile::e8m0_t; + using KScaleDataType = ck_tile::e8m0_t; + using VScaleDataType = ck_tile::e8m0_t; + using PScaleDataType = ck_tile::e8m0_t; + + static constexpr ck_tile::index_t kQKScaleGranularity = 32; + static constexpr ck_tile::index_t kVScaleGranularity = 32; +}; + +template <> struct FmhaFwdTypeConfig { + using QDataType = ck_tile::pk_fp4_t; + using KDataType = ck_tile::pk_fp4_t; + using VDataType = ck_tile::pk_fp4_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = + float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = + ck_tile::pk_fp6x16_t; // data type for A matrix of second gemm (pk_fp4_t + // or pk_fp6x16_t, use pk_fp6x16_t for higher + // precision for free) + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = float; + + using QScaleDataType = ck_tile::e8m0_t; + using KScaleDataType = ck_tile::e8m0_t; + using VScaleDataType = ck_tile::e8m0_t; + using PScaleDataType = ck_tile::e8m0_t; + + static constexpr ck_tile::index_t kQKScaleGranularity = 32; + static constexpr ck_tile::index_t kVScaleGranularity = 32; +}; + struct FmhaMasks { using NoMask = ck_tile::GenericAttentionMask; using GenericMask = ck_tile::GenericAttentionMask; @@ -168,10 +203,15 @@ struct FmhaMasks { // runtime args, some will passed to karg, some will used to compute // grids/blocks struct fmha_fwd_args { + std::string *selected_kernel_name = nullptr; + const void *q_ptr; const void *k_ptr; const void *v_ptr; const void *bias_ptr; // bias or alibi_slope pointer + const void *q_descale_ptr; + const void *k_descale_ptr; + const void *v_descale_ptr; void *rand_val_ptr; void *lse_ptr; void *o_ptr; @@ -228,6 +268,10 @@ struct fmha_fwd_args { const void *cu_seqlen_k_ptr = nullptr; // Cumulative logical (excluding padding) sequence length // array [batch + 1]. (Used with padding) + const void *block_scale_seqstart_q_ptr; + const void *block_scale_seqstart_k_ptr; + const void *seqstart_v_scale_ptr; + const void *sink_ptr; ck_tile::index_t seqlen_q; ck_tile::index_t seqlen_k; @@ -237,11 +281,10 @@ struct fmha_fwd_args { ck_tile::index_t hdim_v; ck_tile::index_t nhead_q; ck_tile::index_t nhead_k; + ck_tile::index_t num_head_q_total = 0; + ck_tile::index_t head_start = 0; float scale_s; - float scale_p; - float scale_o; - float logits_soft_cap; ck_tile::index_t stride_q; @@ -251,6 +294,9 @@ struct fmha_fwd_args { stride_bias; // if alibi, b*h need set this to h, 1*h need set this to 0 ck_tile::index_t stride_randval; ck_tile::index_t stride_o; + ck_tile::index_t stride_q_descale; + ck_tile::index_t stride_k_descale; + ck_tile::index_t stride_v_descale; ck_tile::index_t nhead_stride_q; ck_tile::index_t nhead_stride_k; ck_tile::index_t nhead_stride_v; @@ -258,6 +304,9 @@ struct fmha_fwd_args { ck_tile::index_t nhead_stride_randval; ck_tile::index_t nhead_stride_lse; ck_tile::index_t nhead_stride_o; + ck_tile::index_t nhead_stride_q_descale; + ck_tile::index_t nhead_stride_k_descale; + ck_tile::index_t nhead_stride_v_descale; ck_tile::index_t batch_stride_q; ck_tile::index_t batch_stride_k; ck_tile::index_t batch_stride_v; @@ -265,9 +314,13 @@ struct fmha_fwd_args { ck_tile::index_t batch_stride_randval; ck_tile::index_t batch_stride_lse; ck_tile::index_t batch_stride_o; + ck_tile::index_t batch_stride_q_descale; + ck_tile::index_t batch_stride_k_descale; + ck_tile::index_t batch_stride_v_descale; ck_tile::index_t window_size_left; ck_tile::index_t window_size_right; + ck_tile::index_t sink_size; ck_tile::index_t mask_type; ck_tile::index_t min_seqlen_q; @@ -277,6 +330,9 @@ struct fmha_fwd_args { std::variant, std::pair> drop_seed_offset; + + ck_tile::index_t block_scale_size_q; + ck_tile::index_t block_scale_size_kv; }; struct fmha_fwd_pagedkv_args { @@ -324,6 +380,7 @@ struct fmha_fwd_pagedkv_args { const void *seqstart_q_ptr; const void *seqstart_k_ptr; const void *seqlen_k_ptr; + const void *sink_ptr; ck_tile::index_t seqlen_q; ck_tile::index_t seqlen_k; @@ -361,6 +418,7 @@ struct fmha_fwd_pagedkv_args { ck_tile::index_t window_size_left; ck_tile::index_t window_size_right; + ck_tile::index_t sink_size; ck_tile::index_t mask_type; ck_tile::index_t min_seqlen_q; }; @@ -412,6 +470,7 @@ struct fmha_fwd_splitkv_args { const void *seqstart_q_ptr; const void *seqstart_k_ptr; const void *seqlen_k_ptr; + const void *sink_ptr; ck_tile::index_t seqlen_q; ck_tile::index_t seqlen_k; @@ -457,6 +516,7 @@ struct fmha_fwd_splitkv_args { ck_tile::index_t window_size_left; ck_tile::index_t window_size_right; + ck_tile::index_t sink_size; ck_tile::index_t mask_type; }; @@ -490,6 +550,7 @@ struct fmha_fwd_appendkv_args { const void *cache_batch_idx; // only used if block_table_ptr is nullptr -> // batch mode (kvcache) + const void *sink_ptr; ck_tile::index_t stride_q; ck_tile::index_t stride_k; @@ -513,6 +574,9 @@ struct fmha_batch_prefill_args { const void *k_ptr; const void *v_ptr; const void *bias_ptr; // bias or alibi_slope pointer + const void *q_descale_ptr; + const void *k_descale_ptr; + const void *v_descale_ptr; void *rand_val_ptr; void *lse_ptr; void *o_ptr; @@ -529,6 +593,7 @@ struct fmha_batch_prefill_args { // (kargs.kv_indptr[b + 1] - kargs.kv_indptr[b] - 1) + // kargs.kv_last_page_lens[b] const void *seqstart_q_ptr; + const void *sink_ptr; ck_tile::index_t seqlen_q; ck_tile::index_t seqlen_k; @@ -539,14 +604,26 @@ struct fmha_batch_prefill_args { ck_tile::index_t nhead_q; ck_tile::index_t nhead_k; - // SGLang-style page table - int32_t num_total_pages; - void *kv_indptr; - void *kv_page_indices; -#if 0 // we assume page_block_size=1 for now - void* kv_last_page_lens; - ck_tile::index_t page_block_size; -#endif + // KV cache page table fields (kv_lookup_table selects interpretation): + // - SGLANG_PAGE_TABLE_1D: + // kv_indptr: prefix-sum [batch+1] into kv_page_indices + // kv_page_indices: 1D list of physical page ids, length = num_total_pages + // kv_last_page_lens: per-batch last page lengths [batch] + // - VLLM_BLOCK_TABLE_2D: + // kv_page_indices: block_table [batch, max_blocks_per_seq] (2D) + // batch_stride_block_table: row stride for block_table + // seqlen_k_ptr: per-batch seqlen_k [batch] + int32_t num_total_pages; // total physical pages in KV cache (SGLang/vLLM) + ck_tile::index_t page_block_size; // tokens per page (SGLang/vLLM) + ck_tile::BlockAttentionKVCacheMemoryLayoutEnum + kv_memory_layout; // KV memory layout (SGLang/vLLM) + ck_tile::BlockAttentionKVCacheLookupTableEnum + kv_lookup_table; // lookup table layout selector + void *kv_indptr; // SGLang: prefix-sum; vLLM: unused + void *kv_page_indices; // SGLang: 1D page list; vLLM: block_table 2D + void *kv_last_page_lens; // SGLang: last page lengths; vLLM: unused + void *seqlen_k_ptr; // vLLM: per-batch seqlen_k; SGLang: unused + ck_tile::index_t batch_stride_block_table; // vLLM: row stride; SGLang: unused float scale_s; float scale_p; @@ -578,6 +655,7 @@ struct fmha_batch_prefill_args { ck_tile::index_t window_size_left; ck_tile::index_t window_size_right; + ck_tile::index_t sink_size; ck_tile::index_t mask_type; float p_drop; @@ -586,8 +664,79 @@ struct fmha_batch_prefill_args { std::variant, std::pair> drop_seed_offset; + + // KV_BLOCKSCALE: per-page K/V descales (Q per-tensor, K/V per-page) + // k_descale_ptr/v_descale_ptr are reused for KV_BLOCKSCALE mode: + // k_descale_ptr: [num_block, num_kv_head] - points to k block descale + // v_descale_ptr: [num_block, num_kv_head] - points to v block descale + ck_tile::index_t nblock_stride_kv_block_descale = + 0; // Stride along num_block dimension + ck_tile::index_t nhead_stride_kv_block_descale = + 0; // Stride along num_kv_head dimension }; +// Selects the KV-cache load mode for a batch-prefill dispatch arm. +// GLOBAL_LOAD_LDS: required when (a) the page is smaller than one K/V tile +// so per-page SRD is impossible, AND (b) the SRD voffset arithmetic would +// overflow for either K or V. The hardware computes +// addr = (base[63:32] << 32) | ((base[31:0] + voffset) & 0xFFFFFFFF) +// and the voffset chain in ck_tile is signed int32 (index_t), so it wraps +// once base[31:0] + max_voffset exceeds INT32_MAX (0x7FFFFFFF, ~2GB) - NOT +// 0xFFFFFFFF - even when the KV pool itself is well under 2GB. K and V are +// independently allocated, so each is checked separately; either one +// crossing the bound forces GLOBAL_LOAD_LDS. +// BUFFER_LOAD: every other case - the SGPR-resident SRD path is fastest. +// k_base_ptr/v_base_ptr are the VAs of the first elements of the K and V +// caches; the low 32 bits of each determine whether the voffset addition wraps. +inline ck_tile::BlockAttentionKVCacheLoadModeEnum +fmha_batch_prefill_select_kv_load_mode( + ck_tile::index_t page_block_size, ck_tile::index_t kN0, + ck_tile::index_t num_total_pages, ck_tile::index_t batch_stride_k, + ck_tile::index_t batch_stride_v, ck_tile::index_t element_bytes, + const void *k_base_ptr, const void *v_base_ptr) { + if (page_block_size >= kN0) + return ck_tile::BlockAttentionKVCacheLoadModeEnum::BUFFER_LOAD; + + // Scattered 1D paged KV (page_size=1 LINEAR + SGLANG_PAGE_TABLE_1D): the page + // table holds arbitrary physical-page indices into the whole KV pool, so the + // per-page SRD voffset (physical_page * stride_page_block + within_page) is + // not bounded by this dispatch's num_total_pages and can wrap the signed + // int32 on its own, independent of base[31:0] + pool_bytes. The address-size + // check below cannot see this, so force the 64-bit-safe path for single-token + // pages. + if (page_block_size == 1) + return ck_tile::BlockAttentionKVCacheLoadModeEnum::GLOBAL_LOAD_LDS; + + // Maximum byte offsets that buffer_load will add to K/V base pointers. + // Each page is addressed as page_id * batch_stride * element_bytes. + const auto k_pool_bytes = static_cast(num_total_pages) * + static_cast(batch_stride_k) * + static_cast(element_bytes); + const auto v_pool_bytes = static_cast(num_total_pages) * + static_cast(batch_stride_v) * + static_cast(element_bytes); + + // Low 32 bits of each base VA. K and V are independently allocated and may + // have different low-32 values, so both are checked. Compare against + // INT32_MAX (not UINT32_MAX): the low32 bits themselves are unsigned, but the + // voffset they are added to is signed int32, so the effective address is + // already wrong once base[31:0] + pool crosses 0x7FFFFFFF. A sum in the + // (2GB, 4GB] range still fits in 32 unsigned bits but has passed the signed + // wrap point - that is exactly the band this check must route to global load. + const auto k_lo32 = + static_cast(reinterpret_cast(k_base_ptr)) & + 0xFFFFFFFFULL; + const auto v_lo32 = + static_cast(reinterpret_cast(v_base_ptr)) & + 0xFFFFFFFFULL; + const bool srd_would_overflow = (k_lo32 + k_pool_bytes) > INT32_MAX || + (v_lo32 + v_pool_bytes) > INT32_MAX; + + return srd_would_overflow + ? ck_tile::BlockAttentionKVCacheLoadModeEnum::GLOBAL_LOAD_LDS + : ck_tile::BlockAttentionKVCacheLoadModeEnum::BUFFER_LOAD; +} + template auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) { assert(args.nhead_q % args.nhead_k == 0); @@ -595,33 +744,47 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) { // create group mode kernel arguments if constexpr (FmhaKernel::kIsGroupMode) { return FmhaKernel::MakeKargsImpl( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.rand_val_ptr, + args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.q_descale_ptr, + args.k_descale_ptr, args.v_descale_ptr, args.rand_val_ptr, args.lse_ptr, args.o_ptr, args.seqstart_q_ptr, args.seqstart_k_ptr, - args.seqlen_q_ptr, args.seqlen_k_ptr, args.hdim_q, args.hdim_v, - args.nhead_q, args.nhead_q / args.nhead_k, args.scale_s, args.scale_p, - args.scale_o, args.logits_soft_cap, args.stride_q, args.stride_k, + args.seqlen_q_ptr, args.seqlen_k_ptr, args.block_scale_seqstart_q_ptr, + args.block_scale_seqstart_k_ptr, args.seqstart_v_scale_ptr, + args.hdim_q, args.hdim_v, args.nhead_q, args.nhead_q / args.nhead_k, + args.scale_s, args.logits_soft_cap, args.stride_q, args.stride_k, args.stride_v, args.stride_bias, args.stride_randval, args.stride_o, + args.stride_q_descale, args.stride_k_descale, args.stride_v_descale, args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, args.nhead_stride_bias, args.nhead_stride_randval, - args.nhead_stride_lse, args.nhead_stride_o, args.window_size_left, - args.window_size_right, args.mask_type, args.min_seqlen_q, - args.p_drop, args.s_randval, args.drop_seed_offset, - args.cu_seqlen_q_ptr, args.cu_seqlen_k_ptr); + args.nhead_stride_lse, args.nhead_stride_o, + args.nhead_stride_q_descale, args.nhead_stride_k_descale, + args.nhead_stride_v_descale, args.window_size_left, + args.window_size_right, args.sink_size, args.mask_type, + args.min_seqlen_q, args.p_drop, args.s_randval, args.drop_seed_offset, + args.block_scale_size_q, args.block_scale_size_kv, + args.cu_seqlen_q_ptr, args.cu_seqlen_k_ptr, args.sink_ptr, + args.num_head_q_total, args.head_start); } else { // create batch mode kernel arguments return FmhaKernel::MakeKargsImpl( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.rand_val_ptr, + args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.q_descale_ptr, + args.k_descale_ptr, args.v_descale_ptr, args.rand_val_ptr, args.lse_ptr, args.o_ptr, args.seqlen_q, args.seqlen_k, args.hdim_q, args.hdim_v, args.nhead_q, args.nhead_q / args.nhead_k, args.scale_s, - args.scale_p, args.scale_o, args.logits_soft_cap, args.stride_q, - args.stride_k, args.stride_v, args.stride_bias, args.stride_randval, - args.stride_o, args.nhead_stride_q, args.nhead_stride_k, - args.nhead_stride_v, args.nhead_stride_bias, - args.nhead_stride_randval, args.nhead_stride_lse, args.nhead_stride_o, - args.batch_stride_q, args.batch_stride_k, args.batch_stride_v, - args.batch_stride_bias, args.batch_stride_randval, - args.batch_stride_lse, args.batch_stride_o, args.window_size_left, - args.window_size_right, args.mask_type, args.p_drop, args.s_randval, - args.drop_seed_offset, args.cu_seqlen_q_ptr, args.cu_seqlen_k_ptr); + args.logits_soft_cap, args.stride_q, args.stride_k, args.stride_v, + args.stride_bias, args.stride_randval, args.stride_o, + args.stride_q_descale, args.stride_k_descale, args.stride_v_descale, + args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, + args.nhead_stride_bias, args.nhead_stride_randval, + args.nhead_stride_lse, args.nhead_stride_o, + args.nhead_stride_q_descale, args.nhead_stride_k_descale, + args.nhead_stride_v_descale, args.batch_stride_q, args.batch_stride_k, + args.batch_stride_v, args.batch_stride_bias, + args.batch_stride_randval, args.batch_stride_lse, args.batch_stride_o, + args.batch_stride_q_descale, args.batch_stride_k_descale, + args.batch_stride_v_descale, args.window_size_left, + args.window_size_right, args.sink_size, args.mask_type, args.p_drop, + args.s_randval, args.drop_seed_offset, args.block_scale_size_q, + args.block_scale_size_kv, args.cu_seqlen_q_ptr, args.cu_seqlen_k_ptr, + args.sink_ptr, args.num_head_q_total, args.head_start); } }(); @@ -637,6 +800,55 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) { } } +template +auto fmha_fwd_v3_create_kargs_and_grids(fmha_fwd_args args) { + /// NOTICE: This was borrowed from Aiter. Make sure the selected remap_opt + /// setting truly maximizes the kernel's performance. + int remap_opt = 2; + if (args.mask_type != static_cast(mask_enum::no_mask) && + ((args.nhead_q % 8 != 0) || (16384 < args.seqlen_q))) { + if (65536 <= args.seqlen_q) { + remap_opt = 0; + } else { + remap_opt = 1; + } + } + + auto kargs = [&] { + if constexpr (FmhaKernel::kIsGroupMode) { + return FmhaKernel::MakeKargs( + args.q_ptr, args.k_ptr, args.v_ptr, args.q_descale_ptr, + args.k_descale_ptr, args.v_descale_ptr, args.lse_ptr, args.o_ptr, + args.seqstart_q_ptr, args.seqstart_k_ptr, args.seqlen_q_ptr, + args.seqlen_k_ptr, args.hdim_q, args.hdim_v, args.nhead_q, + args.nhead_q / args.nhead_k, args.scale_s, args.logits_soft_cap, + args.stride_q, args.stride_k, args.stride_v, args.stride_o, + args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, + args.nhead_stride_lse, args.nhead_stride_o, args.window_size_left, + args.window_size_right, args.mask_type, remap_opt, + args.cu_seqlen_q_ptr, args.cu_seqlen_k_ptr); + } else { + return FmhaKernel::MakeKargs( + args.q_ptr, args.k_ptr, args.v_ptr, args.q_descale_ptr, + args.k_descale_ptr, args.v_descale_ptr, args.lse_ptr, args.o_ptr, + args.seqlen_q, args.seqlen_k, args.hdim_q, args.hdim_v, args.nhead_q, + args.nhead_q / args.nhead_k, args.scale_s, args.logits_soft_cap, + args.stride_q, args.stride_k, args.stride_v, args.stride_o, + args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, + args.nhead_stride_lse, args.nhead_stride_o, args.batch_stride_q, + args.batch_stride_k, args.batch_stride_v, args.batch_stride_lse, + args.batch_stride_o, args.window_size_left, args.window_size_right, + args.mask_type, remap_opt, args.cu_seqlen_q_ptr, + args.cu_seqlen_k_ptr); + } + }(); + + dim3 grids = FmhaKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q, + args.hdim_v); + + return ck_tile::make_tuple(kargs, grids); +} + template auto fmha_fwd_pagedkv_create_kargs_and_grids(fmha_fwd_pagedkv_args args) { assert(args.nhead_q % args.nhead_k == 0); @@ -654,8 +866,8 @@ auto fmha_fwd_pagedkv_create_kargs_and_grids(fmha_fwd_pagedkv_args args) { args.stride_o, args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, args.nhead_stride_bias, args.nhead_stride_lse, args.nhead_stride_o, args.batch_stride_k, args.batch_stride_v, - args.window_size_left, args.window_size_right, args.mask_type, - args.min_seqlen_q); + args.window_size_left, args.window_size_right, args.sink_size, + args.mask_type, args.min_seqlen_q, args.sink_ptr); } else { // create batch mode kernel arguments return FmhaKernel::MakeKargs( args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.lse_ptr, @@ -669,7 +881,8 @@ auto fmha_fwd_pagedkv_create_kargs_and_grids(fmha_fwd_pagedkv_args args) { args.nhead_stride_bias, args.nhead_stride_lse, args.nhead_stride_o, args.batch_stride_q, args.batch_stride_k, args.batch_stride_v, args.batch_stride_bias, args.batch_stride_lse, args.batch_stride_o, - args.window_size_left, args.window_size_right, args.mask_type); + args.window_size_left, args.window_size_right, args.sink_size, + args.mask_type, args.sink_ptr); } }(); @@ -706,7 +919,8 @@ auto fmha_fwd_splitkv_create_kargs_and_grids(fmha_fwd_splitkv_args args) { args.batch_stride_k, // only used for paged-kvcache args.batch_stride_v, // only used for paged-kvcache args.split_stride_lse_acc, args.split_stride_o_acc, - args.window_size_left, args.window_size_right, args.mask_type); + args.window_size_left, args.window_size_right, args.sink_size, + args.mask_type, args.sink_ptr); } else { // create batch mode kernel arguments return Kernel::MakeKargs( args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.lse_acc_ptr, @@ -723,7 +937,8 @@ auto fmha_fwd_splitkv_create_kargs_and_grids(fmha_fwd_splitkv_args args) { args.batch_stride_bias, args.batch_stride_lse_acc, args.batch_stride_o_acc, args.split_stride_lse_acc, args.split_stride_o_acc, args.window_size_left, - args.window_size_right, args.mask_type); + args.window_size_right, args.sink_size, args.mask_type, + args.sink_ptr); } }(); @@ -790,46 +1005,59 @@ auto fmha_fwd_appendkv_create_kargs_and_grids(fmha_fwd_appendkv_args args) { template auto fmha_batch_prefill_create_kargs_and_grids(fmha_batch_prefill_args args) { assert(args.nhead_q % args.nhead_k == 0); + using PageTableKargs = typename FmhaKernel::PageBlockTableKargs; + const PageTableKargs page_table = [&]() { + if constexpr (FmhaKernel::kKVLookupTable == + ck_tile::BlockAttentionKVCacheLookupTableEnum:: + SGLANG_PAGE_TABLE_1D) { + return PageTableKargs{ + reinterpret_cast(args.kv_indptr), + reinterpret_cast(args.kv_page_indices), + reinterpret_cast(args.kv_last_page_lens)}; + } else { + return PageTableKargs{ + reinterpret_cast(args.kv_page_indices), + args.batch_stride_block_table, + reinterpret_cast(args.seqlen_k_ptr)}; + } + }(); auto kargs = [&] { // create group mode kernel arguments if constexpr (FmhaKernel::kIsGroupMode) { return FmhaKernel::MakeKargs( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.rand_val_ptr, + args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.q_descale_ptr, + args.k_descale_ptr, args.v_descale_ptr, args.rand_val_ptr, args.lse_ptr, args.o_ptr, args.seqstart_q_ptr, args.hdim_q, args.hdim_v, args.nhead_q, args.nhead_q / args.nhead_k, - args.num_total_pages, args.kv_indptr, args.kv_page_indices, -#if 0 // we assume page_block_size=1 for now - args.kv_last_page_lens, - args.page_block_size, -#endif - args.scale_s, args.scale_p, args.scale_o, args.logits_soft_cap, - args.stride_q, args.stride_k, args.stride_v, args.stride_bias, - args.stride_randval, args.stride_o, args.nhead_stride_q, - args.nhead_stride_k, args.nhead_stride_v, args.nhead_stride_bias, + args.num_total_pages, args.page_block_size, page_table, args.scale_s, + args.scale_p, args.scale_o, args.logits_soft_cap, args.stride_q, + args.stride_k, args.stride_v, args.stride_bias, args.stride_randval, + args.stride_o, args.nhead_stride_q, args.nhead_stride_k, + args.nhead_stride_v, args.nhead_stride_bias, args.nhead_stride_randval, args.nhead_stride_lse, args.nhead_stride_o, args.batch_stride_k, args.batch_stride_v, args.window_size_left, - args.window_size_right, args.mask_type, args.p_drop, args.s_randval, - args.drop_seed_offset); + args.window_size_right, args.sink_size, args.mask_type, args.p_drop, + args.s_randval, args.drop_seed_offset, args.sink_ptr, + args.nblock_stride_kv_block_descale, + args.nhead_stride_kv_block_descale); } else { // create batch mode kernel arguments return FmhaKernel::MakeKargs( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.rand_val_ptr, + args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.q_descale_ptr, + args.k_descale_ptr, args.v_descale_ptr, args.rand_val_ptr, args.lse_ptr, args.o_ptr, args.seqlen_q, args.hdim_q, args.hdim_v, args.nhead_q, args.nhead_q / args.nhead_k, args.num_total_pages, - args.kv_indptr, args.kv_page_indices, -#if 0 // we assume page_block_size=1 for now - args.kv_last_page_lens, - args.page_block_size, -#endif - args.scale_s, args.scale_p, args.scale_o, args.logits_soft_cap, - args.stride_q, args.stride_k, args.stride_v, args.stride_bias, - args.stride_randval, args.stride_o, args.nhead_stride_q, - args.nhead_stride_k, args.nhead_stride_v, args.nhead_stride_bias, - args.nhead_stride_randval, args.nhead_stride_lse, args.nhead_stride_o, - args.batch_stride_q, args.batch_stride_k, args.batch_stride_v, - args.batch_stride_bias, args.batch_stride_randval, - args.batch_stride_lse, args.batch_stride_o, args.window_size_left, - args.window_size_right, args.mask_type, args.p_drop, args.s_randval, - args.drop_seed_offset); + args.page_block_size, page_table, args.scale_s, args.scale_p, + args.scale_o, args.logits_soft_cap, args.stride_q, args.stride_k, + args.stride_v, args.stride_bias, args.stride_randval, args.stride_o, + args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, + args.nhead_stride_bias, args.nhead_stride_randval, + args.nhead_stride_lse, args.nhead_stride_o, args.batch_stride_q, + args.batch_stride_k, args.batch_stride_v, args.batch_stride_bias, + args.batch_stride_randval, args.batch_stride_lse, args.batch_stride_o, + args.window_size_left, args.window_size_right, args.sink_size, + args.mask_type, args.p_drop, args.s_randval, args.drop_seed_offset, + args.sink_ptr, args.nblock_stride_kv_block_descale, + args.nhead_stride_kv_block_descale); } }(); @@ -847,9 +1075,10 @@ template + bool kHasDropout_, ck_tile::BlockAttentionQuantScaleEnum QScaleEnum_, + bool kPadS_, bool kPadSK_, bool kPadD_, bool kPadDv_, + bool kUseTrLoad_, bool kSkipMinSeqlenQ_ = false, + bool kHasSink_ = false> struct fmha_fwd_traits_ { static constexpr ck_tile::index_t HDim = HDim_; using DataType = ck_tile::remove_cvref_t; @@ -867,13 +1096,47 @@ struct fmha_fwd_traits_ { static constexpr auto BiasEnum = BiasEnum_; static constexpr bool kStoreLse = kStoreLse_; static constexpr bool kHasDropout = kHasDropout_; - static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_; + static constexpr auto QScaleEnum = QScaleEnum_; static constexpr bool kPadS = kPadS_; static constexpr bool kPadSK = kPadSK_; static constexpr bool kPadD = kPadD_; static constexpr bool kPadDv = kPadDv_; static constexpr bool kUseTrLoad = kUseTrLoad_; static constexpr bool kSkipMinSeqlenQ = kSkipMinSeqlenQ_; + static constexpr bool kHasSink = kHasSink_; +}; + +template < + ck_tile::index_t HDim_, typename DataType_, bool kIsGroupMode_, + ck_tile::index_t kM0_, ck_tile::index_t kN0_, ck_tile::index_t kK0_, + ck_tile::index_t kN1_, ck_tile::index_t kK1_, + ck_tile::index_t kK0BlockLength_, bool kIsVLayoutRowMajor_, + ck_tile::BlockFmhaPipelineEnum FmhaPipelineEnum_, bool kHasLogitsSoftCap_, + typename FmhaMask_, ck_tile::BlockAttentionBiasEnum BiasEnum_, + bool kStoreLse_, bool kHasDropout_, + ck_tile::BlockAttentionQuantScaleEnum QScaleEnum_, bool kPadS_, + bool kPadSK_, bool kPadD_, bool kPadDv_, bool kUseTrLoad_, + bool kSkipMinSeqlenQ_ = false, bool kHasSink_ = false, + ck_tile::index_t kPageBlockSize_ = 1, + ck_tile::BlockAttentionKVCacheMemoryLayoutEnum kKVMemoryLayout_ = + ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT, + ck_tile::BlockAttentionKVCacheLookupTableEnum kKVLookupTable_ = + ck_tile::BlockAttentionKVCacheLookupTableEnum::SGLANG_PAGE_TABLE_1D, + ck_tile::BlockAttentionKVCacheLoadModeEnum kKVLoadMode_ = + ck_tile::BlockAttentionKVCacheLoadModeEnum::BUFFER_LOAD> +struct fmha_fwd_batch_prefill_traits_ + : public fmha_fwd_traits_ { + static constexpr auto kKVMemoryLayout = kKVMemoryLayout_; + static constexpr auto kKVLookupTable = kKVLookupTable_; + static constexpr ck_tile::index_t kPageBlockSize = kPageBlockSize_; + static constexpr auto kKVLoadMode = kKVLoadMode_; + static_assert(kIsVLayoutRowMajor_, + "Batch prefill only supports row-major V layout"); }; template @@ -887,7 +1150,8 @@ template + bool kPadD_, bool kPadDv_, bool kSkipMinSeqlenQ_ = false, + bool kHasSink_ = false> struct fmha_fwd_pagedkv_traits_ { static constexpr ck_tile::index_t HDim = HDim_; using DataType = ck_tile::remove_cvref_t; @@ -911,6 +1175,7 @@ struct fmha_fwd_pagedkv_traits_ { static constexpr bool kPadD = kPadD_; static constexpr bool kPadDv = kPadDv_; static constexpr bool kSkipMinSeqlenQ = kSkipMinSeqlenQ_; + static constexpr bool kHasSink = kHasSink_; }; template @@ -923,8 +1188,8 @@ template + bool kDoFp8StaticQuant_, bool kIsPagedKV_, bool kHasSink_, + bool kPadS_, bool kPadSK_, bool kPadD_, bool kPadDv_> struct fmha_fwd_splitkv_traits_ { static constexpr ck_tile::index_t HDim = HDim_; using DataType = ck_tile::remove_cvref_t; @@ -947,6 +1212,7 @@ struct fmha_fwd_splitkv_traits_ { static constexpr bool kPadD = kPadD_; static constexpr bool kPadDv = kPadDv_; static constexpr bool kIsPagedKV = kIsPagedKV_; + static constexpr bool kHasSink = kHasSink_; }; template @@ -1022,8 +1288,9 @@ struct fmha_fwd_traits { // BlockAttentionBiasEnum bool has_lse; bool has_dropout; - bool do_fp8_static_quant; + quant_scale_enum qscale_type; bool skip_min_seqlen_q = false; + bool has_sink = false; // TODO: padding check is inside this api }; float fmha_fwd(fmha_fwd_traits, fmha_fwd_args, const ck_tile::stream_config &); @@ -1042,6 +1309,7 @@ struct fmha_fwd_pagedkv_traits { bool use_pagedkv = true; bool do_fp8_static_quant = false; bool skip_min_seqlen_q = false; + bool has_sink = false; // TODO: padding check is inside this api }; @@ -1059,7 +1327,8 @@ struct fmha_fwd_splitkv_traits { bias_enum bias_type; // 0:no bias, 1:elementwise bias, 2:alibi. sync with // BlockAttentionBiasEnum bool has_lse; - bool do_fp8_static_quant; + bool do_fp8_static_quant = false; + bool has_sink = false; // TODO: padding check is inside this api }; float fmha_fwd_splitkv(fmha_fwd_splitkv_traits, fmha_fwd_splitkv_args, @@ -1075,6 +1344,13 @@ struct fmha_fwd_appendkv_traits { float fmha_fwd_appendkv(fmha_fwd_appendkv_traits, fmha_fwd_appendkv_args, const ck_tile::stream_config &); -using fmha_batch_prefill_traits = fmha_fwd_traits; +struct fmha_batch_prefill_traits : public fmha_fwd_traits { + ck_tile::BlockAttentionKVCacheMemoryLayoutEnum kv_memory_layout = + ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT; + ck_tile::BlockAttentionKVCacheLookupTableEnum kv_lookup_table = + ck_tile::BlockAttentionKVCacheLookupTableEnum::SGLANG_PAGE_TABLE_1D; + int page_size = 1; +}; + float fmha_batch_prefill(fmha_batch_prefill_traits, fmha_batch_prefill_args, const ck_tile::stream_config &); diff --git a/Libraries/ComposableKernel/attention/fmha/fmha_fwd_head_grouping.hpp b/Libraries/ComposableKernel/attention/fmha/fmha_fwd_head_grouping.hpp new file mode 100644 index 0000000000..8352d1dfb3 --- /dev/null +++ b/Libraries/ComposableKernel/attention/fmha/fmha_fwd_head_grouping.hpp @@ -0,0 +1,415 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "ck_tile/host.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __linux__ +#include +#endif + +#ifndef CK_TILE_FMHA_ENABLE_HEAD_GROUPING +#define CK_TILE_FMHA_ENABLE_HEAD_GROUPING 1 +#endif + +#if __clang_major__ >= 23 +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions" +#endif +#if CK_TILE_FMHA_ENABLE_HEAD_GROUPING +CK_TILE_DECLARE_ENV_VAR_BOOL(CK_TILE_FMHA_HEAD_GROUP_LOG) +CK_TILE_DECLARE_ENV_VAR_BOOL(CK_TILE_FMHA_DISABLE_HEAD_GROUPING) +CK_TILE_DECLARE_ENV_VAR_UINT64(CK_TILE_FMHA_LLC_CACHE_MB) + +namespace fmha_fwd_head_grouping { + +inline bool log_enabled() { + return ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_FMHA_HEAD_GROUP_LOG)); +} + +inline bool disabled_by_env() { + return ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_FMHA_DISABLE_HEAD_GROUPING)); +} + +inline bool is_decimal_string(const std::string &s) { + if (s.empty()) + return false; + return std::all_of(s.begin(), s.end(), + [](unsigned char c) { return std::isdigit(c) != 0; }); +} + +inline std::optional read_property_value(const std::string &filepath, + const std::string &key) { + std::ifstream fs(filepath); + if (!fs.is_open()) + return std::nullopt; + + std::string k, v; + while (fs >> k >> v) { + if (k == key) { + try { + return std::stoll(v, nullptr, 0); + } catch (...) { + return std::nullopt; + } + } + std::string rest; + std::getline(fs, rest); + } + return std::nullopt; +} + +#if defined(__linux__) + +struct kfd_device_location { + int domain = 0; + int location_id = 0; +}; + +inline std::optional get_current_kfd_location() { + int device = 0; + if (hipGetDevice(&device) != hipSuccess) + return std::nullopt; + + char bdf[64] = {}; + if (hipDeviceGetPCIBusId(bdf, sizeof(bdf), device) == hipSuccess) { + unsigned int domain = 0, bus = 0, dev = 0, fn = 0; + if (std::sscanf(bdf, "%x:%x:%x.%x", &domain, &bus, &dev, &fn) == 4) { + return kfd_device_location{static_cast(domain), + static_cast(((bus & 0xff) << 8) | + ((dev & 0x1f) << 3) | + (fn & 0x7))}; + } + } + + hipDeviceProp_t props{}; + if (hipGetDeviceProperties(&props, device) != hipSuccess) + return std::nullopt; + + return kfd_device_location{props.pciDomainID, + ((props.pciBusID & 0xff) << 8) | + ((props.pciDeviceID & 0x1f) << 3)}; +} + +inline std::optional +find_matching_kfd_node(const kfd_device_location &loc) { + constexpr const char *kKfdNodesDir = "/sys/class/kfd/kfd/topology/nodes"; + DIR *dir = opendir(kKfdNodesDir); + if (dir == nullptr) + return std::nullopt; + + std::optional matched; + while (auto *ent = readdir(dir)) { + const std::string node_name(ent->d_name); + if (!is_decimal_string(node_name)) + continue; + + const std::string prop_path = + std::string(kKfdNodesDir) + "/" + node_name + "/properties"; + const auto location_val = read_property_value(prop_path, "location_id"); + if (!location_val.has_value() || + static_cast(*location_val) != loc.location_id) + continue; + + const auto domain_val = read_property_value(prop_path, "domain"); + if (domain_val.has_value() && static_cast(*domain_val) != loc.domain) + continue; + + matched = node_name; + break; + } + + closedir(dir); + return matched; +} + +inline size_t read_kfd_node_l3_bytes(const std::string &node_name) { + const std::string caches_dir = + "/sys/class/kfd/kfd/topology/nodes/" + node_name + "/caches"; + DIR *dir = opendir(caches_dir.c_str()); + if (dir == nullptr) + return 0; + + size_t l3_kb = 0; + while (auto *ent = readdir(dir)) { + const std::string cache_name(ent->d_name); + if (!is_decimal_string(cache_name)) + continue; + + const std::string prop_path = caches_dir + "/" + cache_name + "/properties"; + const auto level_val = read_property_value(prop_path, "level"); + if (!level_val.has_value() || *level_val != 3) + continue; + + const auto size_val = read_property_value(prop_path, "size"); + if (!size_val.has_value() || *size_val <= 0) + continue; + + l3_kb = std::max(l3_kb, static_cast(*size_val)); + } + + closedir(dir); + return l3_kb * 1024ull; +} + +inline size_t get_kfd_sysfs_llc_cache_bytes() { + const auto loc = get_current_kfd_location(); + if (!loc.has_value()) + return 0; + + const auto node = find_matching_kfd_node(*loc); + if (!node.has_value()) + return 0; + + return read_kfd_node_l3_bytes(*node); +} + +#else + +inline size_t get_kfd_sysfs_llc_cache_bytes() { return 0; } + +#endif + +inline size_t get_default_llc_cache_bytes_for_arch(const std::string &arch); + +inline size_t resolve_llc_cache_bytes_uncached(const std::string &arch) { + // If parsed LLC looks invalidly tiny, ignore it and fallback. + constexpr size_t kMinValidKfdLlcBytes = 32ull * 1024ull; + + const size_t kfd_llc_bytes = get_kfd_sysfs_llc_cache_bytes(); + if (kfd_llc_bytes >= kMinValidKfdLlcBytes) + return kfd_llc_bytes; + + const size_t default_cache_bytes = get_default_llc_cache_bytes_for_arch(arch); + if (default_cache_bytes > 0) + return default_cache_bytes; + + // No default configured -> no grouping. + return 0; +} + +inline bool ck_tile_is_rdna_arch(const std::string &arch) { + return arch.rfind("gfx11", 0) == 0 || arch.rfind("gfx12", 0) == 0; +} + +inline size_t get_default_llc_cache_bytes_for_arch(const std::string &arch) { + if (arch == "gfx1100") + return 96ull * 1024ull * 1024ull; + if (arch == "gfx1101") + return 64ull * 1024ull * 1024ull; + if (arch == "gfx1102") + return 32ull * 1024ull * 1024ull; + if (arch == "gfx1151") + return 32ull * 1024ull * 1024ull; + if (arch == "gfx1200") + return 32ull * 1024ull * 1024ull; + if (arch == "gfx1201") + return 64ull * 1024ull * 1024ull; + return 0; +} + +inline size_t get_llc_cache_bytes(const std::string &arch) { + // resolve once and reuse. + static const size_t resolved_llc_bytes = [&]() -> size_t { + const uint64_t llc_mb = + ck_tile::EnvValue(CK_TILE_ENV(CK_TILE_FMHA_LLC_CACHE_MB)); + if (llc_mb > 0) { + constexpr uint64_t kBytesPerMb = 1024ull * 1024ull; + const uint64_t max_mb_for_size_t = + static_cast(std::numeric_limits::max() / + static_cast(kBytesPerMb)); + + if (llc_mb <= max_mb_for_size_t) + return static_cast(llc_mb * kBytesPerMb); + } + + return resolve_llc_cache_bytes_uncached(arch); + }(); + + return resolved_llc_bytes; +} + +inline std::optional +get_head_group_size(ck_tile::index_t nhead_q, ck_tile::index_t nhead_k, + ck_tile::index_t batch, ck_tile::index_t seqlen_k, + ck_tile::index_t hdim_q, ck_tile::index_t hdim_v, + size_t elem_bytes_k, size_t elem_bytes_v) { + if (disabled_by_env()) + return std::nullopt; + + const std::string arch = ck_tile::get_device_name(); + if (arch.empty() || !ck_tile_is_rdna_arch(arch)) + return std::nullopt; + + const size_t llc_bytes = get_llc_cache_bytes(arch); + if (llc_bytes == 0) + return std::nullopt; + + if (nhead_k <= 0 || nhead_q <= 0 || (nhead_q % nhead_k) != 0) + return std::nullopt; + if (seqlen_k <= 0 || hdim_q <= 0 || hdim_v <= 0 || batch <= 0) + return std::nullopt; + + const size_t kv_bytes_per_head = static_cast(seqlen_k) * + (static_cast(hdim_q) * elem_bytes_k + + static_cast(hdim_v) * elem_bytes_v); + if (kv_bytes_per_head == 0) + return std::nullopt; + + // large LLC GPUs (>= 64MB): slightly more cache-resident grouping + constexpr size_t kLargeLlcThresholdBytes = 64ull * 1024ull * 1024ull; + const bool is_large_llc = llc_bytes >= kLargeLlcThresholdBytes; + const long double llc_utilization = is_large_llc ? 0.85L : 1.0L; + const long double threshold_ratio = is_large_llc ? 1.3L : 1.5L; + const size_t target_llc_bytes = static_cast( + static_cast(llc_bytes) * llc_utilization); + if (target_llc_bytes == 0) + return std::nullopt; + + const size_t total_kv_bytes = + static_cast(nhead_q) * kv_bytes_per_head; + if (static_cast(total_kv_bytes) < + static_cast(target_llc_bytes) * threshold_ratio) + return std::nullopt; + + ck_tile::index_t group = + static_cast(target_llc_bytes / kv_bytes_per_head); + if (group < 1) + group = 1; + + const ck_tile::index_t min_group_size = + std::max(1, nhead_q / 16); + if (group < min_group_size) + group = min_group_size; + + // Cap the number of groups to avoid excessive launch overhead. + constexpr ck_tile::index_t kMaxGroups = 8; + const ck_tile::index_t min_group_for_max_groups = + ck_tile::integer_divide_ceil(nhead_q, kMaxGroups); + if (group < min_group_for_max_groups) + group = min_group_for_max_groups; + + const ck_tile::index_t gqa_ratio = nhead_q / nhead_k; + if (gqa_ratio > 1) { + group = ((group + gqa_ratio - 1) / gqa_ratio) * gqa_ratio; + } + + group = std::min(group, nhead_q); + if (group >= nhead_q) + return std::nullopt; + + return group; +} + +template +inline const void *ptr_offset(const void *base, ck_tile::index_t offset_elems) { + if (base == nullptr) + return nullptr; + return static_cast(reinterpret_cast(base) + + offset_elems); +} + +template +inline void *ptr_offset(void *base, ck_tile::index_t offset_elems) { + if (base == nullptr) + return nullptr; + return static_cast(reinterpret_cast(base) + offset_elems); +} + +template +float run_fwd_head_grouped(const ck_tile::stream_config &sc, + const FmhaFwdTraits &fmha_traits, + const FmhaFwdArgs &base_args_in, + ck_tile::index_t nhead, ck_tile::index_t nhead_k, + ck_tile::index_t group_size_q, + bool use_blockscale_qscale, + RunKernelFn &&run_kernel_fn) { + auto base_args = base_args_in; + base_args.num_head_q_total = nhead; + const ck_tile::index_t gqa_ratio = (nhead_k > 0 ? (nhead / nhead_k) : 1); + const ck_tile::index_t group_sz = std::min(group_size_q, nhead); + const ck_tile::index_t n_groups = + ck_tile::integer_divide_ceil(nhead, group_sz); + + float total_time = 0.0f; + for (ck_tile::index_t head_start = 0; head_start < nhead; + head_start += group_sz) { + const ck_tile::index_t q_heads = std::min(group_sz, nhead - head_start); + const ck_tile::index_t k_head_start = + (gqa_ratio >= 1 ? head_start / gqa_ratio : head_start); + const ck_tile::index_t k_heads = + (gqa_ratio >= 1 ? q_heads / gqa_ratio : q_heads); + + auto args = base_args; + args.nhead_q = q_heads; + args.nhead_k = k_heads; + args.head_start = head_start; + + args.q_ptr = ptr_offset(base_args.q_ptr, + head_start * base_args.nhead_stride_q); + args.k_ptr = ptr_offset(base_args.k_ptr, + k_head_start * base_args.nhead_stride_k); + args.v_ptr = ptr_offset(base_args.v_ptr, + k_head_start * base_args.nhead_stride_v); + args.o_ptr = ptr_offset(base_args.o_ptr, + head_start * base_args.nhead_stride_o); + + args.bias_ptr = ptr_offset( + base_args.bias_ptr, head_start * base_args.nhead_stride_bias); + args.lse_ptr = ptr_offset( + base_args.lse_ptr, head_start * base_args.nhead_stride_lse); + args.rand_val_ptr = ptr_offset( + base_args.rand_val_ptr, head_start * base_args.nhead_stride_randval); + + if (use_blockscale_qscale) { + args.q_descale_ptr = + ptr_offset(base_args.q_descale_ptr, + head_start * base_args.nhead_stride_q_descale); + args.k_descale_ptr = + ptr_offset(base_args.k_descale_ptr, + k_head_start * base_args.nhead_stride_k_descale); + args.v_descale_ptr = + ptr_offset(base_args.v_descale_ptr, + k_head_start * base_args.nhead_stride_v_descale); + } else { + args.q_descale_ptr = base_args.q_descale_ptr; + args.k_descale_ptr = base_args.k_descale_ptr; + args.v_descale_ptr = base_args.v_descale_ptr; + } + + args.sink_ptr = ptr_offset(base_args.sink_ptr, head_start); + + if (log_enabled()) { + const ck_tile::index_t head_end = head_start + q_heads; + std::cout << "[LLC Head Grouping] group " << (head_start / group_sz) + << "/" << n_groups << " heads_q=[" << head_start << ", " + << head_end << ") heads_k=[" << k_head_start << ", " + << (k_head_start + k_heads) << ")" << std::endl; + } + + const float t = run_kernel_fn(fmha_traits, args, sc); + if (t < 0.0f) + return t; + total_time += t; + } + return total_time; +} + +} // namespace fmha_fwd_head_grouping +#endif // CK_TILE_FMHA_ENABLE_HEAD_GROUPING +#if __clang_major__ >= 23 +#pragma clang diagnostic pop +#endif diff --git a/Libraries/ComposableKernel/attention/fmha/fmha_fwd_runner.hpp b/Libraries/ComposableKernel/attention/fmha/fmha_fwd_runner.hpp index 14012f2231..a29e0925b8 100644 --- a/Libraries/ComposableKernel/attention/fmha/fmha_fwd_runner.hpp +++ b/Libraries/ComposableKernel/attention/fmha/fmha_fwd_runner.hpp @@ -1,38 +1,28 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once #include "ck_tile/host.hpp" #include "ck_tile/ref/naive_attention.hpp" -#include "ck_tile/utility/json_dump.hpp" #include "fmha_fwd.hpp" + +#include "fmha_fwd_head_grouping.hpp" + #include "utils.hpp" +#include "ck_tile/utility/json_dump.hpp" + #include -#include +#include + #include + #include + +#include #include +#include #include #include #include @@ -44,6 +34,10 @@ "we should enable fmha_fwd_splitkv() api in order to cooperate with fmha_fwd_appendkv()" #endif +#if __clang_major__ >= 23 +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wlifetime-safety-invalidation" +#endif enum class fwd_result { success, failure, @@ -93,6 +87,18 @@ template <> auto get_elimit(std::string /*init_method*/) { return ck_tile::make_tuple(rtol, atol); } +template <> auto get_elimit(std::string /*init_method*/) { + double rtol = 1e-2; + double atol = 1.8e-1; + return ck_tile::make_tuple(rtol, atol); +} + +template <> auto get_elimit(std::string /*init_method*/) { + double rtol = 1e-1; + double atol = 2.6e-1; + return ck_tile::make_tuple(rtol, atol); +} + int num_splits_heuristic(int batch_nhead_mblocks, int num_SMs, int max_splits) { // If we have enough to almost fill the SMs, then just use 1 split if (batch_nhead_mblocks >= 0.8f * num_SMs) { @@ -141,13 +147,55 @@ int override_num_splits_if_necessary(int batch, int nhead, int max_seqlen_q, const int num_m_blocks = ck_tile::integer_divide_ceil(max_seqlen_q, kM0); if (num_splits < 1 && p_drop == 0.0f) { + // props.multiProcessorCount for >=gfx10 is the number of WGPs (each has 2 + // CUs) + const int num_blocks_per_SM = props.warpSize == 32 ? 4 : 2; return num_splits_heuristic(batch * nhead * num_m_blocks, - props.multiProcessorCount * 2, 128); + props.multiProcessorCount * num_blocks_per_SM, + 128); } return num_splits; } +template +void copy_attention_scores_with_sink( + const ck_tile::HostTensor &s_host_ref, + const ck_tile::HostTensor &sink_host, + ck_tile::HostTensor &s_with_sinks_ref, + ck_tile::index_t nhead, ck_tile::index_t real_seqlen_q, + ck_tile::index_t real_seqlen_k) { + for (auto i_h = 0; i_h < nhead; i_h++) { + for (auto i_r = 0; i_r < real_seqlen_q; i_r++) { + for (auto i_c = 0; i_c < real_seqlen_k; i_c++) { + s_with_sinks_ref(i_h, i_r, i_c) = s_host_ref(i_h, i_r, i_c); + } + // Append sink token at the end of each row + s_with_sinks_ref(i_h, i_r, real_seqlen_k) = sink_host(i_h); + } + } +} + +template struct ScalesConfig { + using QScaleDataType = float; + using KScaleDataType = float; + using VScaleDataType = float; + + static constexpr ck_tile::index_t kQKScaleGranularity = 1; + static constexpr ck_tile::index_t kVScaleGranularity = 1; +}; + +template struct ScalesConfig { + using QScaleDataType = typename TypeConfig::QScaleDataType; + using KScaleDataType = typename TypeConfig::KScaleDataType; + using VScaleDataType = typename TypeConfig::VScaleDataType; + + static constexpr ck_tile::index_t kQKScaleGranularity = + TypeConfig::kQKScaleGranularity; + static constexpr ck_tile::index_t kVScaleGranularity = + TypeConfig::kVScaleGranularity; +}; + template fwd_result fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, @@ -163,11 +211,50 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::index_t page_block_size, bool use_cache_batch_idx, std::string bias_str, float p_drop, uint64_t drop_seed, uint64_t drop_offset, bool drop_prefs, std::string mask_str, - bool squant, bool is_rotary_interleaved, + std::string qscale_str, bool is_rotary_interleaved, ck_tile::index_t num_splits, std::string init_method, - uint32_t seed, int do_validation, - const ck_tile::stream_config &stream_config, - std::optional json = std::nullopt) { + uint32_t seed, int do_validation, int init_sink_value, + int pack_gqa, const ck_tile::stream_config &stream_config, + std::optional json = std::nullopt, + std::string *selected_kernel_name = nullptr) { + using TypeConfig = FmhaFwdTypeConfig; + + if (selected_kernel_name != nullptr) + selected_kernel_name->clear(); + + constexpr bool is_mx = + ck_tile::is_any_of::value; + + using QDataType = typename TypeConfig::QDataType; + using KDataType = typename TypeConfig::KDataType; + using VDataType = typename TypeConfig::VDataType; + using BiasDataType = typename TypeConfig::BiasDataType; + using RandValOutputDataType = typename TypeConfig::RandValOutputDataType; + using LSEDataType = typename TypeConfig::LSEDataType; + using SaccDataType = typename TypeConfig::SaccDataType; + using SMPLComputeDataType = typename TypeConfig::SMPLComputeDataType; + using PDataType = + std::conditional_t; + using OaccDataType = typename TypeConfig::OaccDataType; + using ODataType = typename TypeConfig::ODataType; + + using QScaleDataType = + typename ScalesConfig::QScaleDataType; + using KScaleDataType = + typename ScalesConfig::KScaleDataType; + using VScaleDataType = + typename ScalesConfig::VScaleDataType; + + constexpr ck_tile::index_t kQKScaleGranularity = + ScalesConfig::kQKScaleGranularity; + constexpr ck_tile::index_t kVScaleGranularity = + ScalesConfig::kVScaleGranularity; + + // Note: block_scale_size_q_ and block_scale_size_kv_ should be greater than + // or equal to the compute block size + constexpr ck_tile::index_t block_scale_size_q_ = 128; + constexpr ck_tile::index_t block_scale_size_kv_ = 128; + const std::string data_type = []() { if constexpr (std::is_same_v) return "fp32"; @@ -183,6 +270,10 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, return "fp8bf16"; else if constexpr (std::is_same_v) return "fp8fp32"; + else if constexpr (std::is_same_v) + return "mxfp8"; + else if constexpr (std::is_same_v) + return "mxfp4"; else static_assert(false); }(); @@ -195,6 +286,23 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, return fwd_result::invalid_args; } + if (hdim_q % ck_tile::numeric_traits::PackedSize != 0) { + std::cerr << "hdim_q is made even for fp4 Q data type" << std::endl; + hdim_q = ck_tile::integer_least_multiple( + hdim_q, ck_tile::numeric_traits::PackedSize); + } + if (hdim_q % ck_tile::numeric_traits::PackedSize != 0) { + std::cerr << "hdim_q is made even for fp4 K data type" << std::endl; + hdim_q = ck_tile::integer_least_multiple( + hdim_q, ck_tile::numeric_traits::PackedSize); + } + if (is_mx && !seqlen_kpads.empty() && seqlen_kpads[0] > 0) { + std::cerr << "seqlen_kpads is not supported with MX types. ignoring the " + "'seqlen_kpads' option" + << std::endl; + seqlen_kpads = {-1}; + } + std::mt19937 random_engine(seed != 0 ? seed : std::random_device{}()); auto next_seed = [&random_engine]() { return static_cast(random_engine()); @@ -251,7 +359,7 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, } #if (!(CK_TILE_FMHA_FWD_APPENDKV_API || CK_TILE_FMHA_FWD_SPLITKV_API || \ - CK_TILE_FMHA_FWD_PAGEDKV_API)) + CK_TILE_FMHA_FWD_PAGEDKV_API || CK_TILE_FMHA_FWD_BATCH_PREFILL_API)) if (0 < page_block_size) { std::cerr << "paged-kvcache is not supported. ignoring the " "'page_block_size' option" @@ -259,7 +367,12 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, page_block_size = 0; } #endif - if (!(page_block_size % 128 == 0)) { + // batch_prefill supports flexible page sizes (not just multiples of 128) + const bool need_128_aligned_page = + (CK_TILE_FMHA_FWD_APPENDKV_API || CK_TILE_FMHA_FWD_SPLITKV_API || + CK_TILE_FMHA_FWD_PAGEDKV_API); + if (need_128_aligned_page && 0 < page_block_size && + !(page_block_size % 128 == 0)) { std::cerr << "only paged-kvcache block size divisible by 128 are currently " "supported" << std::endl; @@ -324,6 +437,19 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, mode, batch, seqlen_qs, seqlen_ks, seqlen_qpads, seqlen_kpads, /*seqlen_k_min=*/0 < seqlen_knew ? seqlen_knew : 0, need_append_kvcache, random_engine); + + if (ck_tile::numeric_traits::PackedSize != 0) { + // Ensure that all seqlens are even if V has packed data type + for (auto &s : seqlen_ks) { + s = ck_tile::integer_least_multiple( + s, ck_tile::numeric_traits::PackedSize); + } + for (auto &s : kv_eff_lens_per_batch) { + s = ck_tile::integer_least_multiple( + s, ck_tile::numeric_traits::PackedSize); + } + } + for (ck_tile::index_t wb = 0; wb < batch; ++wb) { if (seqlen_kpads[wb] > 0 && seqlen_kpads[wb] < seqlen_ks[wb]) { std::cerr << "kpad must be greater than or equal to seqlen for k" @@ -359,6 +485,23 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, mask_info mask = mask_info::decode( mask_str, seqlen_qs[0], seqlen_ks[0]); // TODO: we don't need x/y anymore + quant_scale_info qscale = quant_scale_info::decode(qscale_str); + + if (is_mx && qscale.type != quant_scale_enum::mx) { + std::cerr << "The value of qscale_str must be 'mx' for MX data types" + << std::endl; + return fwd_result::invalid_args; + } else if (!is_mx && qscale.type == quant_scale_enum::mx) { + std::cerr << "The value of qscale_str cannot be 'mx' for non-MX data types" + << std::endl; + return fwd_result::invalid_args; + } + if (is_mx && is_v_rowmajor) { + std::cerr << "The value of is_v_rowmajor must be 'false' for MX data types" + << std::endl; + return fwd_result::invalid_args; + } + if (p_drop < 0.0f || p_drop > 1.0f) { std::cerr << "The value of p_drop should be 0~1" << std::endl; return fwd_result::invalid_args; @@ -403,25 +546,17 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, calculate_cumulative(kv_eff_lens_per_batch, cukv_cum); } - using TypeConfig = FmhaFwdTypeConfig; - - using QDataType = typename TypeConfig::QDataType; - using KDataType = typename TypeConfig::KDataType; - using VDataType = typename TypeConfig::VDataType; - using BiasDataType = typename TypeConfig::BiasDataType; - using RandValOutputDataType = typename TypeConfig::RandValOutputDataType; - using LSEDataType = typename TypeConfig::LSEDataType; - using SaccDataType = typename TypeConfig::SaccDataType; - using SMPLComputeDataType = typename TypeConfig::SMPLComputeDataType; - using PDataType = typename TypeConfig::PDataType; - using OaccDataType = typename TypeConfig::OaccDataType; - using ODataType = typename TypeConfig::ODataType; - // accumulation numbers for performance evaluation std::size_t flop = 0, num_byte = 0; auto max_seqlen_q = std::numeric_limits::min(); // we will use max seqlen to decide // grid size + int32_t i_block_scale_q = 0; + int32_t i_block_scale_k = 0; + int32_t i_seqstart_v_scale = 0; + std::vector block_scale_seqstart_q_host = {0}; + std::vector block_scale_seqstart_k_host = {0}; + std::vector seqstart_v_scale_host = {0}; auto max_seqlen_k = std::numeric_limits::min(); { for (ck_tile::index_t wb = 0; wb < batch; ++wb) { @@ -437,15 +572,31 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, if (max_seqlen_k < real_seqlen_k) { max_seqlen_k = real_seqlen_k; } + if (qscale.type == quant_scale_enum::blockscale) { + i_block_scale_q += + ck_tile::integer_divide_ceil(real_seqlen_q, block_scale_size_q_); + i_block_scale_k += + ck_tile::integer_divide_ceil(real_seqlen_k, block_scale_size_kv_); + block_scale_seqstart_q_host.push_back(i_block_scale_q); + block_scale_seqstart_k_host.push_back(i_block_scale_k); + } else if (qscale.type == quant_scale_enum::mx) { + i_seqstart_v_scale += + ck_tile::integer_divide_ceil(real_seqlen_k, kVScaleGranularity); + seqstart_v_scale_host.push_back(i_seqstart_v_scale); + } flop += nhead * (static_cast(2) * mask.get_unmaskarea() * hdim_q + static_cast(2) * mask.get_unmaskarea() * hdim_v); - num_byte += nhead * (sizeof(QDataType) * real_seqlen_q * hdim_q + + num_byte += nhead * (sizeof(QDataType) * real_seqlen_q * hdim_q / + ck_tile::numeric_traits::PackedSize + sizeof(ODataType) * real_seqlen_q * hdim_v); - num_byte += nhead_k * (sizeof(KDataType) * real_seqlen_k * hdim_q + - sizeof(VDataType) * hdim_v * real_seqlen_k); + num_byte += + nhead_k * (sizeof(KDataType) * real_seqlen_k * hdim_q / + ck_tile::numeric_traits::PackedSize + + sizeof(VDataType) * hdim_v * real_seqlen_k / + ck_tile::numeric_traits::PackedSize); } } @@ -457,8 +608,17 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, // legalize num_splits according to other options if (num_splits < 1) { - num_splits = override_num_splits_if_necessary(batch, nhead, max_seqlen_q, - hdim_v, p_drop, num_splits); + int nhead_merged = nhead; + int max_seqlen_q_merged = max_seqlen_q; + // When max_seqlen_q == 1 and multiple head groups are merged + // (kMergeNumHeadGroupsSeqLenQ) then more splits are required + if (bias.type == bias_enum::no_bias && mask.type == mask_enum::no_mask && + max_seqlen_q == 1 && nhead_k < nhead) { + nhead_merged = nhead_k; + max_seqlen_q_merged = max_seqlen_q * (nhead / nhead_k); + } + num_splits = override_num_splits_if_necessary( + batch, nhead_merged, max_seqlen_q_merged, hdim_v, p_drop, num_splits); } if (128 < num_splits) { std::cerr << "num_splits greater than 128 is not supported" << std::endl; @@ -500,8 +660,30 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ? seqstart_k_with_padding_host.back() : seqstart_k_host.back())); + // Compute Pack-GQA dimensions early so buffer allocations use the right sizes + const int nhead_ratio = nhead / nhead_k; + int pack_gqa_nhead = nhead; + int pack_gqa_seqlen_q = shape_seqlen_q; + if (pack_gqa && nhead_ratio > 1 && mask.type == mask_enum::no_mask && + bias.type == bias_enum::no_bias && i_perm && o_perm && + mode == mode_enum::batch && q_eff_lens_per_batch.empty() && + kv_eff_lens_per_batch.empty() && qscale.type != quant_scale_enum::mx) { + pack_gqa_nhead = nhead_k; + pack_gqa_seqlen_q = nhead_ratio * shape_seqlen_q; + } + + const ck_tile::index_t num_block_scale_q = + (mode == mode_enum::batch) + ? ck_tile::integer_divide_ceil(shape_seqlen_q, block_scale_size_q_) + : i_block_scale_q; + const ck_tile::index_t num_block_scale_kv = + (mode == mode_enum::batch) + ? ck_tile::integer_divide_ceil(shape_seqlen_k, block_scale_size_kv_) + : i_block_scale_k; + ck_tile::HostTensor q_host( get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, hdim_q)); + ck_tile::HostTensor sink_host({nhead}); ck_tile::HostTensor k_host( 0 < page_block_size ? get_lengths(i_perm, max_num_page_blocks, nhead_k, page_block_size, @@ -548,29 +730,56 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::HostTensor lse_acc_host( 1 < num_splits || use_kvcache - ? std::array{shape_batch, nhead, num_splits, - shape_seqlen_q} + ? std::array{shape_batch, pack_gqa_nhead, + num_splits, pack_gqa_seqlen_q} : std::array{1, 1, 1, 1}); ck_tile::HostTensor o_acc_host( 1 < num_splits || use_kvcache - ? std::array{shape_batch, nhead, num_splits, - shape_seqlen_q, hdim_v} + ? std::array{shape_batch, pack_gqa_nhead, + num_splits, pack_gqa_seqlen_q, + hdim_v} : std::array{1, 1, 1, 1, 1}); + const ck_tile::index_t hdim_q_scale = + ck_tile::integer_divide_ceil(hdim_q, kQKScaleGranularity); + const ck_tile::index_t shape_seqlen_v_scale = seqstart_v_scale_host.back(); + + ck_tile::HostTensor q_descale_host({1}); + ck_tile::HostTensor k_descale_host({1}); + ck_tile::HostTensor v_descale_host({1}); + if constexpr (is_mx) { + q_descale_host = ck_tile::HostTensor( + get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, hdim_q_scale)); + k_descale_host = ck_tile::HostTensor(get_lengths( + i_perm, shape_batch, nhead_k, shape_seqlen_k, hdim_q_scale)); + v_descale_host = ck_tile::HostTensor(get_lengths( + i_perm, shape_batch, nhead_k, hdim_v, shape_seqlen_v_scale)); + } else if (qscale.type == quant_scale_enum::blockscale) { + q_descale_host = ck_tile::HostTensor( + std::array{shape_batch, nhead, num_block_scale_q}); + k_descale_host = + ck_tile::HostTensor(std::array{ + shape_batch, nhead_k, num_block_scale_kv}); + v_descale_host = + ck_tile::HostTensor(std::array{ + shape_batch, nhead_k, num_block_scale_kv}); + } + // batch mode of lse data layout is [batch, nhead, seqlen_q] // group mode of lse data layout is [nhead, total_seqlen_q] ck_tile::HostTensor lse_host( - lse ? std::array{shape_batch, nhead, shape_seqlen_q} + lse ? std::array{shape_batch, pack_gqa_nhead, + pack_gqa_seqlen_q} : std::array{ 1, 1, 1} /* dummy shape for simplifying code */); - ck_tile::HostTensor o_host( - get_lengths(o_perm, shape_batch, nhead, shape_seqlen_q, hdim_v)); + ck_tile::HostTensor o_host(get_lengths( + o_perm, shape_batch, pack_gqa_nhead, pack_gqa_seqlen_q, hdim_v)); ck_tile::HostTensor randval_host( - p_drop > 0 - ? get_lengths(true, shape_batch, nhead, shape_seqlen_q, max_seqlen_k) - : std::array{1, 1, 1, 1}); + p_drop > 0 ? get_lengths(true, shape_batch, pack_gqa_nhead, + pack_gqa_seqlen_q, max_seqlen_k) + : std::array{1, 1, 1, 1}); ck_tile::HostTensor block_table_host( 0 < page_block_size @@ -580,7 +789,6 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::HostTensor cache_batch_idx_host( use_cache_batch_idx ? std::array{batch} : std::array{1}); - float max_o = 5.0; if (init_method == "ui" || init_method == "0") { ck_tile::FillUniformDistributionIntegerValue{ -3.f, 3.f, next_seed()}(q_host); @@ -594,7 +802,9 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, -3.f, 3.f, next_seed()}(vnew_host); ck_tile::FillUniformDistributionIntegerValue{ -3.f, 3.f, next_seed()}(bias_host); - } else if (init_method == "ni") { + } + + else if (init_method == "ni") { ck_tile::FillNormalDistributionIntegerValue{-3.f, 3.f, next_seed()}(q_host); ck_tile::FillNormalDistributionIntegerValue{-3.f, 3.f, @@ -634,6 +844,26 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::FillTrigValue{}(v_host); ck_tile::FillTrigValue{}(vnew_host); ck_tile::FillTrigValue{}(bias_host); + } else if (init_method == "3") { + float q_dtype_max = + ck_tile::type_convert(ck_tile::numeric::max()); + float k_dtype_max = + ck_tile::type_convert(ck_tile::numeric::max()); + float v_dtype_max = + ck_tile::type_convert(ck_tile::numeric::max()); + + ck_tile::FillUniformDistribution{-q_dtype_max, q_dtype_max, + next_seed()}(q_host); + ck_tile::FillUniformDistribution{-k_dtype_max, k_dtype_max, + next_seed()}(k_host); + ck_tile::FillUniformDistribution{-k_dtype_max, k_dtype_max, + next_seed()}(knew_host); + ck_tile::FillUniformDistribution{-v_dtype_max, v_dtype_max, + next_seed()}(v_host); + ck_tile::FillUniformDistribution{-v_dtype_max, v_dtype_max, + next_seed()}(vnew_host); + ck_tile::FillUniformDistribution{0.f, 1.f, + next_seed()}(bias_host); } if (bias.type == bias_enum::alibi) { auto slopes = ck_tile::get_alibi_slopes(nhead); @@ -649,24 +879,113 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, } } } + if constexpr (is_mx) { + auto gen_scales = [&](auto &scales, auto data, float range) { + using DataType = decltype(data); + using ScaleType = ck_tile::remove_cvref_t; + if constexpr (std::is_same_v) { + const float base = -std::log2( + ck_tile::type_convert(ck_tile::numeric::max())); + // e8m0_t is basically an exponent of float32 + // When scales are applied to tensor values, value * exp2(base - range) + // is around 0.125 and value * exp2(base + range) is around 8 for all + // types (fp8/bf8/fp4) + ck_tile::HostTensor pow2(scales.get_lengths()); + ck_tile::FillUniformDistributionIntegerValue{ + base - range, base + range, next_seed()}(pow2); + scales.ForEach([&](auto &self, const auto &i) { + self(i) = ck_tile::type_convert(std::exp2(pow2(i))); + }); + } else { + static_assert(false); + } + }; + gen_scales(q_descale_host, QDataType{}, 3); + gen_scales(k_descale_host, KDataType{}, 3); + // When P is fp4, only 8 values (0, 0.5, 1, 1.5, 2, 3, 4, 6) are used to + // quantize P. When P is fp6, only 32 values (0 .. 7.5) are used. Too large + // V values can create rare error outliers between host (no quantization) + // and device ("running" FA softmax + quantization), here we reduce max + // value by using smaller range of V scales. + gen_scales( + v_descale_host, VDataType{}, + ck_tile::is_any_of::value + ? 1 + : 3); + } else if (qscale.type == quant_scale_enum::pertensor) { + float q_dtype_max = + ck_tile::type_convert(ck_tile::numeric::max()); + float k_dtype_max = + ck_tile::type_convert(ck_tile::numeric::max()); + float v_dtype_max = + ck_tile::type_convert(ck_tile::numeric::max()); + + float qkv_max = 3.f; + q_descale_host(0) = qkv_max / q_dtype_max; + k_descale_host(0) = qkv_max / k_dtype_max; + v_descale_host(0) = qkv_max / v_dtype_max; + } else if (qscale.type == quant_scale_enum::blockscale) { + float q_dtype_max = + ck_tile::type_convert(ck_tile::numeric::max()); + float k_dtype_max = + ck_tile::type_convert(ck_tile::numeric::max()); + float v_dtype_max = + ck_tile::type_convert(ck_tile::numeric::max()); + + float qkv_max = 3.f; + float max_descale_q = qkv_max / q_dtype_max; + float max_descale_k = qkv_max / k_dtype_max; + float max_descale_v = qkv_max / v_dtype_max; + + ck_tile::FillUniformDistribution{max_descale_q * 0.8f, max_descale_q, + next_seed()}(q_descale_host); + ck_tile::FillUniformDistribution{max_descale_k * 0.8f, max_descale_k, + next_seed()}(k_descale_host); + ck_tile::FillUniformDistribution{max_descale_v * 0.8f, max_descale_v, + next_seed()}(v_descale_host); + } + iota_shuffle(block_table_host.begin(), block_table_host.end(), 0, random_engine); iota_shuffle(cache_batch_idx_host.begin(), cache_batch_idx_host.end(), 0, random_engine); - + if (init_sink_value != 0) { + // Keep the sink small and centered near zero. A sink far above rowmax makes + // lse_new = log(exp(lse_old) + exp(sink)) saturate to sink, so P_sink -> 1 + // while every regular token weight is scaled by exp(lse_old - lse_new) -> + // 0. The output then collapses toward zero and stops depending on which + // K/V, bias or randval tiles were read, which hides tile-window offset bugs + // in the sink path. + ck_tile::FillUniformDistribution{ + -1.f, 1.f, next_seed()}(sink_host); + } ck_tile::DeviceMem q_buf(q_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem k_buf(k_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem v_buf(v_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem sink_buf(sink_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem knew_buf(knew_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem vnew_buf(vnew_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem bias_buf(bias_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem q_descale_buf( + q_descale_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem k_descale_buf( + k_descale_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem v_descale_buf( + v_descale_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem block_scale_seqstart_q_buf( + block_scale_seqstart_q_host.size() * sizeof(int32_t)); + ck_tile::DeviceMem block_scale_seqstart_k_buf( + block_scale_seqstart_k_host.size() * sizeof(int32_t)); + ck_tile::DeviceMem scale_seqstart_v_buf(seqstart_v_scale_host.size() * + sizeof(int32_t)); ck_tile::DeviceMem lse_acc_buf( lse_acc_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem o_acc_buf(o_acc_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem lse_buf(lse_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem o_buf(o_host.get_element_space_size_in_bytes()); - ck_tile::DeviceMem seqstart_q(seqstart_q_host.size() * sizeof(int32_t)); - ck_tile::DeviceMem seqstart_k(seqstart_k_host.size() * sizeof(int32_t)); + ck_tile::DeviceMem seqstart_q_buf(seqstart_q_host.size() * sizeof(int32_t)); + ck_tile::DeviceMem seqstart_k_buf(seqstart_k_host.size() * sizeof(int32_t)); ck_tile::DeviceMem seqstart_q_padded_buf( seqstart_q_with_padding_host.empty() ? 0 @@ -680,11 +999,13 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::DeviceMem seqlen_q_buf( has_group_q_padding ? seqlen_qs.size() * sizeof(int32_t) : 0); // Buffers for key/value per-sequence logical (unpadded) lengths (used in - // batch mode with kvcache or group mode with padding enabled) - ck_tile::DeviceMem seqlen_k_buf((mode == mode_enum::batch && use_kvcache) || - has_group_k_padding - ? seqlen_ks.size() * sizeof(int32_t) - : 0); + // batch mode with kvcache or group mode with padding enabled) batch_prefill + // (group+kvcache) also needs per-batch seqlen_k for VLLM_BLOCK_TABLE_2D + const bool need_seqlen_k_buf = (mode == mode_enum::batch && use_kvcache) || + has_group_k_padding || + (mode == mode_enum::group && use_kvcache); + ck_tile::DeviceMem seqlen_k_buf( + need_seqlen_k_buf ? seqlen_ks.size() * sizeof(int32_t) : 0); ck_tile::DeviceMem cu_seqlen_q_buf( cuq_cum.empty() ? 0 : cuq_cum.size() * sizeof(ck_tile::index_t)); ck_tile::DeviceMem cu_seqlen_kv_buf( @@ -706,89 +1027,22 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, ck_tile::DeviceMem cache_batch_idx_buf( cache_batch_idx_host.get_element_space_size_in_bytes()); - float scale_p = 1.f; - float scale_o = 1.f; - if (squant) { - float q_dtype_max = - ck_tile::type_convert(ck_tile::numeric::max()); - float k_dtype_max = - ck_tile::type_convert(ck_tile::numeric::max()); - float v_dtype_max = - ck_tile::type_convert(ck_tile::numeric::max()); - float p_dtype_max = v_dtype_max; // assume p and v is the same type - // Q tensor - { - float max_value = - ck_tile::type_convert(ck_tile::numeric::min()); - q_host.ForEach([&](auto &self, auto idx) { - float val = ck_tile::type_convert(self(idx)); - if (val > max_value) - max_value = val; - }); - - float scale = q_dtype_max / max_value; - - q_host.ForEach([&](auto &self, auto idx) { - float val = ck_tile::type_convert(self(idx)); - self(idx) = ck_tile::type_convert(val * scale); - }); - scale_s = scale_s / scale; - } - - // K tensor - { - float max_value = - ck_tile::type_convert(ck_tile::numeric::min()); - k_host.ForEach([&](auto &self, auto idx) { - float val = ck_tile::type_convert(self(idx)); - if (val > max_value) - max_value = val; - }); - float scale = k_dtype_max / max_value; - k_host.ForEach([&](auto &self, auto idx) { - float val = ck_tile::type_convert(self(idx)); - self(idx) = ck_tile::type_convert(val * scale); - }); - scale_s = scale_s / scale; - } - - // V tensor - { - float max_value = - ck_tile::type_convert(ck_tile::numeric::min()); - v_host.ForEach([&](auto &self, auto idx) { - float val = ck_tile::type_convert(self(idx)); - if (val > max_value) - max_value = val; - }); - - float scale = k_dtype_max / max_value; - v_host.ForEach([&](auto &self, auto idx) { - float val = ck_tile::type_convert(self(idx)); - self(idx) = ck_tile::type_convert(val * scale); - }); - - scale_o = (1.0 / p_dtype_max) / scale; - } - - scale_p = p_dtype_max; - - if constexpr (std::is_same_v) { - float o_dtype_max = - ck_tile::type_convert(ck_tile::numeric::max()); - scale_o = scale_o * o_dtype_max / max_o; - } - } - q_buf.ToDevice(q_host.data()); k_buf.ToDevice(k_host.data()); v_buf.ToDevice(v_host.data()); + sink_buf.ToDevice(sink_host.data()); knew_buf.ToDevice(knew_host.data()); vnew_buf.ToDevice(vnew_host.data()); bias_buf.ToDevice(bias_host.data()); - seqstart_q.ToDevice(seqstart_q_host.data()); - // Keep logical starts in seqstart_k; pass padded K via separate pointer - seqstart_k.ToDevice(seqstart_k_host.data()); + q_descale_buf.ToDevice(q_descale_host.data()); + k_descale_buf.ToDevice(k_descale_host.data()); + v_descale_buf.ToDevice(v_descale_host.data()); + block_scale_seqstart_q_buf.ToDevice(block_scale_seqstart_q_host.data()); + block_scale_seqstart_k_buf.ToDevice(block_scale_seqstart_k_host.data()); + scale_seqstart_v_buf.ToDevice(seqstart_v_scale_host.data()); + seqstart_q_buf.ToDevice(seqstart_q_host.data()); + // Keep logical starts in seqstart_k_buf; pass padded K via separate pointer + seqstart_k_buf.ToDevice(seqstart_k_host.data()); seqstart_q_padded_buf.ToDevice(seqstart_q_with_padding_host.empty() ? nullptr : seqstart_q_with_padding_host.data()); @@ -797,10 +1051,7 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, cu_seqlen_q_buf.ToDevice(cuq_cum.empty() ? nullptr : cuq_cum.data()); cu_seqlen_kv_buf.ToDevice(cukv_cum.empty() ? nullptr : cukv_cum.data()); seqlen_q_buf.ToDevice(has_group_q_padding ? seqlen_qs.data() : nullptr); - seqlen_k_buf.ToDevice((mode == mode_enum::batch && use_kvcache) || - has_group_k_padding - ? seqlen_ks.data() - : nullptr); + seqlen_k_buf.ToDevice(need_seqlen_k_buf ? seqlen_ks.data() : nullptr); cache_seqlen_k_buf.ToDevice(need_append_kvcache ? cache_seqlen_ks.data() : nullptr); rotary_cos_buf.ToDevice(rotary_cos_host.data()); @@ -830,7 +1081,7 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, std::to_string(seqlen_kpads[0]) + ")")) << ", d:" << hdim_q << "/" << hdim_v << ", scale_s:" << scale_s << ", bias:" << bias << ", p_drop:" << p_drop << ", lse:" << lse - << ", squant:" << squant << ", mask:" << mask + << ", qscale:" << qscale << ", mask:" << mask << ", v:" << (is_v_rowmajor ? "r" : "c"); #if CK_TILE_FMHA_FWD_APPENDKV_API if (0 < rotary_dim) { @@ -909,20 +1160,40 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, traits.has_logits_soft_cap = 0.f < logits_soft_cap; traits.mask_type = mask.type; traits.bias_type = bias.type; + traits.has_sink = (mask.sink > 0 || init_sink_value != 0) ? true : false; traits.has_lse = lse; - traits.do_fp8_static_quant = squant; if constexpr (std::is_same_v>) { traits.has_dropout = (p_drop > 0.0f); + traits.qscale_type = qscale.type; } else if constexpr (std::is_same_v>) { traits.use_pagedkv = (0 < page_block_size); + } else if constexpr (std::is_same_v>) { + traits.has_dropout = (p_drop > 0.0f); + traits.qscale_type = qscale.type; + traits.kv_memory_layout = + ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::LINEAR_LAYOUT; + traits.kv_lookup_table = + ck_tile::BlockAttentionKVCacheLookupTableEnum::VLLM_BLOCK_TABLE_2D; + traits.page_size = page_block_size; } } }; - const auto init_args = [&, k_paddings_ = seqlen_kpads](auto &args) { + // Pack-GQA: print status (dimensions already computed earlier for buffer + // allocation) + if (pack_gqa_nhead != nhead) { + std::cout << "[Pack-GQA] Enabled: nhead " << nhead << "->" << pack_gqa_nhead + << ", seqlen_q " << shape_seqlen_q << "->" << pack_gqa_seqlen_q + << std::endl; + } + + const auto init_args = [&, pack_gqa_nhead_ = pack_gqa_nhead, + pack_gqa_seqlen_q_ = pack_gqa_seqlen_q, + k_paddings_ = seqlen_kpads](auto &args) { /// NOTE: we broadcast bias from [1, 1, seqlen_q, seqlen_k] to [batch, /// nhead, seqlen_q, /// seqlen_k] in this example, hence both the 'batch_stride_bias' & @@ -952,7 +1223,7 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, const ck_tile::index_t stride_o = (o_perm ? hdim_v : nhead * hdim_v); // setup nhead_stride_* arguments const ck_tile::index_t nhead_stride_q = - (i_perm ? shape_seqlen_q * hdim_q : hdim_q); + (i_perm ? pack_gqa_seqlen_q_ * hdim_q : hdim_q); const ck_tile::index_t nhead_stride_k = (0 < page_block_size ? (i_perm ? page_block_size * hdim_q : hdim_q) : (i_perm ? shape_seqlen_k * hdim_q : hdim_q)); @@ -977,15 +1248,20 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, const ck_tile::index_t nhead_stride_bias = (i_perm ? 0 * shape_seqlen_q * max_seqlen_k : 0 * max_seqlen_k); const ck_tile::index_t nhead_stride_randval = - (shape_seqlen_q * max_seqlen_k); - const ck_tile::index_t nhead_stride_lse = shape_seqlen_q; - const ck_tile::index_t nhead_stride_lse_acc = (num_splits * shape_seqlen_q); + (pack_gqa_seqlen_q_ * max_seqlen_k); + const ck_tile::index_t nhead_stride_lse = pack_gqa_seqlen_q_; + const ck_tile::index_t nhead_stride_lse_acc = + (num_splits * pack_gqa_seqlen_q_); const ck_tile::index_t nhead_stride_o_acc = - (num_splits * shape_seqlen_q * hdim_v); + (num_splits * pack_gqa_seqlen_q_ * hdim_v); const ck_tile::index_t nhead_stride_o = - (o_perm ? shape_seqlen_q * hdim_v : hdim_v); + (o_perm ? pack_gqa_seqlen_q_ * hdim_v : hdim_v); + const ck_tile::index_t nhead_stride_q_descale = num_block_scale_q; + const ck_tile::index_t nhead_stride_k_descale = num_block_scale_kv; + const ck_tile::index_t nhead_stride_v_descale = num_block_scale_kv; // setup batch_stride_* arguments - const ck_tile::index_t batch_stride_q = (nhead * shape_seqlen_q * hdim_q); + const ck_tile::index_t batch_stride_q = + (pack_gqa_nhead_ * pack_gqa_seqlen_q_ * hdim_q); const ck_tile::index_t batch_stride_k = (0 < page_block_size ? (nhead_k * page_block_size * hdim_q) : (nhead_k * shape_seqlen_k * hdim_q)); @@ -997,29 +1273,45 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, const ck_tile::index_t batch_stride_bias = (0 * nhead * shape_seqlen_q * max_seqlen_k); const ck_tile::index_t batch_stride_randval = - (nhead * shape_seqlen_q * max_seqlen_k); - const ck_tile::index_t batch_stride_lse = (nhead * shape_seqlen_q); + (pack_gqa_nhead_ * pack_gqa_seqlen_q_ * max_seqlen_k); + const ck_tile::index_t batch_stride_lse = + (pack_gqa_nhead_ * pack_gqa_seqlen_q_); const ck_tile::index_t batch_stride_lse_acc = - (nhead * num_splits * shape_seqlen_q); + (pack_gqa_nhead_ * num_splits * pack_gqa_seqlen_q_); const ck_tile::index_t batch_stride_o_acc = - (nhead * num_splits * shape_seqlen_q * hdim_v); - const ck_tile::index_t batch_stride_o = (nhead * shape_seqlen_q * hdim_v); + (pack_gqa_nhead_ * num_splits * pack_gqa_seqlen_q_ * hdim_v); + const ck_tile::index_t batch_stride_o = + (pack_gqa_nhead_ * pack_gqa_seqlen_q_ * hdim_v); const ck_tile::index_t batch_stride_block_table = (max_num_page_blocks / batch); + const ck_tile::index_t batch_stride_q_descale = num_block_scale_q * nhead; + const ck_tile::index_t batch_stride_k_descale = + num_block_scale_kv * nhead_k; + const ck_tile::index_t batch_stride_v_descale = + num_block_scale_kv * nhead_k; // setup split_stride_* arguments (only used in split-kv kernel) - const ck_tile::index_t split_stride_lse_acc = (shape_seqlen_q); - const ck_tile::index_t split_stride_o_acc = (shape_seqlen_q * hdim_v); + const ck_tile::index_t split_stride_lse_acc = (pack_gqa_seqlen_q_); + const ck_tile::index_t split_stride_o_acc = (pack_gqa_seqlen_q_ * hdim_v); args.q_ptr = q_buf.GetDeviceBuffer(); args.k_ptr = k_buf.GetDeviceBuffer(); args.v_ptr = v_buf.GetDeviceBuffer(); - + if (init_sink_value != 0) + args.sink_ptr = sink_buf.GetDeviceBuffer(); + else + args.sink_ptr = nullptr; args.batch = batch; - args.seqlen_q = shape_seqlen_q; // unused in group mode + args.seqlen_q = + pack_gqa_seqlen_q_; // unused in group mode, or packed seqlen args.hdim_q = hdim_q; args.hdim_v = hdim_v; - args.nhead_q = nhead; + args.nhead_q = pack_gqa_nhead_; args.nhead_k = nhead_k; + if constexpr (std::is_same_v>) { + args.num_head_q_total = pack_gqa_nhead_; + args.head_start = 0; + args.selected_kernel_name = selected_kernel_name; + } args.stride_q = stride_q; args.stride_k = stride_k; @@ -1071,11 +1363,9 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, args.seqlen_k = shape_seqlen_k; // unused in group mode (or kvcache enabled) - args.max_seqlen_q = max_seqlen_q; + args.max_seqlen_q = pack_gqa_seqlen_q_; // use packed seqlen for grid size args.scale_s = scale_s; - args.scale_p = scale_p; - args.scale_o = scale_o; args.logits_soft_cap = logits_soft_cap; @@ -1092,10 +1382,71 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, args.window_size_left = mask.left; args.window_size_right = mask.right; + args.sink_size = mask.sink; args.mask_type = static_cast(mask.type); if constexpr (std::is_same_v>) { + if (qscale.type == quant_scale_enum::pertensor) { + args.q_descale_ptr = q_descale_buf.GetDeviceBuffer(); + args.k_descale_ptr = k_descale_buf.GetDeviceBuffer(); + args.v_descale_ptr = v_descale_buf.GetDeviceBuffer(); + } else if (qscale.type == quant_scale_enum::blockscale) { + args.q_descale_ptr = + reinterpret_cast(q_descale_buf.GetDeviceBuffer()); + args.k_descale_ptr = + reinterpret_cast(k_descale_buf.GetDeviceBuffer()); + args.v_descale_ptr = + reinterpret_cast(v_descale_buf.GetDeviceBuffer()); + + args.block_scale_seqstart_q_ptr = + (mode == mode_enum::group + ? block_scale_seqstart_q_buf.GetDeviceBuffer() + : nullptr); + args.block_scale_seqstart_k_ptr = + (mode == mode_enum::group + ? block_scale_seqstart_k_buf.GetDeviceBuffer() + : nullptr); + + args.nhead_stride_q_descale = nhead_stride_q_descale; + args.nhead_stride_k_descale = nhead_stride_k_descale; + args.nhead_stride_v_descale = nhead_stride_v_descale; + + args.batch_stride_q_descale = batch_stride_q_descale; + args.batch_stride_k_descale = batch_stride_k_descale; + args.batch_stride_v_descale = batch_stride_v_descale; + + args.block_scale_size_q = block_scale_size_q_; + args.block_scale_size_kv = block_scale_size_kv_; + } else if (qscale.type == quant_scale_enum::mx) { + args.q_descale_ptr = q_descale_buf.GetDeviceBuffer(); + args.k_descale_ptr = k_descale_buf.GetDeviceBuffer(); + args.v_descale_ptr = v_descale_buf.GetDeviceBuffer(); + + args.stride_q_descale = + (i_perm ? hdim_q_scale : nhead * hdim_q_scale); + args.stride_k_descale = + (i_perm ? hdim_q_scale : nhead_k * hdim_q_scale); + args.stride_v_descale = + (i_perm ? shape_seqlen_v_scale : nhead_k * shape_seqlen_v_scale); + args.nhead_stride_q_descale = + (i_perm ? shape_seqlen_q * hdim_q_scale : hdim_q_scale); + args.nhead_stride_k_descale = + (i_perm ? shape_seqlen_k * hdim_q_scale : hdim_q_scale); + args.nhead_stride_v_descale = + (i_perm ? hdim_v * shape_seqlen_v_scale : shape_seqlen_v_scale); + if (mode == mode_enum::group) { + args.seqstart_v_scale_ptr = scale_seqstart_v_buf.GetDeviceBuffer(); + } else { + args.batch_stride_q_descale = + (nhead * shape_seqlen_q * hdim_q_scale); + args.batch_stride_k_descale = + (nhead_k * shape_seqlen_k * hdim_q_scale); + args.batch_stride_v_descale = + (nhead_k * hdim_v * shape_seqlen_v_scale); + } + } + args.rand_val_ptr = randval_buf.GetDeviceBuffer(); args.stride_randval = stride_randval; @@ -1121,11 +1472,11 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, args.seqstart_q_ptr = has_group_q_padding && !seqstart_q_with_padding_host.empty() ? seqstart_q_padded_buf.GetDeviceBuffer() - : seqstart_q.GetDeviceBuffer(); + : seqstart_q_buf.GetDeviceBuffer(); args.seqstart_k_ptr = has_group_k_padding && !seqstart_k_with_padding_host.empty() ? seqstart_k_padded_buf.GetDeviceBuffer() - : seqstart_k.GetDeviceBuffer(); + : seqstart_k_buf.GetDeviceBuffer(); // Logical (unpadded) per-sequence lengths, used when padding is // enabled @@ -1186,9 +1537,11 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, args.split_stride_o_acc = split_stride_o_acc; args.seqstart_q_ptr = - (mode == mode_enum::group ? seqstart_q.GetDeviceBuffer() : nullptr); + (mode == mode_enum::group ? seqstart_q_buf.GetDeviceBuffer() + : nullptr); args.seqstart_k_ptr = - (mode == mode_enum::group ? seqstart_k.GetDeviceBuffer() : nullptr); + (mode == mode_enum::group ? seqstart_k_buf.GetDeviceBuffer() + : nullptr); args.seqlen_k_ptr = ((mode == mode_enum::batch && use_kvcache) || 0 <= k_paddings_[0] ? seqlen_k_buf.GetDeviceBuffer() @@ -1206,13 +1559,78 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, : nullptr); args.seqstart_q_ptr = - (mode == mode_enum::group ? seqstart_q.GetDeviceBuffer() : nullptr); + (mode == mode_enum::group ? seqstart_q_buf.GetDeviceBuffer() + : nullptr); args.seqstart_k_ptr = - (mode == mode_enum::group ? seqstart_k.GetDeviceBuffer() : nullptr); + (mode == mode_enum::group ? seqstart_k_buf.GetDeviceBuffer() + : nullptr); args.seqlen_k_ptr = ((mode == mode_enum::batch && use_kvcache) || 0 <= k_paddings_[0] ? seqlen_k_buf.GetDeviceBuffer() : nullptr); + } else if constexpr (std::is_same_v>) { + // Fields already set by the outer else block above: + // bias_ptr, lse_ptr, o_ptr, seqlen_k, max_seqlen_q, scale_s, + // logits_soft_cap, stride_bias/o, nhead/batch stride for bias/lse/o, + // window_size_left/right, sink_size, mask_type. + + // scale_p/scale_o: batch_prefill-specific fields absent from + // fmha_fwd_args. + args.scale_p = 1.f; + args.scale_o = 1.f; + + // Dropout fields: the outer fmha_fwd_args branch sets these; set them + // here for batch_prefill since it takes a separate inner branch. + args.rand_val_ptr = randval_buf.GetDeviceBuffer(); + args.stride_randval = stride_randval; + args.nhead_stride_randval = nhead_stride_randval; + args.batch_stride_randval = batch_stride_randval; + args.p_drop = p_drop; + args.s_randval = s_randval; + if (drop_prefs) + args.drop_seed_offset = + std::make_pair(drop_seed_buf.GetDeviceBuffer(), + drop_offset_buf.GetDeviceBuffer()); + else + args.drop_seed_offset = std::make_pair(drop_seed, drop_offset); + + // Paged KV: LINEAR_LAYOUT + VLLM_BLOCK_TABLE_2D + // block_table_buf: [batch, max_blocks_per_seq] of physical page ids + // seqlen_k_buf: [batch] of per-batch seqlen_k values + args.num_total_pages = max_num_page_blocks; + args.page_block_size = page_block_size; + args.kv_memory_layout = + ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::LINEAR_LAYOUT; + args.kv_lookup_table = + ck_tile::BlockAttentionKVCacheLookupTableEnum::VLLM_BLOCK_TABLE_2D; + args.kv_indptr = nullptr; + args.kv_page_indices = block_table_buf.GetDeviceBuffer(); + args.kv_last_page_lens = nullptr; + args.seqlen_k_ptr = seqlen_k_buf.GetDeviceBuffer(); + args.batch_stride_block_table = batch_stride_block_table; + + // group mode required: seqstart_q is prefix-sum of per-batch seqlen_q + args.seqstart_q_ptr = seqstart_q_buf.GetDeviceBuffer(); + + // batch_prefill LINEAR_LAYOUT strides for runner's K layout + // [max_num_page_blocks, nhead_k, page_block_size, hdim]: + // stride_k = hdim_q (token stride within one head's + // page slice) nhead_stride_k = page_block_size * hdim_q (head + // stride) batch_stride_k = nhead_k * page_block_size * hdim_q (page + // stride, already set) + args.stride_k = hdim_q; + args.nhead_stride_k = page_block_size * hdim_q; + // V is row-major, same layout convention + args.stride_v = hdim_v; + args.nhead_stride_v = page_block_size * hdim_v; + + // descale: not used for fp16/bf16 + args.q_descale_ptr = nullptr; + args.k_descale_ptr = nullptr; + args.v_descale_ptr = nullptr; + args.nblock_stride_kv_block_descale = 0; + args.nhead_stride_kv_block_descale = 0; } } }; @@ -1238,6 +1656,21 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, } auto run_fwd = [&](const ck_tile::stream_config &sc) { +#if CK_TILE_FMHA_FWD_BATCH_PREFILL_API + // batch_prefill: group mode + paged KV, tested against the same CPU + // reference + if (1 == num_splits && use_kvcache && mode == mode_enum::group) { + fmha_batch_prefill_traits bp_traits; + init_traits(bp_traits); + + fmha_batch_prefill_args bp_args; + init_args(bp_args); + + const float ave_time = fmha_batch_prefill(bp_traits, bp_args, sc); + if (ave_time >= 0.0f) + return ave_time; + } +#endif // CK_TILE_FMHA_FWD_BATCH_PREFILL_API #if CK_TILE_FMHA_FWD_PAGEDKV_API if (1 == num_splits && use_kvcache) { fmha_fwd_pagedkv_traits fmha_pagedkv_traits; @@ -1276,7 +1709,68 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, return fmha_fwd(fmha_traits, fmha_args, sc); }; - const float fwd_ave_time = run_fwd(stream_config); + + float fwd_ave_time = -1.0f; +#if CK_TILE_FMHA_ENABLE_HEAD_GROUPING + const bool allow_head_grouping = + !i_perm && !use_kvcache && (num_splits <= 1) && !need_append_kvcache && + (mode == mode_enum::batch || mode == mode_enum::group); + + if (allow_head_grouping) { + if (fmha_fwd_head_grouping::disabled_by_env()) { + if (fmha_fwd_head_grouping::log_enabled()) + std::cout << "[LLC Head Grouping] disabled by env" << std::endl; + } else { + const auto group_size_opt = fmha_fwd_head_grouping::get_head_group_size( + nhead, nhead_k, batch, max_seqlen_k, hdim_q, hdim_v, + sizeof(KDataType), sizeof(VDataType)); + + if (group_size_opt.has_value() && group_size_opt.value() < nhead) { + if (fmha_fwd_head_grouping::log_enabled()) { + const std::string arch = ck_tile::get_device_name(); + const size_t llc_bytes = + fmha_fwd_head_grouping::get_llc_cache_bytes(arch); + const ck_tile::index_t gqa_ratio = + (nhead_k > 0 ? (nhead / nhead_k) : 1); + const ck_tile::index_t group_sz = group_size_opt.value(); + const ck_tile::index_t n_groups = + ck_tile::integer_divide_ceil(nhead, group_sz); + std::cout << "[LLC Head Grouping] enabled" << std::endl; + std::cout << "[LLC Head Grouping] arch=" + << (arch.empty() ? "unknown" : arch) + << " llc_mb=" << (llc_bytes / (1024ull * 1024ull)) + << " nhead_q=" << nhead << " nhead_k=" << nhead_k + << " gqa_ratio=" << gqa_ratio << " group_size=" << group_sz + << " groups=" << n_groups << std::endl; + } + fmha_fwd_traits fmha_traits; + init_traits(fmha_traits); + + fmha_fwd_args fmha_args; + init_args(fmha_args); + + fwd_ave_time = fmha_fwd_head_grouping::run_fwd_head_grouped< + QDataType, KDataType, VDataType, ODataType, BiasDataType, + LSEDataType, RandValOutputDataType>( + stream_config, fmha_traits, fmha_args, nhead, nhead_k, + group_size_opt.value(), qscale.type == quant_scale_enum::blockscale, + [&](const auto &traits, auto &args, const auto &sc) { + return fmha_fwd(traits, args, sc); + }); + } else if (fmha_fwd_head_grouping::log_enabled()) { + std::cout + << "[LLC Head Grouping] skipped (group_size not set or >= nhead)" + << std::endl; + } + } + } else if (fmha_fwd_head_grouping::log_enabled()) { + std::cout << "[LLC Head Grouping] disabled by conditions/layout" + << std::endl; + } +#endif + + if (fwd_ave_time < 0.0f) + fwd_ave_time = run_fwd(stream_config); if (fwd_ave_time < 0.0f) { std::cout << ", not supported yet" << std::flush << std::endl; return fwd_result::no_instance; @@ -1364,25 +1858,38 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, lse_buf.FromDevice(lse_host.data()); randval_buf.FromDevice(randval_host.data()); - constexpr bool supports_squant = + constexpr bool supports_qscale = std::is_same_v || std::is_same_v || std::is_same_v; + float scale_s_host = scale_s; + float scale_p_host = 1.0f; + float scale_o_host = 1.0f; + + if constexpr (!is_mx) { + if (qscale.type == quant_scale_enum::pertensor) { + scale_s_host = scale_s * q_descale_host(0) * k_descale_host(0); + scale_p_host = + ck_tile::type_convert(ck_tile::numeric::max()); + scale_o_host = v_descale_host(0) / scale_p_host; + } + } + auto p_compute_element_func = [&]() { - if constexpr (supports_squant) - return ck_tile::scales{scale_p}; + if constexpr (supports_qscale) + return ck_tile::scales{scale_p_host}; else return ck_tile::identity{}; }(); auto oacc_element_func = [&]() { if constexpr (std::is_same_v && - supports_squant) - return ck_tile::composes(ck_tile::saturates{}, - ck_tile::scales{scale_o}); - else if constexpr (supports_squant) - return ck_tile::scales{scale_o}; + supports_qscale) + return ck_tile::make_composes(ck_tile::saturates{}, + ck_tile::scales{scale_o_host}); + else if constexpr (supports_qscale) + return ck_tile::scales{scale_o_host}; else return ck_tile::identity{}; }(); @@ -1463,7 +1970,8 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, [&](auto &self, auto i) { self(i) = q_host_ref_ro(i); }); } #endif -#if CK_TILE_FMHA_FWD_SPLITKV_API || CK_TILE_FMHA_FWD_PAGEDKV_API +#if CK_TILE_FMHA_FWD_SPLITKV_API || CK_TILE_FMHA_FWD_PAGEDKV_API || \ + CK_TILE_FMHA_FWD_BATCH_PREFILL_API if (0 < page_block_size) { // clang-format off if(i_perm) k_host_ref.ForEach([&](auto& self, auto i) { self(i) = k_host(block_table_host(wb, i[1] / page_block_size), i[0] / nr, i[1] % page_block_size, i[2]); }); @@ -1510,7 +2018,8 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, }); } #endif -#if CK_TILE_FMHA_FWD_SPLITKV_API || CK_TILE_FMHA_FWD_PAGEDKV_API +#if CK_TILE_FMHA_FWD_SPLITKV_API || CK_TILE_FMHA_FWD_PAGEDKV_API || \ + CK_TILE_FMHA_FWD_BATCH_PREFILL_API if (0 < page_block_size) { if (is_v_rowmajor) { // clang-format off @@ -1565,10 +2074,58 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, #endif // reference - ck_tile::reference_batched_gemm( - q_host_ref, k_host_ref, s_host_ref, ck_tile::identity{}, - ck_tile::identity{}, ck_tile::scales(scale_s)); + if constexpr (is_mx) { + ck_tile::HostTensor q_descale_host_ref( + {nhead, real_seqlen_q, hdim_q_scale}); + ck_tile::HostTensor k_descale_host_ref( + {nhead, real_seqlen_k, hdim_q_scale}); + + // clang-format off + if(i_perm) q_descale_host_ref.ForEach([&](auto& self, auto i) { self(i) = q_descale_host(b_idx, i[0], i[1] + query_offset, i[2]); }); + else q_descale_host_ref.ForEach([&](auto& self, auto i) { self(i) = q_descale_host(b_idx, i[1] + query_offset, i[0], i[2]); }); + // clang-format on + + // clang-format off + if(i_perm) k_descale_host_ref.ForEach([&](auto& self, auto i) { self(i) = k_descale_host(cache_b_idx, i[0] / nr, i[1] + key_offset, i[2]); }); + else k_descale_host_ref.ForEach([&](auto& self, auto i) { self(i) = k_descale_host(cache_b_idx, i[1] + key_offset, i[0] / nr, i[2]); }); + // clang-format on + + auto q_host_ref2 = + ck_tile::reference_batched_mx_descale( + q_host_ref, q_descale_host_ref, kQKScaleGranularity); + auto k_host_ref2 = + ck_tile::reference_batched_mx_descale( + k_host_ref, k_descale_host_ref, kQKScaleGranularity); + + ck_tile::reference_batched_gemm( + q_host_ref2, k_host_ref2, s_host_ref, ck_tile::identity{}, + ck_tile::identity{}, ck_tile::scales(scale_s_host)); + } else if (qscale.type == quant_scale_enum::blockscale) { + const ck_tile::index_t q_offset = + (mode == mode_enum::batch) ? 0 : block_scale_seqstart_q_host[wb]; + const ck_tile::index_t k_offset = + (mode == mode_enum::batch) ? 0 : block_scale_seqstart_k_host[wb]; + ck_tile::reference_batched_quant_gemm< + QDataType, KDataType, SaccDataType, SMPLComputeDataType>( + q_host_ref, k_host_ref, s_host_ref, ck_tile::idx_identity{}, + ck_tile::idx_identity{}, [&](auto idx, auto value) { + return value * scale_s * + q_descale_host(b_idx, std::get<0>(idx), + q_offset + std::get<1>(idx) / + block_scale_size_q_) * + k_descale_host(b_idx, std::get<0>(idx) / nr, + k_offset + std::get<2>(idx) / + block_scale_size_kv_); + }); + } else { + ck_tile::reference_batched_gemm( + q_host_ref, k_host_ref, s_host_ref, ck_tile::identity{}, + ck_tile::identity{}, ck_tile::scales(scale_s_host)); + } if (0.f < logits_soft_cap) { ck_tile::reference_unary_elementwise( s_host_ref, ck_tile::make_generic_attention_mask_from_lr_window< - FmhaMasks::GenericMask>(mask.left, mask.right, real_seqlen_q, - real_seqlen_k)); + FmhaMasks::GenericMask>(mask.left, mask.right, mask.sink, + real_seqlen_q, real_seqlen_k)); } else { // if left window size is negative, means causal // else means generic (for current batch) @@ -1647,70 +2204,242 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, s_host_ref, ck_tile::make_generic_attention_mask_from_lr_window< FmhaMasks::CausalMask>( - mask.left, mask.right, real_seqlen_q, real_seqlen_k, - mask.type == mask_enum::mask_top_left)); + mask.left, mask.right, mask.sink, real_seqlen_q, + real_seqlen_k, mask.type == mask_enum::mask_top_left)); else ck_tile::reference_batched_masking( s_host_ref, ck_tile::make_generic_attention_mask_from_lr_window< FmhaMasks::GenericMask>( - mask.left, mask.right, real_seqlen_q, real_seqlen_k, - mask.type == mask_enum::mask_top_left)); + mask.left, mask.right, mask.sink, real_seqlen_q, + real_seqlen_k, mask.type == mask_enum::mask_top_left)); } const ck_tile::HostTensor masked_s_host_ref = s_host_ref; - if (lse) { - ck_tile::reference_batched_softmax( - s_host_ref, p_host_ref, p_compute_element_func, lse_host_ref); + // Softmax is kept in fp32 first, then narrowed to PDataType. The + // per-tensor path below consumes the fp32 copy directly, see the comment + // at its GEMM call. + ck_tile::HostTensor p_host_ref_f32( + {nhead, real_seqlen_q, real_seqlen_k}); + if (init_sink_value != 0) { + // Create extended tensor with sink token + ck_tile::HostTensor s_with_sinks_ref( + {nhead, real_seqlen_q, real_seqlen_k + 1}); + + // Copy original attention scores and append sink values + copy_attention_scores_with_sink(s_host_ref, sink_host, s_with_sinks_ref, + nhead, real_seqlen_q, real_seqlen_k); + + // Compute softmax on extended tensor + ck_tile::HostTensor p_extended( + {nhead, real_seqlen_q, real_seqlen_k + 1}); + + if (lse) { + ck_tile::reference_batched_softmax< + SMPLComputeDataType, SMPLComputeDataType, SMPLComputeDataType>( + s_with_sinks_ref, p_extended, p_compute_element_func, + lse_host_ref); + } else { + ck_tile::reference_batched_softmax< + SMPLComputeDataType, SMPLComputeDataType, SMPLComputeDataType>( + s_with_sinks_ref, p_extended, p_compute_element_func); + } + + // Extract only the original columns (exclude sink token column) + p_host_ref_f32.ForEach([&](auto &self, auto idx) { + self(idx) = p_extended(idx[0], idx[1], idx[2]); + }); } else { - ck_tile::reference_batched_softmax( - s_host_ref, p_host_ref, p_compute_element_func); + // No sink tokens - compute softmax directly + if (lse) { + ck_tile::reference_batched_softmax< + SMPLComputeDataType, SMPLComputeDataType, SMPLComputeDataType>( + s_host_ref, p_host_ref_f32, p_compute_element_func, lse_host_ref); + } else { + ck_tile::reference_batched_softmax< + SMPLComputeDataType, SMPLComputeDataType, SMPLComputeDataType>( + s_host_ref, p_host_ref_f32, p_compute_element_func); + } } + p_host_ref.ForEach([&](auto &self, auto idx) { + self(idx) = ck_tile::type_convert(p_host_ref_f32(idx)); + }); + if (lse) { + ck_tile::HostTensor lse_host_result( + {nhead, real_seqlen_q}); + if (pack_gqa_nhead != nhead) { + lse_host_result.ForEach([&](auto &self, auto idx) { + ck_tile::index_t packed_head = idx[0] / nhead_ratio; + ck_tile::index_t packed_seq = + (idx[0] % nhead_ratio) * real_seqlen_q + idx[1] + query_offset; + self(idx) = lse_host(b_idx, packed_head, packed_seq); + }); + } else { + lse_host_result.ForEach([&](auto &self, auto idx) { + self(idx) = lse_host(b_idx, idx[0], idx[1] + query_offset); + }); + } - if (p_drop > 0) { - ck_tile::HostTensor randval_host_ref( - {nhead, real_seqlen_q, real_seqlen_k}); - ck_tile::reference_batched_dropout_randval(randval_host_ref, wb, - drop_seed, drop_offset); - ck_tile::reference_batched_dropout(p_host_ref, randval_host_ref, - p_undrop_in_uint8_t, rp_undrop); - - ck_tile::HostTensor randval_host_result( - {nhead, real_seqlen_q, real_seqlen_k}); - randval_host_result.ForEach([&](auto &self, const auto &idx) { - self(idx) = - randval_host(b_idx, idx[0], idx[1] + query_offset, idx[2]); - }); - masked_s_host_ref.ForEach([&](const auto &self, const auto &idx) { - // Ignore all masked values in validation check - if (std::isinf(self(idx))) { - randval_host_ref(idx) = 0; - randval_host_result(idx) = 0; - } - }); + // Use smaller rtol/atol as LSE is computed and stored in fp32, so there + // is no precision loss due to conversion bool cur_pass = - ck_tile::check_err(randval_host_result, randval_host_ref, - "DROPOUT RANDVAL Error: Incorrect results!"); + ck_tile::check_err(lse_host_result, lse_host_ref, + "LSE Error: Incorrect results!", 1e-4, 1e-4, + /* allow_infinity_ref = */ true); + pass &= cur_pass; if (!cur_pass) { - break; + std::cerr << "LSE mismatch found at batch: " << wb << std::endl + << "\tseqlen_q: " << real_seqlen_q << std::endl + << "\tseqlen_k: " << real_seqlen_k << std::endl + << "\tseqstart_q: " << seqstart_q_host << std::endl + << "\tseqstart_k: " << seqstart_k_host << std::endl; } } + if (p_drop > 0) { + const bool is_pack_gqa_active = (pack_gqa_nhead != nhead); + if (is_pack_gqa_active) { + const int nhead_ratio_local = nhead / nhead_k; + ck_tile::HostTensor randval_packed_ref( + {pack_gqa_nhead, pack_gqa_seqlen_q, real_seqlen_k}); + ck_tile::reference_batched_dropout_randval(randval_packed_ref, wb, + drop_seed, drop_offset); + + ck_tile::HostTensor randval_host_ref( + {nhead, real_seqlen_q, real_seqlen_k}); + randval_host_ref.ForEach([&](auto &self, const auto &idx) { + ck_tile::index_t h = idx[0]; + ck_tile::index_t s = idx[1]; + ck_tile::index_t k = idx[2]; + ck_tile::index_t packed_head = h / nhead_ratio_local; + ck_tile::index_t packed_seq = + (h % nhead_ratio_local) * real_seqlen_q + s; + self(idx) = randval_packed_ref(packed_head, packed_seq, k); + }); + + ck_tile::reference_batched_dropout(p_host_ref, randval_host_ref, + p_undrop_in_uint8_t, rp_undrop); + + ck_tile::HostTensor randval_host_result( + {pack_gqa_nhead, pack_gqa_seqlen_q, real_seqlen_k}); + randval_host_result.ForEach([&](auto &self, const auto &idx) { + self(idx) = + randval_host(b_idx, idx[0], idx[1] + query_offset, idx[2]); + }); + bool cur_pass = + ck_tile::check_err(randval_host_result, randval_packed_ref, + "DROPOUT RANDVAL Error: Incorrect results!"); + pass &= cur_pass; + if (!cur_pass) { + break; + } + } else { + ck_tile::HostTensor randval_host_ref( + {nhead, real_seqlen_q, real_seqlen_k}); + ck_tile::reference_batched_dropout_randval(randval_host_ref, wb, + drop_seed, drop_offset); + ck_tile::reference_batched_dropout(p_host_ref, randval_host_ref, + p_undrop_in_uint8_t, rp_undrop); + + ck_tile::HostTensor randval_host_result( + {nhead, real_seqlen_q, real_seqlen_k}); + randval_host_result.ForEach([&](auto &self, const auto &idx) { + self(idx) = + randval_host(b_idx, idx[0], idx[1] + query_offset, idx[2]); + }); + masked_s_host_ref.ForEach([&](const auto &self, const auto &idx) { + if (std::isinf(self(idx))) { + randval_host_ref(idx) = 0; + randval_host_result(idx) = 0; + } + }); + bool cur_pass = + ck_tile::check_err(randval_host_result, randval_host_ref, + "DROPOUT RANDVAL Error: Incorrect results!"); + pass &= cur_pass; + if (!cur_pass) { + break; + } + } + } + + if constexpr (is_mx) { + const ck_tile::index_t real_seqlen_v_scale = + seqstart_v_scale_host[wb + 1] - seqstart_v_scale_host[wb]; + const ck_tile::index_t v_scale_offset = + mode == mode_enum::batch ? 0 : seqstart_v_scale_host[wb]; - ck_tile::reference_batched_gemm( - p_host_ref, v_host_ref, o_host_ref, ck_tile::identity{}, - ck_tile::identity{}, oacc_element_func); + ck_tile::HostTensor v_descale_host_ref( + {nhead, hdim_v, real_seqlen_v_scale}); + + // clang-format off + if(i_perm) v_descale_host_ref.ForEach([&](auto& self, auto i) { self(i) = v_descale_host(cache_b_idx, i[0] / nr, i[1], i[2] + v_scale_offset); }); + else v_descale_host_ref.ForEach([&](auto& self, auto i) { self(i) = v_descale_host(cache_b_idx, i[1], i[0] / nr, i[2] + v_scale_offset); }); + // clang-format on + + auto v_host_ref2 = + ck_tile::reference_batched_mx_descale( + v_host_ref, v_descale_host_ref, kVScaleGranularity); + + // P is not quantized and then dequantized here (PDataType = float). + // On host softmax is computed for the whole row of S, while on device + // FA computes softmax and quantizes it in blocks of N0 values. + // Quantization on host would make reference results **less** precise + // than the device results for large seqlen_k! + + ck_tile::reference_batched_gemm( + p_host_ref, v_host_ref2, o_host_ref, ck_tile::identity{}, + ck_tile::identity{}, oacc_element_func); + } else if (qscale.type == quant_scale_enum::blockscale) { + const ck_tile::index_t v_offset = + (mode == mode_enum::batch) ? 0 : block_scale_seqstart_k_host[wb]; + ck_tile::reference_batched_quant_gemm( + p_host_ref, v_host_ref, o_host_ref, ck_tile::idx_identity{}, + [&](auto idx, auto value) { + return ck_tile::type_convert(value) * + v_descale_host(b_idx, std::get<0>(idx) / nr, + v_offset + std::get<2>(idx) / + block_scale_size_kv_); + }, + ck_tile::idx_identity{}); + } else if (qscale.type == quant_scale_enum::pertensor && !(p_drop > 0)) { + // Same reasoning as the mx branch above: quantizing P on the host makes + // the reference *less* precise than the device. The kernel quantizes + // exp(s - m), whose row maximum is 1 and therefore maps onto the top of + // the fp8 range, while the host would quantize the already normalized + // softmax, whose maximum is only 1/l of that range. The two land on + // different fp8 grid points, so with a peaked softmax (narrow sliding + // window) their rounding disagrees by several percent per weight and + // nothing averages it out. + ck_tile::reference_batched_gemm( + p_host_ref_f32, v_host_ref, o_host_ref, ck_tile::identity{}, + ck_tile::identity{}, oacc_element_func); + } else { + ck_tile::reference_batched_gemm( + p_host_ref, v_host_ref, o_host_ref, ck_tile::identity{}, + ck_tile::identity{}, oacc_element_func); + } ck_tile::HostTensor o_host_result( {nhead, real_seqlen_q, hdim_v}); // clang-format off // permute - if(o_perm) o_host_result.ForEach([&](auto& self, auto idx) { self(idx) = o_host(b_idx, idx[0], idx[1] + query_offset, idx[2]); }); + if(pack_gqa_nhead != nhead) + { + // Pack-GQA: o_host uses packed layout, unpack for comparison + o_host_result.ForEach([&](auto& self, auto idx) { + ck_tile::index_t packed_head = idx[0] / nhead_ratio; + ck_tile::index_t packed_seq = (idx[0] % nhead_ratio) * real_seqlen_q + idx[1] + query_offset; + self(idx) = o_host(b_idx, packed_head, packed_seq, idx[2]); + }); + } + else if(o_perm) o_host_result.ForEach([&](auto& self, auto idx) { self(idx) = o_host(b_idx, idx[0], idx[1] + query_offset, idx[2]); }); else o_host_result.ForEach([&](auto& self, auto idx) { self(idx) = o_host(b_idx, idx[1] + query_offset, idx[0], idx[2]); }); // clang-format on - auto [rtol, atol] = get_elimit(init_method); bool cur_pass = ck_tile::check_err( o_host_result, o_host_ref, @@ -1728,32 +2457,10 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, << std::endl << "\tquery_offset used: " << query_offset << std::endl << "\tkey_offset used: " << key_offset << std::endl; - - break; } - if (lse) { - ck_tile::HostTensor lse_host_result( - {nhead, real_seqlen_q}); - lse_host_result.ForEach([&](auto &self, auto idx) { - self(idx) = lse_host(b_idx, idx[0], idx[1] + query_offset); - }); - - cur_pass = - ck_tile::check_err(lse_host_result, lse_host_ref, - "LSE Error: Incorrect results!", rtol, atol, - /* allow_infinity_ref = */ true); - - pass &= cur_pass; - if (!cur_pass) { - std::cerr << "LSE mismatch found at batch: " << wb << std::endl - << "\tseqlen_q: " << real_seqlen_q << std::endl - << "\tseqlen_k: " << real_seqlen_k << std::endl - << "\tseqstart_q: " << seqstart_q_host << std::endl - << "\tseqstart_k: " << seqstart_k_host << std::endl; - - break; - } + if (!pass) { + break; } } @@ -1761,11 +2468,18 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, } if (json) { + const std::string qscale_name = + (qscale.type == quant_scale_enum::no_scale ? "no_scale" + : qscale.type == quant_scale_enum::pertensor ? "pertensor" + : qscale.type == quant_scale_enum::blockscale ? "blockscale" + : qscale.type == quant_scale_enum::kv_blockscale ? "kv_blockscale" + : qscale.type == quant_scale_enum::mx ? "mx" + : "unknown"); dump_fmha_fwd_json_results( *json, data_type, mode == mode_enum::batch ? "batch" : "group", io_layout(i_perm, o_perm), batch, nhead, nhead_k, seqlen_qs[0], seqlen_ks[0], seqlen_kpads[0], hdim_q, hdim_v, scale_s, p_drop, lse, - squant, + qscale_name, bias.type == bias_enum::elementwise_bias ? "elementwise_bias" : (bias.type == bias_enum::alibi ? "alibi" : "no_bias"), @@ -1774,3 +2488,6 @@ fmha_fwd_run(mode_enum mode, ck_tile::index_t batch, ck_tile::index_t nhead, return pass ? fwd_result::success : fwd_result::failure; } +#if __clang_major__ >= 23 +#pragma clang diagnostic pop +#endif diff --git a/Libraries/ComposableKernel/attention/fmha/generate.py b/Libraries/ComposableKernel/attention/fmha/generate.py index 30e94dd7c1..112d52896e 100644 --- a/Libraries/ComposableKernel/attention/fmha/generate.py +++ b/Libraries/ComposableKernel/attention/fmha/generate.py @@ -1,9 +1,11 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -# Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. # generate kernel instances to speed up compilation import argparse +import sys +import importlib.util from enum import IntEnum from pathlib import Path import pkgutil @@ -18,11 +20,23 @@ class HandlerId(IntEnum): WRITE_BLOBS = 1 +def _load_ops_module(importer, module_name: str): + """Load a submodule from codegen.ops, matching Loader.load_module behavior.""" + spec = importer.find_spec(module_name) + loader = spec.loader + if hasattr(loader, "load_module"): + return loader.load_module(module_name) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + loader.exec_module(module) + return module + + # inspect all modules under 'codegen.ops' and register API handlers ops = [] for importer, module_name, _ in pkgutil.iter_modules(codegen.ops.__path__): full_module_name = "%s.%s" % (codegen.ops.__name__, module_name) - ops.append(importer.find_spec(module_name).loader.load_module(module_name)) + ops.append(_load_ops_module(importer, module_name)) unwanted_prefix = "fmha_" handlers = dict( [ @@ -141,7 +155,9 @@ def list_blobs( + " 200-299: Only generate instance for Aiter(mha_varlen_fwd) integration\n" + " 300-399: Only generate instance for Aiter(mha_bwd) integration\n" + " 400-499: Only generate instance for Aiter(mha_varlen_bwd) integration\n" - + " 600-699: Only generate instance for aiter::mha_fwd && aiter::mha_fwd_splitkv && aiter::mha_bwd C++ api integration", + + " 600-699: Only generate instance for aiter::mha_fwd && aiter::mha_fwd_splitkv && aiter::mha_bwd C++ api integration\n" + + " 700: Only generate instance for TransformerEngine integration (fwd + bwd, fp16/bf16 only,\n" + + " invariants: row vlayout, has_lse, no skip/sink/logits/qscale)", ) parser.add_argument( diff --git a/Libraries/ComposableKernel/attention/fmha/mask.hpp b/Libraries/ComposableKernel/attention/fmha/mask.hpp index 3c13e7a55e..82c26135c2 100644 --- a/Libraries/ComposableKernel/attention/fmha/mask.hpp +++ b/Libraries/ComposableKernel/attention/fmha/mask.hpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once @@ -42,6 +23,7 @@ struct mask_info { ck_tile::index_t seqlen_k; ck_tile::index_t y, x; ck_tile::index_t left, right; // FA style SWA left/right + ck_tile::index_t sink; void serialize(std::ostream &os) const { if (type == mask_enum::no_mask) @@ -71,13 +53,14 @@ struct mask_info { ck_tile::index_t window_size = std::stoi(v); ck_tile::index_t left_size = -1; ck_tile::index_t right_size = 0; + ck_tile::index_t sink_size = 0; if (window_size > 0) { left_size = window_size / 2; right_size = window_size - 1 - left_size; } auto r = ck_tile::make_generic_attention_mask_coordinates_from_lr_window( - left_size, right_size, y_total, x_total, t == "xt"); + left_size, right_size, sink_size, y_total, x_total, t == "xt"); tmp.type = t == "xt" ? mask_enum::mask_top_left : mask_enum::mask_bottom_right; @@ -85,65 +68,93 @@ struct mask_info { tmp.x = r.at(ck_tile::number<1>{}); tmp.left = left_size; tmp.right = right_size; + tmp.sink = 0; } else if (t == "t" || t == "b" || t == "g") { auto found_1 = v.find(","); if (found_1 == std::string::npos) { throw std::invalid_argument("invalid mask value: " + str); } - ck_tile::index_t v0 = std::stoi(v.substr(0, found_1)); - ck_tile::index_t v1 = std::stoi(v.substr(found_1 + 1)); + tmp.type = mask_enum::window_generic; + ck_tile::index_t v0 = atoi(v.substr(0, found_1).c_str()); + auto found_2 = v.find(',', found_1 + 1); + ck_tile::index_t v1 = 0; + ck_tile::index_t sink = 0; + // ck_tile::index_t v1 = atoi(v.substr(found_1 + 1).c_str()); + // TODO: some validation if (t == "t") { + if (found_2 != std::string::npos) { + v1 = atoi(v.substr(found_1 + 1, found_2 - found_1 - 1).c_str()); + sink = atoi(v.substr(found_2 + 1).c_str()); + } else { + v1 = atoi(v.substr(found_1 + 1).c_str()); + sink = 0; + } tmp.type = mask_enum::mask_top_left; auto r = ck_tile::make_generic_attention_mask_coordinates_from_lr_window( - v0, v1, y_total, x_total, true); + v0, v1, sink, y_total, x_total, true); tmp.y = r.at(ck_tile::number<0>{}); tmp.x = r.at(ck_tile::number<1>{}); tmp.left = v0; tmp.right = v1; + tmp.sink = sink; } else if (t == "b") { + if (found_2 != std::string::npos) { + v1 = atoi(v.substr(found_1 + 1, found_2 - found_1 - 1).c_str()); + sink = atoi(v.substr(found_2 + 1).c_str()); + } else { + v1 = atoi(v.substr(found_1 + 1).c_str()); + sink = 0; + } tmp.type = mask_enum::mask_bottom_right; auto r = ck_tile::make_generic_attention_mask_coordinates_from_lr_window( - v0, v1, y_total, x_total, false); + v0, v1, sink, y_total, x_total, false); tmp.y = r.at(ck_tile::number<0>{}); tmp.x = r.at(ck_tile::number<1>{}); tmp.left = v0; tmp.right = v1; + tmp.sink = sink; } else if (t == "g") { tmp.type = mask_enum::window_generic; tmp.y = v0; tmp.x = v1; tmp.left = v0; // TODO: don't use this? tmp.right = v1; + tmp.sink = 0; } } else { throw std::invalid_argument("invalid mask value: " + str); } } else if (str == "0") { tmp.type = mask_enum::no_mask; + tmp.left = -1; + tmp.right = -1; + tmp.sink = 0; } else if (str == "1" || str == "t") { tmp.type = mask_enum::mask_top_left; tmp.y = seqlen_q; tmp.x = 1; tmp.left = -1; tmp.right = 0; + tmp.sink = 0; } else if (str == "2" || str == "b") { tmp.type = mask_enum::mask_bottom_right; tmp.y = seqlen_q; tmp.x = seqlen_k - seqlen_q + 1; tmp.left = -1; tmp.right = 0; + tmp.sink = 0; } else { throw std::invalid_argument("invalid mask value: " + str); } return tmp; } - ck_tile::index_t get_unmaskarea() const { + std::size_t get_unmaskarea() const { if (type == mask_enum::no_mask) - return seqlen_q * seqlen_k; - ck_tile::index_t area = 0; + return static_cast(seqlen_q) * seqlen_k; + std::size_t area = 0; for (ck_tile::index_t i_y = 0; i_y < seqlen_q; ++i_y) { ck_tile::index_t x_start = std::max(-y + i_y + 1, static_cast(0)); @@ -155,7 +166,8 @@ struct mask_info { return area; } - friend std::ostream &operator<<(std::ostream &os, const mask_info &mi) { + friend std::ostream &operator<<([[clang::lifetimebound]] std::ostream &os, + const mask_info &mi) { mi.serialize(os); return os; } diff --git a/Libraries/ComposableKernel/attention/fmha/quant.hpp b/Libraries/ComposableKernel/attention/fmha/quant.hpp new file mode 100644 index 0000000000..d1d20ae4ef --- /dev/null +++ b/Libraries/ComposableKernel/attention/fmha/quant.hpp @@ -0,0 +1,68 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/fmha.hpp" + +#if __clang_major__ >= 23 +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions" +#endif +// keep sync with BlockAttentionQuantScaleEnum +enum class quant_scale_enum { + no_scale = 0, + pertensor = 1, + blockscale = 2, + kv_blockscale = 3, // Q per-tensor, K/V per-page block scale + mx = 4, // Microscaling (MX) +}; + +struct quant_scale_info { + quant_scale_enum type; + + void serialize(std::ostream &os) const { + if (type == quant_scale_enum::no_scale) + os << "n"; + else if (type == quant_scale_enum::pertensor) + os << "pt"; + else if (type == quant_scale_enum::blockscale) + os << "bs"; + else if (type == quant_scale_enum::kv_blockscale) + os << "kvbs"; + else if (type == quant_scale_enum::mx) + os << "mx"; + } + + static quant_scale_info decode(std::string str) { + quant_scale_info info{quant_scale_enum::no_scale}; + if (str == "n" || str == "0") { + info.type = quant_scale_enum::no_scale; + } else if (str == "pt" || str == "1") { + info.type = quant_scale_enum::pertensor; + } else if (str == "bs" || str == "2") { + info.type = quant_scale_enum::blockscale; + } else if (str == "kvbs" || str == "3") { + info.type = quant_scale_enum::kv_blockscale; + } else if (str == "mx" || str == "4") { + info.type = quant_scale_enum::mx; + } else { + throw std::invalid_argument("invalid quant scale value: " + str); + } + return info; + } + + friend std::ostream &operator<<(std::ostream &os, + const quant_scale_info &qsi) { + qsi.serialize(os); + return os; + } +}; +#if __clang_major__ >= 23 +#pragma clang diagnostic pop +#endif diff --git a/Libraries/ComposableKernel/attention/fmha/rotary.hpp b/Libraries/ComposableKernel/attention/fmha/rotary.hpp index cd794ec394..b575cf021e 100644 --- a/Libraries/ComposableKernel/attention/fmha/rotary.hpp +++ b/Libraries/ComposableKernel/attention/fmha/rotary.hpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once @@ -27,6 +8,11 @@ #include #include + +#ifndef M_PI // Not there on windows... +#define M_PI 3.141592653589793238462643383279502884 +#endif + #include #include #include diff --git a/Libraries/ComposableKernel/attention/fmha/utils.hpp b/Libraries/ComposableKernel/attention/fmha/utils.hpp index af8234ed99..14b6c09ebc 100644 --- a/Libraries/ComposableKernel/attention/fmha/utils.hpp +++ b/Libraries/ComposableKernel/attention/fmha/utils.hpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/.gitignore b/Libraries/ComposableKernel/attention/fmha_fwd_v3/.gitignore deleted file mode 100644 index f32dd83d6d..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/.gitignore +++ /dev/null @@ -1 +0,0 @@ -ComposableKernel_ck_tile_fmha_fwd_v3 diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/CMakeLists.txt b/Libraries/ComposableKernel/attention/fmha_fwd_v3/CMakeLists.txt deleted file mode 100644 index 33af480b85..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/CMakeLists.txt +++ /dev/null @@ -1,92 +0,0 @@ -# MIT License -# -# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -set(example_name ComposableKernel_ck_tile_fmha_fwd_v3) -set(example_script test_ck_tile_fmha_fwd_v3.sh) - -cmake_minimum_required(VERSION 3.21 FATAL_ERROR) -project(${example_name} LANGUAGES CXX) - -if(CMAKE_SYSTEM_NAME STREQUAL "Windows") - message(STATUS "Composable Kernel examples are supported only on Linux.") - return() -endif() - -include("${CMAKE_CURRENT_LIST_DIR}/../../../../cmake/HipPlatform.cmake") -select_gpu_language() -enable_language(${ROCM_EXAMPLES_GPU_LANGUAGE}) -select_hip_platform() - -include("${CMAKE_CURRENT_LIST_DIR}/../../../../cmake/FilterHIPArchitectures.cmake") -filter_hip_architectures("CK Tile FMHA v3 examples" "gfx950" SHOULD_SKIP) -if(SHOULD_SKIP) - return() -endif() - -if(ROCM_EXAMPLES_GPU_LANGUAGE STREQUAL "CUDA") - message(STATUS "Composable Kernel examples are not supported for CUDA.") - return() -endif() - -include("${CMAKE_CURRENT_LIST_DIR}/../../../../cmake/ROCmPath.cmake") - -find_package(composable_kernel REQUIRED) -include("${CMAKE_CURRENT_LIST_DIR}/../../../../cmake/FindCKTile.cmake") - -find_package(hip REQUIRED) - -add_executable(${example_name} - example_fmha_fwd_v3.cpp - fmha_fwd_v3.cpp - instances/fmha_fwd_v3_d128_bf16_mask.cpp - instances/fmha_fwd_v3_d128_bf16_nmask.cpp - instances/fmha_fwd_v3_d128_fp16_mask.cpp - instances/fmha_fwd_v3_d128_fp16_nmask.cpp -) - -set_source_files_properties( - example_fmha_fwd_v3.cpp - fmha_fwd_v3.cpp - instances/fmha_fwd_v3_d128_bf16_mask.cpp - instances/fmha_fwd_v3_d128_bf16_nmask.cpp - instances/fmha_fwd_v3_d128_fp16_mask.cpp - instances/fmha_fwd_v3_d128_fp16_nmask.cpp - PROPERTIES LANGUAGE ${ROCM_EXAMPLES_GPU_LANGUAGE} -) - -target_compile_features(${example_name} PRIVATE cxx_std_17 hip_std_17) -target_compile_options(${example_name} PRIVATE - -fgpu-flush-denormals-to-zero - -Wno-undefined-func-template - --save-temps -) -target_include_directories(${example_name} PRIVATE ${CMAKE_CURRENT_LIST_DIR}) -target_link_libraries(${example_name} PRIVATE composable_kernel::ck_tile hip::host) - -include(CheckCXXCompilerFlag) -check_cxx_compiler_flag("-mllvm --amdgpu-disable-packed-fp32=1" HAS_DISABLE_PACKED_FP32) -if(HAS_DISABLE_PACKED_FP32) - target_compile_definitions(${example_name} PRIVATE CK_TILE_DISABLE_PACKED_FP32=1) - target_compile_options(${example_name} PRIVATE -mllvm --amdgpu-disable-packed-fp32=1) -endif() - -add_test(${example_name} COMMAND "${CMAKE_CURRENT_LIST_DIR}/${example_script}" "$") diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/README.md b/Libraries/ComposableKernel/attention/fmha_fwd_v3/README.md deleted file mode 100644 index 3022c4990d..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/README.md +++ /dev/null @@ -1,96 +0,0 @@ -# CK Tile Programming Model: Fused Multi-Head Attention Forward v3 example - -## Description - -This example demonstrates how to perform FMHA (fused multi-head attention) -forward pass with CK Tile using the v3 implementation. This version focuses -exclusively on the forward pass and includes support for variable-length -sequences with padding, allowing efficient processing of batches where sequences -have different effective lengths. - -### Supported architectures - -This example works with the `gfx950` architecture only. - -### Application flow - -1. Command line arguments are parsed to configure matrix dimensions and - execution parameters. -2. Buffers for the Q, K, V tensors are created on the host and initialized with - random values. -3. Buffers for all tensors are created on the device. -4. The tensors are copied from the host to the device. -5. Optional variable-length sequence support is configured: If effective - sequence lengths are provided via `q_eff_lens` or `kv_eff_lens`, cumulative - sequence length arrays are computed and uploaded to the device. These arrays - allow the kernel to skip padded regions and only process valid data. -6. CK Tile's `fmha_fwd_v3` kernel is instantiated and launched on the device. -7. If validation is enabled, the results are compared against a custom reference - implementation: For variable-length sequences, the reference computation - processes only the effective (non-padded) portion of each sequence. Padded - regions are zero-initialized and verified to remain zero. -8. All buffers are freed automatically. - -## Key APIs and concepts - -### CK Tile architecture - -The CK Tile framework is built around four key architectural components: - -The CK Tile framework is built around four key architectural components: - -* The **shape** defines the hierarchical tile structure and memory layout. -* The **problem** combines data types with the shape configuration. -* The **pipeline** schedules the sequence of operations for a kernel. -* The **kernel** implements the actual computation using the problem and pipeline definitions. - -For more information on CK Tile terminology, refer to the -[Composable Kernel Glossary](https://rocm.docs.amd.com/projects/composable_kernel/en/latest/reference/Composable-Kernel-Glossary.html). - -### Variable-length sequence support - -The v3 implementation supports efficient handling of variable-length sequences through cumulative sequence length arrays: - -* `cu_seqlen_q_ptr`: Cumulative sequence lengths for queries (size: batch + 1) -* `cu_seqlen_kv_ptr`: Cumulative sequence lengths for keys/values (size: batch + 1) - -When provided, these arrays allow the kernel to: - -* Process only the valid (non-padded) portion of each sequence -* Skip computation on padded regions -* Maintain correct attention masks per sequence - -This is particularly useful for: - -* Batches with sequences of varying lengths -* Reducing unnecessary computation on padding -* Maintaining accuracy when sequences have different effective lengths - -## Used API surface - -### CK Tile - -#### Types - -* `ck_tile::ArgParser` -* `ck_tile::bf16_t` -* `ck_tile::DeviceMem` -* `ck_tile::FillNormalDistribution` -* `ck_tile::fmha_fwd_v3_args` -* `ck_tile::fp16_t` -* `ck_tile::HostTensor` -* `ck_tile::identity` -* `ck_tile::index_t` -* `ck_tile::scales` -* `ck_tile::stream_config` - -#### Functions - -* `ck_tile::check_err` -* `ck_tile::fmha_fwd_v3` -* `ck_tile::make_generic_attention_mask_from_lr_window` -* `ck_tile::make_tuple` -* `ck_tile::reference_batched_gemm` -* `ck_tile::reference_batched_masking` -* `ck_tile::reference_batched_softmax` -* `ck_tile::sqrt` diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/bias.hpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/bias.hpp deleted file mode 100644 index 0b4c728340..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/bias.hpp +++ /dev/null @@ -1,109 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#pragma once - -#include "ck_tile/core.hpp" -#include "ck_tile/ops/fmha.hpp" -#include -#include - -// keep sync with BlockAttentionBiasEnum -enum class bias_enum { - no_bias = 0, - elementwise_bias = 1, - alibi = 2, -}; - -struct bias_info { - bias_enum type; - /* - * simple dispatch logic - * - * if type == elementwise_bias: - * if rank_info == 0: - * bias is 1*1*s*s - * elif rank_info == 1: - * bias is 1*h*s*s - * elif rank_info == 2: - * bias is b*h*s*s - * - * elif type == alibi: - * if rank_info == 0: - * alibi in 1*h - * elif rank_info == 1: - * alibi in b*h - */ - int rank_info; - - void serialize(std::ostream &os) const { - if (type == bias_enum::no_bias) - os << "n"; - else if (type == bias_enum::elementwise_bias) { - os << "e"; - if (rank_info != 0) { - os << "[" << rank_info << "]"; - } - } else if (type == bias_enum::alibi) { - os << "alibi"; - if (rank_info != 0) { - os << "[" << rank_info << "]"; - } - } - } - - static bias_info decode(std::string str) { - bias_info info{bias_enum::no_bias, 0}; - auto found_0 = str.find(':'); - if (found_0 != std::string::npos) { - std::string t = str.substr(0, found_0); - std::string v = str.substr(found_0 + 1); - if (t == "e" || t == "elementwise") { - info.type = bias_enum::elementwise_bias; - info.rank_info = std::stoi(v); - if (info.rank_info < 0 || info.rank_info > 2) - throw std::invalid_argument("invalid bias rank: " + str); - } else if (t == "a" || t == "alibi") { - info.type = bias_enum::alibi; - info.rank_info = std::stoi(v); - if (info.rank_info < 0 || info.rank_info > 1) - throw std::invalid_argument("invalid bias rank: " + str); - } else { - throw std::invalid_argument("invalid bias value: " + str); - } - } else if (str == "0" || str == "n") { - info.type = bias_enum::no_bias; - } else if (str == "1" || str == "e" || str == "elementwise") { - info.type = bias_enum::elementwise_bias; - } else if (str == "2" || str == "a" || str == "alibi") { - info.type = bias_enum::alibi; - } else { - throw std::invalid_argument("invalid bias value: " + str); - } - return info; - } - - friend std::ostream &operator<<(std::ostream &os, const bias_info &bi) { - bi.serialize(os); - return os; - } -}; diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/example_fmha_fwd_v3.cpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/example_fmha_fwd_v3.cpp deleted file mode 100644 index bd55e2c8c8..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/example_fmha_fwd_v3.cpp +++ /dev/null @@ -1,581 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "fmha_fwd.hpp" -#include "fmha_fwd_v3.hpp" -#include "mask.hpp" - -auto parse_cmd_args(int argc, - char *argv[]) -> std::pair { - ck_tile::ArgParser arg_parser; - arg_parser.insert("prec", "fp16", "data type. fp16/bf16") - .insert("b", "2", "batch size") - .insert("h", "8", "num of head, for q") - .insert("h_k", "-1", - "num of head, for k/v, -1 means equal to h\n" - "if not equal to h, then this is GQA/MQA case") - .insert("s", "3328", "seqlen_q") - .insert("s_k", "-1", "seqlen_k, -1 means equal to s") - .insert("d", "128", "head dim for q & k") - .insert("scale_s", "0", - "scale factor of S. 0 means equal to 1/sqrt(hdim)") - .insert("iperm", "0", - "permute input\n" - "if true, will be b*h*s*d, else b*s*h*d") - .insert("operm", "0", "permute output") - .insert("causal", "0", "0: no mask, 1: causal mask") - .insert("v", "1", "0:no verify, 1:verify") - .insert("seed", "11939", - "random seed used for initializing input tensors. 0 for " - "non-deterministic seed") - .insert("warmup", "5", "number of iterations before benchmark the kernel") - .insert("repeat", "30", "number of iterations to benchmark the kernel") - // Optional effective seqlen override (exclude PAD) for batch mode - .insert( - "q_eff_lens", "", - "Batch-mode only: per-batch effective seqlen for Q (exclude PAD).\n" - "Comma-separated list of length 'b'. If empty, no override.") - .insert( - "kv_eff_lens", "", - "Batch-mode only: per-batch effective seqlen for KV (exclude PAD).\n" - "Comma-separated list of length 'b'. If empty, no override."); - - bool result = arg_parser.parse(argc, argv); - return std::make_pair(result, arg_parser); -} - -enum class TensorLayout { - bhsd, - bshd, -}; - -std::ostream &operator<<(std::ostream &stream, TensorLayout layout) { - switch (layout) { - case TensorLayout::bhsd: - return stream << "bhsd"; - case TensorLayout::bshd: - return stream << "bshd"; - default: - return stream << "unknown"; - } -} - -struct Problem { - explicit Problem(const ck_tile::ArgParser &args) { - data_type = args.get_str("prec") == "fp16" - ? ck_tile::fmha_fwd_v3_args::data_type_enum::fp16 - : ck_tile::fmha_fwd_v3_args::data_type_enum::bf16; - batch = args.get_int("b"); - seqlen_q = args.get_int("s"); - seqlen_k = args.get_int("s_k"); - if (seqlen_k < 0) { - seqlen_k = seqlen_q; - } - nhead_q = args.get_int("h"); - nhead_kv = args.get_int("h_k"); - if (nhead_kv < 0) { - nhead_kv = nhead_q; - } - hdim = args.get_int("d"); - softmax_scale = args.get_float("scale_s"); - if (softmax_scale == .0f) - softmax_scale = 1.0 / ck_tile::sqrt(static_cast(hdim)); - - const auto is_causal = args.get_bool("causal"); - if (is_causal) { - mask = mask_info::decode("b:-1,0", seqlen_q, seqlen_k); - } else { - mask = mask_info::decode("0", seqlen_q, seqlen_k); - } - - input_layout = - args.get_int("iperm") == 1 ? TensorLayout::bhsd : TensorLayout::bshd; - output_layout = - args.get_int("operm") == 1 ? TensorLayout::bhsd : TensorLayout::bshd; - q_eff_lens = args.get_int_vec("q_eff_lens"); - kv_eff_lens = args.get_int_vec("kv_eff_lens"); - } - - std::vector get_query_shape() const { - if (input_layout == TensorLayout::bhsd) { - return {batch, nhead_q, seqlen_q, hdim}; - } else { - return {batch, seqlen_q, nhead_q, hdim}; - } - } - - std::vector get_key_shape() const { - if (input_layout == TensorLayout::bhsd) { - return {batch, nhead_kv, seqlen_k, hdim}; - } else { - return {batch, seqlen_k, nhead_kv, hdim}; - } - } - - std::vector get_value_shape() const { - if (input_layout == TensorLayout::bhsd) { - return {batch, nhead_kv, seqlen_k, hdim}; - } else { - return {batch, seqlen_k, nhead_kv, hdim}; - } - } - - std::vector get_output_shape() const { - if (output_layout == TensorLayout::bhsd) { - return {batch, nhead_q, seqlen_q, hdim}; - } else { - return {batch, seqlen_q, nhead_q, hdim}; - } - } - - ck_tile::fmha_fwd_v3_args::data_type_enum data_type; - ck_tile::index_t batch; - ck_tile::index_t seqlen_q; - ck_tile::index_t seqlen_k; - ck_tile::index_t nhead_q; - ck_tile::index_t nhead_kv; - ck_tile::index_t hdim; - float softmax_scale; - mask_info mask; - TensorLayout input_layout; - TensorLayout output_layout; - std::vector q_eff_lens; - std::vector kv_eff_lens; -}; - -struct RunConfig { - explicit RunConfig(const ck_tile::ArgParser &args) { - seed = args.get_uint32("seed"); - if (*seed == 0) { - seed.reset(); - } - - kernel_warmup = args.get_int("warmup"); - kernel_repeat = args.get_int("repeat"); - verify = args.get_bool("v"); - } - - std::optional seed; - int kernel_warmup; - int kernel_repeat; - bool verify; -}; - -template -auto generate_qkv(const Problem &problem, - [[maybe_unused]] std::optional seed = std::nullopt) - -> std::tuple, ck_tile::HostTensor, - ck_tile::HostTensor> { - ck_tile::HostTensor q(problem.get_query_shape()); - ck_tile::HostTensor k(problem.get_key_shape()); - ck_tile::HostTensor v(problem.get_value_shape()); - - ck_tile::FillNormalDistribution{0.f, 3.f, seed}(q); - ck_tile::FillNormalDistribution{0.f, 3.f, seed}(k); - ck_tile::FillNormalDistribution{0.f, 3.f, seed}(v); - - return std::make_tuple(q, k, v); -} - -namespace host { -template -CK_TILE_HOST void fmha_fwd(const ck_tile::HostTensor &q_bshd, - const ck_tile::HostTensor &k_bshd, - const ck_tile::HostTensor &v_bshd, - const mask_info &mask, - ck_tile::HostTensor &o_bshd, - const QElementOp &q_element_op = {}, - const KElementOp &k_element_op = {}, - const VElementOp &v_element_op = {}, - const SAccElementOp &s_acc_element_op = {}) { - const int batch_size = q_bshd.mDesc.get_lengths()[0]; - const int seqlen_q = q_bshd.mDesc.get_lengths()[1]; - const int seqlen_kv = k_bshd.mDesc.get_lengths()[1]; - const int nhead_q = q_bshd.mDesc.get_lengths()[2]; - const int nhead_kv = k_bshd.mDesc.get_lengths()[2]; - const int hdim_qk = q_bshd.mDesc.get_lengths()[3]; - const int hdim_v = v_bshd.mDesc.get_lengths()[3]; - - const int nr = nhead_q / nhead_kv; - - ck_tile::HostTensor q_host_ref({nhead_q, seqlen_q, hdim_qk}); - ck_tile::HostTensor k_host_ref({nhead_q, seqlen_kv, hdim_qk}); - ck_tile::HostTensor v_host_ref({nhead_q, hdim_v, seqlen_kv}); - ck_tile::HostTensor o_host_ref({nhead_q, seqlen_q, hdim_v}); - - ck_tile::HostTensor s_host_ref({nhead_q, seqlen_q, seqlen_kv}); - ck_tile::HostTensor p_host_ref({nhead_q, seqlen_q, seqlen_kv}); - - // do computation for each batch - for (int b = 0; b < batch_size; ++b) { - // copy per-batch data from input tensors - // clang-format off - q_host_ref.ForEach([&](auto& self, auto idx) { self(idx) = q_bshd(b, idx[1], idx[0] , idx[2]); }); - k_host_ref.ForEach([&](auto& self, auto idx) { self(idx) = k_bshd(b, idx[1], idx[0] / nr, idx[2]); }); - v_host_ref.ForEach([&](auto& self, auto idx) { self(idx) = v_bshd(b, idx[2], idx[0] / nr, idx[1]); }); - // clang-format on - ck_tile::reference_batched_gemm( - q_host_ref, k_host_ref, s_host_ref, q_element_op, k_element_op, - s_acc_element_op); - - if (mask.type == mask_enum::no_mask) { - ck_tile::reference_batched_masking( - s_host_ref, FmhaMasks::NoMask{seqlen_q, seqlen_kv}); - } else if (mask.type == mask_enum::window_generic) { - ck_tile::reference_batched_masking( - s_host_ref, ck_tile::make_generic_attention_mask_from_lr_window< - FmhaMasks::GenericMask>(mask.left, mask.right, - seqlen_q, seqlen_kv)); - } else { - // if left window size is negative, means causal - // else means generic (for current batch) - if (mask.left < 0) - ck_tile::reference_batched_masking( - s_host_ref, ck_tile::make_generic_attention_mask_from_lr_window< - FmhaMasks::CausalMask>( - mask.left, mask.right, seqlen_q, seqlen_kv, - mask.type == mask_enum::mask_top_left)); - else - ck_tile::reference_batched_masking( - s_host_ref, ck_tile::make_generic_attention_mask_from_lr_window< - FmhaMasks::GenericMask>( - mask.left, mask.right, seqlen_q, seqlen_kv, - mask.type == mask_enum::mask_top_left)); - } - - ck_tile::reference_batched_softmax( - s_host_ref, p_host_ref, ck_tile::identity{}); - - ck_tile::reference_batched_gemm( - p_host_ref, v_host_ref, o_host_ref, ck_tile::identity{}, v_element_op); - - // copy resulting per-batch data to the output tensor - o_host_ref.ForEach([&](auto &self, auto idx) { - o_bshd(b, idx[1], idx[0], idx[2]) = self(idx); - }); - } -} -} // namespace host - -template -bool run_impl(const Problem &problem, const RunConfig &run_config) { - auto [q, k, v] = generate_qkv(problem, run_config.seed); - - ck_tile::DeviceMem q_buf(q.get_element_space_size_in_bytes()); - ck_tile::DeviceMem k_buf(k.get_element_space_size_in_bytes()); - ck_tile::DeviceMem v_buf(v.get_element_space_size_in_bytes()); - /// FIXME: use correct size for output tensor. just use q size for now since - /// hidm_qk = hdim_v - ck_tile::DeviceMem o_buf(q.get_element_space_size_in_bytes()); - - q_buf.ToDevice(q.data()); - k_buf.ToDevice(k.data()); - v_buf.ToDevice(v.data()); - // Ensure output buffer is zero-initialized so padded regions compare cleanly - o_buf.SetZero(); - - ck_tile::fmha_fwd_v3_args args{}; - - args.data_type = problem.data_type; - args.batch = problem.batch; - args.seqlen_q = problem.seqlen_q; - args.seqlen_k = problem.seqlen_k; - args.nhead_q = problem.nhead_q; - args.nhead_kv = problem.nhead_kv; - args.hdim_qk = problem.hdim; - args.hdim_v = problem.hdim; - args.softmax_scale = problem.softmax_scale; - - args.window_size_left = problem.mask.left; - args.window_size_right = problem.mask.right; - args.mask_type = static_cast(problem.mask.type); - - // bshd: (batch, seqlen_q, nhead_q, hdim) - // bhsd: (batch, nhead_q, seqlen_q, hdim) - args.q_ptr = q_buf.GetDeviceBuffer(); - args.stride_q = problem.input_layout == TensorLayout::bshd - ? problem.nhead_q * problem.hdim - : problem.hdim; - args.nhead_stride_q = problem.input_layout == TensorLayout::bshd - ? problem.hdim - : problem.seqlen_q * problem.hdim; - args.batch_stride_q = problem.seqlen_q * problem.nhead_q * problem.hdim; - - // bshd: (batch, seqlen_k, nhead_kv, hdim) - // bhsd: (batch, nhead_kv, seqlen_k, hdim) - args.k_ptr = k_buf.GetDeviceBuffer(); - args.stride_k = problem.input_layout == TensorLayout::bshd - ? problem.nhead_kv * problem.hdim - : problem.hdim; - args.nhead_stride_k = problem.input_layout == TensorLayout::bshd - ? problem.hdim - : problem.seqlen_k * problem.hdim; - args.batch_stride_k = problem.seqlen_k * problem.nhead_kv * problem.hdim; - - // bshd: (batch, seqlen_k, nhead_kv, hdim) - // bhsd: (batch, nhead_kv, seqlen_k, hdim) - args.v_ptr = v_buf.GetDeviceBuffer(); - args.stride_v = problem.input_layout == TensorLayout::bshd - ? problem.nhead_kv * problem.hdim - : problem.hdim; - args.nhead_stride_v = problem.input_layout == TensorLayout::bshd - ? problem.hdim - : problem.seqlen_k * problem.hdim; - args.batch_stride_v = problem.seqlen_k * problem.nhead_kv * problem.hdim; - - // bshd: (batch, seqlen_q, nhead_q, hdim) - // bhsd: (batch, nhead_q, seqlen_q, hdim) - args.o_ptr = o_buf.GetDeviceBuffer(); - args.stride_o = problem.output_layout == TensorLayout::bshd - ? problem.nhead_q * problem.hdim - : problem.hdim; - args.nhead_stride_o = problem.output_layout == TensorLayout::bshd - ? problem.hdim - : problem.seqlen_q * problem.hdim; - args.batch_stride_o = problem.seqlen_q * problem.nhead_q * problem.hdim; - - // Optional cumulative seqlen overrides (exclude PAD) - const bool has_varlen_q = - !problem.q_eff_lens.empty() && problem.q_eff_lens[0] != -1; - const bool has_varlen_k = - !problem.kv_eff_lens.empty() && problem.kv_eff_lens[0] != -1; - - auto make_effective_vec = [&](const std::vector &opt_vec, - ck_tile::index_t fallback) { - std::vector eff; - if (!opt_vec.empty() && opt_vec[0] != -1) { - eff.assign(opt_vec.begin(), opt_vec.end()); - if (eff.size() < static_cast(problem.batch)) { - eff.resize(problem.batch, eff.back()); - } - } else { - eff.assign(problem.batch, fallback); - } - return eff; - }; - - const auto eff_q_vec = - make_effective_vec(problem.q_eff_lens, problem.seqlen_q); - const auto eff_kv_vec = - make_effective_vec(problem.kv_eff_lens, problem.seqlen_k); - - // Calculate cumulative sums for kernel arguments if varlen is used - std::vector cuq_cum, cukv_cum; - auto calculate_cumulative = - [&](const std::vector &per_batch_vec, - std::vector &cum_vec) { - cum_vec.resize(per_batch_vec.size() + 1); - cum_vec[0] = 0; - for (std::size_t i = 0; i < per_batch_vec.size(); ++i) - cum_vec[i + 1] = cum_vec[i] + per_batch_vec[i]; - }; - - if (has_varlen_q) { - calculate_cumulative(eff_q_vec, cuq_cum); - } - if (has_varlen_k) { - calculate_cumulative(eff_kv_vec, cukv_cum); - } - - ck_tile::DeviceMem cuq_buf( - !cuq_cum.empty() ? cuq_cum.size() * sizeof(ck_tile::index_t) : 0); - ck_tile::DeviceMem cukv_buf( - !cukv_cum.empty() ? cukv_cum.size() * sizeof(ck_tile::index_t) : 0); - cuq_buf.ToDevice(!cuq_cum.empty() ? cuq_cum.data() : nullptr); - cukv_buf.ToDevice(!cukv_cum.empty() ? cukv_cum.data() : nullptr); - args.cu_seqlen_q_ptr = !cuq_cum.empty() - ? reinterpret_cast( - cuq_buf.GetDeviceBuffer()) - : nullptr; - args.cu_seqlen_kv_ptr = !cukv_cum.empty() - ? reinterpret_cast( - cukv_buf.GetDeviceBuffer()) - : nullptr; - - ck_tile::stream_config stream_config{ - nullptr, true, - /*log_level=*/0, run_config.kernel_warmup, run_config.kernel_repeat}; - - auto [result, time] = ck_tile::fmha_fwd_v3(args, stream_config); - if (!result) { - std::cerr << "faild to run fmha_fwd_v3()" << std::endl; - return false; - } - - std::size_t flop = [&] { - if (problem.mask.type == mask_enum::no_mask) { - return 4 * problem.batch * problem.nhead_q * problem.seqlen_q * - problem.seqlen_k * problem.hdim; - } else { - /// FIXME: Use a more accurate method; for now, we’re just dividing the - /// flop by 2. - return 2 * problem.batch * problem.nhead_q * problem.seqlen_q * - problem.seqlen_k * problem.hdim; - } - }(); - float tflops = static_cast(flop) / 1.e9 / time; - - std::cout << "[" << problem.data_type << "|"; - if (problem.input_layout == problem.output_layout) { - std::cout << problem.input_layout; - } else { - std::cout << problem.input_layout << "-" << problem.output_layout; - } - std::cout << "] b:" << problem.batch << ", h:" << problem.nhead_q << "/" - << problem.nhead_kv << ", s:" << problem.seqlen_q << "/" - << problem.seqlen_k << ", d:" << problem.hdim - << ", scale_s:" << problem.softmax_scale - << ", mask:" << problem.mask << std::fixed << ", " - << std::setprecision(3) << time << " ms, " << std::setprecision(2) - << tflops << " TFlops" << std::endl; - - if (!run_config.verify) { - return true; - } - - // transpose tensor descriptors from bhsd to bshd if necessary - if (problem.input_layout != TensorLayout::bshd) { - q = q.transpose({0, 2, 1, 3}); - k = k.transpose({0, 2, 1, 3}); - v = v.transpose({0, 2, 1, 3}); - } - - ck_tile::HostTensor o_ref(problem.get_output_shape()); - if (problem.output_layout != TensorLayout::bshd) { - o_ref = o_ref.transpose({0, 2, 1, 3}); - } - - // If variable lengths are provided, compute per-batch references - // with the effective lengths; else compute a single full reference. - if (has_varlen_q || has_varlen_k) { - // Variable-length aware verification: zero-fill padded region and only - // compute valid part. - o_ref.SetZero(); - - for (int b = 0; b < problem.batch; ++b) { - const ck_tile::index_t seqlen_q_eff = eff_q_vec[b]; - const ck_tile::index_t seqlen_kv_eff = eff_kv_vec[b]; - - if (seqlen_q_eff <= 0 || seqlen_kv_eff <= 0) - continue; - - // Slice current batch from inputs (bshd) and build single-batch tensors - ck_tile::HostTensor q_b( - {1, seqlen_q_eff, problem.nhead_q, problem.hdim}); - ck_tile::HostTensor k_b( - {1, seqlen_kv_eff, problem.nhead_kv, problem.hdim}); - ck_tile::HostTensor v_b( - {1, seqlen_kv_eff, problem.nhead_kv, problem.hdim}); - ck_tile::HostTensor o_b( - {1, seqlen_q_eff, problem.nhead_q, problem.hdim}); - - // Copy effective region - q_b.ForEach([&](auto &self, auto idx) { - // idx: [0, s, h, d] - self(idx) = q(b, idx[1], idx[2], idx[3]); - }); - k_b.ForEach([&](auto &self, auto idx) { - self(idx) = k(b, idx[1], idx[2], idx[3]); - }); - v_b.ForEach([&](auto &self, auto idx) { - self(idx) = v(b, idx[1], idx[2], idx[3]); - }); - - // Compute reference for this batch segment (host::fmha_fwd expects bshd - // tensors) - host::fmha_fwd(q_b, k_b, v_b, problem.mask, o_b, - ck_tile::identity{}, ck_tile::identity{}, - ck_tile::identity{}, - ck_tile::scales{problem.softmax_scale}); - - // Scatter into o_ref's bshd descriptor memory - for (int s = 0; s < seqlen_q_eff; ++s) { - for (int h = 0; h < problem.nhead_q; ++h) { - for (int d = 0; d < problem.hdim; ++d) { - o_ref(b, s, h, d) = o_b(0, s, h, d); - } - } - } - } - } else { - // No varlen override: compute the full reference once - host::fmha_fwd( - q, k, v, problem.mask, o_ref, ck_tile::identity{}, ck_tile::identity{}, - ck_tile::identity{}, ck_tile::scales{problem.softmax_scale}); - } - - ck_tile::HostTensor o(problem.get_output_shape()); - o_buf.FromDevice(o.data()); - - const auto [rtol, atol] = [&] { - if constexpr (std::is_same_v) - return std::make_tuple(1e-3, 1e-3); - else - return std::make_tuple(1e-2, 1e-2); - }(); - return ck_tile::check_err(o, o_ref, std::string("found incorrect results!"), - rtol, atol); -} - -int main(int argc, char *argv[]) { - auto [parse_result, args] = parse_cmd_args(argc, argv); - if (!parse_result) { - std::cerr << "failed to parse command line arguments" << std::endl; - } - - Problem problem(args); - RunConfig run_config(args); - - const auto run = [&] { - if (problem.data_type == ck_tile::fmha_fwd_v3_args::data_type_enum::fp16) { - return run_impl(problem, run_config); - } else { - return run_impl(problem, run_config); - } - }; - - return !run(); -} diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd.hpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd.hpp deleted file mode 100644 index daae03c5d8..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd.hpp +++ /dev/null @@ -1,1080 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#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/epilogue.hpp" -#include "ck_tile/ops/fmha.hpp" - -#include "bias.hpp" -#include "mask.hpp" -#include "rotary.hpp" - -#include -#include -#include - -struct FmhaFwdFp32 {}; - -struct FmhaFwdFp16 {}; - -struct FmhaFwdBf16 {}; - -struct FmhaFwdFp8 {}; - -struct FmhaFwdBf8 {}; - -struct FmhaFwdFp8Fp16 {}; - -struct FmhaFwdFp8Bf16 {}; - -struct FmhaFwdFp8Fp32 {}; - -template struct FmhaFwdTypeConfig; - -template <> struct FmhaFwdTypeConfig { - using QDataType = float; - using KDataType = float; - using VDataType = float; - using BiasDataType = float; - using RandValOutputDataType = uint8_t; - using LSEDataType = - float; // data type for lse(logsumexp L_j = max_j + log(l_j)) - using SaccDataType = float; // data type for first gemm accumulation - using SMPLComputeDataType = float; // data type for reduction, softmax - using PDataType = float; // data type for A matrix of second gemm - using OaccDataType = float; // data type for second gemm accumulation - using ODataType = float; -}; - -template <> struct FmhaFwdTypeConfig { - using QDataType = ck_tile::half_t; - using KDataType = ck_tile::half_t; - using VDataType = ck_tile::half_t; - using BiasDataType = ck_tile::half_t; - using RandValOutputDataType = uint8_t; - using LSEDataType = - float; // data type for lse(logsumexp L_j = max_j + log(l_j)) - using SaccDataType = float; // data type for first gemm accumulation - using SMPLComputeDataType = float; // data type for reduction, softmax - using PDataType = ck_tile::half_t; // data type for A matrix of second gemm - using OaccDataType = float; // data type for second gemm accumulation - using ODataType = ck_tile::half_t; -}; - -template <> struct FmhaFwdTypeConfig { - using QDataType = ck_tile::bf16_t; - using KDataType = ck_tile::bf16_t; - using VDataType = ck_tile::bf16_t; - using BiasDataType = ck_tile::bf16_t; - using RandValOutputDataType = uint8_t; - using LSEDataType = - float; // data type for lse(logsumexp L_j = max_j + log(l_j)) - using SaccDataType = float; // data type for first gemm accumulation - using SMPLComputeDataType = float; // data type for reduction, softmax - using PDataType = ck_tile::bf16_t; // data type for A matrix of second gemm - using OaccDataType = float; // data type for second gemm accumulation - using ODataType = ck_tile::bf16_t; -}; - -template <> struct FmhaFwdTypeConfig { - using QDataType = ck_tile::fp8_t; - using KDataType = ck_tile::fp8_t; - using VDataType = ck_tile::fp8_t; - using BiasDataType = float; - using RandValOutputDataType = uint8_t; - using LSEDataType = - float; // data type for lse(logsumexp L_j = max_j + log(l_j)) - using SaccDataType = float; // data type for first gemm accumulation - using SMPLComputeDataType = float; // data type for reduction, softmax - using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm - using OaccDataType = float; // data type for second gemm accumulation - using ODataType = ck_tile::fp8_t; -}; - -template <> struct FmhaFwdTypeConfig { - using QDataType = ck_tile::bf8_t; - using KDataType = ck_tile::bf8_t; - using VDataType = ck_tile::bf8_t; - using BiasDataType = ck_tile::bf8_t; - using RandValOutputDataType = uint8_t; - using LSEDataType = - float; // data type for lse(logsumexp L_j = max_j + log(l_j)) - using SaccDataType = float; // data type for first gemm accumulation - using SMPLComputeDataType = float; // data type for reduction, softmax - using PDataType = ck_tile::bf8_t; // data type for A matrix of second gemm - using OaccDataType = float; // data type for second gemm accumulation - using ODataType = ck_tile::bf8_t; -}; - -template <> struct FmhaFwdTypeConfig { - using QDataType = ck_tile::fp8_t; - using KDataType = ck_tile::fp8_t; - using VDataType = ck_tile::fp8_t; - using BiasDataType = float; - using RandValOutputDataType = uint8_t; - using LSEDataType = - float; // data type for lse(logsumexp L_j = max_j + log(l_j)) - using SaccDataType = float; // data type for first gemm accumulation - using SMPLComputeDataType = float; // data type for reduction, softmax - using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm - using OaccDataType = float; // data type for second gemm accumulation - using ODataType = ck_tile::bf16_t; -}; - -template <> struct FmhaFwdTypeConfig { - using QDataType = ck_tile::fp8_t; - using KDataType = ck_tile::fp8_t; - using VDataType = ck_tile::fp8_t; - using BiasDataType = float; - using RandValOutputDataType = uint8_t; - using LSEDataType = - float; // data type for lse(logsumexp L_j = max_j + log(l_j)) - using SaccDataType = float; // data type for first gemm accumulation - using SMPLComputeDataType = float; // data type for reduction, softmax - using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm - using OaccDataType = float; // data type for second gemm accumulation - using ODataType = float; -}; - -struct FmhaMasks { - using NoMask = ck_tile::GenericAttentionMask; - using GenericMask = ck_tile::GenericAttentionMask; - using CausalMask = ck_tile::GenericAttentionMask; -}; - -// runtime args, some will passed to karg, some will used to compute -// grids/blocks -struct fmha_fwd_args { - const void *q_ptr; - const void *k_ptr; - const void *v_ptr; - const void *bias_ptr; // bias or alibi_slope pointer - void *rand_val_ptr; - void *lse_ptr; - void *o_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 hdim_q; - ck_tile::index_t hdim_v; - ck_tile::index_t nhead_q; - ck_tile::index_t nhead_k; - - float scale_s; - float scale_p; - float scale_o; - - float logits_soft_cap; - - 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_randval; - ck_tile::index_t stride_o; - 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_randval; - ck_tile::index_t nhead_stride_lse; - ck_tile::index_t nhead_stride_o; - 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_randval; - ck_tile::index_t batch_stride_lse; - ck_tile::index_t batch_stride_o; - - ck_tile::index_t window_size_left; - ck_tile::index_t window_size_right; - ck_tile::index_t mask_type; - ck_tile::index_t min_seqlen_q; - - float p_drop; - bool s_randval; - - std::variant, - std::pair> - drop_seed_offset; -}; - -struct fmha_fwd_pagedkv_args { - const void *q_ptr; - const void *k_ptr; - const void *v_ptr; - const void *bias_ptr; // bias or alibi_slope pointer - void *lse_ptr; - void *o_ptr; - - void *block_table_ptr; - ck_tile::index_t - batch_stride_block_table; // only used if 'block_table_ptr' is not nullptr - ck_tile::index_t - page_block_size; // only used if 'block_table_ptr' is not nullptr - bool is_gappy; // differentiate seqstart_k_ptr usage. only used if - // 'block_table_ptr' is not nullptr. - - const void *cache_batch_idx; - - // the real seqlen_q & seqlen_k are decided by following: - // batch mode: seqlen_q = kargs.seqlen_q - // seqlen_k = kargs.seqlen_k - // group mode: seqlen_q = kargs.seqstart_q_ptr[b + 1] - - // kargs.seqstart_q_ptr[b] - // seqlen_k = kargs.seqstart_k_ptr[b + 1] - - // kargs.seqstart_k_ptr[b] - // or kargs.seqlen_k_ptr[b] - // - // batch mode (kvcache): - // seqlen_q = kargs.seqlen_q - // seqlen_k = kargs.seqlen_k_ptr[b] - // group mode (kvcache): - // seqlen_q = kargs.seqstart_q_ptr[b + 1] - - // kargs.seqstart_q_ptr[b] - // - // when is_gappy=true: - // seqlen_k = kargs.seqlen_k_ptr[b] - // seqstart_k_ptr[b] now store local offset of each batch - // - // when is_gappy=false: - // seqlen_k = kargs.seqstart_k_ptr[b + 1] - - // kargs.seqstart_k_ptr[b] - // or kargs.seqlen_k_ptr[b] - const void *seqstart_q_ptr; - const void *seqstart_k_ptr; - const void *seqlen_k_ptr; - - 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 hdim_q; - ck_tile::index_t hdim_v; - ck_tile::index_t nhead_q; - ck_tile::index_t nhead_k; - - float scale_s; - float scale_p; - float scale_o; - - float logits_soft_cap; - - 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 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_lse; - ck_tile::index_t nhead_stride_o; - 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_lse; - ck_tile::index_t batch_stride_o; - - ck_tile::index_t window_size_left; - ck_tile::index_t window_size_right; - ck_tile::index_t mask_type; - ck_tile::index_t min_seqlen_q; -}; - -struct fmha_fwd_splitkv_args { - const void *q_ptr; - const void *k_ptr; - const void *v_ptr; - const void *bias_ptr; // bias or alibi_slope pointer - void *lse_acc_ptr; - void *o_acc_ptr; - void *lse_ptr; - void *o_ptr; - - void *block_table_ptr; - ck_tile::index_t - batch_stride_block_table; // only used if 'block_table_ptr' is not nullptr - ck_tile::index_t - page_block_size; // only used if 'block_table_ptr' is not nullptr - bool is_gappy; // differentiate seqstart_k_ptr usage. only used if - // 'block_table_ptr' is not nullptr. - - const void *cache_batch_idx; - - // the real seqlen_q & seqlen_k are decided by following: - // batch mode: seqlen_q = kargs.seqlen_q - // seqlen_k = kargs.seqlen_k - // group mode: seqlen_q = kargs.seqstart_q_ptr[b + 1] - - // kargs.seqstart_q_ptr[b] - // seqlen_k = kargs.seqstart_k_ptr[b + 1] - - // kargs.seqstart_k_ptr[b] - // or kargs.seqlen_k_ptr[b] - // - // batch mode (kvcache): - // seqlen_q = kargs.seqlen_q - // seqlen_k = kargs.seqlen_k_ptr[b] - // group mode (kvcache): - // seqlen_q = kargs.seqstart_q_ptr[b + 1] - - // kargs.seqstart_q_ptr[b] - // - // when is_gappy=true: - // seqlen_k = kargs.seqlen_k_ptr[b] - // seqstart_k_ptr[b] now store local offset of each batch - // - // when is_gappy=false: - // seqlen_k = kargs.seqstart_k_ptr[b + 1] - - // kargs.seqstart_k_ptr[b] - // or kargs.seqlen_k_ptr[b] - const void *seqstart_q_ptr; - const void *seqstart_k_ptr; - const void *seqlen_k_ptr; - - 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 hdim_q; - ck_tile::index_t hdim_v; - ck_tile::index_t nhead_q; - ck_tile::index_t nhead_k; - ck_tile::index_t num_splits; - - float scale_s; - float scale_p; - float scale_o; - - float logits_soft_cap; - - 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_acc; - ck_tile::index_t stride_o; - 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_lse; - ck_tile::index_t nhead_stride_lse_acc; - ck_tile::index_t nhead_stride_o_acc; - ck_tile::index_t nhead_stride_o; - 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_lse; - ck_tile::index_t batch_stride_lse_acc; - ck_tile::index_t batch_stride_o_acc; - ck_tile::index_t batch_stride_o; - ck_tile::index_t split_stride_lse_acc; - ck_tile::index_t split_stride_o_acc; - - ck_tile::index_t window_size_left; - ck_tile::index_t window_size_right; - ck_tile::index_t mask_type; -}; - -struct fmha_fwd_appendkv_args { - void *q_ptr; - void *k_ptr; - const void *knew_ptr; - void *v_ptr; - const void *vnew_ptr; - - const void *seqlen_k_ptr; - - ck_tile::index_t seqlen_q; - ck_tile::index_t seqlen_knew; - ck_tile::index_t batch; - ck_tile::index_t hdim_q; - ck_tile::index_t hdim_v; - ck_tile::index_t nhead_q; - ck_tile::index_t nhead_k; - - const void *rotary_cos_ptr; // only used if 'rotary_dim' > 0 - const void *rotary_sin_ptr; // only used if 'rotary_dim' > 0 - ck_tile::index_t rotary_dim; - bool has_mask; - - void *block_table_ptr; - ck_tile::index_t - batch_stride_block_table; // only used if 'block_table_ptr' is not nullptr - ck_tile::index_t - page_block_size; // only used if 'block_table_ptr' is not nullptr - - const void *cache_batch_idx; // only used if block_table_ptr is nullptr -> - // batch mode (kvcache) - - ck_tile::index_t stride_q; - ck_tile::index_t stride_k; - ck_tile::index_t stride_knew; - ck_tile::index_t stride_v; - ck_tile::index_t stride_vnew; - ck_tile::index_t nhead_stride_q; - ck_tile::index_t nhead_stride_k; - ck_tile::index_t nhead_stride_knew; - ck_tile::index_t nhead_stride_v; - ck_tile::index_t nhead_stride_vnew; - ck_tile::index_t batch_stride_q; - ck_tile::index_t batch_stride_k; - ck_tile::index_t batch_stride_knew; - ck_tile::index_t batch_stride_v; - ck_tile::index_t batch_stride_vnew; -}; - -struct fmha_batch_prefill_args { - const void *q_ptr; - const void *k_ptr; - const void *v_ptr; - const void *bias_ptr; // bias or alibi_slope pointer - void *rand_val_ptr; - void *lse_ptr; - void *o_ptr; - - // the real seqlen_q & seqlen_k are decided by following: - // batch mode (kvcache): - // seqlen_q = kargs.seqlen_q - // seqlen_k = kargs.page_block_size * (kargs.kv_indptr[b + 1] - - // kargs.kv_indptr[b] - 1) + - // kargs.kv_last_page_lens[b] - // group mode (kvcache): - // seqlen_q = kargs.seqstart_q_ptr[b + 1] - - // kargs.seqstart_q_ptr[b] seqlen_k = kargs.page_block_size * - // (kargs.kv_indptr[b + 1] - kargs.kv_indptr[b] - 1) + - // kargs.kv_last_page_lens[b] - const void *seqstart_q_ptr; - - 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 hdim_q; - ck_tile::index_t hdim_v; - ck_tile::index_t nhead_q; - ck_tile::index_t nhead_k; - - // SGLang-style page table - int32_t num_total_pages; - void *kv_indptr; - void *kv_page_indices; -#if 0 // we assume page_block_size=1 for now - void* kv_last_page_lens; - ck_tile::index_t page_block_size; -#endif - - float scale_s; - float scale_p; - float scale_o; - - float logits_soft_cap; - - 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_randval; - ck_tile::index_t stride_o; - 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_randval; - ck_tile::index_t nhead_stride_lse; - ck_tile::index_t nhead_stride_o; - 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_randval; - ck_tile::index_t batch_stride_lse; - ck_tile::index_t batch_stride_o; - - ck_tile::index_t window_size_left; - ck_tile::index_t window_size_right; - ck_tile::index_t mask_type; - - float p_drop; - bool s_randval; - - std::variant, - std::pair> - drop_seed_offset; -}; - -template -auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) { - assert(args.nhead_q % args.nhead_k == 0); - auto kargs = [&] { - // create group mode kernel arguments - if constexpr (FmhaKernel::kIsGroupMode) { - return FmhaKernel::MakeKargsImpl( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.rand_val_ptr, - args.lse_ptr, args.o_ptr, args.seqstart_q_ptr, args.seqstart_k_ptr, - args.seqlen_q_ptr, args.seqlen_k_ptr, args.hdim_q, args.hdim_v, - args.nhead_q, args.nhead_q / args.nhead_k, args.scale_s, args.scale_p, - args.scale_o, args.logits_soft_cap, args.stride_q, args.stride_k, - args.stride_v, args.stride_bias, args.stride_randval, args.stride_o, - args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, - args.nhead_stride_bias, args.nhead_stride_randval, - args.nhead_stride_lse, args.nhead_stride_o, args.window_size_left, - args.window_size_right, args.mask_type, args.min_seqlen_q, - args.p_drop, args.s_randval, args.drop_seed_offset, - args.cu_seqlen_q_ptr, args.cu_seqlen_k_ptr); - } else { // create batch mode kernel arguments - return FmhaKernel::MakeKargsImpl( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.rand_val_ptr, - args.lse_ptr, args.o_ptr, args.seqlen_q, args.seqlen_k, args.hdim_q, - args.hdim_v, args.nhead_q, args.nhead_q / args.nhead_k, args.scale_s, - args.scale_p, args.scale_o, args.logits_soft_cap, args.stride_q, - args.stride_k, args.stride_v, args.stride_bias, args.stride_randval, - args.stride_o, args.nhead_stride_q, args.nhead_stride_k, - args.nhead_stride_v, args.nhead_stride_bias, - args.nhead_stride_randval, args.nhead_stride_lse, args.nhead_stride_o, - args.batch_stride_q, args.batch_stride_k, args.batch_stride_v, - args.batch_stride_bias, args.batch_stride_randval, - args.batch_stride_lse, args.batch_stride_o, args.window_size_left, - args.window_size_right, args.mask_type, args.p_drop, args.s_randval, - args.drop_seed_offset, args.cu_seqlen_q_ptr, args.cu_seqlen_k_ptr); - } - }(); - - if constexpr (FmhaKernel::kIsGroupMode) { - dim3 grids = - FmhaKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q, - args.hdim_v, args.seqlen_k_ptr != nullptr); - return ck_tile::make_tuple(kargs, grids); - } else { - dim3 grids = FmhaKernel::GridSize(args.batch, args.nhead_q, - args.max_seqlen_q, args.hdim_v, false); - return ck_tile::make_tuple(kargs, grids); - } -} - -template -auto fmha_fwd_pagedkv_create_kargs_and_grids(fmha_fwd_pagedkv_args args) { - assert(args.nhead_q % args.nhead_k == 0); - auto kargs = [&] { - // create group mode kernel arguments - if constexpr (FmhaKernel::kIsGroupMode) { - return FmhaKernel::MakeKargs( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.lse_ptr, - args.o_ptr, args.seqstart_q_ptr, args.seqstart_k_ptr, - args.seqlen_k_ptr, args.hdim_q, args.hdim_v, args.nhead_q, - args.nhead_q / args.nhead_k, args.block_table_ptr, - args.batch_stride_block_table, args.page_block_size, args.is_gappy, - args.scale_s, args.scale_p, args.scale_o, args.logits_soft_cap, - args.stride_q, args.stride_k, args.stride_v, args.stride_bias, - args.stride_o, args.nhead_stride_q, args.nhead_stride_k, - args.nhead_stride_v, args.nhead_stride_bias, args.nhead_stride_lse, - args.nhead_stride_o, args.batch_stride_k, args.batch_stride_v, - args.window_size_left, args.window_size_right, args.mask_type, - args.min_seqlen_q); - } else { // create batch mode kernel arguments - return FmhaKernel::MakeKargs( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.lse_ptr, - args.o_ptr, args.seqlen_q, args.seqlen_k, args.seqlen_k_ptr, - args.hdim_q, args.hdim_v, args.nhead_q, args.nhead_q / args.nhead_k, - args.block_table_ptr, args.batch_stride_block_table, - args.page_block_size, args.cache_batch_idx, args.scale_s, - args.scale_p, args.scale_o, args.logits_soft_cap, args.stride_q, - args.stride_k, args.stride_v, args.stride_bias, args.stride_o, - args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, - args.nhead_stride_bias, args.nhead_stride_lse, args.nhead_stride_o, - args.batch_stride_q, args.batch_stride_k, args.batch_stride_v, - args.batch_stride_bias, args.batch_stride_lse, args.batch_stride_o, - args.window_size_left, args.window_size_right, args.mask_type); - } - }(); - - // FmhaKernel::PrintParameters(kargs, args.batch); - if constexpr (FmhaKernel::kIsGroupMode) { - dim3 grids = - FmhaKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q, - args.hdim_v, args.seqlen_k_ptr != nullptr); - return ck_tile::make_tuple(kargs, grids); - } else { - dim3 grids = FmhaKernel::GridSize(args.batch, args.nhead_q, - args.max_seqlen_q, args.hdim_v, false); - return ck_tile::make_tuple(kargs, grids); - } -} - -template -auto fmha_fwd_splitkv_create_kargs_and_grids(fmha_fwd_splitkv_args args) { - assert(args.nhead_q % args.nhead_k == 0); - auto kargs = [&] { - // create group mode kernel arguments - if constexpr (Kernel::kIsGroupMode) { - return Kernel::MakeKargs( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.lse_acc_ptr, - args.o_acc_ptr, args.batch, args.seqstart_q_ptr, args.seqstart_k_ptr, - args.seqlen_k_ptr, args.hdim_q, args.hdim_v, args.nhead_q, - args.nhead_q / args.nhead_k, args.num_splits, args.block_table_ptr, - args.batch_stride_block_table, args.page_block_size, args.is_gappy, - args.scale_s, args.scale_p, args.logits_soft_cap, args.stride_q, - args.stride_k, args.stride_v, args.stride_bias, args.stride_o_acc, - args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, - args.nhead_stride_bias, args.nhead_stride_lse_acc, - args.nhead_stride_o_acc, - args.batch_stride_k, // only used for paged-kvcache - args.batch_stride_v, // only used for paged-kvcache - args.split_stride_lse_acc, args.split_stride_o_acc, - args.window_size_left, args.window_size_right, args.mask_type); - } else { // create batch mode kernel arguments - return Kernel::MakeKargs( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.lse_acc_ptr, - args.o_acc_ptr, args.batch, args.seqlen_q, args.seqlen_k, - args.seqlen_k_ptr, args.hdim_q, args.hdim_v, args.nhead_q, - args.nhead_q / args.nhead_k, args.num_splits, args.block_table_ptr, - args.batch_stride_block_table, args.page_block_size, - args.cache_batch_idx, args.scale_s, args.scale_p, - args.logits_soft_cap, args.stride_q, args.stride_k, args.stride_v, - args.stride_bias, args.stride_o_acc, args.nhead_stride_q, - args.nhead_stride_k, args.nhead_stride_v, args.nhead_stride_bias, - args.nhead_stride_lse_acc, args.nhead_stride_o_acc, - args.batch_stride_q, args.batch_stride_k, args.batch_stride_v, - args.batch_stride_bias, args.batch_stride_lse_acc, - args.batch_stride_o_acc, args.split_stride_lse_acc, - args.split_stride_o_acc, args.window_size_left, - args.window_size_right, args.mask_type); - } - }(); - - dim3 grids = - Kernel::GridSize(args.batch, args.nhead_q, args.nhead_k, - args.max_seqlen_q, args.hdim_v, args.num_splits); - - return ck_tile::make_tuple(kargs, grids); -} - -template -auto fmha_fwd_splitkv_combine_create_kargs_and_grids( - fmha_fwd_splitkv_args args) { - assert(args.nhead_q % args.nhead_k == 0); - auto kargs = [&] { - // create group mode kernel argumentszs - if constexpr (Kernel::kIsGroupMode) { - return Kernel::MakeKargs( - args.lse_acc_ptr, args.o_acc_ptr, args.lse_ptr, args.o_ptr, - args.batch, args.seqstart_q_ptr, args.hdim_v, args.num_splits, - args.scale_o, args.stride_o_acc, args.stride_o, - args.nhead_stride_lse_acc, args.nhead_stride_o_acc, - args.nhead_stride_lse, args.nhead_stride_o, args.split_stride_lse_acc, - args.split_stride_o_acc); - } else { // create batch mode kernel arguments - return Kernel::MakeKargs( - args.lse_acc_ptr, args.o_acc_ptr, args.lse_ptr, args.o_ptr, - args.batch, args.seqlen_q, args.hdim_v, args.num_splits, args.scale_o, - args.stride_o_acc, args.stride_o, args.nhead_stride_lse_acc, - args.nhead_stride_o_acc, args.nhead_stride_lse, args.nhead_stride_o, - args.batch_stride_lse_acc, args.batch_stride_o_acc, - args.batch_stride_lse, args.batch_stride_o, args.split_stride_lse_acc, - args.split_stride_o_acc); - } - }(); - - dim3 grids = Kernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q, - args.hdim_v); - - return ck_tile::make_tuple(kargs, grids); -} - -template -auto fmha_fwd_appendkv_create_kargs_and_grids(fmha_fwd_appendkv_args args) { - assert(args.nhead_q % args.nhead_k == 0); - auto kargs = Kernel::MakeKargs( - args.q_ptr, args.k_ptr, args.knew_ptr, args.v_ptr, args.vnew_ptr, - args.seqlen_q, args.seqlen_k_ptr, args.seqlen_knew, args.hdim_q, - args.hdim_v, args.nhead_q, args.nhead_q / args.nhead_k, - args.rotary_cos_ptr, args.rotary_sin_ptr, args.rotary_dim, args.has_mask, - args.block_table_ptr, args.batch_stride_block_table, args.page_block_size, - args.cache_batch_idx, args.stride_q, args.stride_k, args.stride_knew, - args.stride_v, args.stride_vnew, args.nhead_stride_q, args.nhead_stride_k, - args.nhead_stride_knew, args.nhead_stride_v, args.nhead_stride_vnew, - args.batch_stride_q, args.batch_stride_k, args.batch_stride_knew, - args.batch_stride_v, args.batch_stride_vnew); - - dim3 grids = Kernel::GridSize(args.batch, args.nhead_q, args.seqlen_q, - args.seqlen_knew); - - return ck_tile::make_tuple(kargs, grids); -} - -template -auto fmha_batch_prefill_create_kargs_and_grids(fmha_batch_prefill_args args) { - assert(args.nhead_q % args.nhead_k == 0); - auto kargs = [&] { - // create group mode kernel arguments - if constexpr (FmhaKernel::kIsGroupMode) { - return FmhaKernel::MakeKargs( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.rand_val_ptr, - args.lse_ptr, args.o_ptr, args.seqstart_q_ptr, args.hdim_q, - args.hdim_v, args.nhead_q, args.nhead_q / args.nhead_k, - args.num_total_pages, args.kv_indptr, args.kv_page_indices, -#if 0 // we assume page_block_size=1 for now - args.kv_last_page_lens, - args.page_block_size, -#endif - args.scale_s, args.scale_p, args.scale_o, args.logits_soft_cap, - args.stride_q, args.stride_k, args.stride_v, args.stride_bias, - args.stride_randval, args.stride_o, args.nhead_stride_q, - args.nhead_stride_k, args.nhead_stride_v, args.nhead_stride_bias, - args.nhead_stride_randval, args.nhead_stride_lse, args.nhead_stride_o, - args.batch_stride_k, args.batch_stride_v, args.window_size_left, - args.window_size_right, args.mask_type, args.p_drop, args.s_randval, - args.drop_seed_offset); - } else { // create batch mode kernel arguments - return FmhaKernel::MakeKargs( - args.q_ptr, args.k_ptr, args.v_ptr, args.bias_ptr, args.rand_val_ptr, - args.lse_ptr, args.o_ptr, args.seqlen_q, args.hdim_q, args.hdim_v, - args.nhead_q, args.nhead_q / args.nhead_k, args.num_total_pages, - args.kv_indptr, args.kv_page_indices, -#if 0 // we assume page_block_size=1 for now - args.kv_last_page_lens, - args.page_block_size, -#endif - args.scale_s, args.scale_p, args.scale_o, args.logits_soft_cap, - args.stride_q, args.stride_k, args.stride_v, args.stride_bias, - args.stride_randval, args.stride_o, args.nhead_stride_q, - args.nhead_stride_k, args.nhead_stride_v, args.nhead_stride_bias, - args.nhead_stride_randval, args.nhead_stride_lse, args.nhead_stride_o, - args.batch_stride_q, args.batch_stride_k, args.batch_stride_v, - args.batch_stride_bias, args.batch_stride_randval, - args.batch_stride_lse, args.batch_stride_o, args.window_size_left, - args.window_size_right, args.mask_type, args.p_drop, args.s_randval, - args.drop_seed_offset); - } - }(); - - dim3 grids = FmhaKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q, - args.hdim_v); - return ck_tile::make_tuple(kargs, grids); -} - -// this is used to pattern-match internl kernel implementation, not to -// instantiate kernel -template -struct fmha_fwd_traits_ { - static constexpr ck_tile::index_t HDim = HDim_; - using DataType = ck_tile::remove_cvref_t; - static constexpr bool kIsGroupMode = kIsGroupMode_; - static constexpr ck_tile::index_t kM0 = kM0_; - static constexpr ck_tile::index_t kN0 = kN0_; - static constexpr ck_tile::index_t kK0 = kK0_; - static constexpr ck_tile::index_t kN1 = kN1_; - static constexpr ck_tile::index_t kK1 = kK1_; - static constexpr ck_tile::index_t kK0BlockLength = kK0BlockLength_; - static constexpr bool kIsVLayoutRowMajor = kIsVLayoutRowMajor_; - static constexpr auto FmhaPipelineEnum = FmhaPipelineEnum_; - static constexpr bool kHasLogitsSoftCap = kHasLogitsSoftCap_; - using FmhaMask = ck_tile::remove_cvref_t; - static constexpr auto BiasEnum = BiasEnum_; - static constexpr bool kStoreLse = kStoreLse_; - static constexpr bool kHasDropout = kHasDropout_; - static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_; - static constexpr bool kPadS = kPadS_; - static constexpr bool kPadSK = kPadSK_; - static constexpr bool kPadD = kPadD_; - static constexpr bool kPadDv = kPadDv_; - static constexpr bool kUseTrLoad = kUseTrLoad_; - static constexpr bool kSkipMinSeqlenQ = kSkipMinSeqlenQ_; -}; - -template -float fmha_fwd_(const ck_tile::stream_config &, fmha_fwd_args); - -template -struct fmha_fwd_pagedkv_traits_ { - static constexpr ck_tile::index_t HDim = HDim_; - using DataType = ck_tile::remove_cvref_t; - static constexpr bool kIsGroupMode = kIsGroupMode_; - static constexpr ck_tile::index_t kM0 = kM0_; - static constexpr ck_tile::index_t kN0 = kN0_; - static constexpr ck_tile::index_t kK0 = kK0_; - static constexpr ck_tile::index_t kN1 = kN1_; - static constexpr ck_tile::index_t kK1 = kK1_; - static constexpr ck_tile::index_t kK0BlockLength = kK0BlockLength_; - static constexpr bool kIsVLayoutRowMajor = kIsVLayoutRowMajor_; - static constexpr auto FmhaPipelineEnum = FmhaPipelineEnum_; - static constexpr bool kHasLogitsSoftCap = kHasLogitsSoftCap_; - using FmhaMask = ck_tile::remove_cvref_t; - static constexpr auto BiasEnum = BiasEnum_; - static constexpr bool kStoreLse = kStoreLse_; - static constexpr bool kIsPagedKV = kIsPagedKV_; - static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_; - static constexpr bool kPadS = kPadS_; - static constexpr bool kPadSK = kPadSK_; - static constexpr bool kPadD = kPadD_; - static constexpr bool kPadDv = kPadDv_; - static constexpr bool kSkipMinSeqlenQ = kSkipMinSeqlenQ_; -}; - -template -float fmha_fwd_pagedkv_(const ck_tile::stream_config &, fmha_fwd_pagedkv_args); - -template -struct fmha_fwd_splitkv_traits_ { - static constexpr ck_tile::index_t HDim = HDim_; - using DataType = ck_tile::remove_cvref_t; - static constexpr bool kIsGroupMode = kIsGroupMode_; - static constexpr ck_tile::index_t kM0 = kM0_; - static constexpr ck_tile::index_t kN0 = kN0_; - static constexpr ck_tile::index_t kK0 = kK0_; - static constexpr ck_tile::index_t kN1 = kN1_; - static constexpr ck_tile::index_t kK1 = kK1_; - static constexpr ck_tile::index_t kK0BlockLength = kK0BlockLength_; - static constexpr bool kIsVLayoutRowMajor = kIsVLayoutRowMajor_; - static constexpr auto FmhaPipelineEnum = FmhaPipelineEnum_; - static constexpr bool kHasLogitsSoftCap = kHasLogitsSoftCap_; - using FmhaMask = ck_tile::remove_cvref_t; - static constexpr auto BiasEnum = BiasEnum_; - static constexpr bool kStoreLse = kStoreLse_; - static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_; - static constexpr bool kPadS = kPadS_; - static constexpr bool kPadSK = kPadSK_; - static constexpr bool kPadD = kPadD_; - static constexpr bool kPadDv = kPadDv_; - static constexpr bool kIsPagedKV = kIsPagedKV_; -}; - -template -void fmha_fwd_splitkv_oneshot_(const ck_tile::stream_config &, - fmha_fwd_splitkv_args); - -template -std::string fmha_fwd_splitkv_get_name_(); - -template -struct fmha_fwd_splitkv_combine_traits_ { - static constexpr ck_tile::index_t HDim = HDim_; - using DataType = ck_tile::remove_cvref_t; - static constexpr bool kIsGroupMode = kIsGroupMode_; - static constexpr ck_tile::index_t kN1 = kN1_; - static constexpr bool kStoreLse = kStoreLse_; - static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_; - static constexpr bool kPadS = kPadS_; - static constexpr bool kPadDv = kPadDv_; -}; - -template -void fmha_fwd_splitkv_combine_oneshot_(const ck_tile::stream_config &, - fmha_fwd_splitkv_args); - -template -std::string fmha_fwd_splitkv_combine_get_name_(); - -// this is used to pattern-match internl kernel implementation, not to -// instantiate kernel -template -struct fmha_fwd_appendkv_traits_ { - static constexpr ck_tile::index_t HDim = HDim_; - using DataType = ck_tile::remove_cvref_t; - static constexpr ck_tile::index_t kTileSizeS = kTileSizeS_; - static constexpr ck_tile::index_t kTileSizeSk = kTileSizeSk_; - static constexpr ck_tile::index_t kTileSizeD = kTileSizeD_; - static constexpr ck_tile::index_t kTileSizeDv = kTileSizeDv_; - static constexpr bool kIsVLayoutRowMajor = kIsVLayoutRowMajor_; - static constexpr bool kPadS = kPadS_; - static constexpr bool kPadSk = kPadSk_; - static constexpr bool kPadD = kPadD_; - static constexpr bool kPadDv = kPadDv_; - static constexpr auto RotaryEnum = RotaryEnum_; - static constexpr bool kIsPagedKV = kIsPagedKV_; -}; - -template -float fmha_fwd_appendkv_(const ck_tile::stream_config &, - fmha_fwd_appendkv_args); - -template -float fmha_batch_prefill_(const ck_tile::stream_config &, - fmha_batch_prefill_args); - -// This is the public API, will be generated by script -struct fmha_fwd_traits { - int hdim_q; - int hdim_v; - std::string data_type; - bool is_group_mode; - bool is_v_rowmajor; - bool has_logits_soft_cap; - mask_enum mask_type; - bias_enum bias_type; // 0:no bias, 1:elementwise bias, 2:alibi. sync with - // BlockAttentionBiasEnum - bool has_lse; - bool has_dropout; - bool do_fp8_static_quant; - bool skip_min_seqlen_q = false; - // TODO: padding check is inside this api -}; -float fmha_fwd(fmha_fwd_traits, fmha_fwd_args, const ck_tile::stream_config &); - -struct fmha_fwd_pagedkv_traits { - int hdim_q; - int hdim_v; - std::string data_type; - bool is_group_mode; - bool is_v_rowmajor; - bool has_logits_soft_cap; - mask_enum mask_type; - bias_enum bias_type; // 0:no bias, 1:elementwise bias, 2:alibi. sync with - // BlockAttentionBiasEnum - bool has_lse = false; - bool use_pagedkv = true; - bool do_fp8_static_quant = false; - bool skip_min_seqlen_q = false; - // TODO: padding check is inside this api -}; - -float fmha_fwd_pagedkv(fmha_fwd_pagedkv_traits &, fmha_fwd_pagedkv_args &, - const ck_tile::stream_config &); - -struct fmha_fwd_splitkv_traits { - int hdim_q; - int hdim_v; - std::string data_type; - bool is_group_mode; - bool is_v_rowmajor; - bool has_logits_soft_cap; - mask_enum mask_type; - bias_enum bias_type; // 0:no bias, 1:elementwise bias, 2:alibi. sync with - // BlockAttentionBiasEnum - bool has_lse; - bool do_fp8_static_quant; - // TODO: padding check is inside this api -}; -float fmha_fwd_splitkv(fmha_fwd_splitkv_traits, fmha_fwd_splitkv_args, - const ck_tile::stream_config &); - -struct fmha_fwd_appendkv_traits { - int hdim_q; - int hdim_v; - std::string data_type; - bool is_v_rowmajor; - rope_enum rope_type; -}; -float fmha_fwd_appendkv(fmha_fwd_appendkv_traits, fmha_fwd_appendkv_args, - const ck_tile::stream_config &); - -using fmha_batch_prefill_traits = fmha_fwd_traits; -float fmha_batch_prefill(fmha_batch_prefill_traits, fmha_batch_prefill_args, - const ck_tile::stream_config &); diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3.cpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3.cpp deleted file mode 100644 index 9dbb95a69b..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3.cpp +++ /dev/null @@ -1,76 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#include "fmha_fwd_v3.hpp" -#include "fmha_fwd_v3_impl.hpp" -#include "mask.hpp" - -namespace ck_tile { - -std::ostream &operator<<(std::ostream &stream, - const fmha_fwd_v3_args::data_type_enum &data_type) { - switch (data_type) { - case fmha_fwd_v3_args::data_type_enum::fp16: - return stream << "fp16"; - case fmha_fwd_v3_args::data_type_enum::bf16: - return stream << "bf16"; - default: - return stream << "unknown"; - } -} - -std::pair fmha_fwd_v3(const fmha_fwd_v3_args &args, - const stream_config &config) { - if (args.data_type == fmha_fwd_v3_args::data_type_enum::fp16) { - if (args.mask_type == static_cast(mask_enum::no_mask)) { - using kernel_traits = - fmha_fwd_v3_kernel_traits; - - return fmha_fwd_v3_kernel_dispatch(args, config); - } else { - using kernel_traits = - fmha_fwd_v3_kernel_traits; - - return fmha_fwd_v3_kernel_dispatch(args, config); - } - } else if (args.data_type == fmha_fwd_v3_args::data_type_enum::bf16) { - if (args.mask_type == static_cast(mask_enum::no_mask)) { - using kernel_traits = - fmha_fwd_v3_kernel_traits; - - return fmha_fwd_v3_kernel_dispatch(args, config); - } else { - using kernel_traits = - fmha_fwd_v3_kernel_traits; - - return fmha_fwd_v3_kernel_dispatch(args, config); - } - } - - return std::make_pair(false, -1.f); -} - -} // namespace ck_tile diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3.hpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3.hpp deleted file mode 100644 index db67889ebc..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3.hpp +++ /dev/null @@ -1,91 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#pragma once - -#include -#include - -#include "ck_tile/core/numeric/integer.hpp" -#include "ck_tile/host/stream_config.hpp" - -namespace ck_tile { - -struct fmha_fwd_v3_args { - enum class data_type_enum { fp16, bf16 }; - - data_type_enum data_type; - // bool is_varlen; - - index_t batch; - index_t seqlen_q; - index_t seqlen_k; - index_t nhead_q; - index_t nhead_kv; - index_t hdim_qk; - index_t hdim_v; - - float softmax_scale; - - index_t window_size_left; - index_t window_size_right; - index_t mask_type; // should be 0 for no mask; or 2 for causal mask - // (window_size_left < 0 and window_size_right == 0). - - const void *q_ptr; - index_t stride_q; - index_t nhead_stride_q; - index_t batch_stride_q; - - const void *k_ptr; - index_t stride_k; - index_t nhead_stride_k; - index_t batch_stride_k; - - const void *v_ptr; - index_t stride_v; - index_t nhead_stride_v; - index_t batch_stride_v; - - void *o_ptr; - index_t stride_o; - index_t nhead_stride_o; - index_t batch_stride_o; - - // Optional batch-mode cumulative seqlen overrides (exclude PAD) - // If provided, they override per-batch effective lengths to skip tail - // padding. - const ck_tile::index_t *cu_seqlen_q_ptr = nullptr; // [batch+1] - const ck_tile::index_t *cu_seqlen_kv_ptr = nullptr; // [batch+1] -}; - -std::ostream &operator<<(std::ostream &stream, - const fmha_fwd_v3_args::data_type_enum &data_type); - -// return value: -// first = whether the kernel was launched (true = launched, false = skipped) -// second = elapsed time (ms) of the kernel launch, valid only if first == -// true -std::pair fmha_fwd_v3(const fmha_fwd_v3_args &args, - const stream_config &config); - -} // namespace ck_tile diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3_impl.hpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3_impl.hpp deleted file mode 100644 index fd76b28001..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/fmha_fwd_v3_impl.hpp +++ /dev/null @@ -1,170 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#pragma once - -#include - -#include "ck_tile/core/container/sequence.hpp" -#include "ck_tile/core/numeric/bfloat16.hpp" -#include "ck_tile/core/numeric/half.hpp" -#include "ck_tile/host/kernel_launch.hpp" -#include "ck_tile/ops/epilogue/default_2d_epilogue.hpp" -#include "ck_tile/ops/fmha/block/block_masking.hpp" -#include "ck_tile/ops/fmha/kernel/fmha_fwd_v3_kernel.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_fwd_v3_pipeline.hpp" -#include "ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp" -#include "ck_tile/ops/fmha/pipeline/tile_fmha_shape.hpp" -#include "ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp" - -#include "fmha_fwd_v3.hpp" -#include "mask.hpp" - -#define INST_FMHA_FWD_V3_DISPATCH(kernel_traits) \ - template <> \ - std::pair fmha_fwd_v3_kernel_dispatch( \ - const fmha_fwd_v3_args &args, const stream_config &config) { \ - return std::make_pair( \ - true, fmha_fwd_v3_kernel_launch(args, config)); \ - } - -namespace ck_tile { - -template -struct fmha_fwd_v3_problem_traits; - -template <> -struct fmha_fwd_v3_problem_traits { - using qkvp_dtype = ck_tile::half_t; - using acc_dtype = float; - using o_dtype = ck_tile::half_t; - using lse_dtype = float; -}; - -template <> -struct fmha_fwd_v3_problem_traits { - using qkvp_dtype = ck_tile::bf16_t; - using acc_dtype = float; - using o_dtype = ck_tile::bf16_t; - using lse_dtype = float; -}; - -template -struct fmha_fwd_v3_kernel_traits { - static constexpr auto date_type = DataType; - static constexpr bool is_variable_seqlen = IsVariableSeqlen; - static constexpr bool is_masking = IsMasking; - - // M0 N0 K0 N1 K1 - using fmha_block_tile = sequence<256, 32, 128, 128, 32, 128>; - using fmha_warp_gemm_shape = sequence<32, 32, 16>; - using fmha_block_warps = sequence<8, 1, 1>; - - using fmha_shape = - TileFmhaShape; - - using fmha_traits = TileFmhaFwdV3Traits; - - using fmha_mask = GenericAttentionMask; - - using fmha_pipeline_problem = BlockFmhaFwdV3PipelineProblem< - typename fmha_fwd_v3_problem_traits::qkvp_dtype, - typename fmha_fwd_v3_problem_traits::qkvp_dtype, - typename fmha_fwd_v3_problem_traits::qkvp_dtype, - typename fmha_fwd_v3_problem_traits::acc_dtype, - typename fmha_fwd_v3_problem_traits::acc_dtype, - typename fmha_fwd_v3_problem_traits::lse_dtype, - typename fmha_fwd_v3_problem_traits::qkvp_dtype, - typename fmha_fwd_v3_problem_traits::acc_dtype, - typename fmha_fwd_v3_problem_traits::o_dtype, fmha_shape, - IsVariableSeqlen, fmha_mask, fmha_traits>; - - using fmha_pipeline = BlockFmhaFwdV3Pipeline; - - using epilogue = Default2DEpilogue::acc_dtype, - typename fmha_fwd_v3_problem_traits::o_dtype, - true, // kPadM - true, // kPadM - true // UseRawStore - >>; - - using kernel = FmhaFwdV3Kernel; -}; - -template -float fmha_fwd_v3_kernel_launch(const fmha_fwd_v3_args &args, - const stream_config &config) { - /// NOTICE: This was borrowed from Aiter. Make sure the selected remap_opt - /// setting truly maximizes the kernel's performance. - int remap_opt = 2; - if (args.mask_type != static_cast(mask_enum::no_mask) && - ((args.nhead_q % 8 != 0) || (16384 < args.seqlen_q))) { - if (65536 <= args.seqlen_q) { - remap_opt = 0; - } else { - remap_opt = 1; - } - } - - auto kargs = Kernel::MakeKargs( - args.q_ptr, args.k_ptr, args.v_ptr, - nullptr, // lse_ptr - args.o_ptr, args.seqlen_q, args.seqlen_k, args.hdim_qk, args.hdim_v, - args.nhead_q, args.nhead_q / args.nhead_kv, args.softmax_scale, - args.stride_q, args.stride_k, args.stride_v, args.stride_o, - args.nhead_stride_q, args.nhead_stride_k, args.nhead_stride_v, - 0, // nhead_stride_lse - args.nhead_stride_o, args.batch_stride_q, args.batch_stride_k, - args.batch_stride_v, - 0, // batch_stride_lse - args.batch_stride_o, args.window_size_left, args.window_size_right, - args.mask_type, remap_opt, args.cu_seqlen_q_ptr, args.cu_seqlen_kv_ptr); - - dim3 grids = - Kernel::GridSize(args.batch, args.nhead_q, args.seqlen_q, args.hdim_v); - constexpr dim3 blocks = Kernel::BlockSize(); - constexpr index_t kBlockPerCu = Kernel::kBlockPerCu; - - return launch_kernel( - config, make_kernel(Kernel{}, grids, blocks, 0, kargs)); -} - -// return value: -// first = whether the kernel was launched (true = launched, false = skipped) -// second = elapsed time (ms) of the kernel launch, valid only if first == -// true -template -std::pair fmha_fwd_v3_kernel_dispatch(const fmha_fwd_v3_args &args, - const stream_config &config); - -} // namespace ck_tile diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_bf16_mask.cpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_bf16_mask.cpp deleted file mode 100644 index 7a7664acac..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_bf16_mask.cpp +++ /dev/null @@ -1,34 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#include "fmha_fwd_v3.hpp" -#include "fmha_fwd_v3_impl.hpp" - -namespace ck_tile { - -using kernel_traits = - fmha_fwd_v3_kernel_traits; - -INST_FMHA_FWD_V3_DISPATCH(kernel_traits) - -} // namespace ck_tile diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_bf16_nmask.cpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_bf16_nmask.cpp deleted file mode 100644 index 777ef4754e..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_bf16_nmask.cpp +++ /dev/null @@ -1,34 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#include "fmha_fwd_v3.hpp" -#include "fmha_fwd_v3_impl.hpp" - -namespace ck_tile { - -using kernel_traits = - fmha_fwd_v3_kernel_traits; - -INST_FMHA_FWD_V3_DISPATCH(kernel_traits) - -} // namespace ck_tile diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_fp16_mask.cpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_fp16_mask.cpp deleted file mode 100644 index f54e556b12..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_fp16_mask.cpp +++ /dev/null @@ -1,34 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#include "fmha_fwd_v3.hpp" -#include "fmha_fwd_v3_impl.hpp" - -namespace ck_tile { - -using kernel_traits = - fmha_fwd_v3_kernel_traits; - -INST_FMHA_FWD_V3_DISPATCH(kernel_traits) - -} // namespace ck_tile diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_fp16_nmask.cpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_fp16_nmask.cpp deleted file mode 100644 index 4cd91e5545..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/instances/fmha_fwd_v3_d128_fp16_nmask.cpp +++ /dev/null @@ -1,34 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#include "fmha_fwd_v3.hpp" -#include "fmha_fwd_v3_impl.hpp" - -namespace ck_tile { - -using kernel_traits = - fmha_fwd_v3_kernel_traits; - -INST_FMHA_FWD_V3_DISPATCH(kernel_traits) - -} // namespace ck_tile diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/mask.hpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/mask.hpp deleted file mode 100644 index 3c13e7a55e..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/mask.hpp +++ /dev/null @@ -1,162 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#pragma once - -#include -#include - -#include "ck_tile/core.hpp" -#include "ck_tile/ops/fmha.hpp" - -// keep this in sync with ck_tile::GenericAttentionMaskEnum -enum class mask_enum { - no_mask = 0, - mask_top_left, - mask_bottom_right, - window_generic, -}; - -struct mask_info { - mask_enum type; - ck_tile::index_t seqlen_q; - ck_tile::index_t seqlen_k; - ck_tile::index_t y, x; - ck_tile::index_t left, right; // FA style SWA left/right - - void serialize(std::ostream &os) const { - if (type == mask_enum::no_mask) - os << "n"; - else if (type == mask_enum::mask_top_left) - os << "t(" << left << ":" << right << ")"; - else if (type == mask_enum::mask_bottom_right) - os << "b(" << left << ":" << right << ")"; - else { - os << "g(" << y << ":" << x << ")"; - } - } - - static mask_info decode(std::string str, ck_tile::index_t seqlen_q, - ck_tile::index_t seqlen_k) { - ck_tile::index_t x_total = seqlen_k; - ck_tile::index_t y_total = seqlen_q; - mask_info tmp; - tmp.seqlen_q = seqlen_q; - tmp.seqlen_k = seqlen_k; - auto found_0 = str.find(':'); - if (found_0 != std::string::npos) { - std::string t = str.substr(0, found_0); - std::string v = str.substr(found_0 + 1); - if (t == "xt" || t == "xb") { - // xformer style sliding window attn from top-left - ck_tile::index_t window_size = std::stoi(v); - ck_tile::index_t left_size = -1; - ck_tile::index_t right_size = 0; - if (window_size > 0) { - left_size = window_size / 2; - right_size = window_size - 1 - left_size; - } - auto r = - ck_tile::make_generic_attention_mask_coordinates_from_lr_window( - left_size, right_size, y_total, x_total, t == "xt"); - - tmp.type = - t == "xt" ? mask_enum::mask_top_left : mask_enum::mask_bottom_right; - tmp.y = r.at(ck_tile::number<0>{}); - tmp.x = r.at(ck_tile::number<1>{}); - tmp.left = left_size; - tmp.right = right_size; - } else if (t == "t" || t == "b" || t == "g") { - auto found_1 = v.find(","); - if (found_1 == std::string::npos) { - throw std::invalid_argument("invalid mask value: " + str); - } - ck_tile::index_t v0 = std::stoi(v.substr(0, found_1)); - ck_tile::index_t v1 = std::stoi(v.substr(found_1 + 1)); - if (t == "t") { - tmp.type = mask_enum::mask_top_left; - auto r = - ck_tile::make_generic_attention_mask_coordinates_from_lr_window( - v0, v1, y_total, x_total, true); - tmp.y = r.at(ck_tile::number<0>{}); - tmp.x = r.at(ck_tile::number<1>{}); - tmp.left = v0; - tmp.right = v1; - } else if (t == "b") { - tmp.type = mask_enum::mask_bottom_right; - auto r = - ck_tile::make_generic_attention_mask_coordinates_from_lr_window( - v0, v1, y_total, x_total, false); - tmp.y = r.at(ck_tile::number<0>{}); - tmp.x = r.at(ck_tile::number<1>{}); - tmp.left = v0; - tmp.right = v1; - } else if (t == "g") { - tmp.type = mask_enum::window_generic; - tmp.y = v0; - tmp.x = v1; - tmp.left = v0; // TODO: don't use this? - tmp.right = v1; - } - } else { - throw std::invalid_argument("invalid mask value: " + str); - } - } else if (str == "0") { - tmp.type = mask_enum::no_mask; - } else if (str == "1" || str == "t") { - tmp.type = mask_enum::mask_top_left; - tmp.y = seqlen_q; - tmp.x = 1; - tmp.left = -1; - tmp.right = 0; - } else if (str == "2" || str == "b") { - tmp.type = mask_enum::mask_bottom_right; - tmp.y = seqlen_q; - tmp.x = seqlen_k - seqlen_q + 1; - tmp.left = -1; - tmp.right = 0; - } else { - throw std::invalid_argument("invalid mask value: " + str); - } - return tmp; - } - - ck_tile::index_t get_unmaskarea() const { - if (type == mask_enum::no_mask) - return seqlen_q * seqlen_k; - ck_tile::index_t area = 0; - for (ck_tile::index_t i_y = 0; i_y < seqlen_q; ++i_y) { - ck_tile::index_t x_start = - std::max(-y + i_y + 1, static_cast(0)); - ck_tile::index_t x_end = std::min(i_y + x, seqlen_k); - if (x_end > x_start) { - area += (x_end - x_start); - } - } - return area; - } - - friend std::ostream &operator<<(std::ostream &os, const mask_info &mi) { - mi.serialize(os); - return os; - } -}; diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/rotary.hpp b/Libraries/ComposableKernel/attention/fmha_fwd_v3/rotary.hpp deleted file mode 100644 index cd794ec394..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/rotary.hpp +++ /dev/null @@ -1,101 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#pragma once - -#include "ck_tile/core.hpp" -#include "ck_tile/host/host_tensor.hpp" - -#include -#include -#include -#include -#include -#include -#include - -// keep sync with RotaryEmbeddingEnum -enum class rope_enum { - none = 0, - interleaved = 1, - half_rotated = 2, -}; - -template -std::tuple, ck_tile::HostTensor> -generate_rotary_cos_sin(ck_tile::index_t seqlen, ck_tile::index_t rotary_dim, - std::optional seed = std::nullopt) { - // return dummy tensors if we won't apply RoPE at all - if (rotary_dim <= 0) { - ck_tile::HostTensor dummy({1, 1}); - return std::make_tuple(dummy, dummy); - } - - std::mt19937 random_engine(seed.has_value() ? *seed : std::random_device{}()); - std::uniform_real_distribution generator(0.0f, 1.0f); - - const ck_tile::index_t num_rows = seqlen * 2; - const ck_tile::index_t num_cols = rotary_dim / 2; - - using std::begin, std::end; - - ck_tile::HostTensor angle({num_rows, num_cols}); - std::generate(begin(angle), end(angle), - [&] { return generator(random_engine) * 2 * M_PI; }); - - ck_tile::HostTensor cos({num_rows, num_cols}); - std::transform(begin(angle), end(angle), begin(cos), [](float origin_value) { - return ck_tile::type_convert(std::cos(origin_value)); - }); - - ck_tile::HostTensor sin({num_rows, num_cols}); - std::transform(begin(angle), end(angle), begin(sin), [](float origin_value) { - return ck_tile::type_convert(std::sin(origin_value)); - }); - - return std::make_tuple(cos, sin); -} - -template -std::tuple, ck_tile::HostTensor> -slice_rotary_cos_sin(const ck_tile::HostTensor &cos, - const ck_tile::HostTensor &sin, - ck_tile::index_t seqlen_offset, ck_tile::index_t seqlen) { - assert(cos.get_num_of_dimension() == 2 && sin.get_num_of_dimension() == 2); - assert(cos.get_length(0) == sin.get_length(0) && - cos.get_length(1) == sin.get_length(1)); - - assert(static_cast(seqlen_offset + seqlen) <= cos.get_length(0)); - - const ck_tile::index_t num_rows = seqlen; - const ck_tile::index_t num_cols = cos.get_length(1); - - ck_tile::HostTensor cos_pt({num_rows, num_cols}); - cos_pt.ForEach( - [&](auto &self, auto i) { self(i) = cos(i[0] + seqlen_offset, i[1]); }); - - ck_tile::HostTensor sin_pt({num_rows, num_cols}); - sin_pt.ForEach( - [&](auto &self, auto i) { self(i) = sin(i[0] + seqlen_offset, i[1]); }); - - return std::make_tuple(cos_pt, sin_pt); -} diff --git a/Libraries/ComposableKernel/attention/fmha_fwd_v3/test_ck_tile_fmha_fwd_v3.sh b/Libraries/ComposableKernel/attention/fmha_fwd_v3/test_ck_tile_fmha_fwd_v3.sh deleted file mode 100755 index d1a2ebcc87..0000000000 --- a/Libraries/ComposableKernel/attention/fmha_fwd_v3/test_ck_tile_fmha_fwd_v3.sh +++ /dev/null @@ -1,46 +0,0 @@ -#!/bin/sh - -## Copyright © Advanced Micro Devices, Inc. or its affiliates. -## SPDX-License-Identifier: MIT - -# TODO: run this script from CK root or build directory -EXE=$1 -VALID=0 - -for causal in 0 1 ; do -for prec in "fp16" "bf16" ; do -for hdim in 128 ; do -for perm in 0 ; do - -$EXE -prec=$prec -b=32 -h=16 -s=512 -d=$hdim -causal=$causal -iperm=$perm -operm=$perm -v=$VALID -$EXE -prec=$prec -b=16 -h=16 -s=1024 -d=$hdim -causal=$causal -iperm=$perm -operm=$perm -v=$VALID -$EXE -prec=$prec -b=8 -h=16 -s=2048 -d=$hdim -causal=$causal -iperm=$perm -operm=$perm -v=$VALID -$EXE -prec=$prec -b=4 -h=16 -s=4096 -d=$hdim -causal=$causal -iperm=$perm -operm=$perm -v=$VALID -$EXE -prec=$prec -b=2 -h=16 -s=8192 -d=$hdim -causal=$causal -iperm=$perm -operm=$perm -v=$VALID -$EXE -prec=$prec -b=1 -h=16 -s=16384 -d=$hdim -causal=$causal -iperm=$perm -operm=$perm -v=$VALID - -$EXE -prec=$prec -b=1 -h=64 -s=16384 -d=$hdim -causal=$causal -iperm=$perm -operm=$perm -v=$VALID -$EXE -prec=$prec -b=1 -h=16 -h_k=1 -s=65536 -d=$hdim -causal=$causal -iperm=$perm -operm=$perm -v=$VALID -$EXE -prec=$prec -b=1 -h=40 -s=37200 -d=$hdim -causal=$causal -iperm=$perm -operm=$perm -v=$VALID - -done -done -done -done - -# Padding benchmark comparisons for v3 (batch mode only) -# ==== V3 Padding Benchmarks: batch mode (baseline vs low/med/high pad) ==== -prec="fp16" -base_v3_args="-prec=$prec -b=4 -h=16 -d=128 -s=1024 -mask=0 -iperm=0 -operm=0 -v=$VALID" - -# baseline (no pad) -$EXE $base_v3_args - -# low pad (≈90–95% effective) -$EXE $base_v3_args -q_eff_lens=1024,960,992,896 -kv_eff_lens=1024,960,992,896 - -# medium pad (≈60–75% effective) -$EXE $base_v3_args -q_eff_lens=896,768,512,640 -kv_eff_lens=896,768,512,640 - -# high pad (≈30–40% effective) -$EXE $base_v3_args -q_eff_lens=512,384,256,320 -kv_eff_lens=512,384,256,320 diff --git a/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax.cpp b/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax.cpp index 6f5abf13e8..a4c6bca694 100644 --- a/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax.cpp +++ b/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax.cpp @@ -1,38 +1,24 @@ -// MIT License -// -// Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include + +#include + +#include #include #include #include -#include #include #include -#include #include "ck_tile/core.hpp" #include "ck_tile/ops/reduce.hpp" -#include "ck_tile/utility/json_dump.hpp" #include "topk_softmax_api.hpp" +#include "ck_tile/utility/json_dump.hpp" + #if 0 template void dump_host_tensor_2d(const ck_tile::HostTensor& x) diff --git a/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax_api.cpp b/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax_api.cpp index a3c1cc0b22..cf6154f5b7 100644 --- a/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax_api.cpp +++ b/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax_api.cpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #include "topk_softmax_api.hpp" diff --git a/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax_api.hpp b/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax_api.hpp index 71eb58a587..afc38ef940 100644 --- a/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax_api.hpp +++ b/Libraries/ComposableKernel/attention/topk_softmax/topk_softmax_api.hpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once #include "ck_tile/core.hpp" diff --git a/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_api.cpp b/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_api.cpp index f5362d0f3a..8de52e5f58 100644 --- a/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_api.cpp +++ b/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_api.cpp @@ -1,24 +1,6 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + #include "batched_transpose_example.hpp" namespace { @@ -111,18 +93,18 @@ float batched_transpose_dispatch(batched_transpose_kargs &a, // Param Comb: type_size, block_x & y, WarpNum_x & y #define FOREACH_TRANSPOSE_PARAM(F) \ - F(fp8, ck_tile::fp8_t, 64, 64, 1, 1, true, true, 0) \ - F(fp8, ck_tile::fp8_t, 64, 64, 1, 1, false, false, 0) \ - F(fp16, ck_tile::fp16_t, 64, 64, 1, 1, true, true, 0) \ - F(fp16, ck_tile::fp16_t, 64, 64, 1, 1, false, false, 0) \ - F(bf16, ck_tile::bf16_t, 64, 64, 1, 1, true, true, 0) \ - F(bf16, ck_tile::bf16_t, 64, 64, 1, 1, false, false, 0) \ - F(fp8, ck_tile::fp8_t, 64, 64, 1, 1, true, true, 1) \ - F(fp8, ck_tile::fp8_t, 64, 64, 1, 1, false, false, 1) \ - F(fp16, ck_tile::fp16_t, 64, 64, 1, 1, true, true, 1) \ - F(fp16, ck_tile::fp16_t, 64, 64, 1, 1, false, false, 1) \ - F(bf16, ck_tile::bf16_t, 64, 64, 1, 1, true, true, 1) \ - F(bf16, ck_tile::bf16_t, 64, 64, 1, 1, false, false, 1) + F(fp8, ck_tile::fp8_t, 32, 32, 1, 1, true, true, 0) \ + F(fp8, ck_tile::fp8_t, 32, 32, 1, 1, false, false, 0) \ + F(fp16, ck_tile::fp16_t, 32, 32, 1, 1, true, true, 0) \ + F(fp16, ck_tile::fp16_t, 32, 32, 1, 1, false, false, 0) \ + F(bf16, ck_tile::bf16_t, 32, 32, 1, 1, true, true, 0) \ + F(bf16, ck_tile::bf16_t, 32, 32, 1, 1, false, false, 0) \ + F(fp8, ck_tile::fp8_t, 32, 32, 1, 1, true, true, 1) \ + F(fp8, ck_tile::fp8_t, 32, 32, 1, 1, false, false, 1) \ + F(fp16, ck_tile::fp16_t, 32, 32, 1, 1, true, true, 1) \ + F(fp16, ck_tile::fp16_t, 32, 32, 1, 1, false, false, 1) \ + F(bf16, ck_tile::bf16_t, 32, 32, 1, 1, true, true, 1) \ + F(bf16, ck_tile::bf16_t, 32, 32, 1, 1, false, false, 1) // Macro that defines one static function per line #define GEN_TRANSPOSE_FN(SHORT_NAME, REAL_TYPE, BX, BY, WX, WY, PADM, PADN, \ @@ -141,42 +123,42 @@ float batched_transpose(batched_transpose_trait t, batched_transpose_kargs a, ck_tile::stream_config s) { if (t.pipeline == "0") { if (t.type == "fp8") { - if (a.height % 64 == 0 && a.width % 64 == 0) { - return transpose_fn_fp8_64_64_1_1_false_false_v0(a, s); + if (a.height % 32 == 0 && a.width % 32 == 0) { + return transpose_fn_fp8_32_32_1_1_false_false_v0(a, s); } else { - return transpose_fn_fp8_64_64_1_1_true_true_v0(a, s); + return transpose_fn_fp8_32_32_1_1_true_true_v0(a, s); } } else if (t.type == "fp16") { - if (a.height % 64 == 0 && a.width % 64 == 0) { - return transpose_fn_fp16_64_64_1_1_false_false_v0(a, s); + if (a.height % 32 == 0 && a.width % 32 == 0) { + return transpose_fn_fp16_32_32_1_1_false_false_v0(a, s); } else { - return transpose_fn_fp16_64_64_1_1_true_true_v0(a, s); + return transpose_fn_fp16_32_32_1_1_true_true_v0(a, s); } } else if (t.type == "bf16") { - if (a.height % 64 == 0 && a.width % 64 == 0) { - return transpose_fn_bf16_64_64_1_1_false_false_v0(a, s); + if (a.height % 32 == 0 && a.width % 32 == 0) { + return transpose_fn_bf16_32_32_1_1_false_false_v0(a, s); } else { - return transpose_fn_bf16_64_64_1_1_true_true_v0(a, s); + return transpose_fn_bf16_32_32_1_1_true_true_v0(a, s); } } } else if (t.pipeline == "1") { if (t.type == "fp8") { - if (a.height % 64 == 0 && a.width % 64 == 0) { - return transpose_fn_fp8_64_64_1_1_false_false_v1(a, s); + if (a.height % 32 == 0 && a.width % 32 == 0) { + return transpose_fn_fp8_32_32_1_1_false_false_v1(a, s); } else { - return transpose_fn_fp8_64_64_1_1_true_true_v1(a, s); + return transpose_fn_fp8_32_32_1_1_true_true_v1(a, s); } } else if (t.type == "fp16") { - if (a.height % 64 == 0 && a.width % 64 == 0) { - return transpose_fn_fp16_64_64_1_1_false_false_v1(a, s); + if (a.height % 32 == 0 && a.width % 32 == 0) { + return transpose_fn_fp16_32_32_1_1_false_false_v1(a, s); } else { - return transpose_fn_fp16_64_64_1_1_true_true_v1(a, s); + return transpose_fn_fp16_32_32_1_1_true_true_v1(a, s); } } else if (t.type == "bf16") { - if (a.height % 64 == 0 && a.width % 64 == 0) { - return transpose_fn_bf16_64_64_1_1_false_false_v1(a, s); + if (a.height % 32 == 0 && a.width % 32 == 0) { + return transpose_fn_bf16_32_32_1_1_false_false_v1(a, s); } else { - return transpose_fn_bf16_64_64_1_1_true_true_v1(a, s); + return transpose_fn_bf16_32_32_1_1_true_true_v1(a, s); } } } diff --git a/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_example.cpp b/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_example.cpp index 4a7f46377a..7c2e43b26e 100644 --- a/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_example.cpp +++ b/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_example.cpp @@ -1,32 +1,17 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include + +#include + +#include #include #include #include -#include #include #include -#include #include "batched_transpose_example.hpp" @@ -196,7 +181,7 @@ template bool run_batched_transpose(ck_tile::ArgParser args) { return a_; }(); - ck_tile::stream_config sc{nullptr, true, n_warmup, n_repeat}; + ck_tile::stream_config sc{nullptr, true, 0, n_warmup, n_repeat}; auto ms = batched_transpose(trait, karg, sc); diff --git a/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_example.hpp b/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_example.hpp index 5108e911aa..cf6998c4e8 100644 --- a/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_example.hpp +++ b/Libraries/ComposableKernel/basic/batched_transpose/batched_transpose_example.hpp @@ -1,32 +1,16 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + #include "ck_tile/core.hpp" #include "ck_tile/host.hpp" -#include "ck_tile/ops/batched_transpose.hpp" #include "ck_tile/ops/reduce.hpp" -#include +#include "ck_tile/ops/batched_transpose.hpp" + #include +#include + #pragma once struct batched_transpose_trait { diff --git a/Libraries/ComposableKernel/basic/elementwise/elementwise_common.hpp b/Libraries/ComposableKernel/basic/elementwise/elementwise_common.hpp index 5a059cd804..b22d5818ab 100644 --- a/Libraries/ComposableKernel/basic/elementwise/elementwise_common.hpp +++ b/Libraries/ComposableKernel/basic/elementwise/elementwise_common.hpp @@ -1,28 +1,10 @@ -// MIT License -// -// Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT -#include "ck_tile/core/arch/arch.hpp" #include +#include "ck_tile/core/arch/arch.hpp" + auto string_to_datatype(const std::string &datatype) { using PrecVariant = std::variant; diff --git a/Libraries/ComposableKernel/basic/elementwise/elementwise_example.cpp b/Libraries/ComposableKernel/basic/elementwise/elementwise_example.cpp index 02a519491a..464b820bc2 100644 --- a/Libraries/ComposableKernel/basic/elementwise/elementwise_example.cpp +++ b/Libraries/ComposableKernel/basic/elementwise/elementwise_example.cpp @@ -1,28 +1,10 @@ -// MIT License -// -// Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #include "ck_tile/host.hpp" -#include "ck_tile/host/reference/reference_elementwise.hpp" #include "ck_tile/ops/elementwise.hpp" + +#include "ck_tile/host/reference/reference_elementwise.hpp" #include "ck_tile/utility/json_dump.hpp" #include "elementwise_common.hpp" @@ -247,7 +229,7 @@ int main(int argc, char *argv[]) { [&](auto &&x_dt, auto &&y_dt) -> int { using XDataType = std::decay_t; using YDataType = std::decay_t; - return run(arg_parser); + return !run(arg_parser); }, x_prec_variant, y_prec_variant); } catch (const std::exception &e) { diff --git a/Libraries/ComposableKernel/basic/elementwise/elementwise_example_add_4d.cpp b/Libraries/ComposableKernel/basic/elementwise/elementwise_example_add_4d.cpp index 2f121c8d3e..6423aa06ec 100644 --- a/Libraries/ComposableKernel/basic/elementwise/elementwise_example_add_4d.cpp +++ b/Libraries/ComposableKernel/basic/elementwise/elementwise_example_add_4d.cpp @@ -1,28 +1,10 @@ -// MIT License -// -// Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #include "ck_tile/host.hpp" -#include "ck_tile/host/reference/reference_elementwise.hpp" #include "ck_tile/ops/elementwise.hpp" + +#include "ck_tile/host/reference/reference_elementwise.hpp" #include "ck_tile/utility/json_dump.hpp" #include "elementwise_common.hpp" @@ -185,7 +167,7 @@ int main(int argc, char *argv[]) { [&](auto &&x_dt, auto &&y_dt) -> int { using XDataType = std::decay_t; using YDataType = std::decay_t; - return run(arg_parser); + return !run(arg_parser); }, x_prec_variant, y_prec_variant); } catch (const std::exception &e) { diff --git a/Libraries/ComposableKernel/basic/elementwise/elementwise_example_transpose.cpp b/Libraries/ComposableKernel/basic/elementwise/elementwise_example_transpose.cpp index 6814f09d96..b5fbaafdad 100644 --- a/Libraries/ComposableKernel/basic/elementwise/elementwise_example_transpose.cpp +++ b/Libraries/ComposableKernel/basic/elementwise/elementwise_example_transpose.cpp @@ -1,28 +1,10 @@ -// MIT License -// -// Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #include "ck_tile/host.hpp" -#include "ck_tile/host/reference/reference_transpose.hpp" #include "ck_tile/ops/elementwise.hpp" + +#include "ck_tile/host/reference/reference_transpose.hpp" #include "ck_tile/utility/json_dump.hpp" #include "elementwise_common.hpp" @@ -182,7 +164,7 @@ int main(int argc, char *argv[]) { return std::visit( [&](auto &&dt) -> int { using DataType = std::decay_t; - return run(arg_parser); + return !run(arg_parser); }, prec_variant); } catch (const std::exception &e) { diff --git a/Libraries/ComposableKernel/basic/elementwise/elementwise_example_unary.cpp b/Libraries/ComposableKernel/basic/elementwise/elementwise_example_unary.cpp index cecefb5193..21eaea4b08 100644 --- a/Libraries/ComposableKernel/basic/elementwise/elementwise_example_unary.cpp +++ b/Libraries/ComposableKernel/basic/elementwise/elementwise_example_unary.cpp @@ -1,38 +1,21 @@ -// MIT License -// -// Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #include "ck_tile/host.hpp" -#include "ck_tile/host/reference/reference_elementwise.hpp" #include "ck_tile/ops/elementwise.hpp" + +#include "ck_tile/host/reference/reference_elementwise.hpp" #include "ck_tile/utility/json_dump.hpp" #include "elementwise_common.hpp" auto create_args(int argc, char *argv[]) { ck_tile::ArgParser arg_parser; - arg_parser.insert("m", "1024", "m dimension") + arg_parser + .insert("op", "square", "unary operations: square|tanh|fastgelu|convert") + .insert("m", "1024", "m dimension") .insert("n", "1024", "n dimension") .insert("stride", "-1", "stride per row, if -1 then equal to n") .insert("v", "1", "cpu validation or not") - .insert("op", "1", "unary operation, 1: square, 2: convert") .insert("x_prec", "fp16", "input precision") .insert("y_prec", "fp16", "output precision") .insert("warmup", "10", "cold iter") @@ -172,16 +155,9 @@ bool filter_then_run(const ck_tile::ArgParser &arg_parser) { }; bool pass = true; - if constexpr (std::is_same_v && - (std::is_same_v || + if constexpr ((std::is_same_v || std::is_same_v)) { throw_unsupported(); - } else if constexpr (std::is_same_v && - (std::is_same_v || - std::is_same_v)) { - throw_unsupported(); } else { pass = run(arg_parser); } @@ -190,13 +166,18 @@ bool filter_then_run(const ck_tile::ArgParser &arg_parser) { } auto string_to_op(const std::string &op) { - using OpVariant = std::variant; + using OpVariant = std::variant< + ck_tile::element_wise::UnarySquare, ck_tile::element_wise::UnaryConvert, + ck_tile::element_wise::TanH, ck_tile::element_wise::FastGelu>; - if (op == "1") + if (op == "square") return OpVariant{ck_tile::element_wise::UnarySquare{}}; - else if (op == "2") + else if (op == "convert") return OpVariant{ck_tile::element_wise::UnaryConvert{}}; + else if (op == "tanh") + return OpVariant{ck_tile::element_wise::TanH{}}; + else if (op == "fastgelu") + return OpVariant{ck_tile::element_wise::FastGelu{}}; else { throw std::runtime_error("Unsupported unary operation: " + op); } @@ -220,7 +201,7 @@ int main(int argc, char *argv[]) { using XElementwiseOperation = std::decay_t; using XDataType = std::decay_t; using YDataType = std::decay_t; - return filter_then_run( + return !filter_then_run( arg_parser); }, op_variant, x_prec_variant, y_prec_variant); diff --git a/Libraries/ComposableKernel/basic/img2col/image_to_column.cpp b/Libraries/ComposableKernel/basic/img2col/image_to_column.cpp index 2beb832c8f..76edfccd33 100644 --- a/Libraries/ComposableKernel/basic/img2col/image_to_column.cpp +++ b/Libraries/ComposableKernel/basic/img2col/image_to_column.cpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #include #include diff --git a/Libraries/ComposableKernel/basic/img2col/image_to_column.hpp b/Libraries/ComposableKernel/basic/img2col/image_to_column.hpp index c6883fe35a..6b6488aca0 100644 --- a/Libraries/ComposableKernel/basic/img2col/image_to_column.hpp +++ b/Libraries/ComposableKernel/basic/img2col/image_to_column.hpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once diff --git a/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle.cpp b/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle.cpp index 9f04047511..12de2831d2 100644 --- a/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle.cpp +++ b/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle.cpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #include "matrix_core_swizzle.hpp" #include "matrix_core_swizzle_kernel.hpp" diff --git a/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle.hpp b/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle.hpp index e38c5ed15f..5aeeed4da8 100644 --- a/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle.hpp +++ b/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle.hpp @@ -1,27 +1,7 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once - #include "matrix_core_swizzle_kernel.hpp" #include diff --git a/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle_kernel.hpp b/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle_kernel.hpp index b4803fd3fd..7f758ac1db 100644 --- a/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle_kernel.hpp +++ b/Libraries/ComposableKernel/basic/permute/alternative_impl/matrix_core_swizzle_kernel.hpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once diff --git a/Libraries/ComposableKernel/basic/permute/permute.cpp b/Libraries/ComposableKernel/basic/permute/permute.cpp index 028a6f0d84..86ab9dc390 100644 --- a/Libraries/ComposableKernel/basic/permute/permute.cpp +++ b/Libraries/ComposableKernel/basic/permute/permute.cpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #include "permute.hpp" #include "ck_tile/host.hpp" diff --git a/Libraries/ComposableKernel/basic/permute/permute.hpp b/Libraries/ComposableKernel/basic/permute/permute.hpp index d57a301eb5..cbeb15c2b2 100644 --- a/Libraries/ComposableKernel/basic/permute/permute.hpp +++ b/Libraries/ComposableKernel/basic/permute/permute.hpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once diff --git a/Libraries/ComposableKernel/basic/reduce/README.md b/Libraries/ComposableKernel/basic/reduce/README.md index f2656be2ea..bd547ea124 100644 --- a/Libraries/ComposableKernel/basic/reduce/README.md +++ b/Libraries/ComposableKernel/basic/reduce/README.md @@ -5,7 +5,8 @@ This example demonstrates parallel reduction (sum, max, etc.) using the CK Tile programming model, a core operation for normalization, statistics, and aggregation in deep learning. -Given a tensor $\mathbf{X}$ and a reduction axis, the performed computation is $\mathbf{Y} = \sum_i \mathbf{X}_i$. +Given a tensor $\mathbf{X}$ in NHWC layout, the example reduces along the $N$ axis, so the performed computation is +$\mathbf{Y}_{hwc} = \sum_n \mathbf{X}_{nhwc}$. ### Supported architectures @@ -14,12 +15,12 @@ This example works with ### Application flow -1. Command line arguments are parsed to configure matrix dimensions and execution parameters. +1. Command line arguments are parsed to configure tensor dimensions and execution parameters. 2. $\mathbf{X}$ is created on the host and initialized with random floating-point values in the range $[-5, 5]$. 3. Buffers for $\mathbf{X}$ and $\mathbf{Y}$ are created on the device. 4. $\mathbf{X}$ is copied from the host to the device. -5. A custom reduction kernel is instantiated and launched on the device. -6. If validation is enabled the results are compared agains CK Tile's built-in `reference_reduce` function. +5. CK Tile's `ReduceKernel` is instantiated from a `Reduce2dShape` and a `Reduce2dProblem`, and launched on the device. +6. If validation is enabled the results are compared against CK Tile's built-in `reference_reduce` function. 7. All buffers are freed automatically. ## Key APIs and concepts @@ -28,8 +29,6 @@ This example works with The CK Tile framework is built around four key architectural components: -The CK Tile framework is built around four key architectural components: - * The **shape** defines the hierarchical tile structure and memory layout. * The **problem** combines data types with the shape configuration. * The **pipeline** schedules the sequence of operations for a kernel. @@ -51,39 +50,24 @@ shared memory and register accumulation for efficiency. * `ck_tile::ArgParser` * `ck_tile::bf16_t` +* `ck_tile::DataTypeTraits` * `ck_tile::DeviceMem` * `ck_tile::FillUniformDistribution` * `ck_tile::half_t` * `ck_tile::HostTensor` * `ck_tile::index_t` * `ck_tile::number` +* `ck_tile::Reduce2dProblem` +* `ck_tile::Reduce2dShape` +* `ck_tile::ReduceKernel` * `ck_tile::ReduceOp::Add` * `ck_tile::sequence` * `ck_tile::stream_config` #### Functions -* `ck_tile::block_tile_reduce` -* `ck_tile::block_tile_reduce_sync` -* `ck_tile::cast_tile` * `ck_tile::check_err` -* `ck_tile::get_block_id` -* `ck_tile::integer_divide_ceil` * `ck_tile::launch_kernel` -* `ck_tile::load_tile` * `ck_tile::make_kernel` -* `ck_tile::make_naive_tensor_view` -* `ck_tile::make_naive_tensor_view_packed` -* `ck_tile::make_tile_window` * `ck_tile::make_tuple` -* `ck_tile::move_tile_window` -* `ck_tile::reduce_on_sequence` * `ck_tile::reference_reduce` -* `ck_tile::set_tile` -* `ck_tile::store_tile` - -### HIP Runtime - -#### Device symbols - -* `__builtin_amdgcn_readfirstlane` diff --git a/Libraries/ComposableKernel/basic/reduce/reduce.cpp b/Libraries/ComposableKernel/basic/reduce/reduce.cpp index f172019b04..880389c82e 100644 --- a/Libraries/ComposableKernel/basic/reduce/reduce.cpp +++ b/Libraries/ComposableKernel/basic/reduce/reduce.cpp @@ -1,50 +1,22 @@ -// MIT License -// -// Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT -#include "ck_tile/ops/reduce.hpp" #include "ck_tile/host.hpp" + +#include "ck_tile/ops/reduce.hpp" #include "ck_tile/utility/json_dump.hpp" #include -template struct DataTypeTraits; - -template <> struct DataTypeTraits { - static constexpr const char *name = "fp16"; -}; - -template <> struct DataTypeTraits { - static constexpr const char *name = "bf16"; -}; - auto create_args(int argc, char *argv[]) { ck_tile::ArgParser arg_parser; - arg_parser.insert("n", "32", "n dimension") - .insert("h", "7", "h dimension") - .insert("w", "7", "w dimension") - .insert("c", "512", "c dimension") + arg_parser.insert("n", "16", "n dimension") + .insert("h", "64", "h dimension") + .insert("w", "32", "w dimension") + .insert("c", "960", "c dimension") .insert("v", "1", "cpu validation or not") .insert("prec", "fp16", "precision") - .insert("warmup", "5", "cold iter") - .insert("repeat", "20", "hot iter") + .insert("warmup", "20", "cold iter") + .insert("repeat", "100", "hot iter") .insert("json", "0", "0: No Json, 1: Dump Results in Json format") .insert("jsonfile", "reduce.json", "json file name to dump results"); @@ -74,13 +46,13 @@ template bool run(const ck_tile::ArgParser &arg_parser) { // Define reduction specification: constexpr auto kept_dim = - ck_tile::sequence<0, 3>{}; // Which dimension to keep + ck_tile::sequence<1, 2, 3>{}; // Which dimension to keep constexpr auto reduce_dims = - ck_tile::sequence<1, 2>{}; // Which dimensions to reduce + ck_tile::sequence<0>{}; // Which dimensions to reduce ck_tile::HostTensor x_host(problem_shape, strides); - ck_tile::HostTensor y_host_ref({N, C}, {C, 1}); - ck_tile::HostTensor y_host_dev({N, C}, {C, 1}); + ck_tile::HostTensor y_host_ref({H, W, C}, {W * C, C, 1}); + ck_tile::HostTensor y_host_dev({H, W, C}, {W * C, C, 1}); ck_tile::FillUniformDistribution{-5.f, 5.f}(x_host); @@ -90,29 +62,32 @@ template bool run(const ck_tile::ArgParser &arg_parser) { x_buf.ToDevice(x_host.data()); using ReduceOp = ck_tile::ReduceOp::Add; - using BlockWarps = ck_tile::sequence<4, 1>; - using BlockTile = ck_tile::sequence<128, 128>; - using WarpTile = ck_tile::sequence<32, 128>; - using Vector = ck_tile::sequence<8, 8>; + using BlockWarps = ck_tile::sequence<1, 1>; + using BlockTile = ck_tile::sequence<256, 1>; + using WarpTile = ck_tile::sequence<256, 1>; + using ThreadTile = ck_tile::sequence<1, 1>; // cross warp-reduce // using BlockWarps = ck_tile::sequence<2, 2>; // using BlockTile = ck_tile::sequence<2, 1024>; // using WarpTile = ck_tile::sequence<1, 512>; - // using Vector = ck_tile::sequence<1, 8>; + // using ThreadTile = ck_tile::sequence<1, 8>; constexpr ck_tile::index_t kBlockPerCu = 1; - ck_tile::index_t kept_dim_len_prod = N * C; + ck_tile::index_t kept_dim_len_prod = H * W * C; ck_tile::index_t kGridSize = (kept_dim_len_prod + BlockTile::at(ck_tile::number<0>{}) - 1) / BlockTile::at(ck_tile::number<0>{}); std::cout << "grid size " << kGridSize << std::endl; - using Shape = ck_tile::Reduce2dShape; - using Porblem = ck_tile::Reduce2dProblem; + using Shape = + ck_tile::Reduce2dShape; + using Porblem = + ck_tile::Reduce2dProblem; - using Kernel = ck_tile::Reduce; + using Kernel = ck_tile::ReduceKernel; const ck_tile::index_t kBlockSize = Kernel::BlockSize(); // Create input tensor shape and strides auto input_shape = ck_tile::make_tuple(problem_shape[0], problem_shape[1], @@ -120,23 +95,16 @@ template bool run(const ck_tile::ArgParser &arg_parser) { auto input_strides = ck_tile::make_tuple(strides[0], strides[1], strides[2], strides[3]); - if (!Kernel::IsSupportedArgument( - C, input_strides)) // output tensor's continuous dimension and input - // strides - { - throw std::runtime_error("Wrong! Arguments not supported!\n"); - } - float ave_time = launch_kernel(ck_tile::stream_config{nullptr, true, 0, warmup, repeat}, ck_tile::make_kernel( Kernel{}, kGridSize, kBlockSize, 0, static_cast(x_buf.GetDeviceBuffer()), static_cast(y_buf.GetDeviceBuffer()), - input_shape, input_strides, kept_dim, reduce_dims)); + input_shape, input_strides)); std::size_t num_btype = - sizeof(XDataType) * N * C * H * W + sizeof(YDataType) * N * C; + sizeof(XDataType) * N * H * W * C + sizeof(YDataType) * H * W * C; float gb_per_sec = num_btype / 1.E6 / ave_time; @@ -156,7 +124,7 @@ template bool run(const ck_tile::ArgParser &arg_parser) { } if (arg_parser.get_int("json") == 1) { - dump_reduce_json_results( + dump_reduce_json_results( arg_parser.get_str("jsonfile"), N, C, H, W, pass, ave_time, 0, gb_per_sec); } @@ -173,9 +141,7 @@ int main(int argc, char *argv[]) { if (data_type == "fp16") { return run(arg_parser) ? 0 : -2; + } else if (data_type == "bf16") { + return run(arg_parser) ? 0 : -2; } - // else if(data_type == "bf16") - // { - // return run(arg_parser) ? 0 : -2; - // } } diff --git a/Libraries/ComposableKernel/basic/reduce/reduce.hpp b/Libraries/ComposableKernel/basic/reduce/reduce.hpp deleted file mode 100644 index 27e5ca1ce3..0000000000 --- a/Libraries/ComposableKernel/basic/reduce/reduce.hpp +++ /dev/null @@ -1,181 +0,0 @@ -// MIT License -// -// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -#pragma once - -#include -#include -#include -#include - -namespace ck_tile { - -template - typename BlockTile, // block size, seq - typename WarpTile, // warp size, seq - typename Vector> // contiguous pixels(vector size) along seq -struct Reduce2dShape { - static constexpr index_t Block_M = BlockTile::at(number<0>{}); - static constexpr index_t Block_N = BlockTile::at(number<1>{}); - - static constexpr index_t Warp_M = WarpTile::at(number<0>{}); - static constexpr index_t Warp_N = WarpTile::at(number<1>{}); - - static constexpr index_t Vector_M = Vector::at(number<0>{}); - static constexpr index_t Vector_N = Vector::at(number<1>{}); - - static constexpr index_t WarpPerBlock_M = BlockWarps::at(number<0>{}); - static constexpr index_t WarpPerBlock_N = BlockWarps::at(number<1>{}); - - static constexpr index_t ThreadPerWarp_M = Warp_M / Vector_M; - static constexpr index_t ThreadPerWarp_N = Warp_N / Vector_N; - - static constexpr index_t Repeat_M = Block_M / (WarpPerBlock_M * Warp_M); - static constexpr index_t Repeat_N = Block_N / (WarpPerBlock_N * Warp_N); - - static constexpr index_t BlockSize = - ck_tile::get_warp_size() * - reduce_on_sequence(BlockWarps{}, multiplies{}, number<1>{}); -}; - -template -struct Reduce2dProblem { - using XDataType = remove_cvref_t; - using ComputeDataType = remove_cvref_t; - using YDataType = remove_cvref_t; - using BlockShape = remove_cvref_t; - using ReduceOp = ReduceOp_; - - static constexpr bool kNeedCrossLaneSync = BlockShape::ThreadPerWarp_N > 1; - static constexpr bool kNeedCrossWarpSync = BlockShape::WarpPerBlock_N > 1; -}; - -template -struct Reduce { - using Problem = ck_tile::remove_cvref_t; - using Policy = ck_tile::remove_cvref_t; - - using XDataType = ck_tile::remove_cvref_t; - using ComputeDataType = - ck_tile::remove_cvref_t; - using YDataType = ck_tile::remove_cvref_t; - -#if 0 - CK_TILE_DEVICE void operator()(const XDataType* p_x, YDataType* p_y, index_t M, index_t N) - const - { - using S = typename Problem::BlockShape; - - const auto x_m_n = make_naive_tensor_view( - p_x, make_tuple(M, N), make_tuple(N, 1), number{}, number<1>{}); - - const auto y_m = make_naive_tensor_view_packed( - p_y, make_tuple(M), number<1>{}); - - const auto iM = get_block_id() * S::Block_M; - - auto x_window = make_tile_window(x_m_n, - make_tuple(number{}, number{}), - {iM, 0}, - Policy::template MakeXBlockTileDistribution()); - - auto y_window = make_tile_window(y_m, make_tuple(number{}), {iM}); - - const auto f_reduce = [](const auto& v0, const auto& v1) { return v0 + v1; }; - - const XDataType reduce_init_value = 0; - - constexpr auto reduce_dims = sequence<1>{}; - - auto y_compute = decltype(block_tile_reduce( - load_tile(x_window), reduce_dims, f_reduce, reduce_init_value)){}; - - set_tile(y_compute, reduce_init_value); - - index_t num_n_tile_iteration = - __builtin_amdgcn_readfirstlane(integer_divide_ceil(N, S::Block_N)); - - for(int iN = __builtin_amdgcn_readfirstlane(0); iN < num_n_tile_iteration; ++iN) - { - const auto x = load_tile(x_window); - block_tile_reduce(y_compute, x, reduce_dims, f_reduce); - move_tile_window(x_window, {0, S::Block_N}); - } - - block_tile_reduce_sync(y_compute, f_reduce); - - store_tile(y_window, cast_tile(y_compute)); - } -#else - CK_TILE_DEVICE void operator()(const XDataType *p_x, YDataType *p_y, - index_t M, index_t N) const { - using S = typename Problem::BlockShape; - - const auto x_m_n = make_naive_tensor_view( - p_x, make_tuple(M, N), make_tuple(N, 1), number{}, - number<1>{}); - - const auto y_m = make_naive_tensor_view_packed( - p_y, make_tuple(M), number<1>{}); - - const auto iM = get_block_id() * S::Block_M; - - auto x_window = make_tile_window( - x_m_n, make_tuple(number{}, number{}), {iM, 0}, - Policy::template MakeXBlockTileDistribution()); - - auto y_window = - make_tile_window(y_m, make_tuple(number{}), {iM}); - - __shared__ char smem[Policy::template GetSmemSize()]; - - index_t num_n_tile_iteration = - __builtin_amdgcn_readfirstlane(integer_divide_ceil(N, S::Block_N)); - - auto reduce_func = typename Problem::ReduceOp{}; - auto block_reduce2d = Policy::template GetBlockReduce2d(); - auto block_reduce2d_sync = Policy::template GetBlockReduce2dSync(); - auto block_reduce2d_cross_warp_sync = - Policy::template GetBlockReduce2dCrossWarpSync(); - - using XTensorType = decltype(load_tile(x_window)); - auto y_compute = block_reduce2d.template MakeYBlockTile(); - set_tile(y_compute, - reduce_func.template GetIdentityValue()); - - for (int iN = __builtin_amdgcn_readfirstlane(0); iN < num_n_tile_iteration; - ++iN) { - const auto x = load_tile(x_window); - block_reduce2d(x, y_compute, reduce_func); - move_tile_window(x_window, {0, S::Block_N}); - } - - block_reduce2d_sync(y_compute, reduce_func); - block_reduce2d_cross_warp_sync(y_compute, smem, reduce_func); - - store_tile(y_window, cast_tile(y_compute)); - } -#endif -}; - -} // namespace ck_tile diff --git a/Libraries/ComposableKernel/contraction/batched_contraction/CMakeLists.txt b/Libraries/ComposableKernel/contraction/batched_contraction/CMakeLists.txt index 2941d34e4b..a44acb8b94 100644 --- a/Libraries/ComposableKernel/contraction/batched_contraction/CMakeLists.txt +++ b/Libraries/ComposableKernel/contraction/batched_contraction/CMakeLists.txt @@ -1,6 +1,6 @@ # MIT License # -# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal @@ -35,6 +35,12 @@ select_gpu_language() enable_language(${ROCM_EXAMPLES_GPU_LANGUAGE}) select_hip_platform() +include("${CMAKE_CURRENT_LIST_DIR}/../../../../cmake/FilterHIPArchitectures.cmake") +filter_hip_architectures("CK Tile batched contraction example" "gfx908;gfx90a;gfx942;gfx950" SHOULD_SKIP) +if(SHOULD_SKIP) + return() +endif() + if(ROCM_EXAMPLES_GPU_LANGUAGE STREQUAL "CUDA") message(STATUS "Composable Kernel examples are not supported for CUDA.") return() diff --git a/Libraries/ComposableKernel/contraction/batched_contraction/batched_contraction.cpp b/Libraries/ComposableKernel/contraction/batched_contraction/batched_contraction.cpp index 87935f7609..8aebdcb992 100644 --- a/Libraries/ComposableKernel/contraction/batched_contraction/batched_contraction.cpp +++ b/Libraries/ComposableKernel/contraction/batched_contraction/batched_contraction.cpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #include @@ -29,10 +10,12 @@ #include #include "ck_tile/core.hpp" -#include "ck_tile/host.hpp" #include "ck_tile/ops/epilogue.hpp" + #include "ck_tile/ops/gemm.hpp" +#include "ck_tile/host.hpp" + #include "ck_tile/ops/batched_contraction.hpp" #include "contraction_utils.hpp" @@ -46,17 +29,29 @@ template &args, const ck_tile::stream_config &s) { + +#if CK_TILE_USE_WMMA + constexpr ck_tile::index_t M_Tile = 128; + constexpr ck_tile::index_t N_Tile = 128; +#else constexpr ck_tile::index_t M_Tile = 256; constexpr ck_tile::index_t N_Tile = 256; +#endif constexpr ck_tile::index_t K_Tile = 64; constexpr ck_tile::index_t M_Warp = 2; constexpr ck_tile::index_t N_Warp = 2; constexpr ck_tile::index_t K_Warp = 1; +#if CK_TILE_USE_WMMA + constexpr ck_tile::index_t M_Warp_Tile = 16; + constexpr ck_tile::index_t N_Warp_Tile = 16; +#else constexpr ck_tile::index_t M_Warp_Tile = 32; constexpr ck_tile::index_t N_Warp_Tile = 32; - constexpr ck_tile::index_t K_Warp_Tile = 16; +#endif + constexpr ck_tile::index_t K_Warp_Tile = + ck_tile::get_k_warp_tile(); constexpr bool DoubleSmemBuffer = false; @@ -97,77 +92,49 @@ float batched_contraction_impl( ck_tile::GemmPipelineProblem; - using BaseGemmPipeline = UNIVERSAL_GEMM_PIPELINE; + constexpr auto scheduler = GEMM_PIPELINE_SCHEDULER; + + using UniversalGemmProblem = + ck_tile::UniversalGemmPipelineProblem; + + using GemmPipeline = GEMM_PIPELINE; - ck_tile::index_t K_total = 1; - for (ck_tile::index_t i = NumDimG + NumDimM; i < NumDimG + NumDimM + NumDimK; - ++i) { - K_total *= args.A_dims[i]; + using GemmEpilogue = + ck_tile::CShuffleEpilogue>; + + using Kernel = ck_tile::BatchedContractionKernel; + auto kargs = Kernel::MakeKernelArgs(args); + + const dim3 grids = Kernel::GridSize(kargs); + const dim3 blocks = Kernel::GetBlockSize(); + + if (!Kernel::IsSupportedArguments(kargs)) { + throw std::runtime_error( + "Wrong! Arguments not supported! Skipping contraction!\n"); } - const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_total); - const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); - const ck_tile::TailNumber tail_num = - BaseGemmPipeline::GetBlockLoopTailNum(num_loop); - - float ave_time{0}; - - const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { - constexpr bool has_hot_loop_v = has_hot_loop_.value; - constexpr auto tail_number_v = tail_number_.value; - constexpr auto scheduler = GEMM_PIPELINE_SCHEDULER; - constexpr auto memory_operation = - ck_tile::memory_operation_enum::set; // Always set (no atomic_add) - - using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem< - ADataType, BDataType, AccDataType, GemmShape, GemmUniversalTraits, - scheduler, has_hot_loop_v, tail_number_v>; - - using GemmPipeline = GEMM_PIPELINE; - - using GemmEpilogue = - ck_tile::CShuffleEpilogue>; - - using Kernel = - ck_tile::BatchedContractionKernel; - auto kargs = Kernel::MakeKernelArgs(args); - - const dim3 grids = Kernel::GridSize(kargs); - const dim3 blocks = Kernel::GetBlockSize(); - - if (!Kernel::IsSupportedArguments(kargs)) { - throw std::runtime_error( - "Wrong! Arguments not supported! Skipping contraction!\n"); - } - - if (s.log_level_ > 0) { - std::cout << "Launching kernel with args: " << Kernel::GetKernelName() - << '\n' - << "shape: " << GemmShape::GetName() << '\n' - << "problem: " << GemmPipelineProblem::GetName() << '\n' - << "pipeline: " << GemmPipeline::GetName() << '\n' - << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z - << "}" << ", blocks: {" << blocks.x << ", " << blocks.y << ", " - << blocks.z << "}" << std::endl; - } - - auto kernel = - ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs); - - ave_time = ck_tile::launch_kernel(s, kernel); - - return ave_time; - }; - - BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); - - return ave_time; + if (s.log_level_ > 0) { + std::cout << "Launching kernel with args: " << Kernel::GetKernelName() + << '\n' + << "shape: " << GemmShape::GetName() << '\n' + << "problem: " << GemmPipelineProblem::GetName() << '\n' + << "pipeline: " << GemmPipeline::GetName() << '\n' + << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z + << "}" << ", blocks: {" << blocks.x << ", " << blocks.y << ", " + << blocks.z << "}" << std::endl; + } + + auto kernel = + ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs); + + return ck_tile::launch_kernel(s, kernel); } #define HANDLE_CASE(G, M, N, K) \ @@ -194,9 +161,7 @@ float batched_contraction( HANDLE_CASE(2, 1, 1, 1); HANDLE_CASE(2, 2, 2, 1); HANDLE_CASE(1, 2, 1, 1); - HANDLE_CASE(1, 1, 1, 2); HANDLE_CASE(2, 2, 2, 2); - HANDLE_CASE(4, 4, 4, 4); throw std::runtime_error( "Unsupported dimension combination: G=" + std::to_string(num_g_dims) + diff --git a/Libraries/ComposableKernel/contraction/batched_contraction/contraction_utils.hpp b/Libraries/ComposableKernel/contraction/batched_contraction/contraction_utils.hpp index 35d21102c3..2ad6344150 100644 --- a/Libraries/ComposableKernel/contraction/batched_contraction/contraction_utils.hpp +++ b/Libraries/ComposableKernel/contraction/batched_contraction/contraction_utils.hpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once @@ -58,7 +39,80 @@ using AccDataType = ContractionTypes::AccDataType; using EDataType = ContractionTypes::EDataType; using DDataType = ContractionTypes::DDataType; +void print_help(const char *program_name) { + std::cout << "\n"; + std::cout << "Batched Tensor Contraction with element-wise fusion\n"; + std::cout << "E[G,M,N] = element_wise_op(contraction(A[G,M,K], B[G,N,K]), " + "D0, D1, ...)\n"; + std::cout << "(Supports multiple D tensors with configurable element-wise " + "operations)\n\n"; + + std::cout << "Usage: " << program_name << " [OPTIONS]\n\n"; + + std::cout << "Dimension Arguments (comma-separated, no spaces):\n"; + std::cout << " -g_dims= Batch dimensions (default: \"1,2\")\n"; + std::cout + << " -m_dims= M (row) dimensions (default: \"4,256\")\n"; + std::cout + << " -n_dims= N (column) dimensions (default: \"16,128\")\n"; + std::cout << " -k_dims= K (contract) dims (default: \"64\")\n"; + std::cout << " -num_d= Number of D tensors (default: 2, range: " + "0-4)\n\n"; + + std::cout + << "Custom Stride Arguments (for testing non-contiguous tensors):\n"; + std::cout << " -strides_a= A tensor strides (comma-separated, empty = " + "auto)\n"; + std::cout << " -strides_b= B tensor strides (comma-separated, empty = " + "auto)\n"; + std::cout << " -strides_e= E tensor strides (comma-separated, empty = " + "auto)\n"; + std::cout << " -strides_ds= D tensors strides (semicolon-separated, " + "empty = same as E)\n"; + std::cout << " Example: -strides_a=\"32768,128,1\" " + "-strides_ds=\"512,2,1;1024,4,1\"\n\n"; + + std::cout << "Layout Arguments:\n"; + std::cout << " -a_layout= A tensor layout (R=Row-major, " + "C=Column-major, default: \"R\")\n"; + std::cout << " -b_layout= B tensor layout (default: \"C\")\n"; + std::cout << " -e_layout= E tensor layout (default: \"R\")\n\n"; + + std::cout << "Examples:\n"; + std::cout << " Single batch (12 batches of 256x128):\n"; + std::cout + << " " << program_name + << " -g_dims=\"12\" -m_dims=\"256\" -n_dims=\"128\" -k_dims=\"64\"\n\n"; + + std::cout << " 2D batch grid (2x3=6 batches):\n"; + std::cout + << " " << program_name + << " -g_dims=\"2,3\" -m_dims=\"128\" -n_dims=\"128\" -k_dims=\"64\"\n\n"; + + std::cout << " Multi-dimensional (flattened to M=128, N=128, K=128):\n"; + std::cout << " " << program_name + << " -g_dims=\"4\" -m_dims=\"8,16\" -n_dims=\"32,4\" " + "-k_dims=\"16,8\"\n\n"; + + std::cout << "Other Options:\n"; + std::cout << " -v=<0|1> Validation (0=off, 1=on, default: 1)\n"; + std::cout << " -split_k= Split-K value (default: 1)\n"; + std::cout << " -warmup= Warmup iterations (default: 5)\n"; + std::cout << " -repeat= Benchmark iterations (default: 10)\n"; + std::cout << " -log=<0|1> Logging level (default: 1)\n"; + std::cout << " -help Show this help\n\n"; +} + auto create_args(int argc, char *argv[]) { + // Check for --help flag + for (int i = 1; i < argc; ++i) { + std::string arg = argv[i]; + if (arg == "--help" || arg == "-h" || arg == "-help") { + print_help(argv[0]); + std::exit(0); + } + } + ck_tile::ArgParser arg_parser; arg_parser .insert("m_dims", "4,256", @@ -70,12 +124,16 @@ auto create_args(int argc, char *argv[]) { .insert("g_dims", "1,2", "G dimensions separated by comma (e.g., '4,2' for 2D, '2,3,4' " "for 3D)") - .insert("stride_a", "0", - "Custom A tensor leading dimension stride (0 = auto)") - .insert("stride_b", "0", - "Custom B tensor leading dimension stride (0 = auto)") - .insert("stride_e", "0", - "Custom E tensor leading dimension stride (0 = auto)") + .insert("num_d", "2", "Number of D (auxiliary input) tensors") + .insert("strides_a", "", + "A tensor strides (comma-separated, empty = auto/contiguous)") + .insert("strides_b", "", + "B tensor strides (comma-separated, empty = auto/contiguous)") + .insert("strides_e", "", + "E tensor strides (comma-separated, empty = auto/contiguous)") + .insert("strides_ds", "", + "D tensors strides (semicolon-separated for multiple, empty = " + "same as E)") .insert("a_layout", "R", "A tensor data layout - Row by default") .insert("b_layout", "C", "B tensor data layout - Col by default") .insert("e_layout", "R", "E tensor data layout - Row by default") diff --git a/Libraries/ComposableKernel/contraction/batched_contraction/run_batched_contraction_example.inc b/Libraries/ComposableKernel/contraction/batched_contraction/run_batched_contraction_example.inc index e7db122090..e9e9832002 100644 --- a/Libraries/ComposableKernel/contraction/batched_contraction/run_batched_contraction_example.inc +++ b/Libraries/ComposableKernel/contraction/batched_contraction/run_batched_contraction_example.inc @@ -1,30 +1,12 @@ -// MIT License -// -// Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once #include #include #include +#include #include "contraction_utils.hpp" #include "ck_tile/host/reference/reference_batched_contraction.hpp" @@ -64,10 +46,10 @@ float invoke_batched_contraction_kernel( const void* b_full_dims_dev_buf, const std::array& ds_dev_buf, void* e_full_dims_dev_buf, - const std::vector& G_dims, - const std::vector& M_dims, - const std::vector& N_dims, - const std::vector& K_dims, + ck_tile::index_t num_g_dims, + ck_tile::index_t num_m_dims, + ck_tile::index_t num_n_dims, + ck_tile::index_t num_k_dims, const std::vector& A_dims, // [G0,G1,..,M0,M1,..,K0,K1,..] const std::vector& B_dims, // [G0,G1,..,N0,N1,..,K0,K1,..] const std::array, DsDataType::size()>& @@ -98,9 +80,8 @@ float invoke_batched_contraction_kernel( E_strides // E_strides ); - std::cout << "Calling batched_contraction with dimensions: G=" << G_dims.size() - << ", M=" << M_dims.size() << ", N=" << N_dims.size() << ", K=" << K_dims.size() - << std::endl; + std::cout << "Calling batched_contraction with dimensions: G=" << num_g_dims + << ", M=" << num_m_dims << ", N=" << num_n_dims << ", K=" << num_k_dims << std::endl; float ave_time = batched_contraction( args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}, - G_dims.size(), // num_g_dims - M_dims.size(), // num_m_dims - N_dims.size(), // num_n_dims - K_dims.size() // num_k_dims - ); + num_g_dims, + num_m_dims, + num_n_dims, + num_k_dims); return ave_time; } -template +// C++17-compatible helper function to create array of HostTensors +namespace { +template +std::array, NumDTensor> +make_ds_host_tensors_impl(const std::array& descs, + std::index_sequence) +{ + return {ck_tile::HostTensor(descs[Is])...}; +} + +template +std::array, NumDTensor> +make_ds_host_tensors(const std::array& descs) +{ + return make_ds_host_tensors_impl(descs, + std::make_index_sequence{}); +} +} // anonymous namespace + +template int run_batched_contraction_example_with_layouts( int argc, char* argv[], @@ -141,8 +144,6 @@ int run_batched_contraction_example_with_layouts( std::vector N_dims = parse_dimensions(arg_parser.get_str("n_dims")); std::vector K_dims = parse_dimensions(arg_parser.get_str("k_dims")); - constexpr ck_tile::index_t NumDTensor = 2; - ck_tile::index_t G_total = calculate_total_elements(G_dims); ck_tile::index_t M_total = calculate_total_elements(M_dims); ck_tile::index_t N_total = calculate_total_elements(N_dims); @@ -167,13 +168,105 @@ int run_batched_contraction_example_with_layouts( return converted; }; - ck_tile::HostTensorDescriptor a_desc(A_dims); - ck_tile::HostTensorDescriptor b_desc(B_dims); - ck_tile::HostTensorDescriptor e_desc(E_dims); + // Get custom stride arguments + std::string strides_a_str = arg_parser.get_str("strides_a"); + std::string strides_b_str = arg_parser.get_str("strides_b"); + std::string strides_e_str = arg_parser.get_str("strides_e"); + std::string strides_ds_str = arg_parser.get_str("strides_ds"); + + // Create A descriptor with custom or default strides + ck_tile::HostTensorDescriptor a_desc; + if(!strides_a_str.empty()) + { + std::vector custom_a_strides = parse_dimensions(strides_a_str); + if(custom_a_strides.size() != A_dims.size()) + { + throw std::runtime_error("strides_a size must match A_dims size"); + } + std::vector a_strides_size_t(custom_a_strides.begin(), custom_a_strides.end()); + a_desc = ck_tile::HostTensorDescriptor(A_dims, a_strides_size_t); + std::cout << "Using custom strides for A (non-contiguous)" << std::endl; + } + else + { + a_desc = ck_tile::HostTensorDescriptor(A_dims); + } + + // Create B descriptor with custom or default strides + ck_tile::HostTensorDescriptor b_desc; + if(!strides_b_str.empty()) + { + std::vector custom_b_strides = parse_dimensions(strides_b_str); + if(custom_b_strides.size() != B_dims.size()) + { + throw std::runtime_error("strides_b size must match B_dims size"); + } + std::vector b_strides_size_t(custom_b_strides.begin(), custom_b_strides.end()); + b_desc = ck_tile::HostTensorDescriptor(B_dims, b_strides_size_t); + std::cout << "Using custom strides for B (non-contiguous)" << std::endl; + } + else + { + b_desc = ck_tile::HostTensorDescriptor(B_dims); + } + + // Create E descriptor with custom or default strides + ck_tile::HostTensorDescriptor e_desc; + if(!strides_e_str.empty()) + { + std::vector custom_e_strides = parse_dimensions(strides_e_str); + if(custom_e_strides.size() != E_dims.size()) + { + throw std::runtime_error("strides_e size must match E_dims size"); + } + std::vector e_strides_size_t(custom_e_strides.begin(), custom_e_strides.end()); + e_desc = ck_tile::HostTensorDescriptor(E_dims, e_strides_size_t); + std::cout << "Using custom strides for E (non-contiguous)" << std::endl; + } + else + { + e_desc = ck_tile::HostTensorDescriptor(E_dims); + } + // Create D descriptors with custom or default strides (default = same as E) std::array ds_descs; - for(ck_tile::index_t d = 0; d < NumDTensor; ++d) + if(!strides_ds_str.empty()) { - ds_descs[d] = ck_tile::HostTensorDescriptor(Ds_dims[d], e_desc.get_strides()); + // Parse semicolon-separated stride vectors for multiple D tensors + std::vector> all_ds_strides; + std::stringstream ss(strides_ds_str); + std::string d_stride_str; + + while(std::getline(ss, d_stride_str, ';')) + { + all_ds_strides.push_back(parse_dimensions(d_stride_str)); + } + + if(all_ds_strides.size() != NumDTensor) + { + throw std::runtime_error("Number of D stride vectors must match num_d=" + + std::to_string(NumDTensor)); + } + + std::cout << "Using custom strides for D tensors (non-contiguous)" << std::endl; + for(ck_tile::index_t d = 0; d < NumDTensor; ++d) + { + if(all_ds_strides[d].size() != E_dims.size()) + { + throw std::runtime_error("D tensor " + std::to_string(d) + + " stride size must match E_dims size"); + } + std::vector d_strides_size_t(all_ds_strides[d].begin(), + all_ds_strides[d].end()); + ds_descs[d] = ck_tile::HostTensorDescriptor(Ds_dims[d], d_strides_size_t); + } + } + else + { + // Default: use same strides as E + for(ck_tile::index_t d = 0; d < NumDTensor; ++d) + { + ds_descs[d] = ck_tile::HostTensorDescriptor(Ds_dims[d], e_desc.get_strides()); + } } std::vector A_strides = convert_strides(a_desc.get_strides()); @@ -220,11 +313,8 @@ int run_batched_contraction_example_with_layouts( ck_tile::HostTensor<::BDataType> b_full_dims_host(b_desc); ck_tile::HostTensor<::EDataType> e_full_dims_host(e_desc); - std::vector> ds_full_dims_host; - for(int d = 0; d < NumDTensor; ++d) - { - ds_full_dims_host.emplace_back(ck_tile::HostTensor<::DDataType>(ds_descs[d])); - } + // Construct array of HostTensors - C++17 compatible + auto ds_full_dims_host = make_ds_host_tensors<::DDataType, NumDTensor>(ds_descs); ck_tile::FillUniformDistribution<::ADataType>{-5.f, 5.f, std::nullopt}(a_full_dims_host); ck_tile::FillUniformDistribution<::BDataType>{-5.f, 5.f, std::nullopt}(b_full_dims_host); @@ -279,10 +369,10 @@ int run_batched_contraction_example_with_layouts( b_full_dims_dev_buf.GetDeviceBuffer(), ds_ptr_buf, e_full_dims_dev_buf.GetDeviceBuffer(), - G_dims, - M_dims, - N_dims, - K_dims, + G_dims.size(), + M_dims.size(), + N_dims.size(), + K_dims.size(), A_dims, B_dims, Ds_dims, @@ -335,20 +425,25 @@ int run_batched_contraction_example_with_layouts( auto start_time = std::chrono::high_resolution_clock::now(); - calculate_reference_flat_indexing(a_full_dims_host, - b_full_dims_host, - ds_full_dims_host, - e_full_dims_host_ref, - G_total, - M_total, - N_total, - K_total, - CDEElementWise{}); + ck_tile::compute_reference_batched_contraction(a_full_dims_host, + b_full_dims_host, + ds_full_dims_host, + e_full_dims_host_ref, + G_total, + M_total, + N_total, + K_total, + CDEElementWise{}, + G_dims, + M_dims, + N_dims, + K_dims); auto end_time = std::chrono::high_resolution_clock::now(); auto duration = @@ -406,15 +501,45 @@ int run_batched_contraction_example(int argc, char* argv[]) if(!result) return -1; + // Get NumDTensor to dispatch at runtime + const int num_d = arg_parser.get_int("num_d"); + using Row = ck_tile::tensor_layout::gemm::RowMajor; using Col = ck_tile::tensor_layout::gemm::ColumnMajor; std::string a_layout = arg_parser.get_str("a_layout"); std::string b_layout = arg_parser.get_str("b_layout"); + // Runtime dispatch based on num_d value if(a_layout == "R" && b_layout == "C") { - return run_batched_contraction_example_with_layouts(argc, argv, Row{}, Col{}, Row{}, Row{}); + // Dispatch to appropriate template instantiation based on runtime num_d + switch(num_d) + { + case 0: + std::cout << "Running with 0 D tensors" << std::endl; + return run_batched_contraction_example_with_layouts( + argc, argv, Row{}, Col{}, Row{}, Row{}); + case 1: + std::cout << "Running with 1 D tensor" << std::endl; + return run_batched_contraction_example_with_layouts( + argc, argv, Row{}, Col{}, Row{}, Row{}); + case 2: + std::cout << "Running with 2 D tensors" << std::endl; + return run_batched_contraction_example_with_layouts( + argc, argv, Row{}, Col{}, Row{}, Row{}); + case 3: + std::cout << "Running with 3 D tensors" << std::endl; + return run_batched_contraction_example_with_layouts( + argc, argv, Row{}, Col{}, Row{}, Row{}); + case 4: + std::cout << "Running with 4 D tensors" << std::endl; + return run_batched_contraction_example_with_layouts( + argc, argv, Row{}, Col{}, Row{}, Row{}); + default: + throw std::runtime_error("num_d must be between 0 and 4, got: " + + std::to_string(num_d)); + } } else { diff --git a/Libraries/ComposableKernel/convolution/grouped_convolution/CMakeLists.txt b/Libraries/ComposableKernel/convolution/grouped_convolution/CMakeLists.txt index 8d90b5b4f7..85ba09fc76 100644 --- a/Libraries/ComposableKernel/convolution/grouped_convolution/CMakeLists.txt +++ b/Libraries/ComposableKernel/convolution/grouped_convolution/CMakeLists.txt @@ -1,6 +1,6 @@ # MIT License # -# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal @@ -70,17 +70,10 @@ target_compile_options(${example_bwd} PRIVATE -mllvm -enable-noalias-to-md-conve target_link_libraries(${example_bwd} PRIVATE composable_kernel::ck_tile hip::host) # Additional convolution variants -set(example_fwd_large ${example_name}_fwd_large_tensor) set(example_fwd_bias_clamp ${example_name}_fwd_bias_clamp) set(example_bwd_two_stage ${example_name}_bwd_weight_two_stage) set(example_bwd_data ${example_name}_bwd_data) -add_executable(${example_fwd_large} grouped_convolution_forward_large_tensor.cpp) -set_source_files_properties(grouped_convolution_forward_large_tensor.cpp PROPERTIES LANGUAGE ${ROCM_EXAMPLES_GPU_LANGUAGE}) -target_compile_features(${example_fwd_large} PRIVATE cxx_std_17 hip_std_17) -target_compile_options(${example_fwd_large} PRIVATE -mllvm -enable-noalias-to-md-conversion=0) -target_link_libraries(${example_fwd_large} PRIVATE composable_kernel::ck_tile hip::host) - add_executable(${example_fwd_bias_clamp} grouped_convolution_forward_bias_clamp.cpp) set_source_files_properties(grouped_convolution_forward_bias_clamp.cpp PROPERTIES LANGUAGE ${ROCM_EXAMPLES_GPU_LANGUAGE}) target_compile_features(${example_fwd_bias_clamp} PRIVATE cxx_std_17 hip_std_17) diff --git a/Libraries/ComposableKernel/convolution/grouped_convolution/README.md b/Libraries/ComposableKernel/convolution/grouped_convolution/README.md index 6a33312a58..c73217ea4b 100644 --- a/Libraries/ComposableKernel/convolution/grouped_convolution/README.md +++ b/Libraries/ComposableKernel/convolution/grouped_convolution/README.md @@ -6,7 +6,6 @@ This example demonstrates how to perform a grouped convolution operation with CK Tile. Multiple variants are provided: * **Forward pass**: Basic grouped convolution forward operation -* **Forward pass with large tensors**: Optimized for larger tensor dimensions * **Forward pass with bias and clamp**: Includes bias addition and value clamping * **Backward pass (weight)**: Computes gradients with respect to weights @@ -81,7 +80,6 @@ For more information on CK Tile terminology, refer to the * `ck_tile::FillMonotonicSeq` * `ck_tile::FillUniformDistribution` * `ck_tile::GemmPipelineAGmemBGmemCRegV1` -* `ck_tile::GemmPipelineProblem` * `ck_tile::GemmTile1DPartitioner` * `ck_tile::GroupedConvFwdHostArgs` * `ck_tile::GroupedConvolutionForwardKernel` @@ -89,8 +87,6 @@ For more information on CK Tile terminology, refer to the * `ck_tile::half_t` * `ck_tile::HostTensor` * `ck_tile::index_t` -* `ck_tile::integral_constant` -* `ck_tile::memory_operation_enum` * `ck_tile::sequence` * `ck_tile::stream_config` * `ck_tile::tensor_layout::convolution::GKXC` diff --git a/Libraries/ComposableKernel/convolution/grouped_convolution/conv_configs.hpp b/Libraries/ComposableKernel/convolution/grouped_convolution/conv_configs.hpp index 0daa85c1e3..704e936d1e 100644 --- a/Libraries/ComposableKernel/convolution/grouped_convolution/conv_configs.hpp +++ b/Libraries/ComposableKernel/convolution/grouped_convolution/conv_configs.hpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #pragma once @@ -157,7 +138,8 @@ struct ConvConfigComputeV3_WMMA : public ConvConfigBase { static constexpr ck_tile::index_t M_Warp_Tile = 16; static constexpr ck_tile::index_t N_Warp_Tile = 16; - static constexpr ck_tile::index_t K_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = + ck_tile::get_k_warp_tile(); static constexpr bool DoubleSmemBuffer = false; static constexpr ck_tile::GemmPipeline Pipeline = @@ -226,6 +208,26 @@ struct ConvConfigComputeV5 : public ConvConfigBase { static constexpr ck_tile::index_t NumWaveGroups = 2; }; +template +struct ConvConfigComputeV6 : public ConvConfigBase { + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 256; + static constexpr ck_tile::index_t K_Tile = 32; + + static constexpr ck_tile::index_t M_Warp = 2; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 32; + static constexpr ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 16; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::GemmPipeline Pipeline = + ck_tile::GemmPipeline::COMPUTE_V6; + static constexpr ck_tile::index_t NumWaveGroups = 1; +}; + template struct ConvConfigComputeV3_merged_groups : public ConvConfigBase { static constexpr ck_tile::index_t VectorSizeA = 4; @@ -251,6 +253,63 @@ struct ConvConfigComputeV3_merged_groups : public ConvConfigBase { static constexpr ck_tile::index_t NumGroupsToMerge = 2; }; +template +struct ConvConfigComputeAsyncV2 : public ConvConfigBase { + static constexpr ck_tile::index_t M_Tile = 16; + static constexpr ck_tile::index_t N_Tile = 64; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 32; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::GemmPipeline Pipeline = + ck_tile::GemmPipeline::COMPUTE_ASYNC_V2; +}; + +template +struct ConvConfigComputeTDMV1 : public ConvConfigBase { + static constexpr ck_tile::index_t M_Tile = 16; + static constexpr ck_tile::index_t N_Tile = 64; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 32; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::GemmPipeline Pipeline = + ck_tile::GemmPipeline::COMPUTE_TDM_V1; +}; + +template +struct ConvConfigComputeTDMV2 : public ConvConfigBase { + static constexpr ck_tile::index_t M_Tile = 16; + static constexpr ck_tile::index_t N_Tile = 64; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 32; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::GemmPipeline Pipeline = + ck_tile::GemmPipeline::COMPUTE_TDM_V2; +}; + template struct ConvTypeConfig; @@ -271,25 +330,30 @@ struct ConvTypeConfig { using OutDataType = ck_tile::bf16_t; }; -template struct DataTypeTraits; - -template <> struct DataTypeTraits { - static constexpr const char *name = "fp32"; -}; +template struct PipelineTypeTraits; -template <> struct DataTypeTraits { - static constexpr const char *name = "fp16"; +template <> struct PipelineTypeTraits { + template + using GemmPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1< + PipelineProblem, ck_tile::GroupedConvUniversalPipelineAgBgCrPolicy>; + template + using UniversalGemmPipeline = + ck_tile::BaseGemmPipelineAGmemBGmemCRegV1; }; -template <> struct DataTypeTraits { - static constexpr const char *name = "bf16"; +template <> struct PipelineTypeTraits { + template + using GemmPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV2< + PipelineProblem, ck_tile::GroupedConvUniversalPipelineAgBgCrPolicy>; + template + using UniversalGemmPipeline = + ck_tile::BaseGemmPipelineAGmemBGmemCRegV2; }; -template struct PipelineTypeTraits; - template <> struct PipelineTypeTraits { template - using GemmPipeline = ck_tile::GemmPipelineAgBgCrMem; + using GemmPipeline = ck_tile::GemmPipelineAgBgCrMem< + PipelineProblem, ck_tile::GroupedConvUniversalPipelineAgBgCrPolicy>; template using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrMem; @@ -297,7 +361,8 @@ template <> struct PipelineTypeTraits { template <> struct PipelineTypeTraits { template - using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3; + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3< + PipelineProblem, ck_tile::GroupedConvUniversalPipelineAgBgCrPolicy>; template using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; @@ -306,15 +371,17 @@ template <> struct PipelineTypeTraits { template <> struct PipelineTypeTraits { template using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV4; - template - using UniversalGemmPipeline = - ck_tile::BaseGemmPipelineAgBgCrCompV4; }; template <> struct PipelineTypeTraits { template using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV5; +}; + +template <> struct PipelineTypeTraits { + template + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV6; template using UniversalGemmPipeline = - ck_tile::BaseGemmPipelineAgBgCrCompV5; + ck_tile::BaseGemmPipelineAgBgCrCompV6; }; diff --git a/Libraries/ComposableKernel/convolution/grouped_convolution/grouped_convolution_backward_data.cpp b/Libraries/ComposableKernel/convolution/grouped_convolution/grouped_convolution_backward_data.cpp index 2860a2aa04..4f0e087e47 100644 --- a/Libraries/ComposableKernel/convolution/grouped_convolution/grouped_convolution_backward_data.cpp +++ b/Libraries/ComposableKernel/convolution/grouped_convolution/grouped_convolution_backward_data.cpp @@ -1,24 +1,5 @@ -// MIT License -// -// Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT #include @@ -29,8 +10,9 @@ #include #include "ck_tile/host.hpp" -#include "grouped_convolution_backward_data_invoker.hpp" #include "grouped_convolution_utils.hpp" + +#include "grouped_convolution_backward_data_invoker.hpp" #include "run_grouped_convolution_bwd_data_example.inc" template