PRISM: Breaking the O(n) Memory Wall in Long-Context LLM Inference via O(1) Photonic Block Selection

1. 문제 정의 (Problem Statement)

• 메모리 대역폭 병목 현상: 대규모 언어 모델 (LLM) 의 추론 비용은 이제 연산 (Compute) 이 아닌 **메모리 대역폭 (Memory Bandwidth)**에 의해 제한받고 있습니다. 특히 긴 문맥 (Long-context) 에서 토큰을 하나씩 생성할 때마다, 이전 모든 토큰에 대한 키 - 값 (KV) 캐시를 스캔하여 어텐션 점수를 계산해야 합니다.
• O(n) 확장성의 한계: KV 캐시 크기는 문맥 길이 $n$에 비례하여 선형적으로 증가합니다. 현재 GPU 아키텍처 (예: NVIDIA H100) 나 차세대 아키텍처 (Vera Rubin) 에서조차 매 디코딩 단계마다 전체 KV 캐시를 읽는 $O(n)$ 비용이 치명적인 병목이 됩니다.
• 기존 광학 가속기의 한계: 최근의 광학 가속기들은 밀집된 어텐션 계산 (Dense Attention) 에서는 높은 처리량을 보이지만, 긴 문맥을 처리할 때 여전히 KV 캐시 전체를 읽어야 하므로 전자식 접근법과 동일한 $O(n)$ 메모리 확장성 문제를 겪습니다.
• 핵심 통찰: 실제 병목은 어텐션 계산 자체가 아니라, 어떤 KV 블록을 가져와야 할지 결정하는 '블록 선택 (Block Selection)' 단계입니다. 이 단계는 모든 후보 블록에 대한 유사도 검색을 수행해야 하며, 이 과정이 메모리 대역폭을 소모하는 주원인입니다.

2. 방법론 (Methodology)

저자들은 **PRISM (Photonic Ranking via Inner-product Similarity with Microring weights)**이라는 새로운 아키텍처를 제안합니다. 이는 얇은 필름 리튬 니오베이트 (TFLN) 기반의 광학 유사도 엔진으로, 광학 브로드캐스트 앤 웨이트 (Broadcast-and-Weight, B&W) 패러다임을 KV 캐시 블록 선택에 적용합니다.

• 구조적 매칭 (Structural Match):

  1. 브로드캐스트: 쿼리 벡터는 모든 후보 블록에 동일하게 분산됩니다. 이는 광학 분기 (Passive Splitting) 를 통해 자연스럽게 구현됩니다.
  2. 가중치 (Weights): 저장된 블록 시그니처는 quasi-static(준정적) 이며, 마이크로링 공진기 (MRR) 의 전기 - 광학 (Pockels 효과) 프로그래밍으로 빠르게 업데이트됩니다.
  3. 순위 (Rank Order): 정확한 값 계산보다는 순위 (Top-k) 만 중요하므로, 낮은 정밀도 (4~6 비트) 로도 충분합니다.

• 시스템 아키텍처 (5 단계 파이프라인):

  1. 쿼리 인코딩: GPU 가 쿼리 벡터를 생성하고, 이를 WDM(파장 분할 다중화) 레이저 빔에 실어 아날로그 신호로 변환합니다.
  2. 브로드캐스트: 1×N 광학 분할기를 통해 $N$개의 채널로 동일한 쿼리 신호를 동시에 전송합니다.
  3. 가중치 적용: 각 채널의 MRR 배열이 프로그래밍된 시그니처 가중치와 쿼리를 곱합니다 (광학 내적 연산).
  4. 집적 (Integration): 광검출기 (Photodetector) 가 모든 파장의 광력을 합산하여 아날로그 내적 점수 (유사도 점수) 를 생성합니다.
  5. Top-k 선택: 디지털 회로가 가장 높은 점수를 받은 $k$개의 블록 인덱스를 선택하고, 해당 블록만 메모리에서 가져옵니다.

• 시그니처 부호화: KV 블록의 키 벡터 평균 (Mean-key) 을 사용하여 차원을 축소 (예: 128 차원 $\to$ 32 차원) 하고, 이를 MRR 가중치로 매핑합니다. 부호화된 가중치는 균형 광검출 (Balanced Photodetection) 을 통해 음수/양수 값을 직접 표현합니다.

3. 주요 기여 (Key Contributions)

