Update birwkv7.py

Inline Triton kernel and proper FLA arg order

birwkv7.py

@@ -2,6 +2,82 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+_TRITON_AVAILABLE = False
+try:
+    import triton
+    import triton.language as tl
+
+    @triton.jit
+    def _wkv7_fwd_kernel(
+        R, K, V, DECAY, A, O,
+        STATE_OUT, STATE_IN,
+        sab_scale, T,
+        stride_b, stride_t, stride_h,
+        H: tl.constexpr, D: tl.constexpr, BLOCK_D: tl.constexpr,
+        RETURN_STATE: tl.constexpr, HAS_INIT_STATE: tl.constexpr,
+    ):
+        pid = tl.program_id(0)
+        b_idx = pid // H
+        h_idx = pid % H
+        base = b_idx * stride_b + h_idx * stride_h
+
+        di = tl.arange(0, BLOCK_D)
+        dj = tl.arange(0, BLOCK_D)
+        mask_i = di < D
+        mask_j = dj < D
+
+        if HAS_INIT_STATE:
+            s_off = b_idx * (H * D * D) + h_idx * (D * D)
+            state_ptrs = STATE_IN + s_off + di[:, None] * D + dj[None, :]
+            state_mask = mask_i[:, None] & mask_j[None, :]
+            state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32)
+        else:
+            state = tl.zeros((BLOCK_D, BLOCK_D), dtype=tl.float32)
+
+        for t in range(T):
+            t_off = base + t * stride_t
+            kt = tl.load(K + t_off + dj, mask=mask_j, other=0.0).to(tl.float32)
+            vt = tl.load(V + t_off + di, mask=mask_i, other=0.0).to(tl.float32)
+            rt = tl.load(R + t_off + dj, mask=mask_j, other=0.0).to(tl.float32)
+            dt = tl.load(DECAY + t_off + dj, mask=mask_j, other=1.0).to(tl.float32)
+            at = tl.load(A + t_off + dj, mask=mask_j, other=0.0).to(tl.float32)
+
+            sa = tl.sum(state * (-kt)[None, :], axis=1)
+            ka = kt * at
+            sab = sa[:, None] * ka[None, :]
+            state = state * dt[None, :] + sab_scale * sab + vt[:, None] * kt[None, :]
+            state = tl.minimum(tl.maximum(state, -10.0), 10.0)
+
+            out_t = tl.sum(state * rt[None, :], axis=1)
+            tl.store(O + t_off + di, out_t, mask=mask_i)
+
+        if RETURN_STATE:
+            s_off = b_idx * (H * D * D) + h_idx * (D * D)
+            state_ptrs = STATE_OUT + s_off + di[:, None] * D + dj[None, :]
+            state_mask = mask_i[:, None] & mask_j[None, :]
+            tl.store(state_ptrs, state, mask=state_mask)
+
+    def _wkv7_scan_triton(r, decay, k, v, a, sab_scale):
+        B, T, H, D = r.shape
+        r, k, v, decay, a = [x.contiguous() for x in (r, k, v, decay, a)]
+        o = torch.empty_like(r)
+        stride_b, stride_t, stride_h = T * H * D, H * D, D
+        BLOCK_D = triton.next_power_of_2(D)
+        _wkv7_fwd_kernel[(B * H,)](
+            r, k, v, decay, a, o,
+            None, None,
+            float(sab_scale), T,
+            stride_b, stride_t, stride_h,
+            H=H, D=D, BLOCK_D=BLOCK_D,
+            RETURN_STATE=False, HAS_INIT_STATE=False,
+        )
+        return o
+
+    if torch.cuda.is_available():
+        _TRITON_AVAILABLE = True
+except Exception:
+    pass
+
 _FLA_AVAILABLE = False
 try:
     import torch.distributed.tensor as _tdt
@@ -83,13 +159,13 @@ class BiRWKV7Layer(nn.Module):
     def _wkv7_scan_fla(self, r, w, k, v, a, sab_scale):
         B, T, H, D = r.shape
         orig_dtype = r.dtype
-        r, w, k, v, a = [x.bfloat16() for x in (r, w, k, v, a)]
+        r, w, k, v, a = [x.float() for x in (r, w, k, v, a)]
         k_scaled = k * (D ** -0.5)
         w_log = -0.6065306597633104 * torch.sigmoid(w)
         a_sig = torch.sigmoid(a)
         a_fla = -k_scaled
         b_fla = sab_scale * k_scaled * a_sig
-        o, _ = _fla_chunk_rwkv7(r, w_log, k_scaled, v, a_fla, b_fla, scale=1.0)
+        o, _ = _fla_chunk_rwkv7(r, k_scaled, v, a_fla, b_fla, log_w=w_log, scale=1.0, head_first=False)
         return o.to(orig_dtype)
 
     def _wkv7_scan_python(self, r, w, k, v, a, sab_scale):
@@ -120,6 +196,14 @@ class BiRWKV7Layer(nn.Module):
         return torch.stack(outputs, dim=1).to(orig_dtype)
 
     def _wkv7_scan(self, r, w, k, v, a, sab_scale):
+        if _TRITON_AVAILABLE and r.is_cuda:
+            B, T, H, D = r.shape
+            orig_dtype = r.dtype
+            r, w, k, v, a = [x.float() for x in (r, w, k, v, a)]
+            k = k * (D ** -0.5)
+            decay = torch.exp(-0.6065306597633104 * torch.sigmoid(w))
+            a = torch.sigmoid(a)
+            return _wkv7_scan_triton(r, decay, k, v, a, sab_scale).to(orig_dtype)
         if _FLA_AVAILABLE and r.is_cuda:
             return self._wkv7_scan_fla(r, w, k, v, a, sab_scale)
         return self._wkv7_scan_python(r, w, k, v, a, sab_scale)