rwkv-ops 0.1.0__py3-none-any.whl

rwkv_ops/__init__.py

@@ -0,0 +1,26 @@
+import os
+
+KERNEL_TYPE = os.environ.get("KERNEL_TYPE", "triton")
+KERAS_BACKEND = os.environ.get("KERAS_BACKEND")
+BACKEND = os.environ.get("KERNEL_BACKEND")
+
+
+if KERAS_BACKEND is not None:
+    BACKEND = KERAS_BACKEND
+elif BACKEND is not None:
+    os.environ["KERAS_BACKEND"] = BACKEND
+else:
+    import torch
+    import keras
+
+    BACKEND = "torch"
+    os.environ["KERAS_BACKEND"] = BACKEND
+    keras.config.set_backend("torch")
+assert KERNEL_TYPE in ["triton", "cuda", "native"]
+assert BACKEND in ["torch", "jax", "numpy", "tensorflow"]
+from .rwkv7_kernel import get_generalized_delta_rule
+
+generalized_delta_rule, RWKV7_USE_KERNEL = get_generalized_delta_rule(
+    KERNEL_TYPE=KERNEL_TYPE
+)
+rwkv7_op = generalized_delta_rule

rwkv_ops/rwkv7_kernel/__init__.py

@@ -0,0 +1,153 @@
+import keras
+from distutils.util import strtobool
+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"):
+    USE_KERNEL = False
+    if keras.config.backend() == "torch":
+        import torch
+
+        if KERNEL_TYPE.lower() == "triton":
+            from .torch_op import generalized_delta_rule
+
+            USE_KERNEL = True
+        elif KERNEL_TYPE.lower() == "cuda":
+            CHUNK_LEN = 16
+            USE_KERNEL = True
+            from torch.utils.cpp_extension import load
+            import os
+
+            flags = [
+                "-res-usage",
+                f"-D_C_={HEAD_SIZE}",
+                f"-D_CHUNK_LEN_={CHUNK_LEN}",
+                "--use_fast_math",
+                "-O3",
+                "-Xptxas -O3",
+                "--extra-device-vectorization",
+            ]
+            # 获取当前文件的绝对路径
+            current_file_path = os.path.abspath(__file__)
+
+            # 获取当前文件的目录路径
+            current_dir_path = os.path.dirname(current_file_path)
+
+            # 获取上一级目录的路径
+            parent_dir_path = os.path.abspath(
+                os.path.join(current_dir_path, os.path.pardir)
+            )
+            load(
+                name="wind_backstepping",
+                sources=[
+                    os.path.join(parent_dir_path, "cuda_kernel/wkv7_cuda.cu"),
+                    os.path.join(parent_dir_path, "cuda_kernel/wkv7_op.cpp"),
+                ],
+                is_python_module=False,
+                verbose=True,
+                extra_cuda_cflags=flags,
+            )
+
+            class WindBackstepping(torch.autograd.Function):
+                @staticmethod
+                def forward(ctx, w, q, k, v, z, b):
+                    B, T, H, C = w.shape
+                    DTYPE = q.dtype
+                    q = ops.cast(q, "bfloat16")
+                    k = ops.cast(k, "bfloat16")
+                    v = ops.cast(v, "bfloat16")
+                    z = ops.cast(z, "bfloat16")
+                    b = ops.cast(b, "bfloat16")
+                    w = ops.cast(w, "bfloat16")
+                    assert T % CHUNK_LEN == 0
+                    assert all(i.is_contiguous() for i in [w, q, k, v, z, b])
+                    y = torch.empty_like(v)
+                    s = torch.empty(
+                        B, H, T // CHUNK_LEN, C, C, dtype=torch.float32, device=w.device
+                    )
+                    sa = torch.empty(B, T, H, C, dtype=torch.float32, device=w.device)
+                    torch.ops.wind_backstepping.forward(w, q, k, v, z, b, y, s, sa)
+                    ctx.save_for_backward(w, q, k, v, z, b, s, sa)
+                    return ops.cast(y, DTYPE)
+
+                @staticmethod
+                def backward(ctx, dy):
+                    DTYPE = dy.dtype
+                    dy = ops.cast(dy, torch.bfloat16)
+                    dy = dy.contiguous()
+                    assert all(i.dtype == torch.bfloat16 for i in [dy])
+                    assert all(i.is_contiguous() for i in [dy])
+                    w, q, k, v, z, b, s, sa = ctx.saved_tensors
+                    dw, dq, dk, dv, dz, db = [
+                        torch.empty_like(x) for x in [w, q, k, v, z, b]
+                    ]
+                    torch.ops.wind_backstepping.backward(
+                        w, q, k, v, z, b, dy, s, sa, dw, dq, dk, dv, dz, db
+                    )
+                    return (
+                        ops.cast(dw, DTYPE),
+                        ops.cast(dq, DTYPE),
+                        ops.cast(dk, DTYPE),
+                        ops.cast(dv, DTYPE),
+                        ops.cast(dz, DTYPE),
+                        ops.cast(db, DTYPE),
+                    )
+
+            def RUN_CUDA_RWKV7g(q, w, k, v, a, b):
+                B, T, H, C = q.shape
+                q = q.contiguous()
+                w = w.contiguous()
+                k = k.contiguous()
+                v = v.contiguous()
+                a = a.contiguous()
+                b = b.contiguous()
+                return WindBackstepping.apply(w, q, k, v, a, b).view(B, T, H * C)
+
+            def generalized_delta_rule(
+                r: torch.Tensor,
+                w: torch.Tensor,
+                k: torch.Tensor,
+                v: torch.Tensor,
+                a: torch.Tensor,
+                b: torch.Tensor,
+                initial_state: torch.Tensor = None,
+                output_final_state: bool = True,
+                head_first: bool = False,
+                use_chunk: bool = True,
+            ):
+                r = transpose_head(r, head_first)
+                k = transpose_head(k, head_first)
+                v = transpose_head(v, head_first)
+                a = transpose_head(a, head_first)
+                b = transpose_head(b, head_first)
+                w = transpose_head(w, head_first)
+                return RUN_CUDA_RWKV7g(r, w, k, v, a, b), None
+        else:
+            from .native_keras_op import generalized_delta_rule
+
+            USE_KERNEL = False
+    elif keras.config.backend() == "jax":
+        from jax.lib import xla_bridge
+        import jax
+        import os
+
+        if (
+            xla_bridge.get_backend().platform == "gpu"
+            and KERNEL_TYPE.lower() == "triton"
+        ):
+            from .jax_op import generalized_delta_rule
+
+            USE_KERNEL = True
+        else:
+            from .native_keras_op import generalized_delta_rule
+
+    else:
+        from .native_keras_op import generalized_delta_rule
+    return generalized_delta_rule, USE_KERNEL

