[vllm] vLLM, H100에서의 QKNorm+RoPE 커널 최적화: 더 나은 성능을 위한 동적 워크로드 분배

2026년 4월 13일수정: 2026년 4월 13일

PR 링크: vllm-project/vllm#37376 상태: Merged | 변경: +None / -None

들어가며

vLLM은 LLM 추론 속도 향상을 위해 다양한 최적화 기법을 적용하고 있습니다. 그중 하나가 여러 연산을 하나의 CUDA 커널로 융합(fusing)하여 GPU 메모리 접근을 줄이고 연산 효율성을 높이는 것입니다. 이번 PR은 특히 QKNorm과 RoPE(Rotary Positional Embedding) 연산을 융합한 커널의 성능, 특히 NVIDIA H100 GPU에서의 성능을 개선하는 데 초점을 맞추고 있습니다.

기존의 융합 커널은 특정 조건, 특히 큰 배치 사이즈에서 개별 헤드(head)당 하나의 워프(warp)를 할당하는 방식이 비효율적이었습니다. 이로 인해 H100 GPU에서 분리된(unfused) 커널보다 성능이 떨어지는 문제가 발생했습니다. 본 PR은 이러한 문제를 해결하기 위해, 입력 토큰 배치 크기에 따라 워프당 처리하는 헤드 수를 동적으로 조절하는 새로운 융합 커널을 도입했습니다.

코드 변경사항 분석

이번 PR의 핵심 변경사항은 csrc/fused_qknorm_rope_kernel.cu 파일에 새로운 융합 커널 fusedQKNormRopeKernelNTokenHeads를 추가하고, 비동기 복사 유틸리티를 위한 csrc/async_util.cuh 파일을 생성한 것입니다.

1. `csrc/async_util.cuh`: 비동기 복사 유틸리티

새롭게 추가된 async_util.cuh 파일은 CUDA 아키텍처 8.0 (Ampere) 이상에서 지원하는 cp.async 명령어를 활용하여 GPU의 비동기 복사 기능을 추상화합니다. 이는 GPU가 데이터를 로드하는 동안 다른 연산을 수행할 수 있도록 하여 메모리 병목 현상을 완화하는 데 중요한 역할을 합니다.

주요 함수들은 다음과 같습니다:

cp_async_shared_global_16_cg: 16바이트 단위로 공유 메모리에서 전역 메모리로 비동기 복사를 수행합니다.
cp_async_shared_global_ca: 지정된 크기(4, 8, 16바이트)로 공유 메모리에서 전역 메모리로 비동기 복사를 수행합니다. ca는 Cache-Agnostic을 의미하며, 캐시 사용 여부에 관계없이 복사를 수행합니다.
cp_async_commit_group: 비동기 복사 명령 그룹을 커밋합니다.
cp_async_wait_group<n>: n개의 비동기 복사 그룹이 완료되기를 기다립니다.

이 유틸리티들은 H100과 같은 최신 GPU에서 제공하는 하드웨어 가속 기능을 활용하여 데이터 로딩 성능을 극대화하는 기반을 마련합니다.

2. `csrc/fused_qknorm_rope_kernel.cu`: 새로운 융합 커널 `fusedQKNormRopeKernelNTokenHeads`

이 PR의 핵심 로직은 fusedQKNormRopeKernelNTokenHeads 커널에 구현되어 있습니다. 이 커널은 기존 fusedQKNormRopeKernel과 달리, HEADS_PER_WARP라는 템플릿 파라미터를 통해 워프당 처리할 헤드 수를 조절할 수 있습니다.

기존 커널의 문제점:

기존 커널은 1-head per warp 패턴을 사용했습니다. 이는 헤드 수가 적을 때는 효율적일 수 있지만, 토큰 배치가 커질수록 워프의 활용률이 낮아지고 메모리 접근 지연 시간이 두드러지는 문제가 있었습니다. 특히 H100의 경우, cp.async와 같은 하드웨어 기능을 활용하지 못해 성능 저하가 발생했습니다.

새로운 커널의 개선점:

fusedQKNormRopeKernelNTokenHeads 커널은 다음과 같은 방식으로 성능을 개선합니다:

