[vllm] vLLM, FlashInfer BF16 CuTeDSL GEMM 통합으로 저지연 추론 성능 향상

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

PR 링크: vllm-project/vllm#50572 상태: Merged | 변경: +349 / -23

들어가며

최근 대규모 언어 모델(LLM)의 발전과 함께 추론 성능 최적화는 더욱 중요해지고 있습니다. 특히, 모델의 핵심 연산인 행렬 곱셈(GEMM)의 효율성을 높이는 것은 전체 추론 속도에 지대한 영향을 미칩니다. vLLM 프로젝트는 이러한 추론 성능 향상을 위해 지속적으로 노력해왔으며, 이번 PR은 FlashInfer 라이브러리의 최신 BF16 CuTeDSL 저지연 GEMM 커널을 통합하여 unquantized BF16 연산의 성능을 개선하는 것을 목표로 합니다.

이 PR은 특히 SM100 계열 GPU에서 BF16 데이터 타입의 GEMM 연산에 대해 FlashInfer의 CuTeDSL 백엔드를 활용할 수 있도록 지원합니다. 이를 통해 기존 PyTorch의 선형(linear) 연산 대비 더 낮은 지연 시간(low latency)을 달성하고자 합니다. 또한, 특정 조건에 맞는 연산에만 FlashInfer를 적용하고, 그렇지 않은 경우에는 기존 PyTorch 구현으로 안전하게 폴백(fallback)하는 전략을 사용합니다.

코드 분석

이번 PR은 주로 vllm/config/kernel.py, vllm/model_executor/kernels/linear/__init__.py, vllm/model_executor/layers/linear.py, vllm/model_executor/layers/utils.py 파일을 수정하여 FlashInfer BF16 CuTeDSL GEMM을 통합하고 관련 설정을 관리합니다.

1. `vllm/config/kernel.py`: `KernelConfig` 업데이트

KernelConfig 클래스의 linear_backend 옵션 설명이 업데이트되었습니다. 기존에는 양자화된(quantized) 선형 레이어 GEMM 커널에 대한 설명이었으나, 이제는 unquantized BF16 연산을 포함한 전반적인 선형 레이어 GEMM 커널에 대한 설명으로 확장되었습니다. 또한, flashinfer_cutedsl 백엔드가 BF16, NVFP4, MXFP8 데이터 타입을 지원함을 명시했습니다.

Before:

-    linear_backend: LinearBackend = "auto"
-    """Backend for quantized linear layer GEMM kernels. Available options:
+    linear_backend: LinearBackend = "auto"
+    """Backend for linear layer GEMM kernels. Available options:

+    Layer types without an implementation from the requested backend use
+    automatic selection.
+
     - "auto": Automatically select the best backend based on model and hardware
     - "cutlass": Use CUTLASS-based kernels
     - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
-    - "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels (NVFP4, MXFP8)
+    - "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels
+      (BF16, NVFP4, MXFP8)
     - "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels
     - "flashinfer_cudnn": Use FlashInfer with cuDNN kernels
     - "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+)

2. `vllm/model_executor/kernels/linear/init.py`: 백엔드 필터링 로직 개선

_filter_kernels_by_backend 함수가 수정되어, 사용자가 명시적으로 --linear-backend를 설정했을 때 해당 백엔드에 맞는 커널이 없을 경우 에러를 발생시키는 대신 경고 메시지를 출력하고 자동 선택으로 폴백하도록 변경되었습니다. 이는 사용자의 의도를 존중하면서도 호환되지 않는 설정으로 인한 오류를 방지합니다.

Before:

-    return [k for k in kernels if k in backend_kernels]
+    filtered = [kernel for kernel in kernels if kernel in backend_kernels]
+    if not filtered:
+        logger.warning_once(
+            "--linear-backend=%s has no kernel for this linear layer type; "
+            "using automatic selection for that layer type.",
+            backend,
+        )
+        return kernels
+    return filtered

3. `vllm/model_executor/layers/linear.py`: `UnquantizedLinearMethod` 수정

UnquantizedLinearMethod 클래스의 __init__ 메서드에서 _gemm_impl을 초기화할 때 현재 vLLM 설정을 참조하여 dispatch_unquantized_gemm 함수에 linear_backend 인자를 전달하도록 변경되었습니다. 이를 통해 UnquantizedLinearMethod 인스턴스가 생성될 때부터 사용자가 설정한 linear_backend를 인지하고 해당 백엔드에 맞는 GEMM 구현을 사용하게 됩니다.

Before:

 class UnquantizedLinearMethod(LinearMethodBase):
     """Linear method without quantization."""
 
