From 64167feb314a6c8c38bea5eaa19bd91d9f7ecd3b Mon Sep 17 00:00:00 2001
From: "mqliu.leek" <101073612+Estrellajer@users.noreply.github.com>
Date: Sun, 4 Oct 2026 02:36:01 -0700
Subject: [PATCH] [ops] fix: stop NPU GDN Triton kernels recompiling per batch
 length

chunk_gated_delta_rule_fwd_kernel_h, prepare_wy_repr_bwd_kernel and
recompute_w_u_fwd_kernel took NT (total chunk count of the packed batch) as a
tl.constexpr, and chunk_fwd_kernel_o took N (sequence count) as one. Both
change with every packed batch, so each new total length compiled a fresh set
of kernels: ~10 s per GDN forward+backward on 910B2C. Warm caches hide it; a
training run on varlen data sees new lengths every step.

Pass them as runtime ints and add them, with the kkt kernel's NT/B/
TOTAL_TASKS, to do_not_specialize. Outputs and gradients are bit-identical to
before over eight layouts; with a fresh cache, seven of eight new lengths then
run without compiling. A 154-step GRPO run on DocQA-RL goes from 57-65 s to
48-54 s per step.

Same bug class and same technique as fla-org/flash-linear-attention#1207,
which fixed it in fla's own triton_ascend backend; these kernels came from
MindSpeed-MM 39775d87 before that and carry a different set of offending
constexprs, so the fix is reproduced here rather than ported.
---
 .../gated_delta_rule/_ascend/triton/chunk_delta_h.py      | 4 ++--
 .../kernels/gated_delta_rule/_ascend/triton/chunk_o.py    | 4 ++--
 .../_ascend/triton/chunk_scaled_dot_kkt.py                | 2 +-
 .../kernels/gated_delta_rule/_ascend/triton/wy_fast.py    | 8 ++++----
 4 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_delta_h.py b/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_delta_h.py
index 7afcb18825..2ebdc43505 100644
--- a/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_delta_h.py
+++ b/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_delta_h.py
@@ -23,7 +23,7 @@
     configs=get_autotune_config(multibuffer_list=(False,)),
     key=['H', 'K', 'V', 'BT'],
 )
-@triton.jit(do_not_specialize=['T'])
+@triton.jit(do_not_specialize=['T', 'NT'])
 def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
         k,
         v,
@@ -42,7 +42,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
         V: tl.constexpr,
         BT: tl.constexpr,
         BV: tl.constexpr,
-        NT: tl.constexpr,
+        NT,
         USE_G: tl.constexpr,
         USE_GK: tl.constexpr,
         USE_INITIAL_STATE: tl.constexpr,
diff --git a/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_o.py b/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_o.py
index eb54e8b701..5826f0e4a5 100644
--- a/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_o.py
+++ b/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_o.py
@@ -288,7 +288,7 @@ def chunk_bwd_kernel_dv_local(
     'USE_G_GAMMA': lambda args: args['g_gamma'] is not None,
     'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
 })
-@triton.jit(do_not_specialize=['T'])
+@triton.jit(do_not_specialize=['T', 'N'])
 def chunk_fwd_kernel_o(
     q,
     k,
@@ -302,7 +302,7 @@ def chunk_fwd_kernel_o(
     scale,
     T,
     H: tl.constexpr,
-    N: tl.constexpr,
+    N,
     Hg: tl.constexpr,
     K: tl.constexpr,
     V: tl.constexpr,
diff --git a/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_scaled_dot_kkt.py b/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_scaled_dot_kkt.py
index 60a1050a85..6213e8c642 100644
--- a/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_scaled_dot_kkt.py
+++ b/veomni/ops/kernels/gated_delta_rule/_ascend/triton/chunk_scaled_dot_kkt.py
@@ -13,7 +13,7 @@
     'USE_G': lambda args: args['g'] is not None,
     'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
 })
-@triton.jit(do_not_specialize=['T'])
+@triton.jit(do_not_specialize=['T', 'NT', 'B', 'TOTAL_TASKS'])
 def chunk_scaled_dot_kkt_fwd_kernel(
     k,
     g,
diff --git a/veomni/ops/kernels/gated_delta_rule/_ascend/triton/wy_fast.py b/veomni/ops/kernels/gated_delta_rule/_ascend/triton/wy_fast.py
index 374c3d485e..c3afb218b6 100644
--- a/veomni/ops/kernels/gated_delta_rule/_ascend/triton/wy_fast.py
+++ b/veomni/ops/kernels/gated_delta_rule/_ascend/triton/wy_fast.py
@@ -13,7 +13,7 @@
 @triton.heuristics({
     'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
 })
-@triton.jit(do_not_specialize=['T'])
+@triton.jit(do_not_specialize=['T', 'B', 'NT'])
 def prepare_wy_repr_bwd_kernel(
         k,
         v,
@@ -33,7 +33,7 @@ def prepare_wy_repr_bwd_kernel(
         H: tl.constexpr,
         K: tl.constexpr,
         V: tl.constexpr,
-        NT: tl.constexpr,
+        NT,
         BT: tl.constexpr,
         BK: tl.constexpr,
         BV: tl.constexpr,
@@ -145,7 +145,7 @@ def prepare_wy_repr_bwd_kernel(
     'USE_GK': lambda args: args['gk'] is not None,
     'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
 })
-@triton.jit(do_not_specialize=['T'])
+@triton.jit(do_not_specialize=['T', 'B', 'NT'])
 def recompute_w_u_fwd_kernel(
         k,
         v,
@@ -162,7 +162,7 @@ def recompute_w_u_fwd_kernel(
         H: tl.constexpr,
         K: tl.constexpr,
         V: tl.constexpr,
-        NT: tl.constexpr,
+        NT,
         BT: tl.constexpr,
         BK: tl.constexpr,
         BV: tl.constexpr,