[flashinfer] FlashInfer MegaMoE v3 최적화: TP4/TP8 커널 성능 향상 분석
PR 링크: flashinfer-ai/flashinfer#8220 상태: Merged | 변경: +1241 / -508
들어가며
최근 대규모 언어 모델(LLM)의 발전은 모델의 크기뿐만 아니라 추론 속도 향상에 대한 요구를 증대시키고 있습니다. 특히 Mixture-of-Experts (MoE) 모델은 파라미터 수를 늘리면서도 활성화되는 파라미터 수를 제한하여 효율성을 높이는 아키텍처로 주목받고 있습니다. 하지만 이러한 MoE 모델의 복잡성은 GPU 커널 수준에서의 최적화를 더욱 중요하게 만듭니다.
이번 글에서는 flashinfer-ai/flashinfer 레포지토리의 PR(#138)에서 진행된 flashinfer.experimental.mega_moe_v3 파이프라인의 성능 최적화에 대해 심층적으로 분석하고자 합니다. 이 PR은 특히 TP4/TP8 (Tensor Parallelism / Expert Parallelism) 설정에서 특정 토큰 수에 대한 커널 실행 속도를 향상시키는 데 초점을 맞추고 있습니다. 코드 변경 사항을 자세히 살펴보고, 이러한 변경이 왜 성능 향상으로 이어지는지, 그리고 그 교훈은 무엇인지 알아보겠습니다.
코드 분석
이 PR은 주로 deepgemm_mega_moe_v3_pipeline_fp4_..._kernel.cu 파일들의 수정으로 성능을 개선했습니다. 핵심 변경 사항은 공유 메모리(shared memory) 사용량 최적화와 파이프라인 단계(pipeline stages)의 재조정을 통해 데이터 이동 및 동기화 오버헤드를 줄이는 것입니다.
1. 공유 메모리(Shared Memory) 및 파이프라인 단계(Pipeline Stages) 재조정
가장 눈에 띄는 변경은 NUM_PIPE_STAGES의 증가와 여러 공유 메모리 오프셋(offset)의 재조정입니다. 예를 들어, deepgemm_mega_moe_v3_pipeline_fp4_e384_i384_t32_sm103a_kernel.cu 파일에서 NUM_PIPE_STAGES가 11에서 12로 증가했습니다.
Before:
#define NUM_PIPE_STAGES 11
...
#define SMEM_SMEM_OUTPUT_OFF 22528
#define SMEM_SMEM_A_OFF 26624
#define SMEM_SMEM_B_OFF 37888
#define SMEM_SMEM_SFA_OFF 218112
#define SMEM_SMEM_SFB_OFF 223744
...
const int mbar_base = smem + 229760;
#define full_addr (mbar_base + 0)
#define empty_addr (mbar_base + 88)
#define tmem_full_addr (mbar_base + 176)
#define tmem_empty_addr (mbar_base + 192)
#define task_full_addr (mbar_base + 208)
#define task_empty_addr (mbar_base + 240)
#define combine_barriers_addr (mbar_base + 272)
...
uint8_t* smem_output = reinterpret_cast<uint8_t*>(smem_raw + 22528);
const int smem_output_addr = smem + 22528;
uint8_t* smem_a = reinterpret_cast<uint8_t*>(smem_raw + 26624);
const int smem_a_addr = smem + 26624;
uint8_t* smem_b = reinterpret_cast<uint8_t*>(smem_raw + 37888);
const int smem_b_addr = smem + 37888;
unsigned int* smem_sfa = reinterpret_cast<unsigned int*>(smem_raw + 218112);
const int smem_sfa_addr = smem + 218112;
unsigned int* smem_sfb = reinterpret_cast<unsigned int*>(smem_raw + 223744);
const int smem_sfb_addr = smem + 223744;
...
// Mbarrier init (7 pipeline groups, 0 ordered-sequence groups, 82 barriers)
// Mbarriers at smem_raw[229760..230416)
...
// full: 11 barriers, init_count=4
...
// empty: 11 barriers, init_count=1
...
// tmem_full: 2 barriers, init_count=1
...
// tmem_empty: 2 barriers, init_count=512
...
// task_full: 4 barriers, init_count=1
...
// task_empty: 4 barriers, init_count=512
...
// combine_barriers: 48 barriers, init_count=1
...
mbarrier_init(smem + 230408, 1);
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}
}
...
volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230416);
if (warp == 3) {
int _tmem_hold = smem + 230416;
asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(64) : "memory");
__syncwarp();
}
After:
#define NUM_PIPE_STAGES 12
...
#define SMEM_SMEM_OUTPUT_OFF 4096
#define SMEM_SMEM_A_OFF 8192
#define SMEM_SMEM_B_OFF 20480
#define SMEM_SMEM_SFA_OFF 217088
#define SMEM_SMEM_SFB_OFF 223232
...
const int mbar_base = smem + 229760;
#define full_addr (mbar_base + 0)
#define empty_addr (mbar_base + 96)
#define tmem_full_addr (mbar_base + 192)
#define tmem_empty_addr (mbar_base + 208)
#define task_full_addr (mbar_base + 224)
#define task_empty_addr (mbar_base + 256)
#define combine_barriers_addr (mbar_base + 288)
...
uint8_t* smem_output = reinterpret_cast<uint8_t*>(smem_raw + 4096);
const int smem_output_addr = smem + 4096;
uint8_t* smem_a = reinterpret_cast<uint8_t*>(smem_raw + 8192);
const int smem_a_addr = smem + 8192;
uint8_t* smem_b = reinterpret_cast<uint8_t*>(smem_raw + 20480);
const int smem_b_addr = smem + 20480;
unsigned int* smem_sfa = reinterpret_cast<unsigned int*>(smem_raw + 217088);
const int smem_sfa_addr = smem + 217088;
unsigned int* smem_sfb = reinterpret_cast<unsigned int*>(smem_raw + 223232);
const int smem_sfb_addr = smem + 223232;
...
// Mbarrier init (7 pipeline groups, 0 ordered-sequence groups, 84 barriers)
// Mbarriers at smem_raw[229760..230432)
...
// full: 12 barriers, init_count=4
...
// empty: 12 barriers, init_count=1
...
// tmem_full: 2 barriers, init_count=1
mbarrier_init(smem + 229952, 1);
mbarrier_init(smem + 229960, 1);
// tmem_empty: 2 barriers, init_count=512
mbarrier_init(smem + 229968, 512);
mbarrier_init(smem + 229976, 512);
// task_full: 4 barriers, init_count=1
mbarrier_init(smem + 229984, 1);
mbarrier_init(smem + 229992, 1);
mbarrier_init(smem + 230000, 1);
mbarrier_init(smem + 230008, 1);
// task_empty: 4 barriers, init_count=512
mbarrier_init(smem + 230016, 512);
mbarrier_init(smem + 230024, 512);
mbarrier_init(smem + 230032, 512);
mbarrier_init(smem + 230040, 512);
// combine_barriers: 48 barriers, init_count=1
mbarrier_init(smem + 230048, 1);
mbarrier_init(smem + 230056, 1);
mbarrier_init(smem + 230064, 1);
mbarrier_init(smem + 230072, 1);
mbarrier_init(smem + 230080, 1);
mbarrier_init(smem + 230088, 1);
mbarrier_init(smem + 230096, 1);
mbarrier_init(smem + 230104, 1);
mbarrier_init(smem + 230112, 1);
mbarrier_init(smem + 230120, 1);
mbarrier_init(smem + 230128, 1);
mbarrier_init(smem + 230136, 1);
mbarrier_init(smem + 230144, 1);
mbarrier_init(smem + 230152, 1);
mbarrier_init(smem + 230160, 1);
mbarrier_init(smem + 230168, 1);
mbarrier_init(smem + 230176, 1);
mbarrier_init(smem + 230184, 1);
mbarrier_init(smem + 230192, 1);
mbarrier_init(smem + 230200, 1);
mbarrier_init(smem + 230208, 1);
mbarrier_init(smem + 230216, 1);
mbarrier_init(smem + 230224, 1);
mbarrier_init(smem + 230232, 1);
mbarrier_init(smem + 230240, 1);
mbarrier_init(smem + 230248, 1);
mbarrier_init(smem + 230256, 1);
mbarrier_init(smem + 230264, 1);
mbarrier_init(smem + 230272, 1);
mbarrier_init(smem + 230280, 1);
mbarrier_init(smem + 230288, 1);
mbarrier_init(smem + 230296, 1);
mbarrier_init(smem + 230304, 1);
mbarrier_init(smem + 230312, 1);
mbarrier_init(smem + 230320, 1);
mbarrier_init(smem + 230328, 1);
mbarrier_init(smem + 230336, 1);
mbarrier_init(smem + 230344, 1);
mbarrier_init(smem + 230352, 1);
mbarrier_init(smem + 230360, 1);
mbarrier_init(smem + 230368, 1);
mbarrier_init(smem + 230376, 1);
mbarrier_init(smem + 230384, 1);
mbarrier_init(smem + 230392, 1);
mbarrier_init(smem + 230400, 1);
mbarrier_init(smem + 230408, 1);
mbarrier_init(smem + 230416, 1);
mbarrier_init(smem + 230424, 1);
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}
}
...
volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 230432);
if (warp == 3) {
int _tmem_hold = smem + 230432;
asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(64) : "memory");
__syncwarp();
}
이 변경은 다음과 같은 효과를 가져옵니다:
NUM_PIPE_STAGES증가: 파이프라인의 총 단계 수가 늘어나면서, 각 단계에서 처리할 수 있는 데이터의 양이 줄어들거나, 더 많은 동시 작업이 가능해집니다. PR 설명에 따르면, 이는 'dispatch' 단계가 메타데이터를 통합하고 레지스터를 통해 토큰 행을 가져오도록 변경되었기 때문입니다. 이로 인해 'send' 영역이 축소되어 공유 메모리 공간이 확보되었고, 이 확보된 공간이 새로운 파이프라인 단계로 활용됩니다.- 공유 메모리 오프셋 재조정:
SMEM_SMEM_OUTPUT_OFF,SMEM_SMEM_A_OFF,SMEM_SMEM_B_OFF등의 오프셋이 변경되었습니다. 예를 들어SMEM_SMEM_OUTPUT_OFF는 22528에서 4096으로 크게 줄었습니다. 이는 이전에는 불필요하게 크거나 비효율적으로 할당되었던 공유 메모리 영역을 재구성하여, 실제 필요한 데이터만 효율적으로 저장하도록 변경되었음을 시사합니다. 특히, 'dispatch' 영역의 축소로 확보된 공유 메모리가 'output staging'을 위한 새로운 12번째 단계로 사용됩니다. - Mbarrier 초기화 조정:
NUM_PIPE_STAGES증가에 따라mbarrier_init호출 횟수 및 관련 주소(full_addr,empty_addr등)가 조정되었습니다. 이는 파이프라인의 각 단계에서 동기화를 관리하는 메커니즘이 변경되었음을 나타냅니다.
2. TP4 EP4 (1024 tokens) 설정 개선
TP4 EP4 설정에서 1024 토큰 처리 시에도 11단계에서 12단계로 파이프라인이 확장되었습니다. PR 설명에 따르면, 이는 'dispatch'가 각 토큰 행을 두 번의 'half-row' 트립으로 가져오고, FC2 출력 스테이징이 그룹당 8개 행을 저장하도록 변경되었기 때문입니다. 이 역시 공유 메모리 사용을 최적화하여 성능을 향상시킵니다.
3. 기타 최적화
- TP4 EP1 및 TP8 EP1 (32 tokens): 이 설정에서는 런치 셰이프(launch shape)와 호스트 바인딩(host binding)이 유지되었습니다. 이는 해당 설정에서는 큰 변경이 필요 없었거나, 변경 시 오히려 성능 저하가 발생할 수 있음을 의미합니다.
- TP4 EP4 및 TP8 EP8 (32 tokens): 이 설정들은 여전히 하나의 호스트 바인딩을 공유합니다 (228480 B 동적 공유 메모리). 이는 리소스 공유를 통해 효율성을 높이려는 시도로 볼 수 있습니다.
- TP4 EP4 (1024 tokens): 이 설정은 이제 자체 호스트 바인딩을 갖게 되었습니다. 이는 이전보다 더 많은 리소스나 독립적인 관리가 필요함을 나타냅니다.
왜 이게 좋은가?
이 PR의 핵심은 공유 메모리 사용량 최적화와 파이프라인 단계의 효율적인 재구성을 통해 GPU 커널의 성능을 직접적으로 향상시킨 것입니다.
-
실행 시간 단축: PR 설명에 따르면, 변경된 5가지 설정(TP4 EP1, TP8 EP1, TP4 EP4 at 32 tokens, TP4 EP4 at 1024 tokens, TP8 EP8 at 32 tokens)은 각각 1.005x ~ 1.007x의 속도 향상을 보였습니다. 이는 0.32 ~ 1.87 µs의 지연 시간(latency) 감소로 이어집니다.
- TP4 EP1: 1.005x
- TP8 EP1: 1.007x
- TP4 EP4 (32 tokens): 1.006x
- TP4 EP4 (1024 tokens): 1.007x
- TP8 EP8 (32 tokens): 1.004x
이러한 미미해 보이는 속도 향상이 수백만, 수십억 개의 파라미터를 가진 LLM 추론에서는 누적되어 상당한 성능 개선 효과를 가져올 수 있습니다.
-
메모리 대역폭 및 레이턴시 개선: 공유 메모리 사용량을 줄이고 데이터 접근 패턴을 최적화함으로써, GPU 코어는 더 적은 지연 시간으로 필요한 데이터에 접근할 수 있게 됩니다. 이는 메모리 병목 현상을 완화하고 연산 효율성을 높입니다.
-
더 효율적인 커널 설계: 'dispatch' 단계의 메타데이터 폴딩(folding) 및 토큰 행 로딩 방식 변경, FC2 출력 스테이징 개선 등은 복잡한 MoE 연산을 더 작은 단위로 나누고 각 단계를 최적화하는 세밀한 튜닝의 중요성을 보여줍니다.
일반적인 교훈:
- 공유 메모리 최적화의 중요성: GPU 커널 성능에서 공유 메모리는 매우 중요한 자원입니다. 불필요한 할당을 줄이고, 데이터 접근 패턴을 최적화하며, 파이프라이닝과 결합하여 효율적으로 사용하는 것이 핵심입니다.
- 파이프라이닝의 미세 조정: 파이프라인 단계의 수와 각 단계의 작업량을 조절하는 것은 전체 처리량을 결정하는 중요한 요소입니다. 하드웨어 제약(SMEM 크기, 스레드 블록당 최대 공유 메모리 등)과 연산 특성을 고려하여 최적의 균형점을 찾아야 합니다.
- 설정별 최적화: MoE 모델은 다양한 TP/EP 조합과 토큰 길이를 가질 수 있습니다. 각 설정의 특성에 맞춰 커널을 미세 조정하는 것이 전반적인 성능 향상에 필수적입니다.
리뷰 피드백 반영
리뷰어 yyihuang은 이 PR과 관련된 테스트 스크립트 실행을 요청했습니다. 이는 코드 변경이 의도한 대로 작동하는지, 그리고 회귀(regression)가 발생하지 않았는지를 검증하는 표준적인 절차입니다. 이러한 테스트는 PR의 안정성을 보장하는 데 중요한 역할을 합니다.
결론
이번 PR은 FlashInfer의 MegaMoE v3 커널에서 TP4/TP8 설정의 성능을 미세하게나마 향상시키는 중요한 개선을 이루었습니다. 공유 메모리 사용량 최적화와 파이프라인 단계의 재구성을 통해 달성된 1~2%의 속도 향상은 LLM 추론 성능 최적화의 복잡성과 중요성을 다시 한번 강조합니다. 이러한 저수준 최적화는 모델의 효율성을 극대화하고, 더 빠르고 경제적인 AI 서비스 제공의 기반이 됩니다.
참고 자료
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
- [flashinfer] FlashInfer, BF16 활성화 및 MXFP8 가중치에 대한 Cake MegaMoE EP16 백엔드 최적화
- [flashinfer] FlashInfer, B200 GPU에서 Kimi-K3 MoE 라우터 성능 최적화
- [flashinfer] FlashInfer, NVFP4 SVDQuant GEMM 커널 최적화: 아키텍처별 디스패치 도입
- [flashinfer] FlashInfer의 CAKE FP8 GEMM 커널 최적화: 코드 간소화와 성능 향상
- [flashinfer] FlashInfer CAKE FMHA 백엔드 최적화: 코드 중복 제거 및 성능 향상
PR Analysis 의 다른글
- 이전글 [flashinfer] FlashInfer Cake GDN: Blackwell SM107a에서 게이트 스트리밍 프리필 최적화 확장으로 최대 1.12배 성능 향상
- 현재글 : [flashinfer] FlashInfer MegaMoE v3 최적화: TP4/TP8 커널 성능 향상 분석
- 다음글 [flashinfer] FlashInfer: sm_107a 아키텍처를 위한 Fused Decode List Scorer 프리페칭 최적화
댓글