在GPU上加速RWKV6模型的Linear Attention计算

def rwkv6_linear_attention(
    training,
    receptance,
    key,
    value,
    time_decay,
    time_first,
    state,
):
    no_cuda = any(t.device.type != "cuda" for t in [time_decay, time_first, receptance, key, value])
    # Launching the CUDA kernel for just one token will actually be slower (there is no for loop in the CPU version
    # in this case).
    one_token = key.size(1) == 1
    if no_cuda or one_token:
        return rwkv6_linear_attention_cpu(
            receptance, key, value, time_decay, time_first, state
        )
    else:
        ...

def hf_rwkv6_linear_attention_cpu(receptance, key, value, time_decay, time_first, state):
    # For CPU fallback. Will be slower and probably take more memory than the custom CUDA kernel if not executed
    # within a torch.no_grad.
    batch, seq_length, _ = receptance.shape
    num_heads, head_size = time_first.shape
    key = key.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2).transpose(-2, -1) # b, t, h, n -> b, h, t, n -> b, h, n, t
    value = value.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # b, t, h, n -> b, h, t, n
    receptance = receptance.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # b, t, h, n -> b, h, t, n
    time_decay = torch.exp(-torch.exp(time_decay.float())).view(batch, seq_length, num_heads, head_size).permute(0, 2, 3, 1) # b, t, h, n -> b, h, n, t
    time_first = time_first.float().reshape(-1, 1, 1).reshape(num_heads, -1, 1) # h, n -> h * n, 1, 1 -> h, n, 1
    out = torch.zeros_like(key).reshape(batch, seq_length, num_heads, head_size)
for current_index in range(seq_length):
    current_receptance = receptance[:, :, current_index:current_index+1, :]
    current_key = key[:, :, :, current_index:current_index+1]
    current_value = value[:, :, current_index:current_index+1, :]
    current_time_decay = time_decay[:, :, :, current_index:current_index+1]
    attention_output = current_key @ current_value
    out[:, current_index] = (current_receptance @ (time_first * attention_output + state)).squeeze(2)
    with torch.no_grad():
        # attention_output.shape: [b, h, n, 1] x [b, h, 1, n] -> [b, h, n, n]
        # current_time_decay * state: [b, h, n, 1] * [b, h, n, n] ->[b, h, n, n]
        # state.shape: [b, h, n, n]
        state = attention_output + current_time_decay * state 

return out, state

import sys

import torch

from fla.ops.rwkv6.chunk import chunk_rwkv6

from fla.ops.rwkv6.recurrent_fuse import fused_recurrent_rwkv6

from flash_rwkv import rwkv6_cuda_linear_attention

def hf_rwkv6_linear_attention_cpu(receptance, key, value, time_decay, time_first, state):

# For CPU fallback. Will be slower and probably take more memory than the custom CUDA kernel if not executed

# within a torch.no_grad.

batch, seq_length, _ = receptance.shape

num_heads, head_size = time_first.shape

key = key.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2).transpose(-2, -1)

value = value.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2)

receptance = receptance.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2)

time_decay = torch.exp(-torch.exp(time_decay.float())).view(batch, seq_length, num_heads, head_size).permute(0, 2, 3, 1)

time_first = time_first.float().reshape(-1, 1, 1).reshape(num_heads, -1, 1)

out = torch.zeros_like(key).reshape(batch, seq_length, num_heads, head_size)
for current_index in range(seq_length):
    current_receptance = receptance[:, :, current_index:current_index+1, :]
    current_key = key[:, :, :, current_index:current_index+1]
    current_value = value[:, :, current_index:current_index+1, :]
    current_time_decay = time_decay[:, :, :, current_index:current_index+1]
    attention_output = current_key @ current_value
    out[:, current_index] = (current_receptance @ (time_first * attention_output + state)).squeeze(2)
    with torch.no_grad():
        state = attention_output + current_time_decay * state

return out, state

if name == "main":

mode = sys.argv[1]

B = 1

H = 32

L = 54

D = 64

HIDDEN_SIZE = H * D

dtype = torch.float32
if mode == 'hf':
    profile_path = '/bbuf/rwkv_profile_result/hf/'
elif mode == 'recurrent':
    profile_path = '/bbuf/rwkv_profile_result/recurrent/'
elif mode == 'chunk':
    profile_path = '/bbuf/rwkv_profile_result/chunk/'
