Skip to content

Commit 6a82fa7

Browse files
minmengdieminmengdie
authored
MI350 mla ps mode all bf16 cases no longer follow the folding logic and nhead64,1 use m16x4 kernel (#3063)
* mla ps mode support nhead*qseqlen = n*64 through mla_a16w16_qh64_qseqlen1_gqaratio64_v3_ps * modify the gqa_ratio param location * rebase main * fix the native supported conditions * fix the fold condition * fix the dispatch logic * fix the conflict * fix the conflict * Simplify the metadata buffer size calc logic * fix the comments
1 parent ae80f53 commit 6a82fa7

11 files changed

Lines changed: 134 additions & 75 deletions

‎aiter/mla.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -369,12 +369,6 @@ def mla_decode_fwd(
369369
and kv_buffer.dtype == dtypes.fp8
370370
and max_seqlen_q == 2
371371
)
372-
or (
373-
get_gfx() == "gfx950"
374-
and nhead * max_seqlen_q % 128 == 0
375-
and q.dtype == dtypes.bf16
376-
and kv_buffer.dtype == dtypes.bf16
377-
)
378372
or (
379373
get_gfx() == "gfx950"
380374
and nhead == 8
@@ -396,6 +390,11 @@ def mla_decode_fwd(
396390
and kv_buffer.dtype == dtypes.fp8
397391
and max_seqlen_q == 1
398392
)
393+
or (
394+
get_gfx() == "gfx950"
395+
and q.dtype == dtypes.bf16
396+
and kv_buffer.dtype == dtypes.bf16
397+
)
399398
):
400399
# Natively support cases
401400
pass

‎aiter/ops/attention.py‎

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -949,11 +949,10 @@ def get_mla_metadata_info_v1(
949949
)
950950
)
951951