+    def __init__(self) -> None:
+        config = get_current_vllm_config_or_none()
+        linear_backend = (
+            config.kernel_config.linear_backend if config is not None else "auto"
+        )
+        self._gemm_impl = dispatch_unquantized_gemm(linear_backend)
+
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -215,7 +223,7 @@ def apply(
     ) -> torch.Tensor:
         if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
             return linear_batch_invariant(x, layer.weight, bias)
-        return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
+        return self._gemm_impl(layer, x, layer.weight, bias)
 
 
 class LinearBase(PluggableLayer):

4. `vllm/model_executor/layers/utils.py`: FlashInfer BF16 GEMM 통합

이 파일은 가장 핵심적인 변경이 이루어진 곳입니다. FlashInfer의 BF16 CuTeDSL GEMM을 위한 새로운 구현(cuda_flashinfer_bf16_gemm_impl, cuda_flashinfer_bf16_gemm_fake, cuda_flashinfer_bf16_gemm)과 관련 지원 로직(_FlashInferBf16Backend, _can_use_flashinfer_cutedsl_bf16, _get_flashinfer_bf16_backend)이 추가되었습니다.

_can_use_flashinfer_cutedsl_bf16 함수는 SM100 계열 GPU, BF16 데이터 타입, 특정 차원(M <= 32), 정렬 조건 등 FlashInfer CuTeDSL 백엔드를 사용할 수 있는 조건을 검증합니다. 만약 조건이 충족되지 않으면 기존 torch.nn.functional.linear로 폴백합니다.

dispatch_unquantized_gemm 함수는 이제 linear_backend 인자를 받아, 해당 백엔드가 지원되고 조건이 충족되면 FlashInfer BF16 구현을 반환합니다. 지원되지 않거나 조건이 맞지 않으면 경고 메시지를 출력하고 기본 구현으로 돌아갑니다.

추가된 주요 로직 (일부 발췌):

+def _can_use_flashinfer_cutedsl_bf16(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+) -> bool:
+    if not (
+        current_platform.is_cuda() and current_platform.is_device_capability_family(100)
+    ):
+        return False
+    if x.ndim < 1 or weight.ndim != 2:
+        return False
+    if (
+        not x.is_cuda
+        or not weight.is_cuda
+        or x.device != weight.device
+        or x.dtype != torch.bfloat16
+        or weight.dtype != torch.bfloat16
+        or not x.is_contiguous()
+        or not weight.is_contiguous()
+    ):
+        return False
+
+    k = x.shape[-1]
+    n = weight.shape[0]
+    if (
+        k <= 0
+        or n <= 0
+        or weight.shape[1] != k
+        or k % 128 != 0
+        or x.data_ptr() % 32 != 0
+        or weight.data_ptr() % 32 != 0
+    ):
+        return False
+
+    m = x.numel() // k
+    if not 1 <= m <= 32:
+        return False
+    return bias is None or (
+        bias.is_cuda
+        and bias.device == x.device
+        and bias.dtype == torch.bfloat16
+        and bias.ndim == 1
+        and bias.shape[0] == n
+        and bias.is_contiguous()
+    )
+
+
+_FLASHINFER_BF16_BACKENDS = {
+    "flashinfer_cutedsl": _FlashInferBf16Backend(
+        flashinfer_backend="cute-dsl",
+        is_supported=is_flashinfer_cutedsl_bf16_gemm_supported,
+        can_implement=_can_use_flashinfer_cutedsl_bf16,
+    ),
+}
+
+
def _get_flashinfer_bf16_backend(vllm_backend: str) -> _FlashInferBf16Backend:
+    backend_spec = _FLASHINFER_BF16_BACKENDS.get(vllm_backend)
+    if backend_spec is None:
+        supported = ", ".join(sorted(_FLASHINFER_BF16_BACKENDS))
+        raise ValueError(
+            f"Unsupported vLLM FlashInfer BF16 backend {vllm_backend!r}; "
+            f"supported backends: {supported}"
+        )
+    return backend_spec
+
+
def cuda_flashinfer_bf16_gemm_impl(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    pdl: bool,
+    vllm_backend: str,
+) -> torch.Tensor:
+    backend_spec = _get_flashinfer_bf16_backend(vllm_backend)
+    if not backend_spec.can_implement(x, weight, bias):
+        return torch.nn.functional.linear(x, weight, bias)
+
+    k = x.shape[-1]
+    n = weight.shape[0]
+    x_2d = x.view(-1, k)
+    out_2d = flashinfer_bf16_mm(
+        x_2d,
+        weight.t(),
+        bias,
+        pdl,
+        backend_spec.flashinfer_backend,
+    )
+    return out_2d.view(*x.shape[:-1], n)
+
+
def cuda_flashinfer_bf16_gemm_fake(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    pdl: bool,
+    vllm_backend: str,
+) -> torch.Tensor:
+    return x.new_empty((*x.shape[:-1], weight.shape[0]))
+
+
def cuda_flashinfer_bf16_gemm(
+    layer: torch.nn.Module,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    *,
+    vllm_backend: str,
+    pdl: bool,
+) -> torch.Tensor:
+    return torch.ops.vllm.cuda_flashinfer_bf16_gemm(
+        x,
+        weight,
+        bias,
+        pdl,
+        vllm_backend,
+    )
+
+
+direct_register_custom_op(
+    op_name="cuda_flashinfer_bf16_gemm",
+    op_func=cuda_flashinfer_bf16_gemm_impl,
+    fake_impl=cuda_flashinfer_bf16_gemm_fake,
+)
+
+
def dispatch_unquantized_gemm(
+    linear_backend: str = "auto",
+) -> Callable[..., torch.Tensor]:
     if current_platform.is_rocm():
         return rocm_unquantized_gemm
     elif current_platform.is_cpu():
         return cpu_unquantized_gemm
