rwkv-ops 0.1.0__py3-none-any.whl

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_A_bwd.py

@@ -0,0 +1,95 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+import jax_triton as jt
+import jax
+import triton
+from ..triton_kernel.chunk_A_bwd import *
+from ..triton_kernel.utils import is_gather_supported
+from ..get_torch_devices_info import check_shared_mem
+
+
+def chunk_dplr_bwd_dqk_intra(
+    q: jax.Array,
+    k: jax.Array,
+    a: jax.Array,
+    b: jax.Array,
+    gi: jax.Array,
+    ge: jax.Array,
+    dAqk: jax.Array,
+    dAqb: jax.Array,
+    dAak: jax.Array,
+    dAab: jax.Array,
+    dqg: jax.Array,
+    dkg: jax.Array,
+    dag: jax.Array,
+    dbg: jax.Array,
+    dgk_last: jax.Array,
+    scale: float = 1.0,
+    chunk_size: int = 16,
+):
+    B, T, H, K = q.shape
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+    BK = (
+        min(64, triton.next_power_of_2(K))
+        if check_shared_mem()
+        else min(32, triton.next_power_of_2(K))
+    )
+
+    NT = triton.cdiv(T, BT)
+    NK = triton.cdiv(K, BK)
+    grid = (NK, NT, B * H)
+
+    out_shapes = [
+        jax.ShapeDtypeStruct(q.shape, q.dtype),
+        jax.ShapeDtypeStruct(k.shape, k.dtype),
+        jax.ShapeDtypeStruct(a.shape, a.dtype),
+        jax.ShapeDtypeStruct(b.shape, b.dtype),
+        jax.ShapeDtypeStruct(gi.shape, "float32"),
+        jax.ShapeDtypeStruct(gi.shape, "float32"),
+    ]
+
+    dq, dk, da, db, dgk, dgk_offset = jt.triton_call(
+        q,
+        k,
+        a,
+        b,
+        gi,
+        ge,
+        dAqk,
+        dAqb,
+        dAak,
+        dAab,
+        dqg,
+        dkg,
+        dag,
+        dbg,
+        T,
+        scale=scale,
+        H=H,
+        K=K,
+        BT=BT,
+        BC=BT,
+        BK=BK,
+        GATHER_SUPPORTED=is_gather_supported,
+        kernel=chunk_dplr_bwd_kernel_intra,
+        out_shape=out_shapes,
+        grid=grid,
+    )
+
+    def grid(meta):
+        return (NT, triton.cdiv(K, meta["BK"]), B * H)
+
+    dgk_output = jt.triton_call(
+        dgk,
+        dgk_offset,
+        dgk_last,
+        T,
+        H=H,
+        K=K,
+        BT=BT,
+        kernel=chunk_dplr_bwd_dgk_kernel,
+        out_shape=[jax.ShapeDtypeStruct(dgk.shape, "float32")],
+        grid=grid,
+    )[0]
+    return dq, dk, da, db, dgk_output

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_A_fwd.py

@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+
+import jax_triton as jt
+import jax
+import triton
+
+from ..triton_kernel.utils import is_gather_supported
+
+from ..triton_kernel.chunk_A_fwd import *
+
+
+def chunk_dplr_fwd_intra(
+    q: jax.Array,
+    k: jax.Array,
+    a: jax.Array,
+    b: jax.Array,
+    gi: jax.Array,
+    ge: jax.Array,
+    scale: float,
+    chunk_size: int,
+):
+    B, T, H, K = k.shape
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+
+    NT = triton.cdiv(T, BT)
+    shape = [B, T, H, BT]
+    out_shapes = [
+        jax.ShapeDtypeStruct(q.shape, q.dtype),
+        jax.ShapeDtypeStruct(k.shape, q.dtype),
+        jax.ShapeDtypeStruct(a.shape, q.dtype),
+        jax.ShapeDtypeStruct(b.shape, q.dtype),
+        jax.ShapeDtypeStruct(shape, q.dtype),
+        jax.ShapeDtypeStruct(shape, q.dtype),
+        jax.ShapeDtypeStruct(shape, "float32"),
+        jax.ShapeDtypeStruct(shape, "float32"),
+    ]
+    grid = (NT, B, H)
+    BK = triton.next_power_of_2(K)
+    qg, kg, ag, bg, Aqk, Aqb, Aab, Aak = jt.triton_call(
+        q,
+        k,
+        a,
+        b,
+        gi,
+        ge,
+        T,
+        scale=scale,
+        H=H,
+        K=K,
+        BT=BT,
+        BC=BT,
+        BK=BK,
+        GATHER_SUPPORTED=is_gather_supported,
+        kernel=chunk_dplr_fwd_A_kernel_intra_sub_intra,
+        out_shape=out_shapes,
+        grid=grid,
+    )
+    return Aab, Aqk, Aak, Aqb, qg, kg, ag, bg

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_h_bwd.py

