diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/unified_attention.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/unified_attention.py index 4abebc0aefe..1b71c717136 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/unified_attention.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/unified_attention.py @@ -1202,14 +1202,15 @@ def initialize( )[:, None] query_mask_qk = gl.convert_layout(query_mask, cfg.qk_layout) - # Split-KV: carve [tile_start, tile_end) into NUM_SPLITS contiguous - # split_idx==0 / NUM_SPLITS==1 is a no-op + # Split-KV: carve the sequence as reduce_segments does, so every query + # block of a sequence agrees on the split boundaries if NUM_SPLITS > 1: - active_tiles = tile_end - tile_start - tiles_per_split = (active_tiles + NUM_SPLITS - 1) // NUM_SPLITS - split_start = tile_start + split_idx * tiles_per_split + split_len = NUM_SPLITS * cfg.TILE_SIZE + seq_len = context_len + cur_batch_query_len + tiles_per_split = (seq_len + split_len - 1) // split_len + split_start = split_idx * tiles_per_split tile_end = gl.minimum(split_start + tiles_per_split, tile_end) - tile_start = split_start + tile_start = gl.maximum(split_start, tile_start) safe_tile_end = gl.minimum(safe_tile_end, tile_end - 1) safe_tile_end = gl.maximum(safe_tile_end, tile_start) @@ -1432,6 +1433,8 @@ def store_partial( # row-max, so it needs whatever scale the softmax used. if not cfg.USE_SINKS: M = M * self.SM_scale + # no visible key in this split: keep the guarded 0 out of the reduce max + M = gl.where(L > 0.0, M, float("-inf")) layout: gl.constexpr = cfg.pv_layout offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, layout)) offs_d = gl.arange(0, cfg.HEAD_SIZE, layout=gl.SliceLayout(0, layout)) @@ -1904,8 +1907,26 @@ def _unified_attention_gluon_kernel( NUM_SPLITS, ) - # This split owns no tiles if NUM_SPLITS > 1 and pgm.tile_start >= pgm.tile_end: + # past this block's causal edge, but reduce_segments still reads it + if pgm.tile_start * TILE_SIZE < seq_len: + empty = gl.zeros( + [BLOCK_M], dtype=gl.float32, layout=gl.SliceLayout(1, cfg.pv_layout) + ) + pgm.store_partial( + empty, + empty, + gl.zeros([BLOCK_M, HEAD_SIZE], dtype=gl.float32, layout=cfg.pv_layout), + partial_m_ptr, + partial_l_ptr, + partial_acc_ptr, + split_idx, + q_block_local_idx, + cur_batch_in_all_start_index, + kv_head_idx, + cur_batch_query_len, + NUM_SPLITS, + ) return # TILE_SIZE == BLOCK_SIZE: one page per tile (fast path). Otherwise gather diff --git a/aiter/ops/triton/_triton_kernels/attention/unified_attention.py b/aiter/ops/triton/_triton_kernels/attention/unified_attention.py index d7be3904646..a17cd2b7f98 100644 --- a/aiter/ops/triton/_triton_kernels/attention/unified_attention.py +++ b/aiter/ops/triton/_triton_kernels/attention/unified_attention.py @@ -1005,6 +1005,7 @@ def reduce_segments( NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int FP8_MIN: tl.constexpr = float8_info.min, FP8_MAX: tl.constexpr = float8_info.max, + SEGMENT_BLOCK: tl.constexpr = 0, # segments per merge step, 0 = all at once ): query_token_idx = tl.program_id(0) query_head_idx = tl.program_id(1) @@ -1026,9 +1027,13 @@ def reduce_segments( # create masks for subsequent loads act_num_segments = cdiv_fn(seq_len, tiles_per_segment * TILE_SIZE) - segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full( - [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32 + single_tile: tl.constexpr = ( + SEGMENT_BLOCK == 0 or SEGMENT_BLOCK >= NUM_SEGMENTS_PER_SEQ ) + if single_tile: + segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full( + [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32 + ) offs_vd = tl.arange(0, V_HEAD_SIZE_PADDED) if V_HEAD_SIZE_PADDED != V_HEAD_SIZE: @@ -1036,35 +1041,73 @@ def reduce_segments( else: v_dim_mask = tl.full((1,), 1, dtype=tl.int1) - # load segment maxima - segm_offset = ( - query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) - + query_head_idx * NUM_SEGMENTS_PER_SEQ - + tl.arange(0, NUM_SEGMENTS_PER_SEQ) - ) - segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf")) - overall_max = tl.max(segm_max) - - # load and rescale segment exp sums - segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) - segm_expsum = segm_expsum * tl.math.exp2(segm_max - overall_max) - overall_expsum = tl.sum(segm_expsum) - - # load, rescale, and add segment attention outputs - segm_output_offset = ( - query_token_idx.to(tl.int64) - * (num_query_heads * NUM_SEGMENTS_PER_SEQ * V_HEAD_SIZE_PADDED) - + query_head_idx * (NUM_SEGMENTS_PER_SEQ * V_HEAD_SIZE_PADDED) - + tl.arange(0, NUM_SEGMENTS_PER_SEQ)[:, None] * V_HEAD_SIZE_PADDED - + offs_vd[None, :] - ) - segm_output = tl.load( - segm_output_ptr + segm_output_offset, - mask=segm_mask[:, None] & v_dim_mask[None, :], - other=0.0, - ) - segm_output *= tl.math.exp2(segm_max - overall_max)[:, None] - acc_sum = tl.sum(segm_output, axis=0) + if single_tile: + # load segment maxima + segm_offset = ( + query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + query_head_idx * NUM_SEGMENTS_PER_SEQ + + tl.arange(0, NUM_SEGMENTS_PER_SEQ) + ) + segm_max = tl.load( + segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf") + ) + overall_max = tl.max(segm_max) + + # load and rescale segment exp sums + segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) + segm_expsum = segm_expsum * tl.math.exp2(segm_max - overall_max) + overall_expsum = tl.sum(segm_expsum) + + # load, rescale, and add segment attention outputs + segm_output_offset = ( + query_token_idx.to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * V_HEAD_SIZE_PADDED) + + query_head_idx * (NUM_SEGMENTS_PER_SEQ * V_HEAD_SIZE_PADDED) + + tl.arange(0, NUM_SEGMENTS_PER_SEQ)[:, None] * V_HEAD_SIZE_PADDED + + offs_vd[None, :] + ) + segm_output = tl.load( + segm_output_ptr + segm_output_offset, + mask=segm_mask[:, None] & v_dim_mask[None, :], + other=0.0, + ) + segm_output *= tl.math.exp2(segm_max - overall_max)[:, None] + acc_sum = tl.sum(segm_output, axis=0) + else: + segm_base = ( + query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + query_head_idx * NUM_SEGMENTS_PER_SEQ + ) + segm_output_base = query_token_idx.to(tl.int64) * ( + num_query_heads * NUM_SEGMENTS_PER_SEQ * V_HEAD_SIZE_PADDED + ) + query_head_idx * (NUM_SEGMENTS_PER_SEQ * V_HEAD_SIZE_PADDED) + offs_s = tl.arange(0, SEGMENT_BLOCK) + overall_max = tl.full([], float("-inf"), dtype=tl.float32) + overall_expsum = tl.zeros([], dtype=tl.float32) + acc_sum = tl.zeros([V_HEAD_SIZE_PADDED], dtype=tl.float32) + for s0 in range(0, act_num_segments, SEGMENT_BLOCK): + segs = s0 + offs_s + segm_mask = segs < act_num_segments + segm_max = tl.load( + segm_max_ptr + segm_base + segs, mask=segm_mask, other=float("-inf") + ) + new_max = tl.maximum(overall_max, tl.max(segm_max)) + alpha = tl.math.exp2(overall_max - new_max) + weight = tl.math.exp2(segm_max - new_max) + segm_expsum = tl.load( + segm_expsum_ptr + segm_base + segs, mask=segm_mask, other=0.0 + ) + overall_expsum = overall_expsum * alpha + tl.sum(segm_expsum * weight) + segm_output = tl.load( + segm_output_ptr + + segm_output_base + + segs[:, None] * V_HEAD_SIZE_PADDED + + offs_vd[None, :], + mask=segm_mask[:, None] & v_dim_mask[None, :], + other=0.0, + ) + acc_sum = acc_sum * alpha + tl.sum(segm_output * weight[:, None], axis=0) + overall_max = new_max # safely divide by overall_expsum, returning 0.0 if overall_expsum is 0 acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum) diff --git a/aiter/ops/triton/attention/unified_attention.py b/aiter/ops/triton/attention/unified_attention.py index eb3349473d5..b6b62ba2001 100644 --- a/aiter/ops/triton/attention/unified_attention.py +++ b/aiter/ops/triton/attention/unified_attention.py @@ -147,16 +147,40 @@ def _gfx950_gluon_supported(params: _UAParams): ) +def _gfx950_split_params(params: _UAParams): + """Size the split by the gluon grid when a sequence spans several query blocks.""" + if params.all_decode or params.max_seqlen_q * params.num_queries_per_kv <= 128: + return params + config = get_unified_attention_config("attn_3d", params, backend="gluon") + block_m = max(config["BLOCK_M"], triton.next_power_of_2(params.num_queries_per_kv)) + block_q = block_m // params.num_queries_per_kv + num_q_blocks = params.num_tokens // block_q + params.num_seqs + return params._replace(num_2d_prgms=num_q_blocks * params.num_kv_heads) + + +def _gfx950_split_grid(params: _UAParams): + """Decode, and batches too small to fill the GPU, take the split grid.""" + if params.all_decode: + return True + if params.sliding_window > 0: + return False + # unsplit, the 2d grid is faster + config = get_unified_attention_config( + "kv_split", _gfx950_split_params(params), backend="gluon" + ) + return config["NUM_SEGMENTS"] > 1 + + def use_2d_kernel(params: _UAParams, backend: str = "triton"): # if IS_DEVICE_ARCH_GFX12, always use 3D if all_decode and 2D otherwise if IS_DEVICE_ARCH_GFX12: return (params.sliding_window > 0) or (not params.all_decode) # The gfx950 Gluon kernel is one kernel with a decode and a prefill grid, - # and it owns its KV split, so decode always takes the 3d path (a single - # segment simply drops the split axis) and everything else the 2d one. + # and it owns its KV split, so decode and split batches take the 3d path + # (a single segment simply drops the split axis) and the rest the 2d one. if backend == "gluon" and _gfx950_gluon_supported(params): - return not params.all_decode + return not _gfx950_split_grid(params) if params.head_size >= 512 and not get_arch().is_rdna and not params.all_decode: return True @@ -368,8 +392,11 @@ def unified_attention( _unified_attention_2d_triton(params) else: use_gluon_3d = is_3d_gluon_available(params, backend) + split_params = params + if use_gluon_3d and DEVICE_ARCH == "gfx950": + split_params = _gfx950_split_params(params) config = get_unified_attention_config( - "kv_split", params, backend="gluon" if use_gluon_3d else "triton" + "kv_split", split_params, backend="gluon" if use_gluon_3d else "triton" ) NUM_SEGMENTS = config["NUM_SEGMENTS"] # The triton kernels read a shuffled tile as one contiguous page, so @@ -1121,7 +1148,7 @@ def _unified_attention_3d_gfx950( NUM_SEGMENTS, TILE_SIZE, ): - """Decode: one program per (query block, kv head[, KV split]). + """Decode and split batches: one program per (query block, kv head[, KV split]). num_warps is not tuned directly here; BLOCK_M rows are split MFMA_DIM to a warp, so the warp count follows the query group the launch has to cover. @@ -1130,6 +1157,10 @@ def _unified_attention_3d_gfx950( config = get_unified_attention_config("attn_3d", params, backend="gluon") MFMA_DIM = config["MFMA_DIM"] BLOCK_M = max(config["BLOCK_M"], triton.next_power_of_2(params.num_queries_per_kv)) + q_rows = triton.next_power_of_2(params.max_seqlen_q * params.num_queries_per_kv) + if not params.all_decode and q_rows <= 128: + # short queries: one block covers the sequence + BLOCK_M = max(BLOCK_M, q_rows) # An unsplit launch that already leaves CUs idle can afford a wider query # block, which buys back the warps the missing split axis would have used. if NUM_SEGMENTS == 1 and params.num_2d_prgms <= 2 * params.num_sms: diff --git a/aiter/ops/triton/configs/gfx942/triton/attention/unified_attention/DEFAULT.json b/aiter/ops/triton/configs/gfx942/triton/attention/unified_attention/DEFAULT.json index 07b10cfd602..bb4ff26a234 100644 --- a/aiter/ops/triton/configs/gfx942/triton/attention/unified_attention/DEFAULT.json +++ b/aiter/ops/triton/configs/gfx942/triton/attention/unified_attention/DEFAULT.json @@ -145,7 +145,19 @@ "MIN_SEGMENTS": 16, "SEGMENTS_PER_CU": 8, "MAX_SEGMENTS": 128, - "SEGMENT_TILE_MAX": 64 + "SEGMENT_TILE_MAX": 64, + "SEGMENT_BLOCK": 32 + }, + "D_GEQ_256": { + "num_stages": 1, + "waves_per_eu": 2, + "num_warps": 2, + "SMALL_SPLIT_MAX": 8, + "MIN_SEGMENTS": 8, + "SEGMENTS_PER_CU": 16, + "MAX_SEGMENTS": 128, + "SEGMENT_TILE_MAX": 64, + "SEGMENT_BLOCK": 32 }, "any": { "num_stages": 1, diff --git a/aiter/ops/triton/configs/gfx950/gluon/attention/unified_attention/DEFAULT.json b/aiter/ops/triton/configs/gfx950/gluon/attention/unified_attention/DEFAULT.json index ab0d133ccb5..5d2afcf461f 100644 --- a/aiter/ops/triton/configs/gfx950/gluon/attention/unified_attention/DEFAULT.json +++ b/aiter/ops/triton/configs/gfx950/gluon/attention/unified_attention/DEFAULT.json @@ -9,10 +9,13 @@ "attn_3d": [ "D", "SW", + "Q", "DT" ], "kv_split": [ "SW", + "D", + "Q", "DT" ] }, @@ -40,16 +43,7 @@ "num_warps": 4, "MFMA_DIM": 32, "NUM_BUFFERS": 2, - "waves_per_eu": 3, - "TILE_SIZE_MIN": 64, - "TILE_SIZE_MAX": 64 - }, - "D_GEQ_256.SKLT.DT_fp8_fp8": { - "BLOCK_M": 128, - "num_warps": 4, - "MFMA_DIM": 32, - "NUM_BUFFERS": 1, - "waves_per_eu": 1, + "waves_per_eu": 0, "TILE_SIZE_MIN": 64, "TILE_SIZE_MAX": 64 }, @@ -57,8 +51,8 @@ "BLOCK_M": 128, "num_warps": 4, "MFMA_DIM": 32, - "NUM_BUFFERS": 1, - "waves_per_eu": 2, + "NUM_BUFFERS": 2, + "waves_per_eu": 0, "TILE_SIZE_MIN": 64, "TILE_SIZE_MAX": 64 }, @@ -76,12 +70,24 @@ "num_warps": 4, "MFMA_DIM": 32, "NUM_BUFFERS": 2, - "waves_per_eu": 2, + "waves_per_eu": 0, "TILE_SIZE_MIN": 64, "TILE_SIZE_MAX": 64 } }, "attn_3d": { + "D_LEQ_64.Q_GEQ_64.DT_fp8_fp8": { + "BLOCK_M": 128, + "MFMA_DIM": 32, + "NUM_BUFFERS": 2, + "waves_per_eu": 0 + }, + "D_LEQ_64.Q_GEQ_64": { + "BLOCK_M": 128, + "MFMA_DIM": 32, + "NUM_BUFFERS": 2, + "waves_per_eu": 0 + }, "D_LEQ_64.DT_fp8_fp8": { "BLOCK_M": 32, "MFMA_DIM": 32, @@ -108,26 +114,62 @@ "waves_per_eu": 1, "WIDE_BLOCK_M_MULT": 2 }, + "D_GEQ_256.Q_GEQ_64.DT_fp8_fp8": { + "BLOCK_M": 64, + "MFMA_DIM": 32, + "NUM_BUFFERS": 2, + "waves_per_eu": 0 + }, + "D_GEQ_256.Q_GEQ_64": { + "BLOCK_M": 128, + "MFMA_DIM": 32, + "NUM_BUFFERS": 2, + "waves_per_eu": 0 + }, "D_GEQ_256.DT_fp8_fp8": { "BLOCK_M": 32, "MFMA_DIM": 32, - "NUM_BUFFERS": 1, + "NUM_BUFFERS": 2, "waves_per_eu": 1, "WIDE_BLOCK_M_MULT": 2 }, "D_GEQ_256": { "BLOCK_M": 16, "MFMA_DIM": 16, - "NUM_BUFFERS": 1, + "NUM_BUFFERS": 2, "waves_per_eu": 1, "WIDE_BLOCK_M_MULT": 2 }, + "Q_GEQ_64.DT_fp8_fp8": { + "BLOCK_M": 64, + "MFMA_DIM": 32, + "NUM_BUFFERS": 2, + "waves_per_eu": 0 + }, + "Q_GEQ_64": { + "BLOCK_M": 128, + "MFMA_DIM": 32, + "NUM_BUFFERS": 2, + "waves_per_eu": 0 + }, + "Q_GEQ_2.DT_fp8_fp8": { + "BLOCK_M": 32, + "MFMA_DIM": 32, + "NUM_BUFFERS": 2, + "waves_per_eu": 1 + }, "DT_fp8_fp8": { "BLOCK_M": 32, "MFMA_DIM": 32, "NUM_BUFFERS": 2, "waves_per_eu": 2 }, + "Q_GEQ_2": { + "BLOCK_M": 32, + "MFMA_DIM": 32, + "NUM_BUFFERS": 2, + "waves_per_eu": 1 + }, "any": { "BLOCK_M": 16, "MFMA_DIM": 16, @@ -145,27 +187,149 @@ "TILE_SIZE_MIN": 64, "TILE_SIZE_MAX": 64 }, + "Q_LEQ_1.DT_fp8_fp8": { + "SEGMENTS_PER_CU": 4, + "MFMA_DIM": 32, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 4, + "SPLIT_MIN_SHARE": 3, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 + }, + "D_LEQ_64.Q_GEQ_64.DT_fp8_fp8": { + "SEGMENTS_PER_CU": 4, + "MFMA_DIM": 32, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 127, + "SPLIT_MIN_SHARE": 2, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 + }, + "D_LEQ_128.Q_GEQ_64.DT_fp8_fp8": { + "SEGMENTS_PER_CU": 4, + "MFMA_DIM": 32, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 63, + "SPLIT_MIN_SHARE": 2, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 + }, + "D_LEQ_128.Q_GEQ_2.DT_fp8_fp8": { + "SEGMENTS_PER_CU": 4, + "MFMA_DIM": 32, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 4, + "SPLIT_MIN_SHARE": 2, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 + }, + "Q_GEQ_64.DT_fp8_fp8": { + "SEGMENTS_PER_CU": 4, + "MFMA_DIM": 32, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 127, + "SPLIT_MIN_SHARE": 4, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 + }, "DT_fp8_fp8": { "SEGMENTS_PER_CU": 4, "MFMA_DIM": 32, "MIN_SEGMENTS": 1, - "SPLIT_MIN_TILES": 16, + "SPLIT_MIN_TILES": 4, + "SPLIT_MIN_SHARE": 2, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 64 + }, + "Q_LEQ_1": { + "SEGMENTS_PER_CU": 2, + "MFMA_DIM": 16, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 4, "SPLIT_MIN_SHARE": 3, "SEGMENT_TILE_MIN": 64, "SEGMENT_TILE_MAX": 64, "TILE_SIZE_MIN": 64, - "TILE_SIZE_MAX": 64 + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 + }, + "D_LEQ_64.Q_GEQ_64": { + "SEGMENTS_PER_CU": 2, + "MFMA_DIM": 16, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 127, + "SPLIT_MIN_SHARE": 2, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 + }, + "D_LEQ_128.Q_GEQ_64": { + "SEGMENTS_PER_CU": 2, + "MFMA_DIM": 16, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 63, + "SPLIT_MIN_SHARE": 2, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 + }, + "D_LEQ_128.Q_GEQ_2": { + "SEGMENTS_PER_CU": 2, + "MFMA_DIM": 16, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 4, + "SPLIT_MIN_SHARE": 2, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 + }, + "Q_GEQ_64": { + "SEGMENTS_PER_CU": 2, + "MFMA_DIM": 16, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 127, + "SPLIT_MIN_SHARE": 4, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 }, "any": { "SEGMENTS_PER_CU": 2, "MFMA_DIM": 16, "MIN_SEGMENTS": 1, - "SPLIT_MIN_TILES": 16, - "SPLIT_MIN_SHARE": 3, + "SPLIT_MIN_TILES": 4, + "SPLIT_MIN_SHARE": 2, "SEGMENT_TILE_MIN": 64, "SEGMENT_TILE_MAX": 64, "TILE_SIZE_MIN": 64, - "TILE_SIZE_MAX": 64 + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 64 } } } diff --git a/aiter/ops/triton/configs/gfx950/triton/attention/unified_attention/DEFAULT.json b/aiter/ops/triton/configs/gfx950/triton/attention/unified_attention/DEFAULT.json index 348c826d5c2..95c5069fe01 100644 --- a/aiter/ops/triton/configs/gfx950/triton/attention/unified_attention/DEFAULT.json +++ b/aiter/ops/triton/configs/gfx950/triton/attention/unified_attention/DEFAULT.json @@ -246,7 +246,8 @@ "MIN_SEGMENTS": 16, "SEGMENTS_PER_CU": 8, "MAX_SEGMENTS": 128, - "SEGMENT_TILE_MAX": 64 + "SEGMENT_TILE_MAX": 64, + "SEGMENT_BLOCK": 32 }, "D_LEQ_128": { "num_stages": 1, @@ -267,7 +268,8 @@ "MIN_SEGMENTS": 8, "SEGMENTS_PER_CU": 16, "MAX_SEGMENTS": 128, - "SEGMENT_TILE_MAX": 64 + "SEGMENT_TILE_MAX": 64, + "SEGMENT_BLOCK": 32 } }, "kv_split": {