Skip to main content
공개 노트

mKernel: 통신을 기다리지 않는 GPU 커널

Views 3약 4분생성형 AI 활용
Summary

mKernel: 통신을 기다리지 않는 GPU 커널

GPU 클러스터의 병목은 더 이상 계산 성능만으로 설명되지 않습니다. GPU가 빨라질수록 계산과 통신이 번갈아 멈추는 비용이 더 크게 드러납니다. mKernel은 이 문제를 “통신을 더 빠르게 호출하는” 방식이 아니라, 계산과 통신을 같은 CUDA 커널 안에서 함께 다루는 방식으로 접근합니다.

mKernel은 무엇인가

mKernel은 UC Berkeley UCCL 팀이 공개한 다중 GPU·다중 노드용 fused kernel 라이브러리입니다. 정확히는 소규모의 목적 특화 persistent CUDA 커널 라이브러리에 가깝습니다.1

핵심 아이디어는 노드 내부 NVLink 통신, 노드 간 RDMA 통신, dense compute를 하나의 persistent CUDA 커널 안에 배치하는 것입니다.1 persistent 커널은 GPU에 오래 상주하며 작업을 계속 처리하는 커널을 뜻합니다. CPU가 계산 커널과 통신 호출을 단계별로 지시하는 구조보다, GPU 내부에서 더 촘촘하게 파이프라인을 만들 수 있습니다.

무엇이 달라지는가: 통신이 커널 내부로 들어온다

일반적인 GPU 분산 실행에서는 계산과 통신이 별도의 단계처럼 다뤄지는 경우가 많습니다. 예를 들어 먼저 데이터를 모으고, 그다음 GEMM을 수행하고, 다시 결과를 줄이는 식입니다. 이때 각 단계 사이에는 대기 시간이 생깁니다.

mKernel은 이 경계를 줄이려 합니다. NVLink와 RDMA 통신을 같은 persistent 커널 내부에 두고, 타일 또는 청크 단위로 계산과 통신을 겹쳐 실행합니다.1 즉, 필요한 데이터 조각이 도착하면 전체 통신이 끝나기를 기다리지 않고 해당 조각을 바로 계산에 사용할 수 있습니다.

공개된 커널 예시도 이 방향을 보여줍니다. AllGather+GEMM, GEMM+AllReduce, MoE Dispatch+GEMM, Ring Attention, GEMM+ReduceScatter처럼 계산과 통신이 강하게 연결된 패턴을 대상으로 합니다.2

GPU가 직접 RDMA를 진행한다

mKernel의 또 다른 특징은 NCCL이나 NVSHMEM에 의존하지 않고 GPU 주도 RDMA write를 사용한다는 점입니다.1 네트워킹 구현은 libibverbs 기반이며, 통신 백엔드를 처음부터 작성해 이기종 네트워킹 장치에 대응하도록 설계됐습니다.1

이 선택은 mKernel의 성격을 잘 보여줍니다. 통신을 외부 라이브러리에 맡기는 별도 단계로 두기보다, 커널 내부 스케줄링의 일부로 다루겠다는 것입니다. 그래서 mKernel의 핵심은 “어떤 통신 API를 쓰는가”보다 “계산과 통신을 얼마나 잘게 쪼개 겹칠 수 있는가”에 있습니다.

SM specialization: 일부 SM은 계산이 아니라 통신을 맡는다

mKernel은 SM specialization을 사용합니다. CTA가 compute, intra-comm, inter-send, inter-reduce 같은 역할을 자체적으로 나눠 맡고, 각 역할에 배정되는 SM 수는 shape별로 조정할 수 있습니다.1

이 구조는 중요한 판단 기준을 던집니다. 모든 SM을 계산에만 쓰는 것이 항상 최선은 아닐 수 있습니다. 일부 SM이 통신 진행을 맡더라도 전체 파이프라인의 대기 시간이 줄어든다면, end-to-end 성능은 더 좋아질 수 있습니다.

특히 대규모 AI 워크로드에서는 노드 내부 GPU 간 통신과 노드 간 통신이 모두 병목이 될 수 있습니다. mKernel의 intra-kernel overlap은 이 두 영역을 모두 대상으로 합니다.1

mKernel을 볼 때의 실용적 기준

mKernel은 범용 통신 라이브러리의 단순 대체재라기보다, 계산과 통신이 밀접하게 결합된 특정 패턴을 위한 목적 특화 라이브러리로 보는 편이 정확합니다. 따라서 평가할 때도 단순히 “얼마나 빠른가”보다 다음 질문이 더 중요합니다.

  • 내 워크로드가 타일 또는 청크 단위로 계산과 통신을 겹칠 수 있는가

  • NVLink 기반 노드 내부 통신과 RDMA 기반 노드 간 통신이 함께 병목이 되는가

  • 일부 SM을 통신 역할에 배정했을 때 전체 파이프라인 대기가 줄어드는가

  • 사용하는 모델 패턴이 AllGather+GEMM, GEMM+AllReduce, MoE Dispatch+GEMM, Ring Attention, GEMM+ReduceScatter 같은 fused kernel 형태와 맞는가2

mKernel이 보여주는 방향은 분명합니다. GPU 클러스터 성능 최적화의 초점이 “계산 커널을 빠르게 만들기”에서 “계산과 통신을 하나의 실행 파이프라인으로 설계하기”로 이동하고 있다는 것입니다.

1mKernel - UC Berkeley Sky Computing Lab

2GitHub - uccl-project/mKernel

이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.

틸노트는 사용자가 AI를 활용해 노트를 쉽게 작성하거나 편집할 수 있는 서비스입니다.

AI 스퀘어 단톡방에 함께해요

AI 활용과 틸노트 사용 팁을 나누는 오픈채팅방입니다. 입장 비밀번호는 tilnote입니다.

참여하기