elif mode == 'cuda':
    profile_path = '/bbuf/rwkv_profile_result/cuda'
else:
    raise NotImplementedError
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(
        wait=1,
        warmup=1,
        active=1),
    on_trace_ready=torch.profiler.tensorboard_trace_handler(profile_path, worker_name='worker0'),
    record_shapes=True,
    profile_memory=True,  # This will take 1 to 2 minutes. Setting it to False could greatly speedup.
    with_stack=True
) as p:
    for i in range(10):
        q = (torch.randn(B, L, HIDDEN_SIZE).cuda().to(torch.float16)).requires_grad_(True)
        k = (torch.randn(B, L, HIDDEN_SIZE).cuda().to(torch.float16)).requires_grad_(True)
        v = torch.randn(B, L, HIDDEN_SIZE).cuda().to(torch.float16).requires_grad_(True)
        w = torch.nn.functional.logsigmoid(torch.randn(B, L, HIDDEN_SIZE)).cuda().to(torch.float32).requires_grad_(True)
        u = (torch.randn(H, D).cuda().to(torch.float16)).requires_grad_(True)
        state = (torch.randn(B, H, D, D).cuda().to(torch.float32)).requires_grad_(True)
        if mode == 'hf':
            o1, state1 = hf_rwkv6_linear_attention_cpu(q, k, v, w, u, state)
        elif mode =='cuda':
            o2, state2 = rwkv6_cuda_linear_attention(q, k, v, w, u.flatten(), state)
        else:
            batch, seq_length, _ = q.shape
            num_heads, head_size = u.shape
            k = k.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # B, T, H, K -> B, H, T, K
            v = v.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # B, T, H, K - > B, H, T, V
            q = q.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # B, H, T, K
            w = -torch.exp(w.float()).view(batch, seq_length, num_heads, head_size).permute(0, 2, 1, 3) # B, T, H, K -> B, H, T, K
            u = u.float().reshape(num_heads, head_size) # H, K

            if mode == 'recurrent':
                o3, state3 = fused_recurrent_rwkv6(q, k, v, w, u, initial_state=state, scale=1.0, output_final_state=True)
            elif mode == 'chunk':
                o4, state4 = chunk_rwkv6(q, k, v, w, u, initial_state=state, scale=1.0, output_final_state=True)
        p.step()

def rwkv5_linear_attention_cpu(receptance, key, value, time_decay, time_first, state):
    input_dtype = receptance.dtype
    # For CPU fallback. Will be slower and probably take more memory than the custom CUDA kernel if not executed
    # within a torch.no_grad.
    batch, seq_length, hidden_size = receptance.shape
    num_heads, head_size = time_first.shape
    key = key.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2).transpose(-2, -1)
    value = value.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2)
    receptance = receptance.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2)
    time_decay = torch.exp(-torch.exp(time_decay.float())).reshape(-1, 1, 1).reshape(num_heads, -1, 1)
    time_first = time_first.float().reshape(-1, 1, 1).reshape(num_heads, -1, 1)
    out = torch.zeros_like(key).reshape(batch, seq_length, num_heads, head_size)

    for current_index in range(seq_length):
        current_receptance = receptance[:, :, current_index:current_index+1, :]
        current_key = key[:, :, :, current_index:current_index+1]
        current_value = value[:, :, current_index:current_index+1, :]
        attention_output = current_key @ current_value
        out[:, current_index] = (current_receptance @ (time_first * attention_output + state)).squeeze(2)
        with torch.no_grad():
            state = attention_output + time_decay * state

    return out, state

def RUN_FORMULA_1A(B, T, C, H, r, k, v, w, u):
N = C // H
r = r.view(B, T, H, N)
k = k.view(B, T, H, N)
v = v.view(B, T, H, N)
w = w.view(H, N)
u = u.view(H, N)
out = torch.zeros((B, T, H, N), device=DEVICE)

for b in range(B):
    for h in range(H):
        state = torch.zeros((N,N), device=DEVICE).contiguous()
        for t in range(T):
            for i in range(N):
                for j in range(N):
                    x = k[b,t,h,j] * v[b,t,h,i]
                    s = state[i,j]
                    out[b,t,h,i] += r[b,t,h,j] * (u[h,j] * x + s)
                    state[i,j] = s * w[h,j] + x

