AHASD: Asynchronous Heterogeneous Architecture for LLM Adaptive Drafting Speculative Decoding on Mobile Devices

"AHASD: 모바일 기기용 LLM 적응형 드래프팅 speculative decoding 을 위한 비동기 이종 아키텍처" 논문의 상세 기술 요약은 다음과 같습니다.

1. 문제 제기

모바일 기기에서의 대규모 언어 모델 (LLM) 추론은 메모리 대역폭과 연산 능력에 의해 제약받습니다. Speculative Decoding(추측적 디코딩) 은 작은 Draft Language Model(DLM) 을 사용하여 후보 토큰을 생성한 후, 이를 큰 Target Language Model(TLM) 이 배치 단위로 검증함으로써 추론을 가속화하는 기술입니다.

그러나 모바일 NPU-PIM(메모리 내 처리) 시스템에서 Adaptive Drafting(적응형 드래프팅, 맥락에 따라 드래프트 길이가 동적으로 변하는 방식) 을 구현하는 것은 두 가지 중요한 과제에 직면합니다:

1. 동기화에서의 유휴 오버헤드: 전통적인 연산자 수준의 동기식 실행은 고정된 드래프트 길이를 전제로 합니다. 적응형 알고리즘으로 인해 드래프트 길이가 변동할 때, 메모리 바운드인 PIM(드래프팅 담당) 과 연산 바운드인 NPU(검증 담당) 가 동기화를 잃게 됩니다. 이로 인해 한 장치가 다른 장치를 기다리며 유휴 상태로 머무르게 되어 상당한 동기화 오버헤드가 발생합니다.
2. 낮은 신뢰도 드래프트로 인한 연산 낭비: 작업 수준 비동기 실행에서 DLM 은 검증되지 않은 토큰을 기반으로 'look-ahead' 드래프트 생성을 계속할 수 있습니다. 맥락이 약한 경우, 이러한 드래프트는 종종 낮은 수용률을 보입니다. 시스템은 결국 TLM 이 거부할 토큰을 생성하는 데 에너지와 연산 사이클을 낭비하게 됩니다.

2. 방법론: AHASD 아키텍처

저자들은 모바일 단일 NPU-PIM 시스템을 위해 특별히 설계된 AHASD를 제안합니다. 이는 DLM 과 TLM 실행을 독립적으로 수행하도록 분리하면서 하드웨어 수준의 제어를 통해 상호작용을 관리하는 작업 수준 비동기 이종 아키텍처입니다.

핵심 아키텍처 구성 요소

• 작업 수준 비동기 실행:
  • PIM(LPDDR5): 메모리 집약적인 DLM 드래프팅 작업을 처리합니다.
  • NPU: 연산 집약적인 TLM 검증 작업을 처리합니다.
  • 통신: 데이터 전송과 작업 스케줄링을 분리하여 상호 대기 상태를 제거하기 위해 세 가지 장치 간 비동기 큐 (미검증 드래프트 큐, 피드백 큐, 사전 검증 큐) 를 사용합니다.
• 메모리 내 연산 강화:
  • Attention Algorithm Unit(AAU): LPDDR5-PIM 랭크에 직접 통합됩니다. 데이터 경로에서 Softmax, LayerNorm 과 같은 비선형 연산자 및 축소 연산을 수행하여 중간 활성화 값을 NPU 로 전송할 필요가 없게 합니다.
  • Gated Task Scheduling Unit: 특정 메모리 랭크를 게이트 제어하여 PIM 내에서 DLM 과 TLM 작업 간에 마이크로초 이하로 전환을 가능하게 합니다. 이를 통해 PIM 이 드래프팅에서 사전 검증으로 즉시 전환할 수 있습니다.

알고리즘 - 하드웨어 공동 설계 제어

적응형 드래프팅의 특정 과제를 해결하기 위해 AHASD 는 두 가지 새로운 제어 메커니즘을 도입합니다:

1. 엔트로피 - 히스토리 인식 드래프팅 제어 (EDC):

   • 목표: 낮은 신뢰도 예측에 기반한 무효한 look-ahead 드래프팅을 억제합니다.
   • 메커니즘: 드래프트 배치의 평균 softmax 엔트로피와 Leading Length(미검증 배치 수) 를 모니터링합니다.
   • 구현: 하드웨어 모듈이 Pattern History Table(PHT) 과 Local Entropy History Table(LEHT) 을 사용합니다. 엔트로피 패턴과 드래프트 수용 간의 상관관계를 학습합니다. PHT 가 낮은 수용 확률을 예측하면, 시스템은 추가 드래프팅을 중단하여 향후 거부될 가능성이 높은 토큰에 대한 연산 낭비를 방지합니다.

