Token Level Routing Inference System for Edge Devices

https://arxiv.org/pdf/2504.07878

Token Level Routing Inference System for Edge Devices

1. 서론 (Introduction)

대형 언어모델(LLM, Large Language Model)은 문서 요약, 질의응답, 텍스트 생성 등 다양한 자연어 처리 과제에서 혁신적인 성능을 보이고 있다. Meta의 Llama, Google의 Gemma, DeepSeek 등 다양한 시리즈의 LLM들은 실제 서비스 및 연구에서 핵심 역할을 담당하고 있다.그러나 LLM은 수십억~수백억 개의 파라미터로 인해 막대한 연산량과 메모리 사용량을 필요로 하며, 스마트폰, 임베디드 시스템, IoT와 같은 엣지 디바이스에는 직접 적용하기가 매우 어렵다.

이런 환경 제약 속에서 소형 언어모델(SLM, Small Language Model)은 연산 속도와 자원 사용 측면에서 매우 유리해 엣지 환경에 적합하지만, 복잡하거나 맥락이 중요한 문제에서는 품질 저하 및 "환각(hallucination)" 문제(사실과 다른 정보를 생성)가 심각하게 나타난다.이처럼 효율성과 성능 사이의 트레이드오프는 실제 응용에서 큰 제약으로 작용한다.

이를 해결하기 위해 최근에는 협업 디코딩(collaborative decoding) 방식이 주목받고 있다.SLM이 기본적인 토큰 생성을 담당하고, 품질이 중요한(혹은 SLM의 신뢰도가 낮은) 일부 토큰에 대해서만 클라우드에 위치한 LLM이 결과 생성을 보조함으로써, 양쪽의 장점을 결합하려는 시도다.

이 논문에서는 엣지 디바이스에서 소형 모델이 실시간 추론을 수행하면서, 필요할 때만 클라우드의 대형 모델에게 '토큰 단위'로 질의하는 시스템을 제안한다. 실제 CommonsenseQA 데이터셋에서 0.5B 모델을 MacBook M1에서 사용할 때, 전체 토큰 중 7%만 LLM에 업로드해도 SLM 단독 대비 60% 이상의 성능 향상을 보였고, 전체 LLM을 사용하는 것 대비 80% 이상의 비용 절감 효과도 확인하였다.

1. 관련 연구 (Related Work)

2.1. SLM/LLM 트레이드오프 및 엣지 적용

• 대형 모델은 성능은 뛰어나지만, 실제 엣지/모바일 환경에서는 메모리, 연산, 전력 등 제약이 심하다. Jetson Orin Nano, ARM 프로세서, 저전력 IoT 칩 등은 대부분 LLM의 요구사항을 만족시키지 못한다.
• SLM은 빠른 추론과 저자원 사용이 가능하나, 대화의 정확성·복잡성 측면에서 한계가 명확하다.
• 이에 따라 효율성과 품질 사이의 트레이드오프를 해결하기 위해 다양한 '동적 라우팅' 연구가 진행됐다.

2.2. Token-level Routing 기반 협업 디코딩

• 최근 연구들은 입력(혹은 토큰) 단위로 라우팅을 동적으로 결정, 비용은 줄이면서 성능 손실을 최소화하는 방법을 제안한다.
• 대표적인 접근이 CITER와 CO-LLM이다.

(1) CITER (Collaborative Inference with Token-level Routing)

• SLM(빠르지만 정확도가 낮음)과 LLM(느리지만 정확도 높음) 사이에서 각 토큰마다 어떤 모델을 쓸지 결정하는 MLP 기반 라우터 사용
• 이 라우터는 MDP(Markov Decision Process) 기반의 강화학습으로, 성능과 비용의 균형을 학습한다.
• 각 상태(state)는 현재 프롬프트와 이미 생성된 토큰 시퀀스이며, 액션(action)은 SLM/LLM 중 선택이다.
• 라우팅 프리퍼런스(pairwise preference, 어떤 모델을 사용하는 것이 더 나은지)는 실제 생성 결과의 품질에 따라 지정되며, Bradley-Terry 모델과 cross-entropy loss로 라우팅 정책(policy)을 학습한다.