동적 워크로드 분배: 커널은 입력 토큰 배치 크기에 따라 워프당 처리할 헤드 수를 동적으로 결정합니다. 배치 크기가 클수록 워프당 더 많은 헤드를 처리하여 GPU 활용률을 높입니다.
```
int const head_chunks_per_token = (total_qk_heads + HEADS_PER_WARP - 1) / HEADS_PER_WARP;
int const warp_global = blockIdx.x * warpsPerBlock + warpId;
int const tokenIdx = warp_global / head_chunks_per_token;
int const headChunk = warp_global % head_chunks_per_token;
int const first_head = headChunk * HEADS_PER_WARP;
int const num_heads_this_warp = (first_head + HEADS_PER_WARP <= total_qk_heads) ? HEADS_PER_WARP : (total_qk_heads - first_head);
```
위 코드는 각 워프가 처리해야 할 토큰 인덱스(tokenIdx)와 해당 워프가 담당할 헤드 범위(first_head부터 num_heads_this_warp까지)를 계산하는 과정을 보여줍니다. HEADS_PER_WARP 값은 오프라인 벤치마킹을 통해 최적화되었습니다.

비동기 메모리 복사 활용: 새로 추가된 async_util.cuh의 함수들을 사용하여 QKV 데이터와 RoPE에 필요한 cos/sin 값을 공유 메모리(shared memory)로 비동기적으로 로드합니다.

// === Group 0: async load all heads' QKV into smem (issued first). ===
for (int k = 0; k < num_heads_this_warp; ++k) {
  // ... (calculate offsets) ...
  char* smem_dst = this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes;
  cp_async_shared_global_ca(smem_dst, reinterpret_cast<const char*>(&qkv[offThread]), elemSizeBytes);
}
cp_async_commit_group();  // commit group 0 (QKV)

// === Group 1: async load cos/sin into smem (issued second). ===
// ... (calculate cache pointer and copy size) ...
for (int copyId = laneId; copyId < num_copies; copyId += 32) {
  char* smem_ptr = reinterpret_cast<char*>(&smem[warpId * rotary_dim]) + copyId * 16;
  const char* glob_ptr = reinterpret_cast<const char*>(cache_ptr) + copyId * 16;
  cp_async_shared_global_16_cg(smem_ptr, glob_ptr);
}
cp_async_commit_group();  // commit group 1 (cos/sin)

// wait<1>: allow at most 1 pending group (group 1) → group 0 (QKV) is done.
cp_async_wait_group<1>();

위 코드는 QKV 데이터를 먼저 공유 메모리로 비동기 로드하고(cp_async_shared_global_ca), 이어서 cos/sin 값을 로드합니다(cp_async_shared_global_16_cg). cp_async_commit_group()으로 각 그룹을 커밋하고, cp_async_wait_group<1>()을 통해 첫 번째 그룹(QKV 로드)이 완료될 때까지 기다리면서 두 번째 그룹(cos/sin 로드)의 완료를 보장합니다. 이를 통해 데이터 로딩과 연산 간의 겹침을 최대화합니다.

레지스터 재사용: QKV 가중치(weight)는 각 워프 내에서 한 번만 레지스터로 로드되어 여러 헤드 연산에 재사용됩니다.
```
// Preload weights into registers once, reused across all heads.
float q_w[numElemsPerThread];
float k_w[numElemsPerThread];
#pragma unroll
for (int i = 0; i < numElemsPerThread; i++) {
  int const dim = laneId * numElemsPerThread + i;
  q_w[i] = Converter::convert(q_weight[dim]);
  k_w[i] = Converter::convert(k_weight[dim]);
}
```
이 부분은 각 스레드가 담당하는 헤드 차원의 가중치를 레지스터로 미리 로드하여, 루프 내에서 반복적인 전역 메모리 접근을 피하게 해줍니다.

파라미터 forced의 역할:

vllm/_custom_ops.py 파일에서 fused_qknorm_rope_kernel_n_token_heads 함수에 forced 파라미터가 추가되었습니다. 리뷰어 ZJY0516의 질문에 대한 답변에 따르면, 이 파라미터는 새로운 동적 워크로드 분배 메커니즘을 활성화할지 여부를 제어합니다. 기본값 -1은 새로운 메커니즘을 사용하도록 하며, 다른 값은 이전의 1 head per warp 방식을 강제하거나 특정 헤드 수를 지정하는 등 사용자에게 유연성을 제공하기 위함입니다.

