[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/kernels/pa_decode/__init__.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (C) 2025-2026 FlyDSL Project Contributors
+
+"""Object-oriented building blocks for the FP8 paged-attention kernel."""
+
+from functools import cache
+from hashlib import sha256
+from pathlib import Path
+
+
+@cache
+def implementation_cache_tag():
+    """Include operation classes in FlyDSL's persistent compilation cache.
+
+    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 they call, so an operation edit invalidates cached binaries.
+    Paths are relative to kernels/ so the key is stable across installations.
+    """
+    package = Path(__file__).parent
+    shared_helpers = (
+        "buffer_ops.py",
+        "dpp_utils.py",
+        "kernels_common.py",
+        "tensor_shim.py",
+        "utils.py",
+    )
+    sources = sorted(package.glob("*.py")) + [
+        package.parent / name for name in shared_helpers
+    ]
+    digest = sha256()
+    for source in sources:
+        digest.update(source.relative_to(package.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):
+            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):
+                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

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

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (C) 2025-2026 FlyDSL Project Contributors
+
+"""FP8 operand conversion and QK/PV MFMA operations."""
+
+import flydsl.expr as fx
+from flydsl.expr import const_expr, range_constexpr
+from flydsl.expr.typing import T
+
+from .traits import MFMA_ACC_ELEMS
+
+
+class PaDecodeGemm:
+    def __init__(self, traits):
+        self.traits = traits
+
+    def mfma(self, a_ops, b_ops, a_base, b_base, k_packs, acc):
+        if const_expr(self.traits.WIDE_FP8_MFMA):
+            for inst in range_constexpr(k_packs // self.traits.PACKS_PER_MFMA):
+                a_pack = fx.Vector.from_elements(
+                    [
+                        a_ops[a_base + inst * self.traits.PACKS_PER_MFMA + i]
+                        for i in range_constexpr(self.traits.PACKS_PER_MFMA)
+                    ],
+                    dtype=fx.Int64,
+                ).bitcast(fx.Int32)
+                b_pack = fx.Vector.from_elements(
+                    [
+                        b_ops[b_base + inst * self.traits.PACKS_PER_MFMA + i]
+                        for i in range_constexpr(self.traits.PACKS_PER_MFMA)
+                    ],
+                    dtype=fx.Int64,
+                ).bitcast(fx.Int32)
+                acc = fx.rocdl.mfma_scale_f32_16x16x128_f8f6f4(
+                    T.f32x4,
+                    [
+                        a_pack,
+                        b_pack,
+                        acc,
+                        0,
+                        0,
+                        0,
+                        fx.Int32(0x7F7F7F7F),
+                        0,
+                        fx.Int32(0x7F7F7F7F),
+                    ],
+                )
+        else:
+            for pack in range_constexpr(k_packs):
+                acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8(
+                    T.f32x4,
+                    [a_ops[a_base + pack], b_ops[b_base + pack], acc, 0, 0, 0],
+                )
+        return acc
+
+    def fp8_words(self, vf32):
+        n = vf32.shape[0]
+        words = []
+        for i in range_constexpr(n // 4):
+            b = i * 4
+            lo = fx.rocdl.cvt_pk_fp8_f32(T.i32, vf32[b], vf32[b + 1], 0, False)
+            words.append(
+                fx.rocdl.cvt_pk_fp8_f32(T.i32, vf32[b + 2], vf32[b + 3], lo, True)
+            )
+        return fx.Vector.from_elements(words, dtype=fx.Int32)
+
+    def qk(self, k_ops, q_ops, m):
+        frag_Ss = []
+        for a in range_constexpr(self.traits.NCHUNK):
+            acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32)
+            acc = self.mfma(
+                k_ops,
+                q_ops,
+                a * self.traits.N_SUBCHUNKS,
+                m * self.traits.N_SUBCHUNKS,
+                self.traits.N_SUBCHUNKS,
+                acc,
+            )
+            frag_Ss.append(fx.Vector(acc))
+        return frag_Ss
+
+    def pv(self, v_ops, p_ops):
+        acc = fx.Vector.filled(MFMA_ACC_ELEMS, 0.0, fx.Float32)
+        acc = self.mfma(v_ops, p_ops, 0, 0, self.traits.NVOPS, acc)
+        return fx.Vector(acc)