FORGE: 희소 MoE 추론을 위한 절차적 가중치 생성

저희는 Sparse Mixture-of-Experts (MoE) 추론을 위한 절차적 가중치 생성 시스템인 FORGE (Fast On-chip Rotor Generation for Experts)를 소개합니다. 모든 전문가(expert) 가중치 텐서를 HBM에서 저장하고 로드하는 대신, 각 Streaming Multiprocessor 내부에 배치된 소형 가속기를 사용하여 투영당 $O( ext{log}^2 d)$ 개의 로터 매개변수로부터 이를 실시간으로 재생성할 수 있음을 보여줍니다. 레이어당 대체 시, 기가바이트(gigabytes)에 달하는 저장된 전문가 가중치가 1메가바이트(megabyte) 미만의 시드(seeds)로 근사화됩니다. 이는 자기회귀적 (autoregressive) Sparse MoE 디코드의 이중 메모리 병목 현상(dual memory bottleneck)을 해결합니다. 이 병목 현상은 적은 라우팅된 토큰 수에서 전문가당 FFN 연산이 HBM 대역폭에 의해 제한되고, 전체 전문가 영역(aggregate expert footprint)이 모든-대-모든 (all-to-all) 트래픽을 사용하는 다중 노드 전문가 병렬 처리(Expert Parallelism, EP)를 강제하는 경우에 발생합니다. 저희는 이 접근 방식을 Rotor Generation Unit (RGU)으로 구현했습니다. RGU는 Tensor Core 실행과 경쟁하지 않으면서 FFN 가중치 타일을 생성하는 고정 기능 데이터 경로입니다. 다섯 가지 Sparse MoE 모델에서 MoE FFN 레이어의 30%를 대체했을 때, 동일한 GPU 개수를 기준으로 시스템 처리량(system throughput)이 1.3-2.8배 높게 추정되며, 짧은 컨텍스트 길이에서는 사용자당 디코드 지연 시간(per-user decode latency)이 1.10-1.15배 낮습니다. 그 대가로, 세 가지 모델 계열에서 다운스트림 정확도(downstream accuracy)는 로터 대체 없이 미세 조정된 동일한 모델보다 떨어지며, 이 대체 비율이 이러한 트레이드오프를 제어합니다. FORGE는 또한 KV 캐시 용량을 1.2-2.6배 확장합니다. 가중치가 저장되는 비트 너비(bit width)에 구애받지 않으므로 양자화(quantization)와 결합할 수 있습니다. RGU는 SystemVerilog로 완전히 구현 및 검증되었으며, SM 면적 오버헤드는 약 0.9%입니다 (7nm 기준 $0.032 ext{ mm}^2/ ext{SM}$).