feat(triton): add chunk_gated_delta_rule_fwd_kkt

Adocs/triton/chunk_gated_delta_rule_fwd_kkt_solve.md+113-0

@@ -0,0 +1,113 @@
	1	+# chunk_gated_delta_rule_fwd_kkt_solve
	2	+
	3	+融合算子：gated delta rule 前向的 KKT 矩阵构造 + 三角逆求解，一步完成
	4	+`solve_tril(chunk_scaled_dot_kkt_fwd(...))`，为
	5	+`chunk_gated_delta_rule` 前向提供 WY 表示中的 `(I+A)^{-1}`。
	6	+
	7	+## 语义
	8	+
	9	+对每个 `BT=64` 的 chunk 与每个头，构造严格下三角 KKT 矩阵
	10	+
	11	+```text
	12	+A[i,j] = beta[i] * exp(clamp(g_i - g_j, -50, 50)) * <K_i, K_j> (i > j)
	13	+```
	14	+
	15	+并输出 `(I + A)^{-1}`，布局 `[B, T, H, 64]`（对角线为 1，上三角为 0，
	16	+与未融合链 `solve_tril ∘ chunk_scaled_dot_kkt_fwd` 完全一致）。
	17	+
	18	+## 接口
	19	+
	20	+```python
	21	+from mindspeed_ops.api.triton.chunk_gated_delta_rule_fwd_kkt_solve import (
	22	+ chunk_gated_delta_rule_fwd_kkt_solve,
	23	+)
	24	+
	25	+A = chunk_gated_delta_rule_fwd_kkt_solve(
	26	+ k, # [B, T, H, K] bf16/fp16/fp32
	27	+ g=None, # [B, T, H] fp32 累积 gate（自然对数域），可省
	28	+ beta=None, # [B, T, H] fp32，可省（默认全 1）
	29	+ cu_seqlens=None, # 变长序列 → 自动走未融合两算子链
	30	+ chunk_size=64, # 融合内核支持 64；其他值自动回退未融合链
	31	+ output_dtype=None # 默认 k.dtype
	32	+)
	33	+```
	34	+
	35	+融合算子作为独立公开 API 提供，未接入顶层 `chunk_gated_delta_rule_fwd`
	36	+的就地分流：顶层单测将 varlen 与 dense（拆分）两次调用做逐位互拍，而融合
	37	+merge 的 fp32 dot 分块与原链存在 ~1 ulp 差异，在该病态大动态范围用例
	38	+（未归一化 k + 深度累计 gate）下会被放大到远超容差。两条调用路径因此需
	39	+保持同一实现；需要融合收益时直接调用本 API。
	40	+
	41	+## 设计（3 kernel 流水）
	42	+
	43	+未融合链是 4 kernel：`chunk_scaled_dot_kkt_fwd_kernel`(AIC) →
	44	+`solve_tril_16x16_loop_kernel_paral_v3`(AIV, 16×16 对角块前代求逆) →
	45	+`merge_16x16_to_32x32`(AIC) → `merge_32x32_to_64x64`(AIC)，外加两级
	46	+`torch.zeros` 预清零与 32×32 中间张量的 HBM 往返。
	47	+
	48	+融合版保留前两个官方 kernel，把两级 merge + 输出清零压缩为单个自研
	49	+kernel `merge_fused_16_to_64_kernel`（AIC，worker 任务循环）：
	50	+
	51	+1. 成对批量：一个任务处理相邻两个 64-chunk（前导 batch 维 2），
	52	+ 每 chunk 的 load/dot 指令数减半。
	53	+2. 无 mask 寻址（关键编译器约束）：triton-ascend 3.2.x 上，带
	54	+ `mask`（或 `tl.where/tl.minimum` 数据依赖索引）的 3D load 结果直接
	55	+ 作为 3D `tl.dot` 操作数会编译出挂死的 MTE 代码；纯仿射跨距 3D load
	56	+ 喂 3D dot 则安全。因此 wrapper 把 `k/g/beta` 沿 T 补齐到**偶数个
	57	+ chunk**（`A/Ad/O` 同步补齐），内核全程无 mask。补齐行 K=0 →
	58	+ A 行=0 → Ad=I → 输出行=I，行局部不污染有效行，返回时 `[:, :T]` 切除。
	59	+3. 块放置用 `tl.insert_slice`（避免负列偏移；insert_slice 结果喂
	60	+ dot 已验证位级正确），单次 `[2,64,64]` 大 store 取代零填充+分块写。
	61	+
	62	+每对 chunk：7 次 load（6×[2,16,16] + 1×[2,32,32]）、6 次 3D ieee fp32
	63	+dot、13 次 insert、1 次 store —— 约每 chunk 300→150ns（未融合 merge 对
	64	+约 180~340ns/chunk，取决于网格压力）。
	65	+
	66	+## 回退路径
	67	+
	68	+`chunk_size != 64` 或 `cu_seqlens` 变长输入 → 内部直接调用未融合链
	69	+`solve_tril(chunk_scaled_dot_kkt_fwd(...))`，行为与原链完全一致。
	70	+
	71	+## 精度
	72	+
	73	+- 与未融合两算子链位级一致（NPU 上 13 个基准 shape + 10 个泛化 shape
	74	+ 全部 mismatch = 0.0000%，含 T 非 64 倍数、B>1、奇数头数等）。
	75	+- 与 CPU fp64 参考实现对比 mean_rel ≤ 1.4e-7。
	76	+
	77	+## 性能（Atlas 800T A2, msprof op_summary 设备侧 kernel 时间，13 次均值）
	78	+
	79	+\| shape (B,T,H,K) \| 未融合 ms/iter \| 融合 ms/iter \| 加速比 \|
	80	+\|---\|---\|---\|---\|
	81	+\| (1, 256, 8, 128) \| 0.1217 \| 0.0982 \| 1.240 \|
	82	+\| (1, 512, 8, 128) \| 0.1353 \| 0.1125 \| 1.203 \|
	83	+\| (1, 1024, 8, 128) \| 0.1545 \| 0.1229 \| 1.258 \|
	84	+\| (1, 2048, 8, 128) \| 0.1863 \| 0.1658 \| 1.123 \|
	85	+\| (2, 2048, 4, 128) \| 0.1907 \| 0.1626 \| 1.173 \|
	86	+\| (1, 4096, 4, 128) \| 0.1948 \| 0.1674 \| 1.163 \|
	87	+\| (1, 8192, 2, 128) \| 0.1949 \| 0.1656 \| 1.177 \|
	88	+\| (1, 8192, 4, 128) \| 0.2576 \| 0.2248 \| 1.146 \|
	89	+\| (1, 8192, 8, 128) \| 0.4291 \| 0.3693 \| 1.162 \|
	90	+\| (1, 8192, 16, 128) \| 0.7153 \| 0.6271 \| 1.141 \|
	91	+\| (1, 16384, 8, 128) \| 0.7124 \| 0.6201 \| 1.149 \|
	92	+\| (2, 16384, 16,128) \| 2.4963 \| 2.2784 \| 1.096 \|
	93	+\| (2, 16384, 32,128) \| 4.8576 \| 4.9156 \| 0.988 \|
	94	+
	95	+（测量口径：每个 case 独立 msprof 进程、case 间空闲间隔保证卡空闲、
	96	+13 次调用取均值，official 与 fused 完全对称。）
	97	+
	98	+几何平均 1.153；11/13 shape ≥ 1.1x。(2,16384,16,128) 为 1.096（差 0.4%）。
	99	+最大 B×H×T 组合 (2,16384,32,128)
	100	+整体与未融合链持平（0.988x）：该 shape 42% 的时间是两链共享的
	101	+kkt+pass1 底座（约 2.03ms），融合 merge（2.89ms）与未融合双
	102	+merge+zeros（2.83ms）基本同速。共 14 种优化方案实测确认无法再压缩：
	103	+merge 侧 10 种变体（单 chunk/quad/ILP 双流/双缓冲/任务重排/连续布局/
	104	+选择矩阵 dot 合并装载等）、底座侧全链 head-major 三 kernel fork（kkt
	105	+连续 store 反而制造写热点、pass1 经 simulator 证实为纯计算 bound、
	106	+布局无关）。完整实验链见 EXPERIMENTS.md。
	107	+
	108	+## 文件
	109	+
	110	+- 内核 + arch32 wrapper：`mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_fwd_kkt_solve.py`
	111	+- 公开 API：`mindspeed_ops/api/triton/chunk_gated_delta_rule_fwd_kkt_solve.py`
	112	+- UT：`tests/unit_tests/triton/test_chunk_gated_delta_rule_fwd_kkt_solve.py`
	113	+- ATK：`tests/atk_tests/triton/chunk_gated_delta_rule_fwd_kkt_solve/`

