DGX Spark의 높은 연산 성능과 Mac Studio M3 Ultra의 높은 메모리 대역폭을 프리필/디코드 단계로 분리해 결합하고, 레이어별 KV 스트리밍으로 통신 오버헤드를 숨겨 전체 추론 속도를 크게 높이는 방법을 살펴봅니다.
URL: https://blog.exolabs.net/nvidia-dgx-spark/
최근 NVIDIA DGX Spark™ 2대를 얼리 액세스로 받았습니다. NVIDIA는 이를 ‘세계에서 가장 작은 AI 슈퍼컴퓨터’라고 부릅니다. FP16 기준 약 100 TFLOPs의 성능과 273 GB/s로 동작하는 128GB의 CPU‑GPU 코히어런트(coherent) 메모리를 갖추고 있습니다.
EXO를 통해 우리는 이미 M3 Ultra 칩을 탑재한 Apple Mac Studio 클러스터에서 LLM을 실행해 왔습니다. Mac Studio는 819 GB/s의 512GB 통합 메모리를 갖지만, GPU의 FP16 성능은 약 26 TFLOPs 정도입니다.
DGX Spark는 연산이 4배, Mac Studio는 메모리 대역폭이 3배입니다.
그렇다면 둘을 결합하면 어떨까요? DGX Spark가 잘하는 일은 DGX Spark가, Mac Studio가 잘하는 일은 Mac Studio가 맡게 해서, 같은 추론 요청(inference request) 안에서 함께 쓰면 어떨까요?
NVIDIA DGX Spark™ 얼리 액세스 유닛(품질 관리 감독관과 함께)
EXO로 LLM 추론에 사용한 Mac Studio M3 Ultra 스택
사용자 입장에서 보이는 것은 결국 두 숫자로 요약됩니다.
시스템에서 하는 모든 일은 이 두 숫자를 개선하기 위한 것입니다. 둘을 동시에 최적화하기 어려운 이유는, 이 값들이 같은 요청의 서로 다른 두 단계, 즉 프리필(prefill) 과 디코드(decode) 에 의해 좌우되기 때문입니다.
그럼 이 두 단계에서 내부적으로는 무엇이 일어나며, 왜 서로 이렇게 다르게 동작할까요?
그림 1: 요청 라이프사이클에서 프리필 단계(노란색, TTFT 결정) 뒤에 디코드 단계(파란색, TPS 결정)가 이어짐
프리필은 프롬프트를 처리하고, 트랜스포머 각 레이어에 대해 KV 캐시를 구축합니다. KV 캐시는 프롬프트의 각 토큰에 대해 여러 벡터 묶음으로 이루어져 있습니다.
이 벡터들은 프리필 동안 저장되어, 디코드에서 다시 계산할 필요가 없도록 합니다.
컨텍스트가 커질수록, 모든 토큰이 프롬프트 내 다른 모든 토큰을 어텐션해야 하므로 연산량은 프롬프트 길이의 제곱에 비례해 증가합니다(Θ(s²)).
Flash Attention 같은 현대적 기법을 사용하면, 이동되는 데이터 양은 프롬프트 길이에 선형으로 증가하도록 만들 수 있습니다(Θ(s)).
따라서 연산량과 이동 데이터의 비율, 즉 산술 집약도(arithmetic intensity)는 프롬프트 길이에 선형 비례합니다.
이 때문에 큰 컨텍스트에서의 프리필은 연산 병목(compute-bound)이 됩니다.
디코드는 프리필 이후의 자기회귀(auto‑regressive) 루프입니다. 각 스텝마다 지금까지 구축된 전체 KV 캐시를 대상으로 어텐션을 수행해 토큰 하나를 생성합니다.
디코드에서는 행렬‑행렬 곱보다 산술 집약도가 낮은 벡터‑행렬 곱을 수행합니다.
이 때문에 디코드는 메모리 병목(memory-bound)이 됩니다.
단계를 분리하면 하드웨어 선택은 명확해집니다.
한 장치에서 프리필을 하고 다른 장치에서 디코드를 하려면, KV 캐시를 네트워크로 보내야 합니다. 가장 단순한 방식은 프리필을 실행해 끝날 때까지 기다린 뒤, KV 캐시를 전송하고, 그 다음 디코드를 시작하는 것입니다.
그림 2: 프리필(노란색) → KV 전송(초록색) → 디코드(파란색) 순으로 진행되는 단순 분할
이 방식은 두 단계 사이에 통신 비용을 추가합니다. 전송 시간이 너무 크면, 얻는 이득이 사라집니다.
KV 캐시는 마지막에 한 덩어리로 도착할 필요가 없습니다. 레이어별로 도착할 수 있습니다.
레이어 1의 프리필이 완료되는 즉시, 두 가지가 동시에 일어납니다. 레이어 1의 KV 전송이 M3 Ultra로 시작되고, 동시에 DGX Spark에서는 레이어 2의 프리필이 시작됩니다. 각 레이어의 통신은 다음 레이어들의 계산과 겹쳐집니다.
그림 3: 레이어별 파이프라인에서 프리필(노란색)과 KV 전송(초록색)이 레이어를 따라 겹치며 진행. 모든 레이어가 완료되면 즉시 디코드(파란색) 시작.
실제로 EXO는 한 레이어가 처리되는 동안 그 레이어의 KV 벡터를 전송합니다. KV 벡터는 무거운 연산 이전에 계산되기 때문입니다. 통신 오버헤드를 숨기려면, 레이어 처리 시간(t comp)이 KV 전송 시간(t send)보다 크기만 하면 됩니다.
계산 시간은 t comp = F / P 입니다. 여기서 F는 레이어당 FLOPs, P는 장치의 FLOPs/s입니다. 큰 컨텍스트에서 F는 제곱으로 증가합니다: F ∼ c₁ s². 여기서 c₁은 모델 의존 상수입니다.
전송 시간은 t send = D / B 입니다. 여기서 D는 KV 데이터의 비트 수, B는 네트워크 대역폭(비트/s)입니다. KV 캐시는 토큰당 일정 개수의 벡터를 가지므로 D ∼ q·c₂·s 입니다. 여기서 q는 양자화(4비트, 8비트 등), c₂는 모델 의존 상수입니다.
통신을 완전히 숨기려면 전송 시간이 계산 시간보다 작아야 합니다: t send < t comp. 즉 P/B < F/(q·D) ∼ (c₁/c₂)·s/q. DGX Spark가 FP16 100 TFLOPs이고 DGX Spark와 M3 Ultra 사이 링크가 10 GbE(10 Gbps)라면, 비율 P/B = 10,000 입니다. 이는 s > 10,000q/(c₁/c₂) 가 필요하다는 뜻입니다.
상수 K = c₁/c₂는 어텐션 아키텍처에 따라 달라집니다. Llama‑2 7B 같은 구형 MHA(multi-head attention) 모델의 경우 K = 2 입니다. GQA(grouped query attention)를 사용하는 모델에서는 K가 더 큽니다. Llama‑3 8B는 K = 8, Llama‑3 70B와 Qwen‑2.5 72B는 K = 16 입니다.
8비트 KV 스트리밍과 K = 16(Llama‑3 70B)일 때 임계값은 s > 5k 토큰입니다. K = 8(Llama‑3 8B)일 때는 s > 10k 토큰입니다. K = 2(Llama‑2 7B)일 때는 s > 40k 토큰입니다.
Llama‑3.1 8B(FP16)를 8,192 토큰 프롬프트로 실행하고 32 토큰을 생성:
| 구성 | 프리필 시간 | 생성 시간 | 총 시간 | 속도 향상 |
|---|---|---|---|---|
| DGX Spark | 1.47s | 2.87s | 4.34s | 1.9× |
| M3 Ultra Mac Studio | 5.57s | 0.85s | 6.42s | 1.0×(기준) |
| DGX Spark + M3 Ultra | 1.47s | 0.85s | 2.32s | 2.8× |
결합 구성은 양쪽의 장점을 모두 얻습니다. DGX Spark의 빠른 프리필(M3 Ultra 대비 3.8× 빠름)과 M3 Ultra의 빠른 생성(DGX Spark 대비 3.4× 빠름)을 결합해, M3 Ultra 단독 대비 전체적으로 2.8×의 속도 향상을 제공합니다.
분산 프리필/디코드, 레이어별 KV 스트리밍, 하드웨어 인지형 단계 배치까지 모두 EXO에서 자동화되어 있습니다.
EXO를 시작하면, 애드혹(mesh) 네트워크로 연결된 모든 장치를 자동으로 발견하고 각 장치의 연산 처리량, 메모리 대역폭, 메모리 용량, 네트워크 특성을 프로파일링합니다.
모델과 토폴로지가 주어지면, EXO는 어떤 장치가 프리필을 맡고 어떤 장치가 디코드를 맡을지, 레이어 파이프라인을 적용할지, KV를 언제 스트리밍할지, 네트워크 상태가 변할 때 어떻게 적응할지까지 계획합니다. 사용자가 스케줄을 작성하지 않습니다. 임계값을 계산하지도 않습니다. 모델을 실행하기만 하면, EXO가 이질적인(heterogeneous) 클러스터를 빠르게 만드는 방법을 알아서 찾아냅니다.
추론 성능은 더 이상 ‘한 대의 박스’가 할 수 있는 일에 제한되지 않고, 클러스터 전체가 함께 할 수 있는 일에 의해 결정됩니다.
최적화된 추론을 위해 함께 동작하는 NVIDIA DGX Spark와 Mac Studio M3 Ultra
