[FlyDSL] [Perf] pa_decode: launch gfx942 query-split CTAs split-major

The tile grid is (plan capacity, KV heads x query splits). Dispatch walks x
first, so split 1 of task 0 dispatches only after every slot of split 0,
padding slots included. On small grids, where the plan has many more slots
than tasks, the splits ran in sequential waves of mostly empty CTAs. On
large grids a task's splits ran far apart, each re-reading the task's K/V
from HBM.

On gfx942, split kernels now launch (KV heads x query splits, capacity), so a
task's splits dispatch together. They run concurrently and share the K/V
they re-read. The cache key records the choice. BF16 splits forced by the LDS
limit at large head dims only use it when tasks own at least four tiles; with
fewer, the split-major grid measured 0.93-0.97x.

Paired in-process A/B vs 8f9fe39 on MI308X (gfx942):
- 180 MTP grid cases: geomean 1.092x
  - FP8 MTP4/8 batch 1: 1.04-1.45x
  - FP8 sliding-window MTP4 batch 1/16: 1.28-1.63x
  - BF16 D1024 MTP batch >= 16 with ctx >= 8K: 1.10-1.26x
- 111 unit-test cases: geomean 1.018x (MTP4/8 and window cases 1.16-1.20x)
gfx950 binaries are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: fsx950223

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

@@ -66,14 +66,16 @@ def init_thread_mapping(self):
         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"))
+        # A split-major grid puts (KV head, query split) on x, the slot on y.
+        qx, px = ("x", "y") if self.traits.SPLIT_MAJOR else ("y", "x")
+        self.kv_query = fx.Int32(gpu.block_id(qx))
         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
-        self.part = fx.Int32(gpu.block_id("x"))
-        self.partial_slot = self.kv_h * fx.Int32(gpu.grid_dim.x) + self.part
+        self.n_kv = fx.Int32(getattr(gpu.grid_dim, qx)) // self.traits.query_splits
+        self.part = fx.Int32(gpu.block_id(px))
+        self.partial_slot = self.kv_h * fx.Int32(getattr(gpu.grid_dim, px)) + self.part
 
     def init_query_mapping(self):
         self.rgroup = (

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

@@ -451,6 +451,24 @@ def select(
                 or (single_tile_plan and split_workgroups <= 2 * num_compute_units)
             )
         )
+        # gfx942 split CTAs launch split-major, so a task's splits dispatch
+        # together instead of after every slot of the previous split. BF16
+        # head-dim splits forced by LDS only gain when tasks are long.
+        split_major = (
+            "gfx94" in arch
+            and query_splits > 1
+            and (
+                kv_dtype in FP8_KV_DTYPES
+                or head_dim <= 256
+                or (
+                    max_context_length is not None
+                    and max_tiles_per_task(
+                        num_seqs, work_capacity, num_partitions, max_context_length
+                    )
+                    >= 4
+                )
+            )
+        )
         cache_key = (
             head_dim,
             query_group_size,
@@ -472,6 +490,7 @@ def select(
             bool(mask_unowned_kv),
             direct_output,
             gfx942_busy,
+            split_major,
             D256_PIPE,
             stage_pages,
         )
@@ -532,6 +551,7 @@ class PaDecodeTraits:
     WAVES_PER_EU_2: bool
     GFX942_Q_EARLY: bool
     PAGES_EARLY: bool
+    SPLIT_MAJOR: bool
     REUSE_KV_PAGES: bool
     SCALES_BEFORE_CURRENT_V: bool
     M1_SCALE_BEFORE_MASK: bool
@@ -616,6 +636,7 @@ def create(cls, schedule: PaDecodeSchedule) -> PaDecodeTraits:
             mask_unowned_kv,
             DIRECT,
             gfx942_busy,
+            SPLIT_MAJOR,
             D256_PIPE,
             stage_page_capacity,
         ) = cache_key
@@ -927,6 +948,7 @@ def create(cls, schedule: PaDecodeSchedule) -> PaDecodeTraits:
                     or (M_TILES == 2 and kv_dtype in FP8_KV_DTYPES)
                 )
             ),
+            SPLIT_MAJOR=SPLIT_MAJOR,
             GFX942_Q_EARLY=(
                 not is_gfx950 and kv_dtype != MXFP4_KV_DTYPE and not STREAM_KV
             ),

aiter/ops/flydsl/kernels/pa_decode_kernel.py

@@ -441,7 +441,7 @@ def _run_task(seq, start, end, context):
 
         # Skip cleared padding records: no scratch writes or Q loads.
         # This CTA-uniform guard keeps every task barrier convergent.
-        slot = fx.Int32(gpu.block_id("x"))
+        slot = fx.Int32(gpu.block_id("y" if traits.SPLIT_MAJOR else "x"))
         work = fx.recast_iter(fx.Int32, work_info_ptr)
         if const_expr(traits.DIRECT):
             # One task per sequence, no plan: ``work_info`` is the context lengths.
@@ -517,7 +517,11 @@ def pa_decode_tile_launch(
                 work_info,
                 num_seqs,
             ).launch(
-                grid=(work_capacity, num_kv_heads * traits.query_splits, 1),
+                grid=(
+                    (num_kv_heads * traits.query_splits, work_capacity, 1)
+                    if traits.SPLIT_MAJOR
+                    else (work_capacity, num_kv_heads * traits.query_splits, 1)
+                ),
                 block=(traits.BLOCK_THREADS, 1, 1),
                 stream=stream,
             )