2. 시간 인식 사전 검증 제어 (TVC):

   • 목표: 드래프트 고갈을 유발하지 않으면서 NPU 유휴 상태를 방지하기 위해 사전 검증 타이밍을 최적화합니다.
   • 메커니즘: NPU(검증) 와 PIM(드래프팅/사전 검증) 의 런타임 지연을 모델링합니다.
   • 구현: 시스템은 NPU 의 남은 실행 가능 사이클을 예측합니다. PIM 이 NPU 의 남은 실행 창 내에서 가장 초기의 미검증 토큰에 대한 소규모 배치 사전 검증을 완료할 수 있다면, 이를 트리거합니다. 이는 NPU 가 현재 작업을 마친 직후 유효한 토큰을 즉시 사용할 수 있도록 하여 높은 활용도를 유지합니다.

3. 주요 기여

1. 작업 수준 비동기 프레임워크: DLM 을 PIM 에, TLM 을 NPU 에 매핑하여 실행을 분리하는 새로운 방식입니다. 이는 적응형 드래프팅을 위한 연산자 수준 동기식 스케줄링에 내재된 부하 불균형 문제를 극복합니다.
2. 하드웨어 인식 적응형 제어: EDC와 TVC 모듈을 설계하여 온라인 학습과 지연 모델링을 수행함으로써 드래프트 깊이와 사전 검증 타이밍을 동적으로 관리합니다. 이를 통해 낮은 신뢰도 드래프트를 효과적으로 억제하고 NPU 유휴 시간을 제거합니다.
3. 메모리 내 어텐션 처리: LPDDR5-PIM 내에 Attention Algorithm Unit(AAU) 을 통합하여 어텐션 링크를 로컬화하고 칩 간 통신 오버헤드를 줄입니다.
4. 종단 간 구현 및 평가: 시스템은 사이클 정밀 시뮬레이터 (NPU 용 ONNXim 및 PIM 용 SAITPublic-PIMSimulator 사용) 에서 구현 및 평가되어 기존 베이스라인 대비 상당한 성능 향상을 입증했습니다.

4. 실험 결과

저자들은 다양한 LLM(OPT, LLaMA2, PaLM 유사 모델) 과 적응형 드래프팅 알고리즘 (SpecDec++, SVIP, AdaEDL, BanditSpec) 을 사용하여 시뮬레이션된 모바일 플랫폼 (NVIDIA RTX 4090 Laptop GPU 베이스라인, Coral-NPU, LPDDR5-PIM) 에서 AHASD 를 평가했습니다.

• 처리량:
  • GPU 전용 대비: 최대 4.2 배 향상.
  • 최첨단 기술 (SpecPIM, GPU+PIM) 대비: 최대 1.5 배 향상.
• 에너지 효율성:
  • GPU 전용 대비: 최대 5.6 배 향상.
  • 최첨단 기술 (SpecPIM) 대비: 최대 1.24 배 향상.
• 하드웨어 오버헤드:
  • 추가 하드웨어 영역 (EDC, TVC, AAU, 큐) 은 DRAM 다이 영역의 3% 미만(구체적으로 2.68%) 으로, 모바일 배포에 매우 적합합니다.
  • 전력 소비 분석 결과, 베이스라인 NPU+LPDDR5 구성 대비 토큰당 에너지가 약 50% 감소함을 보여주었습니다.

5. 의의

AHASD 는 모바일 LLM 추론 분야에서 중요한 진전을 의미합니다. 경직된 동기식 연산자 수준 접근 방식에서 유연한 작업 수준 비동기 아키텍처로 추측적 디코딩을 전환함으로써, NPU(연산) 와 PIM(메모리 대역폭) 의 상호 보완적 강점을 완전히 활용합니다.

하드웨어 수준 학습 (EDC) 과 지연 인식 스케줄링 (TVC) 의 도입은 동적인 작업 부하 변동이 이전에는 심각한 성능 저하를 초래했던 적응형 드래프팅의 근본적인 비효율성을 해결합니다. 이 연구는 클라우드 오프로딩에 의존하지 않고도 리소스가 제한된 모바일 기기에서 고성능이며 에너지 효율적인 LLM 을 실행할 수 있는 실현 가능한 경로를 제공합니다.