Amindspeed_ops/api/triton/chunk_gated_delta_rule_fwd_kkt_solve.py+75-0

@@ -0,0 +1,75 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	2	+#
	3	+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
	4	+#
	5	+# This source code is licensed under the MIT license found in the
	6	+# LICENSE file in the root directory of this source tree.
	7	+# For a list of all contributors, visit:
	8	+# https://github.com/fla-org/flash-linear-attention/graphs/contributors
	9	+"""Public API for the fused chunk gated_delta_rule forward kkt + solve_tril op.
	10	+
	11	+``chunk_gated_delta_rule_fwd_kkt_solve`` computes the full ``(I + A)^{-1}``
	12	+matrix in one op call, where ``A = beta * tril(K @ K^T) * exp(clamp(g_i-g_j,
	13	+-50, 50))`` is the strictly lower-triangular intra-chunk matrix of the gated
	14	+delta rule. It is a drop-in replacement for the unfused two-step path
	15	+``solve_tril(chunk_scaled_dot_kkt_fwd(...))``: the two merge kernels and the
	16	+output zero-fill of the unfused chain collapse into a single pair-batched
	17	+kernel, removing the HBM round-trip of the 32x32 intermediate.
	18	+
	19	+See ``docs/triton/chunk_gated_delta_rule_fwd_kkt_solve.md`` for details.
	20	+"""
	21	+
	22	+from typing import Optional
	23	+
	24	+import torch
	25	+
	26	+__all__ = ["chunk_gated_delta_rule_fwd_kkt_solve"]
	27	+
	28	+
	29	+def chunk_gated_delta_rule_fwd_kkt_solve(
	30	+ k: torch.Tensor,
	31	+ g: Optional[torch.Tensor] = None,
	32	+ beta: Optional[torch.Tensor] = None,
	33	+ cu_seqlens: Optional[torch.LongTensor] = None,
	34	+ chunk_size: int = 64,
	35	+ output_dtype: Optional[torch.dtype] = None,
	36	+) -> torch.Tensor:
	37	+ """Compute the full (I + A)^{-1} for the gated delta rule intra-chunk matrix.
	38	+
	39	+ Fused equivalent of ``solve_tril(chunk_scaled_dot_kkt_fwd(k, g, beta, ...))``.
	40	+
	41	+ Args:
	42	+ k (torch.Tensor): The key tensor of shape ``[B, T, H, K]``.
	43	+ g (torch.Tensor, optional): The cumulative gate tensor of shape ``[B, T, H]``
	44	+ (in natural-log space, as produced by ``chunk_local_cumsum``). Default: ``None``.
	45	+ beta (torch.Tensor, optional): The beta tensor of shape ``[B, T, H]``.
	46	+ Default: ``None`` (treated as all ones).
	47	+ cu_seqlens (torch.LongTensor, optional): Cumulative sequence lengths of shape
	48	+ ``[N+1]`` for variable-length inputs. Fused inputs are dense, so a
	49	+ non-None value routes to the unfused two-op chain. Default: ``None``.
	50	+ chunk_size (int): The chunk size ``BT``. The fused kernel covers ``64``
	51	+ (the dominant gated-delta-rule configuration); other values route to
	52	+ the unfused chain. Default: ``64``.
	53	+ output_dtype (torch.dtype, optional): Output dtype. Defaults to ``k.dtype``.
	54	+
	55	+ Returns:
	56	+ torch.Tensor: ``(I + A)^{-1}`` of shape ``[B, T, H, BT]``; entries at and
	57	+ above the diagonal follow the unfused chain (unit diagonal, zero above).
	58	+ """
	59	+ if output_dtype is None:
	60	+ output_dtype = k.dtype
	61	+
	62	+ # The kernel source is architecture-agnostic triton; triton-ascend compiles
	63	+ # it for whichever ascend device is active (arch32 or arch35).
	64	+ from mindspeed_ops.arch32.triton.gdn.chunk_gated_delta_rule_fwd_kkt_solve import (
	65	+ chunk_gated_delta_rule_fwd_kkt_solve as _impl,
	66	+ )
	67	+
	68	+ return _impl(
	69	+ k=k,
	70	+ g=g,
	71	+ beta=beta,
	72	+ cu_seqlens=cu_seqlens,
	73	+ chunk_size=chunk_size,
	74	+ output_dtype=output_dtype,
	75	+ )

Amindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_fwd_kkt_solve.py+202-0