1. 광학 브로드캐스트 검색 아키텍처 제안: KV 캐시 블록 선택 작업이 광학 브로드캐스트 - 가중치 패러다임과 구조적으로 일치함을 최초로 규명하고, 이를 TFLN 기반 하드웨어로 구현했습니다. 이를 통해 $O(n)$ 메모리 읽기를 제거하고 $O(1)$ 광학 지연으로 선택을 수행합니다.
2. 하드웨어 손상 모델링 및 검증: 양자화 (4~8 비트), 열 드리프트, 삽입 손실, 광검출기 노이즈 등 실제 하드웨어의 비이상적 요소를 모델링했습니다. 시뮬레이션 결과, 이러한 손상에도 불구하고 Recall@8 이 90% 이상 유지됨을 확인했습니다.
3. 엔드 - 투 - 엔드 정확도 검증 (Needle-in-a-Haystack): Qwen2.5-7B 모델을 사용하여 4K 에서 64K 토큰까지의 문맥에서 PRISM 을 적용한 결과, 전체 어텐션 (Full Attention) 과 동일한 100% 정확도를 달성했습니다. 이는 광학 선택이 모델 성능을 저하시키지 않음을 입증합니다.
4. 이중 확장 이점 (Double-Scaling Advantage): 문맥 길이가 길어질수록 (1) 광학 계산의 이점이 커지고, (2) 긴 거리를 참조해야 하는 '검색 헤드 (Retrieval Heads)'의 비율이 증가함 (Qwen2.5-7B 의 경우 8K 기준 91% 이상) 을 발견했습니다. 이는 광학 가속기의 효율성이 문맥이 길어질수록 기하급수적으로 향상됨을 의미합니다.

4. 주요 결과 (Results)

• 메모리 트래픽 감소: 64K 문맥 ($k=32, B=128$) 에서 16 배, 128K 문맥에서 32 배의 메모리 트래픽 감소를 달성했습니다. 100 만 토큰 문맥에서는 약 244 배 이상의 감소가 예상됩니다.
• 에너지 효율성: GPU 기반 전체 스캔 (Full Scan) 대비 4 차수 (10,000 배) 이상의 에너지 이점을 보입니다.
  • PRISM 선택 에너지: ~2.29 pJ (동적 전력 기준)
  • GPU H100 스캔 에너지: ~16.3 $\mu$J (128K 문맥 기준)
• 지연 시간 (Latency): 광학 선택 단계의 지연은 약 9 ns로, GPU 전체 스캔 (~5 $\mu$s) 대비 약 500 배 빠릅니다. 메모리에서 실제 블록을 가져오는 시간을 포함하더라도 전체 디코딩 지연은 크게 단축됩니다.
• 정확도: 4K~64K 문맥 범위 내에서 NIAH(Needle-in-a-Haystack) 테스트에서 100% 정확도를 유지하며, 128K 이상에서는 모델 자체의 한계로 정확도가 떨어지지만 이는 블록 선택 방식과 무관합니다.

5. 의의 및 결론 (Significance)

• 메모리 벽 (Memory Wall) 극복: PRISM 은 연산 병목이 아닌 메모리 병목 문제를 해결하는 새로운 패러다임을 제시합니다. 전자식 스캔을 제거함으로써 문맥 길이에 무관한 $O(1)$ 선택 지연을 실현했습니다.
• 차세대 LLM 인프라의 방향성: NVIDIA 가 Vera Rubin 아키텍처에서 KV 캐시 관리를 위해 전용 DPU 를 도입한 것처럼, 메모리 관리가 시스템 설계의 최우선 과제가 됨을 시사합니다. PRISM 은 이를 광학 하드웨어로 최적화한 구체적인 솔루션입니다.
• 실용성: 현재 기술 수준 (TFLN 플랫폼) 에서 구현 가능한 규모 (수만 개의 MRR) 로 설계되었으며, 4K 이상의 실용적인 문맥 길이부터 에너지 효율성이 극대화됩니다.
• 광학 컴퓨팅의 새로운 적용 분야: 단순한 행렬 곱셈 가속을 넘어, 대규모 데이터베이스 검색 및 유사도 기반 선택 작업 (Similarity Search) 에 광학 브로드캐스트가 최적의 솔루션임을 입증했습니다.

요약하자면, PRISM 은 긴 문맥 LLM 추론의 가장 큰 병목인 KV 캐시 스캔 문제를 광학 브로드캐스트 기술을 통해 $O(1)$로 해결함으로써, 에너지 효율성과 속도를 획기적으로 개선하고 차세대 LLM 인프라의 핵심 기술로 자리매김할 가능성을 제시한 연구입니다.