[Triton/Gluon] Improving GFX950 Unified Attention for spec. decode and reduce opimizations

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

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,45 +1027,87 @@ 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:
         v_dim_mask = offs_vd < V_HEAD_SIZE
     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)
 

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:

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,