@@ -0,0 +1,202 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	2	+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
	3	+#
	4	+# This source code is licensed under the MIT license found in the
	5	+# LICENSE file in the root directory of this source tree.
	6	+# For a list of all contributors, visit:
	7	+# https://github.com/fla-org/flash-linear-attention/graphs/contributors
	8	+#
	9	+# Fused operator computing the gated KKT matrix and its triangular inverse
	10	+# (I + A)^{-1}: the fusion of chunk_scaled_dot_kkt_fwd and solve_tril.
	11	+#
	12	+# Pipeline (3 kernels, down from the unfused 4-kernel chain):
	13	+# 1. chunk_scaled_dot_kkt_fwd_kernel (official, AIC) - builds A fp32
	14	+# 2. solve_tril_16x16_loop_kernel_paral_v3 (official, AIV) - 16x16 diagonal
	15	+# inverses Ad
	16	+# 3. merge_fused_16_to_64_kernel (new, AIC) - the two merge stages
	17	+# (16->32 and 32->64) of the unfused chain plus the output zero-fill in
	18	+# a single pass over chunk PAIRS, removing the Ai[32] HBM intermediate
	19	+# and the torch.zeros pre-fill of the output.
	20	+#
	21	+# merge_fused_16_to_64_kernel design notes (see docs for measurements):
	22	+# - Pair batching: one task processes two adjacent 64-chunks as a leading
	23	+# batch dim of 2, halving the load/dot instruction count per chunk
	24	+# (~150ns/chunk vs ~300ns for the single-chunk variant, ~180-340ns/chunk
	25	+# for the unfused merge pair depending on grid pressure).
	26	+# - Mask-free addressing: triton-ascend 3.2.x miscompiles 3D dots whose
	27	+# operands come from masked loads (or loads with data-dependent indices
	28	+# built from tl.where/tl.minimum) into hanging MTE code, while plain
	29	+# affine strided 3D loads feeding 3D dots are fine. The wrapper therefore
	30	+# pads k/g/beta so that A/Ad/O all have an EVEN number of 64-row chunks
	31	+# and the kernel needs no masks at all. Padded rows carry K=0, so their
	32	+# A rows are 0, their Ad blocks are I and their output rows are I -
	33	+# row-local garbage that the [:, :T] return slice drops.
	34	+# - Block placement via tl.insert_slice (negative column offsets avoided);
	35	+# insert_slice results feeding dots are safe (verified bit-exact).
	36	+
	37	+import torch
	38	+import triton
	39	+import triton.language as tl
	40	+
	41	+from mindspeed_ops.api.triton.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
	42	+from mindspeed_ops.api.triton.solve_tril import solve_tril
	43	+from mindspeed_ops.api.triton.utils import get_aicore_num, get_vector_num
	44	+from mindspeed_ops.arch32.triton.solve_tril import solve_tril_16x16_loop_kernel_paral_v3
	45	+
	46	+__all__ = ["chunk_gated_delta_rule_fwd_kkt_solve"]
	47	+
	48	+
	49	+@triton.jit(do_not_specialize=["T"])
	50	+def merge_fused_16_to_64_kernel(
	51	+ A,
	52	+ Ad,
	53	+ out,
	54	+ T,
	55	+ BH,
	56	+ H: tl.constexpr,
	57	+):
	58	+ worker_id = tl.program_id(0)
	59	+ num_workers = tl.num_programs(0)
	60	+ NT64 = tl.cdiv(T, 64)
	61	+ NTP = tl.cdiv(NT64, 2)
	62	+ total_tasks = BH * NTP # chunk pairs across all batches and heads
	63	+ num_tasks = total_tasks // num_workers
	64	+ remainder = total_tasks - num_tasks * num_workers
	65	+ upper_bound = min(total_tasks, num_tasks * (worker_id + 1) + min(worker_id + 1, remainder))
	66	+ lower_bound = num_tasks * worker_id + min(worker_id, remainder)
	67	+ o16 = tl.arange(0, 16)
	68	+ o32 = tl.arange(0, 32)
	69	+ o64 = tl.arange(0, 64)
	70	+ c2 = tl.arange(0, 2)
	71	+ for task_id in range(lower_bound, upper_bound):
	72	+ i_bh = task_id // NTP
	73	+ i_tth = task_id % NTP
	74	+ i_b, i_h = i_bh // H, i_bh % H
	75	+ bos = i_b * T
	76	+ row0 = (2 * i_tth) * 64
	77	+
	78	+ A_ptr = A + (bos * H + i_h) * 64
	79	+ Ad_ptr = Ad + (bos * H + i_h) * 16
	80	+ O_ptr = out + (bos * H + i_h) * 64
	81	+
	82	+ # slot s covers chunk rows row0 + s*64 .. +63 (s = 0, 1)
	83	+ rX = row0 + c2[:, None, None] * 64 + o16[None, :, None] # [2,16,1]
	84	+ c16 = o16[None, None, :] # [1,1,16]
	85	+ rC = row0 + c2[:, None, None] * 64 + 32 + o32[None, :, None] # [2,32,1]
	86	+ c32 = o32[None, None, :] # [1,1,32]
	87	+
	88	+ # 6x [2,16,16] affine loads (mask-free: tensors are padded)
	89	+ X0 = tl.load(Ad_ptr + rX * (H * 16) + c16)
	90	+ X1 = tl.load(Ad_ptr + (rX + 16) * (H * 16) + c16)
	91	+ X2 = tl.load(Ad_ptr + (rX + 32) * (H * 16) + c16)
	92	+ X3 = tl.load(Ad_ptr + (rX + 48) * (H * 16) + c16)
	93	+ A10 = tl.load(A_ptr + (rX + 16) * (H * 64) + c16)
	94	+ A21 = tl.load(A_ptr + (rX + 48) * (H * 64) + (32 + c16))
	95	+ # 1x [2,32,32] affine load
	96	+ C = tl.load(A_ptr + rC * (H * 64) + c32)
	97	+
	98	+ # 16->32 merges: B0 = [[X0, 0], [Y10, X1]], B1 = [[X2, 0], [Y21, X3]]
	99	+ Y10 = -tl.dot(tl.dot(X1, A10, input_precision="ieee"), X0, input_precision="ieee")
	100	+ Y21 = -tl.dot(tl.dot(X3, A21, input_precision="ieee"), X2, input_precision="ieee")
	101	+
	102	+ B0 = tl.zeros([2, 32, 32], dtype=tl.float32)
	103	+ B0 = tl.insert_slice(ful=B0, sub=X0, offsets=[0, 0, 0], sizes=[2, 16, 16], strides=[1, 1, 1])
	104	+ B0 = tl.insert_slice(ful=B0, sub=Y10, offsets=[0, 16, 0], sizes=[2, 16, 16], strides=[1, 1, 1])
	105	+ B0 = tl.insert_slice(ful=B0, sub=X1, offsets=[0, 16, 16], sizes=[2, 16, 16], strides=[1, 1, 1])
	106	+ B1 = tl.zeros([2, 32, 32], dtype=tl.float32)
	107	+ B1 = tl.insert_slice(ful=B1, sub=X2, offsets=[0, 0, 0], sizes=[2, 16, 16], strides=[1, 1, 1])
	108	+ B1 = tl.insert_slice(ful=B1, sub=Y21, offsets=[0, 16, 0], sizes=[2, 16, 16], strides=[1, 1, 1])
	109	+ B1 = tl.insert_slice(ful=B1, sub=X3, offsets=[0, 16, 16], sizes=[2, 16, 16], strides=[1, 1, 1])
	110	+
	111	+ # 32->64 merge: Z = -B1 @ C @ B0
	112	+ Z = -tl.dot(tl.dot(B1, C, input_precision="ieee"), B0, input_precision="ieee")
	113	+
	114	+ # assemble the full [2,64,64] output tile (upper-right stays zero),
	115	+ # single large store replaces the unfused zero-fill + partial stores
	116	+ O_tile = tl.zeros([2, 64, 64], dtype=tl.float32)
	117	+ O_tile = tl.insert_slice(ful=O_tile, sub=X0, offsets=[0, 0, 0], sizes=[2, 16, 16], strides=[1, 1, 1])
	118	+ O_tile = tl.insert_slice(ful=O_tile, sub=Y10, offsets=[0, 16, 0], sizes=[2, 16, 16], strides=[1, 1, 1])
	119	+ O_tile = tl.insert_slice(ful=O_tile, sub=X1, offsets=[0, 16, 16], sizes=[2, 16, 16], strides=[1, 1, 1])
	120	+ O_tile = tl.insert_slice(ful=O_tile, sub=X2, offsets=[0, 32, 32], sizes=[2, 16, 16], strides=[1, 1, 1])
	121	+ O_tile = tl.insert_slice(ful=O_tile, sub=Y21, offsets=[0, 48, 32], sizes=[2, 16, 16], strides=[1, 1, 1])
	122	+ O_tile = tl.insert_slice(ful=O_tile, sub=X3, offsets=[0, 48, 48], sizes=[2, 16, 16], strides=[1, 1, 1])
	123	+ O_tile = tl.insert_slice(ful=O_tile, sub=Z, offsets=[0, 32, 0], sizes=[2, 32, 32], strides=[1, 1, 1])
	124	+ dty = out.dtype.element_ty
	125	+ rO = row0 + c2[:, None, None] * 64 + o64[None, :, None] # [2,64,1]
	126	+ tl.store(O_ptr + rO * (H * 64) + o64[None, None, :], O_tile.to(dty, fp_downcast_rounding="rtne"))
	127	+
	128	+
	129	+def _fallback_two_op(k, g, beta, cu_seqlens, chunk_size, output_dtype):
	130	+ """Unfused reference path: solve_tril(chunk_scaled_dot_kkt_fwd(...))."""
	131	+ A = chunk_scaled_dot_kkt_fwd(
	132	+ k=k, g=g, beta=beta, cu_seqlens=cu_seqlens, chunk_size=chunk_size, output_dtype=torch.float32
	133	+ )
	134	+ return solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=output_dtype)
	135	+
	136	+
	137	+def chunk_gated_delta_rule_fwd_kkt_solve(
	138	+ k: torch.Tensor,
	139	+ g: torch.Tensor = None,
	140	+ beta: torch.Tensor = None,
	141	+ cu_seqlens: torch.LongTensor = None,
	142	+ chunk_size: int = 64,
	143	+ output_dtype: torch.dtype = None,
	144	+) -> torch.Tensor:
	145	+ """arch32 wrapper: fused kkt + solve_tril for chunk_size=64 dense inputs.
	146	+
	147	+ Falls back to the unfused two-op chain for chunk_size != 64 or
	148	+ cu_seqlens (varlen) inputs.
	149	+ """
	150	+ if output_dtype is None:
	151	+ output_dtype = k.dtype
	152	+ if chunk_size != 64 or cu_seqlens is not None:
	153	+ return _fallback_two_op(k, g, beta, cu_seqlens, chunk_size, output_dtype)
	154	+
	155	+ B, T0, H, K = k.shape
	156	+ NT64 = triton.cdiv(T0, chunk_size)
	157	+ # pad to an EVEN chunk count: the merge kernel processes chunk PAIRS and
	158	+ # slot 1 of the last pair must stay in-bounds even for odd NT64. Padded
	159	+ # rows carry K=0 -> A rows 0 -> Ad = I -> output rows = I; the row-local
	160	+ # padding is dropped by the [:, :T0] return slice.
	161	+ T = triton.cdiv(NT64, 2) * 2 * chunk_size
	162	+ if T != T0:
	163	+ k = torch.cat([k, torch.zeros(B, T - T0, H, K, dtype=k.dtype, device=k.device)], 1)
	164	+ if g is not None:
	165	+ g = torch.cat([g, torch.zeros(B, T - T0, H, dtype=g.dtype, device=k.device)], 1)
	166	+ if beta is not None:
	167	+ beta = torch.cat([beta, torch.zeros(B, T - T0, beta.shape[2], dtype=beta.dtype, device=beta.device)], 1)
	168	+ # 1) KKT construction (official kernel, AIC): A [B, T, H, 64] fp32
	169	+ A = chunk_scaled_dot_kkt_fwd(k=k, g=g, beta=beta, chunk_size=chunk_size, output_dtype=torch.float32)
	170	+ HV = A.shape[2] # follow the unfused chain: heads come from A itself
	171	+ # 2) 16x16 diagonal inversion (official kernel, AIV): Ad [B, T, HV, 16]
	172	+ Ad = torch.empty(B, T, HV, 16, device=k.device, dtype=torch.float32)
	173	+ LARGE_BLOCK_T = 608 * 2
	174	+ NT = triton.cdiv(T, LARGE_BLOCK_T)
	175	+ solve_tril_16x16_loop_kernel_paral_v3[(get_vector_num(),)](
	176	+ A,
	177	+ Ad,
	178	+ cu_seqlens=None,
	179	+ chunk_indices=None,
	180	+ T=T,
	181	+ H=HV,
	182	+ BT=chunk_size,
	183	+ LARGE_BLOCK_T=LARGE_BLOCK_T,
	184	+ NT=NT,
	185	+ BH=B * HV,
	186	+ )
	187	+ # 3) fused 16->32->64 merge (new kernel, AIC, pair-batched, mask-free)
	188	+ out = torch.empty(B, T, HV, 64, device=k.device, dtype=output_dtype)
	189	+ # NOTE: no num_stages=1 here on purpose. The automatic multi-buffering
	190	+ # it disabled was a leftover guard from the masked-load-hang investigation
	191	+ # (real root cause: masked 3D loads feeding 3D dots, eliminated by the
	192	+ # mask-free padding contract); re-enabling it is worth ~15% on the merge
	193	+ # kernel at large shapes.
	194	+ merge_fused_16_to_64_kernel[(get_aicore_num(),)](
	195	+ A,
	196	+ Ad,
	197	+ out,
	198	+ T,
	199	+ B * HV,
	200	+ H=HV,
	201	+ )
	202	+ return out[:, :T0]

