[flashinfer] FlashInfer MoE 최적화: Ampere 및 Hopper 아키텍처에서 그룹화된 활성화 스테이징 오버랩

2026년 10월 26일수정: 2026년 10월 26일

PR 링크: flashinfer-ai/flashinfer#7368 상태: Merged | 변경: +227 / -77

들어가며

최근 딥러닝 모델, 특히 거대 언어 모델(LLM) 분야에서는 Mixture-of-Experts (MoE) 아키텍처가 큰 주목을 받고 있습니다. MoE는 모델의 파라미터 수를 크게 늘리면서도 추론 시에는 일부 전문가(expert)만 활성화하여 연산 효율성을 높이는 장점이 있습니다. 하지만 MoE 레이어 내에서 각 전문가에게 데이터를 분배하고 처리하는 과정은 복잡하며, 특히 GPU 아키텍처의 특성을 고려한 최적화가 성능에 지대한 영향을 미칩니다.

이번 글에서는 FlashInfer 라이브러리의 Pull Request(PR)에서 제안된 MoE 연산 최적화에 대해 자세히 살펴보겠습니다. 이 PR은 특히 NVIDIA Ampere 및 Hopper 아키텍처 GPU에서 그룹화된(grouped) MoE 프로젝션의 활성화 데이터 로딩 및 처리를 개선하여 레이턴시를 줄이는 것을 목표로 합니다. 기존 코드에서는 활성화 데이터를 수집하는 과정이 연산 파이프라인의 발전을 지연시키는 병목 현상이 있었습니다. 이 PR은 이러한 데이터 수집 작업을 연산과 병렬화하고, 독립적인 작업을 겹치게 함으로써 전체적인 성능을 향상시키고자 합니다.

코드 변경 분석

이번 PR의 핵심은 flashinfer/fused_moe/backends/frost_low_latency/tc/ampere.py와 flashinfer/fused_moe/backends/frost_low_latency/tc/hopper.py 파일의 수정입니다. 두 파일 모두 CuTe DSL을 사용하여 CUDA 커널을 작성하고 있으며, MoE 연산의 핵심인 활성화 데이터 로딩(load_b) 및 커널 실행 로직을 수정했습니다.

1. Ampere 아키텍처 최적화 (`ampere.py`)

Ampere 아키텍처(예: A100 GPU)에서는 K128 타일 크기일 때, 그룹화된 활성화 로딩을 최적화했습니다. 기존 코드에서는 tidx < 32 조건 하에 일반적인 스칼라 로딩 방식을 사용했습니다.