@@ -0,0 +1,78 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+from typing import Optional, Tuple
+
+import jax_triton as jt
+import jax
+import triton
+
+from ..get_jax_devices_info import check_shared_mem
+from ..triton_kernel.chunk_h_bwd import *
+
+
+def chunk_dplr_bwd_dhu(
+    qg: jax.Array,
+    bg: jax.Array,
+    w: jax.Array,
+    gk: jax.Array,
+    h0: jax.Array,
+    dht: Optional[jax.Array],
+    do: jax.Array,
+    dv: jax.Array,
+    chunk_size: int = 64,
+) -> Tuple[jax.Array, jax.Array, jax.Array]:
+    B, T, H, K, V = *qg.shape, do.shape[-1]
+    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
+    BK = triton.next_power_of_2(K)
+    assert BK <= 256, (
+        "current kernel does not support head dimension being larger than 256."
+    )
+    # H100
+    if check_shared_mem("hopper"):
+        BV = 64
+        BC = 64 if K <= 128 else 32
+    elif check_shared_mem("ampere"):  # A100
+        BV = 32
+        BC = 32
+    else:  # Etc: 4090
+        BV = 16
+        BC = 16
+
+    N, NT = B, triton.cdiv(T, BT)
+    BC = min(BT, BC)
+    NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
+    assert NK == 1, (
+        "NK > 1 is not supported because it involves time-consuming synchronization"
+    )
+    dh_shape = (B, NT, H, K, V)
+    out_shapes = [
+        jax.ShapeDtypeStruct(dh_shape, dv.dtype),
+        jax.ShapeDtypeStruct((B, H, K, V), "float32"),
+        jax.ShapeDtypeStruct(dv.shape, dv.dtype),
+    ]
+
+    grid = (NK, NV, N * H)
+    dh, dh0, dv2 = jt.triton_call(
+        qg,
+        bg,
+        w,
+        gk,
+        dht,
+        dv,
+        do,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BC=BC,
+        BK=BK,
+        BV=BV,
+        kernel=chunk_dplr_bwd_kernel_dhu.fn,
+        out_shape=out_shapes,
+        grid=grid,
+        USE_FINAL_STATE_GRADIENT=dht is not None,
+        USE_INITIAL_STATE=h0 is not None,
+    )
+    return dh, dh0, dv2

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_h_fwd.py

