[Triton/Gluon] [FlyDSL] [CI] Refactor PA decode, direct native autotune, and BF16 KV

.github/scripts/aiter_test.sh

@@ -126,17 +126,7 @@ for file in "${sharded_files[@]}"; do
             test_cmd=(env AITER_MLA_DECODE_PERSISTENT_MAX_BATCH=0 timeout 60m python3 "$file")
             ;;
         op_tests/test_flydsl_pa_decode.py)
-            # The CLI sweep is separate from the compact parametrized regression.
-            test_cmd=(
-                timeout 60m
-                bash -c '
-                    set -euo pipefail
-                    test_file=$1
-                    python3 -m pytest -q "$test_file"
-                    python3 "$test_file"
-                '
-                _ "$file"
-            )
+            test_cmd=(timeout 60m python3 -m pytest -q "$file")
             ;;
         op_tests/test_gemm_a6w6.py)
             {

aiter/ops/flydsl/__init__.py

@@ -98,6 +98,7 @@
         "gather_kv_b_proj_flydsl_supported",
     ),
     "pa_decode": (".pa_decode", "pa_decode"),
+    "prepare_pa_decode_plan": (".pa_decode_tuning", "prepare_pa_decode_plan"),
     "gather_kv_b_proj_flydsl_fp8_supported": (
         ".gather_kv_b_proj",
         "gather_kv_b_proj_flydsl_fp8_supported",
@@ -130,6 +131,7 @@
     "gather_kv_b_proj_flydsl_fp8_supported",
     "gather_kv_b_proj_flydsl_supported",
     "pa_decode",
+    "prepare_pa_decode_plan",
 ]
 
 _fused_moe_impl_path = "aiter.ops.flydsl.fused_moe_gfx942:run_flydsl_moe_gfx942_impl"

aiter/ops/flydsl/configs/pa/README.md

@@ -0,0 +1,47 @@
+# FlyDSL paged-attention tuning results
+
+`prepare_pa_decode_plan` in `aiter.ops.flydsl.pa_decode_tuning` stores and loads
+native FlyDSL autotune results in the directory for the GPU running the query:
+
+```text
+aiter/ops/flydsl/configs/pa/<arch>/_pa_decode_autotuner.json
+```
+
+For example, MI300/MI308 results use `gfx942/` and MI350/MI355 results use
+`gfx950/`. FlyDSL creates the directory and JSON file when tuning saves a result.
+Preparing a plan without a cached result uses the `2*CU` default and does not
+create a file.
+
+Run the standalone tuning script from the repository root or an installed AITER
+package. It constructs the requested shape, benchmarks native candidates, and
+saves the winner without requiring an autotune environment variable:
+
+```bash
+python -m aiter.ops.flydsl.pa_decode_tuning \
+    --batch-size 32 --num-query-heads 16 --num-kv-heads 1 --head-dim 128 \
+    --query-length 4 --context-length 16384 --block-size 128 \
+    --kv-dtype fp8 --scale-mode per-token --trans-v 1 \
+    --output-dir /tmp/pa-configs
+```
+
+This saves `/tmp/pa-configs/<arch>/_pa_decode_autotuner.json`. Omit `--output-dir`
+to use the package directory shown above, or the environment override described
+below. Each invocation searches the requested shape and retains cached results
+for other shapes. `--context-length` accepts one length for all batch rows or one
+length per row; lengths include query tokens. Use `--kv-dtype bf16` for unscaled
+BF16 K/V and `--help` for all shape, layout, window, sink, and device options.
+
+Python callers can import `prepare_pa_decode_plan` from
+`aiter.ops.flydsl.pa_decode_tuning` (or `aiter.ops.flydsl`) and pass their actual
+tensors with `tune=True, config_dir="/tmp/pa-configs"` before inference or graph
+capture. Subsequent calls with the same `config_dir` use saved results by default.
+`FLYDSL_AUTOTUNE=1` remains available to force searches through the Python API.
+FlyDSL owns the JSON format, candidate benchmarking, selection, and cache
+persistence. Cache keys retain attention geometry, execution modes, CU count,
+PA source contents, and native toolchain/environment/device fingerprints.
+
+To use a custom cache directory, set `FLYDSL_AUTOTUNE_CACHE_DIR` before importing
+the PA tuning module. This overrides the full cache directory, so the file is
+saved as `<FLYDSL_AUTOTUNE_CACHE_DIR>/_pa_decode_autotuner.json`. Use separate
+directories for parallel tuning processes. Restart inference processes after
+updating results on disk, because native tuners keep loaded configurations in memory.

