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

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:

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:

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:

__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. This file exposes forward and backward functions using PyBind11 and TORCH_LIBRARY:

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, the text stream extension loads as:

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 and Image_rwkv.py pass raw tensors directly to the CUDA kernels. For example, the forward call at line 55 of Text_rwkv.py:

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:

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:

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 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}