SlimWise가 MoE 프리필은 전체로 두고 디코드 전문가만 잘라 처리량을 높인다

무엇인가

MoE 모델은 토큰마다 소수의 전문가만 활성화하지만, 연속 배칭에서는 그 희소성이 배치 단위로 유지되지 않는다. 논문의 식 (1)은 배치 크기 B에 대해 한 스텝에서 접근하는 서로 다른 전문가 수가 m[1-(1-k/m)^B]로 늘어나 결국 전체 풀 m에 수렴함을 보인다. 실제로 Qwen3.6-35B-A3B(256개 중 8개 활성화)를 배치 64로 돌리면 매 디코드 스텝이 60GiB 전문가 풀의 약 75%를 읽고, 평가한 최대 배치에서는 거의 전부를 읽는다. 문제는 이 두 단계의 성격이 다르다는 점이다. 프리필은 프롬프트 토큰을 한꺼번에 처리해 전문가 가중치 재사용이 충분한 계산 바운드 영역이라 풀을 줄여도 처리량 이득이 거의 없고, 디코드는 스텝당 시퀀스마다 토큰 하나만 나와 가중치 로딩 비용을 상각하지 못하는 메모리 바운드 영역이다. 기존 전문가 가지치기는 이 둘에 같은 축소 풀을 적용해, 이득이 없는 프리필의 품질만 깎아 왔다.

어떻게 동작하나

SlimWise는 전문가 용량을 두 단계에 따로 배분한다. 프리필은 전체 전문가 풀로 수행하고, 디코드는 중요도 기준으로 상위 m'개만 남긴 K'_l로 수행한다. 전문가 가지치기는 어텐션 구조와 KV 캐시 형식을 바꾸지 않으므로, 프리필이 만든 KV 캐시를 변환 없이 그대로 프루닝 디코더에 넘길 수 있다(학습이 필요 없는 KV cache handoff). 구현은 두 가지가 등가다. 남긴 전문가와 해당 라우터 항목만 담은 별도 프루닝 체크포인트를 쓰거나, 전체 풀을 유지한 채 라우터에 마스크를 적용한다. 식 (2)처럼 남긴 전문가의 로짓은 그대로 두고 나머지를 큰 음수 τ로 덮으면, top-k 선택과 재정규화가 K'_l 안에서만 일어나 프루닝 모델과 정확히 같은 게이팅이 된다. PD 분리 서빙에서는 디코드 인스턴스가 프루닝 체크포인트만 적재해 제거된 전문가가 쓰던 메모리를 KV 캐시로 돌릴 수 있고, PD 코로케이션 서빙에서는 phase-aware masking으로 프리필은 원본 라우터, 디코드는 마스킹 라우터를 쓴다. 전체 모델, 기존 가지치기, SlimWise가 마스킹 적용 여부만 달리해 같은 서빙 구현을 공유한다.

무엇과 다른가

KV 캐시 핸드오프로도 남는 정확도 손실과 생성 길이 왜곡을 잡기 위해 저비용 증류 단계를 둔다. 100M 토큰만 사용하고 소수 파라미터만 갱신한다. 각 학습 시퀀스에서 분할점 s를 [0,L)에서 균일하게 뽑고, 교사인 전체 모델이 접두부 [0,s)를 처리해 만든 KV 캐시를 학생인 프루닝 디코더가 이어받아 나머지 토큰을 예측한다. 손실은 식 (3)의 0.1·CE + 0.9·KD이고, KD는 교사의 top-64 토큰에 대한 KL 발산에 나머지 확률 질량을 모으는 잔여 버킷을 더한 형태다. 학습 대상은 shared expert, 라우터, 정규화 층으로 Qwen3.6-35B-A3B 전체 파라미터의 0.42%, Gemma 4-26B-A4B의 2.12%에 불과하다. 라우터 학습률은 다른 학습 파라미터의 10분의 1로 낮춘다. 원래 shared expert는 동결하고 down projection을 0으로 초기화한 두 번째 전문가 SE2를 붙여 초기 출력을 보존한 채 가산 보정을 학습시키며, 배포 시에는 중간 폭이 두 배인 단일 전문가로 합쳐 표준 체크포인트로 내보내 별도 서빙 연산이 필요 없다.

어떻게 쓰나