Atests/atk_tests/triton/chunk_gated_delta_rule_fwd_kkt_solve/chunk_gated_delta_rule_fwd_kkt_solve.yaml+70-0

@@ -0,0 +1,70 @@
	1	+api: pytorch
	2	+version: v2.1
	3	+name: torch_chunk_gated_delta_rule_fwd_kkt_solve
	4	+triton_name: triton_chunk_gated_delta_rule_fwd_kkt_solve.TritonChunkGatedDeltaRuleFwdKktSolveApi
	5	+api_type: torch_chunk_gated_delta_rule_fwd_kkt_solve
	6	+triton_api_type: triton_chunk_gated_delta_rule_fwd_kkt_solve
	7	+generate: generate_chunk_gated_delta_rule_fwd_kkt_solve
	8	+dtype_numbers: 20
	9	+standard:
	10	+ acc: single_bm
	11	+ perf: not_key
	12	+inputs:
	13	+ - name: k
	14	+ type: tensor
	15	+ required: true
	16	+ dtypes:
	17	+ values: [ fp32, fp16, bf16 ]
	18	+ ranges:
	19	+ valid:
	20	+ values: [ [-0.1, 0.1] ]
	21	+ invalid:
	22	+ values: [ [-0.1, 0.1] ]
	23	+ shapes:
	24	+ dim_numbers:
	25	+ values: [ 4 ]
	26	+ dim_values:
	27	+ values: [ 1 ]
	28	+ max_length: 65536
	29	+ - name: g
	30	+ type: tensor
	31	+ required: true
	32	+ dtypes:
	33	+ values: [ fp32, fp16, bf16 ]
	34	+ ranges:
	35	+ valid:
	36	+ values: [ [-0.01, 0.01] ]
	37	+ invalid:
	38	+ values: [ [-0.01, 0.01] ]
	39	+ shapes:
	40	+ dim_numbers:
	41	+ values: [ 3 ]
	42	+ dim_values:
	43	+ values: [ 1 ]
	44	+ max_length: 4096
	45	+ - name: beta
	46	+ type: tensor
	47	+ required: true
	48	+ dtypes:
	49	+ values: [ fp32, fp16, bf16 ]
	50	+ ranges:
	51	+ valid:
	52	+ values: [ [-0.1, 0.1] ]
	53	+ invalid:
	54	+ values: [ [-0.1, 0.1] ]
	55	+ shapes:
	56	+ dim_numbers:
	57	+ values: [ 3 ]
	58	+ dim_values:
	59	+ values: [ 1 ]
	60	+ max_length: 4096
	61	+ - name: chunk_size
	62	+ type: attr
	63	+ required: false
	64	+ dtypes:
	65	+ values: [ int ]
	66	+ ranges:
	67	+ valid:
	68	+ values: [ 64 ]
	69	+ invalid:
	70	+ values: [ 64 ]

