Add RWKV-6 and RWKV-7 recurrence operation modules with tests

- Implement RWKV-6 recurrence kernel with multi-head support and variable-length sequence packing.
- Implement RWKV-7 Diagonal + Low-Rank (DPLR) recurrence kernel and its multiplicative parameterization.
- Add unit tests for RWKV-6 and RWKV-7 to ensure output consistency with XLA implementations.
- Include gradient shape tests for RWKV-6 and RWKV-7 to validate backward compatibility.
- Ensure compatibility with Triton and XLA backends for both RWKV-6 and RWKV-7 operations.

Author: erfanzar

README.md

@@ -34,7 +34,7 @@ ejKernel is a production-grade kernel library for JAX that provides highly optim
 
 ### State-of-the-Art Operations
 
-- **20+ Deep Learning Operations**: Flash Attention v2, Ring Attention, Page Attention, Block Sparse, GLA, Lightning, State Space Models (Mamba), and more
+- **25+ Deep Learning Operations**: Flash Attention v2, Ring Attention, Page Attention, Block Sparse, GLA, Lightning, State Space Models (Mamba), RWKV (v4/v6/v7), and more
 - **Memory Efficiency**: Custom VJP implementations with O(N) memory complexity for attention
 - **Distributed Support**: Full shard_map integration for model and data parallelism
 - **Mixed Precision**: Comprehensive dtype support with automatic gradient conversion
@@ -266,6 +266,15 @@ kernel_with_custom_grad.defvjp(kernel_fwd, kernel_bwd)
 | **Prefill Page Attention**       | Page attention prefill phase     | O(N)   | Separate prefill handling                              |
 | **Scaled Dot-Product Attention** | Standard attention               | O(N²)  | Basic reference implementation                         |
 
+### Recurrent Linear Attention (RWKV)
+
+| Operation     | Description                          | Key Features                                       |
+| ------------- | ------------------------------------ | -------------------------------------------------- |
+| **RWKV-4**    | Time-mix recurrence                  | Numerically stable (α,β,ε) state, O(N) memory      |
+| **RWKV-6**    | Multi-head linear attention          | Variable-length packing, reverse mode, O(N) memory |
+| **RWKV-7**    | DPLR (Diagonal + Low-Rank) recurrence| (a,b) parameterization, state-space inspired       |
+| **RWKV-7 Mul**| Multiplicative RWKV-7 variant        | (kk,a) reparameterization for optimized kernels    |
+
 ### Other Operations
 
 | Operation             | Description                          | Use Case                  |
@@ -307,6 +316,10 @@ kernel_with_custom_grad.defvjp(kernel_fwd, kernel_bwd)
 | Prefill Page Attention   | -            | ✅           | ✅              |
 | State Space v1           | -            | -            | ✅              |
 | State Space v2           | -            | -            | ✅              |
+| RWKV-4                   | ✅           | -            | ✅              |
+| RWKV-6                   | ✅           | -            | ✅              |
+| RWKV-7                   | ✅           | -            | ✅              |
+| RWKV-7 Mul               | ✅           | -            | ✅              |
 
 ✅ = Production ready | 🚧 = Under development | - = Not available
 

ejkernel/__init__.py

@@ -41,7 +41,7 @@
 import os as _os
 
 _os.environ["TF_GPU_ALLOCATOR"] = "cuda_malloc_async"
-__version__ = "0.0.40"
+__version__ = "0.0.45"
 
 from . import kernels, modules, types, utils, xla_utils
 from .kernels import Backend, Platform, kernel_registry

ejkernel/kernels/_triton/__init__.py

@@ -45,6 +45,9 @@
 from .ragged_page_attention_v3 import ragged_page_attention_v3
 from .recurrent import recurrent
 from .ring_attention import ring_attention
+from .rwkv4 import rwkv4
+from .rwkv6 import rwkv6
+from .rwkv7 import rwkv7, rwkv7_mul
 from .unified_attention import unified_attention
 
 __all__ = (
@@ -61,5 +64,9 @@
     "recurrent",
     "recurrent_gla",
     "ring_attention",
+    "rwkv4",
+    "rwkv6",
+    "rwkv7",
+    "rwkv7_mul",
     "unified_attention",
 )

ejkernel/kernels/_triton/rwkv4/__init__.py

