[flashinfer] FlashInfer, 초저병렬성 환경에서의 CP 델타 규칙 사전 계산 최적화

2026년 7월 8일수정: 2026년 7월 8일

PR 링크: flashinfer-ai/flashinfer#3866 상태: Merged | 변경: +718 / -23

들어가며

최근 대규모 언어 모델(LLM)의 발전은 모델의 크기와 복잡성을 증가시키면서, 효율적인 추론(inference)을 위한 최적화 기술의 중요성을 더욱 부각시키고 있습니다. 특히, Attention 메커니즘의 계산 비용은 LLM의 성능 병목 현상을 야기하는 주요 원인 중 하나입니다. FlashInfer는 이러한 문제를 해결하기 위해 KV Cache의 효율적인 관리를 위한 다양한 최적화 기법을 도입해왔습니다.

이번 PR(#898)은 FlashInfer의 Continuous Processing (CP) 델타 규칙(delta rule) 사전 계산(precompute) 단계에서 발생하는 성능 문제를 해결하는 데 초점을 맞추고 있습니다. 기존에는 특정 조건, 특히 초저병렬성(ultra-low parallelism) 환경에서 성능 저하가 관찰되었습니다. 이 PR은 이러한 환경에서도 효율적인 계산을 보장하기 위해 CP 청킹(chunking) 휴리스틱을 개선하고, SIMT (Single Instruction, Multiple Threads) 커널을 새롭게 도입했습니다.

본 글에서는 이 PR이 어떤 문제를 해결하려 했는지, 코드 변경사항은 무엇이며, 왜 이러한 변경이 성능 향상으로 이어지는지, 그리고 이 최적화가 가지는 일반적인 교훈은 무엇인지 심층적으로 분석해보겠습니다.

코드 변경사항 분석

이번 PR의 핵심은 flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py 파일의 수정입니다. 특히, 기존의 CPDeltaRuleFixupSm120 클래스를 CPDeltaRuleFixupHmmaSm120으로 변경하고, 새로운 CPDeltaRuleFixupSimtSm120 클래스를 추가한 것이 주목할 만합니다.

1. `CPDeltaRuleFixupHmmaSm120` 클래스 (기존 클래스 리네이밍)

기존의 CPDeltaRuleFixupSm120 클래스는 CPDeltaRuleFixupHmmaSm120으로 이름이 변경되었습니다. 이는 해당 커널이 HMMA (Half-precision Matrix Multiply Accumulate) 연산을 주로 활용함을 명확히 하기 위한 것으로 보입니다. HMMA는 FP16 또는 BF16과 같은 반정밀도 부동소수점 연산을 사용하여 메모리 대역폭을 절약하고 연산 속도를 높이는 데 효과적입니다. 이는 특히 최신 GPU 아키텍처에서 성능 향상을 가져올 수 있습니다.

-class CPDeltaRuleFixupSm120(KeyedCompileMixin):
+class CPDeltaRuleFixupHmmaSm120(KeyedCompileMixin):

2. `CPDeltaRuleFixupSimtSm120` 클래스 신규 도입

가장 중요한 변경사항은 새로운 CPDeltaRuleFixupSimtSm120 클래스의 도입입니다. 이 클래스는 초저병렬성 환경에서의 성능 최적화를 목표로 합니다. SIMT 아키텍처는 GPU의 기본 병렬 처리 모델로, 각 스레드가 동일한 명령어를 실행하지만 다른 데이터를 처리하는 방식입니다. 이 커널은 특히 다음과 같은 기능들을 포함합니다:

init_state_tile: 초기 상태를 로드하거나 로컬 상태로부터 초기화하는 함수입니다. needs_initial_state 플래그에 따라 동작이 달라지며, 이는 다양한 모델 구조에 대한 유연성을 제공합니다.

@cute.jit
def init_state_tile(
    self,
    sState: cute.Tensor,
    gFixedState: cute.Tensor,
    gLocalState: cute.Tensor,
    gInitialState: cute.Tensor,
    col: cutlass.Int32,
) -> cutlass.Int32:
    start = cutlass.Int32(0)
    if cutlass.const_expr(self.needs_initial_state):
        for i in cutlass.range_constexpr(self.rows_per_cta):
            sState[i, col] = gInitialState[i, col]
    else:
        start = cutlass.Int32(1)
        for i in cutlass.range_constexpr(self.rows_per_cta):
            value = gLocalState[i, col, 0]
            sState[i, col] = value
            gFixedState[i, col, 0] = value
    return start

load_transfer_fragment / load_local_state_fragment: 외부 메모리(global memory)로부터 데이터를 로드하는 함수입니다. GPU의 메모리 계층 구조를 효율적으로 활용하기 위해 데이터를 청크(chunk) 단위로 나누어 로드합니다.

@cute.jit
def load_transfer_fragment(
    self,
    rM: cute.Tensor,
    gTransferTile: cute.Tensor,
    chunk_idx: cutlass.Int32,
    k_tile: cutlass.Int32,
):
    for j in cutlass.range_constexpr(16):
        rM[j] = gTransferTile[(j, chunk_idx), (k_tile, 0)]

accumulate_state_fragment: 로드된 데이터와 현재 상태를 곱하고 누적하는 핵심 연산입니다. SIMT 커널에서는 각 스레드가 특정 컬럼(column)을 담당하며, 이 함수는 해당 컬럼에 대한 누적 계산을 수행합니다.

@cute.jit
def accumulate_state_fragment(
    self,
    rAcc: cute.Tensor,
    sState: cute.Tensor,
    rM: cute.Tensor,
    k: cutlass.Int32,
):
    for i in cutlass.range_constexpr(self.rows_per_cta):
        for j in cutlass.range_constexpr(16):
            rAcc[i] = rAcc[i] + sState[i, k + cutlass.Int32(j)] * rM[j]

run_simt_fixup_loop: 전체 CP 델타 규칙 계산 루프를 담당합니다. 청크 단위로 데이터를 로드하고, 상태를 누적하며, 다음 청크로 상태를 전달하는 과정을 반복합니다. 특히, 마지막 K 타일(tile)에서 다음 청크의 데이터를 미리 로드(preload)하여 파이프라인 효율성을 높이는 기법이 사용됩니다.

@cute.jit
def run_simt_fixup_loop(
    self,
    sState: cute.Tensor,
    gTransfer: cute.Tensor,
    gLocalState: cute.Tensor,
    gFixedState: cute.Tensor,
    num_chunks: cutlass.Int32,
    col: cutlass.Int32,
    start: cutlass.Int32,
):
    # ... (중략) ...
    for chunk_idx in cutlass.range(start, num_chunks, unroll=1):
        # ... (중략) ...
        # Last K tile owns the inter-chunk handoff: preload next chunk
        # before the final accumulation, then publish the new state.
        if next_chunk_idx < num_chunks:
            self.load_local_state_fragment(
                rAccNext,
                gLocalState,
                next_chunk_idx,
                col,
            )
            self.load_transfer_fragment(
                rM_next,
                gTransferTile,
                next_chunk_idx,
                cutlass.Int32(0),
            )
        # ... (중략) ...
        cute.arch.sync_threads()
        for i in cutlass.range_constexpr(self.rows_per_cta):
            value = rAcc[i]
            sState[i, col] = value
            gFixedState[i, col, chunk_idx] = value
        # ... (중략) ...

zero_invalid_slots: 계산 범위 밖의 슬롯(slot)들을 0으로 초기화하여 정확성을 보장합니다. 이는 특히 가변 길이 시퀀스(variable length sequence)를 처리할 때 중요합니다.

@cute.jit
def zero_invalid_slots(
    self,
    gFixedState: cute.Tensor,
    gap_len: cutlass.Int32,
    col: cutlass.Int32,
):
    for slot in cutlass.range(0, gap_len, unroll=1):
        for i in cutlass.range_constexpr(self.rows_per_cta):
            gFixedState[i, col, slot] = cutlass.Float32(0.0)

이 CPDeltaRuleFixupSimtSm120 클래스는 __call__ 메서드를 통해 CUDA 커널을 런치하며, grid와 block 크기를 동적으로 설정하여 다양한 GPU 및 워크로드에 맞게 최적화합니다. 특히 min_blocks_per_mp 파라미터는 SM(Streaming Multiprocessor)당 최소 블록 수를 지정하여 낮은 병렬성 환경에서도 GPU 활용률을 높이려는 시도로 보입니다.

3. CP 청킹 휴리스틱 개선

PR 설명에 따르면, CP 청킹 휴리스틱 또한 개선되었습니다. 이는 시퀀스 길이나 GPU의 특성(예: SM 수, 메모리 대역폭)을 고려하여 최적의 청크 크기를 결정하는 로직을 의미합니다. 초저병렬성 환경에서는 작은 청크 크기가 오히려 오버헤드를 증가시킬 수 있으므로, 이러한 환경에 맞는 새로운 휴리스틱이 필요했을 것입니다.

왜 이게 좋은가?

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

초저병렬성 환경에서의 성능 향상: 리뷰어 노트에 제시된 성능 측정 결과는 이 PR의 효과를 명확히 보여줍니다. 예를 들어, RTX PRO 6000 Blackwell Server Edition (SM120)에서 397B/122B 모델, TP4 설정으로 8192x32 시퀀스 길이를 처리할 때, 기존 대비 6.70x의 속도 향상을 기록했습니다. 이는 초저병렬성 환경에서 기존 커널의 비효율성을 성공적으로 개선했음을 시사합니다.
SIMT 커널 도입: SIMT 아키텍처는 GPU의 기본 병렬 처리 모델입니다. CPDeltaRuleFixupSimtSm120 커널은 이러한 SIMT 모델을 활용하여 각 스레드가 효율적으로 데이터를 처리하도록 설계되었습니다. 특히, 스레드 간의 동기화(cute.arch.sync_threads())와 데이터 로딩/누적 로직을 세밀하게 조정하여 파이프라인 효율성을 극대화했습니다.
메모리 접근 패턴 최적화: load_transfer_fragment, load_local_state_fragment 함수들은 데이터를 청크 단위로 로드하며, run_simt_fixup_loop에서는 다음 청크 데이터를 미리 로드하는 기법을 사용합니다. 이는 GPU의 메모리 대역폭을 최대한 활용하고, 데이터 로딩 지연 시간을 숨기는(hiding latency) 데 효과적입니다.
정확성 보장: zero_invalid_slots 함수는 가변 길이 시퀀스 처리 시 발생할 수 있는 부정확성을 방지합니다. 이는 모델의 예측 정확도를 유지하는 데 필수적입니다.
유연성 및 확장성: needs_initial_state와 같은 파라미터, 그리고 HMMA와 SIMT 커널을 분리한 설계는 다양한 모델 아키텍처와 하드웨어 특성에 대한 유연성을 제공합니다. 또한, KeyedCompileMixin을 사용하여 컴파일 캐싱을 지원함으로써 재컴파일 오버헤드를 줄입니다.

일반적인 교훈

하드웨어 특성 이해의 중요성: GPU 아키텍처(SM90, SM120 등)와 병렬 처리 모델(SIMT, HMMA)에 대한 깊은 이해는 최적화의 핵심입니다. 특정 하드웨어에서 성능이 저하되는 원인을 분석하고, 해당 하드웨어에 최적화된 커널을 개발하는 것이 중요합니다.
메모리 계층 구조 활용: GPU의 메모리 대역폭은 종종 연산 성능보다 병목이 됩니다. 데이터를 효율적으로 로드하고, 공유 메모리(shared memory)를 적극적으로 활용하며, 레이턴시를 숨기는 기법(예: 프리페칭)을 적용하는 것이 성능 향상의 지름길입니다.
동적/적응형 최적화: 고정된 최적화 기법보다는 워크로드(시퀀스 길이, 배치 크기, 병렬성 수준)와 하드웨어 특성에 따라 동적으로 최적화 전략을 선택하는 것이 중요합니다. 이 PR에서 SIMT 커널을 추가하고 청킹 휴리스틱을 개선한 것이 좋은 예입니다.
정확성과 성능의 균형: 성능 최적화 과정에서 모델의 정확성이 저하되지 않도록 주의해야 합니다. zero_invalid_slots와 같은 기법은 이러한 균형을 맞추는 데 도움을 줍니다.

References

CUDA C++ Programming Guide - SIMT 모델 및 GPU 프로그래밍에 대한 기본 정보
CUTLASS Documentation - FlashInfer가 사용하는 고성능 GPU 커널 라이브러리
FlashInfer GitHub Repository - 프로젝트 소스 코드 및 관련 문서

참고 자료

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer의 BF16 GEMM 성능 극대화: CUDA Graph와 Cold L2 Cache 도입
현재글 : [flashinfer] FlashInfer, 초저병렬성 환경에서의 CP 델타 규칙 사전 계산 최적화
다음글 [onnxruntime] ONNX Runtime WebGPU: Intel Xe-3LPG를 위한 고성능 GEMM 최적화 분석