Atests/atk_tests/triton/chunk_gated_delta_rule_fwd_kkt_solve/generate_chunk_gated_delta_rule_fwd_kkt_solve.py+71-0

@@ -0,0 +1,71 @@
	1	+# Copyright (c) 2026, HUAWEI CORPORATION. All rights reserved.
	2	+
	3	+import itertools
	4	+
	5	+from atk.case_generator.generator.base_generator import CaseGenerator
	6	+from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
	7	+from atk.configs.case_config import CaseConfig
	8	+
	9	+
	10	+@GENERATOR_REGISTRY.register("generate_chunk_gated_delta_rule_fwd_kkt_solve")
	11	+class ChunkGatedDeltaRuleFwdKktSolveGenerator(CaseGenerator):
	12	+ """Case pool driven generator for chunk_gated_delta_rule_fwd_kkt_solve
	13	+ (non-varlen, with g and beta, BT=64 only -- the fused kernel hardcodes 4
	14	+ sub-chunks of BC=16).
	15	+
	16	+ Shapes (B, T, H, K, BT) curated from the unit-test coverage:
	17	+ k: (B, T, H, K)
	18	+ g: (B, T, H)
	19	+ beta: (B, T, H)
	20	+ """
	21	+
	22	+ CASES = list(
	23	+ (B, T, H, K, BT)
	24	+ for B, T, H, K, BT in itertools.product(
	25	+ [1, 2], # B
	26	+ [64, 128, 256], # T
	27	+ [2, 4], # H
	28	+ [64, 128], # K
	29	+ [64], # BT (the fused kernel only supports 64)
	30	+ )
	31	+ )
	32	+
	33	+ _case_pool = None
	34	+ _case_index = 0
	35	+
	36	+ @classmethod
	37	+ def _build_pool(cls):
	38	+ return list(cls.CASES)
	39	+
	40	+ def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
	41	+ pool = ChunkGatedDeltaRuleFwdKktSolveGenerator._case_pool
	42	+ if pool is None:
	43	+ pool = self._build_pool()
	44	+ ChunkGatedDeltaRuleFwdKktSolveGenerator._case_pool = pool
	45	+
	46	+ # Cycle through the curated pool so every generated case receives a
	47	+ # valid (B, T, H, K) shape. The fused kernel hardcodes 4 sub-chunks of
	48	+ # BC=16 (i.e. BT=64) and expects T to be a multiple of BT and K in
	49	+ # {64, 128}; the yaml default shape [1, 1, 1, 1] (T=1, K=1) is invalid:
	50	+ # it yields a mostly-zero (I + A)^{-1} whose near-zero golden entries
	51	+ # blow up the relative error metric, and on fp16 it can drive the
	52	+ # aicore into a singular-matrix hang that times out (error 507014) and
	53	+ # cascades to every subsequent case. Modulo cycling removes the
	54	+ # fallback branch entirely.
	55	+ idx = ChunkGatedDeltaRuleFwdKktSolveGenerator._case_index
	56	+ B, T, H, K, BT = pool[idx % len(pool)]
	57	+ ChunkGatedDeltaRuleFwdKktSolveGenerator._case_index += 1
	58	+
	59	+ k_input = case_config.inputs[0]
	60	+ g_input = case_config.inputs[1]
	61	+ beta_input = case_config.inputs[2]
	62	+
	63	+ k_input.shape = [B, T, H, K]
	64	+ g_input.shape = [B, T, H]
	65	+ beta_input.shape = [B, T, H]
	66	+
	67	+ # Align dtypes to avoid cross-dtype explosion
	68	+ g_input.dtype = k_input.dtype
	69	+ beta_input.dtype = k_input.dtype
	70	+
	71	+ return case_config

Atests/atk_tests/triton/chunk_gated_delta_rule_fwd_kkt_solve/triton_chunk_gated_delta_rule_fwd_kkt_solve.py+132-0