@@ -0,0 +1,22 @@
+# Copyright 2025 The EasyDeL/ejKernel Author @erfanzar (Erfan Zare Chavoshi).
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Triton backend for RWKV-4 time-mix recurrence."""
+
+from ._interface import rwkv4
+
+__all__ = [
+    "rwkv4",
+]
+

ejkernel/kernels/_triton/rwkv4/_interface.py

@@ -0,0 +1,114 @@
+# Copyright 2025 The EasyDeL/ejKernel Author @erfanzar (Erfan Zare Chavoshi).
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""RWKV-4 recurrent time-mix kernel (Triton)."""
+
+from __future__ import annotations
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+import jaxtyping
+from beartype import beartype
+from jaxtyping import Array, Float
+
+from ..._registry import Backend, Platform, kernel_registry
+from ..._xla.rwkv4 import rwkv4 as xla_rwkv4
+from ._triton_impl_fwd import fwd_triton_impl
+
+
+def _fwd_call(
+    w: Float[Array, "chans"],
+    u: Float[Array, "chans"],
+    k: Float[Array, "batch seq_len chans"],
+    v: Float[Array, "batch seq_len chans"],
+    state: Float[Array, "batch three chans"] | None,
+):
+    state_was_none = state is None
+    if state is None:
+        bsz, _, chans = k.shape
+        alpha0 = jnp.zeros((bsz, chans), dtype=jnp.float32)
+        beta0 = jnp.zeros((bsz, chans), dtype=jnp.float32)
+        eps0 = jnp.full((bsz, chans), -1e30, dtype=jnp.float32)
+        state = jnp.stack([alpha0, beta0, eps0], axis=1)
+
+    w_neg = -jnp.exp(w.astype(jnp.float32))
+    o, final_state = fwd_triton_impl(w_neg, u.astype(jnp.float32), k, v, state.astype(jnp.float32))
+    residual = (w, u, k, v, state, state_was_none)
+    return (o, final_state), residual
+
+
+def _bwd_call(
+    residual,
+    grads,
+):
+    (w, u, k, v, state, state_was_none) = residual
+    do, dstate = grads
+
+    def f(w_, u_, k_, v_, state_):
+        return xla_rwkv4(w_, u_, k_, v_, state_)
+
+    (o_ref, state_ref), vjp = jax.vjp(f, w, u, k, v, state)
+    del o_ref, state_ref
+    dw, du, dk, dv, dstate_in = vjp((do, dstate))
+    if state_was_none:
+        dstate_in = None
+    return dw, du, dk, dv, dstate_in
+
+
+@partial(jax.custom_vjp)
+def _rwkv4(
+    w: Float[Array, "chans"],
+    u: Float[Array, "chans"],
+    k: Float[Array, "batch seq_len chans"],
+    v: Float[Array, "batch seq_len chans"],
+    state: Float[Array, "batch three chans"] | None = None,
+) -> tuple[Float[Array, "batch seq_len chans"], Float[Array, "batch three chans"]]:
+    if state is None:
+        bsz, _, chans = k.shape
+        alpha0 = jnp.zeros((bsz, chans), dtype=jnp.float32)
+        beta0 = jnp.zeros((bsz, chans), dtype=jnp.float32)
+        eps0 = jnp.full((bsz, chans), -1e30, dtype=jnp.float32)
+        state = jnp.stack([alpha0, beta0, eps0], axis=1)
+
+    w_neg = -jnp.exp(w.astype(jnp.float32))
+    return fwd_triton_impl(w_neg, u.astype(jnp.float32), k, v, state.astype(jnp.float32))
+
+
+_rwkv4.defvjp(_fwd_call, _bwd_call)
+
+
+@kernel_registry.register("rwkv4", Platform.TRITON, Backend.GPU)
+@jaxtyping.jaxtyped(typechecker=beartype)
+def rwkv4(
+    w: Float[Array, "chans"],
+    u: Float[Array, "chans"],
+    k: Float[Array, "batch seq_len chans"],
+    v: Float[Array, "batch seq_len chans"],
+    state: Float[Array, "batch three chans"] | None = None,
+) -> tuple[Float[Array, "batch seq_len chans"], Float[Array, "batch three chans"]]:
+    """RWKV-4 time-mix recurrence (Triton GPU implementation).
+
+    Args:
+        w: Time-decay parameter in log space `[C]`.
+        u: Time-mix bias `[C]`.
+        k: Key tensor `[B, T, C]`.
+        v: Value tensor `[B, T, C]`.
+        state: Optional initial state `[B, 3, C]` (alpha, beta, eps).
+
+    Returns:
+        Tuple of (output `[B, T, C]`, final_state `[B, 3, C]`).
+    """
+    return _rwkv4(w, u, k, v, state)

ejkernel/kernels/_triton/rwkv4/_triton_impl_fwd.py