Before:

    @cute.jit
    def load_b(self, inputs, ids, members, sB, tidx, tile, slot):
        if tidx < 32:
            for member in cutlass.range_constexpr(8):
                for segment in cutlass.range_constexpr(self.tile_k // 32):
                    column = tidx + segment * 32
                    value = inputs.element_type(0)
                    if cutlass.const_expr(member < self.group_size):
                        if members[member] >= 0:
                            value = inputs[
                                members[member] // self.input_route_divisor,
                                tile * self.tile_k + column,
                            ]
                    sB[member, column, slot] = value

After: 수정된 코드에서는 cutlass.const_expr(self.tile_k == 128) 조건을 추가하고, 입력 데이터가 16바이트로 정렬(aligned)된 경우 cute.nvgpu.cpasync.CopyG2SOp를 사용하여 128비트(num_bits_per_copy=128)의 벡터화된 비동기 복사(cp.async)를 수행합니다. 이는 데이터 전송 효율을 크게 높입니다. 정렬되지 않은 입력이나 다른 타일 크기의 경우 기존 스칼라 폴백(self.load_b_scalar)을 사용합니다. 또한, 각 워프(warp)가 특정 멤버(expert)의 데이터를 담당하도록 하여 데이터 접근 패턴을 개선했습니다.

    @cute.jit
    def load_b(self, inputs, ids, members, sB, tidx, tile, slot):
        if cutlass.const_expr(self.tile_k == 128):
            # Guard the vectorized copy before constructing aligned copy pointers.
            # Public inputs are 16-byte aligned; retain the scalar device fallback.
            # Row stride and K-tile offsets are multiples of 128 elements.
            if inputs.iterator.toint() % 16 == 0:
                atom = cute.make_copy_atom(
                    cute.nvgpu.cpasync.CopyG2SOp(
                        cache_mode=cute.nvgpu.cpasync.LoadCacheMode.GLOBAL
                    ),
                    inputs.element_type,
                    num_bits_per_copy=128,
                )
                column = (tidx % 16) * 8
                for member in cutlass.range_constexpr(8):
                    if tidx // 16 == member:
                        valid = cutlass.Boolean(False)
                        if cutlass.const_expr(member < self.group_size):
                            valid = members[member] >= 0
                        if valid:
                            offset = (
                                (members[member] // self.input_route_divisor)
                                * inputs.stride[0]
                                + tile * self.tile_k
                                + column
                            )
                            src = cute.make_tensor(
                                cute.make_ptr(
                                    inputs.element_type,
                                    (inputs.iterator + offset).toint(),
                                    cute.AddressSpace.gmem,
                                    assumed_align=16,
                                ),
                                cute.make_layout((8,)),
                            )
                            dst = cute.make_tensor(
                                cute.make_ptr(
                                    inputs.element_type,
                                    (
                                        sB.iterator + sB.layout((member, column, slot))
                                    ).toint(),
                                    cute.AddressSpace.smem,
                                    assumed_align=16,
                                ),
                                cute.make_layout((8,)),
                            )
                            cute.copy(atom, src, dst)
                        else:
                            for element in cutlass.range_constexpr(8):
                                sB[member, column + element, slot] = (
                                    inputs.element_type(0)
                                )
            else:
                self.load_b_scalar(inputs, ids, members, sB, tidx, tile, slot)
        else:
            self.load_b_scalar(inputs, ids, members, sB, tidx, tile, slot)

    @cute.jit
    def load_b_scalar(self, inputs, ids, members, sB, tidx, tile, slot):
        if cutlass.const_expr(self.tile_k == 128):
            # Each warp owns disjoint rows. Member indices stay compile-time constants
            # so register-held route membership never needs dynamic indexing.
            warp = tidx // 32
            lane = tidx % 32
            for member in cutlass.range_constexpr(8):
                if warp == member % 4:
                    for segment in cutlass.range_constexpr(self.tile_k // 32):
                        column = lane + segment * 32
                        value = inputs.element_type(0)
                        if cutlass.const_expr(member < self.group_size):
                            if members[member] >= 0:
                                value = inputs[
                                    members[member] // self.input_route_divisor,
                                    tile * self.tile_k + column,
                                ]
                        sB[member, column, slot] = value
        else:
            if tidx < 32:
                for member in cutlass.range_constexpr(8):
                    for segment in cutlass.range_constexpr(self.tile_k // 32):
                        column = tidx + segment * 32
                        value = inputs.element_type(0)
                        if cutlass.const_expr(member < self.group_size):
                            if members[member] >= 0:
                                value = inputs[
                                    members[member] // self.input_route_divisor,
                                    tile * self.tile_k + column,
                                ]
                        sB[member, column, slot] = value

2. Hopper 아키텍처 최적화 (`hopper.py`)

Hopper 아키텍처(예: H100 GPU)에서는 활성화 데이터 수집 및 가중치 로딩(Weight TMA)의 순서를 조정하고, 이를 병렬화하는 데 중점을 두었습니다.

기존 커널 로직 (개략적): 기존에는 가중치 로딩(Weight TMA)이 완료된 후에 활성화 벡터를 수집하는 방식이었습니다. 이는 GPU의 여러 워프(warp)가 독립적으로 작업할 수 있는 잠재력을 충분히 활용하지 못했습니다.

After: 수정된 코드에서는 다음과 같은 개선이 이루어졌습니다.

활성화 행 분산: 활성화 행(row)을 네 개의 워프로 분산시킵니다.
BF16/FP16: 가중치 TMA(Tensor Memory Access)를 활성화 수집보다 먼저 발행하여 독립적인 데이터 전송이 오버랩되도록 합니다. 이는 mainloop_pipeline.producer_acquire와 mainloop_pipeline.producer_commit을 사용하여 파이프라인 상태를 관리하고, cute.copy 호출 시 tma_bar_ptr을 통해 동기화를 명시적으로 처리합니다.
FP8: 가중치 TMA 이전에 수집하는 기존 순서를 유지합니다.
동기화: 모든 활성화 라이터(writer)는 스토어(store)를 발행하고 소비 전에 동기화합니다. 리뷰어의 코멘트에서 지적된 것처럼, 가중치 TMA 완료만으로는 다른 워프의 활성화 스토어가 게시되지 않으므로, cute.arch.fence_proxy("async.shared", space="cta")와 cute.arch.sync_threads()를 사용하여 명시적인 동기화를 수행합니다.

def kernel(
    self,
    # ... other arguments
    mainloop_producer_state,
    tAgA_mkl,
    tAsA,
    tGgA_mkl,
    tGsA,
    a_mcast_mask,
    tma_atom_a,
    tma_bar_ptr,
    mainloop_pipeline,
):
    # ...
    mainloop_producer_state = pipeline.make_pipeline_state(
        pipeline.PipelineUserType.Producer, self.ab_stage
    )
    for _prefetch_idx in cutlass.range(prefetch_k_tile_cnt, unroll=1):
        tAgA_k = tAgA_mkl[None, mainloop_producer_state.count]
        tAsA_pipe = tAsA[None, mainloop_producer_state.index]
        if warp_idx == 0:
            mainloop_pipeline.producer_acquire(mainloop_producer_state)
            # ... (weight TMA and activation gathering logic)
            if cutlass.const_expr(self.a_dtype != cutlass.Float8E4M3FN):
                cute.copy(
                    tma_atom_a,
                    tAgA_k,
                    tAsA_pipe,
                    tma_bar_ptr=mainloop_pipeline.producer_get_barrier(
                        mainloop_producer_state
                    ),
                    mcast_mask=a_mcast_mask,
                )
                # ... (split_gate logic)
            mainloop_pipeline.producer_commit(mainloop_producer_state)
        cute.arch.sync_threads()
        for bn in cutlass.range_constexpr(8):
            if warp_idx == bn % 4:
                # ... (activation data processing)
        cute.arch.fence_proxy("async.shared", space="cta")
        cute.arch.sync_threads()
        if cutlass.const_expr(self.a_dtype == cutlass.Float8E4M3FN):
            if warp_idx == 0:
                cute.copy(
                    tma_atom_a,
                    # ... (FP8 specific copy logic)
                )
    # ...

또한, 문서(docs/frost_low_latency_moe.md)에도 이러한 변경 사항에 대한 설명이 추가되었습니다. Ampere의 경우 cp.async를 사용한 벡터화된 복사, Hopper의 경우 BF16/FP16에서 가중치 TMA와 활성화 수집의 오버랩, FP8에서의 순서 유지, 그리고 명시적인 동기화 메커니즘에 대한 내용이 포함되었습니다.

왜 이게 좋은가?

이 PR의 핵심 목표는 MoE 레이어의 활성화 데이터 로딩 및 처리를 최적화하여 GPU 커널 실행 시간을 단축하는 것입니다. 특히, 다음과 같은 이유로 좋은 최적화라고 할 수 있습니다.

데이터 로딩 병목 현상 해소: MoE 모델에서는 각 토큰이 소수의 전문가만 활성화하므로, 활성화 데이터의 크기가 작습니다. 이 작은 데이터 조각들을 효율적으로 수집하는 것이 중요합니다. 기존 방식은 이 수집 과정이 GPU의 다른 연산과 겹치지 못하고 지연을 유발했습니다. 이 PR은 비동기 복사(cp.async) 및 TMA(Tensor Memory Access)와 같은 하드웨어 기능을 활용하고, 데이터 로딩과 연산을 병렬화함으로써 이 병목을 크게 완화합니다.
하드웨어 특성 활용: Ampere와 Hopper 아키텍처는 각각 다른 특성과 명령어 세트를 가지고 있습니다. 이 PR은 각 아키텍처에 맞춰 최적화된 방식을 적용했습니다. Ampere에서는 128비트 정렬 복사를 통해 데이터 전송 효율을 높였고, Hopper에서는 BF16/FP16 데이터 타입에 대해 가중치 로딩과 활성화 수집을 오버랩시키는 새로운 파이프라인을 도입했습니다. 이는 GPU 하드웨어의 잠재력을 최대한 활용하는 좋은 예시입니다.
명시적인 동기화: 리뷰어의 피드백에서 지적된 것처럼, GPU 프로그래밍에서는 데이터 종속성 관리가 매우 중요합니다. 특히 여러 워프가 공유 메모리를 사용할 때, 데이터가 준비되기 전에 접근하는 것을 방지해야 합니다. 이 PR은 cute.arch.fence_proxy와 cute.arch.sync_threads를 사용하여 활성화 스토어 발행 및 소비 전에 명시적으로 동기화를 수행함으로써 정확성과 안정성을 확보했습니다.
성능 향상: PR 설명에 포함된 벤치마크 결과는 이러한 최적화의 효과를 명확히 보여줍니다.
- A100 (Ampere): 공유(shared) 라우팅에서 약 15-17%, 비동기(disjoint) 라우팅에서 약 7-8%의 커널 GPU 시간 단축을 보였습니다.
- H100 (Hopper): 공유 라우팅에서 약 13-14%, 비동기 라우팅에서 약 3-4%의 단축을 보였습니다.
특히 H100의 (2,8) 스테이지 구성에서 비동기 라우팅 시 약간의 성능 저하(약 3.5-4.5%)가 관찰되었지만, 이는 특정 구성에서의 트레이드오프이며 다른 구성에서의 상당한 성능 향상을 상쇄할 정도는 아닙니다. 이러한 결과는 MoE 연산의 핵심 경로에서 상당한 성능 개선이 이루어졌음을 시사합니다.

일반적 교훈:

하드웨어별 최적화의 중요성: GPU 아키텍처(Ampere, Hopper 등)의 특성을 깊이 이해하고 이를 활용하는 것이 성능 향상의 핵심입니다.
데이터 로딩 최적화: 딥러닝 모델, 특히 MoE와 같이 데이터 접근 패턴이 복잡한 경우, 데이터 로딩 및 전송 병목을 해결하는 것이 전체 성능에 큰 영향을 미칩니다. cp.async, TMA와 같은 기능을 적극 활용해야 합니다.
명시적인 동기화와 파이프라이닝: 복잡한 병렬 연산에서는 데이터 종속성을 명확히 관리하고, 필요한 경우 명시적인 동기화 메커니즘을 사용하여 정확성을 보장해야 합니다. 파이프라이닝 기법을 통해 독립적인 작업을 오버랩시키는 것은 성능 향상의 중요한 전략입니다.

결론

이번 FlashInfer PR은 MoE 연산의 활성화 데이터 처리 과정을 최적화하여 Ampere 및 Hopper GPU에서 상당한 성능 향상을 달성했습니다. 비동기 데이터 전송, 하드웨어별 최적화, 그리고 명시적인 동기화 메커니즘 도입은 딥러닝 프레임워크에서 GPU 성능을 극대화하는 좋은 사례를 보여줍니다. 이러한 최적화는 더 빠르고 효율적인 LLM 추론을 가능하게 할 것입니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer GDN 커널: SM12x GPU에서 상삼각 행렬 MMA 타일 건너뛰기로 성능 최적화
현재글 : [flashinfer] FlashInfer MoE 최적화: Ampere 및 Hopper 아키텍처에서 그룹화된 활성화 스테이징 오버랩
다음글 없음

[flashinfer] FlashInfer MoE 최적화: Ampere 및 Hopper 아키텍처에서 그룹화된 활성화 스테이징 오버랩

들어가며

코드 변경 분석

1. Ampere 아키텍처 최적화 (`ampere.py`)

2. Hopper 아키텍처 최적화 (`hopper.py`)

왜 이게 좋은가?

결론

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글

들어가며

코드 변경 분석

1. Ampere 아키텍처 최적화 (ampere.py)

2. Hopper 아키텍처 최적화 (hopper.py)

왜 이게 좋은가?

결론

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글

1. Ampere 아키텍처 최적화 (`ampere.py`)

2. Hopper 아키텍처 최적화 (`hopper.py`)