본문으로 건너뛰기

[flashinfer] NVFP4 GEMM 성능의 한계를 넘어서: L2 Weight Prefetch를 통한 파이프라인 최적화

PR 링크: flashinfer-ai/flashinfer#7424 상태: Merged | 변경: +32088 / -199

들어가며

최신 LLM 추론 가속화에서 가장 중요한 과제 중 하나는 데이터 타입의 경량화와 연산 파이프라인의 효율화입니다. NVIDIA의 Blackwell 아키텍처(SM100, SM103)는 NVFP4라는 초저정밀도 데이터 타입을 지원하며, 이를 통해 연산 처리량을 극대화할 수 있는 잠재력을 가지고 있습니다.

FlashInfer 프로젝트의 cake_mm_fp4는 활성화 함수(Activation)를 토큰별로 NVFP4로 양자화한 뒤, 이를 가중치(Weight)와 행렬 곱셈(GEMM)하는 일련의 과정을 수행합니다. 기존 방식에서는 양자화 커널이 완전히 종료된 후 GEMM 커널이 시작되거나, 의존성 대기 시간 동안 하드웨어 자원이 유휴 상태로 남는 문제가 있었습니다.

이번 PR은 'Entry L2 weight prefetch' 기법을 도입하여, 양자화 커널이 실행되는 동안 GEMM에 필요한 가중치 데이터를 미리 L2 캐시로 로드함으로써 이 유휴 시간을 성능 향상의 기회로 전환했습니다.

코드 분석: 무엇이 변했는가?

1. Prefetch 규칙 정의 (cake_backend.py)

가장 먼저 눈에 띄는 변화는 어떤 상황에서 prefetch를 수행할지 결정하는 규칙(CHAIN_EP_RULES)의 추가입니다. 모든 연산에 prefetch를 적용하는 것이 아니라, 성능 이득이 확실한 특정 M(토큰 수) 범위와 아키텍처를 타겟팅합니다.

# Before: prefetch 관련 설정 없음

# After: 아키텍처별 prefetch 적용 범위 정의
CHAIN_EP_STAGES = 4
CHAIN_EP_RULES: dict[str, dict[str, tuple[int, int]]] = {
    "sm_100a": {
        "m128_bF": (128, 128),
        "n16_sk3_l2256b": (17, 32),
        "n8_aF_l2256b": (1, 8),
        # ... 생략 ...
    },
    "sm_103a": {
        "m128_bF": (128, 128),
        "n32_aF_l2256b": (17, 32),
        # ... 생략 ...
    },
}

여기서 CHAIN_EP_STAGES = 4는 GEMM의 첫 번째 타일 연산을 시작하기 전, 가중치의 첫 4개 K-타일을 미리 가져오겠다는 의지를 보여줍니다.

2. Tactic 선택 로직의 고도화 (cake_backend.py)

기존의 prepare_nvfp4_per_token_chain 함수는 단순히 양자화와 GEMM을 순차적으로 준비했습니다. 변경된 코드에서는 chain_gemm_tactic이라는 헬퍼 함수를 통해 prefetch가 가능한 '전략(Tactic)'을 동적으로 선택합니다.

Before:

def prepare_nvfp4_per_token_chain(
    x, global_scale_inv, workspace, b_fp4, b_sf, out, out_scale=None
) -> NVFP4PerTokenChainRunner:
    quantize = prepare_nvfp4_per_token_quantize(
        x, global_scale_inv, workspace, out_scale=out_scale
    )
    gemm = prepare_mm_fp4_per_token(
        workspace.fp4, workspace.sf, b_fp4, b_sf, workspace.scale, out
    )
    return NVFP4PerTokenChainRunner(quantize, gemm)

After:

def prepare_nvfp4_per_token_chain(
    x, global_scale_inv, workspace, b_fp4, b_sf, out, out_scale=None
) -> NVFP4PerTokenChainRunner:
    # ... (중략) ...
    # 현재 연산 환경에 맞는 GEMM Plan을 먼저 계산
    plan = gemm_plan(
        int(workspace.fp4.shape[0]),
        int(b_fp4.shape[0]),
        2 * int(workspace.fp4.shape[1]),
        out.dtype == torch.float16,
        _device_arch(device),
        _sm_count(device),
    )
    # Prefetch가 적용된 Tactic이 있는지 확인하고 적용
    gemm = prepare_mm_fp4_per_token(
        workspace.fp4,
        workspace.sf,
        b_fp4,
        b_sf,
        workspace.scale,
        out,
        tactic=chain_gemm_tactic(plan), # 최적화된 Tactic 주입
    )
    return NVFP4PerTokenChainRunner(quantize, gemm)

chain_gemm_tactic 함수는 내부적으로 ep_stages 파라미터를 추가하여 커널이 컴파일될 때 prefetch 로직이 포함된 바이너리를 선택하도록 유도합니다.

3. JIT 커널 해시 업데이트 (cake_jit.py)

Prefetch 로직이 추가된 커널은 기존 커널과 바이너리 구조가 다르기 때문에, 새로운 커널 해시 값들이 대거 추가되었습니다. 이는 런타임에 최적화된 커널을 즉시 로드할 수 있게 합니다.

# cake_jit.py에 추가된 새로운 커널 엔트리 예시
"05727367662197581d3c": {
    "arg_plan": "gemm",
    "closure_sha256": {
        "sm_100a": "2d66ac73a27ab95f062447b907c166d4b3a9a8f1b96dc16006614128e7b48045",
        "sm_103a": "9cbc854271a2a747171d2980639968308718018e5e16f4e87e04ed92973a75a4",
    },
},

왜 이게 좋은 최적화인가?

1. Latency Hiding (지연 시간 은닉)

일반적인 GPU 커널 체인에서 GEMM은 입력 데이터(양자화 결과물)가 준비될 때까지 기다려야 합니다. 하지만 가중치(Weight)는 이미 메모리에 존재하며 양자화 결과와 무관합니다.

이 PR은 GEMM 커널의 일부 워프(Warp)를 사용하여 양자화 커널이 도는 동안 가중치를 L2 캐시로 미리 끌어올립니다. 결과적으로 양자화가 끝나고 GEMM이 본격적인 연산을 시작할 때, 메모리 로드 대기 시간 없이 즉시 연산 유닛을 가동할 수 있습니다.

2. 구체적인 성능 향상 수치

PR 설명에 포함된 벤치마크 결과에 따르면, NVIDIA B200(SM100) 및 GB300(SM103)에서 다음과 같은 성능 향상이 관찰되었습니다.

  • B200: 1.010x ~ 1.076x (최대 약 7.6% 향상)
  • GB300: 1.012x ~ 1.081x (최대 약 8.1% 향상)

특히 M=8, M=32와 같이 작은 배치 사이즈(Small-M)에서 성능 향상이 두드러지는데, 이는 LLM 서빙 시 디코딩 단계의 성능을 직접적으로 개선할 수 있음을 의미합니다.

3. 정교한 Tactic Selection

무분별한 prefetch는 오히려 캐시 오염(Cache Thrashing)을 일으키거나 불필요한 메모리 대역폭을 점유할 수 있습니다. 이 PR은 CHAIN_EP_RULES를 통해 검증된 Shape에 대해서만 선택적으로 최적화를 적용하는 신중함을 보여줍니다.

결론 및 교훈

이번 FlashInfer의 최적화는 하드웨어의 특성(L2 캐시 구조)과 소프트웨어의 실행 흐름(Kernel Chain)을 깊게 이해했을 때 가능한 결과입니다.

  1. 의존성 분석의 중요성: 데이터 의존성이 없는 피연산자(Weight)를 먼저 처리하여 전체 파이프라인의 Critical Path를 단축했습니다.
  2. 아키텍처 특화 최적화: Blackwell과 같은 최신 GPU의 기능을 십분 활용하기 위해 JIT 컴파일러와 백엔드 로직을 유연하게 설계했습니다.

시니어 엔지니어로서 우리는 단순히 코드를 짜는 것을 넘어, 하드웨어 자원이 단 1마이크로초라도 놀고 있지 않은지 끊임없이 의심해야 한다는 교훈을 얻을 수 있습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글