rwkv-ops 0.6.1__py3-none-any.whl

rwkv_ops/rwkv6_kernel/torch_rwkv_kernel.py

@@ -0,0 +1,305 @@
+import os
+import torch
+from torch.utils.cpp_extension import load
+
+kernel_dir_name = "torch_kernel"
+
+use_rocm = "RWKV_USE_ROCM" in os.environ and os.environ["RWKV_USE_ROCM"] == "1"
+
+
+class RWKVKernelOperator:
+    def __init__(self, head_size, max_sequence_length):
+        current_dir = os.path.dirname(__file__)
+        # current_dir = os.pat
+        if use_rocm:
+            wkv6_cuda = load(
+                name="wkv6",
+                sources=[
+                    os.path.join(current_dir, f"{kernel_dir_name}/wkv6_op.cpp"),
+                    os.path.join(current_dir, f"{kernel_dir_name}/wkv6_cuda.cu"),
+                ],
+                # verbose=True, extra_cuda_cflags=[f"-D_N_={head_size}", f"-D_T_={max_sequence_length}"])
+                verbose=True,
+                extra_cuda_cflags=[
+                    "-fopenmp -ffast-math -munsafe-fp-atomics --gpu-max-threads-per-block=120 -enable-vectorize-compares",
+                    f"-D_N_={head_size}",
+                    f"-D_T_={max_sequence_length}",
+                ],
+            )
+        else:
+            wkv6_cuda = load(
+                name="wkv6",
+                sources=[
+                    os.path.join(current_dir, f"{kernel_dir_name}/wkv6_op.cpp"),
+                    os.path.join(current_dir, f"{kernel_dir_name}/wkv6_cuda.cu"),
+                ],
+                # verbose=True, extra_cuda_cflags=[f"-D_N_={head_size}", f"-D_T_={max_sequence_length}"])
+                verbose=True,
+                extra_cuda_cflags=[
+                    "-res-usage",
+                    "--use_fast_math",
+                    "-O3",
+                    "-Xptxas -O3",
+                    "--extra-device-vectorization",
+                    f"-D_N_={head_size}",
+                    f"-D_T_={max_sequence_length}",
+                ],
+            )
+
+        class RWKV_6(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, B, T, C, H, r, k, v, w, u):
+                if not isinstance(u, torch.Tensor):
+                    u = u.value
+                with torch.no_grad():
+                    assert r.dtype == k.dtype == v.dtype == w.dtype == u.dtype
+                    assert r.dtype in [torch.float32, torch.bfloat16, torch.float16]
+
+                    assert head_size == C // H
+                    ctx.B = B
+                    ctx.T = T
+                    ctx.C = C
+                    ctx.H = H
+                    assert r.is_contiguous()
+                    assert k.is_contiguous()
+                    assert v.is_contiguous()
+                    assert w.is_contiguous()
+                    assert u.is_contiguous()
+                    ctx.save_for_backward(r, k, v, w, u)
+
+                    y_dtype = r.dtype if r.dtype != torch.float16 else torch.float32
+
+                    y = torch.empty(
+                        (B, T, C),
+                        device=r.device,
+                        dtype=y_dtype,
+                        memory_format=torch.contiguous_format,
+                    )  # .uniform_(-100, 100)
+
+                    if r.dtype == torch.float32:
+                        wkv6_cuda.forward_fp32(B, T, C, H, r, k, v, w, u, y)
+                    elif r.dtype == torch.bfloat16:
+                        wkv6_cuda.forward_bf16(B, T, C, H, r, k, v, w, u, y)
+                    else:
+                        wkv6_cuda.forward_fp16(B, T, C, H, r, k, v, w, u, y)
+                    return y
+
+            @staticmethod
+            def backward(ctx, gy):
+                assert gy.is_cuda
+                with torch.no_grad():
+                    assert gy.dtype in [torch.bfloat16, torch.float32]
+                    B = ctx.B
+                    T = ctx.T
+                    C = ctx.C
+                    H = ctx.H
+                    assert gy.is_contiguous()
+                    r, k, v, w, u = ctx.saved_tensors
+                    y_dtype = r.dtype if r.dtype != torch.float16 else torch.float32
+
+                    gr = torch.empty(
+                        (B, T, C),
+                        device=gy.device,
+                        requires_grad=False,
+                        dtype=y_dtype,
+                        memory_format=torch.contiguous_format,
+                    )  # .uniform_(-100, 100)
+                    gk = torch.empty(
+                        (B, T, C),
+                        device=gy.device,
+                        requires_grad=False,
+                        dtype=y_dtype,
+                        memory_format=torch.contiguous_format,
+                    )  # .uniform_(-100, 100)
+                    gv = torch.empty(
+                        (B, T, C),
+                        device=gy.device,
+                        requires_grad=False,
+                        dtype=y_dtype,
+                        memory_format=torch.contiguous_format,
+                    )  # .uniform_(-100, 100)
+                    gw = torch.empty(
+                        (B, T, C),
+                        device=gy.device,
+                        requires_grad=False,
+                        dtype=y_dtype,
+                        memory_format=torch.contiguous_format,
+                    )  # .uniform_(-100, 100)
+                    gu = torch.empty(
+                        (B, C),
+                        device=gy.device,
+                        requires_grad=False,
+                        dtype=y_dtype,
+                        memory_format=torch.contiguous_format,
+                    )  # .uniform_(-100, 100)
+
+                    if r.dtype == torch.float32:
+                        wkv6_cuda.backward_fp32(
+                            B, T, C, H, r, k, v, w, u, gy, gr, gk, gv, gw, gu
+                        )
+                    elif r.dtype == torch.bfloat16:
+                        wkv6_cuda.backward_bf16(
+                            B, T, C, H, r, k, v, w, u, gy, gr, gk, gv, gw, gu
+                        )
+                    else:
+                        wkv6_cuda.backward_fp16(
+                            B, T, C, H, r, k, v, w, u, gy, gr, gk, gv, gw, gu
+                        )
+
+                    gu = torch.sum(gu, 0).view(H, C // H)
+
+                    return (None, None, None, None, gr, gk, gv, gw, gu)
+
+        class RWKV_6_with_state:
+            @staticmethod
+            def apply(B, T, C, H, S, s_map, r, k, v, w, u, s):
+                with torch.no_grad():
+                    assert s_map.dtype == torch.int64, (
+                        "s_map 必须为None 或者是长度为B的，int64类型的数组。"
+                    )
+                    assert (s is None and s_map is None) or (
+                        s is not None and s_map is not None
+                    ), "init_state与s_map必须同时为None 或者同时不为None"
+                    assert (
+                        r.dtype == k.dtype == v.dtype == w.dtype == u.dtype
+                        and r.dtype in [torch.float16, torch.float32, torch.bfloat16]
+                    ), " r, k, v, w, u 必须为fp16 fp32 bf16中的一种 并且类型相同"
+                    if r.dtype in [torch.float32, torch.bfloat16]:
+                        o_dtype = r.dtype
+                    else:
+                        o_dtype = torch.float32
+                    assert (
+                        r.device
+                        == k.device
+                        == v.device
+                        == w.device
+                        == u.device
+                        == s.device
+                        == s_map.device
+                    ), "what kan i say? 请确保r k v w u s s_map在同一设备上，快去检查！"
+
+                    y = torch.empty(
+                        (B, T, C),
+                        device=r.device,
+                        dtype=o_dtype,
+                        memory_format=torch.contiguous_format,
+                    )
+                    ys = torch.empty(
+                        (B, H, head_size, head_size),
+                        device=r.device,
+                        dtype=o_dtype,
+                        memory_format=torch.contiguous_format,
+                    )
+                    # print(ys)
+                    if r.dtype == torch.bfloat16:
+                        wkv6_cuda.forward_with_state_bf16(
+                            B, T, C, H, S, s_map, r, k, v, w, u, s, y, ys
+                        )
+                    elif r.dtype == torch.float32:
+                        wkv6_cuda.forward_with_state_fp32(
+                            B, T, C, H, S, s_map, r, k, v, w, u, s, y, ys
+                        )
+                    else:
+                        wkv6_cuda.forward_with_state_fp16(
+                            B, T, C, H, S, s_map, r, k, v, w, u, s, y, ys
+                        )
+
+                return y, ys
+
+        self.head_size = head_size
+        self.normal_kernenl = RWKV_6
+        self.kernel_with_state = RWKV_6_with_state
+
+    def __call__(
+        self, r, k, v, w, u, with_state=False, init_state=None, state_map=None
+    ):
+        B, T, C = r.shape
+        assert C % self.head_size == 0
+        H = C // self.head_size
+        if not isinstance(u, torch.Tensor):
+            u = u.value
+
+        assert r.is_cuda
+        assert k.is_cuda
+        assert v.is_cuda
+        assert w.is_cuda
+        assert u.is_cuda
+
+        if isinstance(r, torch.Tensor):
+            assert r.device == k.device == v.device == w.device == u.device
+        else:
+            r.get_device() == k.get_device() == v.get_device() == w.get_device() == u.get_device()
+
+        assert r.dtype == k.dtype == v.dtype == w.dtype == u.dtype
+
+        if r.dtype in [torch.float32, torch.bfloat16]:
+            s_dtype = r.dtype
+        else:
+            s_dtype = torch.float32
+
+        is_custom_init = init_state is not None
+
+        if init_state is not None:
+            assert len(init_state.shape) in [3, 4], (
+                "init_state 的形状必须为(state_kinds /*<= Batch_size*/,num_heads,head_size,head_size) 或者(num_heads,head_size,head_size)"
+            )
+            if len(init_state.shape) == 3:
+                init_state = init_state[None, :]
+            assert (
+                init_state.shape[1:] == (H, self.head_size, self.head_size)
+                and init_state.shape[0] <= B
+            ), (
+                "init_state 的形状必须为(state_kinds /*<= Batch_size*/,num_heads,head_size,head_size) 或者(num_heads,head_size,head_size)"
+            )
+
+            assert init_state.dtype == s_dtype, f"init_state的数值类型应为: {s_dtype}"
+            assert init_state.device == r.device
+
+        if state_map is not None:
+            if isinstance(state_map, list):
+                state_map = torch.tensor(state_map, dtype=torch.int64)
+            elif isinstance(state_map, torch.Tensor):
+                assert state_map.dtype in [
+                    torch.int32,
+                    torch.int64,
+                ], "state_map是一个长度为Batch_Size的int64类型的映射数组"
+                state_map = state_map.to(torch.int64)
+            assert state_map.shape == (B,), "state_map的shape必须为(Batch_Size,)"
+            assert state_map.device == r.deivec
+
+        if with_state:
+            if init_state is None:
+                assert state_map is None, (
+                    "您必须在指定了init_state的情况下才能使用state_map"
+                )
+                init_state = torch.zeros((0,), device=r.device, dtype=s_dtype)
+                state_map = torch.zeros((0,), device=r.device, dtype=torch.int64)
+            else:
+                n_state = init_state.shape[0]
+                if state_map is None:
+                    assert n_state == 1 or n_state == B, (
+                        "我无法为您推断state_map的形状，请手动指定。"
+                    )
+                    if n_state == 1:
+                        state_map = torch.tensor(
+                            [0] * B, dtype=torch.int64, device=r.device
+                        )
+                    elif n_state == B:
+                        state_map = torch.tensor(
+                            [i for i in range(B)], dtype=torch.int64, device=r.device
+                        )
+                    else:
+                        assert False, "未实现"
+                else:
+                    assert state_map.shape == (B,), "state_map的形状必须为(batch_size,)"
+                    assert (state_map >= 0).all() and (state_map < n_state).all(), (
+                        f"state_map的取值范围为[0,{n_state})之间的整数，您的输入显然不满足。"
+                    )
+            # print('state map:',state_map)
+            o, ys = self.kernel_with_state.apply(
+                B, T, C, H, is_custom_init, state_map, r, k, v, w, u, init_state
+            )
+            return o, ys
+        else:
+            o = self.normal_kernenl.apply(B, T, C, H, r, k, v, w, u)
+            return o, None

rwkv_ops/rwkv7_kernel/__init__.py

@@ -0,0 +1,113 @@
+import keras
+from distutils.util import strtobool
+import os
+from keras import ops
+
+
+def transpose_head(x, head_first):
+    if head_first:
+        return ops.transpose(x, (0, 2, 1, 3))
+    else:
+        return x
+
+
+def get_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
+    assert HEAD_SIZE % 4 == 0
+    from .native_keras_op import generalized_delta_rule as native_op
+
+    if keras.config.backend() == "torch":
+        import torch
+
+        if torch.cuda.is_available():
+            if KERNEL_TYPE.lower() == "triton":
+                from .torch_op import generalized_delta_rule
+
+                return generalized_delta_rule, generalized_delta_rule, True
+            elif KERNEL_TYPE.lower() == "cuda":
+                from .torch_cuda_kernel.wkv7_torch import (
+                    get_torch_generalized_delta_rule,
+                )
+
+                return get_torch_generalized_delta_rule(HEAD_SIZE) + [False]
+    elif keras.config.backend() == "jax":
+        import jax
+        import os
+
+        if jax.devices()[0].platform == "gpu":
+            if KERNEL_TYPE.lower() == "triton":
+                os.environ["JAX_LOG_COMPUTATION"] = "0"
+                from .jax_op import generalized_delta_rule
+
+                return generalized_delta_rule, native_op, False
+            elif KERNEL_TYPE.lower() == "cuda":
+                from .jax_cuda_kernel.wkv7_jax import get_jax_generalized_delta_rule
+
+                return get_jax_generalized_delta_rule(HEAD_SIZE) + [False]
+    elif keras.config.backend() == "tensorflow":
+        import tensorflow as tf
+
+        if len(tf.config.list_physical_devices("GPU")) > 0:
+            if KERNEL_TYPE.lower() == "cuda" and HEAD_SIZE:
+                try:
+                    from jax.lib import xla_bridge
+
+                    assert xla_bridge.get_backend().platform == "gpu"
+                except:
+                    raise (
+                        "The operation of the TensorFlow kernel depends on the JAX kernel."
+                        "Therefore, it is necessary to ensure that it can be used in JAX, so that TensorFlow can be used."
+                    )
+                print("🎉" * 10)
+                print("Tensorflow CUDA kernel onlt support Forward,not get graident")
+                print("🎉" * 10)
+                from .tf_eager_kernel import get_tf_generalized_delta_rule
+
+                generalized_delta_rule_inference = get_tf_generalized_delta_rule(
+                    HEAD_SIZE
+                )
+                return native_op, generalized_delta_rule_inference, False
+    elif keras.config.backend() == "mlx" and KERNEL_TYPE.lower() == "cuda":
+        from .mlx_op import generalized_delta_rule
+
+        return native_op, generalized_delta_rule, False
+    return native_op, native_op, False
+
+
+def get_rnn_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
+    assert HEAD_SIZE % 4 == 0
+    from .native_keras_op import generalized_delta_rule
+
+    if KERNEL_TYPE == "cuda":
+        if keras.config.backend() == "jax":
+            import jax
+
+            if jax.devices()[0].platform == "gpu":
+                from .jax_cuda_kernel_single.wkv7_single_step_jax import (
+                    get_jax_generalized_delta_rule_single_step,
+                )
+
+                return get_jax_generalized_delta_rule_single_step(HEAD_SIZE)
+        elif keras.config.backend() == "torch":
+            import torch
+
+            if torch.cuda.is_available():
+                from .torch_cuda_kernel_single.wkv7_single_step_torch import (
+                    get_torch_generalized_delta_rule_single_step,
+                )
+
+                return get_torch_generalized_delta_rule_single_step(HEAD_SIZE)
+        elif keras.config.backend() == "tensorflow":
+            import tensorflow as tf
+
+            if len(tf.config.list_physical_devices("GPU")) > 0:
+                try:
+                    import jax
+                except ImportError:
+                    return generalized_delta_rule
+                if jax.devices()[0].platform == "gpu":
+                    from .tf_eager_kernel import (
+                        get_tf_generalized_delta_rule_single_step,
+                    )
+
+                    return get_tf_generalized_delta_rule_single_step(HEAD_SIZE)
+    return generalized_delta_rule

rwkv_ops/rwkv7_kernel/get_jax_devices_info.py

@@ -0,0 +1,220 @@
+import os
+from functools import lru_cache
+from typing import Literal
+import functools
+import triton
+import jax
+import jax.numpy as jnp
+from enum import Enum
+import contextlib
+
+
+@lru_cache(maxsize=None)
+def get_multiprocessor_count(tensor_idx: int = 0) -> int:
+    return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)[
+        "multiprocessor_count"
+    ]
+
+
+@lru_cache(maxsize=None)
+def get_available_device() -> str:
+    try:
+        return triton.runtime.driver.active.get_current_target().backend
+    except BaseException:
+        import warnings
+
+        warnings.warn(
+            ("Triton is not supported on current platform, roll back to CPU."),
+            stacklevel=1,
+        )
+        return "cpu"
+
+
+@lru_cache(maxsize=None)
+def _check_platform() -> Literal["nvidia", "amd", "intel", "musa"]:
+    device = get_available_device()
+    if device == "cuda":
+        return "nvidia"
+    elif device == "hip":
+        return "amd"
+    elif device == "xpu":
+        return "intel"
+    else:
+        return device
+
+
+# For AMD GPUs, the triton backend is 'hip', while for Nvidia GPUs, the triton backend is 'cuda'.
+# However, the torch backend is 'cuda' for both Nvidia and AMD GPUs.
+# Therefore, we need to check the triton backend to determine the actual GPU vendor.
+device = get_available_device() if get_available_device() != "hip" else "cuda"
+
+device_platform = _check_platform()
+
+is_intel = device_platform == "intel"
+is_nvidia = device_platform == "nvidia"
+is_amd = device_platform == "amd"
+
+use_cuda_graph = is_nvidia and os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1"
+
+
+device = get_available_device() if get_available_device() != "hip" else "cuda"
+
+is_intel_a770 = False
+device = jax.devices()
+is_tf32_supported = is_nvidia
+
+
+def get_all_max_shared_memory():
+    return [
+        triton.runtime.driver.active.utils.get_device_properties(i)["max_shared_mem"]
+        for i in range(len(jax.devices()))
+    ]
+
+
+device_shared_mem_list = get_all_max_shared_memory()
+
+
+@lru_cache(maxsize=None)
+def is_triton_shared_mem_enough(
+    max_shared_mem: int = 102400, tensor_idx: int = 0
+) -> bool:
+    max_shared_memory = device_shared_mem_list[tensor_idx]
+    return max_shared_memory >= max_shared_mem
+
+
+device_capacity = is_triton_shared_mem_enough()
+
+
+def _cpu_device_warning():
+    import warnings
+
+    warnings.warn(
+        ("Triton is not supported on current platform, roll back to CPU."), stacklevel=1
+    )
+
+
+def get_all_max_shared_mem():
+    try:
+        return [
+            triton.runtime.driver.active.utils.get_device_properties(i)[
+                "max_shared_mem"
+            ]
+            for i in range(len(jax.devices()))
+        ]
+    except BaseException:
+        _cpu_device_warning()
+        return [-1]
+
+
+class Backend(Enum):
+    ADA = 101376  # RTX 4090
+    AMPERE = 166912  # A100
+    HOPPER = 232448  # H100
+    DEFAULT = 102400  # Default
+
+    @classmethod
+    def get_shared_memory(cls, arch: str) -> int:
+        try:
+            return cls[arch.upper()].value
+        except KeyError:
+            return cls.DEFAULT.value
+
+
+@lru_cache(maxsize=None)
+def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool:
+    try:
+        device_shared_mem_list = get_all_max_shared_mem()
+        max_shared_memory = device_shared_mem_list[tensor_idx]
+        return max_shared_memory >= Backend.get_shared_memory(arch)
+    except Exception:
+        return False
+
+
+def tensor_cache(fn):
+    """
+    A decorator that caches the most recent result of a function with tensor inputs.
+
+    This decorator will store the output of the decorated function for the most recent set of input tensors.
+    If the function is called again with the same input tensors, it will return the cached result.
+
+
+    Args:
+        fn (Callable[..., jax.Array]):
+            The function to be decorated. It should take tensor inputs and return tensor outputs.
+
+    Returns:
+        Callable[..., jax.Array]:
+            A wrapped version of the input function with single-entry caching.
+    """
+    last_args = None
+    last_kwargs = None
+    last_result = None
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        nonlocal last_args, last_kwargs, last_result
+
+        if last_args is not None and last_kwargs is not None:
+            if len(args) == len(last_args) and len(kwargs) == len(last_kwargs):
+                if all(a is b for a, b in zip(args, last_args)) and all(
+                    k in last_kwargs and v is last_kwargs[k] for k, v in kwargs.items()
+                ):
+                    return last_result
+
+        result = fn(*args, **kwargs)
+        last_args, last_kwargs, last_result = args, kwargs, result
+        return result
+
+    return wrapper
+
+
+@tensor_cache
+def prepare_lens(cu_seqlens):
+    return cu_seqlens[1:] - cu_seqlens[:-1]
+
+
+@tensor_cache
+def prepare_chunk_indices(cu_seqlens, chunk_size: int):
+    indices = jnp.concatenate(
+        [
+            jnp.arange(n)
+            for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()
+        ]
+    )
+
+    return jnp.stack([jnp.cumsum(jnp.equal(indices, 0), 0) - 1, indices], 1)
+
+
+def input_guard(fn):
+    """
+    A decorator to make sure all input tensors are contiguous and set the device based on input tensors.
+    """
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        contiguous_args = (i for i in args)
+        contiguous_kwargs = {k: v for k, v in kwargs.items()}
+
+        tensor = None
+        for arg in args:
+            if isinstance(arg, jax.Array):
+                tensor = arg
+                break
+        if tensor is None:
+            for value in kwargs.values():
+                if isinstance(value, jax.Array):
+                    tensor = value
+                    break
+
+        if tensor is not None:
+            ctx = tensor.device
+        else:
+            ctx = contextlib.nullcontext()
+
+        with ctx:
+            return fn(*contiguous_args, **contiguous_kwargs)
+
+    return wrapper
+
+
+is_intel_alchemist = False