@@ -0,0 +1,128 @@
+# Copyright 2025 The EasyDeL/ejKernel Author @erfanzar (Erfan Zare Chavoshi).
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""RWKV-4 forward pass Triton kernel implementation.
+
+This module provides the Triton GPU kernel for RWKV-4 time-mix recurrence.
+The kernel processes sequences in a numerically-stable manner using the
+(alpha, beta, eps) state formulation.
+"""
+
+from __future__ import annotations
+
+import jax
+import triton
+import triton.language as tl
+from jax import numpy as jnp
+from jaxtyping import Array, Float
+
+from ejkernel.callib import cdiv, triton_call
+
+
+@triton.jit
+def _rwkv4_fwd_kernel(
+    w_ptr,  # [-exp(w_raw)], [C]
+    u_ptr,  # [C]
+    k_ptr,  # [B, T, C]
+    v_ptr,  # [B, T, C]
+    state_ptr,  # [B, 3, C]
+    o_ptr,  # [B, T, C]
+    state_out_ptr,  # [B, 3, C]
+    T: tl.constexpr,
+    C: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+):
+    b = tl.program_id(0)
+    c_blk = tl.program_id(1)
+
+    cs = c_blk * BLOCK_C + tl.arange(0, BLOCK_C)
+    cmask = cs < C
+
+    w = tl.load(w_ptr + cs, mask=cmask, other=0.0).to(tl.float32)
+    u = tl.load(u_ptr + cs, mask=cmask, other=0.0).to(tl.float32)
+
+    base_state = (b * 3) * C
+    alpha = tl.load(state_ptr + base_state + 0 * C + cs, mask=cmask, other=0.0).to(tl.float32)
+    beta = tl.load(state_ptr + base_state + 1 * C + cs, mask=cmask, other=0.0).to(tl.float32)
+    eps = tl.load(state_ptr + base_state + 2 * C + cs, mask=cmask, other=-1e30).to(tl.float32)
+
+    base_seq = b * T * C
+    for t in range(0, T):
+        off = base_seq + t * C + cs
+        kt = tl.load(k_ptr + off, mask=cmask, other=0.0).to(tl.float32)
+        vt = tl.load(v_ptr + off, mask=cmask, other=0.0).to(tl.float32)
+
+        ukt = u + kt
+        tau = tl.maximum(ukt, eps)
+        e1a = tl.exp(eps - tau)
+        e2a = tl.exp(ukt - tau)
+        wkv = (e1a * alpha + e2a * vt) / (e1a * beta + e2a)
+        tl.store(o_ptr + off, wkv.to(o_ptr.dtype.element_ty), mask=cmask)
+
+        w_eps = w + eps
+        eps_next = tl.maximum(w_eps, kt)
+        e1b = tl.exp(w_eps - eps_next)
+        e2b = tl.exp(kt - eps_next)
+        alpha = e1b * alpha + e2b * vt
+        beta = e1b * beta + e2b
+        eps = eps_next
+
+    base_state_out = (b * 3) * C
+    tl.store(state_out_ptr + base_state_out + 0 * C + cs, alpha.to(tl.float32), mask=cmask)
+    tl.store(state_out_ptr + base_state_out + 1 * C + cs, beta.to(tl.float32), mask=cmask)
+    tl.store(state_out_ptr + base_state_out + 2 * C + cs, eps.to(tl.float32), mask=cmask)
+
+
+def fwd_triton_impl(
+    w: Float[Array, "chans"],
+    u: Float[Array, "chans"],
+    k: Float[Array, "batch seq_len chans"],
+    v: Float[Array, "batch seq_len chans"],
+    state: Float[Array, "batch three chans"],
+) -> tuple[Float[Array, "batch seq_len chans"], Float[Array, "batch three chans"]]:
+    """Execute RWKV-4 forward pass on GPU via Triton.
+
+    Args:
+        w: Negative exponentiated time-decay `[C]` (already `-exp(w_raw)`).
+        u: Time-mix bias `[C]`.
+        k: Key tensor `[B, T, C]`.
+        v: Value tensor `[B, T, C]`.
+        state: Initial state `[B, 3, C]` (alpha, beta, eps).
+
+    Returns:
+        Tuple of (output `[B, T, C]`, final_state `[B, 3, C]`).
+    """
+    B, T, C = k.shape
+    out_shape = jax.ShapeDtypeStruct(k.shape, v.dtype)
+    state_shape = jax.ShapeDtypeStruct((B, 3, C), jnp.float32)
+
+    BLOCK_C = 128 if C >= 128 else 64 if C >= 64 else 32
+    grid = (B, cdiv(C, BLOCK_C))
+
+    o, state_out = triton_call(
+        w,
+        u,
+        k,
+        v,
+        state,
+        kernel=_rwkv4_fwd_kernel,
+        out_shape=[out_shape, state_shape],
+        name="ejkernel::triton::rwkv4_fwd",
+        grid=grid,
+        T=T,
+        C=C,
+        BLOCK_C=BLOCK_C,
+    )
+    return o, state_out
+

ejkernel/kernels/_triton/rwkv6/__init__.py

@@ -0,0 +1,22 @@
+# Copyright 2025 The EasyDeL/ejKernel Author @erfanzar (Erfan Zare Chavoshi).
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Triton backend for RWKV-6 recurrence."""
+
+from ._interface import rwkv6
+
+__all__ = [
+    "rwkv6",
+]
+