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)