952-
# In sparse mode, each expanded batch has 1 Q token
953952
effective_seqlen_qo = 1 if is_sparse else max_seqlen_qo
954-
max_qo_tiles_per_batch = (
955-
int(math.ceil(effective_seqlen_qo * num_head_qo / 128))
956-
if num_head_qo == 16
953+
max_qo_tiles_per_batch = int(math.ceil(effective_seqlen_qo * num_head_qo / 16))
954+
if (
955+
num_head_qo == 16
957956
or (
958957
get_gfx() in ("gfx942", "gfx950")
959958
and num_head_qo == 128
@@ -968,12 +967,6 @@ def get_mla_metadata_info_v1(
968967
and q_dtype == dtypes.fp8
969968
and is_experimental_enabled()
970969
)
971-
or (
972-
get_gfx() == "gfx950"
973-
and (num_head_qo * effective_seqlen_qo) % 128 == 0
974-
and kv_dtype == dtypes.bf16
975-
and q_dtype == dtypes.bf16
976-
)
977970
or (
978971
get_gfx() == "gfx950"
979972
and num_head_qo == 64
@@ -982,8 +975,22 @@ def get_mla_metadata_info_v1(
982975
and effective_seqlen_qo == 1
983976
)
984977
or use_qseqlen_fold
985-
else int(math.ceil(effective_seqlen_qo * num_head_qo / 16))
986-
)
978+
):
979+
max_qo_tiles_per_batch = int(math.ceil(effective_seqlen_qo * num_head_qo / 128))
980+
elif (
981+
get_gfx() == "gfx950"
982+
and ((num_head_qo * effective_seqlen_qo) >= 128 or num_head_qo > 64)
983+
and kv_dtype == dtypes.bf16
984+
and q_dtype == dtypes.bf16
985+
and num_head_qo != 48
986+
):
987+
if num_head_qo * 2 > 128:
988+
max_qo_tiles_per_batch = effective_seqlen_qo
989+
else:
990+
max_qo_tiles_per_batch = int(
991+
math.ceil(effective_seqlen_qo * num_head_qo / 128)
992+
)
993+
987994
batch_size = batch_size * max_seqlen_qo if is_sparse else batch_size
988995
tile_cnt = batch_size * max_qo_tiles_per_batch
989996

‎csrc/kernels/mla/metadata/v1_2_device.cuh‎

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,16 @@ __launch_bounds__(opus::get_warp_size(), 1) __global__
5959
auto get_num_qo_tiles = [&](const int32_t batch_idx) {
6060
if constexpr(Traits::kQoSplits)
6161
{
62-
const int32_t packed_qo_len = qo_state.get_seqlen(batch_idx) * params.num_heads;
62+
const int32_t seqlen_qo = qo_state.get_seqlen(batch_idx);
63+
// When num_heads * 2 > kPackedQoLenPerWg, each tile can hold at most
64+
// 1 qo position (floor(kPackedQoLenPerWg / num_heads) == 1), so we
65+
// need seqlen_qo tiles. The simple ceil(packed/kPackedQoLenPerWg)
66+
// underestimates this for non-power-of-2 num_heads (e.g., 96, 112).
67+
if(params.num_heads * 2 > Traits::kPackedQoLenPerWg)
68+
{
69+
return seqlen_qo;
70+
}
71+
const int32_t packed_qo_len = seqlen_qo * params.num_heads;
6372
return integer_divide_ceil_power2(
6473
packed_qo_len, Traits::kPackedQoLenPerWg, Traits::kPackedQoLenPerWg_log2);
6574
}
@@ -440,8 +449,8 @@ void get_mla_metadata_v1_2_device(const torch::Tensor& seqlens_qo_indptr, // [ba
440449
torch::Tensor& reduce_final_map,
441450
torch::Tensor& reduce_partial_map)
442451
{
443-
constexpr int32_t kPackedQoLenPerWg = 128;
444-
const hipStream_t stream = at::hip::getCurrentHIPStream();
452+
453+
const hipStream_t stream = at::hip::getCurrentHIPStream();
445454

446455
hipDevice_t dev;
447456
hipDeviceProp_t dev_prop;
@@ -498,12 +507,12 @@ void get_mla_metadata_v1_2_device(const torch::Tensor& seqlens_qo_indptr, // [ba
498507
(max_seqlen_qo == 4)) ||
499508
((arch_id == "gfx950") && (num_heads == 64) && q_is_fp8 && kv_is_fp8 &&
500509
(max_seqlen_qo == 1)) ||
501-
((arch_id == "gfx950") && ((num_heads * max_seqlen_qo) % 128 == 0) && !q_is_fp8 &&
502-
!kv_is_fp8) ||
510+
((arch_id == "gfx950") && !q_is_fp8 && !kv_is_fp8) ||
503511
((arch_id == "gfx942" || arch_id == "gfx950") && (num_heads == 128) && q_is_fp8 &&
504512
kv_is_fp8) ||
505513
hk_mtp_experimental;
506514

515+
507516
const bool use_qseqlen_fold =
508517
!natively_supported && (arch_id == "gfx950") && q_is_fp8 && kv_is_fp8 && (num_heads > 16) &&
509518
((uni_seqlen_qo * (num_heads / 16) == 4) || ((num_heads == 64) && (uni_seqlen_qo == 2)));
@@ -534,12 +543,12 @@ void get_mla_metadata_v1_2_device(const torch::Tensor& seqlens_qo_indptr, // [ba
534543
kv_is_fp8) ||
535544
((arch_id == "gfx942") && (num_heads == 8) && (max_seqlen_qo == 2) && !q_is_fp8 &&
536545
!kv_is_fp8) ||
537-
((arch_id == "gfx950") && ((num_heads * max_seqlen_qo) % 128 == 0) && !q_is_fp8 &&
538-
!kv_is_fp8) ||
546+
((arch_id == "gfx950") && !q_is_fp8 && !kv_is_fp8) ||
539547
hk_mtp_experimental,
540548
__func__,
541549
": only supports #heads in [16, 64, 128], or (#head, uni_seqlen_qo) = (16*N, 1) where "
542-
"N is in [2, 8), or (#head, max_seqlen_qo) = (8, 4) where q and kv are fp8.")
550+
"N is in [2, 8), or (#head, max_seqlen_qo) = (8, 4) where q and kv are fp8, "
551+
"or q and kv are bf16 on gfx950")
543552

544553
int32_t num_splits = max_split_per_batch < 0
545554
? num_clusters
@@ -571,6 +580,14 @@ void get_mla_metadata_v1_2_device(const torch::Tensor& seqlens_qo_indptr, // [ba
571580
params.fixed_over_head_num_blocks = max(1, (16 + page_size - 1) / page_size);
572581
params.tail_done_threshold = max_seqlen_qo;
573582

583+
int32_t kPackedQoLenPerWg = 128;
584+
if ((arch_id == "gfx950") && !q_is_fp8 && !kv_is_fp8 &&
585+
(num_heads * max_seqlen_qo >= 64) && (num_heads <= 64) &&
586+
(((num_heads * max_seqlen_qo) < 128) ||
587+
(num_heads == 48))) {
588+
kPackedQoLenPerWg = 64;
589+
}
590+
574591
// launch kernel
575592
MLA_METADATA_DISPATCHER(
576593
max_seqlen_qo * num_heads_per_head_k,

‎csrc/kernels/mla/metadata/v1_comm.cuh‎

Lines changed: 63 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66
#include "aiter_hip_common.h"
77
#include "custom_all_reduce.cuh"
88
#include "mla.h"
9-
#include "pa.h"
109
#include "opus/opus.hpp"
10+
#include "pa.h"
1111
#include <ATen/hip/HIPContext.h>
1212
#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
1313
#include <torch/python.h>
@@ -28,10 +28,14 @@ __host__ __device__ constexpr T integer_least_multiple(T x, T y)
2828
template <typename T>
2929
__host__ __device__ constexpr T next_power_of_two(T x)
3030
{
31-
if(x <= 1) return 1;
31+
if(x <= 1)
32+
return 1;
3233
--x;
33-
x |= x >> 1; x |= x >> 2; x |= x >> 4;
34-
x |= x >> 8; x |= x >> 16;
34+
x |= x >> 1;
35+
x |= x >> 2;
36+
x |= x >> 4;
37+
x |= x >> 8;
38+
x |= x >> 16;
3539
return x + 1;
3640
}
3741

@@ -138,8 +142,8 @@ __device__ void warp_sort(int32_t* p_batch_idx,
138142
{
139143
const int32_t lane_idx = opus::lane_id();
140144

141-
const int32_t num_batches_padded = integer_least_multiple(
142-
next_power_of_two(num_batches), opus::get_warp_size());
145+
const int32_t num_batches_padded =
146+
integer_least_multiple(next_power_of_two(num_batches), opus::get_warp_size());
143147
const int32_t warp_loops = num_batches_padded / opus::get_warp_size();
144148
int32_t* p_costs = p_workspace;
145149
int32_t* p_indices = p_costs + num_batches_padded;
@@ -247,7 +251,7 @@ __host__ __device__ int32_t cal_packed_causal_kv_len(const int32_t qo_len,
247251
const int kv_len_slop =
248252
integer_divide_ceil((qo_tile_idx + 1) * packed_qo_tile_len, num_heads);
249253
const int sum = kv_len_init + kv_len_slop;
250-
result = (sum < kv_len) ? sum : kv_len;
254+
result = (sum < kv_len) ? sum : kv_len;
251255
}
252256

253257
return result;
@@ -258,9 +262,9 @@ class QoState
258262
{
259263
public:
260264
__device__ explicit QoState(const int32_t uni_seqlen_qo,
261-
const int32_t ori_seqlen_qo,
262-
const int32_t* p_lds_seqlens_qo,
263-
const int32_t* p_seqlens_qo_indptr)
265+
const int32_t ori_seqlen_qo,
266+
const int32_t* p_lds_seqlens_qo,
267+
const int32_t* p_seqlens_qo_indptr)
264268
: uni_seqlen_qo_(uni_seqlen_qo),
265269
ori_seqlen_qo_(ori_seqlen_qo),
266270
p_lds_seqlens_qo_(p_lds_seqlens_qo),
@@ -363,33 +367,53 @@ class QoState
363367
} \
364368
}
365369

366-
#define MLA_METADATA_DISPATCHER( \
367-
MAX_PACKED_SEQLEN_QO, PACKED_QO_LEN_PER_WG, UNI_SEQLEN_QO, TOPK, ...) \
368-
if(((MAX_PACKED_SEQLEN_QO) > 0) && ((MAX_PACKED_SEQLEN_QO) <= PACKED_QO_LEN_PER_WG)) \
369-
{ \
370-
constexpr bool kQoSplits = false; \
371-
if((TOPK) < 0) \
372-
{ \
373-
constexpr bool kIsSparse = false; \
374-
MLA_UNI_SEQLEN_DISPATCHER((UNI_SEQLEN_QO), __VA_ARGS__); \
375-
} \
376-
else \
377-
{ \
378-
constexpr bool kIsSparse = true; \
379-
MLA_UNI_SEQLEN_DISPATCHER((UNI_SEQLEN_QO), __VA_ARGS__); \
380-
} \
381-
} \
382-
else \
383-
{ \
384-
constexpr bool kQoSplits = true; \
385-
if((TOPK) < 0) \
386-
{ \
387-
constexpr bool kIsSparse = false; \
388-
MLA_UNI_SEQLEN_DISPATCHER((UNI_SEQLEN_QO), __VA_ARGS__); \
389-
} \
390-
else \
391-
{ \
392-
constexpr bool kIsSparse = true; \
393-
MLA_UNI_SEQLEN_DISPATCHER((UNI_SEQLEN_QO), __VA_ARGS__); \
394-
} \
370+
#define MLA_PACKED_QO_LEN_PER_WG_CASE(C_PACKED_QO_LEN_PER_WG, ...) \
371+
case C_PACKED_QO_LEN_PER_WG: { \
372+
constexpr int32_t kPackedQoLenPerWg = C_PACKED_QO_LEN_PER_WG; \
373+
__VA_ARGS__; \
374+
break; \
395375
}
376+
377+
#define MLA_PACKED_QO_LEN_PER_WG_DISPATCHER(PACKED_QO_LEN_PER_WG, ...) \
378+
switch(PACKED_QO_LEN_PER_WG) \
379+
{ \
380+
MLA_PACKED_QO_LEN_PER_WG_CASE(128, __VA_ARGS__); \
381+
MLA_PACKED_QO_LEN_PER_WG_CASE(64, __VA_ARGS__); \
382+
MLA_PACKED_QO_LEN_PER_WG_CASE(32, __VA_ARGS__); \
383+
MLA_PACKED_QO_LEN_PER_WG_CASE(16, __VA_ARGS__); \
384+
default: { \
385+
constexpr int32_t kPackedQoLenPerWg = 128; \
386+
__VA_ARGS__; \
387+
break; \
388+
} \
389+
}
390+
391+
#define MLA_METADATA_DISPATCHER( \
392+
MAX_PACKED_SEQLEN_QO, PACKED_QO_LEN_PER_WG, UNI_SEQLEN_QO, TOPK, ...) \
393+
MLA_PACKED_QO_LEN_PER_WG_DISPATCHER( \
394+
PACKED_QO_LEN_PER_WG, \
395+
if(((MAX_PACKED_SEQLEN_QO) > 0) && ((MAX_PACKED_SEQLEN_QO) <= kPackedQoLenPerWg)) { \
396+
constexpr bool kQoSplits = false; \
397+
if((TOPK) < 0) \
398+
{ \
399+
constexpr bool kIsSparse = false; \
400+
MLA_UNI_SEQLEN_DISPATCHER((UNI_SEQLEN_QO), __VA_ARGS__); \
401+
} \
402+
else \
403+
{ \
404+
constexpr bool kIsSparse = true; \
405+
MLA_UNI_SEQLEN_DISPATCHER((UNI_SEQLEN_QO), __VA_ARGS__); \
406+
} \
407+
} else { \
408+
constexpr bool kQoSplits = true; \
409+
if((TOPK) < 0) \
410+
{ \
411+
constexpr bool kIsSparse = false; \
412+
MLA_UNI_SEQLEN_DISPATCHER((UNI_SEQLEN_QO), __VA_ARGS__); \
413+
} \
414+
else \
415+
{ \
416+
constexpr bool kIsSparse = true; \
417+
MLA_UNI_SEQLEN_DISPATCHER((UNI_SEQLEN_QO), __VA_ARGS__); \
418+
} \
419+
})

‎csrc/kernels/mla/reduce.cu‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -883,6 +883,10 @@ __launch_bounds__(Traits::kNumThreads, Traits::kOccupancy) __global__
883883
MLA_REDUCE_CASE_EF(NUM_HEAD, 128, HEAD_DIM, 128, NUM_WG_PER_BH, NAME, __VA_ARGS__) \
884884
MLA_REDUCE_CASE_EF(NUM_HEAD, 128, HEAD_DIM, 512, NUM_WG_PER_BH, NAME, __VA_ARGS__) \
885885
MLA_REDUCE_CASE_EF(NUM_HEAD, 8, HEAD_DIM, 512, NUM_WG_PER_BH, NAME, __VA_ARGS__) \
886+
MLA_REDUCE_CASE_EF(NUM_HEAD, 48, HEAD_DIM, 512, NUM_WG_PER_BH, NAME, __VA_ARGS__) \
887+
MLA_REDUCE_CASE_EF(NUM_HEAD, 80, HEAD_DIM, 512, NUM_WG_PER_BH, NAME, __VA_ARGS__) \
888+
MLA_REDUCE_CASE_EF(NUM_HEAD, 96, HEAD_DIM, 512, NUM_WG_PER_BH, NAME, __VA_ARGS__) \
889+
MLA_REDUCE_CASE_EF(NUM_HEAD, 112, HEAD_DIM, 512, NUM_WG_PER_BH, NAME, __VA_ARGS__) \
886890
else MLA_REDUCE_ERROR(NUM_HEAD, HEAD_DIM, NAME);
887891

888892
#define DISPATCH_MLA_REDUCE_KERNEL( \

‎csrc/py_itfs_cu/asm_mla.cu‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -363,10 +363,14 @@ void mla_decode_stage1_asm_fwd(
363363
}
364364
}
365365

366-
if (arch_id == "gfx950" && q_type == "bf16" && kv_type == "bf16" && persistent && (gqa_ratio* max_seqlen_q % 128 == 0)){
366+
if (arch_id == "gfx950" && q_type == "bf16" && kv_type == "bf16" && persistent && (gqa_ratio * max_seqlen_q >= 128 || gqa_ratio > 64) && gqa_ratio != 48){
367367
config_max_seqlen_q = 4;
368368
config_gqa_ratio = 32;
369-
args.s_Q_Bs = gqa_ratio;
369+
args.s_MQA = gqa_ratio;
370+
} else if (arch_id == "gfx950" && q_type == "bf16" && kv_type == "bf16" && persistent && (gqa_ratio * max_seqlen_q >= 64 || gqa_ratio > 16)){
371+
config_max_seqlen_q = 1;
372+
config_gqa_ratio = 64;
373+
args.s_MQA = gqa_ratio;
370374
}
371375
int lse_flag = (lse != nullptr) ? 1 : 0;
372376
std::string kernelName = get_heuristic_kernel_mla(q_type, kv_type, config_gqa_ratio, ps, prefill, causal, config_max_seqlen_q, arch_id, config_map, lse_flag);

‎hsa/gfx950/mla/mla_a16w16_qh32_qseqlen4_gqaratio32_ps.co‎

-8 Bytes
Binary file not shown.

‎hsa/gfx950/mla/mla_a16w16_qh64_qseqlen1_gqaratio64_lse_v3_ps.co‎

57.8 KB
Binary file not shown.

‎hsa/gfx950/mla/mla_a16w16_qh64_qseqlen1_gqaratio64_v3_ps.co‎

57.2 KB
Binary file not shown.

‎hsa/gfx950/mla/mla_asm.csv‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,3 +34,5 @@ fp8,fp8,8,0,1,0,0,0,_ZN5aiter34mla_a8w8_qh8_qseqlen1_gqaratio8_v3E,mla_a8w8_qh8_
3434
fp8,fp8,8,1,4,0,0,0,_ZN5aiter35mla_a8w8_qh32_qseqlen4_gqaratio8_psE,mla_a8w8_qh32_qseqlen4_gqaratio8_ps.co
3535
fp8,fp8,8,1,4,0,0,1,_ZN5aiter39mla_a8w8_qh32_qseqlen4_gqaratio8_lse_psE,mla_a8w8_qh32_qseqlen4_gqaratio8_lse_ps.co
3636
bf16,bf16,32,1,4,0,0,0,_ZN5aiter38mla_a16w16_qh32_qseqlen4_gqaratio32_psE,mla_a16w16_qh32_qseqlen4_gqaratio32_ps.co
37+
bf16,bf16,64,1,1,0,0,0,_ZN5aiter41mla_a16w16_qh64_qseqlen1_gqaratio64_v3_psE,mla_a16w16_qh64_qseqlen1_gqaratio64_v3_ps.co
38+
bf16,bf16,64,1,1,0,0,1,_ZN5aiter45mla_a16w16_qh64_qseqlen1_gqaratio64_lse_v3_psE,mla_a16w16_qh64_qseqlen1_gqaratio64_lse_v3_ps.co

0 commit comments

Comments
 (0)