[flashinfer] FlashInfer, BF16 활성화 및 MXFP8 가중치에 대한 Cake MegaMoE EP16 백엔드 최적화

PR 링크: flashinfer-ai/flashinfer#5148 상태: Merged | 변경: +10361 / -3480

들어가며

최근 소프트웨어 개발에서 대규모 언어 모델(LLM)의 중요성이 커지면서, 모델의 추론 속도와 효율성을 높이기 위한 다양한 최적화 기법이 연구되고 있습니다. 특히 Mixture-of-Experts (MoE) 모델은 파라미터 수를 늘리면서도 연산량을 효율적으로 관리할 수 있어 주목받고 있습니다.

이번 글에서는 FlashInfer 라이브러리의 실험적인 기능인 CakeMxfp8MegaMoeEp16 백엔드에서 BF16 활성화(activation)와 MXFP8(Mixed-Precision Float 8) 전문가 가중치(expert weights)를 사용할 때의 성능을 최적화한 Pull Request(PR)에 대해 자세히 알아보겠습니다. 이 PR은 기존 백엔드의 성능을 개선하여 더 빠르고 효율적인 MoE 모델 추론을 가능하게 합니다.

코드 분석

이번 PR의 핵심은 flashinfer/experimental/cake_mxfp8_megamoe_ep16 디렉토리 내의 코드를 중심으로 이루어졌습니다. 특히 backend.py 파일과 예제 파일의 변경 사항을 통해 최적화의 내용을 파악할 수 있습니다.

1. 라우팅 로직 변경 (examples/experimental/cake_mxfp8_megamoe_ep16.py)

가장 눈에 띄는 변경 중 하나는 예제 코드에서 라우팅(routing) 방식을 결정하는 로직입니다. 기존에는 owners, groups, first_experts 등을 계산하여 라우팅을 결정했지만, PR 이후에는 route_slots를 계산하고 이를 기반으로 topk_ids를 생성하는 방식으로 변경되었습니다.