rwkv_ops/rwkv7_kernel/get_jax_devices_info.py

@@ -0,0 +1,221 @@
+import os
+from functools import lru_cache
+from typing import Literal
+import functools
+import triton
+import jax
+import jax.numpy as jnp
+
+
+@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()
+
+from enum import Enum
+import contextlib
+
+
+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

rwkv_ops/rwkv7_kernel/get_torch_devices_info.py

@@ -0,0 +1,250 @@
+import functools
+import os
+from functools import lru_cache
+from typing import Literal
+
+import triton
+from packaging import version
+import torch
+
+
+@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"
+
+
+@lru_cache(maxsize=None)
+def check_pytorch_version(version_s: str = "2.4") -> bool:
+    return version.parse(torch.__version__) >= version.parse(version_s)
+
+
+is_intel_a770 = is_intel and "Intel(R) Arc(TM) A" in torch.xpu.get_device_name(0)
+device = get_available_device() if get_available_device() != "hip" else "cuda"
+device_torch_lib = getattr(torch, device)
+if check_pytorch_version("2.4"):
+    device = "cuda" if device == "cpu" else device
+    autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device)
+    autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device)
+
+    def custom_device_ctx(index: int):
+        return device_torch_lib.device(index)
+else:
+    assert device == "cuda", (
+        "Only cuda device is supported for PyTorch version < 2.4.0."
+    )
+    autocast_custom_fwd = device_torch_lib.amp.custom_fwd
+    autocast_custom_bwd = device_torch_lib.amp.custom_bwd
+
+    def custom_device_ctx(index: int):
+        return torch.cuda.device(index)
+
+
+# Nvidia Ampere or newer, haven't check AMD and intel yet.
+is_tf32_supported = is_nvidia and torch.cuda.get_device_capability(0)[0] >= 8
+
+
+def get_all_max_shared_memory():
+    return [
+        triton.runtime.driver.active.utils.get_device_properties(i)["max_shared_mem"]
+        for i in range(device_torch_lib.device_count())
+    ]
+
+
+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()
+from enum import Enum
+import contextlib
+
+
+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(device_torch_lib.device_count())
+        ]
+    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[..., torch.Tensor]):
+            The function to be decorated. It should take tensor inputs and return tensor outputs.
+
+    Returns:
+        Callable[..., torch.Tensor]:
+            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: torch.LongTensor) -> torch.LongTensor:
+    return cu_seqlens[1:] - cu_seqlens[:-1]
+
+
+@tensor_cache
+def prepare_chunk_indices(
+    cu_seqlens: torch.LongTensor, chunk_size: int
+) -> torch.LongTensor:
+    indices = torch.cat(
+        [
+            torch.arange(n)
+            for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()
+        ]
+    )
+    return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens)
+
+
+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 if not isinstance(i, torch.Tensor) else i.contiguous() for i in args
+        )
+        contiguous_kwargs = {
+            k: (v if not isinstance(v, torch.Tensor) else v.contiguous())
+            for k, v in kwargs.items()
+        }
+
+        tensor = None
+        for arg in args:
+            if isinstance(arg, torch.Tensor):
+                tensor = arg
+                break
+        if tensor is None:
+            for value in kwargs.values():
+                if isinstance(value, torch.Tensor):
+                    tensor = value
+                    break
+
+        if tensor is not None:
+            ctx = custom_device_ctx(tensor.device.index)
+        else:
+            ctx = contextlib.nullcontext()
+
+        with ctx:
+            return fn(*contiguous_args, **contiguous_kwargs)
+
+    return wrapper
+
+
+is_intel_alchemist = is_intel and "Intel(R) Arc(TM) A" in torch.xpu.get_device_name(0)

rwkv_ops/rwkv7_kernel/jax_kernel/__init__.py

@@ -0,0 +1,9 @@
+from ..jax_kernel.chunk_A_fwd import *
+from ..jax_kernel.chunk_A_bwd import *
+from ..jax_kernel.chunk_h_fwd import *
+from ..jax_kernel.chunk_h_bwd import *
+from ..jax_kernel.chunk_o_fwd import *
+from ..jax_kernel.chunk_o_bwd import *
+from ..jax_kernel.cumsum import *
+from ..jax_kernel.wy_fast_fwd import *
+from ..jax_kernel.wy_fast_bwd import *