From 02feee30eb685a5cbdee5546f0af30fa064b8c3b Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Wed, 23 Sep 2026 16:59:31 +0000
Subject: [PATCH 1/9] [Triton/Gluon] FlashKDA K2: keep a (segment, head)'s V
 blocks on one XCD

The K2 grid is (cdiv(W, BW), num_segs * H), and every V block of one
(segment, head) re-reads that segment's whole chunk workspace. Workgroups
reach the XCDs round-robin in launch order, x fastest, so with the tuned
BW=16 the eight V blocks of each (segment, head) landed on eight different
XCDs and each pulled the workspace through its own L2.

Measured on MI355X (rocprofv3 FETCH_SIZE/WRITE_SIZE): passes A and C moved
6-8x their compulsory bytes past L2, pinned at ~7 TB/s on every shape --
they were bound by L2-miss traffic, not by the recurrence. Renumbering the
flat program id with the existing remap_xcd, so each XCD owns a contiguous
run of whole (segment, head) groups, brings that traffic to 1.01-1.02x
compulsory and halves those kernels' time at an unchanged grid and tile.

Applied to the Triton segment kernel (passes A and C) and the Gluon pass A,
whose BW=32/64 tiles split V the same way. The remap only changes which
program computes which tile, so results are bit-identical.

FlashKDA end to end, MI355X, bf16, K=V=128:
  gfx950 default route (Gluon K1 + pass A, Triton pass C): 1.24-1.45x
  Triton only (non-gfx950 route):                          1.32-1.61x
  1x512x12:                                                1.01x

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../gfx950/kimi_delta_attn/flash_kda_k2.py    | 12 ++++++++++--
 .../kimi_delta_attn/flash_kda.py              | 19 +++++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
index 5570ad82f95..2865c9f807e 100644
--- a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
+++ b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
@@ -8,6 +8,7 @@
 
 from aiter.ops.triton._triton_kernels.kimi_delta_attn.fast_launch import fast_launch
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
+from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
 
 KW = 8
 KW_BIG = 8
@@ -139,6 +140,7 @@ def k2_ab_fused_gluon(
     B_OP_B: gl.constexpr,
     BLK: gl.constexpr,
     SH_KR: gl.constexpr,
+    NUM_XCDS: gl.constexpr,
 ):
     """Both pass-A recurrences in one launch, sharing every operand load.
 
@@ -150,8 +152,14 @@ def k2_ab_fused_gluon(
     states are then the same width and one program covers a column block of
     both.
     """
-    i_w = gl.program_id(0).to(gl.int64)
-    i_sh = gl.program_id(1).to(gl.int64)
+    # Keeps a (segment, head)'s V blocks on one XCD, so their re-reads of the
+    # chunk workspace share an L2. See the Triton K2 for the full reasoning.
+    n_w = gl.num_programs(0)
+    pid = remap_xcd(
+        gl.program_id(1) * n_w + gl.program_id(0), n_w * gl.num_programs(1), NUM_XCDS
+    )
+    i_w = (pid % n_w).to(gl.int64)
+    i_sh = (pid // n_w).to(gl.int64)
     i_seg = i_sh // H
     i_h = i_sh % H
 
diff --git a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
index 3c36124ef33..379e2fe0c40 100644
--- a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
+++ b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
@@ -71,6 +71,8 @@
     prepare_chunk_indices,
 )
 from aiter.ops.triton.utils._triton import arch_info
+from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
+from aiter.ops.triton.utils.device_info import get_num_xcds
 from aiter.ops.triton.utils.logger import AiterTritonLogger
 from aiter.ops.triton.utils.tuned_config_utils import autotune_configs
 
