[vllm] vLLM NVFP4 커널 선택 최적화: SM120/121 환경에서의 성능 향상

PR 링크: vllm-project/vllm#55170 상태: Merged | 변경: +34 / -1

들어가며

최신 GPU 아키텍처인 Blackwell(SM120/121) 환경에서 vLLM을 사용하여 NVFP4 모델을 구동할 때, 의도치 않게 성능이 낮은 커널이 선택되는 문제가 발견되었습니다. 본 PR은 NVFP4 가속을 지원하는 환경임에도 불구하고, 시스템이 FP4 연산 능력을 활용하지 못하고 Weight-only(W4A16) 커널로 폴백(fallback)되는 현상을 수정하여, 실제 하드웨어 성능을 최대로 끌어올리는 최적화를 다룹니다.

코드 분석

1. 커널 우선순위 재정렬 (vllm/model_executor/kernels/linear/__init__.py)

문제의 핵심은 _POSSIBLE_NVFP4_KERNELS 리스트의 순서였습니다. 기존 코드에서는 FlashInferCuteDslNvFp4W4A16LinearKernel이 W4A4 커널들보다 앞에 위치해 있었습니다.

Before:

_POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
    PlatformEnum.CUDA: [
        FlashInferCuteDslNvFp4LinearKernel,
        FlashInferCuteDslNvFp4W4A16LinearKernel, # 우선순위가 너무 높음
        FlashInferCutlassNvFp4LinearKernel,
        # ...
    ]
}

After:

_POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
    PlatformEnum.CUDA: [
        FlashInferCuteDslNvFp4LinearKernel,
        FlashInferCutlassNvFp4LinearKernel,
        FlashInferB12xNvFp4LinearKernel,
        CutlassNvFp4LinearKernel,
        FlashInferCuteDslNvFp4W4A16LinearKernel, # 후순위로 이동
        # ...
    ]
}

이 변경을 통해 SM120/121 환경에서 첫 번째 커널이 지원되지 않을 경우, 바로 W4A16으로 빠지는 대신 더 효율적인 W4A4 커널들을 먼저 시도하도록 우선순위를 조정했습니다.

2. 회귀 테스트 추가 (tests/kernels/quantization/test_nvfp4_kernel_selection.py)

향후 코드 변경으로 인해 다시 우선순위가 뒤바뀌는 것을 방지하기 위해, W4A4 커널이 W4A16 커널보다 인덱스상 앞에 위치해야 함을 보장하는 테스트 코드를 추가했습니다.

@pytest.mark.parametrize("w4a4_kernel", W4A4_KERNELS_ON_SM12X)
def test_w4a16_kernel_does_not_precede_w4a4_kernels(w4a4_kernel):
    candidates = _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA]
    w4a16_index = candidates.index(FlashInferCuteDslNvFp4W4A16LinearKernel)
    assert candidates.index(w4a4_kernel) < w4a16_index

왜 이게 좋은가

이번 최적화는 단순한 코드 순서 변경만으로도 극적인 성능 향상을 가져왔습니다. RTX PRO 6000 Blackwell 환경에서 측정 결과, 출력 처리량(Output Throughput)은 23.8% 증가하였고, TTFT(Time To First Token)는 36% 단축되었습니다.

핵심 교훈

1. 하드웨어 가속기 활용: FP4 Tensor Core를 사용할 수 있는 환경에서 Weight-only 커널을 사용하는 것은 하드웨어 자원의 낭비입니다. 커널 선택 로직은 하드웨어의 기능을 최대한 활용하는 방향으로 정렬되어야 합니다.
2. 정적 우선순위의 위험성: 커널 선택 리스트가 고정되어 있을 경우, 새로운 아키텍처가 등장할 때 기존의 '안전한 폴백'이 '최적의 경로'를 가로막을 수 있습니다. 우선순위 리스트는 항상 최신 하드웨어의 성능 특성을 반영해야 합니다.
3. 테스트의 중요성: 커널 선택과 같은 인프라 로직은 의존성이 복잡하므로, 이번 PR처럼 명시적인 인덱스 비교 테스트를 통해 의도한 동작을 강제하는 것이 매우 중요합니다.

리뷰 과정에서 언급되었듯이, 일부 환경에서는 기존 커널 선택 로직으로 인해 호스트 메모리 부족(OOM) 문제가 발생하기도 했습니다. 이번 수정은 단순히 성능을 높이는 것을 넘어 시스템의 안정성까지 확보하는 결과를 낳았습니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#tensor-cores

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