From 042a289183e5b47e193df6c10160915a6e8b69f5 Mon Sep 17 00:00:00 2001
From: 1am9trash <1am9trash@gmail.com>
Date: Thu, 1 Oct 2026 06:39:42 +0000
Subject: [PATCH 1/2] Avoid the peeled sparse MLA spill on 64-bit gathers below
 triton 3.8

---
 aiter/ops/triton/attention/pa_decode_sparse.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py
index ef70f9342d4..b506a5c0678 100644
--- a/aiter/ops/triton/attention/pa_decode_sparse.py
+++ b/aiter/ops/triton/attention/pa_decode_sparse.py
@@ -11,6 +11,7 @@
 
 import torch
 import triton
+from packaging.version import Version
 
 from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_mla import (
     _sparse_mla as _sparse_mla_gfx950,
@@ -46,6 +47,9 @@
 except ImportError:
     _HAS_SCALED_UPCAST = False
 
+# Below Triton 3.8 the peeled 64-bit split-K kernel spills ~110 VGPRs to scratch.
+_TRITON_GE_38 = Version(Version(triton.__version__).base_version) >= Version("3.8.0")
+
 _LOGGER = AiterTritonLogger()
 
 
@@ -662,11 +666,15 @@ def _pa_decode_sparse_gfx950_gluon(
     if one_wg_per_cu:
         waves_per_eu = 1
 
-    # Unpeeled is faster at prefill and, on the 64-bit gathers, unless split-K
-    # leaves each program a few tiles. Decode on buffer loads is faster peeled.
+    # Decode on buffer loads is faster peeled; prefill is faster unpeeled.
+    # On the 64-bit gathers unpeeled is faster unless split-K leaves each program
+    # a few tiles, and below Triton 3.8 the peeled kernel spills, so always unpeel.
     row_tiles = max(avg_main, avg_extra) / BLOCK_K
     short_splits = num_splits > 1 and row_tiles <= 4 * num_splits
-    unpeel = num_queries >= _PREFILL_MIN_ROWS if use_buffer_load else not short_splits
+    if use_buffer_load:
+        unpeel = num_queries >= _PREFILL_MIN_ROWS
+    else:
+        unpeel = not (_TRITON_GE_38 and short_splits)
 
     main_splits = num_splits
     if has_extra and avg_main > 0:

From 1cbfb85715bacbb373cbe89cf8a03bde9ace2f94 Mon Sep 17 00:00:00 2001
From: Thomas Wang <1am9trash@gmail.com>
Date: Thu, 1 Oct 2026 18:12:45 +0800
Subject: [PATCH 2/2] Fix version comparison for triton version check

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 aiter/ops/triton/attention/pa_decode_sparse.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py
index b506a5c0678..c7440d3586c 100644
--- a/aiter/ops/triton/attention/pa_decode_sparse.py
+++ b/aiter/ops/triton/attention/pa_decode_sparse.py
@@ -48,7 +48,7 @@
     _HAS_SCALED_UPCAST = False
 
 # Below Triton 3.8 the peeled 64-bit split-K kernel spills ~110 VGPRs to scratch.
-_TRITON_GE_38 = Version(Version(triton.__version__).base_version) >= Version("3.8.0")
+_TRITON_GE_38 = Version(triton.__version__) >= Version("3.8.0")
 
 _LOGGER = AiterTritonLogger()