[flashinfer] FlashInfer, NVFP4 SVDQuant GEMM 커널 최적화: 아키텍처별 디스패치 도입

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

PR 링크: flashinfer-ai/flashinfer#6986 상태: Merged | 변경: +34798 / -1711

들어가며

대규모 언어 모델(LLM)의 발전과 함께, 모델의 추론 속도 향상은 필수적인 과제가 되었습니다. 특히 행렬 곱셈(GEMM) 연산은 LLM 추론의 핵심 병목 구간 중 하나로, 이를 효율적으로 처리하는 것은 전체 성능에 지대한 영향을 미칩니다. FlashInfer는 이러한 요구에 부응하여 고성능 GPU 커널을 제공하는 라이브러리입니다.

이번 PR은 FlashInfer의 flashinfer.mm_nvfp4_svdquant 함수, 특히 backend="cake" 옵션에서 사용되는 NVFP4 SVDQuant GEMM 커널을 최적화하는 데 중점을 둡니다. 이 커널은 out = BF16(alpha * (deq(A) @ deq(B)^T + D @ L1^T) + bias) 연산을 수행하며, A와 B는 NVFP4 형식이고, D, L1, bias, 출력은 BF16 형식입니다. 이전 버전에서는 제한된 수의 커널만 제공되었으나, 이번 PR을 통해 NVIDIA GB200 (sm_100a) 및 GB300 (sm_103a) 아키텍처에 대한 커널 수가 대폭 증가하고, 각 아키텍처에 최적화된 커널을 선택적으로 실행하는 '아키텍처별 디스패치' 기능이 도입되었습니다.

이 글에서는 해당 PR의 코드 변경사항을 분석하고, 이러한 최적화가 왜 성능 향상으로 이어지는지, 그리고 어떤 일반적인 교훈을 얻을 수 있는지 살펴보겠습니다.

코드 분석

이번 PR의 핵심 변경사항은 flashinfer/cake_nvfp4_svdquant_gemm_backend 디렉토리 내의 Python 코드와 빌드 설정 파일에 집중되어 있습니다.

1. `flashinfer/cake_nvfp4_svdquant_gemm_backend/init.py` 변경사항

이 파일은 NVFP4 SVDQuant GEMM 백엔드의 진입점 역할을 하며, 커널 지원 여부 확인 및 실행 로직을 담당합니다. 주요 변경 내용은 다음과 같습니다.

is_supported 함수 수정:

이전에는 _loader.device_arch(_device_index(device))를 호출하여 장치 아키텍처를 가져왔지만, 실제 select_record 함수 호출 시에는 아키텍처 정보를 전달하지 않았습니다. 이번 PR에서는 select_record 함수에 arch 인자를 추가하여, 현재 장치의 아키텍처 정보를 명시적으로 전달하도록 변경되었습니다.

Before:

def is_supported(
    *, m: int, n: int, k: int, rank: int, has_bias: bool, device: torch.device
) -> bool:
    """Whether ``device`` is an exported target and ``(M, N, K, rank, has_bias)``
    has a route (the shape rule of :func:`cake_jit.select_record`)."""
    device = torch.device(device)
    if device.type != "cuda":
        return False
    try:
        _loader.device_arch(_device_index(device))
        select_record(m=m, n=n, k=k, rank=rank, has_bias=has_bias)
        return True
    except (RuntimeError, ValueError):
        return False

After:

def is_supported(
    *, m: int, n: int, k: int, rank: int, has_bias: bool, device: torch.device
) -> bool:
    """Whether ``device`` is an exported target and ``(M, N, K, rank, has_bias)``
    has a route on its architecture (the shape rule of
    :func:`cake_jit.select_record`)."""
    device = torch.device(device)
    if device.type != "cuda":
        return False
    try:
        arch = _loader.device_arch(_device_index(device))
        select_record(m=m, n=n, k=k, rank=rank, has_bias=has_bias, arch=arch)
        return True
    except (RuntimeError, ValueError):
        return False

run 함수 수정:

run 함수에서도 마찬가지로 select_record 호출 시 현재 장치의 아키텍처(arch)를 전달하도록 변경되었습니다. 이를 통해 특정 GPU 아키텍처에 더 적합한 커널을 선택할 수 있게 됩니다.

Before:

name = select_record(m=m, n=n, k=k, rank=rank, has_bias=has_bias)
...
device_index = _device_index(a.device)

After:

device_index = _device_index(a.device)
name = select_record(
    m=m,
    n=n,
    k=k,
    rank=rank,
    has_bias=has_bias,
    arch=_loader.device_arch(device_index),
)
...

이 변경은 select_record 함수가 단순히 입력 텐서의 shape 정보뿐만 아니라, 실행되는 GPU의 아키텍처 정보까지 고려하여 최적의 커널을 선택하도록 로직을 확장한 것입니다.

2. `flashinfer/cake_nvfp4_svdquant_gemm_backend/build.json` 변경사항

이 파일은 컴파일될 CUDA 커널들의 빌드 설정을 정의합니다. 이전에는 각 아키텍처(sm_100a, sm_103a)별로 소수의 커널만 포함되어 있었으나, 이번 PR에서는 커널 수가 크게 증가했습니다.

이전 상태 (일부 발췌):

{
  "arch": "sm_100a",
  "closure_sha256": "54e22f25645ebddead40ca5f3c07ac41ba1cfd80b1a71115a3a3c9fe0deba606",
  "entry": "run",
  "gencode": "-gencode=arch=compute_100a,code=sm_100a",
  "min_cuda": "12.8",
  "name": "cake_nvfp4_svdquant_gemm_3c870d4578b91c2d2310_sm_100a",
  ...
}