Before:

    owners = global_tokens % 16
    groups = (global_tokens // 16) % 4
    first_experts = owners * 32 + groups * 8
    topk_ids = first_experts[:, None] + torch.arange(8, device=device)[None, :]

After:

    route_slots = global_tokens[:, None] * 8 + torch.arange(8, device=device)[None, :]
    # The affine permutation spreads routes across owners while giving every
    # expert the same global load for each supported token count.
    topk_ids = (route_slots * 73 + 19) % 512

이 변경은 라우팅을 좀 더 효율적으로 분산시키고, 각 전문가(expert)가 받는 부하를 균등하게 맞추려는 시도로 보입니다. (route_slots * 73 + 19) % 512와 같은 아핀 변환(affine transformation)은 라우팅 정보를 전문가 ID 공간에 고르게 퍼뜨리는 역할을 합니다. 이는 특정 전문가에게 부하가 집중되는 것을 방지하여 전체적인 성능 향상에 기여할 수 있습니다.

2. MXFP8 양자화 및 스케일 처리 개선 (flashinfer/experimental/cake_mxfp8_megamoe_ep16/backend.py)

MXFP8 가중치를 처리하는 부분에서도 중요한 변경이 있었습니다. 특히 가중치와 스케일(scale)을 양자화하고 패킹(packing)하는 함수들이 수정되었습니다.

• _interleave_gate_up_128 -> _interleave_gate_up_16: 기존에는 128개의 행(row)을 기준으로 가중치를 인터리빙(interleaving)했지만, 변경 후에는 16개의 행을 기준으로 변경되었습니다. 이는 MXFP8 형식에 더 적합하게 데이터를 재구성하여 메모리 접근 패턴을 최적화하려는 의도로 보입니다.

  Before (개념적):

  # _interleave_gate_up_128
  result_blocks = result.view(
      _LOCAL_EXPERTS, _INTERMEDIATE // 128, 128, _HIDDEN
  )
  # ... up, gate 분리 및 복사 ...
  result_blocks[:, :, 0].copy_(up)
  result_blocks[:, :, 1].copy_(gate)
  

  After (개념적):

  # _interleave_gate_up_16
  result_blocks = result.view(experts, intermediate // 16, 2, 16, hidden)
  # ... up, gate 분리 및 복사 ...
  result_blocks[:, :, 0].copy_(gate)
  result_blocks[:, :, 1].copy_(up)
  

• _quantize_mxfp8_block32: 이 함수는 가중치를 MXFP8 형식으로 양자화하는 핵심 로직을 담당합니다. PR에서는 스케일 계산 방식이 개선되었습니다. 기존에는 scale_seed를 계산하고 이를 기반으로 bits와 codes를 생성했지만, 변경 후에는 safe_max, scale_exp를 사용하여 더 안정적인 스케일 값을 계산하고, 이를 통해 codes를 생성합니다. 특히 torch.exp2 대신 torch.pow(2.0, ...)를 사용하여 부동 소수점 연산의 정확도를 높이려는 시도가 엿보입니다.

  Before (스케일 계산 일부):

  scale_seed = torch.clamp(block_max, min=1.0e-7) / 448.0
  bits = scale_seed.contiguous().view(torch.int32)
  codes = ((bits >> 23) & 255) + (((bits & 0x7FFFFF) + 0x7FFFFF) >> 23)
  codes.clamp_(1, 254)
  decoded = torch.exp2(codes.float() - 127.0)
  

  After (스케일 계산 일부):

  safe_max = torch.clamp(block_max, min=1.0e-30)
  scale_exp = torch.ceil(torch.log2(safe_max * (1.0 / 448.0)))
  codes = torch.clamp(scale_exp + 127.0, min=0.0, max=254.0).to(torch.int32)
  codes = torch.where(block_max == 0, torch.zeros_like(codes), codes)
  decoded = torch.pow(2.0, codes.float() - 127.0)
  

• 스케일 패킹 함수 변경: _pack_scale_n256_k128 함수가 _pack_scale_n128_k128로 변경되었습니다. 이는 MXFP8 가중치와 함께 사용되는 스케일 데이터의 레이아웃을 최적화하여 메모리 접근 효율성을 높이는 것으로 보입니다. N128/K128은 특정 차원의 크기를 나타내며, 이는 GPU 메모리 계층 구조와 연산 패턴에 맞춰 조정된 것입니다.

3. 커널 디자인 및 최적화 (README.md 및 backend.py)

PR 설명과 README.md 파일에는 CUDA 커널 수준에서의 다양한 최적화 기법이 언급되어 있습니다.

• [N32, N16, N16] 행 스케줄: 전문가당 최대 64개의 경로(route)를 처리할 수 있는 용량을 제공하며, 실제 오버플로우는 전문가 로드가 32 또는 48개를 초과할 때만 실행됩니다. 이는 불필요한 연산을 줄여 성능을 향상시킵니다.
• MXFP8 가중치 재사용: 변환된 MXFP8 가중치 타일(tile)을 각 행 세그먼트(segment)에서 재사용하여 메모리 로드 횟수를 줄입니다.
• FC1/FC2 작업 분산: 72개의 CTA(Compute Threading Array) 클러스터에 걸쳐 FC1 및 FC2 연산을 분산시키고 균형을 맞춥니다. 각 CTA는 384개의 스레드를 사용하며, 75KiB의 동적 공유 메모리를 활용합니다.
• 디스패치 그룹화: 소스 토큰별로 경로를 그룹화하여, 원격 BF16 은닉 행(hidden row)을 대상 소유자별로 한 번만 가져오고 로컬 소유 경로 슬롯으로 분산시킵니다.
• FC2 결과 직접 반환: FC2 연산 결과를 해당 소스 (token, top-k slot) 버퍼에 직접 반환하여, 중간 저장 및 복사 과정을 제거합니다.
• CTA 리더 동기화 및 레지스터 예산 최적화: 동기화 오버헤드를 줄이고 리소스 사용을 최적화합니다.
• 안정적인 위치 바인딩 및 스트림 브릿징: 반복적인 호스트 준비 작업을 제거합니다.

이러한 커널 수준의 최적화는 GPU 하드웨어의 특성을 최대한 활용하여 연산 효율성을 극대화하는 데 중점을 둡니다.

4. 세션 관리 및 API 동작

• TMA 디스크립터 초기화: TMA(Tensor Memory Access) 디스크립터는 세션 생성 시 한 번만 초기화됩니다. 이후 각 run 호출은 두 개의 커널만 실행합니다: 하나는 모델 커널(FC1, SwiGLU, FC2 등)이고, 다른 하나는 순서대로 Top-K를 줄이는 커널입니다.
• CUDA Graph 미지원: CUDA Graph 캡처는 의도적으로 지원하지 않습니다. 이는 커널 실행의 동적인 특성 때문일 수 있습니다.
• JIT 전용 및 실험적: 이 백엔드는 여전히 JIT(Just-In-Time) 컴파일 전용이며, 실험적인 상태로 유지됩니다. 자동 백엔드 선택이나 AOT(Ahead-Of-Time) 패키징에는 포함되지 않습니다.
• 세션 재사용 제한: 세션은 약 1,491만 번의 포워드 패스 이후 재사용 전에 새로 생성해야 합니다. 이는 내부 카운터의 오버플로우를 방지하기 위함입니다.

왜 이게 좋은가?

이번 PR은 다음과 같은 이유로 좋은 최적화 및 개선이라고 할 수 있습니다.

1. 성능 향상: PR 설명에 포함된 성능 측정 결과에 따르면, 다양한 라우팅 및 토큰 수 구성에서 평균 1.123400x의 기하 평균 속도 향상(geometric mean speedup)을 달성했습니다. 특히 hot-expert 설정에서 최대 1.16x 이상의 속도 향상을 보여줍니다. 이는 BF16 활성화와 MXFP8 가중치를 사용하는 MegaMoE 백엔드의 실제 추론 성능을 크게 개선한 것입니다.

   Routing	Global tokens	Tokens/rank	Reference (ms)	This PR (ms)	Speedup
   balanced	256	16	1.326420	1.162211	1.141290x
   balanced	512	32	1.3156025	1.176719	1.118026x
   balanced	1024	64	1.322886	1.1649745	1.135549x
   hot-expert	256	16	1.321246	1.136386	1.162674x
   hot-expert	512	32	1.310207	1.157282	1.132142x
   hot-expert	1024	64	1.283619	1.217985	1.053887x
   Geometric mean					1.123400x

2. 메모리 및 연산 효율성 증대: MXFP8과 같은 혼합 정밀도 형식을 효과적으로 활용하고, 가중치 재사용, 불필요한 복사 제거, CTA 및 스레드 수준의 작업 분산 등을 통해 메모리 대역폭 사용량과 연산량을 최적화했습니다. 이는 특히 대규모 MoE 모델에서 중요한 요소입니다.

3. 하드웨어 특성 활용: TMA 디스크립터 초기화, CTA 클러스터 설계, 공유 메모리 활용 등은 NVIDIA GPU 아키텍처의 특성을 고려한 최적화입니다. 이를 통해 하드웨어의 잠재력을 최대한 끌어낼 수 있습니다.

4. 코드 품질 및 안정성: 라우팅 유효성 검사(_validate_gathered_routing_capacity, _validate_routing_capacity), 세션 실행 횟수 제한(_MAX_LAUNCH_EPOCH) 등은 코드의 안정성을 높이고 잠재적인 오류를 방지하는 데 기여합니다. 또한, 스케일 계산 로직 개선은 양자화 과정의 수치적 안정성을 향상시킬 수 있습니다.

일반적 교훈

• 혼합 정밀도 활용의 중요성: MXFP8과 같은 혼합 정밀도 데이터 타입을 적극적으로 활용하면 메모리 사용량을 줄이고 연산 속도를 높일 수 있습니다. 다만, 이를 위해서는 데이터 형식에 맞는 양자화, 역양자화, 스케일 관리 로직이 필수적입니다.
• 하드웨어 친화적 커널 설계: GPU의 메모리 계층 구조(레지스터, 공유 메모리, L1/L2 캐시, 전역 메모리)와 병렬 처리 모델(스레드, 워프, CTA)을 깊이 이해하고 커널을 설계하는 것이 성능 향상의 핵심입니다.
• 동적 연산 최적화: MoE와 같이 동적인 라우팅이 발생하는 경우, 라우팅 정보의 분산, 전문가 로드 밸런싱, 불필요한 데이터 이동 최소화가 전체 성능에 큰 영향을 미칩니다.
• 실험적 기능의 체계적 관리: 실험적인 기능은 명확한 API 변경 방지, 명시적 옵트인(opt-in) 요구, 별도 테스트 경로 관리 등을 통해 안정성을 확보하면서 개발을 진행해야 합니다.

리뷰 댓글 분석

제공된 리뷰 댓글은 주로 테스트 실행 및 CI 파이프라인 관련 내용이었습니다. 예를 들어 [yzh119]와 [yyihuang]님이 테스트 실행을 요청하고, [flashinfer-bot]이 CI 파이프라인 상태를 보고하는 내용이 주를 이룹니다. 이는 코드 자체의 복잡성보다는, 실험적인 기능으로서의 철저한 테스트와 검증 과정을 거치고 있음을 시사합니다. 특히 tests/experimental/test_cake_mxfp8_megamoe_ep16.py 파일에 대한 테스트 실행 요청은 해당 백엔드의 기능과 성능이 집중적으로 검증되고 있음을 보여줍니다.

References

• torch.compile: PyTorch의 컴파일 기능으로, FlashInfer와 같은 라이브러리에서 JIT 컴파일을 활용하는 방식과 유사한 맥락에서 참고할 수 있습니다.
• NVIDIA MXFP8 Data Format: NVIDIA에서 제공하는 MXFP8 데이터 형식에 대한 정보 (공식 문서 링크는 변경될 수 있으므로, 일반적인 정보 링크로 대체합니다.)
• FlashInfer Documentation: FlashInfer 라이브러리의 공식 문서 (API 및 백엔드 사용법 참고)

참고 자료

• https://docs.flashinfer.ai/

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