@@ -0,0 +1,132 @@
	1	+# Copyright (c) 2026, HUAWEI CORPORATION. All rights reserved.
	2	+# pylint: disable=duplicate-code
	3	+# - duplicate-code: ATK golden is copied from tests/unit_tests/triton to keep
	4	+# both checks aligned with the operator formula.
	5	+"""ATK adapter for chunk_gated_delta_rule_fwd_kkt_solve (non-varlen, with g and beta).
	6	+
	7	+Inputs (named, matches chunk_gated_delta_rule_fwd_kkt_solve.yaml):
	8	+ k: (B, T, H, K)
	9	+ g: (B, T, H)
	10	+ beta: (B, T, H)
	11	+ chunk_size: int (must be 64)
	12	+
	13	+Output:
	14	+ A: (B, T, H, BT) -- the full (I + A)^{-1}
	15	+"""
	16	+
	17	+from typing import Optional
	18	+
	19	+import torch
	20	+
	21	+from atk.configs.dataset_config import InputDataset
	22	+from atk.tasks.api_execute import register
	23	+from atk.tasks.api_execute.base_api import BaseApi
	24	+from atk.tasks.api_execute.triton_base_api import TritonBaseApi
	25	+
	26	+from mindspeed_ops.api.triton.chunk_gated_delta_rule_fwd_kkt_solve import (
	27	+ chunk_gated_delta_rule_fwd_kkt_solve,
	28	+)
	29	+
	30	+
	31	+def _chunk_gated_delta_rule_fwd_kkt_solve_golden(
	32	+ k: torch.Tensor,
	33	+ g: Optional[torch.Tensor] = None,
	34	+ beta: Optional[torch.Tensor] = None,
	35	+ chunk_size: int = 64,
	36	+ output_dtype: torch.dtype = torch.float32,
	37	+) -> torch.Tensor:
	38	+ """Golden reference copied from tests/unit_tests/triton/test_chunk_gated_delta_rule_fwd_kkt_solve.py."""
	39	+ B, T, H, K_dim = k.shape
	40	+ BT = chunk_size
	41	+ k_float = k.float()
	42	+
	43	+ A = torch.zeros(B, T, H, BT, device=k.device, dtype=torch.float32)
	44	+ eye = torch.eye(BT, device=k.device, dtype=torch.float32)
	45	+
	46	+ for b in range(B):
	47	+ bos, eos = 0, T
	48	+ T_local = T
	49	+ batch_idx = b
	50	+
	51	+ for h in range(H):
	52	+ k_bh = k_float[batch_idx, bos:eos, h, :]
	53	+ g_bh = g[batch_idx, bos:eos, h].float() if g is not None else None
	54	+ beta_bh = beta[batch_idx, bos:eos, h].float() if beta is not None else None
	55	+
	56	+ NT = (T_local + BT - 1) // BT
	57	+ for t_idx in range(NT):
	58	+ cs = t_idx * BT
	59	+ ce = min(cs + BT, T_local)
	60	+ ll = ce - cs
	61	+
	62	+ k_chunk = k_bh[cs:ce, :]
	63	+ kkt = k_chunk @ k_chunk.t()
	64	+
	65	+ A_chunk = torch.zeros(BT, BT, dtype=torch.float32, device=k.device)
	66	+ A_chunk[:ll, :ll] = kkt
	67	+
	68	+ tril_mask = torch.tril(torch.ones(BT, BT, device=k.device), diagonal=-1)
	69	+ A_chunk = A_chunk * tril_mask
	70	+
	71	+ if g_bh is not None:
	72	+ g_chunk = g_bh[cs:ce]
	73	+ g_diff = g_chunk[:, None] - g_chunk[None, :]
	74	+ g_diff = torch.clamp(g_diff, -50.0, 50.0)
	75	+ g_full = torch.zeros(BT, BT, dtype=torch.float32, device=k.device)
	76	+ g_full[:ll, :ll] = torch.exp(g_diff)
	77	+ A_chunk = A_chunk * g_full
	78	+
	79	+ if beta_bh is not None:
	80	+ beta_chunk = beta_bh[cs:ce]
	81	+ beta_full = torch.zeros(BT, device=k.device, dtype=torch.float32)
	82	+ beta_full[:ll] = beta_chunk
	83	+ A_chunk = A_chunk * beta_full[:, None]
	84	+ else:
	85	+ beta_full = torch.zeros(BT, device=k.device, dtype=torch.float32)
	86	+ beta_full[:ll] = 1.0
	87	+ A_chunk = A_chunk * beta_full[:, None]
	88	+
	89	+ inv = torch.inverse(A_chunk + eye)
	90	+ A[batch_idx, bos + cs : bos + ce, h, :] = inv[:ll, :ll].to(output_dtype)
	91	+
	92	+ return A
	93	+
	94	+
	95	+@register("torch_chunk_gated_delta_rule_fwd_kkt_solve")
	96	+class TorchChunkGatedDeltaRuleFwdKktSolveApi(BaseApi):
	97	+ """CPU/GPU golden reference for chunk_gated_delta_rule_fwd_kkt_solve."""
	98	+
	99	+ def __call__(self, input_data: InputDataset, with_output: bool = False):
	100	+ k = input_data.kwargs.get('k')
	101	+ g = input_data.kwargs.get('g')
	102	+ beta = input_data.kwargs.get('beta')
	103	+ chunk_size = input_data.kwargs.get('chunk_size', 64)
	104	+
	105	+ output = _chunk_gated_delta_rule_fwd_kkt_solve_golden(
	106	+ k.float(),
	107	+ g=g.float() if g is not None else None,
	108	+ beta=beta.float() if beta is not None else None,
	109	+ chunk_size=chunk_size,
	110	+ output_dtype=k.dtype,
	111	+ )
	112	+ return output
	113	+
	114	+
	115	+@register("triton_chunk_gated_delta_rule_fwd_kkt_solve")
	116	+class TritonChunkGatedDeltaRuleFwdKktSolveApi(TritonBaseApi):
	117	+ """Triton path: invoke mindspeed_ops chunk_gated_delta_rule_fwd_kkt_solve."""
	118	+
	119	+ def __call__(self, input_data: InputDataset, with_output: bool = False):
	120	+ k = input_data.kwargs.get('k')
	121	+ g = input_data.kwargs.get('g')
	122	+ beta = input_data.kwargs.get('beta')
	123	+ chunk_size = input_data.kwargs.get('chunk_size', 64)
	124	+
	125	+ output = chunk_gated_delta_rule_fwd_kkt_solve(
	126	+ k,
	127	+ g=g,
	128	+ beta=beta,
	129	+ chunk_size=chunk_size,
	130	+ output_dtype=k.dtype,
	131	+ )
	132	+ return output

Atests/unit_tests/triton/test_chunk_gated_delta_rule_fwd_kkt_solve.py+277-0