-    else:
+    elif not current_platform.is_cuda():
         return default_unquantized_gemm
+
+    backend_spec = _FLASHINFER_BF16_BACKENDS.get(linear_backend)
+    if backend_spec is None:
+        return default_unquantized_gemm
+
+    if not backend_spec.is_supported():
+        logger.warning_once(
+            "--linear-backend=%s requested FlashInfer mm_bf16 backend %r, "
+            "but it is unavailable on the current hardware or environment; "
+            "using automatic selection for unquantized linear layers.",
+            linear_backend,
+            backend_spec.flashinfer_backend,
+        )
+        return default_unquantized_gemm
+
+    logger.info_once(
+        "Using FlashInfer %s for eligible unquantized BF16 GEMMs.",
+        backend_spec.flashinfer_backend,
+    )
+    return functools.partial(
+        cuda_flashinfer_bf16_gemm,
+        vllm_backend=linear_backend,
+        pdl=current_platform.is_arch_support_pdl(),
+    )

5. `tests/kernels/test_flashinfer_bf16_gemm.py`: 새로운 테스트 파일 추가

FlashInfer BF16 CuTeDSL GEMM의 정확성을 검증하기 위한 새로운 테스트 파일이 추가되었습니다. 이 테스트는 다양한 m, n, k 차원과 bias, PDL(Post-Dynamic-Loss) 사용 여부에 대해 FlashInfer 구현과 PyTorch의 F.linear 구현 간의 결과가 근사치(close)인지 확인합니다.

주요 테스트 로직:

+def test_flashinfer_bf16_cutedsl_correctness(
+    m: int,
+    n: int,
+    k: int,
+    use_bias: bool,
+    pdl: bool,
+) -> None:
+    torch.manual_seed(0)
+    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.1
+    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.1
+    bias = (
+        torch.randn(n, device="cuda", dtype=torch.bfloat16) * 0.1 if use_bias else None
+    )
+
+    actual = layer_utils.cuda_flashinfer_bf16_gemm_impl(
+        x, weight, bias, pdl, "flashinfer_cutedsl"
+    )
+    expected = F.linear(x, weight, bias)
+
+    torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-1)

왜 이게 좋은가?

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

저지연 GEMM 성능 향상: FlashInfer는 GPU 커널 최적화에 특화된 라이브러리로, 특히 BF16 데이터 타입과 CuTeDSL 백엔드를 활용하여 기존 구현 대비 더 낮은 지연 시간(low latency)을 제공할 수 있습니다. 이는 LLM 추론 시 각 토큰 생성 속도를 높이는 데 직접적으로 기여합니다.
BF16 지원 강화: BF16은 FP16보다 넓은 동적 범위와 FP32에 가까운 정밀도를 제공하면서도 메모리 사용량은 FP32의 절반에 불과하여 LLM에서 널리 사용되는 데이터 타입입니다. FlashInfer의 BF16 GEMM 통합은 이러한 BF16 연산의 효율성을 극대화합니다.
유연한 백엔드 선택: --linear-backend 옵션을 통해 사용자가 명시적으로 flashinfer_cutedsl을 선택할 수 있게 함으로써, 특정 하드웨어 및 모델 구성에 최적화된 커널을 선택할 수 있는 유연성을 제공합니다. 또한, 호환되지 않는 설정에 대한 폴백 메커니즘은 안정성을 보장합니다.
코드 품질 및 유지보수성: 리뷰어의 피드백을 반영하여 관련 유틸리티 함수들을 vllm/utils/flashinfer.py로 이동시키고, 백엔드 키를 직접 사용하는 등 코드 구조를 개선하려는 노력이 있었습니다. (비록 최종 diff에는 반영되지 않았지만, 리뷰 과정에서 논의됨)
테스트 커버리지: 새로운 커널 통합에 맞춰 정확성 테스트(test_flashinfer_bf16_gemm.py)를 추가하여 코드 변경의 신뢰성을 높였습니다.