실험은 Qwen3.6-35B-A3B와 Gemma 4-26B-A4B-it 두 백본에서 routing mass, EAN, REAP 세 가지 프루닝 기준으로 수행했다. 수학은 GSM8K와 MATH-500, 코딩은 HumanEval+와 MBPP+, 도구 사용은 BFCLv4를 쓰고, thinking을 켠 샘플링 디코딩에 응답당 32,768 토큰 예산을 주고 3회 평균 정확도를 보고한다. 증류는 12k 토큰 이하 64k 시퀀스로 8,000 스텝, B200 4장에서 약 6시간이 걸린다. 50% 가지치기에서 학습 없는 KV 캐시 핸드오프만으로 전체 모델 대비 정확도 격차가 여러 설정에서 크게 줄었고, Gemma 4의 REAP 도구 사용(BFCL)에서 특히 두드러졌다.

전제와 한계

저자들이 강조하는 발견은 벤치마크 정확도가 가지치기의 부작용을 가린다는 것이다. REAP 기준으로 HumanEval+의 중앙 출력 길이는 2.3k에서 0.3k 토큰으로 급감하는데 정확도 하락은 완만했고, MATH-500은 반대로 중앙값이 2.4k에서 6.6k로 늘고 90퍼센타일도 크게 증가했다. KV 캐시 핸드오프는 HumanEval+ 중앙값을 약 2.3k로 되돌리지만 MATH-500은 약 5.8k로 남고 MBPP+에서는 오히려 길어진다. 그리디 디코딩에서는 핸드오프 여부와 무관하게 토큰 예산까지 자연 종료 없이 생성된다. EAN 기준 Gemma 4에서는 핸드오프 후 수학·코딩 벤치마크 중앙 길이가 예산에 도달해, 증류 없는 SlimWise가 GSM8K와 MATH-500에서 기존 가지치기보다 낮은 정확도를 기록했다. 75% 가지치기(m'=64)에서 MATH-500 90퍼센타일은 32.8k 예산에 닿았고 증류 후 14.0k로 줄었지만 베이스라인 6.1k보다는 여전히 높다. 25% 가지치기(m'=192)에서는 정확도가 거의 유지되는 대신 MBPP+ 90퍼센타일이 4.4k에서 8.1k로 늘고, SlimWise가 6.7k, 증류가 4.2k로 낮춘다. 교사 KV 캐시로 증류한 학생이 자체 프리필 캐시로만 증류한 학생보다 네 벤치마크 중 세 곳에서 정확도가 높고 꼬리 길이도 짧았다.

처리량은 PD 분리 서빙에서 50% 가지치기(m'=128) 시 전체 모델 대비 1.40~1.81배, 75% 가지치기(m'=64) 시 1.65~2.39배를 기록했고, 사용자당 30·40·50 tokens/s SLO가 엄격할수록 상대 이득이 컸다. 전문가 풀이 줄면 디코드 스텝이 짧아져 같은 SLO에서 더 큰 배치를 쓸 수 있고, 남은 메모리를 KV 캐시로 돌려 최대 배치 자체가 커진다. PD 코로케이션에서는 50%에서 1.36~1.72배, 75%에서 1.57~2.28배로 조금 낮은데, 라우터 마스킹이 디코드 스텝당 약 0.5ms(스텝 시간의 2.5~5%)를 더하기 때문이다. 프리필 처리량은 모든 설정에서 그대로다.

실무 적용을 검토한다면 세 가지를 먼저 확인해야 한다. 첫째, 프리필 비용이 줄지 않으므로 종단간 이득은 전체 서빙 시간 중 디코드가 차지하는 비율에 달려 있다. 둘째, PD 코로케이션 구성은 프리필용 전체 전문가 풀을 계속 상주시켜 모델 가중치 메모리 사용량이 줄지 않는다. 셋째, 정확도 지표만 보면 안 된다. 프루닝 기준과 벤치마크에 따라 출력 길이 분포가 양방향으로 크게 흔들리고, KV 캐시 핸드오프만으로는 과도하게 길어지는 생성을 완전히 해소하지 못해 증류가 필요한 경우가 있다. 저자들이 밝힌 한계도 이 두 가지, 즉 프리필 비용 불변과 PD 코로케이션의 메모리 절감 부재다.