# How Custom CUDA Kernels (wkv6_cuda.cu) Implement the WKV Operation in RWKV-CLIP

> Discover how the wkv6_cuda.cu file implements the WKV operation using custom CUDA kernels. Learn about forward and backward passes, shared memory, and float4 vectorization for linear-time computation.

- Repository: [DeepGlint/rwkv-clip](https://github.com/deepglint/rwkv-clip)
- Tags: internals
- Published: 2026-02-28

---

**The wkv6_cuda.cu file implements the Weighted-Key-Value (WKV) recurrence through three specialized CUDA kernels—one for the forward pass and three for the backward pass—that use shared memory, register arrays, and float4 vectorization to achieve linear-time computation without attention matrices.**

The RWKV-CLIP model from deepglint/rwkv-clip replaces traditional self-attention with a custom WKV (Weighted-Key-Value) recurrence implemented in highly optimized CUDA code. This implementation, found in `model/cuda_text/wkv6_cuda.cu` and `model/cuda_image/wkv6_cuda.cu`, processes sequences in linear time by maintaining running state vectors rather than computing quadratic attention matrices. Understanding these kernels reveals how the model achieves efficient bidirectional encoding for both text and image modalities.

## Kernel Architecture and Launch Configuration

The CUDA kernels operate with a specific launch signature designed for parallel processing across batches and attention heads. The `cuda_forward` function in `wkv6_cuda.cu` accepts the following parameters:

```cpp
void cuda_forward(int B, int T, int C, int H,
                 float *r, float *k, float *v,
                 float *w, float *u, float *y);

```

- **B**: batch size
- **T**: sequence length (time steps)
- **C**: hidden dimension (must equal `H * _N_`)
- **H**: number of heads
- **r, k, v**: recurrent tensors for the three streams
- **w**: learnable decay factors (exponentiated on-the-fly)
- **u**: per-head bias used only in the forward (bidirectional) direction
- **y**: output of the WKV operation

The kernels launch with one block per **(batch, head)** pair and `_N_` threads per block, where `_N_` equals `C / H`. This configuration appears in `model/cuda_text/wkv6_cuda.cu` at lines 48-53:

```cpp
kernel_forward<<<dim3(B * H), dim3(_N_)>>>(...);

```

All kernels copy the per-thread values of `r`, `k`, `u`, and `w` into **shared memory** arrays of length `_N_`, enabling fast cooperative access to the full hidden vector of each head:

```cpp
__shared__ float r[_N_], k[_N_], u[_N_], w[_N_];

```

Meanwhile, the state vectors that accumulate the recurrence are stored in **register arrays** (`float state1[_N_]`) for minimal latency during the many tiny updates per time step.

## Forward Pass Implementation in kernel_forward

The forward kernel processes the sequence bidirectionally through two distinct passes. This design implements true bidirectional encoding without the memory overhead of standard attention mechanisms.

### Future-to-Past Pass

The first pass iterates backwards through time (from future to past) to accumulate initial state. For each time step `t`, the kernel:

1. Loads `w`, `r`, and `k` into shared memory
2. Computes `x = k * v` via element-wise multiplication
3. Updates the running state `s` with exponential decay: `s = s * w + x`
4. Accumulates contributions to the output `y` by weighting the state with `r`

This pass appears in `wkv6_cuda.cu` at lines 23-56.

### Past-to-Future Pass

The second pass iterates forwards through time (from past to future), mirroring the first pass but adding the **bias term `u`** to the computation. The output update follows `y += r * (u * x + s)`, where the `u` term provides the learned per-head bias critical for bidirectional modeling. This logic spans lines 58-95 in the source file.

Both passes employ **float4 vectorization** (`#pragma unroll … float4&`) to process four hidden units per iteration, requiring that `_N_ % 4 == 0` (asserted at lines 50-51).

## Backward Pass Gradient Computation

The backward pass requires three separate kernels to compute gradients efficiently without atomic operations. Each targets specific parameter groups while maintaining the same launch configuration (`dim3(B*H), dim3(_N_)`):

**kernel_backward_111** computes gradients for `r`, `k`, `v`, and `u` (lines 100-166). This kernel traverses the sequence forward while maintaining two sets of register states (`state1`, `state2`) and two cumulative buffers (`scccc*`, `sdddd*`) to propagate gradients through the recurrence efficiently.

**kernel_backward_222** handles the first half of the `w` gradient (lines 188-280). Working backwards through the sequence, it accumulates contributions from future steps using buffers `sa` and `sbbbb` to track the exponential decay derivatives.

**kernel_backward_333** completes the `w` gradient computation (lines 282-380). This kernel works forwards, mirroring the logic of `kernel_backward_222` to capture the bidirectional dependencies of the decay factors.

All three kernels use `__syncthreads()` to synchronize shared memory updates across the `_N_` threads within each block.

## PyTorch Integration via PyBind11

The raw CUDA kernels interface with PyTorch through a thin C++ wrapper in [`model/cuda_text/wkv6_op.cpp`](https://github.com/deepglint/rwkv-clip/blob/main/model/cuda_text/wkv6_op.cpp). This file exposes `forward` and `backward` functions using **PyBind11** and **TORCH_LIBRARY**:

```cpp
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "wkv6_text_birectional forward");
    m.def("backward", &backward, "wkv6_text_birectional backward");
}

```

The Python extensions compile dynamically via `torch.utils.cpp_extension.load`. In [`model/Text_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Text_rwkv.py), the text stream extension loads as:

```python
wkv6_text_cuda = load(
    name="wkv6_text_birectional",
    sources=[
        "model/cuda_text/wkv6_op.cpp",
        "model/cuda_text/wkv6_cuda.cu"
    ],
)

```

High-level calls in [`Text_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/Text_rwkv.py) and [`Image_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/Image_rwkv.py) pass raw tensors directly to the CUDA kernels. For example, the forward call at line 55 of [`Text_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/Text_rwkv.py):

```python
wkv6_text_cuda.forward(B, T, C, H,
                       r.float(), k.float(), v.float(),
                       ew, u.float(), y)

```

And the backward call at line 68 of [`Image_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/Image_rwkv.py):

```python
wkv6_cuda.backward(B, T, C, H,
                   r.float(), k.float(), v.float(),
                   ew, u, gy, gr, gk, gv, gw, gu)

```

## Code Example: Loading and Using the WKV CUDA Extension

To utilize these kernels in a custom model, load the extension and wrap it in an autograd function:

```python
import torch
from torch.utils.cpp_extension import load

# Load the compiled CUDA extension (once per process)

wkv6_text = load(
    name="wkv6_text_birectional",
    sources=[
        "model/cuda_text/wkv6_op.cpp",
        "model/cuda_text/wkv6_cuda.cu",
    ],
)

def wkv_forward(r, k, v, w, u):
    B, T, C = r.shape
    H = C // 64               # _N_ is 64 in the original config

    y = torch.empty_like(r)
    ew = torch.exp(w)         # exponential of w (pre-computed)

    wkv6_text.forward(B, T, C, H,
                      r.float(), k.float(), v.float(),
                      ew, u.float(), y)
    return y

class WKVFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, r, k, v, w, u):
        ctx.save_for_backward(r, k, v, w, u)
        y = wkv_forward(r, k, v, w, u)
        return y

    @staticmethod
    def backward(ctx, gy):
        r, k, v, w, u = ctx.saved_tensors
        B, T, C = r.shape
        H = C // 64
        ew = torch.exp(w)
        
        gr = torch.empty_like(r)
        gk = torch.empty_like(k)
        gv = torch.empty_like(v)
        gu = torch.empty_like(u)
        gw = torch.empty_like(w)
        
        wkv6_text.backward(B, T, C, H,
                           r.float(), k.float(), v.float(),
                           ew, u,
                           gy.float(),
                           gr, gk, gv, gw, gu)
        return gr, gk, gv, gw, gu

```

## Summary

- **Linear-time recurrence**: The WKV operation avoids quadratic attention matrices by maintaining running state vectors updated with exponential decay.
- **Bidirectional processing**: The forward kernel implements true bidirectional encoding through separate future-to-past and past-to-future passes, with the `u` bias term applied only in the forward sweep.
- **Memory hierarchy optimization**: Parameters reside in shared memory for cooperative access, while recurrence states live in registers for low-latency updates.
- **Vectorized execution**: `float4` loads process four elements per instruction, requiring head dimensions divisible by 4.
- **Decomposed gradients**: Three specialized backward kernels compute derivatives for `r/k/v/u`, first-half `w`, and second-half `w` without costly atomic operations.
- **Seamless PyTorch integration**: PyBind11 wrappers in [`wkv6_op.cpp`](https://github.com/deepglint/rwkv-clip/blob/main/wkv6_op.cpp) expose the kernels as Python-callable functions, loaded via `torch.utils.cpp_extension`.

## Frequently Asked Questions

### Why does the forward kernel process the sequence twice?

The future-to-past pass accumulates state information from future tokens without bias, while the past-to-future pass adds the learned per-head bias term `u` and completes the bidirectional context. This two-pass approach enables the WKV operation to capture dependencies in both directions while maintaining linear complexity.

### What is the significance of `_N_` in the kernel configuration?

`_N_` represents the head dimension (calculated as `C / H`), determining the number of threads per block (`dim3(_N_)`). The implementation requires `_N_ % 4 == 0` to support `float4` vectorization, which processes four hidden units per memory instruction for optimal GPU bandwidth utilization.

### How are gradients for the decay factor `w` computed?

The gradients for `w` require bidirectional information through the recurrence. `kernel_backward_222` processes the sequence backwards to accumulate future contributions using buffers `sa` and `sbbbb`, while `kernel_backward_333` processes forwards to capture past dependencies. This split avoids atomic operations and enables efficient parallel gradient computation.

### Can these kernels handle variable sequence lengths?

Yes, the sequence length `T` is passed as a runtime parameter to both forward and backward kernels, allowing dynamic handling of varying input lengths. However, the hidden dimension `C` and head count `H` must satisfy `C = H * _N_` with `_N_` divisible by 4, as these dimensions determine the static shared memory allocation and thread block structure.