[Triton/Gluon] [FlyDSL] [CI] feat(flydsl): Add paged-attention Tile kernel

.github/scripts/aiter_test.sh

@@ -124,6 +124,20 @@ for file in "${sharded_files[@]}"; do
             } | tee -a latest_test.log
             test_cmd=(env AITER_MLA_DECODE_PERSISTENT_MAX_BATCH=0 timeout 60m python3 "$file")
             ;;
+        op_tests/test_flydsl_pa_decode.py)
+            # The CLI sweep does not collect parametrized pytest regressions.
+            test_cmd=(
+                timeout 60m
+                bash -c '
+                    set -euo pipefail
+                    test_file=$1
+                    python3 -m pytest -q \
+                        "${test_file}::test_large_negative_logits_preserve_online_softmax"
+                    python3 "$test_file"
+                '
+                _ "$file"
+            )
+            ;;
         op_tests/test_gemm_a6w6.py)
             {
                 echo "Running tuned dispatch plus every compatible A6W6 ASM kernel"

aiter/ops/attention.py

@@ -21,14 +21,76 @@
 from ..jit.core import compile_ops, is_experimental_enabled
 from ..jit.utils.chip_info import get_cu_num, get_gfx
 
+try:
+    from aiter.ops.flydsl import pa_decode as _pa_decode_flydsl
+except (ImportError, AttributeError, RuntimeError, OSError):
+    _pa_decode_flydsl = None
+
 MD_NAME = "module_attention"
 
+
+def pa_decode_flydsl(
+    output: torch.Tensor,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    context_lengths: torch.Tensor,
+    block_tables: torch.Tensor,
+    softmax_scale: float,
+    query_length: int,
+    max_context_partition_num: int,
+    context_partition_size: int = 256,
+    compute_type: torch.dtype = torch.bfloat16,
+    query_scale: torch.Tensor = None,
+    key_scale: torch.Tensor = None,
+    value_scale: torch.Tensor = None,
+    exp_sums: torch.Tensor = None,
+    max_logits: torch.Tensor = None,
+    temporary_output: torch.Tensor = None,
+    alibi_slopes: torch.Tensor = None,
+    sinks: torch.Tensor = None,
+    sliding_window: int = 0,
+    ps: bool = True,
+) -> None:
+    if _pa_decode_flydsl is None:
+        raise RuntimeError("pa_decode_flydsl requires the `flydsl` package")
+    _pa_decode_flydsl(
+        output,
+        query,
+        key_cache,
+        value_cache,
+        context_lengths,
+        block_tables,
+        softmax_scale,
+        query_length,
+        max_context_partition_num,
+        context_partition_size,
+        compute_type,
+        query_scale,
+        key_scale,
+        value_scale,
+        exp_sums,
+        max_logits,
+        temporary_output,
+        alibi_slopes,
+        sinks,
+        sliding_window,
+        ps,
+    )
+
+
 direct_register_custom_op(
     "pa_decode_gluon",
     pa_decode_gluon,
     ["output", "exp_sums", "max_logits", "temporary_output"],
 )
 
+direct_register_custom_op(
+    "pa_decode_flydsl",
+    pa_decode_flydsl,
+    ["output", "exp_sums", "max_logits", "temporary_output"],
+)
+
 
 def gen_pa_fwd_native_fake(
     # [num_seqs, num_heads, head_size]

aiter/ops/flydsl/__init__.py

@@ -82,6 +82,7 @@
         ".gather_kv_b_proj",
         "gather_kv_b_proj_flydsl_supported",
     ),
+    "pa_decode": (".pa_decode", "pa_decode"),
 }
 
 __all__ = [
@@ -104,6 +105,7 @@
     "flydsl_qk_norm_rope_quant",
     "gather_kv_b_proj_flydsl",
     "gather_kv_b_proj_flydsl_supported",
+    "pa_decode",
 ]
 
 

aiter/ops/flydsl/kernels/dpp_utils.py

@@ -55,3 +55,26 @@ def update_dpp_i32(
         [],
         **kw,
     )
+
+
+def dpp_xor_f32(src, offset: int, **kw):
+    """Return ``src`` from the lane selected by a 16-lane XOR DPP pattern."""
+    from flydsl._mlir.dialects import arith as _arith_dialect
+    from flydsl.expr.typing import T
+
+    src_i32 = _to_ir(src).bitcast(T.i32)
+    if offset == 8:
+        out_i32 = update_dpp_i32(src_i32, src_i32, 280, 0xF, 0xC, False, **kw)
+        out_i32 = update_dpp_i32(out_i32, src_i32, 264, 0xF, 0x3, False, **kw)
+    elif offset == 4:
+        out_i32 = update_dpp_i32(src_i32, src_i32, 276, 0xF, 0xA, False, **kw)
+        out_i32 = update_dpp_i32(out_i32, src_i32, 260, 0xF, 0x5, False, **kw)
+    elif offset == 2:
+        out_i32 = update_dpp_i32(src_i32, src_i32, 78, 0xF, 0xF, False, **kw)
+    elif offset == 1:
+        out_i32 = update_dpp_i32(src_i32, src_i32, 177, 0xF, 0xF, False, **kw)
+    else:
+        raise ValueError(
+            f"dpp_xor_f32 only supports 16-lane offsets 1, 2, 4, 8; got {offset}"
+        )
+    return _arith_dialect.BitcastOp(T.f32, out_i32).result

