Google, Ironwood서 Qwen 3.5 추론 최대 4.7배 가속…JAX·Pallas 최적화 공개

Google이 14일(현지시간) 공식 개발자 블로그를 통해 Ironwood TPU에서 Qwen 3.5-397B-A17B 모델의 추론 성능을 높인 시스템 최적화 결과를 공개했다. 회사는 자체 기준 구성과 비교해 디코드 중심 작업은 약 3.1배, 프리필 중심 작업은 약 4.7배 개선됐다고 밝혔다.

이번 작업은 총 3970억개 매개변수 가운데 토큰당 170억개를 활성화하는 희소 혼합전문가(MoE) 모델을 대상으로 했다. Google은 4개의 물리적 Ironwood 칩이 8개의 논리 장치로 동작하는 단일 호스트에서 vLLM 기반 TPU 추론 엔진을 사용했다. 따라서 발표된 배수는 특정 하드웨어와 소프트웨어 구성에서 얻은 회사 자체 결과로 봐야 한다.

핵심은 모델 구조와 하드웨어 특성에 맞춘 병렬화다. Qwen 3.5의 그룹 쿼리 어텐션 계층은 KV 헤드가 2개인 반면 MoE 계층은 512개 전문가를 사용한다. Google은 이 차이 때문에 일률적인 텐서 병렬화가 비효율적이라고 보고, 어텐션에는 데이터 병렬화, 전문가 계층에는 전문가 병렬화를 적용하는 혼합 구조를 택했다.

장치 사이 통신도 직접 손봤다. 공개된 vLLM TPU 구현에는 토큰 결과를 장치별 데이터 병렬 순위로 돌려보내는 계층형 리듀스-스캐터가 포함됐다. 칩 내부와 칩 사이 전송을 두 단계로 나눠 표준 올리듀스의 통신 부담을 줄이는 방식이다.

연산 커널은 JAX의 Pallas를 이용해 하드웨어에 맞게 작성했다. 비연속적인 KV 캐시를 처리하는 배치형 Ragged Page Attention, 토큰 라우팅을 담당하는 SparseCore 커널, 1차원 합성곱과 Gated DeltaNet 블록을 합친 커널 등이 적용됐다. 중간 활성값을 메모리에 반복해서 쓰고 읽는 과정을 줄이는 데 초점을 맞췄다.

Google이 공개한 동시성 64의 8K 입력·1K 출력 조건에서 프리필 중심 처리량은 칩당 초당 3,707토큰이었다. 회사가 자체 계산한 할인 적용 이론 한계인 칩당 초당 4,500토큰의 82.4% 수준이라는 설명이다. 다만 이 값 역시 모델, 배치, 정밀도와 지연시간 조건에 따라 달라질 수 있다.

이번 공개는 대규모 희소모델의 추론 경쟁이 가속기 사양만으로 결정되지 않는다는 점을 보여준다. 모델의 어텐션·전문가 구조에 맞춘 샤딩, 통신 집합 연산, 메모리 배치와 커널 융합을 함께 설계해야 실제 처리량을 끌어낼 수 있다는 의미다. 공개 PR이 포함돼 구현 방식을 확인할 수 있다는 점도 개발자 생태계에 영향을 줄 것으로 보인다.

출처: Google Developers Blog https://developers.googleblog.com/systems-engineering-playbook-optimizing-qwen-35-397b-moe-on-ironwood-tpu7x/
출처: vLLM Project PR #2679 https://github.com/vllm-project/tpu-inference/pull/2679
출처: vLLM Project PR #3016 https://github.com/vllm-project/tpu-inference/pull/3016