@@ -443,6 +445,7 @@ def _flash_kda_segment_kernel(
     STORE_H_OUT: tl.constexpr,
     STORE_FINAL: tl.constexpr,
     STATE_V_FIRST: tl.constexpr,
+    NUM_XCDS: tl.constexpr,
     CM_OUT: tl.constexpr = "",
     state_cache=None,
     state_indices=None,
@@ -466,8 +469,18 @@ def _flash_kda_segment_kernel(
     With one segment per sequence this degenerates to a plain sequential scan
     and only the third form runs.
     """
-    i_w = tl.program_id(0).to(tl.int64)
-    i_sh = tl.program_id(1).to(tl.int64)
+    # Every V block of one (segment, head) re-reads that segment's whole chunk
+    # workspace. Workgroups reach the XCDs round-robin in launch order, x
+    # fastest, so unmapped those blocks land on different XCDs and each pulls
+    # the workspace through its own L2 -- 6-8x the compulsory traffic at BW=16,
+    # which is what bound this kernel. Renumbered, each XCD owns a contiguous
+    # run of flat ids, so a (segment, head)'s V blocks share an L2.
+    n_w = tl.num_programs(0)
+    pid = remap_xcd(
+        tl.program_id(1) * n_w + tl.program_id(0), n_w * tl.num_programs(1), NUM_XCDS
+    )
+    i_w = (pid % n_w).to(tl.int64)
+    i_sh = (pid // n_w).to(tl.int64)
     i_seg, i_h = i_sh // H, i_sh % H
 
     chunk_base = tl.load(seg_chunk_base + i_seg).to(tl.int64)
@@ -1220,6 +1233,7 @@ def flash_kda_fwd(
         "V": V,
         "C": C,
         "STATE_V_FIRST": state_v_first,
+        "NUM_XCDS": get_num_xcds(),
         "CM_OUT": CM_OUT_STORE,
     }
     paged_args = {
@@ -1276,6 +1290,7 @@ def _launch_k2(*, W, **kw):
                 V=V,
                 C=C,
                 BW=bw,
+                NUM_XCDS=get_num_xcds(),
                 **_g2.build_layouts(nw),
                 num_warps=nw,
             )

From 638ae8b704ae8a8cf94bb961d6bb39fdbc09605c Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Tue, 6 Oct 2026 15:06:33 +0000
Subject: [PATCH 2/9] [Triton/Gluon] KDA: load token i once in
 intra_token_parallel

The token-parallel intra kernel (the default pipeline's path when safe_gate
is off) reloaded token i's q/k/g for every j and every K block. Load them
once outside the j loop, span all of K in one block (BK from K, no longer
tuned) and fold beta into k, following fla. The default config moves from
BH=1 / 4 warps, which idles most of each warp, to BH=4 / 1 warp: within
1.24x of the per-shape best on gfx950.

Default pipeline forward with safe_gate=False on MI355X: 2.0-2.7x
(1x8192x32 2862 -> 1221 us). Aqk/Akk change only in summation order
(output within 2.2e-5 relative).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../kimi_delta_attn/intra_attn.py             | 59 +++++++++----------
 1 file changed, 28 insertions(+), 31 deletions(-)

diff --git a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py
index 420edfcbf2b..e600e5d4add 100644
--- a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py
+++ b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py
@@ -54,18 +54,21 @@ def _gather(src, index, axis, _builder=None):
 @triton.heuristics(
     {
         "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+        # Token i's rows are loaded once, so one block spans all of K.
+        "BK": lambda args: max(triton.next_power_of_2(args["K"]), 16),
     }
 )
 @triton.autotune(
+    # BH=1 idles most of a warp. The default is one config for all shapes,
+    # within 1.24x of the per-shape best on gfx950.
     configs=autotune_configs(
         "CHUNK_DELTA_ATTN",
         [
-            triton.Config({"BH": BH, "BK": BK}, num_warps=nw)
+            triton.Config({"BH": BH}, num_warps=nw, num_stages=3)
             for BH in [1, 2, 4, 8]
-            for BK in [32, 64]
             for nw in [1, 2, 4, 8]
         ],
-        default_config=triton.Config({"BH": 1, "BK": 64}, num_warps=4),
+        default_config=triton.Config({"BH": 4}, num_warps=1, num_stages=3),
     ),
     key=["K", "H", "HV"],
     **autotune_cache_kwargs,
@@ -137,44 +140,38 @@ def chunk_delta_attn_fwd_kernel_intra_token_parallel(
 
     o_hv = i_hg * BH + tl.arange(0, BH)
     o_h = o_hv // G
+    o_k = tl.arange(0, BK)
     m_hv = o_hv < HV
-
+    m_k = o_k < K
+    m_hk = m_hv[:, None] & m_k[None, :]
+
+    # Token i's rows are loaded once, outside the loop over j: re-reading them
+    # per (j, K block) made the kernel load-bound. beta is folded into k here.
+    p_qk = o_h[:, None] * K + o_k[None, :]
+    b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
+    b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
+    p_g = g + i_t * HV * K + o_hv[:, None] * K + o_k[None, :]
+    b_g = tl.load(p_g, mask=m_hk, other=0.0).to(tl.float32)
     p_beta = beta + i_t * HV + o_hv
-    b_beta = tl.load(p_beta, mask=m_hv, other=0.0).to(tl.float32)
+    b_k = b_k * tl.load(p_beta, mask=m_hv, other=0.0).to(tl.float32)[:, None]
 
-    # Accumulate dot products across K-tiles (BK ≤ 64 to avoid slow tl.dot / huge scatter loads)
     for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))):
-        b_Aqk_j = tl.zeros([BH], dtype=tl.float32)
-        b_Akk_j = tl.zeros([BH], dtype=tl.float32)
-        for i_k in range(tl.cdiv(K, BK)):
-            o_k = i_k * BK + tl.arange(0, BK)
-            m_k = o_k < K
-            m_hk = m_hv[:, None] & m_k[None, :]
-            p_qk = o_h[:, None] * K + o_k[None, :]
-
-            b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
-            b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
-            b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
-
-            p_g = g + i_t * HV * K + o_hv[:, None] * K + o_k[None, :]
-            p_gj = g + j * HV * K + o_hv[:, None] * K + o_k[None, :]
-            b_g = tl.load(p_g, mask=m_hk, other=0.0).to(tl.float32)
-            b_gj = tl.load(p_gj, mask=m_hk, other=0.0).to(tl.float32)
-
-            b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0)
-            b_Aqk_j += tl.sum(b_q * b_kgj, axis=1)
-            b_Akk_j += tl.sum(b_k * b_beta[:, None] * b_kgj, axis=1)
-
-        b_Aqk_j *= scale
-        b_Akk_j *= tl.where(j < i_t, 1.0, 0.0)
+        b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
+        p_gj = g + j * HV * K + o_hv[:, None] * K + o_k[None, :]
+        b_gj = tl.load(p_gj, mask=m_hk, other=0.0).to(tl.float32)
+
+        b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0)
+        b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale
+        b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0)
+
         tl.store(
             Aqk + i_t * HV * BT + o_hv * BT + j % BT,
-            b_Aqk_j.to(Aqk.dtype.element_ty),
+            b_Aqk.to(Aqk.dtype.element_ty),
             mask=m_hv,
         )
         tl.store(
             Akk + i_t * HV * BC + o_hv * BC + j - i_ts,
-            b_Akk_j.to(Akk.dtype.element_ty),
+            b_Akk.to(Akk.dtype.element_ty),
             mask=m_hv,
         )
 

From 3201b94bec87758bc7640e8549a1c680ae4afd09 Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Tue, 6 Oct 2026 15:07:42 +0000
Subject: [PATCH 3/9] [Triton/Gluon] FlashKDA: cap segment length at 64 chunks

_choose_chunks_per_seg capped the segment count at 16 over both of its
reasons, so a long few-head sequence got 256-chunk segments whose serial
walk bounds the recurrence. Cap the occupancy target alone at 16 segments
and then cut segments to at most 64 chunks however many that takes. On
gfx950, 64-chunk segments beat longer ones by up to 1.9x on 1x131072x4.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../kimi_delta_attn/flash_kda.py              | 23 +++++++++++--------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
index 379e2fe0c40..619c506b882 100644
--- a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
+++ b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
@@ -827,10 +827,12 @@ def _num_cus(device_index: int = 0) -> int:
     return torch.cuda.get_device_properties(device_index).multi_processor_count
 
 
-# Blocks pass A should end up with, in units of the CU count.
+# Blocks pass A should end up with, in units of the CU count, and the most
+# segments per sequence taken for that reason alone.
 _SEG_TARGET_BLOCKS = 3
 _SEG_MAX_SEGMENTS = 16
-_SEG_MAX_CHUNKS = 32
+# Longest segment, however many segments that takes.
+_SEG_MAX_CHUNKS = 64
 _SEG_MIN_CHUNKS = 64
 
 _SCAN_BV_NARROW = 16
@@ -852,7 +854,10 @@ def _choose_chunks_per_seg(n_chunks_max: int, n_seqs: int, H: int, V: int) -> in
     H=12, V=128 that is 48 blocks against 256 CUs. Segmenting buys blocks by
     turning one pass into three -- pass A, the cross-segment scan, and pass C --
     so what decides the length is the block count it lands on rather than the
-    sequence length on its own.
+    sequence length on its own. Once taken, segments are also kept to
+    ``_SEG_MAX_CHUNKS``: on gfx950, 64-chunk segments beat longer ones by up to
+    1.9x on long few-head sequences (1 x 131072 x 4), where the occupancy target
+    alone gives 256-chunk ones.
     """
     override = os.getenv("CHUNK_DELTA_ATTN_FLASH_KDA_SEG", "").strip()
     if override:
@@ -867,12 +872,12 @@ def _choose_chunks_per_seg(n_chunks_max: int, n_seqs: int, H: int, V: int) -> in
     # up for.
     if n_chunks_max < _SEG_MIN_CHUNKS:
         return n_chunks_max
-    # Enough segments to fill the device, and enough to keep a segment's own
-    # serial walk from becoming the limit on a long sequence. Capped over both,
-    # since the two extra passes scale with the segment count.
-    segs = min(
-        _SEG_MAX_SEGMENTS,
-        max(_SEG_TARGET_BLOCKS * cus / blocks, n_chunks_max / _SEG_MAX_CHUNKS),
+    # Enough segments to fill the device, capped since the two extra passes
+    # scale with the segment count; and however many more it takes to keep a
+    # segment's own serial walk from becoming the limit on a long sequence.
+    segs = max(
+        min(_SEG_MAX_SEGMENTS, _SEG_TARGET_BLOCKS * cus / blocks),
+        n_chunks_max / _SEG_MAX_CHUNKS,
     )
     # To a power of two, the grid the constants were calibrated on.
     return max(1, min(n_chunks_max, 1 << round(math.log2(n_chunks_max / segs))))

From c1397349419e0d8d56f6b0327dc9c31b5e6de0ff Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Tue, 6 Oct 2026 15:08:35 +0000
Subject: [PATCH 4/9] [Triton/Gluon] FlashKDA: LDS-prefetching Gluon K2 and a
 Gluon pass C

Speeds up the Gluon FlashKDA kernels on gfx950 and moves pass C to Gluon:

- K1: sigmoids through exp2 (one v_exp_f32 instead of exp's
  denormal-safe sequence); the gate in a column-per-thread layout, so its
  cumulative sum is in-thread, with tail rows given a zero gate (the last
  and pivot rows are then fixed rows); tail-row masks on the workspace
  dropped; warps split the rows of the [C, K] tiles, so the l2norm sums stay
  in a warp; operands staged through LDS once per product (10 barriers
  instead of 38); the gate load is masked, since buffer loads past the
  tensor's end are not clamped. Launches with waves_per_eu=3 (170 -> 168
  VGPRs, 3 waves/SIMD).
- K1 (Gluon and Triton) writes sigmoid(beta) to a new ws_beta workspace,
  which K2 copies into LDS with the chunk's other tiles.
- Pass A: chunk j+1's tiles are async-copied into a double-buffered LDS
  stage while chunk j computes; one barrier per chunk. The kernel needs
  ~280 VGPRs, so the published schedules carry waves_per_eu=2.
- Pass C: new k2_c_gluon, used for every shape on gfx950. Same prefetch;
  the output projection runs before the state update and is stored from
  the MFMA layout; seg_seq / seg_is_last are read before the loop. It
  supports the paged state_cache read and write-back of the Triton pass C.
  BW=64 unsegmented, BW=128 segmented.

FlashKDA forward on MI355X vs main: 1.43-3.02x (1x8192x12 279 -> 142 us,
1x8192x96 1458 -> 690 us, 1x131072x4 1634 -> 541 us, varlen 1.70-1.99x),
including the two commits before this one. Output within 4.5e-4 relative
of main (K1's l2norm and gate summation order), final state within 3.5e-5.

Tests: test_cases_reach_the_gluon_k2 now checks pass C as well as pass A,
the paged-cache test runs on both K2 routes, and the schedule test covers
the pass C tiles.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../gfx950/kimi_delta_attn/flash_kda_k1.py    | 229 ++++++----
 .../gfx950/kimi_delta_attn/flash_kda_k2.py    | 407 ++++++++++++++++--
 .../kimi_delta_attn/flash_kda.py              | 128 ++++--
 .../attention/chunk_delta_attn/DEFAULT.json   |  14 +-
 .../kimi_delta_attn/test_flash_kda.py         |  71 +--
 5 files changed, 670 insertions(+), 179 deletions(-)

diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
index 1b26e04990f..d90ea8261a4 100644
--- a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
+++ b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
@@ -9,9 +9,12 @@
 from aiter.ops.triton._triton_kernels.kimi_delta_attn.fast_launch import fast_launch
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
 
-_BLK_WARP_K: gl.constexpr = gl.BlockedLayout([1, 8], [8, 8], [1, 2], [1, 0])
+# [C, K] tiles: warps split the rows, so the l2norm row sums stay within a warp
+# (a cross-warp sum costs three barriers).
+_BLK_CK: gl.constexpr = gl.BlockedLayout([1, 8], [4, 16], [2, 1], [1, 0])
 _BLK1: gl.constexpr = gl.BlockedLayout([1], [64], [2], [0])
-_BLK_CC: gl.constexpr = gl.BlockedLayout([1, 1], [4, 16], [2, 1], [1, 0])
+# One whole C = 32 column per thread, for the gate's cumulative sum.
+_BLK_COL: gl.constexpr = gl.BlockedLayout([32, 1], [1, 64], [1, 2], [0, 1])
 
 _MMA_F16: gl.constexpr = gl.amd.AMDMFMALayout(
     version=4, instr_shape=[16, 16, 4], transposed=True, warps_per_cta=[2, 1]
@@ -28,6 +31,10 @@
 _SH_A: gl.constexpr = gl.SwizzledSharedLayout(8, 1, 16, [1, 0])
 _SH_B: gl.constexpr = gl.SwizzledSharedLayout(8, 1, 16, [0, 1])
 _SH_CC_F: gl.constexpr = gl.SwizzledSharedLayout(1, 2, 8, [0, 1])
+_SH_ROWS: gl.constexpr = gl.SwizzledSharedLayout(1, 1, 1, [1, 0])
+_SH_VEC: gl.constexpr = gl.SwizzledSharedLayout(1, 1, 1, [0])
+
+_LOG2E = gl.constexpr(1.4426950408889634)
 
 
 @gluon.jit
@@ -46,28 +53,24 @@ def _exp2(x):
 
 
 @gluon.jit
-def _sigmoid(x):
-    return gl.extra.libdevice.fast_dividef(1.0, 1.0 + _exp(-x.to(gl.float32)))
+def _sigmoid_log2(z):
+    """sigmoid(x) given ``z = -x * log2(e)``.
 
-
-@gluon.jit
-def _l2norm(x):
-    f = x.to(gl.float32)
-    return f * gl.rsqrt(gl.sum(f * f, axis=1) + 1e-6)[:, None]
+    exp2 lowers to a bare v_exp_f32, where exp adds a denormal-range rescale
+    that the sigmoid does not need.
+    """
+    return gl.extra.libdevice.fast_dividef(1.0, 1.0 + gl.exp2(z))
 
 
 @gluon.jit
-def _via_lds(x, shared: gl.constexpr, dot: gl.constexpr):
-    return gl.allocate_shared_memory(x.dtype, x.shape, shared, x).load(dot)
+def _sigmoid(x):
+    return _sigmoid_log2(x.to(gl.float32) * -_LOG2E)
 
 
 @gluon.jit
-def _dot_f32(a, b_op, a_op: gl.constexpr, acc_layout: gl.constexpr, N: gl.constexpr):
-    return gl.amd.cdna4.mfma(
-        gl.convert_layout(a, a_op),
-        b_op,
-        gl.zeros([a.shape[0], N], gl.float32, acc_layout),
-    )
+def _l2norm(x):
+    f = x.to(gl.float32)
+    return f * gl.rsqrt(gl.sum(f * f, axis=1) + 1e-6)[:, None]
 
 
 _k1_prepare_repr = make_kernel_repr(
@@ -89,6 +92,7 @@ def k1_prepare_gluon(
     ws_kr,
     ws_gt,
     ws_inv_mqk,
+    ws_beta,
     cu_seqlens,
     chunk_indices,
     scale,
@@ -105,6 +109,11 @@ def k1_prepare_gluon(
     CM_WS: gl.constexpr = "",
     CM_LOAD: gl.constexpr = ".cg",
 ):
+    """FlashKDA K1 on gfx950; same workspace contract as the Triton prepare kernel.
+
+    Also writes ``sigmoid(beta)`` to ``ws_beta``, which the Gluon K2 copies into
+    LDS with the chunk's other tiles.
+    """
     gl.static_assert(C == 32 and K == 128)
     NUM_DOUBLING: gl.constexpr = BC.bit_length() - 2
     NUM_MERGE: gl.constexpr = (C // BC).bit_length() - 1
@@ -132,8 +141,8 @@ def k1_prepare_gluon(
         return
     actual_len = gl.minimum(C, T_seq - t_off)
 
-    o_c = gl.arange(0, C, layout=gl.SliceLayout(1, _BLK_WARP_K))
-    o_k = gl.arange(0, K, layout=gl.SliceLayout(0, _BLK_WARP_K))
+    o_c = gl.arange(0, C, layout=gl.SliceLayout(1, _BLK_CK))
+    o_k = gl.arange(0, K, layout=gl.SliceLayout(0, _BLK_CK))
     o_k_v = gl.arange(0, K, layout=_BLK1)
     o_i_r = gl.arange(0, C, layout=gl.SliceLayout(1, _MMA_B16))
     o_i_c = gl.arange(0, C, layout=gl.SliceLayout(0, _MMA_B16))
@@ -152,26 +161,59 @@ def k1_prepare_gluon(
     b_k_raw = gl.amd.cdna4.buffer_load(
         ptr=k, offsets=qk_off, mask=m_ck, other=0.0, cache=CM_LOAD
     )
-    if HAS_BIAS:
-        bias = gl.amd.cdna4.buffer_load(
-            ptr=dt_bias, offsets=(i_h * K).to(gl.int32) + o_k
-        )[None, :]
 
+    # The gate is cumulated over rows, so it is computed with each thread owning
+    # a whole column: the scan and the row picks below are then in-thread.
+    o_c_col = gl.arange(0, C, layout=gl.SliceLayout(1, _BLK_COL))
+    o_k_col = gl.arange(0, K, layout=gl.SliceLayout(0, _BLK_COL))
+    g_off = (base * K + o_c_col[:, None].to(gl.int64) * (H * K) + o_k_col[None, :]).to(
+        gl.int32
+    )
+    # Masked: the buffer resource has no size, so rows past the tensor's end are
+    # not clamped by hardware and would fault.
     b_g = gl.amd.cdna4.buffer_load(
-        ptr=g_raw, offsets=qk_off, mask=m_ck, other=0.0, cache=CM_LOAD
+        ptr=g_raw,
+        offsets=g_off,
+        mask=(o_c_col < actual_len)[:, None],
+        other=0.0,
+        cache=CM_LOAD,
     ).to(gl.float32)
+    # sigmoid(e^A * (g + bias)) with -log2(e) * e^A folded into one scalar.
+    rate = _exp(gl.load(A_log + i_h)) * -_LOG2E
     if HAS_BIAS:
-        b_g = b_g + bias
-    b_A = gl.load(A_log + i_h)
-    b_gate = lower_bound * _sigmoid(_exp(b_A) * b_g)
-    log2_e: gl.constexpr = 1.4426950408889634
-    b_gcum = gl.associative_scan(b_gate, 0, _add) * log2_e
-    b_gcum = gl.where(m_ck, b_gcum, 0.0)
-
-    b_g_last = gl.sum(gl.where(o_c[:, None] == actual_len - 1, b_gcum, 0.0), axis=0)
-    b_g_total = _exp2(b_g_last)
+        bias = gl.amd.cdna4.buffer_load(
+            ptr=dt_bias, offsets=(i_h * K).to(gl.int32) + o_k_col
+        )
+        b_z = b_g * rate + (bias * rate)[None, :]
+    else:
+        b_z = b_g * rate
+    # Tail rows get a zero gate (g = 0 alone would give lower_bound / 2, plus the
+    # bias), so the cumulative gate is constant from row actual_len - 1 on: the
+    # last row is row C - 1, and the pivot row min(C / 2, actual_len - 1) is
+    # row C / 2.
+    b_gate = gl.where(
+        o_c_col[:, None] < actual_len,
+        (lower_bound * _LOG2E) * _sigmoid_log2(b_z),
+        0.0,
+    )
+    b_gcum_col = gl.associative_scan(b_gate, 0, _add)
+    b_g_last_col = gl.sum(gl.where(o_c_col[:, None] == C - 1, b_gcum_col, 0.0), axis=0)
+    b_gp_col = gl.sum(gl.where(o_c_col[:, None] == C // 2, b_gcum_col, 0.0), axis=0)
+    # One LDS round trip (one barrier) takes the cumulative gate and its last and
+    # pivot rows to the row layout; separate conversions cost one or two
+    # barriers each.
+    s_gcum = gl.allocate_shared_memory(gl.float32, [C, K], _SH_ROWS, b_gcum_col)
+    s_g_last = gl.allocate_shared_memory(gl.float32, [K], _SH_VEC, b_g_last_col)
+    s_gp = gl.allocate_shared_memory(gl.float32, [K], _SH_VEC, b_gp_col)
+    b_gcum = s_gcum.load(_BLK_CK)
+    b_g_last = s_g_last.load(gl.SliceLayout(0, _BLK_CK))
+    b_gp = s_gp.load(gl.SliceLayout(0, _BLK_CK))
+    b_g_total = _exp2(b_g_last_col)
     b_exp_g = _exp2(b_gcum)
 
+    # Tail rows of q and k load as 0 and stay 0 through the norm. Every factor
+    # they meet below is finite (the cumulative gate is constant over the tail),
+    # so the workspace tiles' tail rows come out zero without a mask.
     b_q = _l2norm(b_q_raw)
     b_k = _l2norm(b_k_raw)
 
@@ -180,20 +222,18 @@ def k1_prepare_gluon(
         gl.int32
     )
     gl.amd.cdna4.buffer_store(
-        gl.where(m_ck, b_k * b_exp_g, 0.0).to(ws_kd.dtype.element_ty),
+        (b_k * b_exp_g).to(ws_kd.dtype.element_ty),
         ws_kd,
         ck_off,
         cache=CM_WS,
     )
     gl.amd.cdna4.buffer_store(
-        gl.where(m_ck, b_q * b_exp_g * scale, 0.0).to(ws_qd.dtype.element_ty),
+        (b_q * b_exp_g * scale).to(ws_qd.dtype.element_ty),
         ws_qd,
         ck_off,
         cache=CM_WS,
     )
-    b_kr_val = gl.where(m_ck, b_k * _exp2(b_g_last[None, :] - b_gcum), 0.0).to(
-        gl.bfloat16
-    )
+    b_kr_val = (b_k * _exp2(b_g_last[None, :] - b_gcum)).to(gl.bfloat16)
     gl.amd.cdna4.buffer_store(
         b_kr_val.to(ws_kr.dtype.element_ty), ws_kr, ck_off, cache=CM_WS
     )
@@ -212,41 +252,43 @@ def k1_prepare_gluon(
             other=0.0,
         ).to(gl.float32)
     )
+    gl.amd.cdna4.buffer_store(
+        b_beta, ws_beta, (ws_idx * C).to(gl.int32) + o_i_r, cache=CM_WS
+    )
 
-    o_mid = gl.minimum(C // 2, actual_len - 1)
-    b_gp = gl.sum(gl.where(o_c[:, None] == o_mid, b_gcum, 0.0), axis=0)
     b_gm = b_gcum - b_gp[None, :]
     b_dec = _exp2(b_gm)
     b_inc = _exp2(-b_gm)
-    b_k_piv = gl.where(m_ck, b_k * b_dec, 0.0).to(gl.bfloat16)
-    b_q_piv = gl.where(m_ck, b_q * b_dec * scale, 0.0).to(gl.bfloat16)
-    b_k_inv = gl.where(m_ck, b_k * b_inc, 0.0).to(gl.bfloat16)
-
-    b_kinv_b = _via_lds(gl.permute(b_k_inv, 1, 0), _SH_B, _B8_16)
+    b_k_piv = (b_k * b_dec).to(gl.bfloat16)
+    b_q_piv = (b_q * b_dec * scale).to(gl.bfloat16)
+    b_k_inv = (b_k * b_inc).to(gl.bfloat16)
+
+    # All three operands go to LDS before any is read back: one barrier for the
+    # lot instead of two per operand.
+    s_kinv = gl.allocate_shared_memory(
+        gl.bfloat16, [K, C], _SH_B, gl.permute(b_k_inv, 1, 0)
+    )
+    s_kpiv = gl.allocate_shared_memory(gl.bfloat16, [C, K], _SH_A, b_k_piv)
+    s_qpiv = gl.allocate_shared_memory(gl.bfloat16, [C, K], _SH_A, b_q_piv)
+    b_kinv_b = s_kinv.load(_B8_16)
 
     b_L = gl.amd.cdna4.mfma(
-        _via_lds(b_k_piv, _SH_A, _A8_16),
-        b_kinv_b,
-        gl.zeros([C, C], gl.float32, _MMA_B16),
+        s_kpiv.load(_A8_16), b_kinv_b, gl.zeros([C, C], gl.float32, _MMA_B16)
     )
     b_L = gl.where(o_i_r[:, None] > o_i_c[None, :], -b_L * b_beta[:, None], 0.0)
 
     b_Mqk = gl.amd.cdna4.mfma(
-        _via_lds(b_q_piv, _SH_A, _A8_16),
-        b_kinv_b,
-        gl.zeros([C, C], gl.float32, _MMA_B16),
+        s_qpiv.load(_A8_16), b_kinv_b, gl.zeros([C, C], gl.float32, _MMA_B16)
     )
     b_Mqk = gl.where(o_i_r[:, None] >= o_i_c[None, :], b_Mqk, 0.0)
 
-    o_r_cc = gl.arange(0, C, layout=gl.SliceLayout(1, _BLK_CC))
-    o_c_cc = gl.arange(0, C, layout=gl.SliceLayout(0, _BLK_CC))
-    cc_off_raw = (
-        ws_idx * (2 * C * C) + o_r_cc[:, None].to(gl.int64) * C + o_c_cc[None, :]
-    ).to(gl.int32)
+    # Mqk and INV are stored from the MFMA layout (four consecutive columns per
+    # thread): converting them for a wider store costs two barriers each.
+    cc_base = (ws_idx * (2 * C * C)).to(gl.int32)
     gl.amd.cdna4.buffer_store(
-        gl.convert_layout(b_Mqk.to(ws_inv_mqk.dtype.element_ty), _BLK_CC),
+        b_Mqk.to(ws_inv_mqk.dtype.element_ty),
         ws_inv_mqk,
-        cc_off_raw + C * C,
+        cc_base + C * C + o_i_r[:, None] * C + o_i_c[None, :],
         cache=CM_WS,
     )
 
@@ -254,14 +296,47 @@ def k1_prepare_gluon(
         b_D = b_L
     else:
         b_D = gl.where(o_i_r[:, None] // BC == o_i_c[None, :] // BC, b_L, 0.0)
-    b_INV = gl.where(o_i_r[:, None] == o_i_c[None, :], 1.0, 0.0) + b_D
-    b_INV = gl.convert_layout(b_INV, _MMA_F16)
-    b_Dp = _dot_f32(b_D, _via_lds(b_D, _SH_CC_F, _BF16), _AF16, _MMA_F16, C)
-    for _ in gl.static_range(NUM_DOUBLING):
-        dp_b = _via_lds(b_Dp, _SH_CC_F, _BF16)
-        b_INV = b_INV + _dot_f32(b_INV, dp_b, _AF16, _MMA_F16, C)
-        b_Dp = _dot_f32(b_Dp, dp_b, _AF16, _MMA_F16, C)
-
+    # I + D; D is strictly lower triangular, so its diagonal is free for the 1s.
+    b_INV = gl.convert_layout(
+        gl.where(o_i_r[:, None] == o_i_c[None, :], 1.0, b_D), _MMA_F16
+    )
+    zero_cc = gl.zeros([C, C], gl.float32, _MMA_F16)
+    # Each product's operands go through LDS once: stored together, one barrier,
+    # then loaded in each operand layout they are used in. Steps alternate
+    # between two buffer pairs, so a step's stores never land on what the
+    # previous step is still reading. Barriers are placed per allocation, so
+    # every operand has its own: two stores into one buffer cost a barrier
+    # between them.
+    s_a0 = gl.allocate_shared_memory(gl.float32, [C, C], _SH_CC_F)
+    s_b0 = gl.allocate_shared_memory(gl.float32, [C, C], _SH_CC_F)
+    s_a1 = gl.allocate_shared_memory(gl.float32, [C, C], _SH_CC_F)
+    s_b1 = gl.allocate_shared_memory(gl.float32, [C, C], _SH_CC_F)
+    s_a1.store(b_D)
+    b_Dp = gl.amd.cdna4.mfma(s_a1.load(_AF16), s_a1.load(_BF16), zero_cc)
+    for i in gl.static_range(NUM_DOUBLING):
+        if i % 2 == 0:
+            s_dp = s_a0
+            s_inv = s_b0
+        else:
+            s_dp = s_a1
+            s_inv = s_b1
+        s_dp.store(b_Dp)
+        s_inv.store(b_INV)
+        dp_b = s_dp.load(_BF16)
+        b_INV = gl.amd.cdna4.mfma(s_inv.load(_AF16), dp_b, b_INV)
+        b_Dp = gl.amd.cdna4.mfma(s_dp.load(_AF16), dp_b, zero_cc)
+
+    # The merge steps keep to fixed buffers: off and INV in the pair the doubling
+    # would use next, inner in the other pair. The first pair was last read
+    # before the previous step's barrier, the second before this step's first.
+    if NUM_DOUBLING % 2 == 0:
+        s_off = s_a0
+        s_inv = s_b0
+        s_inner = s_a1
+    else:
+        s_off = s_a1
+        s_inv = s_b1
+        s_inner = s_a0
     w = BC
     for _ in gl.static_range(NUM_MERGE):
         ne_w = o_i_r[:, None] // w != o_i_c[None, :] // w
@@ -269,22 +344,27 @@ def k1_prepare_gluon(
             m_off = (o_i_r[:, None] // (2 * w) == o_i_c[None, :] // (2 * w)) & ne_w
         else:
             m_off = ne_w
-        b_off = gl.where(m_off, b_L, 0.0)
-        inner = _dot_f32(b_off, _via_lds(b_INV, _SH_CC_F, _BF16), _AF16, _MMA_F16, C)
-        b_INV = b_INV + _dot_f32(
-            b_INV, _via_lds(inner, _SH_CC_F, _BF16), _AF16, _MMA_F16, C
-        )
+        s_off.store(gl.where(m_off, b_L, 0.0))
+        s_inv.store(b_INV)
+        inner = gl.amd.cdna4.mfma(s_off.load(_AF16), s_inv.load(_BF16), zero_cc)
+        s_inner.store(inner)
+        b_INV = gl.amd.cdna4.mfma(s_inv.load(_AF16), s_inner.load(_BF16), b_INV)
         w = 2 * w
 
+    o_f_r = gl.arange(0, C, layout=gl.SliceLayout(1, _MMA_F16))
+    o_f_c = gl.arange(0, C, layout=gl.SliceLayout(0, _MMA_F16))
     gl.amd.cdna4.buffer_store(
-        gl.convert_layout(b_INV.to(ws_inv_mqk.dtype.element_ty), _BLK_CC),
+        b_INV.to(ws_inv_mqk.dtype.element_ty),
         ws_inv_mqk,
-        cc_off_raw,
+        cc_base + o_f_r[:, None] * C + o_f_c[None, :],
         cache=CM_WS,
     )
 
 
 _NUM_WARPS = math.prod(_MMA_F16.warps_per_cta)
+# Holds K1 to 168 VGPRs; left alone it lands at 170 and drops to 2 waves per
+# SIMD (0.85x).
+_WAVES_PER_EU = 3
 
 
 _k1_fast = fast_launch(k1_prepare_gluon)
@@ -302,6 +382,7 @@ def gluon_k1_prepare(
     ws_kr,
     ws_gt,
     ws_inv_mqk,
+    ws_beta,
     cu_seqlens,
     chunk_indices,
     scale,
@@ -329,6 +410,7 @@ def gluon_k1_prepare(
         ws_kr=ws_kr,
         ws_gt=ws_gt,
         ws_inv_mqk=ws_inv_mqk,
+        ws_beta=ws_beta,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
         scale=scale,
@@ -345,4 +427,5 @@ def gluon_k1_prepare(
         CM_WS=CM_WS,
         CM_LOAD=CM_LOAD,
         num_warps=_NUM_WARPS,
+        waves_per_eu=_WAVES_PER_EU,
     )
diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
index 2865c9f807e..17ff000490a 100644
--- a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
+++ b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
@@ -44,16 +44,18 @@ def build_layouts(nw, kw=KW, kw_big=KW_BIG):
         "MMA_B": mma_b,
         "A_OP_B": gl.DotOperandLayout(0, mma_b, kw_big),
         "B_OP_B": gl.DotOperandLayout(1, mma_b, kw_big),
+        # Sources of the global -> LDS copies: 128 (or 32) bits per thread and
+        # no lane replication, as buffer_load_to_shared requires.
         "BLK": gl.BlockedLayout([1, 8], [4, 16], [nw, 1], [1, 0]),
-        "SH_KR": gl.SwizzledSharedLayout(8, 1, 16, [0, 1]),
+        "BLK_CC": gl.BlockedLayout([1, 2], [4, 16], [nw, 1], [1, 0]),
+        # gt and beta: one element per thread.
+        "BLK_1D": gl.BlockedLayout([1], [64], [nw], [0]),
+        "SH_WS": gl.SwizzledSharedLayout(8, 1, 16, [1, 0]),
+        "SH_PLAIN": gl.SwizzledSharedLayout(1, 1, 1, [1, 0]),
+        "SH_1D": gl.SwizzledSharedLayout(1, 1, 1, [0]),
     }
 
 
-@gluon.jit
-def _sigmoid(x):
-    return gl.extra.libdevice.fast_dividef(1.0, 1.0 + gl.exp(-x.to(gl.float32)))
-
-
 @gluon.jit
 def _recur(
     h,
@@ -94,16 +96,71 @@ def _recur(
 
 
 @gluon.jit
-def _kr_operand(
-    kr_raw,
+def _issue_chunk(
+    ws_kd,
+    ws_kr,
+    ws_gt,
+    ws_beta,
+    ws_inv_mqk,
+    v_input,
+    ws_qd,
+    s_kd,
+    s_kr,
+    s_inv,
+    s_v,
+    s_gt,
+    s_beta,
+    s_qd,
+    s_mqk,
+    ws_idx,
+    t0,
+    tok_end,
+    ws_off,
+    cc_off,
+    v_off,
+    o_c_v,
+    o_t,
+    H: gl.constexpr,
     K: gl.constexpr,
+    V: gl.constexpr,
     C: gl.constexpr,
-    SH_KR: gl.constexpr,
-    A_OP: gl.constexpr,
+    WITH_OUTPUT: gl.constexpr,
 ):
-    return gl.allocate_shared_memory(
-        gl.bfloat16, [K, C], SH_KR, gl.permute(kr_raw, 1, 0)
-    ).load(A_OP)
+    """Start fetching one chunk into LDS by async copy, as one commit group.
+
+    ``gt`` and ``beta`` (K1 applied its sigmoid) have fewer elements than
+    threads, and the async copy cannot lower lane replication, so they are
+    copied one element per thread (``o_t``) with the excess threads masked.
+    Tail rows of ``v`` are masked here; ``beta`` of a tail row is harmless
+    because the recurrence zeroes those rows of U anyway. ``WITH_OUTPUT``
+    (pass C) also fetches the output projection's ``qd`` and ``Mqk`` tiles.
+    """
+    tb = t0 * H
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        s_gt, ws_gt + ws_idx * K, o_t, mask=o_t < K
+    )
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        s_beta, ws_beta + ws_idx * C, o_t, mask=o_t < C
+    )
+    ck = ws_idx * (C * K)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(s_kd, ws_kd + ck, ws_off)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(s_kr, ws_kr + ck, ws_off)
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        s_inv, ws_inv_mqk + ws_idx * (2 * C * C), cc_off
+    )
+    gl.amd.cdna4.async_copy.buffer_load_to_shared(
+        s_v,
+        v_input + tb * V,
+        v_off,
+        mask=((t0 + o_c_v) < tok_end)[:, None],
+        other=0.0,
+    )
+    if WITH_OUTPUT:
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(s_qd, ws_qd + ck, ws_off)
+        gl.amd.cdna4.async_copy.buffer_load_to_shared(
+            s_mqk, ws_inv_mqk + ws_idx * (2 * C * C) + C * C, cc_off
+        )
+    gl.amd.cdna4.async_copy.commit_group()
 
 
 _k2_ab_fused_repr = make_kernel_repr(
@@ -117,9 +174,9 @@ def k2_ab_fused_gluon(
     ws_kd,
     ws_kr,
     ws_gt,
+    ws_beta,
     ws_inv_mqk,
     v_input,
-    beta_raw,
     h_out_b,
     h_out_a,
     seg_chunk_base,
@@ -139,7 +196,11 @@ def k2_ab_fused_gluon(
     A_OP_B: gl.constexpr,
     B_OP_B: gl.constexpr,
     BLK: gl.constexpr,
-    SH_KR: gl.constexpr,
+    BLK_CC: gl.constexpr,
+    BLK_1D: gl.constexpr,
+    SH_WS: gl.constexpr,
+    SH_PLAIN: gl.constexpr,
+    SH_1D: gl.constexpr,
     NUM_XCDS: gl.constexpr,
 ):
     """Both pass-A recurrences in one launch, sharing every operand load.
@@ -151,6 +212,12 @@ def k2_ab_fused_gluon(
     serial dependence each one has on its own. Requires K == V, since the two
     states are then the same width and one program covers a column block of
     both.
+
+    Chunk j+1's operands are fetched while chunk j computes: its tiles by async
+    global -> LDS copy into the other half of a double buffer, so the prefetch
+    costs no registers (a register prefetch halves occupancy). Without it the
+    kernel spent a third or more of its wave-cycles waiting on loads issued and
+    consumed in the same iteration.
     """
     # Keeps a (segment, head)'s V blocks on one XCD, so their re-reads of the
     # chunk workspace share an L2. See the Triton K2 for the full reasoning.
@@ -168,52 +235,77 @@ def k2_ab_fused_gluon(
     tok_base = gl.load(seg_tok_base + i_seg).to(gl.int64)
     tok_end = gl.load(seg_tok_end + i_seg).to(gl.int64)
 
-    o_c_ab = gl.arange(0, C, layout=gl.SliceLayout(1, A_OP_B))
-    o_k_ab = gl.arange(0, K, layout=gl.SliceLayout(0, A_OP_B))
-    o_c_a = gl.arange(0, C, layout=gl.SliceLayout(1, A_OP))
-    o_cc_a = gl.arange(0, C, layout=gl.SliceLayout(0, A_OP))
+    BLK_V: gl.constexpr = gl.BlockedLayout(
+        [1, 8], [64 // (BW // 8), BW // 8], [gl.num_warps(), 1], [1, 0]
+    )
     o_c_s = gl.arange(0, C, layout=gl.SliceLayout(1, BLK))
     o_k_s = gl.arange(0, K, layout=gl.SliceLayout(0, BLK))
+    o_r_cc = gl.arange(0, C, layout=gl.SliceLayout(1, BLK_CC))
+    o_c_cc = gl.arange(0, C, layout=gl.SliceLayout(0, BLK_CC))
+    o_c_v = gl.arange(0, C, layout=gl.SliceLayout(1, BLK_V))
+    o_w_v = i_w * BW + gl.arange(0, BW, layout=gl.SliceLayout(0, BLK_V))
+    NT: gl.constexpr = 64 * gl.num_warps()
+    gl.static_assert(NT >= K and NT >= C)
+    o_t = gl.arange(0, NT, layout=BLK_1D)
     o_c_m = gl.arange(0, C, layout=gl.SliceLayout(1, MMA))
     o_k_m = gl.arange(0, K, layout=gl.SliceLayout(1, MMA))
     o_w_m = i_w * BW + gl.arange(0, BW, layout=gl.SliceLayout(0, MMA))
 
-    kd_off = (o_c_ab[:, None] * K + o_k_ab[None, :]).to(gl.int32)
-    inv_off = (o_c_a[:, None] * C + o_cc_a[None, :]).to(gl.int32)
-    kr_off = (o_c_s[:, None] * K + o_k_s[None, :]).to(gl.int32)
-    gt_off = o_k_m.to(gl.int32)
-    beta_off = (i_h + o_c_m * H).to(gl.int32)
-    v_off = (i_h * V + o_c_m[:, None] * (H * V) + o_w_m[None, :]).to(gl.int32)
+    ws_off = (o_c_s[:, None] * K + o_k_s[None, :]).to(gl.int32)
+    cc_off = (o_r_cc[:, None] * C + o_c_cc[None, :]).to(gl.int32)
+    v_off = (i_h * V + o_c_v[:, None] * (H * V) + o_w_v[None, :]).to(gl.int32)
+
+    inv_ty: gl.constexpr = ws_inv_mqk.dtype.element_ty
+    s_kd = gl.allocate_shared_memory(gl.bfloat16, [2, C, K], SH_WS)
+    s_kr = gl.allocate_shared_memory(gl.bfloat16, [2, C, K], SH_WS)
+    s_inv = gl.allocate_shared_memory(inv_ty, [2, C, C], SH_PLAIN)
+    s_v = gl.allocate_shared_memory(v_input.dtype.element_ty, [2, C, BW], SH_PLAIN)
+    # gt / beta live in LDS until their use, so the per-thread copies of the MMA
+    # slice layout are only materialized there.
+    s_gt = gl.allocate_shared_memory(gl.float32, [2, NT], SH_1D)
+    s_beta = gl.allocate_shared_memory(gl.float32, [2, NT], SH_1D)
 
     h_b = gl.zeros([K, BW], gl.float32, MMA)
     h_a = gl.where(o_k_m[:, None] == o_w_m[None, :], 1.0, 0.0)
 
     ws0 = i_h * TOTAL_TILES + chunk_base
-    inv_ty: gl.constexpr = ws_inv_mqk.dtype.element_ty
+    if n_chunks > 0:
+        _issue_chunk(ws_kd, ws_kr, ws_gt, ws_beta, ws_inv_mqk, v_input, None,
+                     s_kd.index(0), s_kr.index(0), s_inv.index(0), s_v.index(0),
+                     s_gt.index(0), s_beta.index(0), None, None, ws0, tok_base,
+                     tok_end, ws_off, cc_off, v_off, o_c_v, o_t, H, K, V, C,
+                     False)  # fmt: skip
 
     for j in range(n_chunks):
-        ws_idx = ws0 + j
-        ck = ws_idx * (C * K)
-        t0 = tok_base + j * C
-        tb = t0 * H
-        m_c = (t0 + o_c_m) < tok_end
+        cur = j % 2
+        # Chunk j has landed in half `cur` (every thread's copies, once past the
+        # barrier the wait carries), and every thread is done reading the other
+        # half, which held chunk j-1, before it is refilled.
+        gl.amd.cdna4.async_copy.wait_group(0)
+        if j + 1 < n_chunks:
+            _issue_chunk(ws_kd, ws_kr, ws_gt, ws_beta, ws_inv_mqk, v_input, None,
+                         s_kd.index(1 - cur), s_kr.index(1 - cur),
+                         s_inv.index(1 - cur), s_v.index(1 - cur),
+                         s_gt.index(1 - cur), s_beta.index(1 - cur), None, None,
+                         ws0 + j + 1, tok_base + (j + 1) * C, tok_end, ws_off,
+                         cc_off, v_off, o_c_v, o_t, H, K, V, C, False)  # fmt: skip
 
-        kd_a = gl.amd.cdna4.buffer_load(ptr=ws_kd + ck, offsets=kd_off)
-        inv_a = gl.amd.cdna4.buffer_load(
-            ptr=ws_inv_mqk + ws_idx * (2 * C * C), offsets=inv_off
+        kd_a = gl.amd.cdna4.async_copy.load_shared_relaxed(s_kd.index(cur), A_OP_B)
+        # ws_kr is [C, K] in memory; the permuted view is the kr^T A operand.
+        kr_a = gl.amd.cdna4.async_copy.load_shared_relaxed(
+            s_kr.index(cur).permute((1, 0)), A_OP
         )
-        gt = gl.amd.cdna4.buffer_load(ptr=ws_gt + ws_idx * K, offsets=gt_off)
-        kr_raw = gl.amd.cdna4.buffer_load(ptr=ws_kr + ck, offsets=kr_off)
-        beta = _sigmoid(
-            gl.amd.cdna4.buffer_load(
-                ptr=beta_raw + tb, offsets=beta_off, mask=m_c, other=0.0
-            )
+        inv_a = gl.amd.cdna4.async_copy.load_shared_relaxed(s_inv.index(cur), A_OP)
+        b_v = gl.amd.cdna4.async_copy.load_shared_relaxed(s_v.index(cur), MMA).to(
+            gl.float32
         )
-        b_v = gl.amd.cdna4.buffer_load(
-            ptr=v_input + tb * V, offsets=v_off, mask=m_c[:, None], other=0.0
-        ).to(gl.float32)
-
-        kr_a = _kr_operand(kr_raw, K, C, SH_KR, A_OP)
+        gt = gl.amd.cdna4.async_copy.load_shared_relaxed(
+            s_gt.index(cur).slice(0, K), gl.SliceLayout(1, MMA)
+        )
+        beta = gl.amd.cdna4.async_copy.load_shared_relaxed(
+            s_beta.index(cur).slice(0, C), gl.SliceLayout(1, MMA)
+        )
+        m_c = (tok_base + j * C + o_c_m) < tok_end
         # Written one after the other so the scheduler has two independent MFMA
         # chains to interleave, which is what covers the serial dependence.
         h_b, _u, _h = _recur(h_b, kd_a, inv_a, kr_a, gt, beta, b_v, m_c, C, BW,
@@ -227,4 +319,229 @@ def k2_ab_fused_gluon(
     gl.amd.cdna4.buffer_store(h_a.to(h_out_a.dtype.element_ty), h_out_a + s_base, s_off)
 
 
+_k2_c_repr = make_kernel_repr(
+    "k2_c_gluon",
+    ["C", "K", "V", "BW", "HAS_H_IN", "STORE_FINAL", "PAGED_CACHE", "PAGED_H_IN"],
+)
+
+
+@gluon.jit(repr=_k2_c_repr)
+def k2_c_gluon(
+    ws_kd,
+    ws_qd,
+    ws_kr,
+    ws_gt,
+    ws_beta,
+    ws_inv_mqk,
+    v_input,
+    out,
+    h_in,
+    final_state,
+    seg_chunk_base,
+    seg_nchunks,
+    seg_tok_base,
+    seg_tok_end,
+    seg_seq,
+    seg_is_last,
+    state_cache,
+    state_indices,
+    has_initial_state,
+    cache_stride,
+    TOTAL_TILES,
+    H: gl.constexpr,
+    K: gl.constexpr,
+    V: gl.constexpr,
+    C: gl.constexpr,
+    BW: gl.constexpr,
+    MMA: gl.constexpr,
+    A_OP: gl.constexpr,
+    B_OP: gl.constexpr,
+    MMA_B: gl.constexpr,
+    A_OP_B: gl.constexpr,
+    B_OP_B: gl.constexpr,
+    BLK: gl.constexpr,
+    BLK_CC: gl.constexpr,
+    BLK_1D: gl.constexpr,
+    SH_WS: gl.constexpr,
+    SH_PLAIN: gl.constexpr,
+    SH_1D: gl.constexpr,
+    HAS_H_IN: gl.constexpr,
+    STORE_FINAL: gl.constexpr,
+    STATE_V_FIRST: gl.constexpr,
+    PAGED_CACHE: gl.constexpr,
+    PAGED_H_IN: gl.constexpr,
+    CM_OUT: gl.constexpr,
+    NUM_XCDS: gl.constexpr,
+):
+    """Pass C (the output pass): the recurrence of one segment from its incoming state.
+
+    Same contract as ``_flash_kda_segment_kernel`` with ``COMPUTE_OUTPUT`` and
+    ``HAS_V`` and without ``STORE_H_OUT``, paged ``state_cache`` included. Chunk
+    j+1's tiles are prefetched into LDS while chunk j computes, as in the fused
+    pass A. The output is stored straight from the MFMA layout: converting it
+    for a wider store costs an LDS round trip and enough scratch to cost
+    occupancy at BW=128.
+    """
+    n_w = gl.num_programs(0)
+    pid = remap_xcd(
+        gl.program_id(1) * n_w + gl.program_id(0), n_w * gl.num_programs(1), NUM_XCDS
+    )
+    i_w = (pid % n_w).to(gl.int64)
+    i_sh = (pid // n_w).to(gl.int64)
+    i_seg = i_sh // H
+    i_h = i_sh % H
+
+    chunk_base = gl.load(seg_chunk_base + i_seg).to(gl.int64)
+    n_chunks = gl.load(seg_nchunks + i_seg)
+    tok_base = gl.load(seg_tok_base + i_seg).to(gl.int64)
+    tok_end = gl.load(seg_tok_end + i_seg).to(gl.int64)
+    # Read up front, not after the loop: past the output stores the compiler
+    # must assume these words may have changed and loads them per lane, which
+    # leaves the final-state store a waterfall loop over a non-uniform pointer.
+    # When STORE_FINAL is false the final-state pointers are null and these
+    # must not be read.
+    if STORE_FINAL or PAGED_CACHE:
+        i_n = gl.load(seg_seq + i_seg).to(gl.int64)
+    if STORE_FINAL:
+        is_last = gl.load(seg_is_last + i_seg)
+    if PAGED_CACHE:
+        slot = gl.load(state_indices + i_n).to(gl.int64)
+
+    BLK_V: gl.constexpr = gl.BlockedLayout(
+        [1, 8], [64 // (BW // 8), BW // 8], [gl.num_warps(), 1], [1, 0]
+    )
+    o_c_s = gl.arange(0, C, layout=gl.SliceLayout(1, BLK))
+    o_k_s = gl.arange(0, K, layout=gl.SliceLayout(0, BLK))
+    o_r_cc = gl.arange(0, C, layout=gl.SliceLayout(1, BLK_CC))
+    o_c_cc = gl.arange(0, C, layout=gl.SliceLayout(0, BLK_CC))
+    o_c_v = gl.arange(0, C, layout=gl.SliceLayout(1, BLK_V))
+    o_w_v = i_w * BW + gl.arange(0, BW, layout=gl.SliceLayout(0, BLK_V))
+    NT: gl.constexpr = 64 * gl.num_warps()
+    gl.static_assert(NT >= K and NT >= C)
+    o_t = gl.arange(0, NT, layout=BLK_1D)
+    o_c_m = gl.arange(0, C, layout=gl.SliceLayout(1, MMA))
+    o_k_m = gl.arange(0, K, layout=gl.SliceLayout(1, MMA))
+    o_w_m = i_w * BW + gl.arange(0, BW, layout=gl.SliceLayout(0, MMA))
+
+    ws_off = (o_c_s[:, None] * K + o_k_s[None, :]).to(gl.int32)
+    cc_off = (o_r_cc[:, None] * C + o_c_cc[None, :]).to(gl.int32)
+    v_off = (i_h * V + o_c_v[:, None] * (H * V) + o_w_v[None, :]).to(gl.int32)
+    o_off = (i_h * V + o_c_m[:, None] * (H * V) + o_w_m[None, :]).to(gl.int32)
+    # K-first [K, V] and V-first [V, K] offsets of this program's state block.
+    s_off = (o_k_m[:, None] * V + o_w_m[None, :]).to(gl.int32)
+    s_off_vk = (o_w_m[None, :] * K + o_k_m[:, None]).to(gl.int32)
+
+    inv_ty: gl.constexpr = ws_inv_mqk.dtype.element_ty
+    out_ty: gl.constexpr = out.dtype.element_ty
+    s_kd = gl.allocate_shared_memory(gl.bfloat16, [2, C, K], SH_WS)
+    s_qd = gl.allocate_shared_memory(gl.bfloat16, [2, C, K], SH_WS)
+    s_kr = gl.allocate_shared_memory(gl.bfloat16, [2, C, K], SH_WS)
+    s_inv = gl.allocate_shared_memory(inv_ty, [2, C, C], SH_PLAIN)
+    s_mqk = gl.allocate_shared_memory(inv_ty, [2, C, C], SH_PLAIN)
+    s_v = gl.allocate_shared_memory(v_input.dtype.element_ty, [2, C, BW], SH_PLAIN)
+    s_gt = gl.allocate_shared_memory(gl.float32, [2, NT], SH_1D)
+    s_beta = gl.allocate_shared_memory(gl.float32, [2, NT], SH_1D)
+
+    if PAGED_H_IN:
+        # A zero-start sequence's slot may hold a previous owner's state: mask
+        # the load so it is never read, not read and then discarded.
+        take = gl.load(has_initial_state + i_n).to(gl.int1)
+        m_h = (o_k_m[:, None] < K) & (o_w_m[None, :] < V) & take
+        h = gl.amd.cdna4.buffer_load(
+            ptr=state_cache + slot * cache_stride + i_h * (V * K),
+            offsets=s_off_vk,
+            mask=m_h,
+            other=0.0,
+        ).to(gl.float32)
+    elif HAS_H_IN:
+        h = gl.amd.cdna4.buffer_load(
+            ptr=h_in + (i_seg * H + i_h) * (K * V), offsets=s_off
+        ).to(gl.float32)
+    else:
+        h = gl.zeros([K, BW], gl.float32, MMA)
+
+    ws0 = i_h * TOTAL_TILES + chunk_base
+    if n_chunks > 0:
+        _issue_chunk(ws_kd, ws_kr, ws_gt, ws_beta, ws_inv_mqk, v_input, ws_qd,
+                     s_kd.index(0), s_kr.index(0), s_inv.index(0), s_v.index(0),
+                     s_gt.index(0), s_beta.index(0), s_qd.index(0),
+                     s_mqk.index(0), ws0, tok_base, tok_end, ws_off, cc_off,
+                     v_off, o_c_v, o_t, H, K, V, C, True)  # fmt: skip
+
+    for j in range(n_chunks):
+        cur = j % 2
+        # As in pass A: chunk j is in half `cur`, and half 1 - cur is free.
+        gl.amd.cdna4.async_copy.wait_group(0)
+        if j + 1 < n_chunks:
+            _issue_chunk(ws_kd, ws_kr, ws_gt, ws_beta, ws_inv_mqk, v_input, ws_qd,
+                         s_kd.index(1 - cur), s_kr.index(1 - cur),
+                         s_inv.index(1 - cur), s_v.index(1 - cur),
+                         s_gt.index(1 - cur), s_beta.index(1 - cur),
+                         s_qd.index(1 - cur), s_mqk.index(1 - cur), ws0 + j + 1,
+                         tok_base + (j + 1) * C, tok_end, ws_off, cc_off, v_off,
+                         o_c_v, o_t, H, K, V, C, True)  # fmt: skip
+
+        t0 = tok_base + j * C
+        kd_a = gl.amd.cdna4.async_copy.load_shared_relaxed(s_kd.index(cur), A_OP_B)
+        inv_a = gl.amd.cdna4.async_copy.load_shared_relaxed(s_inv.index(cur), A_OP)
+        b_v = gl.amd.cdna4.async_copy.load_shared_relaxed(s_v.index(cur), MMA).to(
+            gl.float32
+        )
+        gt = gl.amd.cdna4.async_copy.load_shared_relaxed(
+            s_gt.index(cur).slice(0, K), gl.SliceLayout(1, MMA)
+        )
+        beta = gl.amd.cdna4.async_copy.load_shared_relaxed(
+            s_beta.index(cur).slice(0, C), gl.SliceLayout(1, MMA)
+        )
+        m_c = (t0 + o_c_m) < tok_end
+
+        # _recur's steps, split around the output projection: updating the state
+        # after it keeps the next state and kr out of registers meanwhile
+        # (1.1-1.2x on this kernel).
+        h_op = gl.convert_layout(h.to(gl.bfloat16), B_OP_B)
+        tmp = gl.convert_layout(
+            gl.amd.cdna4.mfma(kd_a, h_op, gl.zeros([C, BW], gl.float32, MMA_B)), MMA
+        )
+        u = gl.where(m_c[:, None], (b_v - tmp) * beta[:, None], 0.0)
+        big_u = gl.amd.cdna4.mfma(
+            inv_a,
+            gl.convert_layout(u.to(inv_ty), B_OP),
+            gl.zeros([C, BW], gl.float32, MMA),
+        )
+
+        # o = qd @ h + Mqk @ U, h being the state entering the chunk.
+        qd_a = gl.amd.cdna4.async_copy.load_shared_relaxed(s_qd.index(cur), A_OP_B)
+        mqk_a = gl.amd.cdna4.async_copy.load_shared_relaxed(s_mqk.index(cur), A_OP)
+        o = gl.convert_layout(
+            gl.amd.cdna4.mfma(qd_a, h_op, gl.zeros([C, BW], gl.float32, MMA_B)), MMA
+        )
+        o = gl.amd.cdna4.mfma(mqk_a, gl.convert_layout(big_u.to(inv_ty), B_OP), o)
+        gl.amd.cdna4.buffer_store(o.to(out_ty), out + t0 * (H * V), o_off,
+                                  mask=m_c[:, None], cache=CM_OUT)  # fmt: skip
+
+        kr_a = gl.amd.cdna4.async_copy.load_shared_relaxed(
+            s_kr.index(cur).permute((1, 0)), A_OP
+        )
+        h = gl.amd.cdna4.mfma(
+            kr_a, gl.convert_layout(big_u.to(gl.bfloat16), B_OP), h * gt[:, None]
+        )
+
+    if STORE_FINAL:  # noqa: SIM102
+        if is_last == 1:
+            if PAGED_CACHE:
+                gl.amd.cdna4.buffer_store(
+                    h,
+                    state_cache + slot * cache_stride + i_h * (V * K),
+                    s_off_vk,
+                )
+            else:
+                f_ptr = final_state + (i_n * H + i_h) * (K * V)
+                f_val = h.to(final_state.dtype.element_ty)
+                if STATE_V_FIRST:
+                    gl.amd.cdna4.buffer_store(f_val, f_ptr, s_off_vk)
+                else:
+                    gl.amd.cdna4.buffer_store(f_val, f_ptr, s_off)
+
+
 k2_ab_fused_fast = fast_launch(k2_ab_fused_gluon)
+k2_c_fast = fast_launch(k2_c_gluon)
diff --git a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
index 619c506b882..4b29ee3506b 100644
--- a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
+++ b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
@@ -152,6 +152,7 @@ def _flash_kda_prepare_kernel(
     ws_kr,
     ws_gt,
     ws_inv_mqk,
+    ws_beta,
     cu_seqlens,
     chunk_indices,
     scale,
@@ -271,6 +272,8 @@ def _flash_kda_prepare_kernel(
 
     p_beta = beta_raw + (bos + t_off) * H + i_h + o_c * H
     b_beta = tl.sigmoid(tl.load(p_beta, mask=m_c, other=0.0).to(tl.float32))
+    # For the Gluon K2, which copies it into LDS with the chunk's other tiles.
+    tl.store(ws_beta + ws_idx * C + o_c, b_beta, cache_modifier=CM_WS)
 
     # The intra-chunk matrices need decay *differences* between two rows of the
     # same chunk. Those are O(1) near the diagonal even when each row's decay from
@@ -795,11 +798,12 @@ def _log_route(k1_gluon: bool, k2_gluon: bool, C: int, K: int, V: int) -> None:
 # allows. The zero floor keeps the pair self-consistent -- a device with no
 # published table takes the wide branch and finds this same schedule there.
 _K2_GLUON_FALLBACK = triton.Config({"BW": 32, "MIN_BLOCKS_PER_CU": 0}, num_warps=2)
+_K2C_GLUON_FALLBACK = triton.Config({"BW": 64}, num_warps=4)
 
 
 @functools.cache
-def _k2_gluon_schedule(W: int, num_segs: int, H: int) -> tuple[int, int]:
-    """``(BW, num_warps)`` for the Gluon K2, which does not autotune.
+def _k2_gluon_schedule(W: int, num_segs: int, H: int) -> tuple[int, int, int]:
+    """``(BW, num_warps, waves_per_eu)`` for the Gluon pass A, which does not autotune.
 
     The warps can only split BW -- the state is the B operand of ``dot(kd, h)``
     and its K axis is the contraction axis -- so a count above ``BW // 16``
@@ -808,18 +812,36 @@ def _k2_gluon_schedule(W: int, num_segs: int, H: int) -> tuple[int, int]:
     ``cdiv(W, BW) * num_segs * H``: the wide tile reads the per-chunk workspace
     half as many times, but only while there are still enough blocks to fill the
     device, so it is taken above a published blocks-per-CU floor.
+
+    ``waves_per_eu`` (0 leaves it to the compiler) is the occupancy hint: the
+    LDS-prefetching kernel wants ~280 VGPRs, and holding it to 256 (a few
+    spills) for 2 waves per SIMD beats letting it drop to one.
     """
-    wide = chunk_delta_attn_tuned_config(
+    cfg = chunk_delta_attn_tuned_config(
         "k2_ab_fused_gluon_wide", _K2_GLUON_FALLBACK, backend="gluon"
     )
-    bw = wide.kwargs["BW"]
+    bw = cfg.kwargs["BW"]
     blocks = (W // bw) * num_segs * H
-    if W % bw == 0 and blocks >= wide.kwargs["MIN_BLOCKS_PER_CU"] * _num_cus():
-        return bw, wide.num_warps
-    narrow = chunk_delta_attn_tuned_config(
-        "k2_ab_fused_gluon_narrow", _K2_GLUON_FALLBACK, backend="gluon"
+    if W % bw != 0 or blocks < cfg.kwargs["MIN_BLOCKS_PER_CU"] * _num_cus():
+        cfg = chunk_delta_attn_tuned_config(
+            "k2_ab_fused_gluon_narrow", _K2_GLUON_FALLBACK, backend="gluon"
+        )
+    return cfg.kwargs["BW"], cfg.num_warps, cfg.kwargs.get("waves_per_eu", 0)
+
+
+@functools.cache
+def _k2c_gluon_schedule(segmented: bool) -> tuple[int, int]:
+    """``(BW, num_warps)`` for the Gluon pass C.
+
+    Wider tiles amortize a block's workspace reads over more of V, but the
+    widest only pays once segmenting has supplied the blocks.
+    """
+    cfg = chunk_delta_attn_tuned_config(
+        "k2_c_gluon_segmented" if segmented else "k2_c_gluon",
+        _K2C_GLUON_FALLBACK,
+        backend="gluon",
     )
-    return narrow.kwargs["BW"], narrow.num_warps
+    return cfg.kwargs["BW"], cfg.num_warps
 
 
 @functools.cache
@@ -1131,6 +1153,8 @@ def flash_kda_fwd(
     # read twice per chunk on K2's serial path. The CUTLASS FlashKDA picks fp16
     # here for the same reason.
     ws_inv_mqk = torch.empty(H * total_tiles, 2 * C, C, dtype=torch.float16, device=dev)
+    # sigmoid(beta) per tile, for the Gluon K2.
+    ws_beta = torch.empty(H * total_tiles, C, dtype=torch.float32, device=dev)
 
     use_gluon_k1 = AITER_FDA_USE_GLUON_K1 and _gluon_k1_usable(C, K)
     if use_gluon_k1:
@@ -1150,6 +1174,7 @@ def flash_kda_fwd(
             ws_kr=ws_kr,
             ws_gt=ws_gt,
             ws_inv_mqk=ws_inv_mqk,
+            ws_beta=ws_beta,
             cu_seqlens=cu_seqlens,
             chunk_indices=chunk_indices,
             scale=scale,
@@ -1178,6 +1203,7 @@ def flash_kda_fwd(
             ws_kr=ws_kr,
             ws_gt=ws_gt,
             ws_inv_mqk=ws_inv_mqk,
+            ws_beta=ws_beta,
             cu_seqlens=cu_seqlens,
             chunk_indices=chunk_indices,
             scale=scale,
@@ -1248,9 +1274,13 @@ def flash_kda_fwd(
         "cache_stride": cache_stride,
     }
 
-    # Only pass A can go to Gluon, so an unsegmented shape reaches none of it.
-    use_gluon_k2 = AITER_FDA_USE_GLUON_K2 and _gluon_k2_usable(C, K, V) and max_segs > 1
+    # Gluon covers pass A and pass C; an unsegmented shape runs only pass C.
+    use_gluon_k2 = AITER_FDA_USE_GLUON_K2 and _gluon_k2_usable(C, K, V)
     _log_route(use_gluon_k1, use_gluon_k2, C, K, V)
+    if use_gluon_k2:
+        from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn import (
+            flash_kda_k2 as _g2,
+        )
 
     def _launch_k2(*, W, **kw):
         launch_args = {
@@ -1271,18 +1301,14 @@ def _launch_k2(*, W, **kw):
         b_seg = torch.empty(num_segs, H, K, V, dtype=torch.float32, device=dev)
         A_seg = torch.empty(num_segs, H, K, K, dtype=torch.bfloat16, device=dev)
         if use_gluon_k2:
-            from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn import (
-                flash_kda_k2 as _g2,
-            )
-
-            bw, nw = _k2_gluon_schedule(V, num_segs, H)
+            bw, nw, wpe = _k2_gluon_schedule(V, num_segs, H)
             _g2.k2_ab_fused_fast[(triton.cdiv(V, bw), num_segs * H)](
                 ws_kd=ws_kd,
                 ws_kr=ws_kr,
                 ws_gt=ws_gt,
+                ws_beta=ws_beta,
                 ws_inv_mqk=ws_inv_mqk,
                 v_input=v,
-                beta_raw=beta,
                 h_out_b=b_seg,
                 h_out_a=A_seg,
                 seg_chunk_base=seg_chunk_base,
@@ -1298,6 +1324,7 @@ def _launch_k2(*, W, **kw):
                 NUM_XCDS=get_num_xcds(),
                 **_g2.build_layouts(nw),
                 num_warps=nw,
+                waves_per_eu=wpe,
             )
         else:
             for buf, width, identity, has_v in (
@@ -1339,20 +1366,57 @@ def _launch_k2(*, W, **kw):
         h_in = None if paged else h0
 
     # Pass C: re-run each segment from its true incoming state, writing outputs.
-    _launch_k2(
-        v_input=v,
-        out=o,
-        h_in=h_in,
-        h_out=None,
-        final_state=final_state,
-        W=V,
-        INIT_IDENTITY=False,
-        HAS_H_IN=h_in is not None or paged,
-        HAS_V=True,
-        COMPUTE_OUTPUT=True,
-        STORE_H_OUT=False,
-        STORE_FINAL=store_final,
-        **(paged_args if paged else {}),
-    )
+    if use_gluon_k2:
+        bw, nw = _k2c_gluon_schedule(max_segs > 1)
+        _g2.k2_c_fast[(triton.cdiv(V, bw), num_segs * H)](
+            ws_kd=ws_kd,
+            ws_qd=ws_qd,
+            ws_kr=ws_kr,
+            ws_gt=ws_gt,
+            ws_beta=ws_beta,
+            ws_inv_mqk=ws_inv_mqk,
+            v_input=v,
+            out=o,
+            h_in=h_in,
+            final_state=final_state,
+            seg_chunk_base=seg_chunk_base,
+            seg_nchunks=seg_nchunks,
+            seg_tok_base=seg_tok_base,
+            seg_tok_end=seg_tok_end,
+            seg_seq=seg_seq,
+            seg_is_last=seg_is_last,
+            **paged_args,
+            TOTAL_TILES=total_tiles,
+            H=H,
+            K=K,
+            V=V,
+            C=C,
+            BW=bw,
+            **_g2.build_layouts(nw),
+            HAS_H_IN=h_in is not None,
+            STORE_FINAL=store_final,
+            STATE_V_FIRST=state_v_first,
+            PAGED_CACHE=paged,
+            PAGED_H_IN=paged and h_in is None,
+            CM_OUT=CM_OUT_STORE,
+            NUM_XCDS=get_num_xcds(),
+            num_warps=nw,
+        )
+    else:
+        _launch_k2(
+            v_input=v,
+            out=o,
+            h_in=h_in,
+            h_out=None,
+            final_state=final_state,
+            W=V,
+            INIT_IDENTITY=False,
+            HAS_H_IN=h_in is not None or paged,
+            HAS_V=True,
+            COMPUTE_OUTPUT=True,
+            STORE_H_OUT=False,
+            STORE_FINAL=store_final,
+            **(paged_args if paged else {}),
+        )
 
     return o, final_state
diff --git a/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json b/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json
index 000214a6f9e..60a71a250cd 100644
--- a/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json
+++ b/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json
@@ -2,10 +2,20 @@
     "k2_ab_fused_gluon_wide": {
         "BW": 64,
         "num_warps": 4,
-        "MIN_BLOCKS_PER_CU": 2
+        "MIN_BLOCKS_PER_CU": 2,
+        "waves_per_eu": 2
     },
     "k2_ab_fused_gluon_narrow": {
         "BW": 32,
-        "num_warps": 2
+        "num_warps": 2,
+        "waves_per_eu": 2
+    },
+    "k2_c_gluon": {
+        "BW": 64,
+        "num_warps": 4
+    },
+    "k2_c_gluon_segmented": {
+        "BW": 128,
+        "num_warps": 4
     }
 }
diff --git a/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py b/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
index 9f96311a588..b6966fcf5e9 100644
--- a/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
+++ b/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
@@ -424,8 +424,10 @@ def test_published_k2_schedules_can_split_their_tile():
     for W in (64, 128, 256):
         for num_segs in (1, 2, 8, 64, 512):
             for H in (1, 4, 12, 64):
-                reached.add(_flash_kda._k2_gluon_schedule(W, num_segs, H))
+                reached.add(_flash_kda._k2_gluon_schedule(W, num_segs, H)[:2])
     assert len(reached) > 1, f"only one schedule reachable: {reached}"
+    # Pass C splits its tile the same way.
+    reached |= {_flash_kda._k2c_gluon_schedule(seg) for seg in (False, True)}
     for bw, nw in sorted(reached):
         assert bw % 16 == 0, f"BW={bw} is not a whole number of MFMA tiles"
         assert nw <= bw // 16, f"BW={bw} cannot be split {nw} ways"
@@ -463,9 +465,8 @@ def _varlen(lens):
         },
     ),
     "weak gate": lambda: (make_inputs(1, 512, 4), {"lower_bound": -0.01}),
-    # Only a segmented schedule has a pass A, and pass A is the only thing the
-    # Gluon K2 is routed to, so these are the cases where K2's route is visible
-    # at all -- see test_cases_reach_the_gluon_k2.
+    # Only a segmented schedule has a pass A, so these are the cases where the
+    # Gluon pass A is reached at all -- see test_cases_reach_the_gluon_k2.
     "segmented": lambda: (
         make_inputs(1, 1024, 4),
         {"output_final_state": True, "chunks_per_seg": 4},
@@ -520,9 +521,10 @@ def test_routes_agree(case):
 
     Both sides write the same ABI and nothing downstream is told which one ran,
     so a divergence here is a bug in whichever kernel moved rather than a
-    tolerance to widen. Measured across these cases: K2's two implementations
-    agree to the bit, and K1's differ at 3e-4 to 1e-3 on the output and under
-    3e-5 on the state, so the bounds are really about K1.
+    tolerance to widen. Measured across these cases: K1's two implementations
+    differ at 3e-4 to 1e-3 on the output and under 3e-5 on the state, and K2's
+    pass C by an ulp or two of the output, its output projection being ordered
+    differently.
     """
     args, kw = _CASES[case]()
     with _route(k1=False, k2=False):
@@ -538,10 +540,9 @@ def test_routes_agree(case):
 def test_cases_reach_the_gluon_k2():
     """The cases above have to exercise the route they are comparing.
 
-    ``use_gluon_k2`` is tested inside ``max_segs > 1``, so an unsegmented shape
-    runs the Triton K2 whichever way the flag is set and test_routes_agree is
-    comparing it against itself. Every non-segmented case here was in exactly
-    that position, and nothing in the assertions would have said so.
+    Pass C goes to Gluon on every shape, pass A only on segmented ones, so an
+    unsegmented case says nothing about pass A; nothing in the assertions of
+    test_routes_agree would say which kernels a case actually reached.
     """
     if not _flash_kda._gluon_k2_usable(FLASH_KDA_CHUNK, K_DIM, K_DIM):
         pytest.skip("this arch or tile shape never routes K2 to Gluon")
@@ -549,30 +550,38 @@ def test_cases_reach_the_gluon_k2():
         flash_kda_k2 as _g2,
     )
 
-    reached = set()
-    saved = _g2.k2_ab_fused_fast
+    reached = {"k2_ab_fused_fast": set(), "k2_c_fast": set()}
+    saved = {name: getattr(_g2, name) for name in reached}
 
     class _Counting:
+        def __init__(self, name):
+            self.name = name
+
         def __getitem__(self, grid):
-            inner = saved[grid]
+            inner = saved[self.name][grid]
 
             def launch(**kw):
-                reached.add(current)
+                reached[self.name].add(current)
                 return inner(**kw)
 
             return launch
 
-    _g2.k2_ab_fused_fast = _Counting()
+    for name in reached:
+        setattr(_g2, name, _Counting(name))
     try:
         for current, make in _CASES.items():
             args, kw = make()
             with _route(k1=True, k2=True):
                 run_flash(*args, **kw)
     finally:
-        _g2.k2_ab_fused_fast = saved
+        for name, fn in saved.items():
+            setattr(_g2, name, fn)
 
-    want = {name for name in _CASES if name.startswith("segmented")}
-    assert want <= reached, f"never reached the Gluon K2: {sorted(want - reached)}"
+    want_a = {name for name in _CASES if name.startswith("segmented")}
+    missing_a = want_a - reached["k2_ab_fused_fast"]
+    missing_c = set(_CASES) - reached["k2_c_fast"]
+    assert not missing_a, f"never reached the Gluon pass A: {sorted(missing_a)}"
+    assert not missing_c, f"never reached the Gluon pass C: {sorted(missing_c)}"
 
 
 # A weak gate is the only setting that exposes the intra-chunk inverse. At the
@@ -816,9 +825,16 @@ def _kimi_kwargs(q, k, v, g, beta, A_log, dt_bias, scale, **kw):
     }
 
 
+@pytest.mark.parametrize("gluon_k2", [False, True])
 @pytest.mark.parametrize("T,chunks_per_seg", [(256, 0), (1024, 4)])
-def test_paged_cache_matches_dense(T, chunks_per_seg):
-    """In-kernel paged I/O matches gather into a dense V-first state."""
+def test_paged_cache_matches_dense(T, chunks_per_seg, gluon_k2):
+    """In-kernel paged I/O matches gather into a dense V-first state.
+
+    On both K2 routes: the two pass C kernels each carry their own paged loads
+    and stores.
+    """
+    if gluon_k2 and not _flash_kda._gluon_k2_usable(FLASH_KDA_CHUNK, K_DIM, K_DIM):
+        pytest.skip("this arch or tile shape never routes K2 to Gluon")
     H = 4
     q, k, v, g, beta, A_log, dt_bias, scale = make_inputs(1, T, H)
     h0 = torch.randn(1, H, K_DIM, K_DIM, device=device, dtype=torch.float32) * 0.1
@@ -835,7 +851,6 @@ def test_paged_cache_matches_dense(T, chunks_per_seg):
         "state_v_first": True,
         "chunks_per_seg": chunks_per_seg,
     }
-    o_dense, ht = flash_kda_fwd(**common, initial_state=h0, output_final_state=True)
     cache, indices, _storage = _paged_pool(1, H, h0)
     ptr_before = cache.data_ptr()
     out = torch.empty_like(v)
@@ -848,11 +863,13 @@ def test_paged_cache_matches_dense(T, chunks_per_seg):
         "state_indices": indices,
         "has_initial_state": torch.ones(1, device=device, dtype=torch.bool),
     }
-    # First paged launch autotunes a new PAGED_CACHE specialization; compare
-    # after that, not against a trial config.
-    flash_kda_fwd(**paged_kw)
-    cache[indices] = h0
-    o_paged, ht_paged = flash_kda_fwd(**paged_kw)
+    with _route(k2=gluon_k2):
+        o_dense, ht = flash_kda_fwd(**common, initial_state=h0, output_final_state=True)
+        # First paged launch autotunes a new PAGED_CACHE specialization; compare
+        # after that, not against a trial config.
+        flash_kda_fwd(**paged_kw)
+        cache[indices] = h0
+        o_paged, ht_paged = flash_kda_fwd(**paged_kw)
     assert ht_paged is None
     assert o_paged.data_ptr() == out.data_ptr()
     assert cache.data_ptr() == ptr_before, "paged cache must not be packed/copied"

From f10e7842bf36b4be70e5997ff711954a82ea1f2f Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Tue, 6 Oct 2026 16:02:23 +0000
Subject: [PATCH 5/9] [Triton/Gluon] FlashKDA: skip ws_beta off the Gluon K2
 route; name K2 kernels by config

- ws_beta is only read by the Gluon K2, so it is now allocated only when
  K2 takes that route, and both K1s store it only then (STORE_BETA). The
  K2 route is resolved before K1 launches for this.
- The Gluon pass A and pass C names now carry their warp count, through a
  NUM_WARPS constexpr that must equal the launch's num_warps (num_warps
  itself is not a constexpr and would name as NONE), and pass C's also
  carries STATE_V_FIRST.
- test_fast_launch covers the pass C launcher too.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../gfx950/kimi_delta_attn/flash_kda_k1.py    | 15 ++++++----
 .../gfx950/kimi_delta_attn/flash_kda_k2.py    | 29 +++++++++++++++----
 .../kimi_delta_attn/flash_kda.py              | 24 ++++++++++-----
 .../kimi_delta_attn/test_fast_launch.py       |  1 +
 4 files changed, 50 insertions(+), 19 deletions(-)

diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
index d90ea8261a4..02e5534fa86 100644
--- a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
+++ b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
@@ -75,7 +75,7 @@ def _l2norm(x):
 
 _k1_prepare_repr = make_kernel_repr(
     "k1_prepare_gluon",
-    ["C", "K", "BC", "IS_VARLEN", "HAS_BIAS"],
+    ["C", "K", "BC", "IS_VARLEN", "HAS_BIAS", "STORE_BETA"],
 )
 
 
@@ -108,11 +108,12 @@ def k1_prepare_gluon(
     HAS_BIAS: gl.constexpr,
     CM_WS: gl.constexpr = "",
     CM_LOAD: gl.constexpr = ".cg",
+    STORE_BETA: gl.constexpr = False,
 ):
     """FlashKDA K1 on gfx950; same workspace contract as the Triton prepare kernel.
 
-    Also writes ``sigmoid(beta)`` to ``ws_beta``, which the Gluon K2 copies into
-    LDS with the chunk's other tiles.
+    With ``STORE_BETA`` also writes ``sigmoid(beta)`` to ``ws_beta``, which the
+    Gluon K2 copies into LDS with the chunk's other tiles.
     """
     gl.static_assert(C == 32 and K == 128)
     NUM_DOUBLING: gl.constexpr = BC.bit_length() - 2
@@ -252,9 +253,10 @@ def k1_prepare_gluon(
             other=0.0,
         ).to(gl.float32)
     )
-    gl.amd.cdna4.buffer_store(
-        b_beta, ws_beta, (ws_idx * C).to(gl.int32) + o_i_r, cache=CM_WS
-    )
+    if STORE_BETA:
+        gl.amd.cdna4.buffer_store(
+            b_beta, ws_beta, (ws_idx * C).to(gl.int32) + o_i_r, cache=CM_WS
+        )
 
     b_gm = b_gcum - b_gp[None, :]
     b_dec = _exp2(b_gm)
@@ -424,6 +426,7 @@ def gluon_k1_prepare(
         BC=BC,
         IS_VARLEN=cu_seqlens is not None,
         HAS_BIAS=dt_bias is not None,
+        STORE_BETA=ws_beta is not None,
         CM_WS=CM_WS,
         CM_LOAD=CM_LOAD,
         num_warps=_NUM_WARPS,
diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
index 17ff000490a..a1cda53a4b6 100644
--- a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
+++ b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
@@ -165,7 +165,7 @@ def _issue_chunk(
 
 _k2_ab_fused_repr = make_kernel_repr(
     "k2_ab_fused_gluon",
-    ["C", "K", "V", "BW"],
+    ["C", "K", "V", "BW", "NUM_WARPS"],
 )
 
 
@@ -202,6 +202,7 @@ def k2_ab_fused_gluon(
     SH_PLAIN: gl.constexpr,
     SH_1D: gl.constexpr,
     NUM_XCDS: gl.constexpr,
+    NUM_WARPS: gl.constexpr,
 ):
     """Both pass-A recurrences in one launch, sharing every operand load.
 
@@ -235,8 +236,10 @@ def k2_ab_fused_gluon(
     tok_base = gl.load(seg_tok_base + i_seg).to(gl.int64)
     tok_end = gl.load(seg_tok_end + i_seg).to(gl.int64)
 
+    # The launch's warp count, as a constexpr so it can name the kernel.
+    gl.static_assert(NUM_WARPS == gl.num_warps())
     BLK_V: gl.constexpr = gl.BlockedLayout(
-        [1, 8], [64 // (BW // 8), BW // 8], [gl.num_warps(), 1], [1, 0]
+        [1, 8], [64 // (BW // 8), BW // 8], [NUM_WARPS, 1], [1, 0]
     )
     o_c_s = gl.arange(0, C, layout=gl.SliceLayout(1, BLK))
     o_k_s = gl.arange(0, K, layout=gl.SliceLayout(0, BLK))
@@ -244,7 +247,7 @@ def k2_ab_fused_gluon(
     o_c_cc = gl.arange(0, C, layout=gl.SliceLayout(0, BLK_CC))
     o_c_v = gl.arange(0, C, layout=gl.SliceLayout(1, BLK_V))
     o_w_v = i_w * BW + gl.arange(0, BW, layout=gl.SliceLayout(0, BLK_V))
-    NT: gl.constexpr = 64 * gl.num_warps()
+    NT: gl.constexpr = 64 * NUM_WARPS
     gl.static_assert(NT >= K and NT >= C)
     o_t = gl.arange(0, NT, layout=BLK_1D)
     o_c_m = gl.arange(0, C, layout=gl.SliceLayout(1, MMA))
@@ -321,7 +324,18 @@ def k2_ab_fused_gluon(
 
 _k2_c_repr = make_kernel_repr(
     "k2_c_gluon",
-    ["C", "K", "V", "BW", "HAS_H_IN", "STORE_FINAL", "PAGED_CACHE", "PAGED_H_IN"],
+    [
+        "C",
+        "K",
+        "V",
+        "BW",
+        "HAS_H_IN",
+        "STORE_FINAL",
+        "STATE_V_FIRST",
+        "PAGED_CACHE",
+        "PAGED_H_IN",
+        "NUM_WARPS",
+    ],
 )
 
 
@@ -372,6 +386,7 @@ def k2_c_gluon(
     PAGED_H_IN: gl.constexpr,
     CM_OUT: gl.constexpr,
     NUM_XCDS: gl.constexpr,
+    NUM_WARPS: gl.constexpr,
 ):
     """Pass C (the output pass): the recurrence of one segment from its incoming state.
 
@@ -407,8 +422,10 @@ def k2_c_gluon(
     if PAGED_CACHE:
         slot = gl.load(state_indices + i_n).to(gl.int64)
 
+    # The launch's warp count, as a constexpr so it can name the kernel.
+    gl.static_assert(NUM_WARPS == gl.num_warps())
     BLK_V: gl.constexpr = gl.BlockedLayout(
-        [1, 8], [64 // (BW // 8), BW // 8], [gl.num_warps(), 1], [1, 0]
+        [1, 8], [64 // (BW // 8), BW // 8], [NUM_WARPS, 1], [1, 0]
     )
     o_c_s = gl.arange(0, C, layout=gl.SliceLayout(1, BLK))
     o_k_s = gl.arange(0, K, layout=gl.SliceLayout(0, BLK))
@@ -416,7 +433,7 @@ def k2_c_gluon(
     o_c_cc = gl.arange(0, C, layout=gl.SliceLayout(0, BLK_CC))
     o_c_v = gl.arange(0, C, layout=gl.SliceLayout(1, BLK_V))
     o_w_v = i_w * BW + gl.arange(0, BW, layout=gl.SliceLayout(0, BLK_V))
-    NT: gl.constexpr = 64 * gl.num_warps()
+    NT: gl.constexpr = 64 * NUM_WARPS
     gl.static_assert(NT >= K and NT >= C)
     o_t = gl.arange(0, NT, layout=BLK_1D)
     o_c_m = gl.arange(0, C, layout=gl.SliceLayout(1, MMA))
diff --git a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
index 4b29ee3506b..f0a76c92689 100644
--- a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
+++ b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
@@ -122,6 +122,7 @@
     {
         "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
         "HAS_BIAS": lambda args: args["dt_bias"] is not None,
+        "STORE_BETA": lambda args: args["ws_beta"] is not None,
     }
 )
 @triton.autotune(
@@ -170,6 +171,7 @@ def _flash_kda_prepare_kernel(
     HAS_BIAS: tl.constexpr,
     CM_QKG: tl.constexpr = "",
     CM_WS: tl.constexpr = "",
+    STORE_BETA: tl.constexpr = False,
 ):
     """Per-chunk prepare: decayed q/k, gate total, Mqk, and (I - L)^-1."""
     i_t = tl.program_id(0).to(tl.int64)
@@ -273,7 +275,8 @@ def _flash_kda_prepare_kernel(
     p_beta = beta_raw + (bos + t_off) * H + i_h + o_c * H
     b_beta = tl.sigmoid(tl.load(p_beta, mask=m_c, other=0.0).to(tl.float32))
     # For the Gluon K2, which copies it into LDS with the chunk's other tiles.
-    tl.store(ws_beta + ws_idx * C + o_c, b_beta, cache_modifier=CM_WS)
+    if STORE_BETA:
+        tl.store(ws_beta + ws_idx * C + o_c, b_beta, cache_modifier=CM_WS)
 
     # The intra-chunk matrices need decay *differences* between two rows of the
     # same chunk. Those are O(1) near the diagonal even when each row's decay from
@@ -1153,10 +1156,18 @@ def flash_kda_fwd(
     # read twice per chunk on K2's serial path. The CUTLASS FlashKDA picks fp16
     # here for the same reason.
     ws_inv_mqk = torch.empty(H * total_tiles, 2 * C, C, dtype=torch.float16, device=dev)
-    # sigmoid(beta) per tile, for the Gluon K2.
-    ws_beta = torch.empty(H * total_tiles, C, dtype=torch.float32, device=dev)
-
+    # Gluon covers pass A and pass C; an unsegmented shape runs only pass C.
     use_gluon_k1 = AITER_FDA_USE_GLUON_K1 and _gluon_k1_usable(C, K)
+    use_gluon_k2 = AITER_FDA_USE_GLUON_K2 and _gluon_k2_usable(C, K, V)
+    _log_route(use_gluon_k1, use_gluon_k2, C, K, V)
+    # sigmoid(beta) per tile, which only the Gluon K2 reads; the Triton K2
+    # applies the sigmoid to the raw beta itself.
+    ws_beta = (
+        torch.empty(H * total_tiles, C, dtype=torch.float32, device=dev)
+        if use_gluon_k2
+        else None
+    )
+
     if use_gluon_k1:
         from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn.flash_kda_k1 import (
             gluon_k1_prepare,
@@ -1274,9 +1285,6 @@ def flash_kda_fwd(
         "cache_stride": cache_stride,
     }
 
-    # Gluon covers pass A and pass C; an unsegmented shape runs only pass C.
-    use_gluon_k2 = AITER_FDA_USE_GLUON_K2 and _gluon_k2_usable(C, K, V)
-    _log_route(use_gluon_k1, use_gluon_k2, C, K, V)
     if use_gluon_k2:
         from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn import (
             flash_kda_k2 as _g2,
@@ -1323,6 +1331,7 @@ def _launch_k2(*, W, **kw):
                 BW=bw,
                 NUM_XCDS=get_num_xcds(),
                 **_g2.build_layouts(nw),
+                NUM_WARPS=nw,
                 num_warps=nw,
                 waves_per_eu=wpe,
             )
@@ -1400,6 +1409,7 @@ def _launch_k2(*, W, **kw):
             PAGED_H_IN=paged and h_in is None,
             CM_OUT=CM_OUT_STORE,
             NUM_XCDS=get_num_xcds(),
+            NUM_WARPS=nw,
             num_warps=nw,
         )
     else:
diff --git a/op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py b/op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py
index 57c033bdf1c..1ecc39b198c 100644
--- a/op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py
+++ b/op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py
@@ -134,6 +134,7 @@ def _wrapped():
         fk._seg_scan_fast,
         flash_kda_k1._k1_fast,
         flash_kda_k2.k2_ab_fused_fast,
+        flash_kda_k2.k2_c_fast,
     )
     return [c for c in candidates if isinstance(c, fast_launch._FastLaunch)]
 

From 0952d961a77e8b299e4542a7add3a681b670145b Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Wed, 7 Oct 2026 19:59:25 +0000
Subject: [PATCH 6/9] [Triton/Gluon] FlashKDA: keep small unsegmented pass C on
 Triton; bench without safe_gate

- On unsegmented shapes with few (sequence, head) pairs, the serial walk
  over the chunks sets pass C's time, and the Triton kernel's lighter
  iteration at BW=16 beats the Gluon one, which cannot go narrower than
  BW=32: 1.08-1.18x for Triton at <= 32 pairs on MI355X, while Gluon wins
  from 40 pairs up (to 2.2x). The gfx950 k2_c_gluon schedule now publishes
  a MIN_BLOCKS_PER_CU floor, as pass A's does, below which pass C runs on
  Triton; ws_beta is then skipped too. Segmented shapes keep the Gluon
  pass C.
- bench_chunk_delta_attn.py gains --no-safe-gate, which takes the
  token-parallel intra kernel.

bench_flash_kda.py default sweep vs main on MI355X: 1.04-2.61x (1x256x12
and 1x1664x12 were 0.91x and 0.94x before this change).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../kimi_delta_attn/flash_kda.py              | 39 +++++++++++++++----
 .../attention/chunk_delta_attn/DEFAULT.json   |  3 +-
 .../triton/bench_chunk_delta_attn.py          | 13 ++++++-
 .../kimi_delta_attn/test_flash_kda.py         | 29 ++++++++++++--
 4 files changed, 71 insertions(+), 13 deletions(-)

diff --git a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
index f0a76c92689..27269609e87 100644
--- a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
+++ b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
@@ -93,6 +93,10 @@
 _GLUON_BOTH: bool = _GLUON_SEL in ("1", "true", "yes", "on", "all")
 AITER_FDA_USE_GLUON_K1: bool = _GLUON_BOTH or _GLUON_SEL == "k1"
 AITER_FDA_USE_GLUON_K2: bool = _GLUON_BOTH or _GLUON_SEL == "k2"
+# Whether a Gluon K2 still leaves small unsegmented shapes to the Triton pass C,
+# which is faster there (see _k2c_gluon_schedule). Tests clear it to pin pass C
+# to Gluon.
+_K2C_ROUTE_BY_SIZE: bool = True
 
 DEVICE_ARCH = arch_info.get_arch()
 _GLUON_ARCH_OK: bool = DEVICE_ARCH == "gfx950"
@@ -801,7 +805,7 @@ def _log_route(k1_gluon: bool, k2_gluon: bool, C: int, K: int, V: int) -> None:
 # allows. The zero floor keeps the pair self-consistent -- a device with no
 # published table takes the wide branch and finds this same schedule there.
 _K2_GLUON_FALLBACK = triton.Config({"BW": 32, "MIN_BLOCKS_PER_CU": 0}, num_warps=2)
-_K2C_GLUON_FALLBACK = triton.Config({"BW": 64}, num_warps=4)
+_K2C_GLUON_FALLBACK = triton.Config({"BW": 64, "MIN_BLOCKS_PER_CU": 0}, num_warps=4)
 
 
 @functools.cache
@@ -833,8 +837,8 @@ def _k2_gluon_schedule(W: int, num_segs: int, H: int) -> tuple[int, int, int]:
 
 
 @functools.cache
-def _k2c_gluon_schedule(segmented: bool) -> tuple[int, int]:
-    """``(BW, num_warps)`` for the Gluon pass C.
+def _k2c_gluon_schedule(segmented: bool) -> tuple[int, int, float]:
+    """``(BW, num_warps, MIN_BLOCKS_PER_CU)`` for the Gluon pass C.
 
     Wider tiles amortize a block's workspace reads over more of V, but the
     widest only pays once segmenting has supplied the blocks.
@@ -844,7 +848,23 @@ def _k2c_gluon_schedule(segmented: bool) -> tuple[int, int]:
         _K2C_GLUON_FALLBACK,
         backend="gluon",
     )
-    return cfg.kwargs["BW"], cfg.num_warps
+    return cfg.kwargs["BW"], cfg.num_warps, cfg.kwargs.get("MIN_BLOCKS_PER_CU", 0)
+
+
+def _k2c_gluon_usable(segmented: bool, W: int, num_segs: int, H: int) -> bool:
+    """Whether the Gluon pass C beats the Triton one on this grid.
+
+    With few (segment, head) pairs the serial walk over a segment's chunks is
+    what sets the time, and there the Triton kernel's lighter iteration at
+    BW=16 wins; the Gluon one cannot go narrower than BW=32. With more pairs
+    the Gluon kernel's prefetch and fewer workspace re-reads take over. The
+    published ``MIN_BLOCKS_PER_CU`` floor marks the crossover: unsegmented on
+    MI355X it lies between 32 and 40 pairs (1.08-1.18x for Triton below it,
+    up to 2.2x for Gluon above).
+    """
+    bw, _, floor = _k2c_gluon_schedule(segmented)
+    blocks = triton.cdiv(W, bw) * num_segs * H
+    return not _K2C_ROUTE_BY_SIZE or blocks >= floor * _num_cus()
 
 
 @functools.cache
@@ -1160,11 +1180,14 @@ def flash_kda_fwd(
     use_gluon_k1 = AITER_FDA_USE_GLUON_K1 and _gluon_k1_usable(C, K)
     use_gluon_k2 = AITER_FDA_USE_GLUON_K2 and _gluon_k2_usable(C, K, V)
     _log_route(use_gluon_k1, use_gluon_k2, C, K, V)
+    # Unsegmented, there is one segment per sequence, so pass C's route is
+    # known here; segmented, pass A runs on Gluon whichever route pass C takes.
+    use_gluon_k2c = use_gluon_k2 and (segmented or _k2c_gluon_usable(False, V, N, H))
     # sigmoid(beta) per tile, which only the Gluon K2 reads; the Triton K2
     # applies the sigmoid to the raw beta itself.
     ws_beta = (
         torch.empty(H * total_tiles, C, dtype=torch.float32, device=dev)
-        if use_gluon_k2
+        if use_gluon_k2c
         else None
     )
 
@@ -1375,8 +1398,10 @@ def _launch_k2(*, W, **kw):
         h_in = None if paged else h0
 
     # Pass C: re-run each segment from its true incoming state, writing outputs.
-    if use_gluon_k2:
-        bw, nw = _k2c_gluon_schedule(max_segs > 1)
+    if use_gluon_k2c and segmented:
+        use_gluon_k2c = _k2c_gluon_usable(True, V, num_segs, H)
+    if use_gluon_k2c:
+        bw, nw, _ = _k2c_gluon_schedule(segmented)
         _g2.k2_c_fast[(triton.cdiv(V, bw), num_segs * H)](
             ws_kd=ws_kd,
             ws_qd=ws_qd,
diff --git a/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json b/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json
index 60a71a250cd..a8fd078a0e1 100644
--- a/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json
+++ b/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json
@@ -12,7 +12,8 @@
     },
     "k2_c_gluon": {
         "BW": 64,
-        "num_warps": 4
+        "num_warps": 4,
+        "MIN_BLOCKS_PER_CU": 0.3
     },
     "k2_c_gluon_segmented": {
         "BW": 128,
diff --git a/op_tests/op_benchmarks/triton/bench_chunk_delta_attn.py b/op_tests/op_benchmarks/triton/bench_chunk_delta_attn.py
index 2250a638e5a..45a6c96b90e 100644
--- a/op_tests/op_benchmarks/triton/bench_chunk_delta_attn.py
+++ b/op_tests/op_benchmarks/triton/bench_chunk_delta_attn.py
@@ -14,6 +14,9 @@
 # Single shape: B=2 T=4096 H=16 K=64 V=64
 python bench_chunk_delta_attn.py --shape 2 4096 16 64 64
 
+# Without safe_gate (unbounded gate), which takes the token-parallel intra kernel
+python bench_chunk_delta_attn.py --no-safe-gate
+
 # Save CSV
 python bench_chunk_delta_attn.py -o
 """
@@ -141,8 +144,8 @@ def fn(
                 use_gate_in_kernel=True,
                 A_log=A_log,
                 dt_bias=dt_bias,
-                lower_bound=-5.0,
-                safe_gate=True,
+                lower_bound=None if args.no_safe_gate else -5.0,
+                safe_gate=not args.no_safe_gate,
                 use_qk_l2norm_in_kernel=True,
                 use_beta_sigmoid_in_kernel=True,
             )
@@ -190,6 +193,12 @@ def parse_args(args=None):
         metavar=("B", "T", "H", "K", "V"),
         help="Single shape to benchmark instead of the default sweep.",
     )
+    parser.add_argument(
+        "--no-safe-gate",
+        action="store_true",
+        help="Run without safe_gate (and without lower_bound), which routes the "
+        "intra-chunk step to the token-parallel kernel instead of the sub-chunk one.",
+    )
     parser.add_argument(
         "-o",
         action="store_true",
diff --git a/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py b/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
index b6966fcf5e9..31de29179fc 100644
--- a/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
+++ b/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
@@ -76,19 +76,26 @@ def _route(k1: bool | None = None, k2: bool | None = None):
     flags, so a test that wants the other route sets the flags rather than the
     variable. Both default to on wherever the tile shape and the arch allow it,
     which leaves the Triton kernels unreached by every test that does not come
-    through here.
+    through here. ``k2=True`` also pins pass C to Gluon on the small
+    unsegmented shapes the dispatch would otherwise leave to Triton.
     """
-    saved = _flash_kda.AITER_FDA_USE_GLUON_K1, _flash_kda.AITER_FDA_USE_GLUON_K2
+    saved = (
+        _flash_kda.AITER_FDA_USE_GLUON_K1,
+        _flash_kda.AITER_FDA_USE_GLUON_K2,
+        _flash_kda._K2C_ROUTE_BY_SIZE,
+    )
     if k1 is not None:
         _flash_kda.AITER_FDA_USE_GLUON_K1 = k1
     if k2 is not None:
         _flash_kda.AITER_FDA_USE_GLUON_K2 = k2
+        _flash_kda._K2C_ROUTE_BY_SIZE = not k2
     try:
         yield
     finally:
         (
             _flash_kda.AITER_FDA_USE_GLUON_K1,
             _flash_kda.AITER_FDA_USE_GLUON_K2,
+            _flash_kda._K2C_ROUTE_BY_SIZE,
         ) = saved
 
 
@@ -427,7 +434,7 @@ def test_published_k2_schedules_can_split_their_tile():
                 reached.add(_flash_kda._k2_gluon_schedule(W, num_segs, H)[:2])
     assert len(reached) > 1, f"only one schedule reachable: {reached}"
     # Pass C splits its tile the same way.
-    reached |= {_flash_kda._k2c_gluon_schedule(seg) for seg in (False, True)}
+    reached |= {_flash_kda._k2c_gluon_schedule(seg)[:2] for seg in (False, True)}
     for bw, nw in sorted(reached):
         assert bw % 16 == 0, f"BW={bw} is not a whole number of MFMA tiles"
         assert nw <= bw // 16, f"BW={bw} cannot be split {nw} ways"
@@ -584,6 +591,22 @@ def launch(**kw):
     assert not missing_c, f"never reached the Gluon pass C: {sorted(missing_c)}"
 
 
+def test_small_unsegmented_pass_c_stays_on_triton():
+    """The published floor sends few-pair unsegmented shapes to the Triton pass C.
+
+    Below it the serial walk sets the time and the Triton kernel's narrower tile
+    wins; above it, and whenever segmenting has supplied the blocks, Gluon does.
+    """
+    if not _flash_kda._gluon_k2_usable(FLASH_KDA_CHUNK, K_DIM, K_DIM):
+        pytest.skip("this arch or tile shape never routes K2 to Gluon")
+    if _flash_kda._k2c_gluon_schedule(False)[2] == 0:
+        pytest.skip("this arch publishes no pass C floor")
+    usable = _flash_kda._k2c_gluon_usable
+    assert not usable(False, K_DIM, 1, 12)
+    assert usable(False, K_DIM, 8, 16)
+    assert usable(True, K_DIM, 16, 4)
+
+
 # A weak gate is the only setting that exposes the intra-chunk inverse. At the
 # Kimi default lower_bound=-5 the decay damps L to almost nothing, so even an
 # outright wrong inverse stays inside bf16 noise; these three catch it.

From 5b6f66a0452a0dc12694c9635ac737f423ddce2a Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Tue, 13 Oct 2026 16:39:50 +0000
Subject: [PATCH 7/9] [Triton/Gluon] FlashKDA: publish the Gluon K1 occupancy
 hint in the gfx950 config

K1's waves_per_eu=3 was pinned in Python while pass A's occupancy hints are
published in configs/gfx950/gluon/.../chunk_delta_attn/DEFAULT.json. Move it
there (k1_prepare_gluon) and load it through chunk_delta_attn_tuned_config,
so retuning it is a config change. The value holds K1 to ~168 VGPRs, i.e. 3
waves per SIMD; left alone it lands at 170 and drops to 2 (0.85x).

The warp count stays fixed by K1's layouts: the published entry restates it,
and an entry naming another count is rejected rather than launched against
them. Devices without an entry get no hint. Compiled K1 on gfx950 is
unchanged (waves_per_eu=3, 164 VGPRs).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../gfx950/kimi_delta_attn/flash_kda_k1.py    | 35 ++++++++++++++++---
 .../attention/chunk_delta_attn/DEFAULT.json   |  4 +++
 .../kimi_delta_attn/test_flash_kda.py         | 23 ++++++++++++
 3 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
index 02e5534fa86..b5aed32527c 100644
--- a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
+++ b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
@@ -1,11 +1,16 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
+import functools
 import math
 
+import triton
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
 
+from aiter.ops.triton._triton_kernels.kimi_delta_attn.chunk_delta_attn_utils import (
+    chunk_delta_attn_tuned_config,
+)
 from aiter.ops.triton._triton_kernels.kimi_delta_attn.fast_launch import fast_launch
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
 
@@ -364,9 +369,28 @@ def k1_prepare_gluon(
 
 
 _NUM_WARPS = math.prod(_MMA_F16.warps_per_cta)
-# Holds K1 to 168 VGPRs; left alone it lands at 170 and drops to 2 waves per
-# SIMD (0.85x).
-_WAVES_PER_EU = 3
+# No occupancy hint where none is published; the warp count is fixed by the
+# layouts above.
+_K1_FALLBACK = triton.Config({"waves_per_eu": 0}, num_warps=_NUM_WARPS)
+
+
+@functools.cache
+def _k1_schedule() -> tuple[int, int]:
+    """``(num_warps, waves_per_eu)`` for the Gluon K1, from the published config.
+
+    Only ``waves_per_eu`` is tunable: the layouts are built for ``_NUM_WARPS``,
+    so a published entry naming another warp count is rejected rather than
+    launched against them.
+    """
+    cfg = chunk_delta_attn_tuned_config(
+        "k1_prepare_gluon", _K1_FALLBACK, backend="gluon"
+    )
+    if cfg.num_warps != _NUM_WARPS:
+        raise ValueError(
+            f"k1_prepare_gluon is built for num_warps={_NUM_WARPS}, "
+            f"but the published config sets {cfg.num_warps}"
+        )
+    return cfg.num_warps, cfg.kwargs.get("waves_per_eu", 0)
 
 
 _k1_fast = fast_launch(k1_prepare_gluon)
@@ -400,6 +424,7 @@ def gluon_k1_prepare(
     CM_WS="",
     CM_LOAD=".cg",
 ):
+    num_warps, waves_per_eu = _k1_schedule()
     return _k1_fast[(TOTAL_TILES if cu_seqlens is not None else NT, B * H)](
         q=q,
         k=k,
@@ -429,6 +454,6 @@ def gluon_k1_prepare(
         STORE_BETA=ws_beta is not None,
         CM_WS=CM_WS,
         CM_LOAD=CM_LOAD,
-        num_warps=_NUM_WARPS,
-        waves_per_eu=_WAVES_PER_EU,
+        num_warps=num_warps,
+        waves_per_eu=waves_per_eu,
     )
diff --git a/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json b/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json
index a8fd078a0e1..19e92631f0b 100644
--- a/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json
+++ b/aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json
@@ -1,4 +1,8 @@
 {
+    "k1_prepare_gluon": {
+        "num_warps": 2,
+        "waves_per_eu": 3
+    },
     "k2_ab_fused_gluon_wide": {
         "BW": 64,
         "num_warps": 4,
diff --git a/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py b/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
index 31de29179fc..8d2394fd36f 100644
--- a/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
+++ b/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
@@ -18,6 +18,7 @@
 
 import pytest
 import torch
+import triton
 
 from aiter.ops.triton._triton_kernels.kimi_delta_attn import chunk_delta_attn_fwd
 from aiter.ops.triton._triton_kernels.kimi_delta_attn import chunk_fwd as _chunk_fwd
@@ -440,6 +441,28 @@ def test_published_k2_schedules_can_split_their_tile():
         assert nw <= bw // 16, f"BW={bw} cannot be split {nw} ways"
 
 
+def test_k1_schedule_comes_from_the_config(monkeypatch):
+    """K1's occupancy hint is read from the published config, and a published
+    warp count its layouts were not built for is rejected, not launched."""
+    if not _flash_kda._gluon_k1_usable(FLASH_KDA_CHUNK, K_DIM):
+        pytest.skip("this arch never routes K1 to Gluon")
+    from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn import (
+        flash_kda_k1 as _g1,
+    )
+
+    assert _g1._k1_schedule()[0] == _g1._NUM_WARPS
+
+    wrong = triton.Config({"waves_per_eu": 0}, num_warps=2 * _g1._NUM_WARPS)
+    monkeypatch.setattr(_g1, "chunk_delta_attn_tuned_config", lambda *a, **kw: wrong)
+    _g1._k1_schedule.cache_clear()
+    try:
+        with pytest.raises(ValueError, match="num_warps"):
+            _g1._k1_schedule()
+    finally:
+        monkeypatch.undo()
+        _g1._k1_schedule.cache_clear()
+
+
 def _varlen(lens):
     return torch.tensor(
         [0] + list(torch.tensor(lens).cumsum(0)), device=device, dtype=torch.long

From 2da6f512277002eb7a789c5e774b11d8616d2ee3 Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Fri, 23 Oct 2026 20:32:11 +0000
Subject: [PATCH 8/9] =?UTF-8?q?[Triton/Gluon]=20KDA:=20address=20review=20?=
 =?UTF-8?q?=E2=80=94=20shared=20Gluon=20helpers,=20align=20wrapper=20with?=
 =?UTF-8?q?=20#6609?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Move sigmoid_log2/fast_sigmoid/sum_combine/l2norm into
  _gluon_kernels/common/utils.py (same definitions as #6609); k1 uses them.
- Align flash_kda.py with #6609: _k2_gluon_schedule(..., kernel) returns a
  Config, use_gluon_k2a/k2c split, _GLUON_K2C_ARCH_OK, unpaged_args; gfx950
  k2 gets gluon_k2_ab_fused/gluon_k2_c launchers mirroring gfx1250's.
- Tests: arch-generic Gluon module imports, pass-C checks gated on arch.
- Fix remap_xcd import path (moved to _triton_kernels/common in #6505).
- Trim verbose comments.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../ops/triton/_gluon_kernels/common/utils.py |  22 +++
 .../gfx950/kimi_delta_attn/flash_kda_k1.py    | 123 ++++--------
 .../gfx950/kimi_delta_attn/flash_kda_k2.py    | 187 ++++++++++++++----
 .../kimi_delta_attn/flash_kda.py              | 172 +++++++---------
 .../kimi_delta_attn/intra_attn.py             |   9 +-
 .../kimi_delta_attn/test_fast_launch.py       |  18 +-
 .../kimi_delta_attn/test_flash_kda.py         |  87 ++++----
 7 files changed, 328 insertions(+), 290 deletions(-)

diff --git a/aiter/ops/triton/_gluon_kernels/common/utils.py b/aiter/ops/triton/_gluon_kernels/common/utils.py
index 2b306b3bf6f..75ae2b3ebef 100644
--- a/aiter/ops/triton/_gluon_kernels/common/utils.py
+++ b/aiter/ops/triton/_gluon_kernels/common/utils.py
@@ -31,6 +31,28 @@ def sigmoid(x):
     return 0.5 + 0.5 * libdevice.tanh(0.5 * x)
 
 
+@gluon.jit
+def sigmoid_log2(z):
+    """sigmoid(x) given ``z = -x * log2(e)``."""
+    return gl.extra.libdevice.fast_dividef(1.0, 1.0 + gl.exp2(z))
+
+
+@gluon.jit
+def fast_sigmoid(x):
+    return sigmoid_log2(x.to(gl.float32) * -math.log2(math.e))
+
+
+@gluon.jit
+def sum_combine(a, b):
+    return a + b
+
+
+@gluon.jit
+def l2norm(x):
+    f = x.to(gl.float32)
+    return f * gl.rsqrt(gl.sum(f * f, axis=1) + 1e-6)[:, None]
+
+
 @gluon.jit
 def mx_e8m0_scale(amax, EXP_OFFSET: gl.constexpr):
     """Biased e8m0 scale byte from fp32 block amax; EXP_OFFSET is 2 for e2m1, 8 for e4m3."""
diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
index b5aed32527c..c3aae647f1e 100644
--- a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
+++ b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
@@ -8,14 +8,19 @@
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
 
+from aiter.ops.triton._gluon_kernels.common.utils import (
+    fast_sigmoid,
+    l2norm,
+    sigmoid_log2,
+    sum_combine,
+)
 from aiter.ops.triton._triton_kernels.kimi_delta_attn.chunk_delta_attn_utils import (
     chunk_delta_attn_tuned_config,
 )
 from aiter.ops.triton._triton_kernels.kimi_delta_attn.fast_launch import fast_launch
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
 
-# [C, K] tiles: warps split the rows, so the l2norm row sums stay within a warp
-# (a cross-warp sum costs three barriers).
+# [C, K] tiles: warps split the rows, so the l2norm row sums stay within a warp.
 _BLK_CK: gl.constexpr = gl.BlockedLayout([1, 8], [4, 16], [2, 1], [1, 0])
 _BLK1: gl.constexpr = gl.BlockedLayout([1], [64], [2], [0])
 # One whole C = 32 column per thread, for the gate's cumulative sum.
@@ -42,42 +47,6 @@
 _LOG2E = gl.constexpr(1.4426950408889634)
 
 
-@gluon.jit
-def _add(a, b):
-    return a + b
-
-
-@gluon.jit
-def _exp(x):
-    return gl.exp(x.to(gl.float32))
-
-
-@gluon.jit
-def _exp2(x):
-    return gl.exp2(x.to(gl.float32))
-
-
-@gluon.jit
-def _sigmoid_log2(z):
-    """sigmoid(x) given ``z = -x * log2(e)``.
-
-    exp2 lowers to a bare v_exp_f32, where exp adds a denormal-range rescale
-    that the sigmoid does not need.
-    """
-    return gl.extra.libdevice.fast_dividef(1.0, 1.0 + gl.exp2(z))
-
-
-@gluon.jit
-def _sigmoid(x):
-    return _sigmoid_log2(x.to(gl.float32) * -_LOG2E)
-
-
-@gluon.jit
-def _l2norm(x):
-    f = x.to(gl.float32)
-    return f * gl.rsqrt(gl.sum(f * f, axis=1) + 1e-6)[:, None]
-
-
 _k1_prepare_repr = make_kernel_repr(
     "k1_prepare_gluon",
     ["C", "K", "BC", "IS_VARLEN", "HAS_BIAS", "STORE_BETA"],
@@ -115,11 +84,8 @@ def k1_prepare_gluon(
     CM_LOAD: gl.constexpr = ".cg",
     STORE_BETA: gl.constexpr = False,
 ):
-    """FlashKDA K1 on gfx950; same workspace contract as the Triton prepare kernel.
-
-    With ``STORE_BETA`` also writes ``sigmoid(beta)`` to ``ws_beta``, which the
-    Gluon K2 copies into LDS with the chunk's other tiles.
-    """
+    """FlashKDA K1 on gfx950, same workspace contract as the Triton prepare kernel.
+    ``STORE_BETA`` also writes ``sigmoid(beta)`` to ``ws_beta`` for the Gluon K2."""
     gl.static_assert(C == 32 and K == 128)
     NUM_DOUBLING: gl.constexpr = BC.bit_length() - 2
     NUM_MERGE: gl.constexpr = (C // BC).bit_length() - 1
@@ -168,15 +134,13 @@ def k1_prepare_gluon(
         ptr=k, offsets=qk_off, mask=m_ck, other=0.0, cache=CM_LOAD
     )
 
-    # The gate is cumulated over rows, so it is computed with each thread owning
-    # a whole column: the scan and the row picks below are then in-thread.
+    # One column per thread, so the gate's scan and row picks are in-thread.
     o_c_col = gl.arange(0, C, layout=gl.SliceLayout(1, _BLK_COL))
     o_k_col = gl.arange(0, K, layout=gl.SliceLayout(0, _BLK_COL))
     g_off = (base * K + o_c_col[:, None].to(gl.int64) * (H * K) + o_k_col[None, :]).to(
         gl.int32
     )
-    # Masked: the buffer resource has no size, so rows past the tensor's end are
-    # not clamped by hardware and would fault.
+    # Masked: rows past the tensor's end are not clamped by hardware.
     b_g = gl.amd.cdna4.buffer_load(
         ptr=g_raw,
         offsets=g_off,
@@ -185,7 +149,7 @@ def k1_prepare_gluon(
         cache=CM_LOAD,
     ).to(gl.float32)
     # sigmoid(e^A * (g + bias)) with -log2(e) * e^A folded into one scalar.
-    rate = _exp(gl.load(A_log + i_h)) * -_LOG2E
+    rate = gl.exp(gl.load(A_log + i_h).to(gl.float32)) * -_LOG2E
     if HAS_BIAS:
         bias = gl.amd.cdna4.buffer_load(
             ptr=dt_bias, offsets=(i_h * K).to(gl.int32) + o_k_col
@@ -193,35 +157,30 @@ def k1_prepare_gluon(
         b_z = b_g * rate + (bias * rate)[None, :]
     else:
         b_z = b_g * rate
-    # Tail rows get a zero gate (g = 0 alone would give lower_bound / 2, plus the
-    # bias), so the cumulative gate is constant from row actual_len - 1 on: the
-    # last row is row C - 1, and the pivot row min(C / 2, actual_len - 1) is
-    # row C / 2.
+    # Tail rows get a zero gate, so the cumulative gate does not change after
+    # the last valid row. Rows C - 1 and C / 2 give the last and pivot values.
     b_gate = gl.where(
         o_c_col[:, None] < actual_len,
-        (lower_bound * _LOG2E) * _sigmoid_log2(b_z),
+        (lower_bound * _LOG2E) * sigmoid_log2(b_z),
         0.0,
     )
-    b_gcum_col = gl.associative_scan(b_gate, 0, _add)
+    b_gcum_col = gl.associative_scan(b_gate, 0, sum_combine)
     b_g_last_col = gl.sum(gl.where(o_c_col[:, None] == C - 1, b_gcum_col, 0.0), axis=0)
     b_gp_col = gl.sum(gl.where(o_c_col[:, None] == C // 2, b_gcum_col, 0.0), axis=0)
-    # One LDS round trip (one barrier) takes the cumulative gate and its last and
-    # pivot rows to the row layout; separate conversions cost one or two
-    # barriers each.
+    # One LDS round trip (one barrier) for all three, not one conversion each.
     s_gcum = gl.allocate_shared_memory(gl.float32, [C, K], _SH_ROWS, b_gcum_col)
     s_g_last = gl.allocate_shared_memory(gl.float32, [K], _SH_VEC, b_g_last_col)
     s_gp = gl.allocate_shared_memory(gl.float32, [K], _SH_VEC, b_gp_col)
     b_gcum = s_gcum.load(_BLK_CK)
     b_g_last = s_g_last.load(gl.SliceLayout(0, _BLK_CK))
     b_gp = s_gp.load(gl.SliceLayout(0, _BLK_CK))
-    b_g_total = _exp2(b_g_last_col)
-    b_exp_g = _exp2(b_gcum)
+    b_g_total = gl.exp2(b_g_last_col)
+    b_exp_g = gl.exp2(b_gcum)
 
-    # Tail rows of q and k load as 0 and stay 0 through the norm. Every factor
-    # they meet below is finite (the cumulative gate is constant over the tail),
-    # so the workspace tiles' tail rows come out zero without a mask.
-    b_q = _l2norm(b_q_raw)
-    b_k = _l2norm(b_k_raw)
+    # Tail rows of q and k load as 0 and all factors below are finite, so the
+    # tail rows store as 0 without a mask.
+    b_q = l2norm(b_q_raw)
+    b_k = l2norm(b_k_raw)
 
     ws_idx = i_h * TOTAL_TILES + g_tile
     ck_off = (ws_idx * (C * K) + o_c[:, None].to(gl.int64) * K + o_k[None, :]).to(
@@ -239,7 +198,7 @@ def k1_prepare_gluon(
         ck_off,
         cache=CM_WS,
     )
-    b_kr_val = (b_k * _exp2(b_g_last[None, :] - b_gcum)).to(gl.bfloat16)
+    b_kr_val = (b_k * gl.exp2(b_g_last[None, :] - b_gcum)).to(gl.bfloat16)
     gl.amd.cdna4.buffer_store(
         b_kr_val.to(ws_kr.dtype.element_ty), ws_kr, ck_off, cache=CM_WS
     )
@@ -250,7 +209,7 @@ def k1_prepare_gluon(
         cache=CM_WS,
     )
 
-    b_beta = _sigmoid(
+    b_beta = fast_sigmoid(
         gl.amd.cdna4.buffer_load(
             ptr=beta_raw,
             offsets=(base.to(gl.int32) + o_i_r * H),
@@ -264,14 +223,13 @@ def k1_prepare_gluon(
         )
 
     b_gm = b_gcum - b_gp[None, :]
-    b_dec = _exp2(b_gm)
-    b_inc = _exp2(-b_gm)
+    b_dec = gl.exp2(b_gm)
+    b_inc = gl.exp2(-b_gm)
     b_k_piv = (b_k * b_dec).to(gl.bfloat16)
     b_q_piv = (b_q * b_dec * scale).to(gl.bfloat16)
     b_k_inv = (b_k * b_inc).to(gl.bfloat16)
 
-    # All three operands go to LDS before any is read back: one barrier for the
-    # lot instead of two per operand.
+    # All three operands go to LDS before any is read back: one barrier for all.
     s_kinv = gl.allocate_shared_memory(
         gl.bfloat16, [K, C], _SH_B, gl.permute(b_k_inv, 1, 0)
     )
@@ -289,8 +247,7 @@ def k1_prepare_gluon(
     )
     b_Mqk = gl.where(o_i_r[:, None] >= o_i_c[None, :], b_Mqk, 0.0)
 
-    # Mqk and INV are stored from the MFMA layout (four consecutive columns per
-    # thread): converting them for a wider store costs two barriers each.
+    # Stored from the MFMA layout; a wider-store conversion costs two barriers.
     cc_base = (ws_idx * (2 * C * C)).to(gl.int32)
     gl.amd.cdna4.buffer_store(
         b_Mqk.to(ws_inv_mqk.dtype.element_ty),
@@ -308,12 +265,8 @@ def k1_prepare_gluon(
         gl.where(o_i_r[:, None] == o_i_c[None, :], 1.0, b_D), _MMA_F16
     )
     zero_cc = gl.zeros([C, C], gl.float32, _MMA_F16)
-    # Each product's operands go through LDS once: stored together, one barrier,
-    # then loaded in each operand layout they are used in. Steps alternate
-    # between two buffer pairs, so a step's stores never land on what the
-    # previous step is still reading. Barriers are placed per allocation, so
-    # every operand has its own: two stores into one buffer cost a barrier
-    # between them.
+    # Steps alternate between two buffer pairs, so the stores of a step do not
+    # overwrite what the previous step still reads.
     s_a0 = gl.allocate_shared_memory(gl.float32, [C, C], _SH_CC_F)
     s_b0 = gl.allocate_shared_memory(gl.float32, [C, C], _SH_CC_F)
     s_a1 = gl.allocate_shared_memory(gl.float32, [C, C], _SH_CC_F)
@@ -334,8 +287,7 @@ def k1_prepare_gluon(
         b_Dp = gl.amd.cdna4.mfma(s_dp.load(_AF16), dp_b, zero_cc)
 
     # The merge steps keep to fixed buffers: off and INV in the pair the doubling
-    # would use next, inner in the other pair. The first pair was last read
-    # before the previous step's barrier, the second before this step's first.
+    # would use next, inner in the other pair.
     if NUM_DOUBLING % 2 == 0:
         s_off = s_a0
         s_inv = s_b0
@@ -369,19 +321,14 @@ def k1_prepare_gluon(
 
 
 _NUM_WARPS = math.prod(_MMA_F16.warps_per_cta)
-# No occupancy hint where none is published; the warp count is fixed by the
-# layouts above.
+# The warp count is fixed by the layouts above.
 _K1_FALLBACK = triton.Config({"waves_per_eu": 0}, num_warps=_NUM_WARPS)
 
 
 @functools.cache
 def _k1_schedule() -> tuple[int, int]:
-    """``(num_warps, waves_per_eu)`` for the Gluon K1, from the published config.
-
-    Only ``waves_per_eu`` is tunable: the layouts are built for ``_NUM_WARPS``,
-    so a published entry naming another warp count is rejected rather than
-    launched against them.
-    """
+    """``(num_warps, waves_per_eu)`` for the Gluon K1. Only ``waves_per_eu`` is
+    tunable; a config naming another warp count is rejected."""
     cfg = chunk_delta_attn_tuned_config(
         "k1_prepare_gluon", _K1_FALLBACK, backend="gluon"
     )
diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
index a1cda53a4b6..e1d53c7ba75 100644
--- a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
+++ b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
@@ -3,12 +3,14 @@
 
 import functools
 
+import triton
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
 
+from aiter.ops.triton._triton_kernels.common.pid_preprocessing import remap_xcd
 from aiter.ops.triton._triton_kernels.kimi_delta_attn.fast_launch import fast_launch
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
-from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
+from aiter.ops.triton.utils.device_info import get_num_xcds
 
 KW = 8
 KW_BIG = 8
@@ -44,8 +46,7 @@ def build_layouts(nw, kw=KW, kw_big=KW_BIG):
         "MMA_B": mma_b,
         "A_OP_B": gl.DotOperandLayout(0, mma_b, kw_big),
         "B_OP_B": gl.DotOperandLayout(1, mma_b, kw_big),
-        # Sources of the global -> LDS copies: 128 (or 32) bits per thread and
-        # no lane replication, as buffer_load_to_shared requires.
+        # Global -> LDS copy sources: no lane replication, as the copy requires.
         "BLK": gl.BlockedLayout([1, 8], [4, 16], [nw, 1], [1, 0]),
         "BLK_CC": gl.BlockedLayout([1, 2], [4, 16], [nw, 1], [1, 0]),
         # gt and beta: one element per thread.
@@ -126,14 +127,10 @@ def _issue_chunk(
     C: gl.constexpr,
     WITH_OUTPUT: gl.constexpr,
 ):
-    """Start fetching one chunk into LDS by async copy, as one commit group.
-
-    ``gt`` and ``beta`` (K1 applied its sigmoid) have fewer elements than
-    threads, and the async copy cannot lower lane replication, so they are
-    copied one element per thread (``o_t``) with the excess threads masked.
-    Tail rows of ``v`` are masked here; ``beta`` of a tail row is harmless
-    because the recurrence zeroes those rows of U anyway. ``WITH_OUTPUT``
-    (pass C) also fetches the output projection's ``qd`` and ``Mqk`` tiles.
+    """Start the async copy of one chunk into LDS, as one commit group.
+
+    ``gt`` and ``beta`` are copied one element per thread (``o_t``), excess
+    threads masked. ``WITH_OUTPUT`` (pass C) also fetches ``qd`` and ``Mqk``.
     """
     tb = t0 * H
     gl.amd.cdna4.async_copy.buffer_load_to_shared(
@@ -212,16 +209,10 @@ def k2_ab_fused_gluon(
     independent, so the scheduler interleaves them, and that is what covers the
     serial dependence each one has on its own. Requires K == V, since the two
     states are then the same width and one program covers a column block of
-    both.
-
-    Chunk j+1's operands are fetched while chunk j computes: its tiles by async
-    global -> LDS copy into the other half of a double buffer, so the prefetch
-    costs no registers (a register prefetch halves occupancy). Without it the
-    kernel spent a third or more of its wave-cycles waiting on loads issued and
-    consumed in the same iteration.
+    both. Chunk j+1 is async-copied into the other half of an LDS double buffer
+    while chunk j computes.
     """
-    # Keeps a (segment, head)'s V blocks on one XCD, so their re-reads of the
-    # chunk workspace share an L2. See the Triton K2 for the full reasoning.
+    # Keeps a (segment, head)'s V blocks on one XCD, sharing an L2.
     n_w = gl.num_programs(0)
     pid = remap_xcd(
         gl.program_id(1) * n_w + gl.program_id(0), n_w * gl.num_programs(1), NUM_XCDS
@@ -263,8 +254,6 @@ def k2_ab_fused_gluon(
     s_kr = gl.allocate_shared_memory(gl.bfloat16, [2, C, K], SH_WS)
     s_inv = gl.allocate_shared_memory(inv_ty, [2, C, C], SH_PLAIN)
     s_v = gl.allocate_shared_memory(v_input.dtype.element_ty, [2, C, BW], SH_PLAIN)
-    # gt / beta live in LDS until their use, so the per-thread copies of the MMA
-    # slice layout are only materialized there.
     s_gt = gl.allocate_shared_memory(gl.float32, [2, NT], SH_1D)
     s_beta = gl.allocate_shared_memory(gl.float32, [2, NT], SH_1D)
 
@@ -281,9 +270,7 @@ def k2_ab_fused_gluon(
 
     for j in range(n_chunks):
         cur = j % 2
-        # Chunk j has landed in half `cur` (every thread's copies, once past the
-        # barrier the wait carries), and every thread is done reading the other
-        # half, which held chunk j-1, before it is refilled.
+        # Chunk j is in half `cur`, and every thread is done with half 1 - cur.
         gl.amd.cdna4.async_copy.wait_group(0)
         if j + 1 < n_chunks:
             _issue_chunk(ws_kd, ws_kr, ws_gt, ws_beta, ws_inv_mqk, v_input, None,
@@ -294,7 +281,7 @@ def k2_ab_fused_gluon(
                          cc_off, v_off, o_c_v, o_t, H, K, V, C, False)  # fmt: skip
 
         kd_a = gl.amd.cdna4.async_copy.load_shared_relaxed(s_kd.index(cur), A_OP_B)
-        # ws_kr is [C, K] in memory; the permuted view is the kr^T A operand.
+        # The permuted [C, K] view is the kr^T A operand.
         kr_a = gl.amd.cdna4.async_copy.load_shared_relaxed(
             s_kr.index(cur).permute((1, 0)), A_OP
         )
@@ -388,14 +375,11 @@ def k2_c_gluon(
     NUM_XCDS: gl.constexpr,
     NUM_WARPS: gl.constexpr,
 ):
-    """Pass C (the output pass): the recurrence of one segment from its incoming state.
+    """Pass C: one segment's recurrence from its incoming state, writing outputs.
 
     Same contract as ``_flash_kda_segment_kernel`` with ``COMPUTE_OUTPUT`` and
-    ``HAS_V`` and without ``STORE_H_OUT``, paged ``state_cache`` included. Chunk
-    j+1's tiles are prefetched into LDS while chunk j computes, as in the fused
-    pass A. The output is stored straight from the MFMA layout: converting it
-    for a wider store costs an LDS round trip and enough scratch to cost
-    occupancy at BW=128.
+    ``HAS_V`` and without ``STORE_H_OUT``, paged ``state_cache`` included.
+    Prefetches like pass A. The output is stored from the MFMA layout.
     """
     n_w = gl.num_programs(0)
     pid = remap_xcd(
@@ -410,11 +394,8 @@ def k2_c_gluon(
     n_chunks = gl.load(seg_nchunks + i_seg)
     tok_base = gl.load(seg_tok_base + i_seg).to(gl.int64)
     tok_end = gl.load(seg_tok_end + i_seg).to(gl.int64)
-    # Read up front, not after the loop: past the output stores the compiler
-    # must assume these words may have changed and loads them per lane, which
-    # leaves the final-state store a waterfall loop over a non-uniform pointer.
-    # When STORE_FINAL is false the final-state pointers are null and these
-    # must not be read.
+    # Read before the output stores, after which the compiler reloads them per
+    # lane (a waterfall loop on the final-state store). Null unless needed.
     if STORE_FINAL or PAGED_CACHE:
         i_n = gl.load(seg_seq + i_seg).to(gl.int64)
     if STORE_FINAL:
@@ -460,8 +441,7 @@ def k2_c_gluon(
     s_beta = gl.allocate_shared_memory(gl.float32, [2, NT], SH_1D)
 
     if PAGED_H_IN:
-        # A zero-start sequence's slot may hold a previous owner's state: mask
-        # the load so it is never read, not read and then discarded.
+        # A zero-start sequence's slot may hold a stale state: never read it.
         take = gl.load(has_initial_state + i_n).to(gl.int1)
         m_h = (o_k_m[:, None] < K) & (o_w_m[None, :] < V) & take
         h = gl.amd.cdna4.buffer_load(
@@ -512,9 +492,8 @@ def k2_c_gluon(
         )
         m_c = (t0 + o_c_m) < tok_end
 
-        # _recur's steps, split around the output projection: updating the state
-        # after it keeps the next state and kr out of registers meanwhile
-        # (1.1-1.2x on this kernel).
+        # _recur's steps, with the state update moved after the output
+        # projection to keep the next state and kr out of registers meanwhile.
         h_op = gl.convert_layout(h.to(gl.bfloat16), B_OP_B)
         tmp = gl.convert_layout(
             gl.amd.cdna4.mfma(kd_a, h_op, gl.zeros([C, BW], gl.float32, MMA_B)), MMA
@@ -562,3 +541,127 @@ def k2_c_gluon(
 
 k2_ab_fused_fast = fast_launch(k2_ab_fused_gluon)
 k2_c_fast = fast_launch(k2_c_gluon)
+
+
+def gluon_k2_ab_fused(
+    ws_kd,
+    ws_kr,
+    ws_gt,
+    ws_beta,
+    ws_inv_mqk,
+    v_input,
+    h_out_b,
+    h_out_a,
+    seg_chunk_base,
+    seg_nchunks,
+    seg_tok_base,
+    seg_tok_end,
+    TOTAL_TILES,
+    H,
+    K,
+    V,
+    C,
+    num_segs,
+    BW,
+    num_warps,
+    waves_per_eu=0,
+):
+    return k2_ab_fused_fast[(triton.cdiv(V, BW), num_segs * H)](
+        ws_kd=ws_kd,
+        ws_kr=ws_kr,
+        ws_gt=ws_gt,
+        ws_beta=ws_beta,
+        ws_inv_mqk=ws_inv_mqk,
+        v_input=v_input,
+        h_out_b=h_out_b,
+        h_out_a=h_out_a,
+        seg_chunk_base=seg_chunk_base,
+        seg_nchunks=seg_nchunks,
+        seg_tok_base=seg_tok_base,
+        seg_tok_end=seg_tok_end,
+        TOTAL_TILES=TOTAL_TILES,
+        H=H,
+        K=K,
+        V=V,
+        C=C,
+        BW=BW,
+        **build_layouts(num_warps),
+        NUM_XCDS=get_num_xcds(),
+        NUM_WARPS=num_warps,
+        num_warps=num_warps,
+        waves_per_eu=waves_per_eu,
+    )
+
+
+def gluon_k2_c(
+    ws_kd,
+    ws_qd,
+    ws_kr,
+    ws_gt,
+    ws_beta,
+    ws_inv_mqk,
+    v_input,
+    out,
+    h_in,
+    final_state,
+    seg_chunk_base,
+    seg_nchunks,
+    seg_tok_base,
+    seg_tok_end,
+    seg_seq,
+    seg_is_last,
+    state_cache,
+    state_indices,
+    has_initial_state,
+    cache_stride,
+    TOTAL_TILES,
+    H,
+    K,
+    V,
+    C,
+    num_segs,
+    BW,
+    num_warps,
+    STORE_FINAL,
+    STATE_V_FIRST,
+    CM_OUT="",
+):
+    paged = state_cache is not None
+    return k2_c_fast[(triton.cdiv(V, BW), num_segs * H)](
+        ws_kd=ws_kd,
+        ws_qd=ws_qd,
+        ws_kr=ws_kr,
+        ws_gt=ws_gt,
+        ws_beta=ws_beta,
+        ws_inv_mqk=ws_inv_mqk,
+        v_input=v_input,
+        out=out,
+        h_in=h_in,
+        final_state=final_state,
+        seg_chunk_base=seg_chunk_base,
+        seg_nchunks=seg_nchunks,
+        seg_tok_base=seg_tok_base,
+        seg_tok_end=seg_tok_end,
+        seg_seq=seg_seq,
+        seg_is_last=seg_is_last,
+        state_cache=state_cache,
+        state_indices=state_indices,
+        has_initial_state=has_initial_state,
+        cache_stride=cache_stride,
+        TOTAL_TILES=TOTAL_TILES,
+        H=H,
+        K=K,
+        V=V,
+        C=C,
+        BW=BW,
+        HAS_H_IN=h_in is not None,
+        STORE_FINAL=STORE_FINAL,
+        STATE_V_FIRST=STATE_V_FIRST,
+        PAGED_CACHE=paged,
+        PAGED_H_IN=paged and h_in is None,
+        **build_layouts(num_warps),
+        NUM_XCDS=get_num_xcds(),
+        NUM_WARPS=num_warps,
+        CM_OUT=CM_OUT,
+        num_warps=num_warps,
+    )
diff --git a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
index f3d74c66f51..211e0c8ef4c 100644
--- a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
+++ b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
@@ -56,6 +56,7 @@
 import triton
 import triton.language as tl
 
+from aiter.ops.triton._triton_kernels.common.pid_preprocessing import remap_xcd
 from aiter.ops.triton._triton_kernels.kimi_delta_attn.chunk_delta_attn_utils import (
     CHUNK_DELTA_ATTN_TRITON_AUTOTUNE,
     autotune_cache_kwargs,
@@ -71,7 +72,6 @@
     prepare_chunk_indices,
 )
 from aiter.ops.triton.utils._triton import arch_info
-from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
 from aiter.ops.triton.utils.device_info import get_num_xcds
 from aiter.ops.triton.utils.logger import AiterTritonLogger
 from aiter.ops.triton.utils.tuned_config_utils import autotune_configs
@@ -93,13 +93,13 @@
 _GLUON_BOTH: bool = _GLUON_SEL in ("1", "true", "yes", "on", "all")
 AITER_FDA_USE_GLUON_K1: bool = _GLUON_BOTH or _GLUON_SEL == "k1"
 AITER_FDA_USE_GLUON_K2: bool = _GLUON_BOTH or _GLUON_SEL == "k2"
-# Whether a Gluon K2 still leaves small unsegmented shapes to the Triton pass C,
-# which is faster there (see _k2c_gluon_schedule). Tests clear it to pin pass C
-# to Gluon.
+# Leave small unsegmented shapes to the Triton pass C (see _k2c_gluon_usable).
+# Tests clear it to pin pass C to Gluon.
 _K2C_ROUTE_BY_SIZE: bool = True
 
 DEVICE_ARCH = arch_info.get_arch()
 _GLUON_ARCH_OK: bool = DEVICE_ARCH == "gfx950"
+_GLUON_K2C_ARCH_OK: bool = DEVICE_ARCH == "gfx950"
 
 _LOGGED_ROUTES: set = set()
 
@@ -278,7 +278,7 @@ def _flash_kda_prepare_kernel(
 
     p_beta = beta_raw + (bos + t_off) * H + i_h + o_c * H
     b_beta = tl.sigmoid(tl.load(p_beta, mask=m_c, other=0.0).to(tl.float32))
-    # For the Gluon K2, which copies it into LDS with the chunk's other tiles.
+    # For the Gluon K2.
     if STORE_BETA:
         tl.store(ws_beta + ws_idx * C + o_c, b_beta, cache_modifier=CM_WS)
 
@@ -480,12 +480,9 @@ def _flash_kda_segment_kernel(
     With one segment per sequence this degenerates to a plain sequential scan
     and only the third form runs.
     """
-    # Every V block of one (segment, head) re-reads that segment's whole chunk
-    # workspace. Workgroups reach the XCDs round-robin in launch order, x
-    # fastest, so unmapped those blocks land on different XCDs and each pulls
-    # the workspace through its own L2 -- 6-8x the compulsory traffic at BW=16,
-    # which is what bound this kernel. Renumbered, each XCD owns a contiguous
-    # run of flat ids, so a (segment, head)'s V blocks share an L2.
+    # Every V block of a (segment, head) re-reads its chunk workspace; remapped,
+    # they land on one XCD and share its L2 instead of each pulling it through
+    # their own.
     n_w = tl.num_programs(0)
     pid = remap_xcd(
         tl.program_id(1) * n_w + tl.program_id(0), n_w * tl.num_programs(1), NUM_XCDS
@@ -810,61 +807,42 @@ def _log_route(k1_gluon: bool, k2_gluon: bool, C: int, K: int, V: int) -> None:
 
 
 @functools.cache
-def _k2_gluon_schedule(W: int, num_segs: int, H: int) -> tuple[int, int, int]:
-    """``(BW, num_warps, waves_per_eu)`` for the Gluon pass A, which does not autotune.
-
-    The warps can only split BW -- the state is the B operand of ``dot(kd, h)``
-    and its K axis is the contraction axis -- so a count above ``BW // 16``
-    duplicates columns instead of splitting them, which bounds the pairs the
-    config may publish. Which of the two a shape runs is decided by the grid,
-    ``cdiv(W, BW) * num_segs * H``: the wide tile reads the per-chunk workspace
-    half as many times, but only while there are still enough blocks to fill the
-    device, so it is taken above a published blocks-per-CU floor.
-
-    ``waves_per_eu`` (0 leaves it to the compiler) is the occupancy hint: the
-    LDS-prefetching kernel wants ~280 VGPRs, and holding it to 256 (a few
-    spills) for 2 waves per SIMD beats letting it drop to one.
-    """
-    cfg = chunk_delta_attn_tuned_config(
-        "k2_ab_fused_gluon_wide", _K2_GLUON_FALLBACK, backend="gluon"
+def _k2_gluon_schedule(
+    W: int, num_segs: int, H: int, kernel: str = "k2_ab_fused_gluon"
+) -> triton.Config:
+    """The schedule of the Gluon K2 pass ``kernel``, which does not autotune:
+    ``BW``, ``num_warps`` and ``waves_per_eu``. The wide tile is taken while the
+    grid stays above the published blocks-per-CU floor."""
+    wide = chunk_delta_attn_tuned_config(
+        f"{kernel}_wide", _K2_GLUON_FALLBACK, backend="gluon"
     )
-    bw = cfg.kwargs["BW"]
+    bw = wide.kwargs["BW"]
     blocks = (W // bw) * num_segs * H
-    if W % bw != 0 or blocks < cfg.kwargs["MIN_BLOCKS_PER_CU"] * _num_cus():
-        cfg = chunk_delta_attn_tuned_config(
-            "k2_ab_fused_gluon_narrow", _K2_GLUON_FALLBACK, backend="gluon"
-        )
-    return cfg.kwargs["BW"], cfg.num_warps, cfg.kwargs.get("waves_per_eu", 0)
+    if W % bw == 0 and blocks >= wide.kwargs["MIN_BLOCKS_PER_CU"] * _num_cus():
+        return wide
+    return chunk_delta_attn_tuned_config(
+        f"{kernel}_narrow", _K2_GLUON_FALLBACK, backend="gluon"
+    )
 
 
 @functools.cache
-def _k2c_gluon_schedule(segmented: bool) -> tuple[int, int, float]:
-    """``(BW, num_warps, MIN_BLOCKS_PER_CU)`` for the Gluon pass C.
-
-    Wider tiles amortize a block's workspace reads over more of V, but the
-    widest only pays once segmenting has supplied the blocks.
-    """
-    cfg = chunk_delta_attn_tuned_config(
+def _k2c_gluon_schedule(segmented: bool) -> triton.Config:
+    """The schedule of the gfx950 Gluon pass C: ``BW``, ``num_warps`` and the
+    ``MIN_BLOCKS_PER_CU`` floor below which the Triton pass C is faster."""
+    return chunk_delta_attn_tuned_config(
         "k2_c_gluon_segmented" if segmented else "k2_c_gluon",
         _K2C_GLUON_FALLBACK,
         backend="gluon",
     )
-    return cfg.kwargs["BW"], cfg.num_warps, cfg.kwargs.get("MIN_BLOCKS_PER_CU", 0)
 
 
 def _k2c_gluon_usable(segmented: bool, W: int, num_segs: int, H: int) -> bool:
-    """Whether the Gluon pass C beats the Triton one on this grid.
-
-    With few (segment, head) pairs the serial walk over a segment's chunks is
-    what sets the time, and there the Triton kernel's lighter iteration at
-    BW=16 wins; the Gluon one cannot go narrower than BW=32. With more pairs
-    the Gluon kernel's prefetch and fewer workspace re-reads take over. The
-    published ``MIN_BLOCKS_PER_CU`` floor marks the crossover: unsegmented on
-    MI355X it lies between 32 and 40 pairs (1.08-1.18x for Triton below it,
-    up to 2.2x for Gluon above).
-    """
-    bw, _, floor = _k2c_gluon_schedule(segmented)
-    blocks = triton.cdiv(W, bw) * num_segs * H
+    """Whether the Gluon pass C beats the Triton one on this grid. With few
+    (segment, head) pairs the serial walk dominates and the Triton kernel's
+    narrower BW=16 tile wins."""
+    sched = _k2c_gluon_schedule(segmented)
+    blocks = triton.cdiv(W, sched.kwargs["BW"]) * num_segs * H
+    floor = sched.kwargs.get("MIN_BLOCKS_PER_CU", 0)
     return not _K2C_ROUTE_BY_SIZE or blocks >= floor * _num_cus()
 
 
@@ -873,11 +851,10 @@ def _num_cus(device_index: int = 0) -> int:
     return torch.cuda.get_device_properties(device_index).multi_processor_count
 
 
-# Blocks pass A should end up with, in units of the CU count, and the most
-# segments per sequence taken for that reason alone.
+# Blocks pass A should end up with, in units of the CU count.
 _SEG_TARGET_BLOCKS = 3
 _SEG_MAX_SEGMENTS = 16
-# Longest segment, however many segments that takes.
+# Longest segment, regardless of segment count.
 _SEG_MAX_CHUNKS = 64
 _SEG_MIN_CHUNKS = 64
 
@@ -900,10 +877,8 @@ def _choose_chunks_per_seg(n_chunks_max: int, n_seqs: int, H: int, V: int) -> in
     H=12, V=128 that is 48 blocks against 256 CUs. Segmenting buys blocks by
     turning one pass into three -- pass A, the cross-segment scan, and pass C --
     so what decides the length is the block count it lands on rather than the
-    sequence length on its own. Once taken, segments are also kept to
-    ``_SEG_MAX_CHUNKS``: on gfx950, 64-chunk segments beat longer ones by up to
-    1.9x on long few-head sequences (1 x 131072 x 4), where the occupancy target
-    alone gives 256-chunk ones.
+    sequence length on its own. Segments are also kept to ``_SEG_MAX_CHUNKS``,
+    which wins on long few-head sequences.
     """
     override = os.getenv("CHUNK_DELTA_ATTN_FLASH_KDA_SEG", "").strip()
     if override:
@@ -918,9 +893,8 @@ def _choose_chunks_per_seg(n_chunks_max: int, n_seqs: int, H: int, V: int) -> in
     # up for.
     if n_chunks_max < _SEG_MIN_CHUNKS:
         return n_chunks_max
-    # Enough segments to fill the device, capped since the two extra passes
-    # scale with the segment count; and however many more it takes to keep a
-    # segment's own serial walk from becoming the limit on a long sequence.
+    # Enough segments to fill the device (capped, as the extra passes scale with
+    # the count), and enough to bound a segment's serial walk.
     segs = max(
         min(_SEG_MAX_SEGMENTS, _SEG_TARGET_BLOCKS * cus / blocks),
         n_chunks_max / _SEG_MAX_CHUNKS,
@@ -1177,18 +1151,22 @@ def flash_kda_fwd(
     # read twice per chunk on K2's serial path. The CUTLASS FlashKDA picks fp16
     # here for the same reason.
     ws_inv_mqk = torch.empty(H * total_tiles, 2 * C, C, dtype=torch.float16, device=dev)
-    # Gluon covers pass A and pass C; an unsegmented shape runs only pass C.
+
     use_gluon_k1 = AITER_FDA_USE_GLUON_K1 and _gluon_k1_usable(C, K)
     use_gluon_k2 = AITER_FDA_USE_GLUON_K2 and _gluon_k2_usable(C, K, V)
-    _log_route(use_gluon_k1, use_gluon_k2, C, K, V)
-    # Unsegmented, there is one segment per sequence, so pass C's route is
-    # known here; segmented, pass A runs on Gluon whichever route pass C takes.
-    use_gluon_k2c = use_gluon_k2 and (segmented or _k2c_gluon_usable(False, V, N, H))
-    # sigmoid(beta) per tile, which only the Gluon K2 reads; the Triton K2
-    # applies the sigmoid to the raw beta itself.
+    # Only a segmented shape has a pass A. Pass C's route is final here when
+    # unsegmented; segmented, it is settled once the segment count is known.
+    use_gluon_k2a = use_gluon_k2 and segmented
+    use_gluon_k2c = (
+        use_gluon_k2
+        and _GLUON_K2C_ARCH_OK
+        and (segmented or _k2c_gluon_usable(False, V, N, H))
+    )
+    _log_route(use_gluon_k1, use_gluon_k2a or use_gluon_k2c, C, K, V)
+    # sigmoid(beta) per tile, read only by the Gluon K2.
     ws_beta = (
         torch.empty(H * total_tiles, C, dtype=torch.float32, device=dev)
-        if use_gluon_k2c
+        if use_gluon_k2a or use_gluon_k2c
         else None
     )
 
@@ -1308,21 +1286,20 @@ def flash_kda_fwd(
         "has_initial_state": has_initial_state,
         "cache_stride": cache_stride,
     }
+    unpaged_args = {
+        "state_cache": None,
+        "state_indices": None,
+        "has_initial_state": None,
+        "cache_stride": 0,
+    }
 
-    if use_gluon_k2:
+    if use_gluon_k2a or use_gluon_k2c:
         from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn import (
             flash_kda_k2 as _g2,
         )
 
     def _launch_k2(*, W, **kw):
-        launch_args = {
-            **common,
-            "state_cache": None,
-            "state_indices": None,
-            "has_initial_state": None,
-            "cache_stride": 0,
-            **kw,
-        }
+        launch_args = {**common, **unpaged_args, **kw}
         return _segment_fast[
             lambda meta, _w=W: (triton.cdiv(_w, meta["BW"]), num_segs * H)
         ](W=W, **launch_args)
@@ -1332,9 +1309,9 @@ def _launch_k2(*, W, **kw):
         # parallel across segments. Neither writes outputs.
         b_seg = torch.empty(num_segs, H, K, V, dtype=torch.float32, device=dev)
         A_seg = torch.empty(num_segs, H, K, K, dtype=torch.bfloat16, device=dev)
-        if use_gluon_k2:
-            bw, nw, wpe = _k2_gluon_schedule(V, num_segs, H)
-            _g2.k2_ab_fused_fast[(triton.cdiv(V, bw), num_segs * H)](
+        if use_gluon_k2a:
+            sched = _k2_gluon_schedule(V, num_segs, H)
+            _g2.gluon_k2_ab_fused(
                 ws_kd=ws_kd,
                 ws_kr=ws_kr,
                 ws_gt=ws_gt,
@@ -1352,12 +1329,10 @@ def _launch_k2(*, W, **kw):
                 K=K,
                 V=V,
                 C=C,
-                BW=bw,
-                NUM_XCDS=get_num_xcds(),
-                **_g2.build_layouts(nw),
-                NUM_WARPS=nw,
-                num_warps=nw,
-                waves_per_eu=wpe,
+                num_segs=num_segs,
+                BW=sched.kwargs["BW"],
+                num_warps=sched.num_warps,
+                waves_per_eu=sched.kwargs.get("waves_per_eu", 0),
             )
         else:
             for buf, width, identity, has_v in (
@@ -1402,8 +1377,8 @@ def _launch_k2(*, W, **kw):
     if use_gluon_k2c and segmented:
         use_gluon_k2c = _k2c_gluon_usable(True, V, num_segs, H)
     if use_gluon_k2c:
-        bw, nw, _ = _k2c_gluon_schedule(segmented)
-        _g2.k2_c_fast[(triton.cdiv(V, bw), num_segs * H)](
+        sched = _k2c_gluon_schedule(segmented)
+        _g2.gluon_k2_c(
             ws_kd=ws_kd,
             ws_qd=ws_qd,
             ws_kr=ws_kr,
@@ -1420,23 +1395,18 @@ def _launch_k2(*, W, **kw):
             seg_tok_end=seg_tok_end,
             seg_seq=seg_seq,
             seg_is_last=seg_is_last,
-            **paged_args,
+            **(paged_args if paged else unpaged_args),
             TOTAL_TILES=total_tiles,
             H=H,
             K=K,
             V=V,
             C=C,
-            BW=bw,
-            **_g2.build_layouts(nw),
-            HAS_H_IN=h_in is not None,
+            num_segs=num_segs,
+            BW=sched.kwargs["BW"],
+            num_warps=sched.num_warps,
             STORE_FINAL=store_final,
             STATE_V_FIRST=state_v_first,
-            PAGED_CACHE=paged,
-            PAGED_H_IN=paged and h_in is None,
             CM_OUT=CM_OUT_STORE,
-            NUM_XCDS=get_num_xcds(),
-            NUM_WARPS=nw,
-            num_warps=nw,
         )
     else:
         _launch_k2(
diff --git a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py
index e600e5d4add..78394c09c3d 100644
--- a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py
+++ b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py
@@ -54,13 +54,12 @@ def _gather(src, index, axis, _builder=None):
 @triton.heuristics(
     {
         "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
-        # Token i's rows are loaded once, so one block spans all of K.
+        # One block spans all of K.
         "BK": lambda args: max(triton.next_power_of_2(args["K"]), 16),
     }
 )
 @triton.autotune(
-    # BH=1 idles most of a warp. The default is one config for all shapes,
-    # within 1.24x of the per-shape best on gfx950.
+    # BH=1 idles most of a warp.
     configs=autotune_configs(
         "CHUNK_DELTA_ATTN",
         [
@@ -145,8 +144,8 @@ def chunk_delta_attn_fwd_kernel_intra_token_parallel(
     m_k = o_k < K
     m_hk = m_hv[:, None] & m_k[None, :]
 
-    # Token i's rows are loaded once, outside the loop over j: re-reading them
-    # per (j, K block) made the kernel load-bound. beta is folded into k here.
+    # Token i's rows are loaded once, outside the loop over j; beta is folded
+    # into k.
     p_qk = o_h[:, None] * K + o_k[None, :]
     b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
     b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
diff --git a/op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py b/op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py
index e56f658dfd5..fd4101364f1 100644
--- a/op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py
+++ b/op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py
@@ -17,6 +17,8 @@
   alignment is one of the things the compiler was told it could assume.
 """
 
+import importlib
+
 import pytest
 import torch
 
@@ -124,20 +126,18 @@ def test_a_second_shape_does_not_reuse_the_first_entry():
 
 def _wrapped():
     """Every fast_launch-wrapped kernel the pipeline can reach."""
-    from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn import (
-        flash_kda_k1,
-        flash_kda_k2,
-    )
     from aiter.ops.triton._triton_kernels.kimi_delta_attn import flash_kda as fk
 
-    candidates = (
+    candidates = [
         fk._prepare_fast,
         fk._segment_fast,
         fk._seg_scan_fast,
-        flash_kda_k1._k1_fast,
-        flash_kda_k2.k2_ab_fused_fast,
-        flash_kda_k2.k2_c_fast,
-    )
+    ]
+    if fk._GLUON_ARCH_OK:
+        pkg = f"aiter.ops.triton._gluon_kernels.{fk.DEVICE_ARCH}.kimi_delta_attn"
+        k1 = importlib.import_module(f"{pkg}.flash_kda_k1")
+        k2 = importlib.import_module(f"{pkg}.flash_kda_k2")
+        candidates += [k1._k1_fast, k2.k2_ab_fused_fast, getattr(k2, "k2_c_fast", None)]
     return [c for c in candidates if isinstance(c, fast_launch._FastLaunch)]
 
 
diff --git a/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py b/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
index 712a18f5a55..e7113a76fd1 100644
--- a/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
+++ b/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
@@ -14,6 +14,7 @@
 """
 
 import contextlib
+import importlib
 import math
 
 import pytest
@@ -78,8 +79,7 @@ def _route(k1: bool | None = None, k2: bool | None = None):
     flags, so a test that wants the other route sets the flags rather than the
     variable. Both default to on wherever the tile shape and the arch allow it,
     which leaves the Triton kernels unreached by every test that does not come
-    through here. ``k2=True`` also pins pass C to Gluon on the small
-    unsegmented shapes the dispatch would otherwise leave to Triton.
+    through here. ``k2=True`` also pins pass C to Gluon on small shapes.
     """
     saved = (
         _flash_kda.AITER_FDA_USE_GLUON_K1,
@@ -442,20 +442,25 @@ def test_published_k2_schedules_can_split_their_tile():
     for W in (64, 128, 256):
         for num_segs in (1, 2, 8, 64, 512):
             for H in (1, 4, 12, 64):
-                reached.add(_flash_kda._k2_gluon_schedule(W, num_segs, H)[:2])
+                sched = _flash_kda._k2_gluon_schedule(W, num_segs, H)
+                reached.add((sched.kwargs["BW"], sched.num_warps))
     assert len(reached) > 1, f"only one schedule reachable: {reached}"
     # Pass C splits its tile the same way.
-    reached |= {_flash_kda._k2c_gluon_schedule(seg)[:2] for seg in (False, True)}
+    if _flash_kda._GLUON_K2C_ARCH_OK:
+        for segmented in (False, True):
+            sched = _flash_kda._k2c_gluon_schedule(segmented)
+            reached.add((sched.kwargs["BW"], sched.num_warps))
     for bw, nw in sorted(reached):
         assert bw % 16 == 0, f"BW={bw} is not a whole number of MFMA tiles"
         assert nw <= bw // 16, f"BW={bw} cannot be split {nw} ways"
 
 
 def test_k1_schedule_comes_from_the_config(monkeypatch):
-    """K1's occupancy hint is read from the published config, and a published
-    warp count its layouts were not built for is rejected, not launched."""
-    if not _flash_kda._gluon_k1_usable(FLASH_KDA_CHUNK, K_DIM):
-        pytest.skip("this arch never routes K1 to Gluon")
+    """K1's occupancy hint comes from the config; a wrong warp count is rejected."""
+    if _flash_kda.DEVICE_ARCH != "gfx950" or not _flash_kda._gluon_k1_usable(
+        FLASH_KDA_CHUNK, K_DIM
+    ):
+        pytest.skip("only the gfx950 Gluon K1 reads its schedule itself")
     from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn import (
         flash_kda_k1 as _g1,
     )
@@ -507,8 +512,7 @@ def _varlen(lens):
         },
     ),
     "weak gate": lambda: (make_inputs(1, 512, 4), {"lower_bound": -0.01}),
-    # Only a segmented schedule has a pass A, so these are the cases where the
-    # Gluon pass A is reached at all -- see test_cases_reach_the_gluon_k2.
+    # Only a segmented shape has a pass A.
     "segmented": lambda: (
         make_inputs(1, 1024, 4),
         {"output_final_state": True, "chunks_per_seg": 4},
@@ -563,10 +567,7 @@ def test_routes_agree(case):
 
     Both sides write the same ABI and nothing downstream is told which one ran,
     so a divergence here is a bug in whichever kernel moved rather than a
-    tolerance to widen. Measured across these cases: K1's two implementations
-    differ at 3e-4 to 1e-3 on the output and under 3e-5 on the state, and K2's
-    pass C by an ulp or two of the output, its output projection being ordered
-    differently.
+    tolerance to widen.
     """
     args, kw = _CASES[case]()
     with _route(k1=False, k2=False):
@@ -580,20 +581,19 @@ def test_routes_agree(case):
 
 
 def test_cases_reach_the_gluon_k2():
-    """The cases above have to exercise the route they are comparing.
-
-    Pass C goes to Gluon on every shape, pass A only on segmented ones, so an
-    unsegmented case says nothing about pass A; nothing in the assertions of
-    test_routes_agree would say which kernels a case actually reached.
-    """
+    """The cases above have to exercise the route they are comparing."""
     if not _flash_kda._gluon_k2_usable(FLASH_KDA_CHUNK, K_DIM, K_DIM):
         pytest.skip("this arch or tile shape never routes K2 to Gluon")
-    from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn import (
-        flash_kda_k2 as _g2,
+    _g2 = importlib.import_module(
+        "aiter.ops.triton._gluon_kernels."
+        f"{_flash_kda.DEVICE_ARCH}.kimi_delta_attn.flash_kda_k2"
     )
 
-    reached = {"k2_ab_fused_fast": set(), "k2_c_fast": set()}
-    saved = {name: getattr(_g2, name) for name in reached}
+    launchers = ["k2_ab_fused_fast"]
+    if _flash_kda._GLUON_K2C_ARCH_OK:
+        launchers.append("k2_c_fast")
+    reached = {name: set() for name in launchers}
+    saved = {name: getattr(_g2, name) for name in launchers}
 
     class _Counting:
         def __init__(self, name):
@@ -608,7 +608,7 @@ def launch(**kw):
 
             return launch
 
-    for name in reached:
+    for name in launchers:
         setattr(_g2, name, _Counting(name))
     try:
         for current, make in _CASES.items():
@@ -616,25 +616,25 @@ def launch(**kw):
             with _route(k1=True, k2=True):
                 run_flash(*args, **kw)
     finally:
-        for name, fn in saved.items():
-            setattr(_g2, name, fn)
+        for name in launchers:
+            setattr(_g2, name, saved[name])
 
-    want_a = {name for name in _CASES if name.startswith("segmented")}
-    missing_a = want_a - reached["k2_ab_fused_fast"]
-    missing_c = set(_CASES) - reached["k2_c_fast"]
-    assert not missing_a, f"never reached the Gluon pass A: {sorted(missing_a)}"
-    assert not missing_c, f"never reached the Gluon pass C: {sorted(missing_c)}"
+    want = {name for name in _CASES if name.startswith("segmented")}
+    missed = sorted(want - reached["k2_ab_fused_fast"])
+    assert not missed, f"never reached the Gluon pass A: {missed}"
+    if "k2_c_fast" in reached:
+        missed = sorted(set(_CASES) - reached["k2_c_fast"])
+        assert not missed, f"never reached the Gluon pass C: {missed}"
 
 
 def test_small_unsegmented_pass_c_stays_on_triton():
-    """The published floor sends few-pair unsegmented shapes to the Triton pass C.
-
-    Below it the serial walk sets the time and the Triton kernel's narrower tile
-    wins; above it, and whenever segmenting has supplied the blocks, Gluon does.
-    """
-    if not _flash_kda._gluon_k2_usable(FLASH_KDA_CHUNK, K_DIM, K_DIM):
-        pytest.skip("this arch or tile shape never routes K2 to Gluon")
-    if _flash_kda._k2c_gluon_schedule(False)[2] == 0:
+    """The published floor sends few-pair unsegmented shapes to the Triton pass C."""
+    if not (
+        _flash_kda._GLUON_K2C_ARCH_OK
+        and _flash_kda._gluon_k2_usable(FLASH_KDA_CHUNK, K_DIM, K_DIM)
+    ):
+        pytest.skip("this arch or tile shape never routes pass C to Gluon")
+    if not _flash_kda._k2c_gluon_schedule(False).kwargs.get("MIN_BLOCKS_PER_CU"):
         pytest.skip("this arch publishes no pass C floor")
     usable = _flash_kda._k2c_gluon_usable
     assert not usable(False, K_DIM, 1, 12)
@@ -888,11 +888,8 @@ def _kimi_kwargs(q, k, v, g, beta, A_log, dt_bias, scale, **kw):
 @pytest.mark.parametrize("gluon_k2", [False, True])
 @pytest.mark.parametrize("T,chunks_per_seg", [(256, 0), (1024, 4)])
 def test_paged_cache_matches_dense(T, chunks_per_seg, gluon_k2):
-    """In-kernel paged I/O matches gather into a dense V-first state.
-
-    On both K2 routes: the two pass C kernels each carry their own paged loads
-    and stores.
-    """
+    """In-kernel paged I/O matches gather into a dense V-first state, on both
+    K2 routes."""
     if gluon_k2 and not _flash_kda._gluon_k2_usable(FLASH_KDA_CHUNK, K_DIM, K_DIM):
         pytest.skip("this arch or tile shape never routes K2 to Gluon")
     H = 4

From 097849da359b05984b8a07bf0433d797dea97063 Mon Sep 17 00:00:00 2001
From: Meekail Zain <zainmeekail@gmail.com>
Date: Mon, 26 Oct 2026 21:41:17 +0000
Subject: [PATCH 9/9] =?UTF-8?q?[Triton/Gluon]=20KDA:=20address=20review=20?=
 =?UTF-8?q?=E2=80=94=20published=20schedules,=20resolved-route=20log,=20la?=
 =?UTF-8?q?unch-level=20route=20tests?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- K1 and pass C read their Gluon schedules through _published_gluon_schedule
  and fall back to Triton when none is published (no Python fallback). The
  K1 schedule moves to the wrapper; gluon_k1_prepare rejects a wrong warp count.
- Segment cap moves to the gfx950 Triton config (flash_kda_segments:
  MAX_CHUNKS=64, HARD_CAP). Other archs keep the previous 32-chunk soft cap.
- intra_attn token-parallel default (BH=4, 1 warp) moves to the gfx950 Triton
  config; other archs keep the previous BH=1, 4-warp default.
- Use get_num_sms() (honors CU_NUM) for routing and segmentation.
- Log the resolved K1, pass A and pass C routes after pass C is settled.
- Pass A and pass C share _solve_u/_update (AMDGCN unchanged).
- Tests: count actual pass-C launches for size routing, cover the segment cap
  and unpublished-schedule fallback, test K1 warp validation via flash_kda_fwd.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .../gfx950/kimi_delta_attn/flash_kda_k1.py    |  30 +---
 .../gfx950/kimi_delta_attn/flash_kda_k2.py    |  54 ++++--
 .../kimi_delta_attn/flash_kda.py              |  93 ++++++----
 .../kimi_delta_attn/intra_attn.py             |   8 +-
 .../attention/chunk_delta_attn/DEFAULT.json   |   9 +
 .../kimi_delta_attn/test_flash_kda.py         | 161 ++++++++++++------
 6 files changed, 227 insertions(+), 128 deletions(-)

diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
index c3aae647f1e..0d8a9a7dc01 100644
--- a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
+++ b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k1.py
@@ -1,10 +1,8 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
-import functools
 import math
 
-import triton
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
 
@@ -14,9 +12,6 @@
     sigmoid_log2,
     sum_combine,
 )
-from aiter.ops.triton._triton_kernels.kimi_delta_attn.chunk_delta_attn_utils import (
-    chunk_delta_attn_tuned_config,
-)
 from aiter.ops.triton._triton_kernels.kimi_delta_attn.fast_launch import fast_launch
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
 
@@ -320,24 +315,8 @@ def k1_prepare_gluon(
     )
 
 
-_NUM_WARPS = math.prod(_MMA_F16.warps_per_cta)
 # The warp count is fixed by the layouts above.
-_K1_FALLBACK = triton.Config({"waves_per_eu": 0}, num_warps=_NUM_WARPS)
-
-
-@functools.cache
-def _k1_schedule() -> tuple[int, int]:
-    """``(num_warps, waves_per_eu)`` for the Gluon K1. Only ``waves_per_eu`` is
-    tunable; a config naming another warp count is rejected."""
-    cfg = chunk_delta_attn_tuned_config(
-        "k1_prepare_gluon", _K1_FALLBACK, backend="gluon"
-    )
-    if cfg.num_warps != _NUM_WARPS:
-        raise ValueError(
-            f"k1_prepare_gluon is built for num_warps={_NUM_WARPS}, "
-            f"but the published config sets {cfg.num_warps}"
-        )
-    return cfg.num_warps, cfg.kwargs.get("waves_per_eu", 0)
+_NUM_WARPS = math.prod(_MMA_F16.warps_per_cta)
 
 
 _k1_fast = fast_launch(k1_prepare_gluon)
@@ -370,8 +349,13 @@ def gluon_k1_prepare(
     B,
     CM_WS="",
     CM_LOAD=".cg",
+    num_warps=_NUM_WARPS,
+    waves_per_eu=0,
 ):
-    num_warps, waves_per_eu = _k1_schedule()
+    if num_warps != _NUM_WARPS:
+        raise ValueError(
+            f"k1_prepare_gluon is built for num_warps={_NUM_WARPS}, got {num_warps}"
+        )
     return _k1_fast[(TOTAL_TILES if cu_seqlens is not None else NT, B * H)](
         q=q,
         k=k,
diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
index e1d53c7ba75..099c629f263 100644
--- a/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
+++ b/aiter/ops/triton/_gluon_kernels/gfx950/kimi_delta_attn/flash_kda_k2.py
@@ -58,12 +58,10 @@ def build_layouts(nw, kw=KW, kw_big=KW_BIG):
 
 
 @gluon.jit
-def _recur(
+def _solve_u(
     h,
     kd_a,
     inv_a,
-    kr_a,
-    gt,
     beta,
     v,
     m_c,
@@ -76,6 +74,7 @@ def _recur(
     INV_TY: gl.constexpr,
     HAS_V: gl.constexpr,
 ):
+    """U = inv @ (beta * (v - kd @ h)), and h as a dot operand."""
     h_op = gl.convert_layout(h.to(gl.bfloat16), B_OP_B)
     tmp = gl.convert_layout(
         gl.amd.cdna4.mfma(kd_a, h_op, gl.zeros([C, BW], gl.float32, MMA_B)), MMA
@@ -90,10 +89,39 @@ def _recur(
     big_u = gl.amd.cdna4.mfma(
         inv_a, gl.convert_layout(u.to(INV_TY), B_OP), gl.zeros([C, BW], gl.float32, MMA)
     )
-    h_next = gl.amd.cdna4.mfma(
+    return big_u, h_op
+
+
+@gluon.jit
+def _update(h, kr_a, big_u, gt, B_OP: gl.constexpr):
+    """The next state: diag(gt) h + kr^T U."""
+    return gl.amd.cdna4.mfma(
         kr_a, gl.convert_layout(big_u.to(gl.bfloat16), B_OP), h * gt[:, None]
     )
-    return h_next, big_u, h_op
+
+
+@gluon.jit
+def _recur(
+    h,
+    kd_a,
+    inv_a,
+    kr_a,
+    gt,
+    beta,
+    v,
+    m_c,
+    C: gl.constexpr,
+    BW: gl.constexpr,
+    MMA: gl.constexpr,
+    B_OP: gl.constexpr,
+    MMA_B: gl.constexpr,
+    B_OP_B: gl.constexpr,
+    INV_TY: gl.constexpr,
+    HAS_V: gl.constexpr,
+):
+    big_u, h_op = _solve_u(h, kd_a, inv_a, beta, v, m_c, C, BW, MMA, B_OP,
+                           MMA_B, B_OP_B, INV_TY, HAS_V)  # fmt: skip
+    return _update(h, kr_a, big_u, gt, B_OP), big_u, h_op
 
 
 @gluon.jit
@@ -494,16 +522,8 @@ def k2_c_gluon(
 
         # _recur's steps, with the state update moved after the output
         # projection to keep the next state and kr out of registers meanwhile.
-        h_op = gl.convert_layout(h.to(gl.bfloat16), B_OP_B)
-        tmp = gl.convert_layout(
-            gl.amd.cdna4.mfma(kd_a, h_op, gl.zeros([C, BW], gl.float32, MMA_B)), MMA
-        )
-        u = gl.where(m_c[:, None], (b_v - tmp) * beta[:, None], 0.0)
-        big_u = gl.amd.cdna4.mfma(
-            inv_a,
-            gl.convert_layout(u.to(inv_ty), B_OP),
-            gl.zeros([C, BW], gl.float32, MMA),
-        )
+        big_u, h_op = _solve_u(h, kd_a, inv_a, beta, b_v, m_c, C, BW, MMA, B_OP,
+                               MMA_B, B_OP_B, inv_ty, True)  # fmt: skip
 
         # o = qd @ h + Mqk @ U, h being the state entering the chunk.
         qd_a = gl.amd.cdna4.async_copy.load_shared_relaxed(s_qd.index(cur), A_OP_B)
@@ -518,9 +538,7 @@ def k2_c_gluon(
         kr_a = gl.amd.cdna4.async_copy.load_shared_relaxed(
             s_kr.index(cur).permute((1, 0)), A_OP
         )
-        h = gl.amd.cdna4.mfma(
-            kr_a, gl.convert_layout(big_u.to(gl.bfloat16), B_OP), h * gt[:, None]
-        )
+        h = _update(h, kr_a, big_u, gt, B_OP)
 
     if STORE_FINAL:  # noqa: SIM102
         if is_last == 1:
diff --git a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
index 211e0c8ef4c..02dbc019054 100644
--- a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
+++ b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py
@@ -72,7 +72,7 @@
     prepare_chunk_indices,
 )
 from aiter.ops.triton.utils._triton import arch_info
-from aiter.ops.triton.utils.device_info import get_num_xcds
+from aiter.ops.triton.utils.device_info import get_num_sms, get_num_xcds
 from aiter.ops.triton.utils.logger import AiterTritonLogger
 from aiter.ops.triton.utils.tuned_config_utils import autotune_configs
 
@@ -770,7 +770,12 @@ def flash_kda_supported(
 
 
 def _gluon_k1_usable(C: int, K: int) -> bool:
-    return _GLUON_ARCH_OK and C == FLASH_KDA_CHUNK and K == FLASH_KDA_K
+    return (
+        _GLUON_ARCH_OK
+        and C == FLASH_KDA_CHUNK
+        and K == FLASH_KDA_K
+        and _published_gluon_schedule("k1_prepare_gluon") is not None
+    )
 
 
 def _gluon_k2_usable(C: int, K: int, V: int) -> bool:
@@ -783,15 +788,21 @@ def _gluon_k2_usable(C: int, K: int, V: int) -> bool:
     )
 
 
-def _log_route(k1_gluon: bool, k2_gluon: bool, C: int, K: int, V: int) -> None:
-    key = (k1_gluon, k2_gluon, C, K, V)
+def _log_route(k1: bool, k2a: bool | None, k2c: bool, C: int, K: int, V: int) -> None:
+    """Log the resolved routes once each; ``k2a`` is None when there is no pass A."""
+    key = (k1, k2a, k2c, C, K, V)
     if key in _LOGGED_ROUTES:
         return
     _LOGGED_ROUTES.add(key)
+
+    def name(gluon):
+        return "none" if gluon is None else "gluon" if gluon else "triton"
+
     _LOGGER.info(
-        "flash_kda route: K1=%s K2=%s (C=%d K=%d V=%d arch=%s)",
-        "gluon" if k1_gluon else "triton",
-        "gluon" if k2_gluon else "triton",
+        "flash_kda route: K1=%s K2 pass A=%s pass C=%s (C=%d K=%d V=%d arch=%s)",
+        name(k1),
+        name(k2a),
+        name(k2c),
         C,
         K,
         V,
@@ -803,7 +814,15 @@ def _log_route(k1_gluon: bool, k2_gluon: bool, C: int, K: int, V: int) -> None:
 # allows. The zero floor keeps the pair self-consistent -- a device with no
 # published table takes the wide branch and finds this same schedule there.
 _K2_GLUON_FALLBACK = triton.Config({"BW": 32, "MIN_BLOCKS_PER_CU": 0}, num_warps=2)
-_K2C_GLUON_FALLBACK = triton.Config({"BW": 64, "MIN_BLOCKS_PER_CU": 0}, num_warps=4)
+_UNPUBLISHED = triton.Config({})
+
+
+@functools.cache
+def _published_gluon_schedule(name: str) -> triton.Config | None:
+    """The published Gluon schedule ``name``, or None. K1 and pass C have no
+    Python fallback: without a published schedule they are not routed to."""
+    cfg = chunk_delta_attn_tuned_config(name, _UNPUBLISHED, backend="gluon")
+    return None if cfg is _UNPUBLISHED else cfg
 
 
 @functools.cache
@@ -818,21 +837,18 @@ def _k2_gluon_schedule(
     )
     bw = wide.kwargs["BW"]
     blocks = (W // bw) * num_segs * H
-    if W % bw == 0 and blocks >= wide.kwargs["MIN_BLOCKS_PER_CU"] * _num_cus():
+    if W % bw == 0 and blocks >= wide.kwargs["MIN_BLOCKS_PER_CU"] * get_num_sms():
         return wide
     return chunk_delta_attn_tuned_config(
         f"{kernel}_narrow", _K2_GLUON_FALLBACK, backend="gluon"
     )
 
 
-@functools.cache
-def _k2c_gluon_schedule(segmented: bool) -> triton.Config:
+def _k2c_gluon_schedule(segmented: bool) -> triton.Config | None:
     """The schedule of the gfx950 Gluon pass C: ``BW``, ``num_warps`` and the
     ``MIN_BLOCKS_PER_CU`` floor below which the Triton pass C is faster."""
-    return chunk_delta_attn_tuned_config(
-        "k2_c_gluon_segmented" if segmented else "k2_c_gluon",
-        _K2C_GLUON_FALLBACK,
-        backend="gluon",
+    return _published_gluon_schedule(
+        "k2_c_gluon_segmented" if segmented else "k2_c_gluon"
     )
 
 
@@ -841,22 +857,29 @@ def _k2c_gluon_usable(segmented: bool, W: int, num_segs: int, H: int) -> bool:
     (segment, head) pairs the serial walk dominates and the Triton kernel's
     narrower BW=16 tile wins."""
     sched = _k2c_gluon_schedule(segmented)
+    if sched is None:
+        return False
     blocks = triton.cdiv(W, sched.kwargs["BW"]) * num_segs * H
     floor = sched.kwargs.get("MIN_BLOCKS_PER_CU", 0)
-    return not _K2C_ROUTE_BY_SIZE or blocks >= floor * _num_cus()
-
-
-@functools.cache
-def _num_cus(device_index: int = 0) -> int:
-    return torch.cuda.get_device_properties(device_index).multi_processor_count
+    return not _K2C_ROUTE_BY_SIZE or blocks >= floor * get_num_sms()
 
 
 # Blocks pass A should end up with, in units of the CU count.
 _SEG_TARGET_BLOCKS = 3
 _SEG_MAX_SEGMENTS = 16
-# Longest segment, regardless of segment count.
-_SEG_MAX_CHUNKS = 64
 _SEG_MIN_CHUNKS = 64
+# Where none is published: segments of at most 32 chunks, unless that takes more
+# than _SEG_MAX_SEGMENTS of them.
+_SEG_LIMITS_FALLBACK = triton.Config({"MAX_CHUNKS": 32, "HARD_CAP": False})
+
+
+@functools.cache
+def _seg_limits() -> tuple[int, bool]:
+    """``(MAX_CHUNKS, HARD_CAP)``: the longest segment, and whether that bound
+    holds past ``_SEG_MAX_SEGMENTS``."""
+    cfg = chunk_delta_attn_tuned_config("flash_kda_segments", _SEG_LIMITS_FALLBACK)
+    return cfg.kwargs["MAX_CHUNKS"], cfg.kwargs["HARD_CAP"]
+
 
 _SCAN_BV_NARROW = 16
 _SCAN_BV_WIDE = 32
@@ -864,7 +887,7 @@ def _num_cus(device_index: int = 0) -> int:
 
 def _scan_bv(n_seqs: int, H: int, V: int) -> tuple[int, int]:
     # Not a fill-the-CUs knob: under 16 always loses; this asks only if 16 overshoots.
-    if (V // _SCAN_BV_NARROW) * n_seqs * H > _num_cus():
+    if (V // _SCAN_BV_NARROW) * n_seqs * H > get_num_sms():
         return _SCAN_BV_WIDE, 2
     return _SCAN_BV_NARROW, 4
 
@@ -877,15 +900,15 @@ def _choose_chunks_per_seg(n_chunks_max: int, n_seqs: int, H: int, V: int) -> in
     H=12, V=128 that is 48 blocks against 256 CUs. Segmenting buys blocks by
     turning one pass into three -- pass A, the cross-segment scan, and pass C --
     so what decides the length is the block count it lands on rather than the
-    sequence length on its own. Segments are also kept to ``_SEG_MAX_CHUNKS``,
-    which wins on long few-head sequences.
+    sequence length on its own. Segments are also kept to the published
+    ``MAX_CHUNKS`` (see ``_seg_limits``).
     """
     override = os.getenv("CHUNK_DELTA_ATTN_FLASH_KDA_SEG", "").strip()
     if override:
         val = int(override)
         return n_chunks_max if val <= 0 else val
     blocks = n_seqs * H * max(1, V // 32)
-    cus = _num_cus()
+    cus = get_num_sms()
     # There has to be idle capacity to absorb the extra passes.
     if blocks > cus:
         return n_chunks_max
@@ -895,10 +918,11 @@ def _choose_chunks_per_seg(n_chunks_max: int, n_seqs: int, H: int, V: int) -> in
         return n_chunks_max
     # Enough segments to fill the device (capped, as the extra passes scale with
     # the count), and enough to bound a segment's serial walk.
-    segs = max(
-        min(_SEG_MAX_SEGMENTS, _SEG_TARGET_BLOCKS * cus / blocks),
-        n_chunks_max / _SEG_MAX_CHUNKS,
-    )
+    max_chunks, hard_cap = _seg_limits()
+    by_length = n_chunks_max / max_chunks
+    if not hard_cap:
+        by_length = min(_SEG_MAX_SEGMENTS, by_length)
+    segs = max(min(_SEG_MAX_SEGMENTS, _SEG_TARGET_BLOCKS * cus / blocks), by_length)
     # To a power of two, the grid the constants were calibrated on.
     return max(1, min(n_chunks_max, 1 << round(math.log2(n_chunks_max / segs))))
 
@@ -1162,7 +1186,6 @@ def flash_kda_fwd(
         and _GLUON_K2C_ARCH_OK
         and (segmented or _k2c_gluon_usable(False, V, N, H))
     )
-    _log_route(use_gluon_k1, use_gluon_k2a or use_gluon_k2c, C, K, V)
     # sigmoid(beta) per tile, read only by the Gluon K2.
     ws_beta = (
         torch.empty(H * total_tiles, C, dtype=torch.float32, device=dev)
@@ -1171,6 +1194,7 @@ def flash_kda_fwd(
     )
 
     if use_gluon_k1:
+        k1_sched = _published_gluon_schedule("k1_prepare_gluon")
         from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn.flash_kda_k1 import (
             gluon_k1_prepare,
         )
@@ -1202,6 +1226,8 @@ def flash_kda_fwd(
             B=B,
             CM_WS=CM_STORE,
             CM_LOAD=CM_LOAD,
+            num_warps=k1_sched.num_warps,
+            waves_per_eu=k1_sched.kwargs.get("waves_per_eu", 0),
         )
     else:
         _prepare_fast[(total_tiles if cu_seqlens is not None else NT, B * H)](
@@ -1376,6 +1402,9 @@ def _launch_k2(*, W, **kw):
     # Pass C: re-run each segment from its true incoming state, writing outputs.
     if use_gluon_k2c and segmented:
         use_gluon_k2c = _k2c_gluon_usable(True, V, num_segs, H)
+    _log_route(
+        use_gluon_k1, use_gluon_k2a if segmented else None, use_gluon_k2c, C, K, V
+    )
     if use_gluon_k2c:
         sched = _k2c_gluon_schedule(segmented)
         _g2.gluon_k2_c(
diff --git a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py
index 78394c09c3d..05599c145b1 100644
--- a/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py
+++ b/aiter/ops/triton/_triton_kernels/kimi_delta_attn/intra_attn.py
@@ -23,6 +23,7 @@
     IS_GATHER_SUPPORTED,
     IS_TF32_SUPPORTED,
     autotune_cache_kwargs,
+    chunk_delta_attn_tuned_config,
     exp2,
     input_guard,
 )
@@ -59,7 +60,6 @@ def _gather(src, index, axis, _builder=None):
     }
 )
 @triton.autotune(
-    # BH=1 idles most of a warp.
     configs=autotune_configs(
         "CHUNK_DELTA_ATTN",
         [
@@ -67,7 +67,11 @@ def _gather(src, index, axis, _builder=None):
             for BH in [1, 2, 4, 8]
             for nw in [1, 2, 4, 8]
         ],
-        default_config=triton.Config({"BH": 4}, num_warps=1, num_stages=3),
+        # The fallback is the untuned pre-gfx950 default.
+        default_config=chunk_delta_attn_tuned_config(
+            "chunk_delta_attn_fwd_kernel_intra_token_parallel",
+            triton.Config({"BH": 1}, num_warps=4, num_stages=3),
+        ),
     ),
     key=["K", "H", "HV"],
     **autotune_cache_kwargs,
diff --git a/aiter/ops/triton/configs/gfx950/triton/attention/chunk_delta_attn/DEFAULT.json b/aiter/ops/triton/configs/gfx950/triton/attention/chunk_delta_attn/DEFAULT.json
index 16fe84c1292..6b1da55e4db 100644
--- a/aiter/ops/triton/configs/gfx950/triton/attention/chunk_delta_attn/DEFAULT.json
+++ b/aiter/ops/triton/configs/gfx950/triton/attention/chunk_delta_attn/DEFAULT.json
@@ -42,5 +42,14 @@
                 "num_stages": 2
             }
         ]
+    },
+    "chunk_delta_attn_fwd_kernel_intra_token_parallel": {
+        "BH": 4,
+        "num_warps": 1,
+        "num_stages": 3
+    },
+    "flash_kda_segments": {
+        "MAX_CHUNKS": 64,
+        "HARD_CAP": true
     }
 }
diff --git a/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py b/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
index e7113a76fd1..96ee436b968 100644
--- a/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
+++ b/op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py
@@ -101,6 +101,41 @@ def _route(k1: bool | None = None, k2: bool | None = None):
         ) = saved
 
 
+@contextlib.contextmanager
+def _count_launches(module, names):
+    """Count the launches of ``module``'s fast launchers ``names``."""
+    calls = dict.fromkeys(names, 0)
+    saved = {name: getattr(module, name) for name in names}
+
+    class _Counting:
+        def __init__(self, name):
+            self.name = name
+
+        def __getitem__(self, grid):
+            inner = saved[self.name][grid]
+
+            def launch(**kw):
+                calls[self.name] += 1
+                return inner(**kw)
+
+            return launch
+
+    for name in names:
+        setattr(module, name, _Counting(name))
+    try:
+        yield calls
+    finally:
+        for name in names:
+            setattr(module, name, saved[name])
+
+
+def _gluon_k2_module():
+    return importlib.import_module(
+        "aiter.ops.triton._gluon_kernels."
+        f"{_flash_kda.DEVICE_ARCH}.kimi_delta_attn.flash_kda_k2"
+    )
+
+
 @contextlib.contextmanager
 def _k2_tuner_search_space():
     """Hand K2's autotuner the candidates a tuning build gives it.
@@ -446,36 +481,49 @@ def test_published_k2_schedules_can_split_their_tile():
                 reached.add((sched.kwargs["BW"], sched.num_warps))
     assert len(reached) > 1, f"only one schedule reachable: {reached}"
     # Pass C splits its tile the same way.
-    if _flash_kda._GLUON_K2C_ARCH_OK:
-        for segmented in (False, True):
-            sched = _flash_kda._k2c_gluon_schedule(segmented)
+    for segmented in (False, True):
+        sched = _flash_kda._k2c_gluon_schedule(segmented)
+        if sched is not None:
             reached.add((sched.kwargs["BW"], sched.num_warps))
     for bw, nw in sorted(reached):
         assert bw % 16 == 0, f"BW={bw} is not a whole number of MFMA tiles"
         assert nw <= bw // 16, f"BW={bw} cannot be split {nw} ways"
 
 
-def test_k1_schedule_comes_from_the_config(monkeypatch):
-    """K1's occupancy hint comes from the config; a wrong warp count is rejected."""
-    if _flash_kda.DEVICE_ARCH != "gfx950" or not _flash_kda._gluon_k1_usable(
-        FLASH_KDA_CHUNK, K_DIM
-    ):
-        pytest.skip("only the gfx950 Gluon K1 reads its schedule itself")
-    from aiter.ops.triton._gluon_kernels.gfx950.kimi_delta_attn import (
-        flash_kda_k1 as _g1,
+def _patch_gluon_schedule(monkeypatch, replace):
+    """Make ``_published_gluon_schedule`` return ``replace[name]`` for those names."""
+    published = _flash_kda._published_gluon_schedule
+    monkeypatch.setattr(
+        _flash_kda,
+        "_published_gluon_schedule",
+        lambda name: replace[name] if name in replace else published(name),
     )
 
-    assert _g1._k1_schedule()[0] == _g1._NUM_WARPS
 
-    wrong = triton.Config({"waves_per_eu": 0}, num_warps=2 * _g1._NUM_WARPS)
-    monkeypatch.setattr(_g1, "chunk_delta_attn_tuned_config", lambda *a, **kw: wrong)
-    _g1._k1_schedule.cache_clear()
-    try:
-        with pytest.raises(ValueError, match="num_warps"):
-            _g1._k1_schedule()
-    finally:
-        monkeypatch.undo()
-        _g1._k1_schedule.cache_clear()
+def test_gluon_k1_rejects_a_wrong_warp_count(monkeypatch):
+    """K1's layouts fix its warp count, so a config with another one must fail."""
+    if not _flash_kda._gluon_k1_usable(FLASH_KDA_CHUNK, K_DIM):
+        pytest.skip("this arch never routes K1 to Gluon")
+    good = _flash_kda._published_gluon_schedule("k1_prepare_gluon")
+    wrong = triton.Config(dict(good.kwargs), num_warps=2 * good.num_warps)
+    _patch_gluon_schedule(monkeypatch, {"k1_prepare_gluon": wrong})
+    with _route(k1=True), pytest.raises(ValueError, match="num_warps"):
+        run_flash(*make_inputs(1, 64, 2))
+
+
+def test_unpublished_gluon_schedules_fall_back_to_triton(monkeypatch):
+    """Without a published schedule, K1 and pass C take their Triton kernels."""
+    unpublished = ("k1_prepare_gluon", "k2_c_gluon", "k2_c_gluon_segmented")
+    _patch_gluon_schedule(monkeypatch, dict.fromkeys(unpublished))
+    args = make_inputs(1, 1024, 4)
+    o_ref, _ = run_reference(*args)
+    with (
+        _route(k1=True, k2=True),
+        _count_launches(_flash_kda, ["_prepare_fast", "_segment_fast"]) as calls,
+    ):
+        o, _ = run_flash(*args, chunks_per_seg=4)
+    assert calls == {"_prepare_fast": 1, "_segment_fast": 1}
+    assert rel_err(o, o_ref) < 2e-2
 
 
 def _varlen(lens):
@@ -584,40 +632,19 @@ def test_cases_reach_the_gluon_k2():
     """The cases above have to exercise the route they are comparing."""
     if not _flash_kda._gluon_k2_usable(FLASH_KDA_CHUNK, K_DIM, K_DIM):
         pytest.skip("this arch or tile shape never routes K2 to Gluon")
-    _g2 = importlib.import_module(
-        "aiter.ops.triton._gluon_kernels."
-        f"{_flash_kda.DEVICE_ARCH}.kimi_delta_attn.flash_kda_k2"
-    )
-
     launchers = ["k2_ab_fused_fast"]
     if _flash_kda._GLUON_K2C_ARCH_OK:
         launchers.append("k2_c_fast")
     reached = {name: set() for name in launchers}
-    saved = {name: getattr(_g2, name) for name in launchers}
-
-    class _Counting:
-        def __init__(self, name):
-            self.name = name
-
-        def __getitem__(self, grid):
-            inner = saved[self.name][grid]
-
-            def launch(**kw):
-                reached[self.name].add(current)
-                return inner(**kw)
-
-            return launch
-
-    for name in launchers:
-        setattr(_g2, name, _Counting(name))
-    try:
+    with _count_launches(_gluon_k2_module(), launchers) as calls:
         for current, make in _CASES.items():
             args, kw = make()
+            before = dict(calls)
             with _route(k1=True, k2=True):
                 run_flash(*args, **kw)
-    finally:
-        for name in launchers:
-            setattr(_g2, name, saved[name])
+            for name in launchers:
+                if calls[name] > before[name]:
+                    reached[name].add(current)
 
     want = {name for name in _CASES if name.startswith("segmented")}
     missed = sorted(want - reached["k2_ab_fused_fast"])
@@ -627,19 +654,47 @@ def launch(**kw):
         assert not missed, f"never reached the Gluon pass C: {missed}"
 
 
-def test_small_unsegmented_pass_c_stays_on_triton():
+@pytest.mark.parametrize(
+    "B, H, chunks_per_seg, gluon",
+    [(1, 4, None, False), (4, 12, None, True), (1, 4, 4, True)],
+    ids=["small", "large", "segmented"],
+)
+def test_pass_c_route_follows_size(monkeypatch, B, H, chunks_per_seg, gluon):
     """The published floor sends few-pair unsegmented shapes to the Triton pass C."""
     if not (
         _flash_kda._GLUON_K2C_ARCH_OK
         and _flash_kda._gluon_k2_usable(FLASH_KDA_CHUNK, K_DIM, K_DIM)
     ):
         pytest.skip("this arch or tile shape never routes pass C to Gluon")
-    if not _flash_kda._k2c_gluon_schedule(False).kwargs.get("MIN_BLOCKS_PER_CU"):
+    sched = _flash_kda._k2c_gluon_schedule(False)
+    if sched is None or not sched.kwargs.get("MIN_BLOCKS_PER_CU"):
         pytest.skip("this arch publishes no pass C floor")
-    usable = _flash_kda._k2c_gluon_usable
-    assert not usable(False, K_DIM, 1, 12)
-    assert usable(False, K_DIM, 8, 16)
-    assert usable(True, K_DIM, 16, 4)
+    if chunks_per_seg is None:
+        assert _flash_kda._k2c_gluon_usable(False, K_DIM, B, H) == gluon
+    monkeypatch.setattr(_flash_kda, "AITER_FDA_USE_GLUON_K2", True)
+    monkeypatch.setattr(_flash_kda, "_K2C_ROUTE_BY_SIZE", True)
+    args = make_inputs(B, 1024 if chunks_per_seg else 64, H)
+    with (
+        _count_launches(_gluon_k2_module(), ["k2_c_fast"]) as gluon_calls,
+        _count_launches(_flash_kda, ["_segment_fast"]) as triton_calls,
+    ):
+        run_flash(*args, chunks_per_seg=chunks_per_seg)
+    assert gluon_calls["k2_c_fast"] == int(gluon)
+    assert triton_calls["_segment_fast"] == int(not gluon)
+
+
+def test_segment_length_cap(monkeypatch):
+    """A hard cap bounds every segment; the fallback's soft cap yields to the
+    segment-count cap on long sequences."""
+    monkeypatch.delenv("CHUNK_DELTA_ATTN_FLASH_KDA_SEG", raising=False)
+    choose = _flash_kda._choose_chunks_per_seg
+    long_seq = 64 * _flash_kda._SEG_MAX_SEGMENTS * 32
+    monkeypatch.setattr(_flash_kda, "_seg_limits", lambda: (64, True))
+    for n_chunks in (64, 512, 4096, long_seq):
+        for H in (1, 4, 12):
+            assert choose(n_chunks, 1, H, K_DIM) <= 64, (n_chunks, H)
+    monkeypatch.setattr(_flash_kda, "_seg_limits", lambda: (32, False))
+    assert choose(long_seq, 1, 1, K_DIM) == long_seq // _flash_kda._SEG_MAX_SEGMENTS
 
 
 # A weak gate is the only setting that exposes the intra-chunk inverse. At the