변경 후 상태 (일부 발췌):

{
  "arch": "sm_100a",
  "closure_sha256": "cbadb417875f1279242e2f9506be0c09062c97bd42ba95f66ba8f4a26c23103a",
  "entry": "run",
  "gencode": "-gencode=arch=compute_100a,code=sm_100a",
  "min_cuda": "12.8",
  "name": "cake_nvfp4_svdquant_gemm_0bd3a4870db1dc91bce4_sm_100a",
  ...
},
{
  "arch": "sm_103a",
  "closure_sha256": "4310843c4e4865a60e5b262f97af4e35d37b27daf5779ee0b43caaff61757ccd",
  "entry": "run",
  "gencode": "-gencode=arch=compute_103a,code=sm_103a",
  "min_cuda": "12.9",
  "name": "cake_nvfp4_svdquant_gemm_0bd3a4870db1dc91bce4_sm_103a",
  ...
},
{
  "arch": "sm_100a",
  "closure_sha256": "863051b50bc22e14a3a61db6f8fd436b065d139eb8377da6aa622d3370923b94",
  "entry": "run",
  "gencode": "-gencode=arch=compute_100a,code=sm_100a",
  "min_cuda": "12.8",
  "name": "cake_nvfp4_svdquant_gemm_1d78f2bd5722fb02b75a_sm_100a",
  ...
}
// ... (총 34개의 커널 정의) ...

이전에는 각 아키텍처별로 15개의 커널이 제공되었으나, 이제는 총 34개의 커널이 제공됩니다. 또한, 각 커널은 특정 아키텍처(sm_100a 또는 sm_103a)에 명시적으로 연결됩니다. 이는 build.json 파일에 정의된 각 라이브러리 항목이 특정 GPU 아키텍처에 최적화된 커널을 나타냄을 의미합니다.

이러한 변경은 select_record 함수가 특정 입력 shape뿐만 아니라, 실행 환경의 GPU 아키텍처까지 고려하여 가장 성능이 좋은 커널을 선택할 수 있도록 기반을 마련합니다. 즉, 동일한 입력 shape라도 sm_100a와 sm_103a에서 서로 다른 최적화된 커널이 실행될 수 있습니다.

왜 이게 좋은가?

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

성능 향상: PR 설명에 포함된 벤치마크 결과는 이 최적화가 실제 성능 향상으로 이어졌음을 명확히 보여줍니다. 예를 들어, NVIDIA GB200 (sm_100a)에서 M=129, N=3072, K=3072, rank=32, bias=1인 경우, 이전 Cake 백엔드 대비 1.065배의 속도 향상(6.34µs -> 5.95µs)을 달성했습니다. GB300 (sm_103a)에서도 유사한 성능 향상이 관찰되며, 대부분의 경우 1.01배에서 1.08배 사이의 속도 향상을 보입니다.

이러한 성능 향상은 다음과 같은 요인들의 복합적인 결과입니다:
- 커널 수 증가: 이전 15개에서 34개로 커널 수가 늘어나면서, 더 다양한 입력 shape와 파라미터 조합에 대해 최적화된 커널을 사용할 수 있게 되었습니다.
- 아키텍처별 튜닝: 각 GPU 아키텍처(sm_100a, sm_103a)의 특성에 맞춰 커널을 재튜닝함으로써, 해당 아키텍처에서 더 높은 성능을 끌어낼 수 있습니다.
- 아키텍처별 디스패치: select_record 함수가 GPU 아키텍처를 고려하여 커널을 선택하도록 함으로써, 각 GPU에 가장 적합한 커널이 실행되도록 보장합니다.
일반화된 최적화 전략: 이 PR은 특정 하드웨어 아키텍처에 대한 최적화가 얼마나 중요한지를 잘 보여줍니다. GPU 아키텍처는 계속 발전하며, 각 아키텍처는 고유한 성능 특성과 명령어셋을 가집니다. 따라서, 특정 아키텍처에 맞춰 커널을 세밀하게 튜닝하고, 해당 아키텍처에서만 실행되도록 디스패치하는 전략은 고성능 컴퓨팅 라이브러리에서 매우 효과적입니다.

이는 단순히 코드를 일반화하는 것을 넘어, 하드웨어의 잠재력을 최대한 활용하기 위한 엔지니어링적 접근 방식입니다. LLM과 같이 연산 집약적인 워크로드에서는 이러한 미세한 최적화가 누적되어 상당한 성능 차이를 만들어낼 수 있습니다.
코드 유지보수성 및 확장성: build.json 파일에 커널 빌드 설정을 명확하게 분리하고, Python 코드에서는 아키텍처 정보를 전달하여 커널을 선택하는 방식은 코드의 구조를 깔끔하게 유지합니다. 새로운 아키텍처가 지원되거나 기존 아키텍처에 대한 커널 튜닝이 필요할 때, build.json 파일을 수정하고 해당 커널을 컴파일하는 방식으로 비교적 쉽게 확장할 수 있습니다.

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer의 Kimi-K3 AttnRes 연산자 최적화: SM100/SM103를 위한 성능 향상
현재글 : [flashinfer] FlashInfer, NVFP4 SVDQuant GEMM 커널 최적화: 아키텍처별 디스패치 도입
다음글 [flashinfer] Blackwell 아키텍처 최적화: FlashInfer의 Cake RMSNorm 성능 개선 분석