aiter/ops/flydsl/kernels/pa_decode_reduce.py

@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (C) 2025-2026 FlyDSL Project Contributors
+
+"""Partitioned-softmax reduction kernel for FlyDSL paged attention."""
+
+from functools import lru_cache
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl.expr.typing import T
+
+
+@lru_cache(maxsize=256)
+def compile_pa_decode_ps_reduce(
+    *,
+    max_context_partition_num: int,
+    head_size: int,
+    output_dtype_str: str,
+    logits_dtype_str: str,
+    sink_dtype_str: str,
+    use_sinks: bool,
+):
+    """Build the partitioned-softmax reduction used by ``pa_decode``.
+
+    Each wave reduces the partition statistics independently, then reuses the
+    per-partition weights to combine one output element per thread.  The host
+    API limits the partition count to one wave, so the LDS path from the more
+    general Gluon implementation is unnecessary here.
+    """
+    if not 1 <= max_context_partition_num <= 64:
+        raise ValueError(
+            "max_context_partition_num must be in [1, 64], "
+            f"got {max_context_partition_num}"
+        )
+    if head_size <= 0 or head_size > 1024 or head_size % 64:
+        raise ValueError(
+            f"head_size must be a multiple of 64 in [64, 1024], got {head_size}"
+        )
+
+    dtype_map = {
+        "f32": fx.Float32,
+        "f16": fx.Float16,
+        "bf16": fx.BFloat16,
+    }
+    try:
+        output_dtype = dtype_map[output_dtype_str]
+        logits_dtype = dtype_map[logits_dtype_str]
+        sink_dtype = dtype_map[sink_dtype_str]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported FlyDSL dtype: {exc.args[0]!r}") from exc
+
+    warp_size = 64
+    log2e = 1.4426950408889634
+    reduce_width = (
+        1
+        if max_context_partition_num == 1
+        else 1 << ((max_context_partition_num - 1).bit_length())
+    )
+    reduce_shuffle_offsets = [
+        offset for offset in (32, 16, 8, 4, 2, 1) if offset < reduce_width
+    ]
+
+    @flyc.kernel(known_block_size=[head_size, 1, 1])
+    def pa_decode_ps_reduce_kernel(
+        output_ptr: fx.Pointer,
+        exp_sums_ptr: fx.Pointer,
+        max_logits_ptr: fx.Pointer,
+        logits_ptr: fx.Pointer,
+        sink_token_ptr: fx.Pointer,
+        stride_output_bs: fx.Int32,
+        stride_output_len: fx.Int32,
+        stride_output_kv_head: fx.Int32,
+        stride_output_group_size: fx.Int32,
+        stride_exp_sums_seq: fx.Int32,
+        stride_exp_sums_head: fx.Int32,
+        stride_exp_sums_part: fx.Int32,
+        stride_logits_seq: fx.Int32,
+        stride_logits_head: fx.Int32,
+        stride_logits_part: fx.Int32,
+        stride_logits_group: fx.Int32,
+        query_group_size: fx.Int32,
+    ):
+        tid = fx.thread_idx.x
+        batch_idx = fx.block_idx.x
+        kv_head_idx = fx.block_idx.y
+        eqgs_idx = fx.block_idx.z
+
+        output = fx.recast_iter(output_dtype, output_ptr)
+        exp_sums = fx.recast_iter(fx.Float32, exp_sums_ptr)
+        max_logits = fx.recast_iter(fx.Float32, max_logits_ptr)
+        logits = fx.recast_iter(logits_dtype, logits_ptr)
+        if fx.const_expr(use_sinks):
+            sink_token = fx.recast_iter(sink_dtype, sink_token_ptr)
+
+        zero_f = fx.Float32(0.0)
+        one_f = fx.Float32(1.0)
+        neg_inf = fx.Float32(float("-inf"))
+        c_log2e = fx.Float32(log2e)
+        zero_i = fx.Int32(0)
+        c_warp_size = fx.Int32(warp_size)
+        c_wave_mask = fx.Int32(warp_size - 1)
+        c_part_num = fx.Int32(max_context_partition_num)
+        c_reduce_width = fx.Int32(reduce_width)
+        c_four = fx.Int32(4)
+        c_qgs = query_group_size
+        lane = tid & c_wave_mask
+        group_idx = eqgs_idx % c_qgs
+
+        def _wave_reduce_max(value):
+            reduced = value
+            for offset in reduce_shuffle_offsets:
+                reduced = reduced.maximumf(
+                    reduced.shuffle_xor(fx.Int32(offset), c_warp_size)
+                )
+            return reduced
+
+        def _wave_reduce_sum(value):
+            reduced = value
+            for offset in reduce_shuffle_offsets:
+                reduced = fx.Float32(
+                    reduced.addf(
+                        reduced.shuffle_xor(fx.Int32(offset), c_warp_size),
+                        fastmath="fast",
+                    )
+                )
+            return reduced
+
+        lane_in_range = lane < c_part_num
+        lane_in_reduce = lane < c_reduce_width
+        part_sum = zero_f
+        part_max = neg_inf
+        if lane_in_reduce:
+            part_idx = lane_in_range.select(lane, zero_i)
+            stats_offset = (
+                batch_idx * stride_exp_sums_seq
+                + kv_head_idx * stride_exp_sums_head
+                + part_idx * stride_exp_sums_part
+                + eqgs_idx
+            )
+            loaded_sum = fx.Float32(exp_sums[stats_offset])
+            loaded_max = fx.Float32(max_logits[stats_offset])
+            part_sum = lane_in_range.select(loaded_sum, zero_f)
+            part_max = lane_in_range.select(loaded_max, neg_inf)
+
+        global_max = _wave_reduce_max(part_max)
+        safe_global_max = (global_max > neg_inf).select(global_max, zero_f)
+        part_scale = (part_max > neg_inf).select(
+            fx.exp2((part_max - safe_global_max) * c_log2e, fastmath="fast"),
+            zero_f,
+        )
+        scaled_sum = part_sum * part_scale
+        global_exp_sum = _wave_reduce_sum(scaled_sum)
+        if fx.const_expr(use_sinks):
+            sink_value = fx.Float32(sink_token[kv_head_idx * c_qgs + group_idx])
+            sink_scale = (global_max > neg_inf).select(
+                fx.exp2(
+                    (sink_value - safe_global_max) * c_log2e,
+                    fastmath="fast",
+                ),
+                zero_f,
+            )
+            global_exp_sum = global_exp_sum + sink_scale
+        safe_global_exp_sum = (global_exp_sum > zero_f).select(global_exp_sum, one_f)
+        weight_local = scaled_sum / safe_global_exp_sum
+        weight_local_i32 = weight_local.bitcast(fx.Int32)
+
+        acc = zero_f
+        for part_idx in fx.range_constexpr(max_context_partition_num):
+            c_part_idx = fx.Int32(part_idx)
+            weight_i32 = fx.Int32(
+                fx.rocdl.ds_bpermute(
+                    T.i32,
+                    c_part_idx * c_four,
+                    weight_local_i32,
+                )
+            )
+            weight = weight_i32.bitcast(fx.Float32)
+            logits_offset = (
+                batch_idx * stride_logits_seq
+                + kv_head_idx * stride_logits_head
+                + c_part_idx * stride_logits_part
+                + eqgs_idx * stride_logits_group
+                + tid
+            )
+            part_logits = fx.Float32(logits[logits_offset])
+            acc = acc + part_logits * weight
+
+        query_idx = eqgs_idx // c_qgs
+        output_offset = (
+            batch_idx * stride_output_bs
+            + query_idx * stride_output_len
+            + kv_head_idx * stride_output_kv_head
+            + group_idx * stride_output_group_size
+            + tid
+        )
+        output[output_offset] = acc.to(output_dtype)
+
+    @flyc.jit
+    def launch_pa_decode_ps_reduce_kernel(
+        output: fx.Pointer,
+        exp_sums: fx.Pointer,
+        max_logits: fx.Pointer,
+        logits: fx.Pointer,
+        sink_token: fx.Pointer,
+        stride_output_bs: fx.Int32,
+        stride_output_len: fx.Int32,
+        stride_output_kv_head: fx.Int32,
+        stride_output_group_size: fx.Int32,
+        stride_exp_sums_seq: fx.Int32,
+        stride_exp_sums_head: fx.Int32,
+        stride_exp_sums_part: fx.Int32,
+        stride_logits_seq: fx.Int32,
+        stride_logits_head: fx.Int32,
+        stride_logits_part: fx.Int32,
+        stride_logits_group: fx.Int32,
+        query_seq_len: fx.Int32,
+        query_group_size: fx.Int32,
+        batch_size: fx.Int32,
+        num_kv_heads: fx.Int32,
+        stream: fx.Stream,
+    ):
+        pa_decode_ps_reduce_kernel(
+            output,
+            exp_sums,
+            max_logits,
+            logits,
+            sink_token,
+            stride_output_bs,
+            stride_output_len,
+            stride_output_kv_head,
+            stride_output_group_size,
+            stride_exp_sums_seq,
+            stride_exp_sums_head,
+            stride_exp_sums_part,
+            stride_logits_seq,
+            stride_logits_head,
+            stride_logits_part,
+            stride_logits_group,
+            query_group_size,
+        ).launch(
+            grid=(batch_size, num_kv_heads, query_seq_len * query_group_size),
+            block=(head_size, 1, 1),
+            stream=stream,
+        )
+
+    return {
+        "launch": launch_pa_decode_ps_reduce_kernel,
+        "kernel": pa_decode_ps_reduce_kernel,
+    }