StreamIndex: 메모리 제한적 압축 희소 주의 (CSA) 를 위한 스트리밍 Top-k
요약
본 기술 기사는 메모리 제한적인 압축 희소 주의(CSA)를 위한 스트리밍 Top-k 메커니즘인 StreamIndex를 소개합니다. 기존 CSA 구현체는 전체 중간 점수 텐서를 물질화해야 하므로 긴 시퀀스 길이에서 GPU의 HBM 용량을 초과하는 메모리 부족(OOM) 문제를 일으킵니다. StreamIndex는 이러한 문제를 해결하기 위해, 전체 중간 값을 물질화하지 않고 청크 기반으로 Top-k를 처리하고 병합하는 드라이버를 제공하며, 이를 통해 훨씬 긴 시퀀스 길이에서도 효율적으로 작동할 수 있음을 입증했습니다.
핵심 포인트
- StreamIndex는 메모리 제한적인 CSA 파이프라인을 위한 Triton 구현입니다.
- 기존 방식은 전체 중간 점수 텐서 물질화로 인해 긴 시퀀스 길이(S=65,536 이상)에서 OOM 문제가 발생합니다.
- StreamIndex는 청크 기반의 Top-k 드라이버를 사용하여 메모리 사용량을 크게 줄이고 더 긴 시퀀스 길이까지 처리할 수 있습니다.
- 제안된 방법은 기존 주의 커널(TileLang 등)과 호환되며, 높은 재현율(recall)을 유지하면서 효율적인 실행 시간을 보여줍니다.
DeepSeek-V3.2 과 V4 는 Compressed Sparse Attention (CSA) 를 도입했습니다: CSA 는 학습된 점수 투영 (압축된 키에 대한) 인덱서를 통해 점수를 산출하고, 쿼리당 Top-k 을 선택하며, 희소 주의 커널은 오직 해당 항목만 읽습니다. 공개 CSA 구현체는 [B, S, H_I, T] FP32 점수 텐서를 Top-k 감소를 전에 물질화 (materialize) 합니다. H_I=64 인덱서 헤드와 V4-Flash 압축 비율 m=4 를 사용할 때, 그 중간 값은 시퀀스 길이 S=65,536 에서 256 GB 가 되어 단일 GPU 의 고대역폭 메모리 (HBM) 예산을 초과합니다. 우리는 StreamIndex 를 제시하며, 이는 CSA 파이프라인의 Triton 구현이며, 핵심 구성 요소는 전체 중간 값을 물질화하지 않는 chunked partition-merge Top-k 드라이버입니다. 단일 NVIDIA H200 에서 인덱서 단계 (레이어) 수준에서 합성적이지만 현실적인 V4 형식의 입력을 사용할 때, 물질화 경로는 S=65,536 에서 메모리 부족 (OOM) 을 발생하며, StreamIndex 는 동일한 인덱서를 S=1,048,576 까지 6.21 GB 피크 HBM 으로 실행합니다. 작은 S 에서 두 구현 모두 적합할 때 Set-overlap recall 은 물질화 기준값과 비트 정확도 (bit-exact) 를 가집니다. 세 가지 5 점 디자인 스페이스 스윕 (chunk size, key-tile size, top-k) 을 통해 평균 recall 은 1.0000 로 반올림되며, 모든 셀에서 최소 recall 은 0.9980 이상입니다. chunked 드라이버는 TileLang 의 파이프라인 주의 커널과 조화됩니다: S=262,144 와 V4-Flash 차원을 사용할 때, 물질화 인덱서와 TileLang 주의 OOM 을 발생시키지만, chunked 인덱서와 동일한 주의 커널은 1.97 초에 18.56 GB 피크로 실행됩니다. 우리의 기여는 인덱서 단계에 집중하며, 더 빠른 주의 커널이나 실제 체크포인트 엔드 투 엔드 동작에 대한 주장은 없습니다. 코드: https://github.com/RightNow-AI/StreamIndex.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.LG의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기