@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+from typing import Optional, Tuple
+
+import jax_triton as jt
+import jax
+import triton
+
+from ..get_jax_devices_info import check_shared_mem
+from ..triton_kernel.chunk_h_fwd import *
+
+
+def chunk_dplr_fwd_h(
+    kg: jax.Array,
+    v: jax.Array,
+    w: jax.Array,
+    u: jax.Array,
+    bg: jax.Array,
+    gk: jax.Array,
+    initial_state: Optional[jax.Array] = None,
+    output_final_state: bool = False,
+    chunk_size: int = 64,
+) -> Tuple[jax.Array, jax.Array]:
+    B, T, H, K, V = *kg.shape, u.shape[-1]
+    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
+
+    N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
+    BK = triton.next_power_of_2(K)
+    assert BK <= 256, "current kernel does not support head dimension larger than 256."
+    # H100 can have larger block size
+
+    if check_shared_mem("hopper"):
+        BV = 64
+        BC = 64 if K <= 128 else 32
+    elif check_shared_mem("ampere"):  # A100
+        BV = 32
+        BC = 32
+    else:
+        BV = 16
+        BC = 16
+
+    BC = min(BT, BC)
+    NK = triton.cdiv(K, BK)
+    NV = triton.cdiv(V, BV)
+    assert NK == 1, (
+        "NK > 1 is not supported because it involves time-consuming synchronization"
+    )
+
+    out_shapes = [
+        jax.ShapeDtypeStruct((B, NT, H, K, V), kg.dtype),
+        jax.ShapeDtypeStruct([N, H, K, V], "float32"),
+        jax.ShapeDtypeStruct(u.shape, u.dtype),
+    ]
+    grid = (NK, NV, N * H)
+    if initial_state is None:
+        initial_state = jax.numpy.zeros([N, H, K, V], "float32")
+    h, final_state, v_new = jt.triton_call(
+        kg,
+        v,
+        w,
+        bg,
+        u,
+        gk,
+        initial_state,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BC=BC,
+        BK=BK,
+        BV=BV,
+        kernel=chunk_dplr_fwd_kernel_h.fn,
+        out_shape=out_shapes,
+        grid=grid,
+        STORE_FINAL_STATE=True,
+        USE_INITIAL_STATE=True,
+    )
+    return h, v_new, final_state

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_o_bwd.py

@@ -0,0 +1,150 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+from typing import Tuple
+
+import jax_triton as jt
+import jax
+import triton
+
+from ..get_torch_devices_info import check_shared_mem
+from ..triton_kernel.chunk_o_bwd import *
+
+
+def chunk_dplr_bwd_dv(
+    A_qk: jax.Array,
+    kg: jax.Array,
+    do: jax.Array,
+    dh: jax.Array,
+    chunk_size: int = 64,
+) -> jax.Array:
+    B, T, H, K, V = *kg.shape, do.shape[-1]
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+
+    NT = triton.cdiv(T, BT)
+
+    def grid(meta):
+        return (triton.cdiv(V, meta["BV"]), NT, B * H)
+
+    dv = jt.triton_call(
+        A_qk,
+        kg,
+        do,
+        dh,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        kernel=chunk_dplr_bwd_kernel_dv,
+        out_shape=jax.ShapeDtypeStruct(do.shape, do.dtype),
+        grid=grid,
+    )
+    return dv
+
+
+def chunk_dplr_bwd_o(
+    k: jax.Array,
+    b: jax.Array,
+    v: jax.Array,
+    v_new: jax.Array,
+    gk: jax.Array,
+    do: jax.Array,
+    h: jax.Array,
+    dh: jax.Array,
+    dv: jax.Array,
+    w: jax.Array,
+    chunk_size: int = 64,
+    scale: float = 1.0,
+) -> Tuple[jax.Array, jax.Array, jax.Array]:
+    B, T, H, K, V = *w.shape, v.shape[-1]
+
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+    NT = triton.cdiv(T, BT)
+
+    BK = (
+        min(triton.next_power_of_2(K), 64)
+        if check_shared_mem()
+        else min(triton.next_power_of_2(K), 32)
+    )
+    BV = (
+        min(triton.next_power_of_2(V), 64)
+        if check_shared_mem()
+        else min(triton.next_power_of_2(K), 32)
+    )
+    NK = triton.cdiv(K, BK)
+    grid = (NK, NT, B * H)
+
+    out_shapes = [
+        jax.ShapeDtypeStruct(k.shape, k.dtype),
+        jax.ShapeDtypeStruct(k.shape, k.dtype),
+        jax.ShapeDtypeStruct(w.shape, w.dtype),
+        jax.ShapeDtypeStruct(b.shape, b.dtype),
+        jax.ShapeDtypeStruct([B, NT, H, K], "float32"),
+    ]
+    dq, dk, dw, db, dgk_last = jt.triton_call(
+        v,
+        v_new,
+        h,
+        do,
+        dh,
+        w,
+        dv,
+        gk,
+        k,
+        b,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+        kernel=chunk_dplr_bwd_o_kernel,
+        out_shape=out_shapes,
+        grid=grid,
+    )
+    return dq, dk, dw, db, dgk_last
+
+
+def chunk_dplr_bwd_dAu(
+    v: jax.Array,
+    v_new: jax.Array,
+    do: jax.Array,
+    A_qb: jax.Array,
+    scale: float,
+    chunk_size: int = 64,
+) -> jax.Array:
+    B, T, H, V = v.shape
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+    NT = triton.cdiv(T, BT)
+
+    if check_shared_mem("ampere"):  # A100
+        BV = min(triton.next_power_of_2(V), 128)
+    elif check_shared_mem("ada"):  # 4090
+        BV = min(triton.next_power_of_2(V), 64)
+    else:
+        BV = min(triton.next_power_of_2(V), 32)
+
+    grid = (NT, B * H)
+    out_shapes = [
+        jax.ShapeDtypeStruct([B, T, H, BT], "float32"),
+        jax.ShapeDtypeStruct([B, T, H, BT], "float32"),
+        jax.ShapeDtypeStruct(v_new.shape, v_new.dtype),
+    ]
+    dA_qk, dA_qb, dv_new = jt.triton_call(
+        v,
+        do,
+        v_new,
+        A_qb,
+        T,
+        scale=scale,
+        H=H,
+        V=V,
+        BT=BT,
+        BV=BV,
+        grid=grid,
+        out_shape=out_shapes,
+        kernel=chunk_dplr_bwd_kernel_dAu,
+    )
+    return dv_new, dA_qk, dA_qb

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_o_fwd.py