aiter/ops/flydsl/kernels/pa_decode/__init__.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (C) 2025-2026 FlyDSL Project Contributors
+
+"""Object-oriented building blocks for FP8/BF16 paged-attention kernels."""
+
+from functools import cache
+from hashlib import sha256
+from pathlib import Path
+
+
+@cache
+def implementation_cache_tag():
+    """Cover composed PA sources in compilation and autotuning caches.
+
+    FlyDSL follows function closures, but does not traverse the classes held by
+    a composed pipeline. Hash their sources once per process, including the
+    shared helpers, kernel, planner, reducer and host wrapper, so changes
+    invalidate cached binaries and tuning results. Paths are relative to
+    flydsl/ so the key is stable across installations.
+    """
+    package = Path(__file__).parent
+    shared_helpers = (
+        "buffer_ops.py",
+        "dpp_utils.py",
+        "kernels_common.py",
+        "pa_decode_kernel.py",
+        "pa_decode_plan.py",
+        "pa_decode_reduce.py",
+        "tensor_shim.py",
+        "utils.py",
+    )
+    sources = (
+        sorted(package.glob("*.py"))
+        + [package.parent / name for name in shared_helpers]
+        + [package.parent.parent / "pa_decode.py"]
+    )
+    digest = sha256()
+    for source in sources:
+        digest.update(source.relative_to(package.parent.parent).as_posix().encode())
+        digest.update(source.read_bytes())
+    return digest.hexdigest()

