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

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,179 @@
+# 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 ..tensor_shim import ptr_buf_tensor
+from .traits import LOG2E, MFMA_MNK, WAVE
+
+
+class PaDecodeContext:
+    def __init__(
+        self,
+        traits,
+        shared_storage,
+        output_ptr,
+        pmax_ptr,
+        psum_ptr,
+        pout_ptr,
+        query_ptr,
+        key_cache_ptr,
+        value_cache_ptr,
+        block_tables_ptr,
+        context_lengths_ptr,
+        key_scale_ptr,
+        value_scale_ptr,
+        sinks_ptr,
+        max_blocks_per_seq,
+        stride_ks_block,
+        stride_ks_head,
+        stride_o_row,
+        stride_o_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.output_ptr = output_ptr
+        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.context_lengths_ptr = context_lengths_ptr
+        self.key_scale_ptr = key_scale_ptr
+        self.value_scale_ptr = value_scale_ptr
+        self.sinks_ptr = sinks_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_o_row = stride_o_row
+        self.stride_o_head = stride_o_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 = fx.Int32(gpu.block_id("x"))
+        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.part = fx.Int32(gpu.block_id("z"))  # context partition handled by this CTA
+        self.n_kv = fx.Int32(gpu.grid_dim.y) // self.traits.query_splits
+        if const_expr(self.traits.use_work_plan):
+            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.
+                self.seq = self.planned_seq
+                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 = self.seq  # Packed work slot, shared by all KV heads.
+                self.seq = self.planned_seq
+                self.partial_slot = self.kv_h * fx.Int32(gpu.grid_dim.x) + self.part
+        else:
+            self.partial_slot = (
+                self.seq * self.n_kv + self.kv_h
+            ) * self.traits.NP + 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):
+        if const_expr(self.traits.use_work_plan):
+            self.context_len = self.planned_context
+        else:
+            ctx_buf = ptr_buf_tensor(self.context_lengths_ptr, fx.Int32)
+            ctx_tiled = fx.logical_divide(ctx_buf, fx.make_layout(1, 1))
+            ctx_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32)
+            ctx_reg = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32)
+            fx.copy(ctx_copy_atom, fx.slice(ctx_tiled, (None, self.seq)), ctx_reg)
+            self.context_len = fx.Int32(fx.Vector(fx.memref_load_vec(ctx_reg))[0])
+
+    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
+        if const_expr(self.traits.use_work_plan):
+            self.part_start = self.planned_start
+            self.part_end = self.planned_end
+        else:
+            num_tiles = (
+                self.context_len + self.traits.TILE_TOK - 1
+            ) // self.traits.TILE_TOK
+            tiles_per_part = (num_tiles + self.traits.NP - 1) // self.traits.NP
+            self.part_start = self.part * tiles_per_part
+            part_end_raw = self.part_start + tiles_per_part
+            self.part_end = (part_end_raw < num_tiles).select(part_end_raw, num_tiles)
+
+    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,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (C) 2025-2026 FlyDSL Project Contributors
+
+"""Normalize PA output and store direct results or 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.output = fx.recast_iter(self.traits.Q_DTYPE, self.ctx.output_ptr)
+        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,
+            )
+        if const_expr(self.traits.DIRECT_SINKS):
+            self.sink_token = fx.recast_iter(self.traits.SINK_DTYPE, self.ctx.sinks_ptr)
+
+    def _emit(self, o_norm, sub, query_idx, query_head, global_row):
+        if const_expr(self.traits.NP == 1 and not self.traits.use_work_plan):
+            out_offset = (
+                (
+                    self.ctx.seq * self.traits.query_length
+                    + self.ctx.query_begin
+                    + query_idx
+                )
+                * self.ctx.stride_o_row
+                + query_head * self.ctx.stride_o_head
+                + sub * self.traits.OP_ELEMS
+            )
+            fx.ptr_store(o_norm, fx.add_offset(self.output, out_offset))
+        elif 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
+            qi_e = row // self.traits.query_group_size
+            gs_head_e = row - qi_e * self.traits.query_group_size
+            qh = self.ctx.kv_h * self.traits.query_group_size + gs_head_e
+            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.DIRECT_SINKS):
+                # Sinks affect only the denominator. Compare in natural-logit
+                # units before LOG2E scaling to avoid overflow for finite sinks.
+                sink_value = fx.Float32(self.sink_token[qh])
+                row_max = o_final[self.traits.m_slot(m)] * fx.Float32(1.0 / LOG2E)
+                total_max = fx.maxnumf(row_max, sink_value)
+                safe_max = (total_max > self.ctx.NEG_INF).select(
+                    total_max, self.ctx.ZERO_F
+                )
+                kv_mass = (l_row > self.ctx.ZERO_F).select(
+                    fx.exp2((row_max - safe_max) * fx.Float32(LOG2E), fastmath="fast"),
+                    self.ctx.ZERO_F,
+                )
+                # +inf suppresses all KV output; -inf disables this head's
+                # sink. Select before exp2 to avoid the +inf - +inf case.
+                sink_shift = (sink_value == safe_max).select(
+                    self.ctx.ZERO_F, sink_value - safe_max
+                )
+                sink_mass = fx.exp2(sink_shift * fx.Float32(LOG2E), fastmath="fast")
+                denominator = l_row * kv_mass + sink_mass
+                safe_denominator = (denominator > self.ctx.ZERO_F).select(
+                    denominator, fx.Float32(1.0)
+                )
+                inv_l = kv_mass * fx.Float32(rcp_f32(safe_denominator))
+            if const_expr(self.traits.per_token_kv):
+                o_scale = inv_l
+            elif const_expr(self.traits.SCALAR_FP8_DECODE):
+                # Direct P conversion needs no compensating 1/FP8_MAX.
+                o_scale = inv_l * self.ctx.v_scale_f
+            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, qi_e, qh, global_row)
+
+            if const_expr(  # noqa: SIM102
+                self.traits.NP > 1 or self.traits.use_work_plan
+            ):
+                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