MoE 추론 WideEP와 KV 캐시 인식 라우팅 — vLLM·llm-d 개발자가 짚은 스케줄링 문제

핵심 메시지

• 일반 클라우드 요청과 달리 LLM 추론은 요청마다 길이가 제각각이고 KV 캐시라는 상태를 남기기 때문에, 기존 로드밸런서 방식으로는 하드웨어 성능을 다 쓰지 못한다.
• 부하만 보는 라우팅과 캐시 위치만 보는 라우팅 모두 한쪽으로 치우치며, 둘을 함께 고려해야 대기 시간과 처리량이 동시에 산다.
• MoE 모델이 커지면서 텐서 병렬만으로는 쪼갤 수 없게 됐고, 어텐션은 데이터 병렬로 전문가 계층은 전문가 병렬로 나누는 구성이 표준이 되고 있다.
• 장비를 넓게 펼치는 이유는 전문가 연산 때문이 아니라, 랭크당 확보할 수 있는 KV 캐시 공간이 급격히 늘어나기 때문이다.
• 다중 턴 에이전트형 워크로드 실험에서 캐시를 아는 라우팅은 무작위 라우팅보다 출력 토큰 처리량을 크게 끌어올렸다.

쉽게 이해하기

파이토치 콘퍼런스 무대에 선 두 발표자는 각각 레드햇과 IBM 리서치에서 대규모 추론 서빙을 맡고 있고, 오픈소스 추론 엔진 vLLM과 그 위에서 클러스터 규모 배포를 담당하는 llm-d 프로젝트를 함께 만들고 있다. llm-d는 쿠버네티스 위에서 돌아가는 분산 LLM 추론 프레임워크로, 특정 가속기 벤더에 묶이지 않고 IBM·구글·레드햇·미스트랄·코어위브·엔비디아 등이 함께 참여하는 오픈소스 프로젝트라고 소개됐다.

발표의 출발점은 'LLM 서빙이 왜 일반 클라우드 워크로드와 다른가'였다. 전통적인 웹 요청은 대체로 균일해서 어느 서버로 보내도 비슷하게 처리된다. 반면 LLM 요청은 입력이 아주 길거나 짧고 출력도 제각각이며, 무엇보다 KV 캐시라는 상태를 남긴다. 클래식 클라우드는 상태가 없지만 AI는 상태가 있다는 것이 핵심 차이다. 그래서 어느 복제본에 어떤 캐시가 남아 있는지 모르는 채로 요청을 흩뿌리면, 이미 계산해 둔 결과를 버리고 같은 일을 다시 하게 된다.

발표자들은 여기서 두 축을 든다. 하나는 대기열 길이나 가속기 메모리 여유를 보는 부하 인식이고, 다른 하나는 이 요청의 앞부분이 어느 복제본에 캐시돼 있는지를 보는 접두 캐시 인식이다. 캐시 위치만 쫓으면 특정 노드에 쏠림이 생기고, 부하만 보면 캐시를 낭비한다. 둘을 균형 있게 결합해야 한다는 것이다. 이들이 보여준 간단한 실험에서는, 하드웨어 사양상 충분히 감당 가능한 워크로드인데도 단순한 로드밸런싱을 쓰면 첫 토큰 지연과 처리량이 무너졌고, 캐시를 아는 스케줄링을 쓰자 이론적으로 기대할 수 있는 수준에 도달했다. 발표자 스스로 수치는 참고용으로 봐 달라고 단서를 달았다.