(2) CO-LLM (Learning to Defer and Collaborate Efficiently)

• 토큰 단위의 'defer(위임)' 정책을 학습하는 방법론.
• SLM과 LLM이 토큰별로 누가 더 잘 예측하는지를 나타내는 약한 지도 신호(weak supervision, pseudo label)를 통해 deferral 정책을 빠르게 초기화한다.
• 추론 시, 특정 확률(η)보다 deferral 신호가 높으면 SLM이 아닌 LLM으로 위임하여 토큰을 생성.
• 복잡한 추론, 도메인 특화 문제에서 좋은 성능을 보인다.

1. 방법론 및 시스템 설계 (Methodology & System Overview)

3.1. 시스템 아키텍처 및 동작 흐름

Token Routing System은 다음과 같이 세 가지 주요 모듈로 구성되며, 각 모듈은 클라우드(서버)와 엣지(단말) 환경에서 협력적으로 동작한다.

1. 서버(클라우드) LLM 서빙 모듈
   • 대규모 언어모델(예: Qwen2.5-32B)이 배포되어 있으며, SGLang 백엔드를 활용해 효율적이고 유연한 추론 요청을 처리한다.
   • 토큰 라우터가 “이 토큰은 SLM이 아니라 LLM이 생성해야 한다”고 판단한 경우, 현재까지 생성된 컨텍스트(입력 프롬프트, 생성된 토큰, 상태 정보)를 서버 LLM에 전송한다.
   • 서버는 해당 토큰을 LLM으로 생성해 반환하며, 이후 컨텍스트가 엣지 SLM에 다시 전달된다.
2. 엣지 디바이스 SLM 추론 모듈
   • 경량 소형 언어모델(SLM, 예: Qwen2.5-0.5B)을 ONNX 포맷으로 변환/배포하여 모바일·PC 등 엣지 디바이스에서 빠르게 추론할 수 있다.
   • 추론 과정에서 각 토큰별로 마지막 레이어의 hidden state를 실시간으로 산출하며, 해당 값을 라우터에 전달한다.
3. 토큰 라우팅 선택 모듈
   • SLM에서 생성된 각 토큰의 마지막 레이어 hidden state를 입력으로 받아, MLP 기반 confidence 라우터가 해당 토큰을 SLM이 생성할지, LLM으로 라우팅할지 판단한다.
   • confidence(신뢰도) 점수가 임계값(threshold)보다 낮으면, 해당 토큰은 서버 LLM으로 라우팅된다.

시스템의 주요 특성 및 처리 흐름

• 토큰 라우팅 구조의 핵심은 SLM이 문장 생성의 대부분을 담당하되, “품질에 민감하거나 불확실한 토큰”만 LLM이 담당하도록 분기시키는 것이다.
• 단일 inference 요청 동안 여러 번 prefill 및 decode가 반복될 수 있으며,LLM 호출 후에는 전체 시퀀스 컨텍스트를 다시 SLM에 전달하여 이어서 생성해야 한다.
• 기존의 ONNX/LLM serving 엔진들은 이런 토큰 단위 반복 라우팅 시나리오에 맞춰 설계되어 있지 않아,kv-cache(키-밸류 캐시) 관리가 중요하며, ONNX 그래프를 직접 수정해 마지막 hidden state를 추가로 노출시켜야 라우터와 연동이 가능하다.
• 각 토큰별로 라우팅 결정, SLM/LLM 간 컨텍스트 전환, 내부 상태값(hidden state, attention 등) 전송 등 세부적 프로토콜이 동작한다.

1. 실험 및 평가 (Experiments & Results)

4.1 시스템 성능 평가 (System Throughput)

실험 환경

