[Triton/Gluon] KDA: faster FlashKDA and token-parallel intra forward on gfx950

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,
 )
@@ -54,18 +55,23 @@ def _gather(src, index, axis, _builder=None):
 @triton.heuristics(
     {
         "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+        # One block spans all of K.
+        "BK": lambda args: max(triton.next_power_of_2(args["K"]), 16),
     }
 )
 @triton.autotune(
     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),
+        # 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,
@@ -137,44 +143,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; 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)
+    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,
         )
 

aiter/ops/triton/configs/gfx950/gluon/attention/chunk_delta_attn/DEFAULT.json

@@ -1,11 +1,26 @@
 {
+    "k1_prepare_gluon": {
+        "num_warps": 2,
+        "waves_per_eu": 3
+    },
     "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,
+        "MIN_BLOCKS_PER_CU": 0.3
+    },
+    "k2_c_gluon_segmented": {
+        "BW": 128,
+        "num_warps": 4
     }
 }

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
     }
 }

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,
             )
@@ -187,6 +190,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",

op_tests/triton_tests/kimi_delta_attn/test_fast_launch.py

@@ -138,7 +138,6 @@ def _wrapped():
         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")
-        # Pass C has a Gluon kernel on gfx1250 only.
         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)]