diff --git a/aiter/ops/mha.py b/aiter/ops/mha.py index 5c9db14959a..7e0379e2853 100644 --- a/aiter/ops/mha.py +++ b/aiter/ops/mha.py @@ -1293,7 +1293,7 @@ def _flash_attn_forward( swa = (window_size_left > 0) or (window_size_right > 0) def is_fmha_v3_fp8(): - ret = get_gfx() == "gfx942" + ret = get_gfx() in ("gfx942", "gfx950") ret = ret and (hdim_q == 128) ret = ret and (q.dtype == dtypes.fp8) ret = ret and ( @@ -1319,6 +1319,11 @@ def can_impl_fmha_v3_fwd(): ret = ret and (not swa) ret = ret and (q.dtype == dtypes.bf16 or is_fmha_v3_fp8()) ret = ret and (cu_seqlens_q is None and cu_seqlens_kv is None) + # FP8 ASM kernels assemble the GQA-shift from a fixed log2 table + # (1,2,4,8,16); arbitrary divisor ratios route to CK. + if is_fmha_v3_fp8(): + gqa_ratio = nhead_q // nhead_k + ret = ret and ((gqa_ratio & (gqa_ratio - 1)) == 0) return ret q, k, v = [maybe_contiguous(x) for x in (q, k, v)] @@ -2074,7 +2079,7 @@ def _flash_attn_varlen_forward( swa = (window_size_left > 0) or (window_size_right > 0) def is_fmha_v3_fp8(): - ret = get_gfx() == "gfx942" + ret = get_gfx() in ("gfx942", "gfx950") ret = ret and (hdim_q == 128) ret = ret and (q.dtype == dtypes.fp8) ret = ret and ( @@ -2100,6 +2105,11 @@ def can_impl_fmha_v3_fwd(): ret = ret and (not swa) ret = ret and (q.dtype == dtypes.bf16 or is_fmha_v3_fp8()) ret = ret and logits_soft_cap == 0.0 + # FP8 ASM kernels assemble the GQA-shift from a fixed log2 table + # (1,2,4,8,16); arbitrary divisor ratios route to CK. + if is_fmha_v3_fp8(): + gqa_ratio = nhead_q // nhead_k + ret = ret and ((gqa_ratio & (gqa_ratio - 1)) == 0) return ret q, k, v = [maybe_contiguous(x) for x in (q, k, v)] diff --git a/csrc/cpp_itfs/mha_fwd.cu b/csrc/cpp_itfs/mha_fwd.cu index be7f94b8486..1673fa6a05b 100644 --- a/csrc/cpp_itfs/mha_fwd.cu +++ b/csrc/cpp_itfs/mha_fwd.cu @@ -183,6 +183,10 @@ std::tuple get_grid_dim(const mha_fwd_args& a, int ts_qo, const s { tg_div = 1; // do not merge the head and tail in seqlen_q direction } + if(arch_id == "gfx950") + { + tg_div = 1; // v2 asm kernel processes one Q tile per workgroup + } // batch int gdx = ((a.seqlen_q + ts_qo - 1) / ts_qo + tg_div - 1) / tg_div; int gdy = a.nhead_q; diff --git a/hsa/gfx950/fmha_v3_fwd/fmha_fwd.csv b/hsa/gfx950/fmha_v3_fwd/fmha_fwd.csv index cabc8f8b274..f418e10cd0f 100644 --- a/hsa/gfx950/fmha_v3_fwd/fmha_fwd.csv +++ b/hsa/gfx950/fmha_v3_fwd/fmha_fwd.csv @@ -6,4 +6,8 @@ bf16,128,128,2,1,0,256,64,_ZN5aiter32fmha_fwd_hd128_bf16_causal_groupE,fwd_hd128 bf16,192,128,0,0,0,128,128,_ZN5aiter25fmha_fwd_hd192_hd128_bf16E,fwd_hd192_hd128_bf16.co bf16,192,128,2,0,0,128,128,_ZN5aiter32fmha_fwd_hd192_hd128_bf16_causalE,fwd_hd192_hd128_bf16_causal.co bf16,192,128,0,1,0,128,128,_ZN5aiter31fmha_fwd_hd192_hd128_bf16_groupE,fwd_hd192_hd128_bf16_group.co -bf16,192,128,2,1,0,128,128,_ZN5aiter38fmha_fwd_hd192_hd128_bf16_causal_groupE,fwd_hd192_hd128_bf16_causal_group.co \ No newline at end of file +bf16,192,128,2,1,0,128,128,_ZN5aiter38fmha_fwd_hd192_hd128_bf16_causal_groupE,fwd_hd192_hd128_bf16_causal_group.co +fp8bf16,128,128,0,0,0,256,128,_ZN5aiter24fmha_fwd_hd128_fp8_gfx950E,fwd_hd128_fp8.co +fp8bf16,128,128,2,0,0,256,128,_ZN5aiter31fmha_fwd_hd128_fp8_causal_gfx950E,fwd_hd128_fp8_causal.co +fp8bf16,128,128,0,1,0,256,128,_ZN5aiter30fmha_fwd_hd128_fp8_group_gfx950E,fwd_hd128_fp8_group.co +fp8bf16,128,128,2,1,0,256,128,_ZN5aiter37fmha_fwd_hd128_fp8_causal_group_gfx950E,fwd_hd128_fp8_causal_group.co \ No newline at end of file diff --git a/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8.co b/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8.co new file mode 100755 index 00000000000..50cf62c3d25 Binary files /dev/null and b/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8.co differ diff --git a/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_causal.co b/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_causal.co new file mode 100755 index 00000000000..1c7b59f8d57 Binary files /dev/null and b/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_causal.co differ diff --git a/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_causal_group.co b/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_causal_group.co new file mode 100755 index 00000000000..2c5efb49315 Binary files /dev/null and b/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_causal_group.co differ diff --git a/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_group.co b/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_group.co new file mode 100755 index 00000000000..90122e1840b Binary files /dev/null and b/hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_group.co differ