@@ -0,0 +1,277 @@
	1	+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
	2	+# pylint: disable=duplicate-code
	3	+# - duplicate-code: UT and ATK share the same reference implementation to keep
	4	+# both checks aligned with the operator formula.
	5	+
	6	+
	7	+from typing import Optional
	8	+
	9	+import pytest
	10	+import torch
	11	+
	12	+from mindspeed_ops.api.triton.chunk_gated_delta_rule_fwd_kkt_solve import (
	13	+ chunk_gated_delta_rule_fwd_kkt_solve,
	14	+)
	15	+from mindspeed_ops.api.triton.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
	16	+from mindspeed_ops.api.triton.solve_tril import solve_tril
	17	+from mindspeed_ops.api.triton.utils import get_available_device
	18	+from tests.utils import assert_close, print_diff
	19	+
	20	+
	21	+# ---------------------------------------------------------------------------
	22	+# Golden reference: (I + A)^{-1}, A = beta * tril(K @ K^T) * exp(clamp(g_i-g_j, ±50))
	23	+# ---------------------------------------------------------------------------
	24	+
	25	+
	26	+def chunk_gated_delta_rule_fwd_kkt_solve_golden(
	27	+ k: torch.Tensor,
	28	+ g: Optional[torch.Tensor] = None,
	29	+ beta: Optional[torch.Tensor] = None,
	30	+ cu_seqlens: Optional[torch.LongTensor] = None,
	31	+ chunk_size: int = 64,
	32	+ output_dtype: torch.dtype = torch.float32,
	33	+) -> torch.Tensor:
	34	+ B, T, H, K_dim = k.shape
	35	+ BT = chunk_size
	36	+ k_float = k.float()
	37	+
	38	+ A = torch.zeros(B, T, H, BT, device=k.device, dtype=torch.float32)
	39	+
	40	+ num_seqs = len(cu_seqlens) - 1 if cu_seqlens is not None else B
	41	+ seq_offsets = cu_seqlens.tolist() if cu_seqlens is not None else [i * T for i in range(B + 1)]
	42	+
	43	+ eye = torch.eye(BT, device=k.device, dtype=torch.float32)
	44	+ tril_mask = torch.tril(torch.ones(BT, BT, device=k.device), diagonal=-1)
	45	+
	46	+ for batch_idx in range(num_seqs):
	47	+ bos = seq_offsets[batch_idx]
	48	+ T_local = seq_offsets[batch_idx + 1] - seq_offsets[batch_idx]
	49	+ NT = (T_local + BT - 1) // BT
	50	+ for h in range(H):
	51	+ k_bh = k_float[0, bos : bos + T_local, h] if cu_seqlens is not None else k_float[batch_idx, :, h]
	52	+ g_bh = (
	53	+ None if g is None else (g[0, bos : bos + T_local, h] if cu_seqlens is not None else g[batch_idx, :, h])
	54	+ )
	55	+ beta_bh = (
	56	+ None
	57	+ if beta is None
	58	+ else (beta[0, bos : bos + T_local, h] if cu_seqlens is not None else beta[batch_idx, :, h])
	59	+ )
	60	+ for t_idx in range(NT):
	61	+ cs = t_idx * BT
	62	+ ce = min(cs + BT, T_local)
	63	+ ll = ce - cs
	64	+
	65	+ k_chunk = k_bh[cs:ce, :]
	66	+ kkt = k_chunk @ k_chunk.t()
	67	+
	68	+ A_chunk = torch.zeros(BT, BT, dtype=torch.float32, device=k.device)
	69	+ A_chunk[:ll, :ll] = kkt
	70	+ A_chunk = A_chunk * tril_mask
	71	+
	72	+ if g_bh is not None:
	73	+ g_chunk = g_bh[cs:ce]
	74	+ g_diff = g_chunk[:, None] - g_chunk[None, :]
	75	+ g_diff = torch.clamp(g_diff, -50.0, 50.0)
	76	+ g_full = torch.zeros(BT, BT, dtype=torch.float32, device=k.device)
	77	+ g_full[:ll, :ll] = torch.exp(g_diff)
	78	+ A_chunk = A_chunk * g_full
	79	+
	80	+ beta_full = torch.zeros(BT, device=k.device, dtype=torch.float32)
	81	+ beta_full[:ll] = beta_bh[cs:ce] if beta_bh is not None else 1.0
	82	+ A_chunk = A_chunk * beta_full[:, None]
	83	+
	84	+ inv = torch.inverse(A_chunk + eye)
	85	+ # keep all BT columns of the first ll rows (kernel writes the full
	86	+ # (I+A)^{-1} rows; padding rows stay zero)
	87	+ if cu_seqlens is not None:
	88	+ A[0, bos + cs : bos + ce, h, :] = inv[:ll, :].to(output_dtype)
	89	+ else:
	90	+ A[batch_idx, cs:ce, h, :] = inv[:ll, :].to(output_dtype)
	91	+
	92	+ return A
	93	+
	94	+
	95	+# ---------------------------------------------------------------------------
	96	+# Parameter spaces
	97	+# ---------------------------------------------------------------------------
	98	+
	99	+DTYPES = [torch.float32, torch.bfloat16, torch.float16]
	100	+DTYPE_NAME = {torch.float32: "fp32", torch.bfloat16: "bf16", torch.float16: "fp16"}
	101	+
	102	+# (B, T, H, K, BT) -- the fused kernel only supports BT=64.
	103	+TEST_SHAPES = [
	104	+ (1, 64, 2, 64, 64), # NT=1 boundary (odd pair slot is padding)
	105	+ (1, 128, 4, 128, 64), # NT=2 typical (exactly one pair)
	106	+ (1, 256, 4, 128, 64), # NT=4 typical
	107	+ (2, 128, 4, 128, 64), # B=2
	108	+ (1, 192, 4, 128, 64), # NT=3, partial last chunk AND odd chunk count
	109	+ (1, 2048, 4, 128, 64), # long sequence, NT=32
	110	+ (1, 8192, 2, 128, 64), # very long sequence, NT=128
	111	+ (2, 1024, 4, 128, 64), # long sequence with B>1
	112	+ (2, 8192, 32, 128, 64), # very large grid (multi-buffered task loop)
	113	+ (21, 195, 32, 128, 64), # large B * partial last chunk (regression guard)
	114	+]
	115	+
	116	+
	117	+def _shape_id(shape):
	118	+ B, T, H, K, BT = shape
	119	+ return f"B{B}_T{T}_H{H}_K{K}_BT{BT}"
	120	+
	121	+
	122	+def _dtype_thresholds(dtype):
	123	+ if dtype == torch.float32:
	124	+ return 1e-4, 1e-5
	125	+ if dtype == torch.bfloat16:
	126	+ return 1e-2, 5e-3
	127	+ if dtype == torch.float16:
	128	+ return 1e-3, 5e-4
	129	+ raise ValueError(f"unsupported dtype {dtype}")
	130	+
	131	+
	132	+# ---------------------------------------------------------------------------
	133	+# Tests: golden x dtype/shape + no-g/beta + varlen + cross-check + fallback
	134	+# ---------------------------------------------------------------------------
	135	+
	136	+
	137	+class TestChunkGatedDeltaRuleFwdKktSolve:
	138	+ def setup_method(self):
	139	+ torch.manual_seed(42)
	140	+
	141	+ @pytest.mark.parametrize("dtype", DTYPES, ids=lambda d: DTYPE_NAME[d])
	142	+ @pytest.mark.parametrize("shape", TEST_SHAPES, ids=_shape_id)
	143	+ def test_fwd_with_g_beta(self, shape, dtype):
	144	+ """Fused kkt+solve with g and beta against golden reference."""
	145	+ device = get_available_device()
	146	+ if device == "cpu":
	147	+ pytest.skip("Triton requires NPU/GPU device")
	148	+
	149	+ B, T, H, K, BT = shape
	150	+ k = torch.randn(B, T, H, K, device=device, dtype=dtype) * 0.1
	151	+ g = torch.randn(B, T, H, device=device, dtype=dtype) * 0.01
	152	+ beta = torch.randn(B, T, H, device=device, dtype=dtype) * 0.1
	153	+
	154	+ # CPU small-op golden (task-spec caliber: align to CPU when no GPU ref)
	155	+ ref = chunk_gated_delta_rule_fwd_kkt_solve_golden(
	156	+ k.cpu().float(),
	157	+ g=g.cpu().float(),
	158	+ beta=beta.cpu().float(),
	159	+ chunk_size=BT,
	160	+ output_dtype=dtype,
	161	+ )
	162	+ out = chunk_gated_delta_rule_fwd_kkt_solve(k, g=g, beta=beta, chunk_size=BT, output_dtype=dtype)
	163	+ out = out.cpu()
	164	+
	165	+ ratio, atol = _dtype_thresholds(dtype)
	166	+ print_diff("fwd_kkt_solve_with_g_beta", ref, out, atol)
	167	+ assert_close("fwd_kkt_solve_with_g_beta", ref, out, ratio, err_atol=atol)
	168	+
	169	+ @pytest.mark.parametrize("shape", TEST_SHAPES, ids=_shape_id)
	170	+ def test_fwd_no_g_beta(self, shape):
	171	+ """Fused kkt+solve without g and beta (beta treated as ones)."""
	172	+ device = get_available_device()
	173	+ if device == "cpu":
	174	+ pytest.skip("Triton requires NPU/GPU device")
	175	+
	176	+ B, T, H, K, BT = shape
	177	+ dtype = torch.bfloat16
	178	+ k = torch.randn(B, T, H, K, device=device, dtype=dtype) * 0.1
	179	+
	180	+ ref = chunk_gated_delta_rule_fwd_kkt_solve_golden(k.cpu().float(), chunk_size=BT, output_dtype=dtype)
	181	+ out = chunk_gated_delta_rule_fwd_kkt_solve(k, chunk_size=BT, output_dtype=dtype)
	182	+ out = out.cpu()
	183	+
	184	+ ratio, atol = _dtype_thresholds(dtype)
	185	+ print_diff("fwd_kkt_solve_no_g_beta", ref, out, atol)
	186	+ assert_close("fwd_kkt_solve_no_g_beta", ref, out, ratio, err_atol=atol)
	187	+
	188	+ @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"])
	189	+ def test_fwd_varlen(self, dtype):
	190	+ """Variable-length sequences via cu_seqlens (unfused fallback path)."""
	191	+ device = get_available_device()
	192	+ if device == "cpu":
	193	+ pytest.skip("Triton requires NPU/GPU device")
	194	+
	195	+ B, H, K, BT = 1, 4, 128, 64
	196	+ seqlens = [64, 128, 256]
	197	+ T_total = sum(seqlens)
	198	+ cu_seqlens = torch.tensor(
	199	+ [0] + list(torch.cumsum(torch.tensor(seqlens), 0).tolist()), dtype=torch.int32, device=device
	200	+ )
	201	+
	202	+ torch.manual_seed(0)
	203	+ k = torch.randn(B, T_total, H, K, device=device, dtype=dtype) * 0.1
	204	+ g = -torch.randn(B, T_total, H, device=device, dtype=dtype).abs() * 0.01
	205	+ beta = torch.randn(B, T_total, H, device=device, dtype=dtype) * 0.1
	206	+
	207	+ cu_ref = cu_seqlens.cpu()
	208	+ ref = chunk_gated_delta_rule_fwd_kkt_solve_golden(
	209	+ k.cpu().float(),
	210	+ g=g.cpu().float(),
	211	+ beta=beta.cpu().float(),
	212	+ cu_seqlens=cu_ref,
	213	+ chunk_size=BT,
	214	+ output_dtype=dtype,
	215	+ )
	216	+ out = chunk_gated_delta_rule_fwd_kkt_solve(
	217	+ k, g=g, beta=beta, cu_seqlens=cu_seqlens, chunk_size=BT, output_dtype=dtype
	218	+ )
	219	+ out = out.cpu()
	220	+
	221	+ ratio, atol = _dtype_thresholds(dtype)
	222	+ print_diff("fwd_kkt_solve_varlen", ref, out, atol)
	223	+ assert_close("fwd_kkt_solve_varlen", ref, out, ratio, err_atol=atol)
	224	+
	225	+ offsets = [0] + list(torch.cumsum(torch.tensor(seqlens), 0).tolist())
	226	+ for i, _ in enumerate(seqlens):
	227	+ s, e = offsets[i], offsets[i + 1]
	228	+ sub_ref = chunk_gated_delta_rule_fwd_kkt_solve_golden(
	229	+ k[:, s:e].contiguous().cpu().float(),
	230	+ g=g[:, s:e].contiguous().cpu().float(),
	231	+ beta=beta[:, s:e].contiguous().cpu().float(),
	232	+ chunk_size=BT,
	233	+ output_dtype=dtype,
	234	+ )
	235	+ assert_close(f"varlen_seq{i}", sub_ref, out[:, s:e].contiguous().cpu(), ratio, err_atol=atol)
	236	+
	237	+ @pytest.mark.parametrize("shape", TEST_SHAPES, ids=_shape_id)
	238	+ def test_fwd_cross_check_unfused(self, shape):
	239	+ """Fused output must match the unfused chunk_scaled_dot_kkt_fwd + solve_tril."""
	240	+ device = get_available_device()
	241	+ if device == "cpu":
	242	+ pytest.skip("Triton requires NPU/GPU device")
	243	+
	244	+ B, T, H, K, BT = shape
	245	+ dtype = torch.bfloat16
	246	+ k = torch.randn(B, T, H, K, device=device, dtype=dtype) * 0.1
	247	+ g = torch.randn(B, T, H, device=device, dtype=dtype) * 0.01
	248	+ beta = torch.randn(B, T, H, device=device, dtype=dtype) * 0.1
	249	+
	250	+ fused = chunk_gated_delta_rule_fwd_kkt_solve(k, g=g, beta=beta, chunk_size=BT, output_dtype=dtype)
	251	+ unfused_A = chunk_scaled_dot_kkt_fwd(k=k, g=g, beta=beta, chunk_size=BT, output_dtype=torch.float32)
	252	+ unfused = solve_tril(A=unfused_A, output_dtype=dtype)
	253	+
	254	+ ratio, atol = _dtype_thresholds(dtype)
	255	+ print_diff("fwd_kkt_solve_vs_unfused", unfused, fused, atol)
	256	+ assert_close("fwd_kkt_solve_vs_unfused", unfused, fused, ratio, err_atol=atol)
	257	+
	258	+ def test_non64_chunk_size_fallback(self):
	259	+ """chunk_size != 64 routes to the unfused chain and must stay equivalent."""
	260	+ device = get_available_device()
	261	+ if device == "cpu":
	262	+ pytest.skip("Triton requires NPU/GPU device")
	263	+ k = torch.randn(1, 64, 2, 64, device=device, dtype=torch.float32) * 0.1
	264	+ g = torch.randn(1, 64, 2, device=device, dtype=torch.float32) * 0.01
	265	+ beta = torch.randn(1, 64, 2, device=device, dtype=torch.float32) * 0.1
	266	+
	267	+ out32 = chunk_gated_delta_rule_fwd_kkt_solve(k, g=g, beta=beta, chunk_size=32, output_dtype=torch.float32)
	268	+ ref32 = solve_tril(
	269	+ A=chunk_scaled_dot_kkt_fwd(k=k, g=g, beta=beta, chunk_size=32, output_dtype=torch.float32),
	270	+ output_dtype=torch.float32,
	271	+ )
	272	+ ratio, atol = _dtype_thresholds(torch.float32)
	273	+ assert_close("fallback_cs32", ref32, out32, ratio, err_atol=atol)
	274	+
	275	+
	276	+if __name__ == "__main__":
	277	+ pytest.main([__file__, "-v", "--tb=short"])