다음 주제는 프리필과 디코드를 분리하는 구성이다. 요청이 들어오면 분리 여부를 판단해 프리필 전용 인스턴스에서 입력을 처리하고, 계산된 KV 캐시를 디코드 인스턴스로 끌어온다. 이 전송에는 RDMA 기반 전송이 쓰이고 복사 없이 캐시에서 곧바로 상대 캐시로 넣기 때문에 메모리 부담이 거의 없다고 설명됐다. 다만 분리가 항상 이득은 아니다. 사용자당 초당 토큰 같은 반응성과 GPU당 처리량을 함께 놓고 볼 때 중간 구간에서 뚜렷한 이득이 나오고, 한쪽만 극단적으로 최적화하면 오히려 합쳐 놓은 배포가 이긴다. 프리필과 디코드의 개수뿐 아니라 각각의 크기도 따로 조정해야 하며, 보통은 큰 디코드 하나에 작은 프리필 여러 개를 붙이는 형태가 잘 맞는다고 한다. 디코드 쪽이 텐서 병렬을 늘릴수록 더 잘 확장되고 KV 캐시 공간도 넓어지는데, 요청이 오래 머무는 쪽이 디코드이기 때문이다.

마지막이 본 주제인 넓은 전문가 병렬이다. 최근 큰 모델은 대부분 MoE 구조로, 토큰마다 라우터가 몇 개의 전문가를 골라 보낸다. 전문가가 잘게 쪼개질수록 성능이 좋아진다는 관찰과 공통 정보를 담는 공유 전문가 덕분에 모델은 점점 희소해졌고, 그 결과 각 전문가의 내부 차원이 너무 작아져 텐서 병렬로 쪼개는 방식이 성립하지 않게 됐다. 대신 어텐션은 데이터 병렬로 복제하고 전문가 계층은 여러 장비에 펼친 뒤 희소한 전체 교환 연산으로 토큰을 주고받는다. 여기에 특정 전문가에 부하가 몰리는 문제를 온라인으로 재조정하는 기법과, 순전파를 두 개의 마이크로배치로 쪼개 계산과 통신을 겹치는 기법이 덧붙는다. 후자는 통신 비용이 큰 장비에서 큰 효과를 내지만, 통신이 훨씬 싼 최신 장비에서는 실익이 줄어든다고 한다.

주요 인사이트

• 왜 그렇게 넓게 펼치느냐는 질문에 대한 답이 흥미롭다. 전문가 연산 자체가 아니라, 가중치를 여러 랭크로 흩어 놓을수록 랭크당 확보되는 KV 캐시 토큰 수가 선형을 넘어 늘어나기 때문이다. 특히 요청이 오래 머무는 디코드 쪽에서 이 효과가 크다.
• MoE 추론은 동시 요청 수를 굶주릴 만큼 요구한다. 발표자들은 동시 요청을 64에서 256으로 올릴 때 두 배 이상 빨라졌고, 1024 언저리에 가서야 한계에 닿았다고 밝혔다. 다만 그 지점은 실시간 서비스라기보다 배치 추론 영역에 가깝고, 동시성과 긴 컨텍스트 사이에는 맞바꿈이 있다.
• 쿠버네티스 관점에서는 새로운 형태의 배포가 등장한다. 여러 노드에 걸친 하나의 배포지만 데이터 병렬 때문에 진입점이 여러 개다. 하나의 모델인데 입구가 여럿이고, 어느 입구로 보내느냐에 따라 각 노드의 KV 캐시가 달라진다. 표준 파드 모델로는 이걸 표현하기 어려워 게이트웨이 확장 규격 자체를 손봐야 했다.
• 실험 워크로드 설계가 현실적이다. 입력 1만 토큰에 출력 2천 토큰짜리 추론형 응답, 이후 턴마다 500토큰 남짓의 짧은 입력이 다섯 턴 이어지는 형태로, 오늘날 에이전트 사용 패턴과 닮았다. 캐시를 아는 라우팅은 접두 캐시 적중률을 높였고, 무작위 라우팅에서는 프리필이 밀리면서 디코드가 일감을 못 받아 굶는 현상이 나타났다.
• 앞으로의 과제로는 AMD GPU 지원, 4분의 1초까지 걸리는 전문가 재조정 단계의 비동기화, 프리필과 디코드 인스턴스를 동시에 정해야 하는 제약 완화, 입력 인코더까지 분리하는 구성, 그리고 실행 중인 디코드 인스턴스를 늘리고 줄이는 탄력적 전문가 병렬이 언급됐다.

자주 묻는 질문

관련 AI 소식