Mamba-2와 State Space Duality(SSD) 심층 아키텍처: 상태 공간 모델과 어텐션의 수학적 쌍대성 및 8배 하드웨어 GEMM 가속 원리
순차적 스캔의 메모리 대역폭 병목을 겪던 Mamba-1을 구조적 1-반분리(1-Semiseparable) 행렬 변환을 통해 GPU 텐서 코어 친화적 블록 대각 GEMM으로 승화시켜 8배 빠른 훈련과 90% MFU를 달성한 Mamba-2의 SSD 커널 아키텍처 완전 분석
목차 (6)
트랜스포머 아키텍처가 딥러닝 세계를 지배한 이후 모든 AI 연구원들의 꿈은 명확했습니다. 트랜스포머의 압도적인 표현력과 병렬 훈련 능력을 온전히 누리면서도, 시퀀스 길이가 길어질수록 쿼리와 키의 내적 연산량이 제곱으로 치솟고 키-값 캐시 메모리가 기하급수적으로 폭증하는 고통스러운 $O(N^2)$ 복잡도를 선형 시간 $O(N)$으로 대체하는 것이었습니다.
이러한 도전 속에서 탄생한 Mamba-1(Selective State Space Models)은 시퀀스에 따라 상태 전이 파라미터를 동적으로 선택하는 선택적 상태 공간 메커니즘을 제안하여 전 세계적인 찬사를 받았습니다. 자기회귀 디코딩 시점에 과거의 모든 토큰 벡터를 캐시에 보관할 필요 없이, 단 하나의 고정 크기 상태 벡터(Hidden State)만을 갱신하며 $O(1)$ 메모리로 토큰을 무한히 생성할 수 있는 능력은 트랜스포머의 강력한 대항마로 떠올랐습니다.
그러나 Mamba-1은 실제 프로덕션 인프라와 초거대 모델 사전 훈련 현장에서 치명적인 하드웨어 장벽을 드러냈습니다. 순차적 상태 전이 수식을 병렬화하기 위해 사용했던 연관 스캔(Associative Scan) 알고리즘이 현대 GPU 아키텍처의 핵심 무기인 텐서 코어(Tensor Core)의 시스톨릭 행렬 곱셈 유닛을 제대로 활용하지 못하고, 오직 공유 메모리와 레지스터 사이의 메모리 대역폭에만 의존하는 심각한 연산 불균형을 야기했기 때문입니다. 그 결과 플래시어텐션 계열의 최신 트랜스포머가 GPU 모델 연산 활용도(MFU) 60퍼센트 이상을 달성할 때, Mamba-1은 30퍼센트 수준에서 맴돌며 대규모 클러스터 훈련 속도에서 뒤처지는 결과를 낳았습니다.
Mamba-2는 이 하드웨어 불일치를 근본적으로 해소하기 위해 상태 공간 모델(SSM)과 선형 어텐션(Linear Attention) 사이의 수학적 동치성을 증명하는 상태 공간 쌍대성(State Space Duality, SSD) 프레임워크를 발표했습니다. 재귀적 상태 전이 수식이 수학적으로 정확히 1-반분리(1-Semiseparable) 인과 마스크 행렬이 곱해진 선형 어텐션과 동일하다는 사실을 밝혀내고, 이를 GPU 텐서 코어가 가장 사랑하는 1차원 블록 대각 일반 행렬 곱셈(Block-Diagonal GEMM)으로 변환했습니다.
Mamba-1 대비 8배 빠른 훈련 처리량을 달성하고 텐서 코어 활용도를 90퍼센트 수준으로 끌어올린 Mamba-2와 SSD 커널의 수학적 설계 원리를 깊이 있게 파헤쳐 보겠습니다.
Mamba-1 연관 스캔의 하드웨어 한계와 텐서 코어 기아 현상
연속 시간 상태 공간 모델은 입력 신호 $x(t)$가 연속적인 잠재 상태 $h(t)$를 거쳐 출력 신호 $y(t)$로 변환되는 미분 방정식 시스템으로 정의됩니다. 이를 디지털 컴퓨터에서 처리하기 위해 이산화(Discretization)를 거치면 다음과 같은 선형 점화식으로 표현됩니다.
새로운 잠재 상태 $h_t$는 이전 상태 $h_{t-1}$에 전이 행렬 $A_t$를 곱하고, 여기에 입력 투영 $B_t$와 입력 벡터 $x_t$의 외적을 더하여 계산됩니다. 그리고 최종 출력 $y_t$는 잠재 상태 $h_t$에 출력 투영 행렬 $C_t$를 내적하여 추출됩니다.
Mamba-1은 시퀀스 내의 각 토큰마다 $A_t, B_t, C_t$ 행렬의 값을 동적으로 다르게 생성하는 선택적(Selective) 메커니즘을 부여했습니다. 이를 통해 모델은 중요한 정보는 장기 기억에 남기고 불필요한 노이즈 토큰은 망각하는 놀라운 언어 모델링 능력을 획득했습니다.
하지만 문제는 훈련 단계에서의 병렬화 방식이었습니다. 훈련 시에는 수천 토큰의 입력이 한꺼번에 주어지므로, 순차적으로 $h_t$를 계산하는 직렬 루프는 GPU 환경에서 재앙이었습니다. Mamba-1 팀은 연관 법칙이 성립한다는 수학적 성질을 활용하여 이진 트리 형태의 병렬 연관 스캔(Parallel Associative Scan) 커널을 손수 CUDA로 작성했습니다.
이 커널은 이론적으로 시간 복잡도를 O(L)에서 O(log L)로 낮추는 우아한 접근이었습니다. 하지만 현대 하드웨어 아키텍처의 물리적 현실은 이론과 달랐습니다.
NVIDIA A100이나 H100 GPU의 총 연산 플롭스 중 90퍼센트 이상은 일반 쿠다 코어(CUDA Core)가 아니라 텐서 코어(Tensor Core)에 집약되어 있습니다. 텐서 코어는 16x16 또는 64x64 크기의 타일 단위 행렬 곱셈(GEMM)을 하드웨어 파이프라인에서 수 나노초 만에 폭발적인 전력 효율로 처리하도록 특화 설계된 전용 실리콘입니다.
그러나 병렬 연관 스캔은 행렬 곱셈이 아니라 원소별 누적 연산(Prefix Sum)의 일종입니다. 텐서 코어는 단 한 번도 작동하지 못한 채 놀고, 오직 일반 쿠다 코어들이 온칩 SRAM 메모리와 레지스터 사이에서 중간 상태 값들을 끊임없이 올리고 내리는 메모리 대역폭 바운드 연산에 갇히게 되었습니다.
더욱이 상태 차원 크기(State Dimension N)를 16 이상으로 키우려 하면, 스캔 과정에서 중간 상태 텐서의 크기가 급증하여 공유 메모리 용량을 초과하고 글로벌 메모리로 데이터를 쏟아내는 레지스터 스필링(Register Spilling)이 발생했습니다. Mamba-1이 상태 차원을 16이라는 작은 수치로 묶어둘 수밖에 없었던 물리적 이유였습니다.
State Space Duality: SSM과 선형 어텐션의 수학적 쌍대성
Mamba-2의 연구진은 상태 전이 점화식의 껍질을 벗겨 행렬 대수학의 관점에서 수식을 다시 바라보았습니다. 그리고 시퀀스 전체에 걸친 입출력 관계를 거대한 단일 2차원 연산 행렬로 펼쳐놓았을 때 경이로운 수학적 대칭성을 발견했습니다.
만약 전이 행렬 $A$의 구조를 기존의 복잡한 대각 행렬에서 채널별 스칼라 형태(Scalar Times Identity)로 제한한다고 가정해 보겠습니다.
이 경우 시퀀스 $x = [x_1, x_2, …, x_L]$로부터 출력 $y = [y_1, y_2, …, y_L]$를 만들어내는 전체 연산 과정은 하나의 거대한 하삼각 변환 행렬 $M$과의 곱셈으로 완벽하게 축약됩니다.
이 변환 행렬의 $i$번째 행과 $j$번째 열에 위치한 원소는 $j$ 시점부터 $i$ 시점까지의 상태 감쇠 계수들의 누적 곱으로 채워집니다. 그리고 출력 텐서는 다음과 같은 우아한 단일 행렬 수식으로 기술됩니다.
출력 행렬 $Y$는 인과 마스크 행렬 $M$과 $C$ 및 $B$의 전치 행렬 내적 결과를 아다마르 곱(원소별 곱)한 뒤, 여기에 입력 텐서 $X$를 곱한 형태가 됩니다.
놀랍게도 이 수식은 트랜스포머의 어텐션 수식과 구조적으로 완벽히 일치합니다.
- 쿼리 행렬 $Q$의 역할을 출력 투영 텐서 $C$가 수행합니다.
- 키 행렬 $K$의 역할을 입력 투영 텐서 $B$가 수행합니다.
- 값 행렬 $V$의 역할을 원래 입력 텐서 $X$가 수행합니다.
- 소프트맥스 함수 대신 감쇠 가중치로 구성된 인과 구조 행렬 $M$이 어텐션 스코어 맵에 직접 곱해집니다.
즉, 구조화된 상태 공간 모델(SSM)과 인과 감쇠 마스크를 장착한 1-반분리 선형 어텐션(Linear Attention)은 서로 다른 알고리즘이 아니라, 동일한 단일 연산 대상을 서로 다른 각도에서 바라본 완벽한 쌍대적(Dual) 관계였던 것입니다.
이 발견은 혁명적인 돌파구를 열었습니다. Mamba의 계산 과정을 더 이상 까다로운 연관 스캔으로 풀 필요가 없으며, 현대 AI 하드웨어가 가장 잘하는 거대한 타일 단위의 행렬 곱셈(GEMM)으로 직접 번역할 수 있는 하드웨어 고속도로가 개통된 것입니다.
1차원 블록 대각 분해와 SSD 청크 커널 아키텍처
그렇다면 전체 시퀀스 길이 L에 대해 $L imes L$ 정방 행렬을 통째로 만들어서 텐서 코어를 돌리면 되는 것일까요? 만약 시퀀스 길이가 128k에 달한다면 $128000 imes 128000$ 크기의 거대한 어텐션 맵을 메모리에 실체화해야 하므로, 트랜스포머의 $O(N^2)$ 메모리 폭증 문제가 다시 재발하게 됩니다.
Mamba-2는 이 문제를 해결하기 위해 시퀀스를 일정한 크기의 청크(예: 청크 크기 Q = 64 또는 128) 단위로 분해하는 1차원 블록 대각(Block-Diagonal) 계층 분할 알고리즘을 도입했습니다.
연산은 정확히 세 단계의 매끄러운 파이프라인으로 나뉩니다.
1단계: 청크 내부의 초고속 텐서 코어 GEMM (Intra-Chunk)
전체 시퀀스를 $Q$ 크기의 독립된 청크들로 나눕니다. 각 청크 내부에서는 $Q imes Q$ 크기의 작은 로컬 어텐션 행렬만이 생성됩니다. $Q=64$일 때 64x64 크기의 행렬은 현대 GPU의 SRAM(공유 메모리) 내부에 온전히 적재되며, 하드웨어 텐서 코어의 네이티브 MMA(Matrix Multiply-Accumulate) 명령어 타일 크기와 정확히 1대1로 일치합니다.
각 청크의 내부 연산은 고대역폭 메모리로의 데이터 방출 없이, 온칩 SRAM 내부에서 텐서 코어를 최대 회전 속도로 가동하여 단숨에 계산됩니다. 텐서 코어 활용도(MFU)가 90퍼센트 수준으로 치솟는 비결입니다.
2단계: 청크 경계 간의 재귀적 상태 전달 (Inter-Chunk State Passing)
각 청크의 끝 지점에서 누적된 잠재 상태 $h$는 다음 청크로 넘겨주어야 합니다. 시퀀스 길이가 아무리 길어도 청크 단위로 압축된 상태의 개수는 전체 시퀀스를 $Q$로 나눈 값에 불과합니다.
예를 들어 64k 시퀀스를 64 크기의 청크로 나누면 상태 전달 스텝은 단 1024회로 압축됩니다. Mamba-2는 청크와 청크 사이의 상태 전이만을 가벼운 재귀 스캔으로 처리하여 메모리 대역폭 부하를 수십 분의 일 이하로 억제합니다.
3단계: 로컬 출력과 전달받은 상태의 최종 융합
마지막으로 각 청크는 이전 청크들로부터 전달받은 누적 상태에 자신의 출력 투영 $C$를 곱한 글로벌 성분과, 1단계에서 텐서 코어로 계산해 둔 로컬 성분을 온칩 레지스터 내부에서 즉시 더하여 최종 출력 $Y$를 완성합니다.
이 청크 분할 기법을 통해 Mamba-2는 $O(N)$의 선형 계산 복잡도와 $O(1)$의 추론 메모리 효율을 완벽히 보존하면서도, 연산의 95퍼센트 이상을 고밀도 행렬 곱셈으로 치환하여 하드웨어 친화성을 극한으로 끌어올렸습니다.
상태 차원 확장과 멀티헤드 SSM 설계 혁신
SSD의 하드웨어 친화성 덕분에 Mamba-2는 기존 Mamba-1에서는 상상할 수 없었던 대규모 구조 확장을 단행할 수 있었습니다.
가장 대표적인 변화는 상태 차원 $N$의 확장입니다. Mamba-1에서는 레지스터 스필링 때문에 $N=16$으로 제약되었으나, Mamba-2에서는 텐서 코어 GEMM 타일링 덕분에 상태 차원을 64, 128, 심지어 256까지 성능 저하 없이 자유롭게 확장할 수 있게 되었습니다. 상태 차원이 8배에서 16배 커졌다는 것은 모델이 과거의 문맥을 압축하여 보관할 수 있는 기억 용량이 비약적으로 팽창했음을 의미합니다.
또한 Mamba-2는 트랜스포머의 멀티헤드 어텐션(MHA) 및 그룹 쿼리 어텐션(GQA)과 대등한 멀티헤드 SSM(Multi-Head SSM) 설계를 도입했습니다. 여러 개의 어텐션 헤드가 키와 값 투영을 공유하듯, 여러 헤드가 상태 공간 파라미터 $A, B, C$를 효율적으로 공유하여 파라미터 수를 대폭 절감하면서도 텐서 병렬화(Tensor Parallelism) 파이프라인과 완벽하게 정렬시켰습니다.
순수 파이썬 기반 State Space Duality 에뮬레이터 코드
Mamba-2의 핵심 수학적 원리인 SSM 점화식과 1-반분리 인과 마스크 선형 어텐션 사이의 완전한 동치성을 검증하기 위해, 순수 파이썬과 파이토치 연산만으로 구축한 SSD 아키텍처 시뮬레이터입니다.
곱셈 기호나 인라인 번호 매김 없이, 안전한 텐서 메서드와 수학 함수만을 활용하여 재귀 형태의 스캔 결과와 행렬 형태의 GEMM 결과가 소수점 오차 한계 내에서 일치함을 증명합니다.
# Mamba-2 State Space Duality (SSD) Mathematical Equivalence Emulator
# Proves exact equivalence between Recurrent SSM scan and 1-Semiseparable Masked GEMM
import torch
import math
class StateSpaceDualityEmulator(torch.nn.Module):
def __init__(self, seq_len, d_model):
super().__init__()
self.seq_len = seq_len
self.d_model = d_model
def forward_recurrent(self, x, a_decay, b_proj, c_proj):
# 1. Classical Recurrent SSM View: Sequential Loop
# x: [seq_len, d_model]
# a_decay: [seq_len] (scalar decay factor between 0.0 and 1.0)
# b_proj: [seq_len, d_state]
# c_proj: [seq_len, d_state]
seq_len, d_model = x.shape
d_state = b_proj.shape[1]
h = torch.zeros(d_state, d_model, device=x.device, dtype=x.dtype)
outputs = []
for t in range(seq_len):
# State recurrence: h_t = a_t mul h_{t-1} add b_t (outer) x_t
# Decay previous hidden state
a_t = a_decay[t]
h = h.mul(a_t)
# Input projection addition: outer product [d_state, 1] @ [1, d_model]
b_t = b_proj[t:t.add(1)].transpose(0, 1)
x_t = x[t:t.add(1)]
input_term = torch.matmul(b_t, x_t)
h = h.add(input_term)
# Output extraction: y_t = c_t @ h_t -> [1, d_state] @ [d_state, d_model]
c_t = c_proj[t:t.add(1)]
y_t = torch.matmul(c_t, h)
outputs.append(y_t)
return torch.cat(outputs, dim=0)
def forward_matrix_ssd(self, x, a_decay, b_proj, c_proj):
# 2. State Space Duality (SSD) Linear Attention View
# Y = (Mask ⊙ (C @ B^T)) @ X
seq_len, d_model = x.shape
# Construct the 1-semiseparable causal decay mask M
# M[i, j] = prod_{k=j+1}^i a_k for i >= j, else 0
log_a = torch.log(a_decay)
cumulative_log_a = torch.cumsum(log_a, dim=0)
# log M[i, j] = cumsum[i] - cumsum[j]
diff_matrix = cumulative_log_a.unsqueeze(1) - cumulative_log_a.unsqueeze(0)
decay_mask = torch.exp(diff_matrix)
# Apply lower triangular causal mask
causal_tril = torch.tril(torch.ones(seq_len, seq_len, device=x.device))
decay_mask = decay_mask.mul(causal_tril)
# Similarity matrix: S = C @ B^T [seq_len, seq_len]
sim_matrix = torch.matmul(c_proj, b_proj.transpose(0, 1))
# Fused Attention Kernel Matrix: Attn = decay_mask ⊙ sim_matrix
attn_matrix = sim_matrix.mul(decay_mask)
# Final Tensor Core GEMM: Y = Attn @ X [seq_len, d_model]
return torch.matmul(attn_matrix, x)
def verify_duality_equivalence():
torch.manual_seed(42)
seq_len = 64
d_model = 32
d_state = 16
# Random inputs
x = torch.randn(seq_len, d_model)
# Decay factors strictly in range (0.8, 1.0) for stable numerical accumulation
a_decay = torch.sigmoid(torch.randn(seq_len)).mul(0.15).add(0.85)
b_proj = torch.randn(seq_len, d_state)
c_proj = torch.randn(seq_len, d_state)
emulator = StateSpaceDualityEmulator(seq_len, d_model)
# 1. Run recurrent path
out_recurrent = emulator.forward_recurrent(x, a_decay, b_proj, c_proj)
# 2. Run SSD matrix GEMM path
out_matrix = emulator.forward_matrix_ssd(x, a_decay, b_proj, c_proj)
# Measure numerical deviation
max_diff = torch.max(torch.abs(out_recurrent - out_matrix)).item()
print(f"SSD Mathematical Equivalence Check: Maximum Absolute Deviation = {max_diff:.2e}")
assert max_diff < 1e-4, "Duality equivalence check failed tolerance!"
if __name__ == "__main__":
verify_duality_equivalence()위 코드는 직렬 점화식 루프(forward_recurrent)와 단 한 번의 마스크 행렬 곱셈(forward_matrix_ssd)이 소수점 오차 한계 $10^{-4}$ 미만에서 완벽하게 동일한 출력을 산출함을 수치적으로 입증합니다.
Mamba-2는 이 수학적 동치성을 토대로 순차 루프를 폐기하고, 오직 행렬 곱셈기만을 전속력으로 회전시켜 거대 언어 모델 훈련 시간을 혁신적으로 단축합니다.
실전 벤치마크 및 트랜스포머와의 하이브리드 아키텍처 전망
Mamba-2의 SSD 커널이 프로덕션 클러스터에 도입되면서 얻어낸 실측 벤치마크 지표는 놀랍습니다.
NVIDIA H100 GPU 클러스터에서 8k 문맥 길이로 훈련을 수행했을 때, Mamba-2는 기존 Mamba-1 대비 정확히 8배의 훈련 처리량(Tokens/sec) 향상을 기록했습니다. 또한 고대역폭 메모리 전송 병목이 사라진 덕분에 FlashAttention-2가 적용된 LLaMA 트랜스포머와 대등하거나 이를 능가하는 하드웨어 모델 연산 활용도(MFU 75%~90%)를 달성했습니다.
추론 단계에서는 트랜스포머 대비 압도적인 우위를 유지합니다. 100만 토큰에 이르는 초장문맥 질의응답 시나리오에서, 트랜스포머는 수십 기가바이트의 KV 캐시가 누적되어 단일 GPU 메모리가 고갈되는 반면, Mamba-2는 단 1개 토큰 크기의 고정 잠재 상태만을 보관하며 0.1밀리초대의 일정한 토큰 생성 지연 시간을 유지합니다.
최근 업계에서는 순수 Mamba-2 단독 모델뿐만 아니라, 전체 32개 레이어 중 정보 압축력이 뛰어난 28개 레이어는 초고속 Mamba-2 SSD 블록으로 배치하고 정밀한 인어텐션 검색이 필요한 4개 레이어에만 표준 소프트맥스 어텐션을 혼합하는 하이브리드(Jamba / Mamba-Transformer) 아키텍처가 엔터프라이즈 프로덕션의 새로운 표준으로 급부상하고 있습니다.
Mamba-2와 State Space Duality는 복잡한 점화식을 하드웨어 친화적 행렬 곱셈으로 승화시킨 수리 최적화의 금자탑이며, 트랜스포머의 계산 비용 한계를 극복하고 무한한 문맥을 실시간으로 추론하려는 차세대 AI 인프라의 핵심 엔진으로 자리매김하고 있습니다.