LLM 추론에서 TP·CP·EP 병렬화를 적용한 방법

요약

Meta가 LLM 실시간 서빙(TTFT 350ms, TTIT 25ms 목표)을 위해 세 병렬화를 조합. Tensor parallelism은 DDA allreduce로 통신 레이턴시 O(N)→O(1), Context parallelism은 ring attention(Pass-KV/Pass-Q)로 1M 토큰을 H100 단일 호스트 1분 이내·10M 토큰을 32대 H100 1분 이내, Expert parallelism은 MoE의 two-shot all-to-all로 전문가 분산. AMD MI300X에서 RCCL 대비 decode 1050%·prefill 1030% 개선.

내용

LLM 추론은 두 단계로 나뉨. prefill은 입력 프롬프트(수천 토큰까지)를 처리해 transformer 레이어별 KV 캐시를 만드는 단계로, attention 연산이 시퀀스 길이에 제곱으로 늘어 compute-bound. decode는 KV 캐시를 갱신하면서 토큰을 한 개씩 생성하는 단계로, 모델 가중치와 KV 캐시 메모리 I/O가 attention 시간을 압도해 memory-bandwidth-bound. 두 단계의 병목 자체가 다르다는 점이 출발점.

목표 지표는 TTFT(Time-to-first-token) 350ms 이내, TTIT(Time-to-incremental-token) 25ms 이내. 모델이 단일 GPU에 안 들어갈 정도로 커지고 컨텍스트가 1M·10M 토큰까지 늘면서, 디바이스 다수에 걸쳐 어떻게 쪼갤지(병렬화)와 그 사이의 통신을 어떻게 줄일지가 서빙 성능을 결정. 본문은 Meta가 운용 중인 세 갈래 병렬화와 그 통신 최적화 기법 설명.

해결 / 접근

Tensor Parallelism (TP) — 모델 레이어 분할 + DDA

• attention 블록·MLP 레이어를 작은 독립 블록으로 샤딩해 여러 디바이스에 배치 → 단일 디바이스로 안 되는 throughput 확보
• 병목: allreduce 통신이 E2E 레이턴시의 최대 30% 차지
• DDA(direct data access) flat 알고리즘 — 각 rank가 다른 rank의 메모리를 직접 읽어 로컬 reduce, 데이터 교환량을 O(n)→O(n²)로 늘리는 대가로 레이턴시를 O(N)→O(1)
• DDA tree 알고리즘 — allreduce를 reduce-scatter + all-gather 두 단계로 쪼개고 각 단계에 direct data access 적용, ring과 같은 데이터 이동량으로 상수 레이턴시 (조금 큰 메시지 사이즈에 적합)
• AMD MI300X에서 NVIDIA H100과 전체 성능 동등, RCCL 대비 decode(작은 메시지) 1050%·prefill 1030% 개선, TTIT 약 10% 감소

Context Parallelism (CP) — 장문 컨텍스트 분산 (ring attention)

• 장문 추론의 세 부담: dense attention FLOPs 제곱 증가(compute) / KV 캐시 컨텍스트에 선형 증가(memory) / 호스트 다수 병렬화 시 통신 레이턴시 증가(communication)
• attention 모듈에 ring attention 두 변종 구현
  • Pass-KV: 입력 토큰을 CP rank로 쪼개고 각 rank가 자기 몫의 Q·K·V를 계산, 이후 K·V 텐서를 rank 간 교환
  • Pass-Q: 동일하지만 Q 텐서를 교환
• fast attention 커널과 결합한 결과
  • H100 단일 호스트에서 1M 토큰 1분 이내
  • H100 32대 분산 추론에서 10M 토큰 1분 이내
  • Llama 3 405B 기준 16노드 CP에서 128K 토큰 prefill 3.8초, 1M 토큰 prefill 77초 (거의 선형 스케일링)

Expert Parallelism (EP) — MoE 전문가 분산

• MoE의 expert 수가 단일 호스트에 다 못 들어가는 환경 대상
• two-shot all-to-all 통신으로 라우팅에 따라 DP rank ↔ EP rank 사이에서 토큰 교환
• all-to-all이 E2E 레이턴시의 1030% 차지, 특히 decode 메시지(100KB2MB)에서 부담
• 최적화 방향
  • Dynamic all-to-all: 데이터를 서브 청크로 쪼개 원격 이웃으로 전송
  • Persistent all-to-all: 메모리 핸들 교환·네트워크 로드 밸런싱·CPU 오버헤드로 인한 슬로다운 직접 처치

결과 / 참고

• TTFT 350ms·TTIT 25ms 목표 아래에서 세 병렬화 + 통신 최적화로 장문·MoE까지 서빙 가능 영역 확장
• DDA: AMD MI300X RCCL 대비 decode 1050% / prefill 1030%, TTIT 10% 감소
• CP: H100 단일 1M 토큰 1분 이내 / H100 32대 10M 토큰 1분 이내 / Llama 3 405B 128K 3.8초·1M 77초
• 다음 방향: N-D parallelism(CP·PP·EP·TP를 노드 간에 + 별도 DP) + disaggregated inference로 prefill/decode 티어 분리 — compute-heavy 하드웨어를 prefill에, memory-bandwidth-heavy 하드웨어를 decode에 이종 배치
• 관련 기술 과제: cloud fabric 설계, fused kernel(통신을 연산 커널에 통합), device-initiated kernel(CPU 오버헤드 축소)
• 출처: Meta Engineering 블로그, ring attention 논문(arxiv:2411.01783)