• 엣지 디바이스: MacBook Pro (M1, CPU 기반, ONNX 엔진에서 SLM 및 라우터 구동)
• 서버: NVIDIA A100 2장 (Qwen2.5-32B, SGLang backend, tensor parallelism 활성화)

실험 설계 및 과정

1. 데이터셋
   • CommonsenseQA에서 100개의 객관식 문제를 무작위로 추출
2. 추론 세팅
   • 각 샘플마다 최대 100 토큰까지 생성
   • SLM(Qwen2.5-0.5B)은 엣지(로컬)에서, LLM(Qwen2.5-32B)은 클라우드(서버)에서 구동
3. 라우터 임계값(Threshold) 조정
   • MLP 라우터의 임계값을 0.4, 0.5, 0.6, 0.7, 0.72, 0.76, 0.8, 0.9로 변화시키며 실험
   • 임계값이 낮을수록 SLM이 모든 토큰을 생성
   • 임계값이 높아질수록 confidence score가 낮은(불확실한) 토큰이 LLM으로 라우팅됨
4. 성능 측정 지표
   • TTFT (Time To First Token): 최초 토큰 생성까지 걸린 시간 (SLM의 prefill 시간 반영)
   • TBT (Time Between Tokens): 토큰 간 생성 시간 (LLM 호출/통신이 늘어나면 증가)
   • Routing Number: LLM에 라우팅된 토큰 수
   • 전체 추론 소요 시간
   • 정확도(Accuracy)
5. 네트워크·시스템 조건
   • 최악의 배포 환경 시뮬레이션:
     • 클라이언트-서버 간 네트워크 지연 170ms
     • LLM 응답 대기 시간 0.9초
     • ONNX 엔진이 kv-cache 미지원 → LLM 호출 시마다 전체 시퀀스 재-prefill 필요(추가 지연 발생)

결과 요약

• 임계값이 0.3 이하면 거의 모든 토큰이 SLM에서 생성, 약 4토큰/초의 처리 속도 기록
• 임계값이 올라갈수록 일부 토큰이 LLM에 라우팅 → 네트워크, LLM 응답 등으로 latency(TBT) 가 증가
• LLM 호출이 늘어날수록 전체 소요 시간도 늘어남 (최악 환경 기준)

4.2 품질 평가 (Response Quality)

실험 구성

• SLM: Qwen2.5-0.5B (엣지/로컬), LLM: Qwen2.5-32B (서버)
• 라우터 임계값을 조정하며, 각 임계값별로 성능 평가

주요 결과 및 해석

• 임계값 0.3 이하:
  • 대부분 SLM이 토큰 생성
  • 정확도 약 50% (CommonsenseQA 랜덤 추측 기준 20%보다 월등히 높음)
• 임계값 증가
  • LLM 라우팅 비율 및 정확도 모두 상승
  • 토큰의 7%만 LLM에 라우팅해도, SLM 단독 대비 60% 이상 정확도 향상
  • LLM 전체 사용 대비 통신·추론 비용 대폭 절감
• LLM 호출이 너무 잦으면 latency가 과도하게 늘어나므로,실제로는 임계 값을 0.7~0.8 수준에서 운영하는 것이 성능·비용 trade-off에 가장 적합함.

결론 (Conclusion)

본 논문은 토큰 단위 라우팅 기반 협업 디코딩 구조를 실질적으로 엣지 디바이스-클라우드 연동 환경에서 구현했다.중요 토큰의 작은 하위 집합을 추론을 위해 클라우드의 대규모 모델로 라우팅함으로써

시스템은 낮은 추론 지연 시간을 유지하면서 에지 모델의 성능을 크게 향상시킨다. 이 아키텍처는 기기 내 배포가 필요하지만 모델 성능에 큰 영향을 미치지 않는 시나리오에 적합하다. 실험 결과, CommonsenseQA 데이터셋에서 토큰의 7%만 대형 모델로 라우팅해도 소형 모델의 정확도가 60% 이상 향상되는 것으로 나타났다.