@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+
+import jax_triton as jt
+import jax
+import triton
+
+from ..triton_kernel.chunk_o_fwd import *
+
+
+def chunk_dplr_fwd_o(
+    qg: jax.Array,
+    v: jax.Array,
+    v_new: jax.Array,
+    A_qk: jax.Array,
+    A_qb: jax.Array,
+    h: jax.Array,
+    chunk_size: int = 64,
+) -> jax.Array:
+    B, T, H, K, V = *qg.shape, v.shape[-1]
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+
+    NT = triton.cdiv(T, BT)
+
+    def grid(meta):
+        return (triton.cdiv(V, meta["BV"]), NT, B * H)
+
+    o = jt.triton_call(
+        qg,
+        v,
+        v_new,
+        A_qk,
+        A_qb,
+        h,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        kernel=chunk_dplr_fwd_kernel_o,
+        out_shape=jax.ShapeDtypeStruct(v.shape, v.dtype),
+        grid=grid,
+    )
+    return o

rwkv_ops/rwkv7_kernel/jax_kernel/cumsum.py

@@ -0,0 +1,34 @@
+from ..triton_kernel.cumsum import *
+import jax_triton as jt
+import jax
+import triton
+
+
+def chunk_rwkv6_fwd_cumsum(
+    g: jax.Array,
+    chunk_size: int,
+) -> jax.Array:
+    B, T, H, S = g.shape
+    BT = chunk_size
+    NT = triton.cdiv(T, BT)
+
+    out_shapes = [
+        jax.ShapeDtypeStruct(g.shape, "float32"),
+        jax.ShapeDtypeStruct(g.shape, "float32"),
+    ]
+
+    def grid(meta):
+        return (triton.cdiv(meta["S"], meta["BS"]), NT, B * H)
+
+    gi, ge = jt.triton_call(
+        g,
+        T,
+        H=H,
+        S=S,
+        BT=BT,
+        grid=grid,
+        kernel=chunk_rwkv6_fwd_cumsum_kernel,
+        out_shape=out_shapes,
+    )
+
+    return gi, ge

rwkv_ops/rwkv7_kernel/jax_kernel/wy_fast_bwd.py