왜 이게 좋은가?

성능 향상

PR 설명에 포함된 테스트 결과는 이 최적화가 H100 GPU에서 상당한 성능 향상을 가져왔음을 보여줍니다.

전반적인 성능:
- Request throughput (req/s): 7.84 (unfused) vs 7.81 (fused baseline) vs 7.82 (fused new). 큰 차이는 없지만, 기존 융합 커널보다 약간 개선되었습니다.
- Output token throughput (tok/s): 1003.09 (unfused) vs 999.42 (fused baseline) vs 1000.73 (fused new). 역시 미미한 개선을 보입니다.
- Mean TTFT (ms): 113.55 (unfused) vs 116.39 (fused baseline) vs 112.31 (fused new). 새로운 융합 커널이 가장 낮은 TTFT(Time To First Token)를 기록했습니다.
- P99 TTFT (ms): 179.22 (unfused) vs 199.65 (fused baseline) vs 180.09 (fused new). 기존 융합 커널 대비 P99 TTFT가 크게 개선되었습니다.
**세부 벤치마크 결과 (EricccYang 제공):
- fused new vs fused baseline 비교 시, Mean TTFT는 약 2% 감소, Median TTFT도 약 1-3% 감소했습니다.
- 특히, TP=4, BF16, Qwen3-Coder-30B-A3B 모델 환경에서 Rate 5 (100 prompts/rep) 테스트 시 Mean TTFT가 93.73ms에서 61.23ms로 34.7% 감소하고, P99 TTFT는 489.53ms에서 106.91ms로 78.2% 감소하는 놀라운 결과를 보여주었습니다. 이는 작은 배치에서는 큰 차이가 없지만, 배치 크기가 커질수록 새로운 커널의 효율성이 극대화됨을 시사합니다.
Nsys 분석:
- Nsys 프로파일링 결과, 새로운 커널(fused new)은 H100에서 더 낮은 최대 지연 시간(약 39us)을 기록하며, 이는 더 효율적인 워프 활용과 비동기 복사 덕분입니다.
- 리뷰어 ProExpertProg의 언급처럼, 이 최적화는 특히 큰 배치 사이즈에서 지연 시간을 크게 줄일 수 있습니다. 비록 작은 배치에서는 전반적인 처리량 증가가 미미하더라도, TTFT 개선은 사용자 경험에 직접적인 영향을 미칩니다.

일반적인 교훈

하드웨어 특성 활용: 최신 GPU 아키텍처(예: H100의 SM9.0)가 제공하는 새로운 기능(예: cp.async)을 적극적으로 활용하면 성능을 크게 향상시킬 수 있습니다.
동적 워크로드 분배: 고정된 워크로드 분배 방식은 특정 시나리오에서 비효율적일 수 있습니다. 입력 데이터의 특성(예: 배치 크기, 시퀀스 길이)에 따라 워크로드를 동적으로 조절하는 것은 GPU 활용률을 높이는 효과적인 방법입니다.
메모리 계층 구조 최적화: 비동기 메모리 복사(cp.async)와 공유 메모리 활용을 통해 전역 메모리 접근 지연 시간을 숨기는 것은 LLM과 같은 메모리 집약적인 워크로드에서 성능 병목을 해결하는 데 필수적입니다.
벤치마킹의 중요성: 최적화된 파라미터(예: HEADS_PER_WARP 값)는 오프라인 벤치마킹을 통해 결정되어야 합니다. 다양한 워크로드와 하드웨어에서 철저한 테스트를 거쳐야 실제 환경에서의 성능을 보장할 수 있습니다.

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

PR Analysis 의 다른글

이전글 [vllm] vLLM의 XPU 가속을 위한 MXFP8 GEMM 커널 도입 분석
현재글 : [vllm] vLLM, H100에서의 QKNorm+RoPE 커널 최적화: 더 나은 성능을 위한 동적 워크로드 분배
다음글 [cpython] CPython JIT 최적화: 키워드 및 바운드 메서드 호출 성능 개선