aiter/ops/flydsl/kernels/pa_decode/context.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (C) 2025-2026 FlyDSL Project Contributors
+
+"""Per-CTA operands, thread ownership and partition bounds for PA decode.
+
+Initialization stays explicit so the pipeline controls the order of global loads.
+Loop-carried softmax and output values are passed as SSA values, not stored here.
+"""
+
+import flydsl.expr as fx
+from flydsl.expr import arith, const_expr, gpu, range_constexpr
+
+from .traits import LOG2E, MFMA_MNK, WAVE
+
+
+class PaDecodeContext:
+    def __init__(
+        self,
+        traits,
+        shared_storage,
+        pmax_ptr,
+        psum_ptr,
+        pout_ptr,
+        query_ptr,
+        key_cache_ptr,
+        value_cache_ptr,
+        block_tables_ptr,
+        key_scale_ptr,
+        value_scale_ptr,
+        max_blocks_per_seq,
+        stride_ks_block,
+        stride_ks_head,
+        stride_q_row,
+        stride_q_head,
+        num_sequences,
+        planned_seq,
+        planned_start,
+        planned_end,
+        planned_context,
+    ):
+        self.traits = traits
+        self.shared_storage = shared_storage
+        self.pmax_ptr = pmax_ptr
+        self.psum_ptr = psum_ptr
+        self.pout_ptr = pout_ptr
+        self.query_ptr = query_ptr
+        self.key_cache_ptr = key_cache_ptr
+        self.value_cache_ptr = value_cache_ptr
+        self.block_tables_ptr = block_tables_ptr
+        self.key_scale_ptr = key_scale_ptr
+        self.value_scale_ptr = value_scale_ptr
+        self.max_blocks_per_seq = max_blocks_per_seq
+        self.stride_ks_block = stride_ks_block
+        self.stride_ks_head = stride_ks_head
+        self.stride_q_row = stride_q_row
+        self.stride_q_head = stride_q_head
+        self.num_sequences = num_sequences
+        self.planned_seq = planned_seq
+        self.planned_start = planned_start
+        self.planned_end = planned_end
+        self.planned_context = planned_context
+
+    def init_thread_mapping(self):
+        """Resolve the physical grid to a sequence, KV head and partial slot."""
+        self.tid = fx.Int32(gpu.thread_id("x"))
+        self.warp = self.tid // WAVE  # 0..NWARP-1
+        self.lane = self.tid - self.warp * WAVE  # 0..63
+        self.seq = self.planned_seq
+        self.kv_query = fx.Int32(gpu.block_id("y"))
+        self.kv_h = self.kv_query // self.traits.query_splits
+        self.query_begin = (
+            self.kv_query % self.traits.query_splits
+        ) * self.traits.QUERIES_PER_CTA
+        self.n_kv = fx.Int32(gpu.grid_dim.y) // self.traits.query_splits
+        if const_expr(self.traits.batch_first_plan_grid):
+            self.part = fx.Int32(
+                fx.Uint32(gpu.block_id("x")) * fx.Uint32(gpu.grid_dim.z)
+                + fx.Uint32(gpu.block_id("z"))
+            )
+            # Physical x groups packed slots, not sequences.
+            capacity = fx.Int32(fx.Uint32(gpu.grid_dim.x) * fx.Uint32(gpu.grid_dim.z))
+            self.partial_slot = self.kv_h * capacity + self.part
+        else:
+            self.part = fx.Int32(gpu.block_id("x"))
+            self.partial_slot = self.kv_h * fx.Int32(gpu.grid_dim.x) + self.part
+
+    def init_query_mapping(self):
+        self.rgroup = (
+            self.lane // MFMA_MNK
+        )  # 0..3: quarter-wave (paired with warp -> query row)
+        self.lane16 = (
+            self.lane - self.rgroup * MFMA_MNK
+        )  # 0..15: this row's head-dim chunk index
+        self.qh_local = (
+            self.warp * 4 + self.rgroup
+        )  # 0..15: this thread's query row within an M-tile
+
+    def init_context_length(self):
+        self.context_len = self.planned_context
+
+    def init_partition_bounds(self):
+        self.num_pages = (
+            self.context_len + self.traits.block_size - 1
+        ) // self.traits.block_size  # pages this sequence really owns
+        self.part_start = self.planned_start
+        self.part_end = self.planned_end
+
+    def init_scales(self, key_scale, value_scale):
+        if const_expr(self.traits.per_token_kv or self.traits.is_bf16_kv):
+            self.scale_qk = fx.Float32(self.traits.softmax_scale * LOG2E)
+        else:
+            self.scale_qk = fx.Float32(self.traits.softmax_scale * LOG2E) * fx.Float32(
+                key_scale
+            )
+            self.v_scale_f = fx.Float32(value_scale)
+        self.NEG_INF = fx.Float32(float("-inf"))
+        self.ZERO_F = fx.Float32(0.0)
+        # Finite or -inf scores permit nnan's bare max instructions.
+        # Do not set ninf: -inf is the mask sentinel.
+        self.fm_nnan = arith.FastMathFlags.nnan
+
+    def init_causal_offsets(self):
+        # Distance to the newest query keeps causal bounds tile-relative.
+        if const_expr(self.traits.QUERIES_PER_CTA == 1):
+            self.causal_offset = [
+                self.traits.query_length - 1 - self.query_begin
+                for _m in range_constexpr(self.traits.M_TILES)
+            ]
+        else:
+            self.causal_offset = [
+                self.traits.query_length
+                - 1
+                - self.query_begin
+                - (m * MFMA_MNK + self.lane16) // self.traits.query_group_size
+                for m in range_constexpr(self.traits.M_TILES)
+            ]
+            if const_expr(self.traits.CTA_ROWS % MFMA_MNK != 0):
+                # Padded rows still need nonnegative causal offsets.
+                self.causal_offset = [
+                    (offset > 0).select(offset, 0) for offset in self.causal_offset
+                ]