성능 수치:

PR 설명에 따르면, TP16 Qwen3.8 모델 서빙 벤치마크에서 다음과 같은 결과가 나왔습니다:

8K 입력 / 1K 출력, 동시성 1, MTP 없음:
- 10/10 요청 완료
- 114.65 출력 토큰/초 (Output Tokens/s)
- 평균 TPOT (Time Per Output Token): 7.800 ms

이 수치는 이전 대비 상당한 성능 향상을 시사합니다. (정확한 비교 대상 수치가 없어 직접적인 성능 향상률을 명시하기는 어렵지만, 높은 처리량과 낮은 지연 시간을 보여줍니다.)

일반적 교훈:

최신 라이브러리 활용: FlashInfer와 같이 GPU 커널 최적화에 특화된 라이브러리를 적극적으로 통합하는 것은 추론 성능 향상의 핵심입니다.
데이터 타입별 최적화: BF16과 같은 주요 데이터 타입에 대한 최적화는 LLM 성능에 큰 영향을 미칩니다.
조건부 적용 및 폴백: 새로운 최적화 기법을 도입할 때는 항상 특정 조건(하드웨어, 데이터 타입, 텐서 크기 등)을 만족하는 경우에만 적용하고, 그렇지 않은 경우에는 안정적인 기존 구현으로 폴백하는 전략이 중요합니다.
철저한 테스트: 새로운 커널 통합 시에는 반드시 정확성 및 성능 테스트를 포함하여 변경 사항의 신뢰성을 확보해야 합니다.

리뷰 피드백 반영

리뷰 과정에서 Isotr0py는 다음과 같은 유의미한 피드백을 제공했습니다:

_can_use_flashinfer_cutedsl_bf16 함수와 같은 유틸리티를 vllm/utils/flashinfer.py로 이동시키는 것을 제안했습니다. 이는 코드의 모듈성과 재사용성을 높이는 좋은 제안입니다.
_FLASHINFER_BF16_BACKENDS 딕셔너리의 키를 vLLM 백엔드 이름(flashinfer_cutedsl)으로 직접 사용하는 것을 제안하여, 룩업 로직을 단순화할 수 있음을 시사했습니다.
dispatch_unquantized_gemm 함수가 linear_backend 인자를 직접 받도록 하여, 초기화 시점에 백엔드를 결정하는 방식을 제안했습니다. 이는 UnquantizedLinearMethod의 __init__에서 linear_backend를 명시적으로 전달하는 변경으로 이어졌습니다.
_filter_kernels_by_backend 함수에서 커널 매칭 실패 시 warning_once를 사용하도록 제안했습니다. 이는 사용자에게 유용한 정보를 제공하면서도 과도한 로그 출력을 방지합니다.

이러한 피드백들은 코드의 구조, 명확성, 안정성을 향상시키는 데 기여했습니다.

References

참고 자료

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

PR Analysis 의 다른글

이전글 [vllm] vLLM, CUDA 네이티브 SwiGLU 커널 도입으로 Humming MoE 성능 1.4% 향상
현재글 : [vllm] vLLM, FlashInfer BF16 CuTeDSL GEMM 통합으로 저지연 추론 성능 향상
다음글 [Liger-Kernel] Liger-Kernel의 Fused Linear Cross Entropy 성능 최적화: C=16 전략

[vllm] vLLM, FlashInfer BF16 CuTeDSL GEMM 통합으로 저지연 추론 성능 향상

들어가며

코드 분석

1. `vllm/config/kernel.py`: `KernelConfig` 업데이트

2. `vllm/model_executor/kernels/linear/init.py`: 백엔드 필터링 로직 개선

3. `vllm/model_executor/layers/linear.py`: `UnquantizedLinearMethod` 수정

4. `vllm/model_executor/layers/utils.py`: FlashInfer BF16 GEMM 통합

5. `tests/kernels/test_flashinfer_bf16_gemm.py`: 새로운 테스트 파일 추가

왜 이게 좋은가?

리뷰 피드백 반영

References

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글