Skip to content

Update to new QoLA/CK-JIT/AITER #709

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ipanfilo wants to merge 1 commit into
base: dev
Choose a base branch
from
Draft

Update to new QoLA/CK-JIT/AITER #709

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion 3rdparty/QoLA
Open in desktop
Submodule QoLA updated 7 files
+12 −12 patches/aiter/0002-hsa-only-aiter-hip-common.patch
+0 −114 patches/aiter/0003-ck-te-receipt-700.patch
+16 −14 patches/aiter/0005-bwd-ondevice-workspace-prepare.patch
+6 −6 patches/aiter/0007-mha-bwd-v3-api-check-dry-run.patch
+371 −0 patches/aiter/0008-mha-fwd-native-splitkv.patch
+1 −1 qola/cpp_itfs/qola_mha_fwd.cu
+1 −1 qola/cpp_itfs/qola_mha_fwd.h
2 changes: 1 addition & 1 deletion 3rdparty/ck_jit
Open in desktop
Submodule ck_jit updated 2 files
+40 −17 ck_jit_build.py
+100 −0 codegen_jit_hints_22ee91463.patch
2 changes: 1 addition & 1 deletion transformer_engine/common/ck_fused_attn/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,7 +222,7 @@ add_library(ck_fused_attn SHARED ${ck_fused_attn_SOURCES})
set(CK_FUSED_ATTN_COMPILE_OPTIONS)
list(APPEND CK_FUSED_ATTN_COMPILE_OPTIONS
-DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=${CK_FUSED_ATTN_FLOAT_TO_BFLOAT16_DEFAULT}
-DENABLE_CK=1 -DFAV_NATIVE_ON=1)
-DENABLE_CK=1 -DFA_WITH_NATIVE_SPLITKV=1)

# Public QoLA headers ship alongside the .so libs in ${__AITER_MHA_PATH}/../include
# (emitted by qola.cli build, or copied from the QoLA build dir above for the
Expand Down
5 changes: 3 additions & 2 deletions transformer_engine/common/ck_fused_attn/qola_manifest.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
[qola]
aiter_commit = "7a24cd87525834fa9aeaa021e6ee80ab7de028f9" # pinned AITER submodule commit
aiter_commit = "a50a62aaa292427e00190dd7ba58e39a6db4e61e" # pinned AITER submodule commit
namespace = "te"
rocm_versions = ["7.2"]
rocm_versions = ["7.14"]

[build]
architectures = ["gfx950", "gfx942"]
Expand All @@ -12,6 +12,7 @@ mode = "cpp_itfs"
receipt = 700
drop_srcs = ["mha_fwd_split.cu", "mha_fwd_batch_prefill.cu"]
drop_directions = ["fwd_splitkv", "batch_prefill"]
flags_extra_cc = ["'-DFA_WITH_NATIVE_SPLITKV=1'"]

[[modules]]
name = "libmha_bwd"
Expand Down
4 changes: 2 additions & 2 deletions transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@ hipError_t ck_attn_fwd(const CKAttnFwdArgs& args, hipStream_t stream){
}

int ck_attn_fwd_num_splits(const CKAttnFwdArgs& args){
#if FAV_NATIVE_ON
#if FA_WITH_NATIVE_SPLITKV
aiter::mha_fwd_args fmha_args = build_fwd_fmha_args(args);
return QOLA_NS(mha_fwd_calculate_num_splits)(fmha_args);
#else
Expand All @@ -276,7 +276,7 @@ int ck_attn_fwd_num_splits(const CKAttnFwdArgs& args){
}

size_t ck_attn_fwd_workspace_size(const CKAttnFwdArgs& args){
#if FAV_NATIVE_ON
#if FA_WITH_NATIVE_SPLITKV
aiter::mha_fwd_args fmha_args = build_fwd_fmha_args(args);
return QOLA_NS(mha_fwd_workspace_size)(fmha_args);
#else
Expand Down
Loading