M5 Ultra 256 GB와 M3 Ultra 512 GB를 결합하여 디분산된 prefill/decode 구현하기

제가 해결하려는 문제는 이렇습니다. 저는 M3 Ultra 512 GB를 가지고 있는데, 여전히 메모리가 이 기기를 소유하는 이유이지만, 긴 프롬프트에 대한 prefill 속도가 느립니다. 17.5k 토큰짜리 프롬프트를 처리하는 데 첫 번째 토큰이 나오기까지 41초가 걸립니다. M5 Ultra는 prefill을 개선하지만 현재 최대 256 GB로 출시되었고, M5 Ultra 512 GB는 아직 존재하지 않습니다. 따라서 M3를 교체하기보다는, M5 Ultra 256 GB를 추가하여 작업을 분할했습니다. 즉, M5가 prefill을 담당하고, 캐시 상태는 Thunderbolt 5 RDMA를 통해 전송되며, M3가 자신의 512 GB 풀을 사용하여 decode를 수행합니다.

단일 노드 수치부터 보겠습니다. GLM-5.3-Flash (4비트, MTP)를 사용하고, 프롬프트는 17.5k 토큰, 출력 토큰은 256개로 설정했습니다.

M5는 prefill에서는 약 2.4배 빠르지만, decode에서는 약 1.4배 정도만 빠릅니다. 그리고 decode가 더 큰 메모리(더 많은 컨텍스트, 더 많은 동시 대화)의 이점을 얻는 단계입니다. 이러한 비대칭성이 바로 분할 작업의 핵심 논거입니다.

M5의 행렬 유닛이 대략 5배인데 왜 '단지' 2.4배에 그칠까요? 제가 프로파일링해 봤습니다. MLX에서 순수 matmul을 측정했을 때: bf16은 118 vs 24 TFLOPS, 4비트는 108 vs 24입니다 (양자화된 커널이 bf16으로 역양자화되어 동일한 유닛에서 실행되므로; 8비트가 더 빠르지는 않습니다). 하지만 GLM-5.3-Flash의 prefill은 matmul의 절반 정도에 불과합니다.

구성 요소별로, 2,048 토큰 청크로 8,192 토큰을 처리할 때:

실제 matmul이 필요한 부분은 34배의 속도 향상을 얻습니다. 반면, 그렇지 않은 부분(희소 어텐션의 per-query top-2048 gather 커널, 재귀 스캔, 전문가 라우팅)은 1.31.8배 정도에 그치는데, 이는 대략 메모리 대역폭 비율(558 vs 382 GB/s 스트리밍)과 유사합니다.

M5는 전처리 과정(prefill)의 거의 절반을 매트릭스 유닛을 전혀 사용할 수 없는 커널에서 소비합니다. 이로 인해 두 가지 개선점이 도출되었습니다. 첫째, 더 큰 전처리 청크: 8,192 토큰 청크에서는 M5가 1,188 tok/s를 기록하며 (+14%, MoE 전문가 매트멀(matmuls)이 대역폭 제한을 벗어날 만큼 충분한 토큰을 얻게 됩니다). 반면 M3는 4%의 성능 향상을 보였습니다. 둘째, 희소 어텐션 커널: 이 커널은 스레드 그룹당 하나의 쿼리를 처리하고 해당 쿼리의 2,048개의 키를 스칼라 내적(scalar dot products)으로 수집합니다. 따라서 주변 쿼리 간에 수집된 키를 공유하고 매트릭스 유닛을 사용하는 블록-희소 버전에서 다음 큰 M5의 이점이 나타납니다. 모든 사람이 걱정하는 선형 어텐션은 고작 6%입니다. 결과를 분산 처리한 경우, 동일한 프롬프트로 전처리(prefill)를 M5에서 수행하면 988 tok/s가 나오고, 캐시를 M3로 핸드오프(hand-off)하는 데는 두 TB5 케이블을 통해 총 38ms 동안 358MB의 데이터가 전송됩니다. 이로써 종단 간(end to end) 첫 토큰 생성까지 걸리는 시간은 M3에서만 디코딩할 때 약 40초가 걸리던 것과 비교하여 약 18.5초로 단축됩니다. 이후 M3는 36.8 tok/s로 디코딩하며, 출력 결과는 M5 자체의 성능과 동일합니다. M5는 요청 후 18초 만에 다시 자유로워져 다음 프롬프트를 전처리할 수 있고, 그동안 M3가 디코딩을 수행합니다. 이 모델에서 핸드오프가 저렴한 이유: GLM-5.3-Flash는 고정 크기 상태를 가진 34개의 선형 어텐션(KDA) 레이어와 총 20KB/토큰에 불과한 11개의 희소 어텐션 레이어를 가지고 있습니다. 이 속도로 M3의 여유 공간인 약 300GB는 최대 1,500만 개의 캐시된 대화 토큰을 저장할 수 있습니다. 스택: MLX의 JACCL 백엔드(Thunderbolt를 통한 RDMA, 별도의 컨트롤러에 연결된 두 케이블, 합산 약 150 Gb/s), 모델용 mlx-vlm 0.7.6 및 캐시 객체를 세 가지 dtype으로 패킹된 배열로 직렬화하여 mx.distributed.send/recv를 통해 전송하는 Python 프로토타입인 MTP drafter입니다. 제한 사항: 전처리 노드는 전체 모델을 필요로 하므로, 이 방식은 256GB에 맞는 모델에만 적용됩니다. 그보다 큰 모델의 경우 여전히 512GB 단독 사용 또는 실제 512GB M5가 필요합니다. 요청 하나당 디코딩 속도는 M5 단독이 더 빠릅니다 (47 vs 37 tok/s).

긴 프롬프트에서의 첫 토큰 시간, 요청이 진행 중일 때의 처리량(throughput), 그리고 512 GB 풀을 유지하는 것이 이점입니다. 제가 시간을 절약한 것들 중에서 여러분에게도 도움이 될 만한 점들이 있습니다: Apple의 Thunderbolt RDMA는 신뢰할 수 없는 연결된 큐 쌍(queue pairs)만 제공하지만, 전송(send)은 수신자가 버퍼를 게시(post)할 때까지 단순히 차단됩니다. 흐름 제어(flow control)가 필요하지 않습니다. 상위 JACCL은 완료 상태를 절대 확인하지 않으며, 송신자/수신자 크기 불일치는 memmove에서 SIGBUS로 이어집니다. Mac GPU는 약 17초 동안 유휴 상태였다가 첫 작업 후에는 약 1.25초가 걸립니다. 불안정한 링크처럼 보였지만 그렇지 않았습니다. Keep-alive 스레드가 이를 해결합니다. mx.set_wired_limit을 권장 작동 집합(working set)으로 설정하는 것이 필수적이며, 그렇지 않으면 모델 페이징 및 디코딩 속도가 9~17 토큰/초로 떨어집니다. 다른 분들도 이런 방식으로 Mac을 조합하고 계신가요, 아니면 이 모델군에서 prefill/decode 분할 작업을 exo나 mlx-lm으로 진행하고 계신가요? 중복 작업을 하는 것보다 경험을 공유하는 것이 낫겠습니다. /u/Usual_Tackle5892 제출 [링크] [댓글]