Meta의 추천 시스템 추론 최적화

Meta는 추천 시스템 추론에서 요청별로 같은 사용자 임베딩을 후보마다 복제하던 병목을 없애기 위해 In-Kernel Broadcast Optimization (IKBO) 를 도입했다. 기존 방식에서는 후보 수가 늘수록 복제 비용이 선형으로 커진다. 핵심은 브로드캐스트를 텐서 복제가 아니라 데이터 레이아웃 문제로 보고, 커널 내부에서 직접 흡수하는 것이다.

IKBO는 요청 전체에 공통인 Request-Only (RO) 임베딩과 후보별 Non-Request-Only (NRO) 임베딩을 분리한다. 추론 런타임은 복제 텐서를 만들지 않고 후보-사용자 매핑만 넘기며, 상호작용 연산이 필요한 지점에서만 두 배치 크기를 그대로 처리한다. 기존의 system-level broadcast나 ROO 같은 net-splitting이 복제를 우회하는 데 그쳤다면, IKBO는 기본 연산 계층에서 broadcast 자체를 없앤다.

최적화 유형은 두 가지다.

• Type I: 분해 가능한 연산은 RO와 NRO를 먼저 독립적으로 계산하고 마지막에만 결합해 메모리 대역폭과 연산을 줄인다.
• Type II: 커널 내부에서 broadcast를 처리해 불필요한 데이터 이동을 제거하고 IO 병목을 줄인다.

이 적용은 kernels, compilation specification, inference runtime의 세 층을 모두 건드린다. ML compiler는 연산자별 동적 shape 범위를 알아야 맞는 커널을 고를 수 있다. 모델 정의에 IKBO 커널을 직접 넣는 direct adoption 과, inference 단계에서 기존 연산자를 IKBO 버전으로 바꾸는 inference-time transformation 이 모두 가능하며, 후보/사용자 비율이 1보다 큰 training에서도 같은 커널을 그대로 쓸 수 있다. 이 구조는 초기 랭킹부터 후기 랭킹까지 GPU와 MTIA 전반에 배포됐고, 코디자인 모델에서는 계산 집약적 net latency를 최대 2/3 줄였다. 또한 LLM-scale 모델을 광고에 서빙하는 Meta Adaptive Ranking Model의 확장성 기반으로도 쓰인다.

커널 심층 사례 1인 Linear Compression Embedding (LCE) 에서는 동일한 사용자 임베딩을 후보마다 복제하지 않고 K축을 기준으로 사용자와 후보를 분리한 뒤, 두 개의 GEMM을 각자 자연스러운 배치 크기로 실행하고 작은 압축 결과만 broadcast한다. 후보/사용자 비율이 약 70:1인 설정에서 사용자 배치는 1024 -> 약 15로 줄었고, 실행 시간은 1.944ms -> 1.389ms로 28.5% 감소했다.

이후 메모리 정렬 문제를 해결하기 위해 K를 8의 배수로 패딩해 cp.async가 4바이트 단위로 쪼개지던 비효율을 없앴다. H100 SXM5에서 네 단계의 progressive co-design을 거치며 ~4x 속도를 냈고, 마지막에는 TLX를 활용한 warp-specialized multi-stage fusion으로 마무리됐다.

커널 심층 사례 2인 Flash Attention에서는 IKBO가 커널을 IO-bound에서 compute-bound로 전환했고 621 BF16 TFLOPs를 기록했다. 비코디자인 CuTeDSL FA4 Hopper 기준으로는 커널만 비교해 2.4x, broadcast까지 포함한 전체 기준으로 6.4x 처리량 향상을 냈다.