TIDE: 효율적이고 손실 없는 MoE Diffusion LLM 추론을 위한 I/O-인식 Expert Offload

리소스 제약 환경에서 MoE-dLLM의 대용량 파라미터를 효율적으로 운용하려면 GPU-CPU 간의 I/O 비용과 CPU 계산 병목을 줄여야 한다. TIDE는 expert activations의 시간적 안정성을 이용해 interval-based refresh를 도입하고, 모델 정확도에 영향을 주지 않으면서 추론 속도를 높인다. LLaDA2.0-mini와 LLaDA2.0-flash에서 최대 1.4×, 1.5×의 처리량 개선이 보고된다.

단계별 디노이징으로 구성된 MoE-dLLM의 각 denoising step은 활성화되는 전문가 집합이 시간에 따라 큰 폭으로 달라지지 않는다. 즉, 인접한 denoising step 간 전문가 할당은 높은 유사성을 가지며, 5단계 수준에서도 cosine 유사도가 대체로 0.95 이상이다. 이 시간적 안정성을 활용해 interval τ 간격으로 GPU-메모리에 상주하는 전문가 세트를 재생성하고, skipped steps에서는 현 배치를 재사용한다. 이로써 GPU-메모리의 활용도와 토큰 처리율을 유지하면서 지속적인 전문가 교체에 따른 I/O 오버헤드를 줄일 수 있다. 최적의 τ는 LatI/O(τ)와 LatCPU(τ)의 합을 최소화하도록 MP를 통해 오프라인으로 결정한다. 비동기 토큰 라우팅과 hit 기반 재배치를 도입해 CPU와 GPU 간 병렬화를 극대화한다.

• 문제 정의: MoE 모델은 L개의 sparse FFN 계층과 총 E개의 전문가를 갖고, 한 토큰당 활성화되는 k명의 전문가가 존재한다. GPU 메모리에는 B명의 GPU experts가 상주하고, 나머지는 host 메모리에 있다. LatFFN은 LatGPU(KGPU)와 LatCPU(KCPU)의 최대값 혹은 합으로 표현된다. K와 B의 관계에 따라 LatFFN이 달라진다.
• 관찰: MoE-dLLM에서 단계 간 expert activation은 시간적으로 안정적이며, 인접 단계 간의 유사도가 높다(dt ≈ 상수). 이로 인해 특정 Window 내에서 expert 배치를 재사용할 수 있다.
• TIDE 설계: decoding을 τ 구간으로 나누어 refresh steps에서 GPU-resident expert를 top-B로 재배치하고, skipped steps에서는 현재 배치를 유지한다. 최적의 τ를 MP로 결정하고 offline으로 해결한다. Algorithm 1은 GPU/CPU expert 간 비동기 마이그레이션과 토큰 라우팅 흐름을 정의한다.
• 최적화 문제: LatI/O(τ) ≈ CI/O · B · T/τ · [1 − (1 − d)τ], LatCPU(τ) ≈ CCU · T · B · f(τ) 이고, 총 Lat(τ) = LatI/O(τ) + LatCPU(τ)이다. 이를 최소화하는 τ를 구한다(식 7).
• 구현: CPU 속도·I/O 대역폭 프로파일링으로 상수 CI/O, CCU를 추정하고 Greedy 탐색으로 최적 τ를 구한다. 모든 내용은 inference 과정에서 오버헤드 없이 offline으로 수행된다.
• 실행 정책: 토큰이 미스일 경우 비동기로 CPU에서 처리하도록 하고, 결과를 동기화한다. Lossless하게 동작하며 출력은 기존 GPU-only 실행과 동일하다.

• 메인 벤치마크: Table 1에 따라 LLaDA2.0-mini에서 256-64-10GB 구성은 Fiddler 1.81, Mixtral-Offload 1.69, TIDE 2.11 token/s를 달성. 128-18GB 구성은 1.74, 1.76, 2.36; 1024-64-10GB 구성은 1.79, 1.45, 1.89; 128-18GB 구성은 1.80, 1.91, 2.44. LLaDA2.0-flash에선 256-32-30GB에서 0.95, 1.01, 1.25; 64-55GB에서 1.14, 1.35, 1.73; 1024-32-30GB에서 0.89, 1.01, 1.24; 64-55GB에서 1.05, 1.16, 1.45가 관측되었다.
• τ에 따른 Ablation: Table 2는 τ=1, Random τ, Optimal τ에 대한 각 조합별 처리량을 제시한다. LLaDA2.0-mini의 (32,64)에서 1.79(τ=1), 1.62(Random), 1.89(Optimal); (32,128)에서 1.91, 2.14, 2.44; (64,64)에서 1.71, 1.65, 1.90; (64,128)에서 2.14, 2.24, 2.40이다. LLaDA2.0-flash의 경우 (32,32)에서 1.01, 0.85, 1.24; (64,32)에서 1.16, 1.23, 1.45; (64,64)에서 1.12, 0.95, 1.32; (64,64)에서 1.27, 1.32, 1.49이다.
• 추가 분석: Figures 4–5에서 τ 증가에 따른 Miss Rate 증가와 CPU-GPU 마이그레이션 비용의 상호 작용을 시각화한다. 최적 τ는 3~6 구간에서 미세 조정되며, τ가 커질수록 I/O 비용의 1/τ 감소 효과가 CPU 비용 증가를 상쇄한다. 실험적으로 τ*를 오프라인으로 결정하는 MP 접근법이 Random τ에 비해 최대 1.4×의 속도up를 제공한다. block size 및 GPU 예산에 따라 TIDE의 이득은 일관되게 유지된다.

1. Expert Activation 패턴 분석은 블록 내부 수준에 국한된다. 블록 단위 이상의 전역적 유사성 분석은 아직 미수집이다. 2) 평가가 한정된 하드웨어 플랫폼에서 수행되었다. AMD GPU와 ARM CPU 등 다양한 아키텍처에서의 일반화는 필요하다. 3) 현재 설정은 단일 GPU-CPU 시스템에 국한되며 다중 GPU/다중 노드 확장은 향후 연구과제다.