필사 모드: Mamba 与 State Space Model 论文深度解析：从选择性 SSM 到 Mamba-2 的 Transformer 替代架构

import torch
import torch.nn as nn
import numpy as np

def make_hippo_matrix(N: int) -> torch.Tensor:
    """生成 HiPPO-LegS 矩阵。

    该矩阵的设计目标是用勒让德多项式基对过去的输入
    做最优近似，有利于捕捉长期依赖关系。

    Args:
        N: 状态维度 (State dimension)

    Returns:
        A: (N, N) HiPPO 矩阵
    """
    P = np.sqrt(2 * np.arange(N) + 1)
    A = np.zeros((N, N))
    for i in range(N):
        for j in range(N):
            if i > j:
                A[i, j] = P[i] * P[j]
            elif i == j:
                A[i, j] = i + 1
            # i < j: 0
    A = -A  # 为保证稳定性取负
    return torch.tensor(A, dtype=torch.float32)


class S4Layer(nn.Module):
    """S4(Structured State Spaces)层的核心实现。

    结合了 HiPPO 初始化的 A 矩阵与基于卷积的并行计算。
    """

    def __init__(self, d_model: int, state_dim: int = 64, seq_len: int = 1024):
        super().__init__()
        self.d_model = d_model
        self.state_dim = state_dim
        self.seq_len = seq_len

        # HiPPO 初始化
        A = make_hippo_matrix(state_dim)
        self.A_log = nn.Parameter(torch.log(torch.clamp(-A.diagonal(), min=1e-4)))
        self.B = nn.Parameter(torch.randn(state_dim, 1) * 0.01)
        self.C = nn.Parameter(torch.randn(1, state_dim) * 0.01)
        self.log_dt = nn.Parameter(torch.rand(d_model).uniform_(-4, -1))

        # skip connection
        self.D = nn.Parameter(torch.ones(d_model))

    def _compute_kernel(self, L: int) -> torch.Tensor:
        """预先计算 SSM 卷积核。"""
        dt = torch.exp(self.log_dt)  # (d_model,)
        A = -torch.exp(self.A_log)   # (state_dim,) 对角元素

        # ZOH 离散化 (假设 A 为对角矩阵)
        dtA = dt.unsqueeze(-1) * A.unsqueeze(0)  # (d_model, state_dim)
        A_bar = torch.exp(dtA)

        # 核计算: K[k] = C @ A_bar^k @ B_bar
        # 用高效的 Vandermonde 乘法实现 O(N*L) 计算
        powers = torch.arange(L, device=A.device).float()
        # A_bar^k = exp(k * dtA)
        kernel = torch.einsum(
            'dn,dn,nl->dl',
            self.C.squeeze(0).expand(self.d_model, -1),
            dt.unsqueeze(-1) * self.B.squeeze(-1).unsqueeze(0),
            torch.exp(dtA.unsqueeze(-1) * powers.unsqueeze(0).unsqueeze(0))
            .squeeze()
            .T
        )

        return kernel  # (d_model, L)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """卷积模式的 forward pass。

        Args:
            x: (batch, seq_len, d_model) 输入序列

        Returns:
            y: (batch, seq_len, d_model) 输出序列
        """
        B, L, D = x.shape
        kernel = self._compute_kernel(L)  # (d_model, L)

        # 基于 FFT 的卷积
        x_perm = x.permute(0, 2, 1)  # (B, D, L)
        k_f = torch.fft.rfft(kernel, n=2 * L)  # zero-pad
        x_f = torch.fft.rfft(x_perm, n=2 * L)
        y = torch.fft.irfft(x_f * k_f, n=2 * L)[..., :L]

        # skip connection
        y = y + self.D.unsqueeze(0).unsqueeze(-1) * x_perm
        return y.permute(0, 2, 1)

import torch
import torch.nn as nn
import torch.nn.functional as F

class SelectiveSSM(nn.Module):
    """Mamba 选择性 State Space Model 的核心实现。

    把 B、C、delta 变成输入的函数，
    从而实现 content-aware 的序列处理。
    """

    def __init__(self, d_model: int, d_state: int = 16, d_conv: int = 4):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state

        # A 矩阵：对角结构，在对数空间中参数化
        A = torch.arange(1, d_state + 1, dtype=torch.float32)
        self.A_log = nn.Parameter(torch.log(A).unsqueeze(0).expand(d_model, -1))

        # 输入依赖的 B、C 投影
        self.B_proj = nn.Linear(d_model, d_state, bias=False)
        self.C_proj = nn.Linear(d_model, d_state, bias=False)

        # delta（步长）投影
        self.dt_proj = nn.Linear(d_model, d_model, bias=True)

        # skip connection
        self.D = nn.Parameter(torch.ones(d_model))

    def selective_scan(
        self,
        x: torch.Tensor,
        dt: torch.Tensor,
        B: torch.Tensor,
        C: torch.Tensor,
    ) -> torch.Tensor:
        """选择性扫描算法（顺序实现）。

        实际的 Mamba 使用针对 GPU SRAM 优化的 CUDA 内核，
        这里为了清楚展示算法逻辑而采用顺序实现。

        Args:
            x: (batch, seq_len, d_model) 输入
            dt: (batch, seq_len, d_model) 离散化步长
            B: (batch, seq_len, d_state) 输入投影
            C: (batch, seq_len, d_state) 输出投影

        Returns:
            y: (batch, seq_len, d_model) 输出
        """
        batch, seq_len, d_model = x.shape
        d_state = B.shape[-1]

        # A 离散化: A_bar = exp(dt * A)
        A = -torch.exp(self.A_log)  # (d_model, d_state)
        dt_A = torch.einsum('bld,dn->bldn', dt, A)  # (B, L, D, N)
        A_bar = torch.exp(dt_A)

        # B 离散化: B_bar = dt * B
        dt_B = torch.einsum('bld,bln->bldn', dt, B)  # (B, L, D, N)

        # 顺序扫描
        h = torch.zeros(batch, d_model, d_state, device=x.device)
        outputs = []

        for t in range(seq_len):
            # h_t = A_bar_t * h_{t-1} + B_bar_t * x_t
            h = A_bar[:, t] * h + dt_B[:, t] * x[:, t].unsqueeze(-1)
            # y_t = C_t @ h_t
            y_t = torch.einsum('bdn,bn->bd', h, C[:, t])
            outputs.append(y_t)

        y = torch.stack(outputs, dim=1)  # (B, L, D)
        return y

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """选择性 SSM 的 forward pass。

        Args:
            x: (batch, seq_len, d_model)

        Returns:
            y: (batch, seq_len, d_model)
        """
        # 计算输入依赖的参数
        B = self.B_proj(x)                          # (B, L, N)
        C = self.C_proj(x)                          # (B, L, N)
        dt = F.softplus(self.dt_proj(x))            # (B, L, D), 保证为正数

        # 选择性扫描
        y = self.selective_scan(x, dt, B, C)

        # skip connection
        y = y + self.D.unsqueeze(0).unsqueeze(0) * x

        return y

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass


@dataclass
class MambaConfig:
    d_model: int = 768
    d_state: int = 16
    d_conv: int = 4
    expand: int = 2
    n_layers: int = 24
    vocab_size: int = 50257
    dropout: float = 0.0


class MambaBlock(nn.Module):
    """Mamba 块：整合选择性 SSM、卷积、门控的结构。

    用一个块取代 Transformer 的 Attention + FFN。
    扩展比率(expand)决定内部维度，
    默认值 2 意味着 d_inner = 2 * d_model。
    """

    def __init__(self, config: MambaConfig):
        super().__init__()
        self.d_model = config.d_model
        self.d_state = config.d_state
        self.d_conv = config.d_conv
        d_inner = config.d_model * config.expand

        # 输入投影：分成两路(z, x)
        self.in_proj = nn.Linear(config.d_model, d_inner * 2, bias=False)

        # 1D 卷积：捕捉局部上下文
        self.conv1d = nn.Conv1d(
            in_channels=d_inner,
            out_channels=d_inner,
            kernel_size=config.d_conv,
            padding=config.d_conv - 1,
            groups=d_inner,  # depthwise convolution
            bias=True,
        )

        # 选择性 SSM 参数
        self.B_proj = nn.Linear(d_inner, config.d_state, bias=False)
        self.C_proj = nn.Linear(d_inner, config.d_state, bias=False)
        self.dt_proj = nn.Linear(d_inner, d_inner, bias=True)

        # A 矩阵：对角结构
        A = torch.arange(1, config.d_state + 1, dtype=torch.float32)
        self.A_log = nn.Parameter(
            torch.log(A).unsqueeze(0).expand(d_inner, -1).clone()
        )
        self.D = nn.Parameter(torch.ones(d_inner))

        # 输出投影
        self.out_proj = nn.Linear(d_inner, config.d_model, bias=False)
        self.norm = nn.RMSNorm(config.d_model)

    def _selective_scan(self, x, dt, B, C):
        """选择性扫描（顺序实现）。实际使用 CUDA 内核。"""
        batch, seq_len, d_inner = x.shape
        d_state = B.shape[-1]

        A = -torch.exp(self.A_log)
        dt_A = torch.einsum('bld,dn->bldn', dt, A)
        A_bar = torch.exp(dt_A)
        dt_B_x = torch.einsum('bld,bln->bldn', dt * x, B)

        h = torch.zeros(batch, d_inner, d_state, device=x.device)
        outputs = []
        for t in range(seq_len):
            h = A_bar[:, t] * h + dt_B_x[:, t]
            y_t = torch.einsum('bdn,bn->bd', h, C[:, t])
            outputs.append(y_t)

        y = torch.stack(outputs, dim=1)
        return y

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Mamba 块的 forward pass。

        Args:
            x: (batch, seq_len, d_model)

        Returns:
            output: (batch, seq_len, d_model)
        """
        residual = x
        x = self.norm(x)

        # 分成两路
        xz = self.in_proj(x)
        x_branch, z = xz.chunk(2, dim=-1)

        # 卷积路径：捕捉局部模式
        x_conv = x_branch.transpose(1, 2)  # (B, D_inner, L)
        x_conv = self.conv1d(x_conv)[:, :, :x.shape[1]]  # causal trim
        x_conv = x_conv.transpose(1, 2)    # (B, L, D_inner)
        x_conv = F.silu(x_conv)

        # 选择性 SSM
        B = self.B_proj(x_conv)
        C = self.C_proj(x_conv)
        dt = F.softplus(self.dt_proj(x_conv))

        y = self._selective_scan(x_conv, dt, B, C)
        y = y + self.D.unsqueeze(0).unsqueeze(0) * x_conv

        # 门控：与 z 支路逐元素相乘
        y = y * F.silu(z)

        # 输出投影 + 残差连接
        output = self.out_proj(y)
        return output + residual

