Low Latency Communication Kernels with NVSHMEM

대부분의 multi-GPU 학습/추론 코드는 NCCL 을 통한다 — torch.distributed.all_reduce, broadcast, all_to_all. NCCL 은 잘 도는 추상이지만 — 한 가지 자리에서 일관되게 부족하다. 작은 메시지의 latency. 그 자리를 NVSHMEM 이 메운다.

강의 자막이 실패해 발표자의 정확한 워딩은 못 잡았지만, distributed systems 영역에서 NVSHMEM 의 입장은 일관된다.

1. NCCL 은 collective 단위 추상 — AllReduce 한 번이 한 launch. CPU 가 launch 하고, GPU 가 받고, 끝나면 CPU 에 신호. 한 step의 collective 에는 좋다.
2. 그런데 collective 호출이 50개씩 쌓이면 — 각 호출마다 launch overhead 와 host-device 동기화 비용이 누적. 학습 한 step 안에서 작은 collective 가 많을수록 성능이 깎인다.
3. NVSHMEM 은 한 단계 아래 — 커널 안에서 직접 remote GPU 의 메모리에 put/get. CPU 가 다음 launch 를 안 기다려도 되는 구조. 커널 fusion 으로 통신을 같이 묶을 수 있다.

NCCL 이 잘 받쳐주는 메시지 크기는 보통 수 MB 이상이다. 그 아래의 자리는 — production 영역에서도 — 자주 등장한다.

이 자리들의 공통점 — 한 step 안에 통신 호출이 많거나, 한 호출의 메시지가 작거나, 통신과 compute 가 같은 커널 안에서 묶여야. NCCL 의 launch-per-collective 모델이 자연스럽지 않은 자리들.

NVSHMEM 의 모델은 PGAS (Partitioned Global Address Space). 모든 GPU 가 “symmetric heap” 이라는 공유 주소 공간을 가지고, 각자 자기 부분에만 직접 메모리를 쓸 수 있고, 다른 GPU 의 부분에는 put / get 으로만 접근.

# CPU 가 launch
ncclAllReduce(send, recv, count, ...)
# 모든 GPU 가 같은 collective 에 참여
# CPU 가 끝을 기다리거나 stream sync

// 커널 안에서 직접
__global__ void my_kernel(...) {
  // 다른 GPU 의 symmetric heap 위치에
  // 직접 데이터를 쓴다
  nvshmem_float_put(
    dest,     // remote address
    src,      // local data
    nelems,
    target_pe // PE = processing element (GPU)
  );
}

이 차이가 만드는 효과는 세 가지로 정리된다.

• launch overhead 제거 — 한 통신당 ~10 us 의 cudaLaunchKernel 비용이 사라짐. 작은 메시지에서 결정적.
• compute-comm overlap 의 자연스러움 — 같은 커널 안에서 한 thread 는 compute, 다른 thread 는 put. CUDA stream 차원의 overlap 보다 fine-grained.
• persistent kernel 의 가능성 — 한 커널이 한 번만 launch 되고 안에서 여러 step 을 돈다. CPU 동기화가 step 별로 안 필요.

강의에서 발표자가 보여줬을 가장 핵심 그래프 — 같은 hardware, 같은 message, NCCL vs NVSHMEM 의 latency 비교. 정확한 수치는 시점/구현에 민감 (확인 필요), 패턴은 일관됨.

이 패턴이 가리키는 의사결정 — “메시지 크기가 1 MB 이하이거나 한 step 안에 collective 호출이 많을 때 NVSHMEM”. 두 조건 모두 LLM 학습/추론에서 자주 등장한다.

NVSHMEM 위에서 AllReduce 를 짜는 방법은 여러 가지다. 메시지 크기에 따라 다른 알고리즘이 자연스럽다.

NVSHMEM 의 가치 — 같은 코드 안에서 알고리즘을 바꾸기 쉽다. one-shot AllReduce 한 번 짜본 후 ring 으로 변경하는 게 — NCCL 의 wrapping 보다 — primitives 차원에서 자연스럽다.

이 통합이 NVSHMEM 의 큰 가치 — 코드 수준에서는 모든 PE 가 평등해 보인다. 실제로는 NVLink 의 GPU pair 와 IB 너머의 GPU pair 가 latency 차이가 큰데, 같은 API 로 다룬다. 알고리즘 선택은 사용자가 “이 PE 들이 어디 있는지” 의 hint 로 가능.

• DeepSeek DeepEP — MoE 모델의 expert all-to-all. NVSHMEM 위에 fine-grained dispatch / combine 커널. NCCL 의 latency 한계를 넘은 가장 분명한 사례.
• Megatron-LM 의 일부 영역 — sequence parallel 의 small all-gather, tensor parallel 의 일부 영역에서 — NVSHMEM 백엔드 옵션 — 확인 필요.
• vLLM / SGLang 의 추론 collective — TP 추론의 작은 AllReduce. 한 step 안에 호출 횟수가 많아 NVSHMEM 채택의 자연스러운 자리.
• NVIDIA NVSHMEM-X — NCCL 과 NVSHMEM 의 통합 시도. 확인 필요 — release 시점에 따라 다름.
• 학계의 communication-compute fusion 연구 — 한 커널 안에서 GEMM 과 AllReduce 를 묶는 연구. NVSHMEM 의 커널 안 통신이 핵심.

강의에서 Singhania 가 강조했을 입장 — NVSHMEM 은 일반 사용자가 매일 만지는 도구는 아니다. NCCL 이 잘 도는 자리에서는 NCCL 이 답. 그러나 NCCL 이 못 받쳐주는 자리에서는 NVSHMEM 외 대안이 거의 없다.

손에 새기기 — 실습 시퀀스

1. NVSHMEM hello world — 2-GPU 환경에서 PE 0 이 PE 1 의 메모리에 put 한 번. 실제로 도는지 검증.
2. latency 측정 — 1 KB ~ 16 MB 의 메시지 크기에 대해 NCCL 과 NVSHMEM 의 AllReduce latency 측정. cross-over point 직접 확인.
3. one-shot AllReduce 짜기 — § 06 의 one-shot 알고리즘을 NVSHMEM 으로 직접 구현. ring 과 비교.
4. compute-comm fusion — GEMM 의 일부 결과를 같은 커널 안에서 다른 GPU 로 put. 시간 비교.
5. DeepEP 코드 읽기 — DeepSeek 의 DeepEP repo 의 dispatch.cu / combine.cu 한 번 펼친다. NVSHMEM 의 production 활용 패턴.

• 강의의 정확한 latency 수치 — § 02, § 05 의 막대 그래프는 외부 자료 + 일반 지식 기반의 개념적 형태. 강의 안의 실제 측정 hardware 와 정확한 수치는 영상 직접 확인 필요.
• 스피커의 background — Prajwal Singhania 의 정확한 affiliation, NVSHMEM 활용 사례. 영상의 자기소개 부분 확인.
• NVSHMEM-X 의 release — NCCL 과 NVSHMEM 의 통합 — 발표 시점 기준 release 상태.
• PyTorch 통합 status — torch.distributed 가 NVSHMEM backend 를 받는지. 직접 사용해야 하는지.
• DeepEP 의 정확한 구현 디테일 — dispatch / combine 의 알고리즘, signal 의 사용 방식.
• NVL72 의 실측 — Blackwell NVL72 환경에서의 NVSHMEM AllReduce latency. 발표 시점에 데이터가 있다면 확인.