@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+from typing import Tuple
+
+import jax_triton as jt
+import jax
+import triton
+
+
+from ..get_torch_devices_info import check_shared_mem
+from ..triton_kernel.wy_fast_bwd import *
+
+
+def chunk_dplr_bwd_wy(
+    A_ab_inv: jax.Array,
+    A_ak: jax.Array,
+    v: jax.Array,
+    ag: jax.Array,
+    dw: jax.Array,
+    du: jax.Array,
+    dv0: jax.Array,
+    chunk_size: int = 16,
+) -> Tuple[jax.Array, jax.Array, jax.Array]:
+    B, T, H, K, V = *dw.shape, du.shape[-1]
+    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
+
+    NT = triton.cdiv(T, BT)
+    BK = min(triton.next_power_of_2(K), 64)
+    BV = (
+        min(triton.next_power_of_2(V), 64)
+        if check_shared_mem()
+        else min(triton.next_power_of_2(V), 32)
+    )
+    grid = (NT, B * H)
+    out_shapes = [
+        jax.ShapeDtypeStruct(A_ak.shape, "float32"),
+        jax.ShapeDtypeStruct(A_ab_inv.shape, "float32"),
+        jax.ShapeDtypeStruct(v.shape, v.dtype),
+        jax.ShapeDtypeStruct(ag.shape, ag.dtype),
+    ]
+    dA_ak, dA_ab, dv, dag = jt.triton_call(
+        A_ab_inv,
+        A_ak,
+        ag,
+        v,
+        dw,
+        du,
+        dv0,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+        grid=grid,
+        kernel=prepare_wy_repr_bwd_kernel,
+        out_shape=out_shapes,
+    )
+    return dA_ab, dA_ak, dv, dag

rwkv_ops/rwkv7_kernel/jax_kernel/wy_fast_fwd.py

@@ -0,0 +1,86 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+from typing import Tuple
+
+import jax_triton as jt
+import jax
+import triton
+
+
+from ..triton_kernel.wy_fast_fwd import *
+
+
+def wu_fwd(
+    ag: jax.Array,
+    v: jax.Array,
+    A_ak: jax.Array,
+    A_ab_inv: jax.Array,
+    chunk_size: int,
+) -> Tuple[jax.Array, jax.Array]:
+    B, T, H, K, V = *ag.shape, v.shape[-1]
+    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
+
+    NT = triton.cdiv(T, BT)
+    BK = min(triton.next_power_of_2(K), 64)
+    BV = min(triton.next_power_of_2(V), 64)
+
+    out_shapes = [
+        jax.ShapeDtypeStruct(v.shape, v.dtype),
+        jax.ShapeDtypeStruct(ag.shape, ag.dtype),
+    ]
+    grid = (NT, B * H)
+    w, u = jt.triton_call(
+        ag,
+        v,
+        A_ab_inv,
+        A_ak,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+        grid=grid,
+        kernel=wu_fwd_kernel,
+        out_shape=out_shapes,
+    )
+    return w, u
+
+
+def prepare_wy_repr_fwd(
+    ag: jax.Array,
+    v: jax.Array,
+    A_ak: jax.Array,
+    A_ab: jax.Array,
+    chunk_size: int = 64,
+) -> Tuple[jax.Array, jax.Array, jax.Array]:
+    B, T, H, _ = ag.shape
+    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
+
+    NT = triton.cdiv(T, BT)
+    BC = min(BT, 32)
+    fwd_fn = (
+        prepare_wy_repr_fwd_kernel_chunk64
+        if BT == 64
+        else prepare_wy_repr_fwd_kernel_chunk32
+    )
+    grid = (NT, B * H)
+    A_ab_inv = jt.triton_call(
+        A_ab,
+        T,
+        H=H,
+        BT=BT,
+        BC=BC,
+        grid=grid,
+        kernel=fwd_fn,
+        out_shape=jax.ShapeDtypeStruct(A_ab.shape, A_ab.dtype),
+    )
+    w, u = wu_fwd(ag=ag, v=v, A_ak=A_ak, A_ab_inv=A_ab_inv, chunk_size=BT)
+    return w, u, A_ab_inv
+
+
+fwd_prepare_wy_repr = prepare_wy_repr_fwd
+
+fwd_wu = wu_fwd