본문으로 건너뛰기

[flashinfer] Blackwell과 그 너머를 위한 최적화: FlashInfer의 Chain-Aware State Walk 개선 분석

PR 링크: flashinfer-ai/flashinfer#8521 상태: Merged | 변경: +400 / -35

들어가며

최근 대규모 언어 모델(LLM)의 컨텍스트 윈도우가 비약적으로 확장됨에 따라, 효율적인 어텐션 메커니즘과 상태 업데이트(State Walk)의 중요성이 그 어느 때보다 커졌습니다. FlashInfer 라이브러리는 이러한 흐름 속에서 고성능 커널을 제공하는 핵심 프로젝트 중 하나입니다.

이번에 분석할 PR(#8521)은 비결정론적(non-deterministic) KCP(Kernel-based Context Parallelism) GDN 훈련 패키지의 성능을 최적화하는 내용을 담고 있습니다. 핵심 문제는 "GPU의 SM(Streaming Multiprocessor) 자원을 어떻게 하면 가장 효율적으로 채울 것인가?"입니다. 기존에는 단순히 SM 개수와 워크로드의 크기를 비교하여 커널 설정을 결정했지만, 이 PR은 실제 데이터의 '체인 길이(Chain Length)'를 고려한 Chain-aware 전략을 도입하여 Blackwell(B200) 및 차세대 아키텍처에서 유의미한 성능 향상을 이끌어냈습니다.

코드 분석: 무엇이 어떻게 바뀌었나?

이 PR의 핵심은 cake_kcp.py 파일에 구현된 호스트 측의 커널 런칭 규칙 변경입니다. 특히 Forward와 Backward 과정에서 사용되는 vh(Value Width) 선택 로직이 훨씬 정교해졌습니다.

1. LPT Makespan 기반의 vh=64 판단 로직 도입

기존에는 vh를 결정할 때 단순히 전체 워크로드가 SM 하나에 들어가는지(one-wave)만을 확인했습니다. 변경 후에는 LPT(Longest Processing Time) 스케줄링 개념을 도입하여, 두 개의 CTA(Cooperative Thread Array)를 사용하는 그리드가 가장 긴 세그먼트 내에서 처리가 가능한지를 판단합니다.

Before (기존 로직의 개념적 형태):

# 단순히 n_walk와 SM count를 비교
def fwdh_vh(n_walk, num_heads, sm_count):
    return 64 if int(n_walk) * int(num_heads) <= int(sm_count) else 128

After (변경 후):

def walk_vh64_fits(
    chains: Sequence[tuple[int, int]], num_heads: int, sm_count: int
) -> bool:
    # LPT makespan을 계산하여 2-CTA 그리드가 가장 긴 체인 내에 수용되는지 확인
    two = 2 * int(num_heads)
    makespan = cake_kcp_plan.lpt_makespan(
        [(c, two * n) for c, n in chains], 2 * int(sm_count)
    )
    return makespan <= max(c for c, _ in chains)

def fwdh_vh(..., fits: bool = False) -> int:
    # fits 인자가 추가되어 더 유연한 결정이 가능해짐
    if forced is not None: return int(forced)
    return 64 if int(n_walk) * int(num_heads) <= int(sm_count) or fits else 128

이 변경을 통해 vh=64 모드(세그먼트 및 헤드당 2개의 CTA 사용)가 더 넓은 범위의 워크로드에서 활성화될 수 있게 되었습니다. 이는 GPU 자원을 더 촘촘하게 사용하여 유휴 SM을 줄이는 효과를 줍니다.

2. 아키텍처별 Adjoint State Walk 최적화 (dhu_vh)

Backward 과정인 Adjoint State Walk(kcp_dhu)에서도 아키텍처별 특성을 반영한 세밀한 튜닝이 이루어졌습니다. 특히 Blackwell(sm_100a) 아키텍처에서 긴 체인을 처리할 때의 특이 케이스를 처리하는 로직이 추가되었습니다.

After (cake_kcp.py):

def dhu_vh(
    ...,
    tape: bool = True,
    arch: str | None = None,
    long_chain: bool = False,
) -> int:
    # sm_100a, sm_103a 등 특정 아키텍처에 대한 분기 처리
    two_wave = 64 if int(n_walk) * int(num_heads) * 2 <= 4 * int(sm_count) else 128
    if long_chain and arch in DHU_LONG_CHAIN_ARCHES:
        return two_wave
    if not tape or arch is None or arch in DHU_TAPE_VH64_ARCHES:
        return 64
    return two_wave

여기서 DHU_LONG_CHAIN_MIN_CHUNKS = 400과 같은 임계값을 설정하여, Blackwell에서 세그먼트가 매우 길 경우(400개 이상의 청크) 2-CTA 그리드가 오히려 늦게 끝날 수 있는 상황을 방지합니다. 이는 하드웨어의 특성을 소프트웨어 스케줄링 레벨에서 완벽하게 제어하려는 시니어 엔지니어의 디테일을 보여줍니다.

왜 이게 좋은 최적화인가?

1. 하드웨어 활용도(Utilization) 극대화

이 최적화의 핵심은 Makespan 분석입니다. GPU 커널 최적화에서 단순히 "병렬성을 높인다"는 것보다 중요한 것은 "가장 늦게 끝나는 작업(Critical Path)을 얼마나 빨리 끝내는가"입니다. walk_vh64_fits 함수는 이를 수학적으로 모델링하여, 병렬도를 높였을 때(2-CTA) 전체 실행 시간이 줄어드는 경우에만 해당 모드를 선택하도록 합니다.

2. 성능 수치로 증명된 결과

PR 설명에 포함된 벤치마크 결과는 놀랍습니다.

  • B200 (sm_100a): Forward 연산에서 최대 1.114x (약 11%)의 성능 향상을 보였습니다.
  • GB300 (sm_103a): Forward에서 최대 1.116x, Forward+Backward 통합 연산에서 1.062x의 향상을 달성했습니다.
  • R200 (sm_107a): 역시 Forward에서 1.109x의 성능 이득을 확인했습니다.

특히 주목할 점은 Bitwise Equality를 유지했다는 것입니다. 결과값의 정확도에 단 1비트의 오차도 허용하지 않으면서 순수하게 스케줄링과 자원 할당만으로 이뤄낸 성과입니다.

3. 리뷰어 피드백과 검증 전략

리뷰 과정에서 Single-rank proxy를 활용한 테스트 전략이 돋보였습니다. 실제 다중 노드 환경에서의 통신 노이즈를 배제하고, 순수하게 특정 랭크의 워크로드만을 추출하여 단일 GPU에서 재현(Proxy)함으로써 최적화의 순수 효과를 정밀하게 측정했습니다. 이는 복잡한 분산 환경 최적화에서 매우 권장되는 접근 방식입니다.

결론 및 교훈

FlashInfer의 이번 PR은 고성능 라이브러리가 단순히 커널 코드(CUDA C++)만 잘 짜는 것이 아니라, 호스트 측의 런타임 결정 로직(Python)이 얼마나 중요한지를 잘 보여줍니다.

  1. Workload-Awareness: 데이터의 형상(Shape)뿐만 아니라 내부 구조(Chain length)를 분석하여 커널을 선택해야 합니다.
  2. Architecture-Specific Tuning: 모든 GPU 아키텍처에 동일한 규칙을 적용하는 대신, Blackwell과 같은 최신 칩의 특성을 별도로 관리해야 최상의 성능을 얻을 수 있습니다.
  3. Mathematical Modeling: LPT Makespan과 같은 스케줄링 이론을 실제 코드에 녹여내어 휴리스틱을 과학적 근거로 대체했습니다.

시니어 소프트웨어 엔지니어로서, 이러한 '디테일의 힘'이 모여 전체 시스템의 경쟁력을 만든다는 점을 다시 한번 상기하게 되는 훌륭한 사례였습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글