[flashinfer] FlashInfer, Blackwell GPU를 위한 최적화된 Dense Projection GEMM 도입

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

PR 링크: flashinfer-ai/flashinfer#5753 상태: Merged | 변경: +90147 / -0

들어가며

최근 머신러닝 모델, 특히 대규모 언어 모델(LLM)의 발전은 GPU 하드웨어의 성능 한계를 끊임없이 시험하고 있습니다. 이러한 모델의 핵심 연산 중 하나인 행렬 곱셈(GEMM)은 병목 현상의 주요 원인이 될 수 있습니다. 본 블로그 글에서는 FlashInfer 라이브러리의 최신 Pull Request(PR)에서 제안된 flashinfer.experimental.dense_projection_gemm 기능에 대해 자세히 살펴보겠습니다. 이 PR은 NVIDIA의 최신 Blackwell 아키텍처(SM100/SM107)를 타겟으로 하여 Dense Projection GEMM 연산을 최적화하는 것을 목표로 합니다. 이를 통해 GLM-5.2와 같은 모델에서 요구하는 복잡한 연산을 더욱 효율적으로 처리할 수 있게 됩니다.

이 PR은 단순히 새로운 기능을 추가하는 것을 넘어, 코드 생성, JIT 컴파일, 런타임 최적화 등 다양한 측면에서 상당한 개선을 이루었습니다. 본 글에서는 코드 변경 사항을 상세히 분석하고, 이러한 변경이 왜 성능 향상으로 이어지는지, 그리고 기술적으로 어떤 의미를 가지는지 심층적으로 탐구할 것입니다.

코드 분석

이번 PR의 핵심은 flashinfer.experimental.dense_projection_gemm 모듈의 도입과 이를 지원하기 위한 다양한 코드베이스의 수정입니다. 주요 변경 사항은 다음과 같습니다.

1. 벤치마크 스크립트 추가 (`benchmarks/bench_cake_dense_projection_gemm.py`)

새로운 최적화된 연산을 검증하고 성능을 측정하기 위한 벤치마크 스크립트가 추가되었습니다. 이 스크립트는 다양한 GLM-5.2 관련 연산(projection, batched MLA, router)에 대해 FlashInfer의 새로운 구현과 PyTorch의 torch.matmul을 비교합니다.

# benchmarks/bench_cake_dense_projection_gemm.py
# ... (스크립트 내용) ...

def row_table(T_values, ops, f32_grads, mla, router):
    """label -> row spec, in the order the rows are reported."""
    rows = {}
    for T in T_values:
        for label in PROJECTION_ROWS:
            for op in ops:
                dtypes = ("bf16",) if op == "fwd" or not f32_grads else ("bf16", "f32")
                for dt in dtypes:
                    rows[f"{label}_{op}_{dt}_t{T}"] = dict(
                        family="proj", row=label, op=op, out_dtype=dt, T=T
                    )
        # ... (MLA 및 Router 관련 로직) ...
    return rows

def make_inputs(spec, seed, device):
    # ... (입력 생성 로직) ...

def cake_arm(spec, inp):
    """The prepared Cake launch of the row (allocation-free replays)."""
    if spec["family"] == "router":
        prepared = cake_backend.prepare_router_fp32_gemm(inp["A"], inp["B"], inp["out"])
    elif spec["family"] == "proj" and spec["op"] == "wgrad":
        prepared = cake_backend.prepare_projection_wgrad(inp["G"], inp["X"], inp["out"])
    else:
        prepared = cake_backend.dense_projection_gemm(inp["A"], inp["B"], inp["out"])
    return prepared

def eager_arm(spec, inp):
    """The convenience entry point of the row (the cached planned launch; ``out`` is passed, so a cache
    hit allocates nothing)."""
    # ... (eager_arm 구현) ...