aiter/ops/flydsl/kernels/pa_decode/op_epilog.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (C) 2025-2026 FlyDSL Project Contributors
+
+"""Normalize PA output and store packed partition partials."""
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl.expr import const_expr, gpu, range_constexpr
+
+from ..tensor_shim import buf_base_i64, buf_copy_store, ptr_buf_tensor
+from ..utils import rcp_f32
+from .traits import LOG2E, MFMA_MNK
+
+
+class PaDecodeEpilogue:
+    def __init__(self, ctx):
+        self.ctx = ctx
+        self.traits = ctx.traits
+
+    def init_outputs(self):
+        self.pmax = fx.recast_iter(fx.Float32, self.ctx.pmax_ptr)
+        self.psum = fx.recast_iter(fx.Float32, self.ctx.psum_ptr)
+        self.pout = fx.recast_iter(self.traits.Q_DTYPE, self.ctx.pout_ptr)
+        if const_expr(self.traits.buffer_plan_output):
+            # Widen before multiplying: packed storage can exceed 2 GiB.
+            if const_expr(self.traits.batch_first_plan_grid):
+                pout_slot = fx.Int64(self.ctx.kv_h) * fx.Int64(
+                    gpu.grid_dim.x
+                ) * fx.Int64(gpu.grid_dim.z) + fx.Int64(self.ctx.part)
+            else:
+                pout_slot = fx.Int64(self.ctx.kv_h) * fx.Int64(
+                    gpu.grid_dim.x
+                ) + fx.Int64(self.ctx.part)
+            pout_slot_bytes = self.traits.TOTAL_ROWS * self.traits.head_dim * 2
+            pout_base = buf_base_i64(self.ctx.pout_ptr) + pout_slot * fx.Int64(
+                pout_slot_bytes
+            )
+            self.pout_buffer = ptr_buf_tensor(
+                pout_base,
+                self.traits.Q_DTYPE,
+                n=self.traits.TOTAL_ROWS * self.traits.head_dim,
+                unit_elems=self.traits.OP_ELEMS,
+                # BF16 views may be only 2-byte aligned, even for 8-byte stores.
+                unit_stride=1,
+                num_records_bytes=pout_slot_bytes,
+            )
+
+    def _emit(self, o_norm, sub, global_row):
+        if const_expr(self.traits.buffer_plan_output):
+            # Row guards keep the full 8-byte store within this slot.
+            pout_offset = global_row * self.traits.head_dim + sub * self.traits.OP_ELEMS
+            buf_copy_store(
+                self.pout_buffer,
+                pout_offset,
+                o_norm,
+                elem=self.traits.Q_DTYPE,
+                unit_elems=self.traits.OP_ELEMS,
+                cache_modifier=0,
+            )
+        else:
+            base = self.ctx.partial_slot * self.traits.TOTAL_ROWS + global_row
+            pout_offset = base * self.traits.head_dim + sub * self.traits.OP_ELEMS
+            fx.ptr_store(o_norm, fx.add_offset(self.pout, pout_offset))
+
+    @flyc.jit
+    def store(self, o_final):
+        # Keep pointer writes as SSA operands of the row/warp guards.
+        pmax = self.pmax
+        psum = self.psum
+        # Each lane stores four head-dim values for one query row.
+        inv_fp8 = fx.Float32(1.0 / self.traits.FP8_MAX)
+        for m in range_constexpr(self.traits.M_TILES):
+            # Flat (mtp, gqa) query-row for this lane.
+            row = m * MFMA_MNK + self.ctx.lane16
+            global_row = self.ctx.query_begin * self.traits.query_group_size + row
+            l_row = o_final[self.traits.l_slot(m)]
+            safe_l = (l_row > self.ctx.ZERO_F).select(l_row, fx.Float32(1.0))
+            inv_l = fx.Float32(rcp_f32(safe_l))
+            if const_expr(self.traits.per_token_kv or self.traits.is_bf16_kv):
+                o_scale = inv_l
+            else:
+                o_scale = inv_l * (self.ctx.v_scale_f * inv_fp8)
+            o_scale_b = fx.Vector.from_elements(
+                [o_scale], dtype=fx.Float32
+            ).broadcast_to(self.traits.OP_ELEMS)
+
+            for vh in range_constexpr(self.traits.VHE_CHUNKS):
+                o_slot = self.traits.o_slot(m, vh)
+                o_norm = (o_final[o_slot] * o_scale_b).to(self.traits.Q_DTYPE)
+                head_base = (
+                    vh * (self.traits.NWARP * MFMA_MNK)
+                    + self.ctx.warp * MFMA_MNK
+                    + self.ctx.rgroup * self.traits.OP_ELEMS
+                )
+                sub = head_base // self.traits.OP_ELEMS
+                if row < self.traits.CTA_ROWS:
+                    self._emit(o_norm, sub, global_row)
+
+            if self.ctx.warp == 0 and self.ctx.rgroup == 0:
+                base = self.ctx.partial_slot * self.traits.TOTAL_ROWS + global_row
+                if row < self.traits.CTA_ROWS:
+                    # The shared reducer expects maxima in natural-log units.
+                    pmax[base] = o_final[self.traits.m_slot(m)] * fx.Float32(
+                        1.0 / LOG2E
+                    )
+                    psum[base] = l_row