return out.view(B, T, C)

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

{

dim3 threadsPerBlock( min(BC, 32) );

assert(B * C % threadsPerBlock.x == 0);

dim3 numBlocks(B * C / threadsPerBlock.x);

kernel_forward<<<numBlocks, threadsPerBlock>>>(B, T, C, H, r, k, v, w, u, y);

}

template <typename F>

global void kernel_forward(const int B, const int T, const int C, const int H,

const F *restrict const _r, const F *restrict const _k, const F *restrict const _v, const F *restrict const _w, const F *restrict const _u,

F *restrict const _y)

{

const int idx = blockIdx.x * blockDim.x + threadIdx.x;

const int _b = idx / C;

const int _h = (idx / N) % H;

const int _i = idx % N;

const int _o0 = _b*T*C + _h*N;
const int _o1 = _h*N;
const F *__restrict__ const k = _k + _o0;
const F *__restrict__ const v = _v + _o0 + _i;
const F *__restrict__ const r = _r + _o0;
F *__restrict__ const y = _y + _o0 + _i;

float state[N] = {0};   

for (int __t = 0; __t &lt; T; __t++)
{
    const int _t = __t*C;
    const F vv = v[_t];

    for (int _j = 0; _j &lt; N; _j++) 
    {
        const int j = _t + _j;
        const int m = _o1 + _j;

        const float x = k[j] * vv;
        const float s = state[_j];
        
        atomicAdd(y + _t, r[j] * (_u[m] * x + s));
        state[_j] = s * _w[m] + x;
    }
}

}

template <typename F>

__global__ void kernel_forward(const int B, const int T, const int C, const int H,

const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const F *__restrict__ _u,

F *__restrict__ const _y)

{

const int b = blockIdx.x / H;

const int h = blockIdx.x % H;

const int i = threadIdx.x;

_w += h*_N_;

_u += h*_N_;

__shared__ float r[_N_], k[_N_], u[_N_], w[_N_];
float state[_N_] = {0};

__syncthreads();
u[i] = float(_u[i]);
w[i] = float(_w[i]);
__syncthreads();

for (int t = b*T*C + h*_N_ + i; t &lt; (b+1)*T*C + h*_N_ + i; t += C)
{
    __syncthreads();
    r[i] = float(_r[t]);
    k[i] = float(_k[t]);
    __syncthreads();

    const float v = float(_v[t]);
    float y = 0;

    #pragma unroll
    for (int j = 0; j &lt; _N_; j+=4)
    {
        const float4&amp; r_ = (float4&amp;)(r[j]);
        const float4&amp; k_ = (float4&amp;)(k[j]);
        const float4&amp; w_ = (float4&amp;)(w[j]);
        const float4&amp; u_ = (float4&amp;)(u[j]);
        float4&amp; s = (float4&amp;)(state[j]);
        float4 x;

        x.x = k_.x * v;
        x.y = k_.y * v;
        x.z = k_.z * v;
        x.w = k_.w * v;

        y += r_.x * (u_.x * x.x + s.x);
        y += r_.y * (u_.y * x.y + s.y);
        y += r_.z * (u_.z * x.z + s.z);
        y += r_.w * (u_.w * x.w + s.w);

        s.x = s.x * w_.x + x.x;
        s.y = s.y * w_.y + x.y;
        s.z = s.z * w_.z + x.z;
        s.w = s.w * w_.w + x.w;
    }
    _y[t] = F(y);
}

}

def naive_recurrent_rwkv6(

q,

k,

v,

w,

u,

initial_state=None,

output_final_state=False

):

orig_dtype = q.dtype

q, k, v, w, u = map(lambda x: x.float(), (q, k, v, w, u))

batch_size, n_heads, seq_len, d_head_k = q.shape

_, _, _, d_head_v = v.shape

h = torch.zeros(batch_size, n_heads, d_head_k, d_head_v, dtype=torch.float32, device=q.device)

o = torch.zeros_like(v)

if initial_state is not None:
    h += initial_state

for i in range(seq_len):
    q_i = q[:, :, i, :]
    k_i = k[:, :, i]
    v_i = v[:, :, i, :]
    w_i = w[:, :, i].exp()
    kv_i = k_i[..., None] * v_i[..., None, :]
    o_i = (h + u[None, ..., None] * kv_i) * q_i[..., None]
    o[:, :, i] = o_i.sum(-2)
    h = h * w_i[..., None] + kv_i
