[Paper] Stream-CQSA: Flexible Workload Scheduling을 통한 Attention Computation의 Out-of-Memory 방지

Source: arXiv - 2604.20819v1

번역을 위해 원본 텍스트를 제공해 주시겠어요? 텍스트를 받으면 그대로 마크다운 형식과 코드 블록을 유지하면서 한국어로 번역해 드리겠습니다.

개요

논문 **“Stream‑CQSA: Flexible Workload Scheduling을 통한 Attention Computation의 Out‑of‑Memory 방지”**는 현대 대형 언어 모델(LLM)의 핵심 병목 현상인 매우 긴 시퀀스를 처리할 때 정확한 자체 어텐션의 2차 메모리 폭증 문제를 다룹니다. 어텐션을 독립적인 하위 연산들의 집합으로 재구성하여 실시간으로 스트리밍할 수 있게 함으로써, 저자들은 근사화나 비용이 많이 드는 GPU 간 통신 없이도 단일 GPU에서 수십억 토큰 입력에 대해 정확한 어텐션을 실행할 수 있음을 보여줍니다.

주요 기여

• CQS Divide operation – 순환 쿼럼 집합(CQS) 이론에서 파생된 새로운 분해 방법으로, 전체 시퀀스 어텐션을 수학적으로 독립적인 서브시퀀스 작업으로 분할하여 원래 어텐션 행렬을 정확히 복원할 수 있게 합니다.
• Stream‑CQSA framework – 메모리 적응형 스케줄러로, 어텐션 작업량을 동적으로 분할하여 사용자가 지정한 GPU 메모리 예산에 맞추고, 어텐션을 스트리밍 가능한 파이프라인으로 전환합니다.
• Hardware‑agnostic execution – 이 접근 방식은 단일 디바이스에서 동작하며 다중 GPU 샤딩이나 커스텀 커널에 의존하지 않습니다; 기존 트랜스포머 라이브러리 위에 레이어링할 수 있습니다.
• Empirical validation – 실험을 통해 메모리 사용량에 따라 예측 가능한 선형 확장을 보여주며, 10억 토큰을 초과하는 시퀀스에서도 정확한 어텐션을 성공적으로 실행함을 입증했습니다. 이 모든 과정에서 어텐션의 수학적 정의는 동일하게 유지됩니다.

Methodology

1. Cyclic Quorum Set (CQS) Theory – 쿼럼 시스템에서 요소들의 부분집합(쿼럼)은 제어된 방식으로 교차합니다. 저자들은 이 아이디어를 어텐션의 query, key, value 텐서에 적용하여 시퀀스를 겹치는 블록으로 나누는 순환 쿼럼을 구성합니다.
2. CQS Divide – 쿼럼을 이용해 전체 어텐션 행렬 (A = \text{softmax}(QK^\top)V)를 더 작은 어텐션 서브‑매트릭스들의 합으로 표현합니다. 각 서브‑매트릭스는 메모리에 들어갈 수 있는 서브‑시퀀스에 대해 계산됩니다. 쿼럼이 모든 위치 쌍을 정확히 한 번씩 포함하도록 설계되었기 때문에, 서브 결과들을 다시 결합하면 단일 연산과 동일한 출력을 얻을 수 있습니다.
3. Streaming Scheduler – Stream‑CQSA는 각 서브 어텐션을 큐의 작업으로 취급합니다. 스케줄러는 GPU 메모리 사용량을 모니터링하고, 필요에 따라 작업을 스트리밍 입·출력하며 중간 텐서를 호스트 메모리로 스와핑합니다. 각 서브 작업이 독립적이므로 작업 간 통신이 필요하지 않습니다.
4. Implementation – 저자들은 파이토치 위에 파이프라인을 구축했으며, 표준 텐서 연산과 CUDA 스트림을 사용했습니다. 추가로 필요한 것은 사용 가능한 메모리 예산에 따라 블록 크기를 결정하는 가벼운 컨트롤러뿐입니다.

Results & Findings

• Predictable memory scaling – 메모리 사용량은 전체 시퀀스 길이가 아니라 선택한 블록 크기에 따라 선형적으로 증가합니다.
• Zero approximation error – 분해가 수학적으로 정확하기 때문에, 하위 모델 성능(예: 언어 모델링 벤치마크의 퍼플렉시티)이 단순 전체‑어텐션 실행과 동일합니다.
• No extra hardware – 모든 실험은 단일 NVIDIA A100 (40 GB) 또는 24 GB RTX‑3090에서도 실행되어, 이 방법이 대부분의 연구실 및 고성능 워크스테이션에서도 실용적임을 보여줍니다.

Practical Implications

• Long‑document processing – 법률 계약 분석, 과학 논문 요약, 코드베이스 이해와 같은 애플리케이션이 이제 전체 문서(수백 메가바이트)를 트랜스포머에 잘라내기나 청크별 휴리스틱 없이 입력할 수 있습니다.
• Cost‑effective scaling – 기업은 긴 컨텍스트에 대한 추론 시 다중 GPU 클러스터를 피할 수 있어 하드웨어 비용과 엔지니어링 복잡성을 모두 줄일 수 있습니다.
• Plug‑and‑play integration – Stream‑CQSA가 표준 텐서 연산과 함께 작동하므로 기존 트랜스포머 코드베이스(예: Hugging Face Transformers, DeepSpeed)는 어텐션 모듈을 교체하고 스케줄러 래퍼를 추가함으로써 쉽게 도입할 수 있습니다.
• Enabling new research – 방대한 컨텍스트에서 어텐션 패턴을 연구하는 연구자들(예: emergent reasoning, 장거리 의존성 탐색)은 이제 메모리 제한 없이 정확한 도구를 사용할 수 있습니다.

Limitations & Future Work

• Throughput trade‑off – 스트리밍은 GPU와 호스트 메모리 사이에 추가 데이터 이동을 발생시켜, 짧은 시퀀스에서 완전 GPU 내 구현에 비해 초당 토큰 처리량이 낮아질 수 있습니다.
• Scheduler overhead – 현재 블록 크기 지정에 대한 휴리스틱은 단순합니다; 보다 정교한 메모리 예측 모델을 사용하면 성능을 더 최적화할 수 있습니다.
• Extension to sparse/approximate attention – 이 방법은 이미 정확성을 달성하지만, 기존의 sparse‑attention 커널과 결합하면 초장 시퀀스에 대해 더욱 높은 속도 향상을 얻을 수 있습니다.
• Multi‑device orchestration – 저자들은 여러 GPU 또는 TPU에 걸친 협조 스트리밍을 탐색하여 단일 디바이스의 계산 용량을 초과하는 작업을 처리하면서도 정확성 보장을 유지하려고 합니다.

Stream‑CQSA는 “out‑of‑memory” 오류를 단단한 장벽에서 조정 가능한 자원 노브로 바꾸어, 개발자들이 정확성을 희생하지 않고도 진정한 장기 컨텍스트 어텐션을 활용할 수 있는 길을 열어줍니다.

저자

• Yiming Bian
• Joshua M. Akey

논문 정보

• arXiv ID: 2604.20819v1
• Categories: cs.LG, cs.DC
• Published: 2026년 4월 22일
• PDF: PDF 다운로드