def torch_arm(spec, inp):
    ""``torch.matmul`` on the same views: ..."""
    # ... (torch_arm 구현) ...

def measure_row(label, spec, args):
    inp = make_inputs(spec, SEED, args.device)
    entry = dict(row=label, spec=spec, flops=inp["flops"])
    try:
        prepared = cake_arm(spec, inp)
        entry.update(
            template=prepared.template,
            module=prepared.module_name,
            grid=list(prepared.grid),
        )
        # ... (측정 로직) ...
    except Exception as e:
        entry["error"] = str(e)
    return entry

# ... (메인 실행 로직) ...

이 스크립트는 PROJECTION_ROWS, MLA_ROWS, ROUTER_K, ROUTER_N 등의 상수를 정의하고, 다양한 연산(fwd, dgrad, wgrad)과 데이터 타입(bf16, f32)에 대한 행(row)을 생성합니다. make_inputs 함수는 각 행에 대한 무작위 입력 데이터를 생성하며, cake_arm, eager_arm, torch_arm 함수는 각각 FlashInfer의 최적화된 구현, 편의 함수, 그리고 PyTorch의 기본 구현을 나타냅니다. measure_row 함수는 각 행에 대해 GPU 시간을 측정하고 정확도를 검증합니다.

2. Cake JIT Registry (`cake_jit.py`)

이 PR은 JIT(Just-In-Time) 컴파일러의 레지스트리 구조를 v2 스키마로 업데이트했습니다. 이전 버전보다 훨씬 간결해진 구조는 ARG_PLANS, MODULES, KERNELS 섹션으로 구성됩니다. 특히 ARG_PLANS는 고유한 인자 계획을 SHA-256 해시값으로 키잉하여 저장함으로써 중복을 제거하고 레지스트리 크기를 크게 줄였습니다.

--- a/flashinfer/experimental/dense_projection_gemm/cake_jit.py
+++ b/flashinfer/experimental/dense_projection_gemm/cake_jit.py
@@ -1,11 +1,4 @@
-# schema v1
-ARG_PLANS = [
-    # (sha256, arg_plan_idx, arg_plan_str)
-    ("00000000000000000000000000000000", 0, "(T, K, N, K, N, K, N, K, N, K, N, K, N, K, N, K, N, K, N)"),
-]
-
-MODULES = [
-    # (sha256, module_name, kernel_translation_unit_name, arg_plan_idx, arch, arch_variant)
-]
-
-KERNELS = [
-    # (sha256, kernel_translation_unit_name, arch, arch_variant, smem_bytes, sm_per_cu, ncu_per_sm, nwarps_per_sm, nthreads_per_warp)
-]
+ARG_PLANS = []
+MODULES = []
+KERNELS = []

이전 버전의 3803줄에서 297줄로 줄어든 것은 코드 관리의 효율성을 크게 높였음을 시사합니다. 이는 각 연산의 인자 계획을 효율적으로 저장하고 검색하는 데 중점을 둔 결과입니다.

3. Cake Backend Planner (`cake_backend.py`)

cake_backend.py 파일에는 새로운 JIT 컴파일러를 위한 플래너 로직이 포함되어 있습니다. 이 플래너는 109개의 규칙 테이블(SM100용 51개, SM107용 58개)을 사용하여 최적의 커널 실행 계획을 결정합니다. 또한, 아키텍처별 기본값과 규칙 병합 로직을 포함하여 Cake 런처의 동작을 모방합니다.

특히, 캐싱된 편의 경로(dense_projection_gemm, projection_*, router_*)는 계획된 실행의 LRU(Least Recently Used) 캐시를 유지하여 반복적인 호출에 대한 오버헤드를 줄입니다. KeyError 가드와 캐시 크기 제한은 이러한 캐싱 메커니즘의 안정성을 보장합니다.

# flashinfer/experimental/dense_projection_gemm/cake_backend.py
# ... (planner 관련 로직) ...

class Planner:
    # ...
    def __init__(self):
        # ... (rule table 초기화)
        self.cached_launches = {}
        self.cache_bound = 128

    def plan(self, A, B, out, **kwargs):
        # ... (계획 로직)
        key = (A.shape, B.shape, out.shape, ...)
        if key in self.cached_launches:
            return self.cached_launches[key]
        # ... (새로운 계획 생성 및 캐싱)
        self.cached_launches[key] = prepared
        if len(self.cached_launches) > self.cache_bound:
            # LRU eviction
            pass
        return prepared

# ... (convenience functions using Planner)

4. 템플릿 폴딩 (Template Folding)

이 PR은 템플릿 폴딩 기법을 사용하여 커널 코드의 중복을 줄였습니다. template_members.py에 정의된 15개의 단일 축 멤버(14개의 불리언 노브와 1개의 int01 축)를 사용하여 30개의 개별 커널 인스턴스화를 15개의 템플릿화된 유닛으로 통합했습니다. 이를 통해 K1 커널 번역 유닛의 수가 84개에서 69개로 줄었습니다.

이 과정에서 48개의 쌍 관계에 대한 폴딩 검증이 수행되었으며, 대부분의 경우 SASS(Streaming Assembler) 수준에서 동일함이 확인되었습니다. 일부는 SM 아키텍처가 다르거나 NVCC 컴파일러의 결함으로 인해 폴딩되지 않았지만, 이는 후속 작업으로 추적될 예정입니다.

5. CUDA 커널 코드 (`csrc/cake_dense_projection_gemm/sm_100a/` 및 `sm_107a/`)

실제 CUDA 커널 코드는 이 PR에서 변경되지 않았습니다. 이는 이전 수정 사항에서 이미 최적화된 상태이며, 이번 PR은 주로 호스트 측 바인딩, JIT 레지스트리, 플래너 등 시스템적인 측면에 집중했음을 의미합니다. 각 모듈은 공유 호스트 헤더와 커널 번역 유닛을 통해 얇은 호스트 바인딩을 제공합니다.

왜 이게 좋은가?

이 PR은 여러 측면에서 중요한 성능 개선과 코드 품질 향상을 가져옵니다.

성능 향상: Blackwell 아키텍처(SM100/SM107)에 특화된 최적화를 통해 Dense Projection GEMM 연산의 속도를 높입니다. PR 설명에 따르면, 이전 개정판에서도 B200 GPU에서 122개 벤치마크 행 중 59개가 베이스라인보다 빨랐으며, 84개는 2% 이내의 성능을 보였습니다. 이번 PR은 이러한 성능을 유지하거나 개선하는 것을 목표로 합니다.
코드 중복 감소 및 관리 용이성: 템플릿 폴딩 기법을 사용하여 유사한 커널 코드의 중복을 줄였습니다. 30개의 개별 커널 인스턴스화가 15개의 템플릿 유닛으로 통합되면서 코드베이스의 크기가 줄고 유지보수가 용이해졌습니다. 이는 개발 및 디버깅 시간을 단축하는 데 기여합니다.
JIT 레지스트리 최적화: JIT 레지스트리 스키마 v2를 도입하여 파일 크기를 3803줄에서 297줄로 대폭 줄였습니다. 이는 커널 검색 및 로딩 시간을 단축하여 런타임 성능에 긍정적인 영향을 미칩니다.
효율적인 런타임 계획: 캐싱된 편의 경로와 LRU 캐시를 통해 반복적인 연산 계획에 드는 오버헤드를 줄였습니다. 이는 특히 동적인 워크로드에서 성능 일관성을 유지하는 데 중요합니다.
일반화된 교훈: 이 PR은 다음과 같은 일반적인 최적화 교훈을 제공합니다.
- 하드웨어 특화 최적화: 최신 GPU 아키텍처의 기능을 최대한 활용하는 커널 코드를 작성하는 것이 중요합니다.
- 코드 재사용 및 추상화: 템플릿 메타프로그래밍과 같은 기법을 사용하여 코드 중복을 줄이고 유지보수성을 높여야 합니다.
- 효율적인 런타임 시스템: JIT 컴파일, 커널 계획 캐싱 등은 딥러닝 프레임워크의 성능에 결정적인 영향을 미칩니다.
- 철저한 벤치마킹 및 검증: 새로운 최적화의 효과를 입증하기 위해 포괄적인 벤치마크와 정확도 검증이 필수적입니다.

리뷰 댓글 분석

제공된 리뷰 댓글은 주로 CI 파이프라인 실행 및 테스트 결과에 대한 자동화된 알림입니다. flashinfer-bot이 테스트 실행을 트리거하고 결과를 보고하는 과정이 반복적으로 나타납니다. 몇몇 댓글에서는 테스트 실패 또는 인프라 문제로 인해 파이프라인이 중단되었음을 알리고 있습니다. 예를 들어, tests.gemm.test_mm_mxfp8 또는 tests.gdn.test_prefill_cudnn_backend와 같은 특정 테스트 모듈에서 실패가 보고되었습니다. 또한, CUDA 버전별 호환성 문제나 인프라 자체의 불안정성으로 인한 실패도 관찰됩니다.

이러한 리뷰 댓글은 PR이 아직 안정화 단계에 있으며, 특히 다양한 GPU 아키텍처와 CUDA 버전에서의 호환성 및 정확성 검증이 더 필요함을 시사합니다. 하지만, 본 PR의 핵심 변경 사항인 dense_projection_gemm 자체의 기능 구현 및 최적화 로직은 이러한 테스트 실패와 직접적인 관련이 없어 보이며, 주로 호스트 측 코드와 JIT 시스템에 집중되어 있습니다. 향후 이러한 테스트 실패를 해결하는 과정에서 라이브러리 전체의 안정성이 더욱 향상될 것으로 기대됩니다.

References

torch.compile - PyTorch의 컴파일 기능 (이 PR의 맥락과 유사한 최적화 기법을 다룸)
NVIDIA Blackwell Architecture - 타겟 하드웨어 아키텍처에 대한 정보
Cake JIT - FlashInfer가 사용하는 JIT 컴파일 시스템 관련 문서 (Cutlass 라이브러리 내)
flashinfer.experimental.dense_projection_gemm - 해당 실험적 기능의 Python 바인딩 (PR에서 직접적으로 수정된 부분은 아니지만, 관련 기능)

참고 자료

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

PR Analysis 의 다른글

이전글 [onnxruntime] ONNX Runtime WebGPU: AMD GPU를 위한 Subgroup Matrix MatMulNBits 최적화
현재글 : [flashinfer] FlashInfer, Blackwell GPU를 위한 최적화된 Dense Projection GEMM 도입
다음글 없음

댓글

관련 포스트

PR Analysis 의 다른글