return o.to(orig_dtype)

B = 4

H = 4

L = 1024

D = 100

dtype = torch.float32

q = (torch.randn(B, H, L, D).cuda().to(dtype)).requires_grad_(True)

k = (torch.randn(B, H, L, D).cuda().to(dtype)).requires_grad_(True)

v = torch.randn(B, H, L, D).cuda().to(dtype).requires_grad_(True)

w = torch.nn.functional.logsigmoid(torch.randn(B, H, L, D)).cuda().to(torch.float32).requires_grad_(True)

u = (torch.randn(H, D).cuda().to(dtype)).requires_grad_(True)

do = torch.rand_like(v).cuda()

o = naive_recurrent_rwkv6(q, k, v, w, u)

# if scale is None, use d_head_qk ** -0.5 by default. Otherwise specify the scale yourself. e.g. scale = 1.0

def fused_recurrent_rwkv6(

r: torch.Tensor,

k: torch.Tensor,

v: torch.Tensor,

w: torch.Tensor,

u: torch.Tensor,

scale: int = -1,

initial_state: torch.Tensor = None,

output_final_state: bool = False,

causal: bool = True

) -> Tuple[torch.Tensor, torch.Tensor]:

r"""

Args:

r (torch.Tensor):

reception of shape (B, H, T, K). Alias: q, query in linear attention.

k (torch.Tensor):

keys of shape (B, H, T, K)

v (torch.Tensor):

values of shape (B, H, T, V)

w (torch.Tensor):

data-dependent decays of shape (B, H, T, K) in log space! Alias: g.

u (torch.Tensor):

bonus of shape (H, K)

scale (Optional[int]):

Scale factor for the RWKV6 attention scores.

If not provided, it will default to 1 / sqrt(K). Default: None.

initial_state (Optional[torch.Tensor]):

Initial state of shape (B, H, K, V). Default: None.

output_final_state (Optional[bool]):

Whether to output the final state of shape (B, H, K, V). Default: False.

"""

if scale == -1:

scale = r.shape[-1] ** -0.5

if initial_state is not None:

initial_state = initial_state.detach()

o, final_state = FusedRecurrentRWKV6Function.apply(r, k, v, w, u, scale, initial_state, output_final_state)

return o, final_state

class FusedRecurrentRWKV6Function(torch.autograd.Function):

@staticmethod
@contiguous
@custom_fwd
def forward(ctx, r, k, v, w, u, scale=None, initial_state=None, output_final_state=False, reverse=False):
    # alias
    q = r
    batch_size, n_heads, seq_len, d_head_qk = q.shape
    d_head_v = v.shape[-1]
    # default scale
    if scale is None:
        scale = d_head_qk ** -0.5

    BK, BV = min(triton.next_power_of_2(d_head_qk), 32), min(triton.next_power_of_2(d_head_v), 32)
    NK, NV = triton.cdiv(d_head_qk, BK), triton.cdiv(d_head_v, BV)
    num_stages = 1
    num_warps = 1

    o = q.new_empty(NK, batch_size, n_heads, seq_len,
                    d_head_v, dtype=torch.float32)

    if output_final_state:
        final_state = q.new_empty(batch_size, n_heads, d_head_qk, d_head_v)
    else:
        final_state = None

    grid = (NV, NK, batch_size * n_heads)
    fused_recurrent_rwkv6_fwd_kernel[grid](
        q, k, v, w, u, o, initial_state, final_state,
        q.stride(1), q.stride(2), q.stride(3),
        v.stride(1), v.stride(2), v.stride(3),
        batch_size, n_heads, seq_len, scale,
        DK=d_head_qk, DV=d_head_v, BK=BK, BV=BV,
        USE_INITIAL_STATE=initial_state is not None,
        STORE_FINAL_STATE=final_state is not None,
        REVERSE=reverse,
        num_warps=num_warps,
        num_stages=num_stages
    )

    o = o.sum(0)
    ctx.save_for_backward(q, k, v, w, u, initial_state, o)
    ctx.scale = scale
    ctx.reverse = reverse
    # we do not need the gradient of the final state from the next chunk
    # similiar to Trunctated BPTT
    if final_state is not None:
        final_state = final_state.detach()
    return o.to(q.dtype), final_state