[flashinfer] FlashInfer의 Cake Ulysses All-to-All 최적화: World Size 2 성능 향상 및 Symmetric Handshake 도입

PR 링크: flashinfer-ai/flashinfer#7566 상태: Merged | 변경: +830 / -818

들어가며

대규모 언어 모델(LLM)의 학습과 추론에서 Sequence Parallelism은 긴 문맥(Long Context)을 처리하기 위한 필수적인 기술입니다. 그 중에서도 Ulysses 방식은 Attention 헤드를 가로질러 데이터를 분산시키는 All-to-All 통신을 핵심으로 합니다.

FlashInfer는 이러한 통신을 최적화하기 위해 cake_ulysses 백엔드를 사용합니다. 기존의 Cake 구현체는 NVLink를 통해 피어(Peer)의 등록된 출력 버퍼에 직접 데이터를 쓰는 'Direct Route'를 제공해왔으나, 특정 상황(특히 World Size 2)에서의 효율성과 핸드셰이크(Handshake) 로직의 복잡성이라는 숙제가 남아 있었습니다.

이번 PR은 World Size 2에서의 CTA(Cooperative Thread Array) 스케줄링을 재설계하고, 통신 동기화를 위한 핸드셰이크 방식을 대칭형(Symmetric)으로 개선하여 성능과 안정성을 동시에 잡았습니다.

코드 분석: 커널 업데이트와 빌드 구성 변경

이번 변경사항의 핵심 로직은 cake 라이브러리의 내부 커널 업데이트에 있으며, FlashInfer 레포지토리에서는 이를 반영하기 위한 build.json 및 바이너리 해시 업데이트가 이루어졌습니다.

1. Cake Revision 업데이트

가장 먼저 눈에 띄는 변화는 최적화된 커널이 포함된 새로운 Cake 리비전을 참조하도록 변경된 점입니다.

// flashinfer/cake_ulysses_backend/build.json

   "generator": {
-    "cake_revision": "b19b4eb91758889d661a816b11adc0a92d2a8b97",
+    "cake_revision": "9bd936722a2f9af4f4054da183b44d2080e6b65e",
     "cuda": "13.4.92",
-    "rendered_at": "2026-10-27T00:11:54Z"
+    "rendered_at": "2026-10-27T05:57:01Z"
   },

이 리비전 업데이트를 통해 bf16, f16, f32 데이터 타입과 World Size 2, 4, 6, 8에 대응하는 24개의 Direct Route 프로그램이 모두 갱신되었습니다.

2. 커널 바이너리 및 체크섬 변경

최적화된 커널이 적용됨에 따라 각 아키텍처(sm_100a, sm_103a)별 공유 라이브러리의 해시값이 변경되었습니다. 이는 단순한 코드 수정을 넘어 컴파일된 바이너리 수준에서의 최적화가 반영되었음을 의미합니다.

     {
       "arch": "sm_100a",
-      "closure_sha256": "23ce8f16e3330d4542364c4501557ce8ace15b9c44c350101800fe7c461647d2",
+      "closure_sha256": "46a2bdbe19e5b4eb78d740a5b00a3bd7f34c72c29efa8187ae7735e4d708ca19",
       "entry": "run",
       // ...
       "sha256": {
         "csrc/cake_ulysses_a2a_09847afbb3404357a34a_binding.cu": "6fc3d172628aebf25ab75fac27a45e57e1fb38f3bfd1b22914719945118ed889",
-        "csrc/cake_ulysses_a2a_09847afbb3404357a34a_kernel.cu": "48e44a830376e28332a8f7c5edccb6cd454d9fa3f7adfa7f0b20f13ea50c32a3"
+        "csrc/cake_ulysses_a2a_09847afbb3404357a34a_kernel.cu": "ce43535f1e1abc52d347c0533c2ddf0bdfe6a213ee545bd28207ed23ebf9016c"
       },

왜 이게 좋은 최적화인가?

1. World Size 2에서의 비대칭 CTA 할당 (Load Balancing)

기존 방식에서는 World Size 2 환경에서 각 GPU가 자신의 블록을 복사하는 작업과 상대방(Peer)에게 데이터를 보내는 작업 사이의 부하 균형이 완벽하지 않았습니다.

이번 PR에서는 128개의 CTA가 피어의 블록을 스트리밍하도록 설정하고, 20개의 CTA가 자신의 블록을 복사하도록 역할을 분담했습니다.

• 이유: NVLink를 통한 원격 쓰기(Remote Write)는 로컬 메모리 복사보다 지연 시간이 길고 대역폭 관리가 까다롭습니다. 따라서 더 많은 CTA를 원격 통신에 배치함으로써 전체적인 커널 실행 시간을 단축시켰습니다.

2. Symmetric Handshake (대칭형 핸드셰이크)

통신 시작(Entry)과 완료(Landed)를 확인하는 핸드셰이크 로직이 개선되었습니다.

• 기존: 각 CTA가 처리하는 청크(Chunk) 수에 의존하여 동기화를 수행했습니다.
• 개선: CTA별 청크 수와 무관하게 대칭적인 핸드셰이크 구조를 가집니다.
• 이점: 로직이 단순해지면 하드웨어 수준에서의 동기화 오버헤드가 줄어들고, 특히 데이터 정렬(Alignment)이 맞지 않거나 특정 랭크에 보낼 데이터가 없는 예외 상황에서도 안정적으로 동작합니다.

3. 정렬되지 않은 데이터(Unaligned Input) 지원

리뷰어인 @yyihuang이 요청한 테스트(test_cake_ulysses_a2a.py)를 통해 확인할 수 있듯이, 이번 최적화는 입력 데이터가 메모리 경계에 정렬되지 않은 경우에도 안정성을 보장합니다. 이는 실제 서비스 환경에서 가변 길이 입력을 처리할 때 매우 중요한 요소입니다.

성능 수치 분석

NVIDIA Blackwell 아키텍처(B200, B300)에서의 벤치마크 결과는 다음과 같습니다.

• B300 SXM6 (World Size 2): Scatter/Gather Heads 작업에서 약 1.03x (3%)의 속도 향상을 보였습니다 (134us -> 130us).
• B200 (World Size 2): 최대 1.039x (약 4%)의 성능 향상이 관찰되었습니다.

World Size가 커질수록(4, 6, 8) 성능 향상폭은 미미하지만(0.1%~0.5%), 가장 빈번하게 사용되는 최소 단위 병렬화인 World Size 2에서 유의미한 이득을 얻었다는 점이 핵심입니다.

결론

이번 PR은 분산 통신 커널에서 CTA 스케줄링 최적화와 동기화 프로토콜의 단순화가 얼마나 중요한지를 잘 보여줍니다. 특히 하드웨어의 특성(NVLink vs Local Copy)을 고려하여 워크로드를 비대칭적으로 분배한 결정이 실제 성능 향상으로 이어졌습니다.

고성능 라이브러리 개발자라면, 단순히 모든 코어에 작업을 균등하게 나누는 것이 아니라 병목 지점(Remote Write)에 더 많은 자원을 투입하는 전략을 고민해볼 필요가 있습니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html
• https://pytorch.org/docs/stable/distributed.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.