기술 요약: LYNX – 동적 배치 인식 전문가 선택을 통한 MoE 추론 효율성 실현

1. 문제 제기

현대 기초 모델들은 파라미터 수와 계산 비용을 분리하기 위해 점점 더 전문가 혼합 (Mixture-of-Experts, MoE) 아키텍처를 활용하고 있습니다. MoE 는 토큰당 소수의 파라미터만 활성화합니다 (예: Qwen2-57B 의 570 억 개 중 14 개). 그러나 이러한 효율성 약속은 지연 시간과 처리량 간의 긴장 관계로 인해 프로덕션 서빙 환경에서는 깨집니다.

높은 처리량을 달성하기 위해 추론 시스템은 요청을 배치합니다. 그러나 배치 내의 각 토큰이 독립적으로 상위-k 개의 전문가를 선택하기 때문에, 고대역폭 메모리 (HBM) 에서 로드해야 하는 전문가 집합은 배치 크기와 구성에 따라 증가합니다. 지연 시간 서비스 수준 목표 (SLO) 를 충족하기 위해 필요한 중간 규모의 배치 크기에서, 이 동적이고 데이터에 의존하는 메모리 접근 패턴은 MoE 디코딩을 메모리 대역폭 병목 상태로 만듭니다. 결과적으로 지연 시간은 전체 배치에서 활성화된 고유 전문가의 총 수에 따라 증가하며, 종종 전체 전문가 풀을 포화시킵니다 (예: 단 8 개의 다양한 요청으로 구성된 배치도 Qwen2-57B 의 64 개 전문가를 모두 활성화할 수 있음). 이는 MoE 의 희소성 이점을 무효화하여, 동적 전문가 이동의 추가 오버헤드가 있음에도 불구하고 동등한 용량의 밀집 모델과 유사한 성능을 내게 만듭니다.

정적 가지치기, 양자화, 또는 클러스터링과 같은 기존 솔루션들은 오프라인 보정에 의존하며 모델을 영구적으로 변경하므로 작업 부하 변화에 유연하지 못합니다. 동적 토큰 단위 접근 방식은 배치 수준에서 가져오는 전문가의 총합을 줄이지 않기 때문에 병목 현상을 해결하지 못합니다.

2. 방법론: LYNX

LYNX 는 모델 아키텍처를 수정하거나 오프라인 보정이 필요 없이 MoE 추론의 메모리 대역폭 병목 현상을 완화하도록 설계된 작업 부하 무관 (workload-agnostic) 시스템입니다. 이는 MoE 학습 역학에서 도출된 핵심 통찰력을 활용합니다: 로드 밸런싱 손실은 2 차 선택에 대한 신뢰도가 낮을지라도 라우터가 토큰을 전문가들 사이로 광범위하게 분배하도록 강제합니다. 이는 라우터가 진정한 토큰 - 전문가 친화도보다는 학습 제약 조건으로 인해 전문가를 선택하는 배치 수준 중복성을 생성합니다.

LYNX 는 이 중복성을 새로운 AffinityBinning 기술과 런타임에서 완전히 작동하는 3 가지 구성 요소 파이프라인을 통해 활용합니다:

2.1 핵심 원칙

이 설계는 다음 세 가지 경험적 관찰에 기반합니다:

1. 위상 민감성: MoE 는 프리필 (prefill) 단계에서는 계산 병목이지만, 디코딩 단계에서는 메모리 대역폭 병목입니다. 전문가 재매핑은 디코딩 단계에서만 유익합니다. 또한 프리필 단계는 전문가 충실도에 매우 민감한 반면, 디코딩 단계는 재할당에 대해 견고합니다.
2. 토큰 신뢰도: 라우터 로짓 (최상위 순위 전문가 로짓 간의 차이에서 유래) 에서 도출된 라우터 신뢰도 점수는 어떤 토큰 - 전문가 매핑이 중요한지 정확하게 식별합니다. 낮은 신뢰도를 가진 토큰은 정확도 저하를 최소화하면서 재매핑될 수 있습니다.
3. 전문가 순위 계층: 최상위 순위 전문가 (rank-0) 가 출력 품질을 지배합니다. 상위-k 개 중 하위 순위 전문가들은 높은 중복성을 보입니다. 재할당에 대한 민감도는 레이어마다 다릅니다.

2.2 시스템 구성 요소

LYNX 는 위상 인식 최적화기와 세 가지 경량 CUDA 커널 구성 요소를 통해 서빙 엔진 (예: vLLM) 과 통합됩니다:

• 위상 인식 최적화기: 현재 반복이 메모리 병목 (디코딩) 인지 계산 병목 (프리필) 인지 결정합니다. 불필요한 오버헤드를 피하기 위해 프리필 또는 청크형 프리필 단계에서는 LYNX 로직을 우회합니다.
• 신뢰도 분석기: 모든 레이어에서 라우터 로짓을 가로챕니다. AffinityBinning을 사용하여 각 토큰의 라우터 신뢰도를 이진 (bin) 으로 분할합니다. 이진 폭과 개수는 모델의 희소성 비율 ($k/N$) 에 의해서만 결정되므로, 시스템은 자체 보정되며 작업 부하와 무관합니다. 이는 할당된 전문가와의 친화도가 낮은 토큰을 식별합니다.
• 적응형 전문가 스코어러: 각 전문가에 대해 배치 전체 점수를 계산합니다. 이는 신뢰도 분석기가 할당한 이진 값에 기반한 지수 가중치 방식을 사용합니다. 이를 통해 고신뢰도 토큰이 선호하는 전문가들은 비례하여 훨씬 높은 점수를 받습니다. 이는 배치에 필요한 최소한의 "중요" 전문가 집합을 결정합니다.
• 전문가 재매핑기: 친화도가 낮은 토큰 - 전문가 할당을 생존한 중요 전문가 집합으로 재지시합니다. 중요한 점은 LYNX 가 상위-k 활성화 의미론을 유지한다는 것입니다. 모든 토큰은 여전히 정확히 $k$개의 전문가를 활성화하지만, 배치는 더 작고 통합된 고유 전문가 집합을 활성화합니다.

2.3 구현

LYNX 는 vLLM 위에 플러그 앤 플레이 모듈로 구현되었습니다. 핵심 로직은 라우터 출력을 가로채는 4 개의 퓨즈드 Triton/CUDA 커널로 실현됩니다. 이 설계는 중간 텐서 데이터 이동을 제거하고 커널 시작 오버헤드를 줄여, 시스템이 임계 경로에서 작동하면서도 지연 시간을 무시할 수 있을 정도로 (<4% 오버헤드) 증가시키지 않도록 보장합니다.

3. 주요 기여

1. 최초의 작업 부하 무관 시스템: LYNX 는 보정 데이터나 영구적인 모델 변경 없이 MoE 추론의 메모리 대역폭 병목 현상을 완화하는 최초의 시스템입니다.
2. AffinityBinning 기술: 아키텍처 희소성 비율에 기반하여 라우터 신뢰도를 이산화하는 새로운 방법으로, 모든 MoE 아키텍처에 대한 자동 적응을 가능하게 합니다.
3. 배치 수준 최적화: 이전의 토큰 단위 동적 접근 방식과 달리, LYNX 는 HBM 에서 가져오는 고유 전문가의 총 수를 줄이기 위해 배치 수준에서 작동하여 메모리 대역폭 병목 현상을 직접 해결합니다.
4. 상호 보완성: LYNX 는 오프로딩 및 양자화를 포함한 기존 최적화와 함께 작동하도록 설계되었습니다.

4. 실험 결과

LYNX 는 4 개 계열 (Qwen, DeepSeek, Mixtral, Llama) 의 최신 모델과 8 가지 벤치마크 (코드 생성, 수학, 추론, 비전) 에서 평가되었습니다.

• 지연 시간 감소: LYNX 는 베이스라인 대비 최대 1.30 배 낮은 지연 시간(출력 토큰당 시간) 을 달성합니다. 분리된 서빙 시나리오 (프리필과 디코딩이 분리된 경우) 에서는 감소폭이 최대 1.35 배에 달합니다.
• 정확도 영향: 시스템은 최악의 경우 1% 미만의 정확도 손실만 발생합니다. 주목할 점은 LYNX 가 평균적으로 정확도를 향상시키는 경우가 많다는 것입니다 (예: DeepSeek-Coder-V2 의 HumanEval 에서 +1.9%). 이는 학습 시간 로드 밸런싱이 LYNX 가 수정하는 비최적의 전문가 선택을 강제함을 시사합니다.
• 처리량: 고정된 SLO 예산 (예: 25ms 또는 50ms) 하에서 LYNX 는 시스템 처리량을 최대 2.1 배 향상시킵니다.
• 호환성:
  • 오프로딩: 오프로딩 기술 (예: Fiddler) 과 결합할 때, LYNX 는 성능을 31% 향상시킵니다.
  • 양자화: INT4 양자화 위에 적용될 때, LYNX 는 추가 **10%**의 속도 향상을 제공합니다.
  • 병렬화: LYNX 는 전문가 병렬화 (EP) 및 텐서 병렬화 (TP) 위에 가산적인 처리량 향상을 제공합니다.
• 견고성: 성능 향상은 다양한 배치 크기 (중간 규모에서 포화됨) 와 버스트형 도착률을 가진 실제 세계 추적 데이터에서 일관되게 나타납니다.

5. 의의

본 논문은 GPU 연산 능력이 메모리 대역폭 성장보다 계속 앞설수록 MoE 디코딩의 산술 강도 격차가 벌어져 메모리 병목 현상이 점점 더 중요해질 것이라고 주장합니다. LYNX 는 배치 수준에서 전문가가 선택되는 방식을 근본적으로 변경함으로써 이를 해결합니다. 학습 시간 로드 밸런싱으로 인해 도입된 중복성을 활용함으로써, LYNX 는 모델 재학습이나 오프라인 프로파일링 없이도 MoE 아키텍처가 미래 하드웨어에서 효율성 약속을 실현할 수 있게 합니다. 이 시스템은 동적이고 런타임 인식형 전문가 관리가 현재 MoE 서빙 패러다임에 내재된 메모리 대역폭 제약으로부터 처리량을 효과적으로 분리할 수 있음을 보여줍니다.