[vllm] vLLM 기술 딥다이브: CUTLASS를 활용한 NVFP4 Linear 커널의 Batch Invariance 최적화

PR 링크: vllm-project/vllm#39912 상태: Merged | 변경: +210 / -49

들어가며

LLM(대규모 언어 모델) 서빙 엔진인 vLLM에서 Batch Invariance(배치 불변성)는 매우 중요한 속성입니다. 이는 추론 시 배치 크기(Batch Size)가 달라지더라도 개별 요청에 대한 수치적 결과가 동일하게 유지되어야 함을 의미합니다. 만약 배치 크기에 따라 결과값이 미세하게 달라진다면, 디버깅이 어려워지고 모델의 결정론적(Deterministic) 동작을 보장할 수 없게 됩니다.

최근 NVIDIA Blackwell 아키텍처(SM100, SM120)에서 도입된 NVFP4(4-bit Floating Point)는 높은 연산 효율을 제공하지만, 이를 활용하는 커널이 배치 크기에 따라 다른 타일링(Tiling) 전략이나 스케줄링을 사용하게 되면 수치적 오차가 발생할 가능성이 큽니다.

이번 포스트에서는 vLLM 프로젝트의 최근 PR을 통해, CUTLASS 라이브러리를 사용하여 NVFP4 Linear 커널에 명시적인 Batch Invariance를 구현하고, Triton 커널 대비 성능을 대폭 향상시킨 기법을 분석해 보겠습니다.

---

코드 분석: 무엇이 바뀌었는가?

1. 명시적인 Tile Scheduler 도입 및 설정 고정

기존 코드에서는 CUTLASS의 KernelScheduleAuto에 의존하여 시스템이 자동으로 스케줄러를 선택하게 했습니다. 하지만 배치 불변성을 보장하려면 배치 크기(M dimension)에 관계없이 동일한 타일링과 스케줄링 로직이 적용되어야 합니다.

Before:

struct sm100_fp4_config_default {
  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
  using TileShape = Shape<_256, _256, _256>;
  // ...
};

After:

struct sm100_fp4_config_default {
  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
  // 명시적인 PersistentScheduler 도입
  using TileScheduler = cutlass::gemm::PersistentScheduler;
  using TileShape = Shape<_256, _256, _256>;
  // ...
};

변경 후에는 cutlass::gemm::PersistentScheduler를 명시적으로 지정했습니다. PersistentScheduler는 SM(Streaming Multiprocessor)에 타일을 지속적으로 할당하여 배치 크기가 작더라도 대규모 배치와 동일한 타일링 구조를 유지하는 데 유리합니다.

2. Batch Invariant 모드 전용 디스패치 로직

런타임에 vllm_is_batch_invariant() 플래그를 확인하여, 해당 모드가 활성화된 경우 M 크기에 상관없이 가장 큰 타일 설정을 사용하는 sm100_fp4_config_default로 강제 고정합니다.

Before (nvfp4_scaled_mm_kernels.cu):

// M 크기에 따라 동적으로 설정을 변경함
if (mp2 <= 16) {
    runGemm<Fp4GemmSm100<sm100_fp4_config_M16, OutType>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else if (mp2 <= 256) {
    runGemm<Fp4GemmSm100<sm100_fp4_config_M256, OutType>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else {
    runGemm<Fp4GemmSm100<sm100_fp4_config_default, OutType>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
}

After (nvfp4_scaled_mm_kernels.cu):

if (vllm::vllm_is_batch_invariant()) {
    using BiGemm = Fp4GemmSm100<sm100_fp4_config_default, OutType>;
    // 스케줄러가 변경되어 수치적 불변성이 깨지는 것을 방지하는 안전장치
    static_assert(
        cute::is_same_v<typename BiGemm::TileScheduler,
                        cutlass::gemm::PersistentScheduler>,
        "batch_invariant requires a persistent tile scheduler; stream-K or "
        "split-K would break numerical invariance");
    runGemm<BiGemm>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
    return;
}
// ... 기존의 동적 디스패치 로직

여기서 static_assert를 사용한 점이 인상적입니다. 향후 다른 개발자가 성능 최적화를 위해 Stream-K나 Split-K 같은 스케줄러로 변경할 경우, 배치 불변성이 깨질 수 있음을 컴파일 타임에 경고하도록 설계되었습니다.

---

왜 이게 좋은 최적화인가?

1. Triton 대비 압도적인 성능

PR 작성자가 공유한 벤치마크 결과에 따르면, CUTLASS 기반의 Batch Invariant 커널은 기존 Triton 구현체보다 모든 시나리오에서 우수한 성능을 보였습니다.

• Decode (Batch 128):
  • Triton: 334.39 TFLOP/s
  • CUTLASS (Batch-invariant): 426.94 TFLOP/s (약 27% 향상)
• Prefill (4096x8192):
  • Triton: 1321.99 TFLOP/s
  • CUTLASS (Batch-invariant): 1963.86 TFLOP/s (약 48% 향상)

일반적으로 Batch Invariance를 강제하면 최적의 타일 크기를 사용하지 못해 성능 손실이 발생할 것으로 예상되지만, CUTLASS의 고도로 최적화된 라이브러리 구조 덕분에 Triton의 유연한 커널보다 훨씬 높은 처리량을 기록했습니다.

2. 수치적 안정성과 결정론적 결과

Floating-point 연산은 결합 법칙이 성립하지 않습니다 ((a + b) + c != a + (b + c)). 따라서 연산 순서가 바뀌면 결과가 달라집니다. PersistentScheduler와 고정된 TileShape를 사용함으로써, 배치 크기가 1이든 128이든 동일한 순서로 부분합(Partial Sum)이 계산되도록 보장합니다. 이는 모델의 신뢰성을 높이는 핵심 요소입니다.

3. 유지보수성 (Defensive Programming)

리뷰 과정에서 jzakrzew는 "배치 불변성 모드가 일반 모드와 다른 코드 경로를 타게 되므로, 이에 대한 회귀 테스트(Regression Test)가 반드시 필요하다"고 강조했습니다. 비록 CUTLASS 자체는 검증된 라이브러리지만, vllm의 통합 로직(Integration Logic)에서 발생할 수 있는 실수를 방지하기 위해 전용 유닛 테스트를 추가한 점은 시니어 엔지니어다운 면모를 보여줍니다.

---

결론

이번 PR은 단순히 성능을 높이는 것을 넘어, "성능과 결정론적 동작 사이의 균형"을 어떻게 잡아야 하는지 잘 보여주는 사례입니다.

1. 명시적 설정: 자동화된 최적화(Auto)에만 의존하지 않고, 특정 요구사항(Batch Invariance)을 위해 설정을 명시적으로 제어했습니다.
2. 안전장치: static_assert를 통해 미래의 코드 변경이 의도치 않게 핵심 기능을 망가뜨리지 않도록 방어했습니다.
3. 검증: 벤치마크를 통해 Triton 대비 우위를 증명하고, 전용 테스트 케이스를 통해 로직의 정확성을 확보했습니다.

Blackwell과 같은 차세대 하드웨어에서 NVFP4와 같은 새로운 데이터 타입을 다룰 때, 이러한 정교한 커널 제어 기법은 LLM 엔진의 경쟁력을 결정짓는 중요한 요소가 될 것입니다.

참고 자료

• https://github.com/NVIDIA/cutlass
• https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#floating-point-instructions-fp4
• https://pytorch.org/docs/stable/generated/torch.mm.html

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