RWKV-7 "Goose" 架构分析 — 超越 Transformer 的线性时间模型

# Standard Attention: O(n^2) 时间, O(n) 内存
# output = softmax(Q @ K^T / sqrt(d)) @ V

# Linear Attention: O(n) 时间, O(d^2) 内存（常数）
# output = (Q @ (K^T @ V)) -- 改变矩阵乘法顺序

# RWKV-7: O(n) 时间, O(d^2) 内存（常数）
# + Dynamic State Evolution 将表达能力最大化

# TC0 限制示例：状态追踪（state tracking）问题
# 输入："A 在房间1。A 移动到房间2。B 在房间1。A 在哪里？"
# TC0 模型在理论上无法追踪这类状态变化

import torch
import torch.nn as nn

class RWKV7_TimeMix(nn.Module):
    """RWKV-7 核心：Dynamic State Evolution"""

    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        # 输入依赖的参数生成器
        self.W_r = nn.Linear(d_model, d_model)  # Receptance
        self.W_k = nn.Linear(d_model, d_model)  # Key
        self.W_v = nn.Linear(d_model, d_model)  # Value
        self.W_a = nn.Linear(d_model, d_model)  # State transition
        self.W_g = nn.Linear(d_model, d_model)  # Gate

        # Dynamic decay 参数
        self.w_decay = nn.Parameter(torch.randn(n_heads, self.d_head))

        # State evolution matrix 可学习参数
        self.A_log = nn.Parameter(torch.randn(n_heads, self.d_head, self.d_head))

    def forward(self, x, state=None):
        B, T, C = x.shape
        H = self.n_heads
        D = self.d_head

        r = self.W_r(x).view(B, T, H, D)
        k = self.W_k(x).view(B, T, H, D)
        v = self.W_v(x).view(B, T, H, D)
        a = torch.sigmoid(self.W_a(x).view(B, T, H, D))
        g = torch.sigmoid(self.W_g(x).view(B, T, H, D))

        if state is None:
            state = torch.zeros(B, H, D, D, device=x.device)

        outputs = []
        for t in range(T):
            # Dynamic State Evolution：核心创新
            # 以往：state = decay * state + k^T @ v（固定 decay）
            # RWKV-7：state = A(x_t) @ state + k^T @ v（动态转移矩阵）

            # 输入依赖的转移矩阵
            A_t = self._compute_transition(a[:, t], state)

            # State update
            kv = torch.einsum('bhd,bhe->bhde', k[:, t], v[:, t])
            state = A_t * state + kv

            # Output 计算
            out = torch.einsum('bhd,bhde->bhe', r[:, t], state)
            out = out * g[:, t]
            outputs.append(out)

        output = torch.stack(outputs, dim=1).view(B, T, C)
        return output, state

    def _compute_transition(self, a_t, state):
        """根据输入动态生成转移矩阵"""
        # a_t: [B, H, D] - 从当前输入导出的转移参数
        # 这正是 RWKV-7 突破 TC0 的关键所在
        decay = torch.exp(-torch.exp(self.w_decay))
        A = decay.unsqueeze(-1) * torch.eye(
            self.d_head, device=a_t.device
        ).unsqueeze(0)

        # 输入依赖的修正项
        A = A + torch.einsum('bhd,bhe->bhde', a_t, a_t) * \
            torch.exp(self.A_log).unsqueeze(0)

        return A

# RWKV-6（以往）：state_{t+1} = diag(w) * state_t + k_t * v_t^T
# -> 对角矩阵乘法：每个维度独立地进行 decay
# -> 处于 TC0 范围内：有限深度无法完成状态追踪

# RWKV-7（新）：state_{t+1} = A(x_t) * state_t + k_t * v_t^T
# -> A(x_t) 是输入依赖的转移矩阵（非对角）
# -> 维度之间可以相互作用：能够解决状态追踪问题
# -> 表达能力突破了 TC0

