From 1aacfa54c2f8dda7f40e224d67c0aa8403cf0ff2 Mon Sep 17 00:00:00 2001 From: "Mehmet (Cagri) Kaymak" Date: Thu, 8 Oct 2026 19:59:13 +0000 Subject: [PATCH 1/9] [Triton/Gluon] [GFX950] Unified attention: split KV for short multi-token batches, chunked segment reduce Fixes the gfx950 Gluon slowdowns reported in #6221 (speculative-decode verify and decode with few KV heads). - Batches of short queries (max_seqlen_q * queries_per_kv <= 128, no sliding window) take the split-KV grid when the kv_split config splits them. The query block is widened to hold every query of a sequence, so the split ranges match what reduce_segments assumes. Before, every such batch took the unsplit 2d grid: at batch 1 with 2 KV heads that is 2 programs walking the whole context. - store_partial writes -inf as the max of a row that saw no key in its split. - reduce_segments takes SEGMENT_BLOCK and merges the segments in chunks of that size with an online softmax. The single [NUM_SEGMENTS, head] tile spills at head size 256 and 512. 0 (default) keeps the single-tile code. - gfx950 configs: kv_split caps the split at 128 for decode and 64 otherwise (vLLM full-graph capture passes max_model_len as max_seqlen_k, so the uncapped fp8 split reached 512); D >= 256 decode uses the KV double buffer; the Triton reduce section sets SEGMENT_BLOCK 32 for D >= 256. --- .../gfx950/attention/unified_attention.py | 4 + .../attention/unified_attention.py | 92 +++++++--- .../ops/triton/attention/unified_attention.py | 38 +++- .../attention/unified_attention/DEFAULT.json | 35 +++- .../attention/unified_attention/DEFAULT.json | 6 +- .../attention/test_unified_attention.py | 165 ++++++++++++++++++ 6 files changed, 303 insertions(+), 37 deletions(-) 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..36abad22ba0 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/unified_attention.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/unified_attention.py @@ -1432,6 +1432,10 @@ def store_partial( # row-max, so it needs whatever scale the softmax used. if not cfg.USE_SINKS: M = M * self.SM_scale + # A row with no causally visible key in this split (a short multi-token + # query whose split ends past the row's position) keeps the guarded 0 + # max; keep it out of the reduce's 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)) diff --git a/aiter/ops/triton/_triton_kernels/attention/unified_attention.py b/aiter/ops/triton/_triton_kernels/attention/unified_attention.py index d7be3904646..bf4e34c07df 100644 --- a/aiter/ops/triton/_triton_kernels/attention/unified_attention.py +++ b/aiter/ops/triton/_triton_kernels/attention/unified_attention.py @@ -1005,6 +1005,8 @@ def reduce_segments( NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int FP8_MIN: tl.constexpr = float8_info.min, FP8_MAX: tl.constexpr = float8_info.max, + # segments merged per loop step; 0 merges all NUM_SEGMENTS_PER_SEQ in one tile + SEGMENT_BLOCK: tl.constexpr = 0, ): query_token_idx = tl.program_id(0) query_head_idx = tl.program_id(1) @@ -1026,9 +1028,6 @@ 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 - ) offs_vd = tl.arange(0, V_HEAD_SIZE_PADDED) if V_HEAD_SIZE_PADDED != V_HEAD_SIZE: @@ -1036,35 +1035,74 @@ def reduce_segments( else: v_dim_mask = tl.full((1,), 1, dtype=tl.int1) - # load segment maxima - segm_offset = ( + segm_base = ( 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) + 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) - # 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) + if SEGMENT_BLOCK == 0 or SEGMENT_BLOCK >= NUM_SEGMENTS_PER_SEQ: + segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full( + [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32 + ) - # 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) + # load segment maxima + segm_offset = segm_base + 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 = ( + segm_output_base + + 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: + # online merge over the written segments only, so the cost follows the + # context rather than NUM_SEGMENTS_PER_SEQ + 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..b750aad2787 100644 --- a/aiter/ops/triton/attention/unified_attention.py +++ b/aiter/ops/triton/attention/unified_attention.py @@ -59,6 +59,9 @@ # LDS the gfx950 KV double buffer may occupy; a wider tile falls back to one buffer. _GFX950_LDS_BUDGET = 160 * 1024 +# Widest query block the gfx950 split-KV grid takes for a multi-token batch. +_GFX950_SPLIT_MAX_BLOCK_M = 128 + def _is_gluon_available(): return any(supported in DEVICE_ARCH for supported in _GLUON_SUPPORTED_ARCHS) @@ -147,16 +150,36 @@ def _gfx950_gluon_supported(params: _UAParams): ) +def _gfx950_split_grid(params: _UAParams): + """Whether the gfx950 Gluon launch takes the split-capable (decode) grid. + + Decode does. So does a batch of short queries such as a speculative decoding + verify step, when one query block can hold every query of a sequence and the + launch is small enough to split: each block then walks the whole context, and + the KV split keeps a small batch from leaving most CUs idle. Unsplit, the 2d + grid's wider query block is faster. + """ + if params.all_decode: + return True + if ( + params.sliding_window > 0 + or params.max_seqlen_q * params.num_queries_per_kv > _GFX950_SPLIT_MAX_BLOCK_M + ): + return False + config = get_unified_attention_config("kv_split", 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 short queries 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 @@ -1121,7 +1144,7 @@ def _unified_attention_3d_gfx950( NUM_SEGMENTS, TILE_SIZE, ): - """Decode: one program per (query block, kv head[, KV split]). + """Decode and short queries: 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 +1153,13 @@ 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)) + if not params.all_decode: + # one query block per sequence: the split carves the KV range from + # seq_len, which is what reduce_segments assumes + BLOCK_M = max( + BLOCK_M, + triton.next_power_of_2(params.max_seqlen_q * params.num_queries_per_kv), + ) # 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/gfx950/gluon/attention/unified_attention/DEFAULT.json b/aiter/ops/triton/configs/gfx950/gluon/attention/unified_attention/DEFAULT.json index ab0d133ccb5..4ecfae798df 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 @@ -13,6 +13,7 @@ ], "kv_split": [ "SW", + "Q", "DT" ] }, @@ -111,14 +112,14 @@ "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 }, @@ -145,6 +146,18 @@ "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": 16, + "SPLIT_MIN_SHARE": 3, + "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, @@ -154,7 +167,20 @@ "SEGMENT_TILE_MIN": 64, "SEGMENT_TILE_MAX": 64, "TILE_SIZE_MIN": 64, - "TILE_SIZE_MAX": 64 + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 64 + }, + "Q_LEQ_1": { + "SEGMENTS_PER_CU": 2, + "MFMA_DIM": 16, + "MIN_SEGMENTS": 1, + "SPLIT_MIN_TILES": 16, + "SPLIT_MIN_SHARE": 3, + "SEGMENT_TILE_MIN": 64, + "SEGMENT_TILE_MAX": 64, + "TILE_SIZE_MIN": 64, + "TILE_SIZE_MAX": 64, + "MAX_SEGMENTS": 128 }, "any": { "SEGMENTS_PER_CU": 2, @@ -165,7 +191,8 @@ "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": { diff --git a/op_tests/triton_tests/attention/test_unified_attention.py b/op_tests/triton_tests/attention/test_unified_attention.py index ca14159e273..cb9d310453e 100644 --- a/op_tests/triton_tests/attention/test_unified_attention.py +++ b/op_tests/triton_tests/attention/test_unified_attention.py @@ -738,6 +738,171 @@ def test_triton_unified_attn( ) +# Decode split wide enough that reduce_segments merges the segments in chunks +# (reduce config SEGMENT_BLOCK < NUM_SEGMENTS). +@pytest.mark.parametrize( + "backend, head_size", [("triton", 512), ("triton", 256), ("gluon", 256)] +) +@torch.inference_mode() +def test_unified_attn_chunked_reduce(backend: str, head_size: int) -> None: + if DEVICE_ARCH != "gfx950": + pytest.skip("gfx950 reduce config") + seq_lens = [(1, 8192), (1, 1025), (1, 3000), (1, 37)] + ( + query, + key_cache_orig, + value_cache_orig, + key_cache, + value_cache, + sinks, + output, + cu_query_lens, + kv_lens, + max_query_len, + max_kv_len, + scale, + window_size, + block_tables, + *_, + ) = generate_data( + seq_lens=seq_lens, + num_blocks=1024, + block_size=64, + head_size=head_size, + num_heads=(16, 1), + device="cuda", + ) + unified_attention( + q=query, + k=key_cache, + v=value_cache, + out=output, + cu_seqlens_q=cu_query_lens, + seqused_k=kv_lens, + max_seqlen_q=max_query_len, + max_seqlen_k=max_kv_len, + softmax_scale=scale, + causal=True, + window_size=window_size, + block_table=block_tables, + softcap=0, + q_descale=None, + k_descale=None, + v_descale=None, + sinks=sinks, + backend=backend, + ) + ref_output = ref_paged_attn( + query=query, + key_cache=key_cache_orig, + value_cache=value_cache_orig, + query_lens=[x[0] for x in seq_lens], + kv_lens=[x[1] for x in seq_lens], + block_tables=block_tables, + scale=scale, + out_dtype=torch.bfloat16, + sinks=sinks, + ) + torch.testing.assert_close( + output.float(), ref_output.float(), atol=1.5e-2, rtol=1e-2 + ) + + +# Speculative-decode verify batches: short queries over a long context take the +# split-KV grid. kv_len 1025 puts the last key alone in a split that the first +# query of its sequence cannot see. +@pytest.mark.parametrize("num_heads", [(32, 2), (8, 1)]) +@pytest.mark.parametrize("head_size", [128, 256]) +@pytest.mark.parametrize( + "q_dtype, kv_dtype", + [(torch.bfloat16, torch.bfloat16), (e4m3_dtype, e4m3_dtype)], +) +@torch.inference_mode() +def test_gluon_unified_attn_spec_verify( + num_heads: tuple[int, int], + head_size: int, + q_dtype: torch.dtype, + kv_dtype: torch.dtype, +) -> None: + if DEVICE_ARCH != "gfx950": + pytest.skip("gfx950 gluon split-KV grid") + seq_lens = [(2, 1025), (4, 8192), (3, 3000), (1, 523)] + query_lens = [x[0] for x in seq_lens] + ( + query, + key_cache_orig, + value_cache_orig, + key_cache, + value_cache, + sinks, + output, + cu_query_lens, + kv_lens, + max_query_len, + max_kv_len, + scale, + window_size, + block_tables, + _maybe_quant_query, + _query_scales, + q_descale, + k_descale, + v_descale, + output_scale, + ) = generate_data( + seq_lens=seq_lens, + num_blocks=2048, + block_size=64, + head_size=head_size, + num_heads=num_heads, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + device="cuda", + ) + for use_sinks in (False, True): + unified_attention( + q=query, + k=key_cache, + v=value_cache, + out=output, + cu_seqlens_q=cu_query_lens, + seqused_k=kv_lens, + max_seqlen_q=max_query_len, + max_seqlen_k=max_kv_len, + softmax_scale=scale, + causal=True, + window_size=window_size, + block_table=block_tables, + softcap=0, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + sinks=sinks if use_sinks else None, + backend="gluon", + ) + ref_output = ref_paged_attn( + query=query, + key_cache=key_cache_orig, + value_cache=value_cache_orig, + query_lens=query_lens, + kv_lens=[x[1] for x in seq_lens], + block_tables=block_tables, + scale=scale, + out_dtype=torch.bfloat16, + sliding_window=None, + soft_cap=None, + sinks=sinks if use_sinks else None, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + output_scale=output_scale, + ) + atol, rtol = (1.5e-1, 1.5e-1) if q_dtype == e4m3_dtype else (1.5e-2, 1e-2) + torch.testing.assert_close( + output.float(), ref_output.float(), atol=atol, rtol=rtol + ) + + @pytest.mark.parametrize( "seq_lens", [ From 04678497799ee6291e811e48c51fcdd91f96f352 Mon Sep 17 00:00:00 2001 From: "Mehmet (Cagri) Kaymak" Date: Thu, 8 Oct 2026 20:28:41 +0000 Subject: [PATCH 2/9] [Triton] Unified attention: chunked segment reduce on gfx942, keep the single-tile reduce as on main - gfx942 Triton reduce config: SEGMENT_BLOCK 32 for D > 128, as on gfx950. A D_LEQ_128 entry with the previous values keeps D <= 128 on the single tile. - reduce_segments: the single-tile branch keeps main's op order, so an arch without SEGMENT_BLOCK compiles to the same instructions as before. --- .../attention/unified_attention.py | 38 +++++++++++-------- .../attention/unified_attention/DEFAULT.json | 14 ++++++- 2 files changed, 36 insertions(+), 16 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/unified_attention.py b/aiter/ops/triton/_triton_kernels/attention/unified_attention.py index bf4e34c07df..ac86a6a3496 100644 --- a/aiter/ops/triton/_triton_kernels/attention/unified_attention.py +++ b/aiter/ops/triton/_triton_kernels/attention/unified_attention.py @@ -1028,6 +1028,13 @@ def reduce_segments( # create masks for subsequent loads act_num_segments = cdiv_fn(seq_len, tiles_per_segment * TILE_SIZE) + 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: @@ -1035,21 +1042,13 @@ def reduce_segments( else: v_dim_mask = tl.full((1,), 1, dtype=tl.int1) - 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) - - if SEGMENT_BLOCK == 0 or SEGMENT_BLOCK >= NUM_SEGMENTS_PER_SEQ: - segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full( - [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32 - ) - + if single_tile: # load segment maxima - segm_offset = segm_base + tl.arange(0, NUM_SEGMENTS_PER_SEQ) + 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") ) @@ -1062,7 +1061,9 @@ def reduce_segments( # load, rescale, and add segment attention outputs segm_output_offset = ( - 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) + tl.arange(0, NUM_SEGMENTS_PER_SEQ)[:, None] * V_HEAD_SIZE_PADDED + offs_vd[None, :] ) @@ -1076,6 +1077,13 @@ def reduce_segments( else: # online merge over the written segments only, so the cost follows the # context rather than NUM_SEGMENTS_PER_SEQ + 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) 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..49495a8d2a4 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,6 +145,17 @@ "MIN_SEGMENTS": 16, "SEGMENTS_PER_CU": 8, "MAX_SEGMENTS": 128, + "SEGMENT_TILE_MAX": 64, + "SEGMENT_BLOCK": 32 + }, + "D_LEQ_128": { + "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 }, "any": { @@ -155,7 +166,8 @@ "MIN_SEGMENTS": 8, "SEGMENTS_PER_CU": 16, "MAX_SEGMENTS": 128, - "SEGMENT_TILE_MAX": 64 + "SEGMENT_TILE_MAX": 64, + "SEGMENT_BLOCK": 32 } }, "kv_split": { From c5ee0f6a2bdaa4ac754f2f70545f7d5b57db863e Mon Sep 17 00:00:00 2001 From: "Mehmet (Cagri) Kaymak" Date: Thu, 8 Oct 2026 20:46:37 +0000 Subject: [PATCH 3/9] [Triton] [GFX942] Unified attention: chunked reduce through a D_GEQ_256 entry, leave any as on main Replaces the D_LEQ_128 entry and the SEGMENT_BLOCK on any with a D_GEQ_256 reduce entry (the previous any values plus SEGMENT_BLOCK 32), so head sizes below 256 keep the single-tile reduce from main. --- .../triton/attention/unified_attention/DEFAULT.json | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 49495a8d2a4..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 @@ -148,7 +148,7 @@ "SEGMENT_TILE_MAX": 64, "SEGMENT_BLOCK": 32 }, - "D_LEQ_128": { + "D_GEQ_256": { "num_stages": 1, "waves_per_eu": 2, "num_warps": 2, @@ -156,7 +156,8 @@ "MIN_SEGMENTS": 8, "SEGMENTS_PER_CU": 16, "MAX_SEGMENTS": 128, - "SEGMENT_TILE_MAX": 64 + "SEGMENT_TILE_MAX": 64, + "SEGMENT_BLOCK": 32 }, "any": { "num_stages": 1, @@ -166,8 +167,7 @@ "MIN_SEGMENTS": 8, "SEGMENTS_PER_CU": 16, "MAX_SEGMENTS": 128, - "SEGMENT_TILE_MAX": 64, - "SEGMENT_BLOCK": 32 + "SEGMENT_TILE_MAX": 64 } }, "kv_split": { From e422976a3fdc6d92421efce6a1e5ab803bf03c29 Mon Sep 17 00:00:00 2001 From: "Mehmet (Cagri) Kaymak" Date: Thu, 8 Oct 2026 20:53:11 +0000 Subject: [PATCH 4/9] [Triton/Gluon] [GFX950] Unified attention: move the split-grid query block bound into _gfx950_split_grid, drop the new tests --- .../ops/triton/attention/unified_attention.py | 7 +- .../attention/test_unified_attention.py | 165 ------------------ 2 files changed, 3 insertions(+), 169 deletions(-) diff --git a/aiter/ops/triton/attention/unified_attention.py b/aiter/ops/triton/attention/unified_attention.py index b750aad2787..b740f1f29de 100644 --- a/aiter/ops/triton/attention/unified_attention.py +++ b/aiter/ops/triton/attention/unified_attention.py @@ -59,9 +59,6 @@ # LDS the gfx950 KV double buffer may occupy; a wider tile falls back to one buffer. _GFX950_LDS_BUDGET = 160 * 1024 -# Widest query block the gfx950 split-KV grid takes for a multi-token batch. -_GFX950_SPLIT_MAX_BLOCK_M = 128 - def _is_gluon_available(): return any(supported in DEVICE_ARCH for supported in _GLUON_SUPPORTED_ARCHS) @@ -161,9 +158,11 @@ def _gfx950_split_grid(params: _UAParams): """ if params.all_decode: return True + # widest query block the split grid takes, the 2d grid's + max_block_m = 128 if ( params.sliding_window > 0 - or params.max_seqlen_q * params.num_queries_per_kv > _GFX950_SPLIT_MAX_BLOCK_M + or params.max_seqlen_q * params.num_queries_per_kv > max_block_m ): return False config = get_unified_attention_config("kv_split", params, backend="gluon") diff --git a/op_tests/triton_tests/attention/test_unified_attention.py b/op_tests/triton_tests/attention/test_unified_attention.py index cb9d310453e..ca14159e273 100644 --- a/op_tests/triton_tests/attention/test_unified_attention.py +++ b/op_tests/triton_tests/attention/test_unified_attention.py @@ -738,171 +738,6 @@ def test_triton_unified_attn( ) -# Decode split wide enough that reduce_segments merges the segments in chunks -# (reduce config SEGMENT_BLOCK < NUM_SEGMENTS). -@pytest.mark.parametrize( - "backend, head_size", [("triton", 512), ("triton", 256), ("gluon", 256)] -) -@torch.inference_mode() -def test_unified_attn_chunked_reduce(backend: str, head_size: int) -> None: - if DEVICE_ARCH != "gfx950": - pytest.skip("gfx950 reduce config") - seq_lens = [(1, 8192), (1, 1025), (1, 3000), (1, 37)] - ( - query, - key_cache_orig, - value_cache_orig, - key_cache, - value_cache, - sinks, - output, - cu_query_lens, - kv_lens, - max_query_len, - max_kv_len, - scale, - window_size, - block_tables, - *_, - ) = generate_data( - seq_lens=seq_lens, - num_blocks=1024, - block_size=64, - head_size=head_size, - num_heads=(16, 1), - device="cuda", - ) - unified_attention( - q=query, - k=key_cache, - v=value_cache, - out=output, - cu_seqlens_q=cu_query_lens, - seqused_k=kv_lens, - max_seqlen_q=max_query_len, - max_seqlen_k=max_kv_len, - softmax_scale=scale, - causal=True, - window_size=window_size, - block_table=block_tables, - softcap=0, - q_descale=None, - k_descale=None, - v_descale=None, - sinks=sinks, - backend=backend, - ) - ref_output = ref_paged_attn( - query=query, - key_cache=key_cache_orig, - value_cache=value_cache_orig, - query_lens=[x[0] for x in seq_lens], - kv_lens=[x[1] for x in seq_lens], - block_tables=block_tables, - scale=scale, - out_dtype=torch.bfloat16, - sinks=sinks, - ) - torch.testing.assert_close( - output.float(), ref_output.float(), atol=1.5e-2, rtol=1e-2 - ) - - -# Speculative-decode verify batches: short queries over a long context take the -# split-KV grid. kv_len 1025 puts the last key alone in a split that the first -# query of its sequence cannot see. -@pytest.mark.parametrize("num_heads", [(32, 2), (8, 1)]) -@pytest.mark.parametrize("head_size", [128, 256]) -@pytest.mark.parametrize( - "q_dtype, kv_dtype", - [(torch.bfloat16, torch.bfloat16), (e4m3_dtype, e4m3_dtype)], -) -@torch.inference_mode() -def test_gluon_unified_attn_spec_verify( - num_heads: tuple[int, int], - head_size: int, - q_dtype: torch.dtype, - kv_dtype: torch.dtype, -) -> None: - if DEVICE_ARCH != "gfx950": - pytest.skip("gfx950 gluon split-KV grid") - seq_lens = [(2, 1025), (4, 8192), (3, 3000), (1, 523)] - query_lens = [x[0] for x in seq_lens] - ( - query, - key_cache_orig, - value_cache_orig, - key_cache, - value_cache, - sinks, - output, - cu_query_lens, - kv_lens, - max_query_len, - max_kv_len, - scale, - window_size, - block_tables, - _maybe_quant_query, - _query_scales, - q_descale, - k_descale, - v_descale, - output_scale, - ) = generate_data( - seq_lens=seq_lens, - num_blocks=2048, - block_size=64, - head_size=head_size, - num_heads=num_heads, - q_dtype=q_dtype, - kv_dtype=kv_dtype, - device="cuda", - ) - for use_sinks in (False, True): - unified_attention( - q=query, - k=key_cache, - v=value_cache, - out=output, - cu_seqlens_q=cu_query_lens, - seqused_k=kv_lens, - max_seqlen_q=max_query_len, - max_seqlen_k=max_kv_len, - softmax_scale=scale, - causal=True, - window_size=window_size, - block_table=block_tables, - softcap=0, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - sinks=sinks if use_sinks else None, - backend="gluon", - ) - ref_output = ref_paged_attn( - query=query, - key_cache=key_cache_orig, - value_cache=value_cache_orig, - query_lens=query_lens, - kv_lens=[x[1] for x in seq_lens], - block_tables=block_tables, - scale=scale, - out_dtype=torch.bfloat16, - sliding_window=None, - soft_cap=None, - sinks=sinks if use_sinks else None, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - output_scale=output_scale, - ) - atol, rtol = (1.5e-1, 1.5e-1) if q_dtype == e4m3_dtype else (1.5e-2, 1e-2) - torch.testing.assert_close( - output.float(), ref_output.float(), atol=atol, rtol=rtol - ) - - @pytest.mark.parametrize( "seq_lens", [ From 6347babd3d4bffc4cac2e788e79ede0ce4603f4b Mon Sep 17 00:00:00 2001 From: "Mehmet (Cagri) Kaymak" Date: Thu, 8 Oct 2026 20:54:30 +0000 Subject: [PATCH 5/9] Shorten the unified attention comments --- .../gfx950/attention/unified_attention.py | 4 +--- .../attention/unified_attention.py | 6 ++---- aiter/ops/triton/attention/unified_attention.py | 15 ++++----------- 3 files changed, 7 insertions(+), 18 deletions(-) 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 36abad22ba0..fb13b52d5ce 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/unified_attention.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/unified_attention.py @@ -1432,9 +1432,7 @@ def store_partial( # row-max, so it needs whatever scale the softmax used. if not cfg.USE_SINKS: M = M * self.SM_scale - # A row with no causally visible key in this split (a short multi-token - # query whose split ends past the row's position) keeps the guarded 0 - # max; keep it out of the reduce's max. + # 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)) diff --git a/aiter/ops/triton/_triton_kernels/attention/unified_attention.py b/aiter/ops/triton/_triton_kernels/attention/unified_attention.py index ac86a6a3496..a1ccbbbff8e 100644 --- a/aiter/ops/triton/_triton_kernels/attention/unified_attention.py +++ b/aiter/ops/triton/_triton_kernels/attention/unified_attention.py @@ -1005,8 +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, - # segments merged per loop step; 0 merges all NUM_SEGMENTS_PER_SEQ in one tile - SEGMENT_BLOCK: tl.constexpr = 0, + 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) @@ -1075,8 +1074,7 @@ def reduce_segments( segm_output *= tl.math.exp2(segm_max - overall_max)[:, None] acc_sum = tl.sum(segm_output, axis=0) else: - # online merge over the written segments only, so the cost follows the - # context rather than NUM_SEGMENTS_PER_SEQ + # one big tile spills at large head sizes; this also reads only written segments segm_base = ( query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + query_head_idx * NUM_SEGMENTS_PER_SEQ diff --git a/aiter/ops/triton/attention/unified_attention.py b/aiter/ops/triton/attention/unified_attention.py index b740f1f29de..fb6bcd8f97c 100644 --- a/aiter/ops/triton/attention/unified_attention.py +++ b/aiter/ops/triton/attention/unified_attention.py @@ -148,23 +148,17 @@ def _gfx950_gluon_supported(params: _UAParams): def _gfx950_split_grid(params: _UAParams): - """Whether the gfx950 Gluon launch takes the split-capable (decode) grid. - - Decode does. So does a batch of short queries such as a speculative decoding - verify step, when one query block can hold every query of a sequence and the - launch is small enough to split: each block then walks the whole context, and - the KV split keeps a small batch from leaving most CUs idle. Unsplit, the 2d - grid's wider query block is faster. - """ + """Decode and short-query batches (e.g. spec-decode verify) take the split grid.""" if params.all_decode: return True - # widest query block the split grid takes, the 2d grid's + # one query block has to cover a sequence's queries; 128 is the 2d BLOCK_M max_block_m = 128 if ( params.sliding_window > 0 or params.max_seqlen_q * params.num_queries_per_kv > max_block_m ): return False + # unsplit, the 2d grid is faster config = get_unified_attention_config("kv_split", params, backend="gluon") return config["NUM_SEGMENTS"] > 1 @@ -1153,8 +1147,7 @@ def _unified_attention_3d_gfx950( MFMA_DIM = config["MFMA_DIM"] BLOCK_M = max(config["BLOCK_M"], triton.next_power_of_2(params.num_queries_per_kv)) if not params.all_decode: - # one query block per sequence: the split carves the KV range from - # seq_len, which is what reduce_segments assumes + # one query block per sequence, so the split ranges match reduce_segments BLOCK_M = max( BLOCK_M, triton.next_power_of_2(params.max_seqlen_q * params.num_queries_per_kv), From 684eb1e51f2a91cd49a85985bb6f85a0f1d10a16 Mon Sep 17 00:00:00 2001 From: "Mehmet (Cagri) Kaymak" Date: Thu, 8 Oct 2026 20:55:28 +0000 Subject: [PATCH 6/9] Drop a comment in reduce_segments --- aiter/ops/triton/_triton_kernels/attention/unified_attention.py | 1 - 1 file changed, 1 deletion(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/unified_attention.py b/aiter/ops/triton/_triton_kernels/attention/unified_attention.py index a1ccbbbff8e..a17cd2b7f98 100644 --- a/aiter/ops/triton/_triton_kernels/attention/unified_attention.py +++ b/aiter/ops/triton/_triton_kernels/attention/unified_attention.py @@ -1074,7 +1074,6 @@ def reduce_segments( segm_output *= tl.math.exp2(segm_max - overall_max)[:, None] acc_sum = tl.sum(segm_output, axis=0) else: - # one big tile spills at large head sizes; this also reads only written segments segm_base = ( query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + query_head_idx * NUM_SEGMENTS_PER_SEQ From 8a508194ff3a69b50c0ef112982adf6b39f0c0af Mon Sep 17 00:00:00 2001 From: "Mehmet (Cagri) Kaymak" Date: Thu, 8 Oct 2026 22:39:32 +0000 Subject: [PATCH 7/9] Tune gfx950 gluon unified attention configs for fp8 2d and short-query splits --- .../attention/unified_attention/DEFAULT.json | 66 +++++++++++++++---- 1 file changed, 52 insertions(+), 14 deletions(-) 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 4ecfae798df..fa2b5ce6ad7 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,12 @@ "attn_3d": [ "D", "SW", + "Q", "DT" ], "kv_split": [ "SW", + "D", "Q", "DT" ] @@ -46,20 +48,20 @@ "TILE_SIZE_MAX": 64 }, "D_GEQ_256.SKLT.DT_fp8_fp8": { - "BLOCK_M": 128, - "num_warps": 4, + "BLOCK_M": 64, + "num_warps": 2, "MFMA_DIM": 32, - "NUM_BUFFERS": 1, - "waves_per_eu": 1, + "NUM_BUFFERS": 2, + "waves_per_eu": 0, "TILE_SIZE_MIN": 64, "TILE_SIZE_MAX": 64 }, "D_GEQ_256.DT_fp8_fp8": { - "BLOCK_M": 128, - "num_warps": 4, + "BLOCK_M": 64, + "num_warps": 2, "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 }, @@ -123,12 +125,24 @@ "waves_per_eu": 1, "WIDE_BLOCK_M_MULT": 2 }, + "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, @@ -150,7 +164,7 @@ "SEGMENTS_PER_CU": 4, "MFMA_DIM": 32, "MIN_SEGMENTS": 1, - "SPLIT_MIN_TILES": 16, + "SPLIT_MIN_TILES": 4, "SPLIT_MIN_SHARE": 3, "SEGMENT_TILE_MIN": 64, "SEGMENT_TILE_MAX": 64, @@ -158,12 +172,24 @@ "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 + }, "DT_fp8_fp8": { "SEGMENTS_PER_CU": 4, "MFMA_DIM": 32, "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, @@ -174,7 +200,7 @@ "SEGMENTS_PER_CU": 2, "MFMA_DIM": 16, "MIN_SEGMENTS": 1, - "SPLIT_MIN_TILES": 16, + "SPLIT_MIN_TILES": 4, "SPLIT_MIN_SHARE": 3, "SEGMENT_TILE_MIN": 64, "SEGMENT_TILE_MAX": 64, @@ -182,12 +208,24 @@ "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 + }, "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, From 07d82b0f5b49a6e2bb2ce268d9568467b03dacce Mon Sep 17 00:00:00 2001 From: "Mehmet (Cagri) Kaymak" Date: Fri, 9 Oct 2026 00:27:24 +0000 Subject: [PATCH 8/9] Split the KV of gfx950 gluon prefill when the grid is too small to fill the GPU --- .../gfx950/attention/unified_attention.py | 33 +++++-- .../ops/triton/attention/unified_attention.py | 43 +++++---- .../attention/unified_attention/DEFAULT.json | 88 ++++++++++++++++++- 3 files changed, 138 insertions(+), 26 deletions(-) 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 fb13b52d5ce..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) @@ -1906,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/attention/unified_attention.py b/aiter/ops/triton/attention/unified_attention.py index fb6bcd8f97c..b6b62ba2001 100644 --- a/aiter/ops/triton/attention/unified_attention.py +++ b/aiter/ops/triton/attention/unified_attention.py @@ -147,19 +147,27 @@ 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 short-query batches (e.g. spec-decode verify) take the split grid.""" + """Decode, and batches too small to fill the GPU, take the split grid.""" if params.all_decode: return True - # one query block has to cover a sequence's queries; 128 is the 2d BLOCK_M - max_block_m = 128 - if ( - params.sliding_window > 0 - or params.max_seqlen_q * params.num_queries_per_kv > max_block_m - ): + if params.sliding_window > 0: return False # unsplit, the 2d grid is faster - config = get_unified_attention_config("kv_split", params, backend="gluon") + config = get_unified_attention_config( + "kv_split", _gfx950_split_params(params), backend="gluon" + ) return config["NUM_SEGMENTS"] > 1 @@ -169,7 +177,7 @@ def use_2d_kernel(params: _UAParams, backend: str = "triton"): 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 and short queries take the 3d path + # 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 _gfx950_split_grid(params) @@ -384,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 @@ -1137,7 +1148,7 @@ def _unified_attention_3d_gfx950( NUM_SEGMENTS, TILE_SIZE, ): - """Decode and short queries: 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. @@ -1146,12 +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)) - if not params.all_decode: - # one query block per sequence, so the split ranges match reduce_segments - BLOCK_M = max( - BLOCK_M, - triton.next_power_of_2(params.max_seqlen_q * 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/gfx950/gluon/attention/unified_attention/DEFAULT.json b/aiter/ops/triton/configs/gfx950/gluon/attention/unified_attention/DEFAULT.json index fa2b5ce6ad7..f3d23b8c08b 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 @@ -43,7 +43,7 @@ "num_warps": 4, "MFMA_DIM": 32, "NUM_BUFFERS": 2, - "waves_per_eu": 3, + "waves_per_eu": 0, "TILE_SIZE_MIN": 64, "TILE_SIZE_MAX": 64 }, @@ -79,12 +79,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, @@ -111,6 +123,18 @@ "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, @@ -125,6 +149,18 @@ "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, @@ -172,6 +208,18 @@ "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, @@ -184,6 +232,18 @@ "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, @@ -208,6 +268,18 @@ "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, @@ -220,6 +292,18 @@ "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, From 1ca41053a0a3ba8fa64866017a20b0d7764f2f03 Mon Sep 17 00:00:00 2001 From: "Mehmet (Cagri) Kaymak" Date: Fri, 9 Oct 2026 01:23:13 +0000 Subject: [PATCH 9/9] Pick gfx950 gluon fp8 D=256 2d and D=64 prefill split configs that hold on Triton 3.7 --- .../attention/unified_attention/DEFAULT.json | 37 +++++++++++++------ 1 file changed, 26 insertions(+), 11 deletions(-) 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 f3d23b8c08b..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 @@ -47,18 +47,9 @@ "TILE_SIZE_MIN": 64, "TILE_SIZE_MAX": 64 }, - "D_GEQ_256.SKLT.DT_fp8_fp8": { - "BLOCK_M": 64, - "num_warps": 2, - "MFMA_DIM": 32, - "NUM_BUFFERS": 2, - "waves_per_eu": 0, - "TILE_SIZE_MIN": 64, - "TILE_SIZE_MAX": 64 - }, "D_GEQ_256.DT_fp8_fp8": { - "BLOCK_M": 64, - "num_warps": 2, + "BLOCK_M": 128, + "num_warps": 4, "MFMA_DIM": 32, "NUM_BUFFERS": 2, "waves_per_eu": 0, @@ -208,6 +199,18 @@ "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, @@ -268,6 +271,18 @@ "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,