import torch
import torch.nn as nn
import torch.nn.functional as F

class SSDBlock(nn.Module):
    """Mamba-2 的 Structured State Space Duality(SSD)块。

    把序列切分成多个块，块内部用矩阵乘法(并行)处理，
    块之间用 SSM 递归处理的混合算法。
    """

    def __init__(
        self,
        d_model: int = 768,
        d_state: int = 128,
        n_heads: int = 8,
        chunk_size: int = 256,
    ):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.n_heads = n_heads
        self.chunk_size = chunk_size
        self.head_dim = d_model // n_heads

        # 多头结构：Q(C)、K(B)、V(x) 投影
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, d_state * n_heads, bias=False)
        self.v_proj = nn.Linear(d_model, d_model, bias=False)
        self.dt_proj = nn.Linear(d_model, n_heads, bias=True)

        # A 矩阵（按头的标量）
        self.A_log = nn.Parameter(torch.log(torch.arange(1, n_heads + 1, dtype=torch.float32)))

        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.norm = nn.RMSNorm(d_model)

    def _chunk_scan(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        dt: torch.Tensor,
    ) -> torch.Tensor:
        """基于分块的 SSD 扫描。

        Args:
            Q: (batch, n_heads, seq_len, head_dim)  -- 承担 C 的角色
            K: (batch, n_heads, seq_len, d_state)   -- 承担 B 的角色
            V: (batch, n_heads, seq_len, head_dim)  -- 承担 x 的角色
            dt: (batch, n_heads, seq_len)            -- 承担 delta 的角色

        Returns:
            y: (batch, n_heads, seq_len, head_dim)
        """
        B, H, L, D = Q.shape
        C = self.chunk_size

        # 把序列切分成块
        n_chunks = (L + C - 1) // C
        pad_len = n_chunks * C - L
        if pad_len > 0:
            Q = F.pad(Q, (0, 0, 0, pad_len))
            K = F.pad(K, (0, 0, 0, pad_len))
            V = F.pad(V, (0, 0, 0, pad_len))
            dt = F.pad(dt, (0, pad_len))

        # 按块 reshape
        Q = Q.view(B, H, n_chunks, C, D)
        K = K.view(B, H, n_chunks, C, -1)
        V = V.view(B, H, n_chunks, C, D)
        dt = dt.view(B, H, n_chunks, C)

        # 块内部：attention-like 矩阵乘法（可并行处理）
        # 在块内部计算 Q 与 K 的相似度矩阵
        # 生成带 decay 的 causal mask
        A = -torch.exp(self.A_log)  # (n_heads,)
        # 计算累积衰减率
        dt_cumsum = dt.cumsum(dim=-1)  # (B, H, n_chunks, C)

        # 块内部 attention 矩阵
        # M[i,j] = exp(A * (dt_cumsum[i] - dt_cumsum[j])) for i >= j
        decay_diff = dt_cumsum.unsqueeze(-1) - dt_cumsum.unsqueeze(-2)  # (B,H,nc,C,C)
        decay = torch.exp(A.view(1, H, 1, 1, 1) * decay_diff)
        causal_mask = torch.tril(torch.ones(C, C, device=Q.device))
        decay = decay * causal_mask

        # 计算块内部输出
        attn = torch.einsum('bhncqd,bhnckd->bhncqk', Q, K) * decay
        y_intra = torch.einsum('bhncqk,bhnckd->bhncqd', attn, V)

        # 最终 reshape
        y = y_intra.view(B, H, n_chunks * C, D)
        return y[:, :, :L]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """SSD 块的 forward pass。"""
        residual = x
        x = self.norm(x)
        B, L, D = x.shape

        # 多头投影
        Q = self.q_proj(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        K = self.k_proj(x).view(B, L, self.n_heads, self.d_state).transpose(1, 2)
        V = self.v_proj(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        dt = F.softplus(self.dt_proj(x)).transpose(1, 2)  # (B, H, L)

        # SSD 扫描
        y = self._chunk_scan(Q, K, V, dt)

        # 输出投影
        y = y.transpose(1, 2).contiguous().view(B, L, D)
        return self.out_proj(y) + residual

# Mamba 安装与基本用法

# 1. 安装 (需要 CUDA 11.8 以上)
pip install mamba-ssm

# 2. 安装 causal-conv1d 依赖
pip install causal-conv1d>=1.2.0

# 3. 使用预训练模型 (Hugging Face)
pip install transformers

# 4. 用预训练 Mamba 模型生成文本
python -c "
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
import torch

# 加载 Mamba-2.8B 模型
model = MambaLMHeadModel.from_pretrained(
    'state-spaces/mamba-2.8b',
    device='cuda',
    dtype=torch.float16,
)
model.eval()

# 分词器 (兼容 GPT-NeoX 分词器)
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')

# 文本生成
prompt = 'State Space Models are'
input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to('cuda')
output = model.generate(
    input_ids=input_ids,
    max_length=100,
    temperature=0.7,
    top_p=0.9,
)
print(tokenizer.decode(output[0]))
"

import torch
import time
from dataclasses import dataclass

@dataclass
class BenchmarkResult:
    model_name: str
    seq_len: int
    batch_size: int
    latency_ms: float
    memory_mb: float
    tokens_per_sec: float


def benchmark_inference(
    model: torch.nn.Module,
    model_name: str,
    seq_lengths: list[int],
    batch_size: int = 1,
    n_warmup: int = 5,
    n_measure: int = 20,
    device: str = "cuda",
) -> list[BenchmarkResult]:
    """在不同序列长度下测量推理性能。

    Args:
        model: 待测模型
        model_name: 模型名称 (用于结果展示)
        seq_lengths: 待测的序列长度列表
        batch_size: 批大小
        n_warmup: 预热迭代次数
        n_measure: 测量迭代次数
        device: 运行设备

    Returns:
        按序列长度分组的基准测试结果列表
    """
    model = model.to(device).eval()
    results = []

    for seq_len in seq_lengths:
        input_ids = torch.randint(
            0, 50257, (batch_size, seq_len), device=device
        )

        # 清理 GPU 缓存
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()

        # 预热
        with torch.no_grad():
            for _ in range(n_warmup):
                _ = model(input_ids)
        torch.cuda.synchronize()

        # 测量
        latencies = []
        with torch.no_grad():
            for _ in range(n_measure):
                torch.cuda.synchronize()
                start = time.perf_counter()
                _ = model(input_ids)
                torch.cuda.synchronize()
                end = time.perf_counter()
                latencies.append((end - start) * 1000)

        avg_latency = sum(latencies) / len(latencies)
        peak_memory = torch.cuda.max_memory_allocated() / 1024 / 1024
        tokens_per_sec = (batch_size * seq_len) / (avg_latency / 1000)

        results.append(BenchmarkResult(
            model_name=model_name,
            seq_len=seq_len,
            batch_size=batch_size,
            latency_ms=avg_latency,
            memory_mb=peak_memory,
            tokens_per_sec=tokens_per_sec,
        ))

        print(
            f"[{model_name}] seq_len={seq_len:>6d} | "
            f"latency={avg_latency:>8.2f}ms | "
            f"memory={peak_memory:>8.1f}MB | "
            f"throughput={tokens_per_sec:>10.0f} tok/s"
        )

    return results


def compare_models(results_dict: dict[str, list[BenchmarkResult]]):
    """打印模型间的性能比较表格。"""
    print("\n" + "=" * 80)
    print(f"{'模型':<20} {'序列长度':>10} {'延迟(ms)':>10} "
          f"{'内存(MB)':>12} {'吞吐量(tok/s)':>15}")
    print("=" * 80)

    for model_name, results in results_dict.items():
        for r in results:
            print(
                f"{r.model_name:<20} {r.seq_len:>10d} "
                f"{r.latency_ms:>10.2f} {r.memory_mb:>12.1f} "
                f"{r.tokens_per_sec:>15.0f}"
            )
        print("-" * 80)