class RWKV7Block(nn.Module):
    """RWKV-7 基本 block"""

    def __init__(self, d_model, n_heads):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.time_mix = RWKV7_TimeMix(d_model, n_heads)
        self.channel_mix = RWKV7_ChannelMix(d_model)

    def forward(self, x, state=None):
        # Time Mixing（token 间的相互作用）
        h, state = self.time_mix(self.ln1(x), state)
        x = x + h

        # Channel Mixing（承担 FFN 的角色）
        x = x + self.channel_mix(self.ln2(x))

        return x, state


class RWKV7_ChannelMix(nn.Module):
    """channel mixing（SwiGLU 变体）"""

    def __init__(self, d_model, expand=3.5):
        super().__init__()
        hidden = int(d_model * expand)
        self.W_key = nn.Linear(d_model, hidden)
        self.W_value = nn.Linear(hidden, d_model)
        self.W_gate = nn.Linear(d_model, hidden)

    def forward(self, x):
        k = self.W_key(x)
        v = self.W_value(torch.relu(k) ** 2)  # squared ReLU
        g = torch.sigmoid(self.W_gate(x))
        return v * g

class RWKV7_Inference:
    """推理时以 RNN 模式运行：每个 token 需要 O(d^2) 计算量"""

    def __init__(self, model):
        self.model = model
        self.states = None  # [n_layers, n_heads, d_head, d_head]

    def generate_token(self, token_id):
        x = self.model.embed(token_id)

        for i, block in enumerate(self.model.blocks):
            x, self.states[i] = block(x.unsqueeze(0).unsqueeze(0),
                                       self.states[i])
            x = x.squeeze(0).squeeze(0)

        logits = self.model.head(self.model.ln_out(x))
        return logits

    # 内存占用：与序列长度无关，恒定不变！
    # Transformer：KV Cache 占用 O(n * d) 内存
    # RWKV-7：State 占用 O(d^2) 固定内存

# RWKV-7 在多语言场景中表现尤为出色
# 使用 100+ 种语言训练，在英语以外的语言上达到 SOTA 级水平

# 在韩语/日语/中文基准测试中
# 相较同尺寸的 Transformer 表现更优
# -> 多语言分词效率 + 长上下文处理能力

# 按序列长度对比推理成本
#
# Sequence Length | Transformer | RWKV-7
# 1K             | 1x          | 0.8x
# 4K             | 4x          | 0.8x
# 16K            | 16x         | 0.8x
# 64K            | 64x         | 0.8x
# 1M             | OOM         | 0.8x
#
# RWKV-7 的成本与序列长度无关，始终恒定！

# pip install rwkv torch transformers

from rwkv.model import RWKV
from rwkv.utils import PIPELINE

# 加载模型
model = RWKV(
    model='/path/to/RWKV-7-World-3B',
    strategy='cuda fp16'  # GPU FP16
)

pipeline = PIPELINE(model, "rwkv_vocab_v20230424")

# 文本生成
context = "Kubernetes VPA 的优点是"
result = pipeline.generate(
    context,
    token_count=200,
    temperature=0.8,
    top_p=0.9
)
print(result)

# RWKV-7 vLLM 服务（实验性支持）
pip install vllm>=0.6.0

python -m vllm.entrypoints.openai.api_server \
  --model RWKV/rwkv-7-world-3b \
  --tokenizer RWKV/rwkv-7-world-3b \
  --dtype float16 \
  --max-model-len 32768

# 下载 RWKV-7 GGUF 模型后
ollama create rwkv7 -f Modelfile

# Modelfile 示例：
# FROM rwkv-7-world-3b-q4_k_m.gguf
# PARAMETER temperature 0.7
# PARAMETER num_ctx 32768

ollama run rwkv7

# 适合 RWKV-7 的场景：
# 1. 处理超长上下文（100K+ token）
# 2. 边缘设备推理（常数内存）
# 3. 多语言服务（韩语/日语/中文等）
# 4. 实时流式处理（每个 token 耗时恒定）

# Transformer 仍然更优的场景：
# 1. 需要极致性能时（尤其是英语）
# 2. 短上下文（1K 以下）
# 3. 需要利用现有生态系统/工具