[onnxruntime] ONNX Runtime CUDA Paged Attention: BF16/INT8 추론 21배 가속 최적화 분석

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

PR 링크: microsoft/onnxruntime#33538 상태: Merged | 변경: +348 / -62

들어가며

최근 대규모 언어 모델(LLM)의 추론 성능 최적화는 매우 중요한 과제입니다. 특히 Paged Attention과 같은 기술은 KV 캐시 메모리 관리를 효율적으로 하여 긴 컨텍스트 길이에서도 메모리 사용량을 줄이고 처리량을 높이는 데 기여합니다. 여기에 더해, 추론 속도를 획기적으로 개선하는 기술 중 하나가 바로 Speculative Decoding입니다. Speculative Decoding은 작은 모델로 다음 토큰을 미리 예측하고, 큰 모델은 이를 검증하는 방식으로 작동하며, 이 검증 단계의 효율성이 전체 추론 속도에 큰 영향을 미칩니다.

이번에 분석할 ONNX Runtime의 GitHub PR은 이러한 Speculative Decoding의 검증 단계에서 특정 모델 구성(BF16 활성화, INT8 Paged KV 캐시, head_size = 256, group_size = 6 또는 16)이 비효율적인 경로로 폴백(fallback)되던 문제를 해결하고, 이를 21배 가속한 최적화에 대한 내용입니다. 이 PR은 Ornith-1.5-122B와 같은 모델의 추론 성능에 직접적인 영향을 미칩니다.

문제점 분석: 비효율적인 Speculative Verification

이 PR이 해결하고자 하는 핵심 문제는 다음과 같습니다. BF16 쿼리 및 출력, INT8 Paged KV 캐시, head_size = 256, 그리고 group_size = 6 또는 16을 사용하는 모델의 Speculative Verification 단계에서, 기존 xqa_spec_dec_candidate 커널이 이 조합을 지원하지 않아 비효율적인 경로로 폴백되었습니다. 구체적으로 두 가지 주요 비효율이 발생했습니다.

GatherAndExpandPagedKVCache 호출: INT8로 양자화된 KV 캐시를 BF16으로 역양자화(dequantize)하고 스테이징 버퍼로 읽어들이는 과정에서 3.68 ms/step의 오버헤드가 발생했습니다.
flash_fwd_kernel varlen 폴백: GatherAndExpandPagedKVCache를 통해 준비된 데이터를 처리하기 위해 일반적인 flash_fwd_kernel varlen 경로로 폴백되었으며, 이 과정에서 11.49 ms/step이 소요되었습니다.

결과적으로, 32K 컨텍스트에서 12개 레이어의 어텐션 검증에 총 15.17 ms/step이 소요되는 심각한 성능 병목이 있었습니다.

핵심 최적화: H256 BF16-query INT8-cache Paged Speculative XQA

이 PR의 핵심은 BF16 쿼리 및 출력, INT8 Paged KV 캐시, head_size = 256, group_size = 6 또는 16 조합을 위한 새로운 Paged XQA(eXtreme Q-Attention) 특수화 커널인 H256::LaunchXQAPagedSpecDecInt8KernelBF16을 추가하는 것입니다. 이 커널은 INT8 Paged KV 캐시를 직접 사용하여 GatherAndExpandPagedKVCache 단계를 완전히 제거하고, 멀티 블록 split-KV (M_TILESIZE = 32) 방식을 통해 효율적으로 어텐션을 계산합니다.

파일별 코드 변경 분석

1. `onnxruntime/contrib_ops/cuda/bert/paged_attention.cc`: Speculative XQA 후보 조건 확장

가장 중요한 변경 사항 중 하나는 PagedAttention 연산자의 ComputeInternal 메서드에서 xqa_spec_dec_candidate가 활성화되는 조건을 확장한 것입니다. 이전에는 MLFloat16 (FP16) 쿼리 타입과 group_size == 6으로 제한되었던 조건에 BF16 쿼리 타입과 INT8 캐시, head_size == 256, 그리고 group_size가 6 또는 16인 경우를 추가했습니다.

Before:

  const bool xqa_spec_dec_candidate =
      decode_eligible && has_metadata_bounds && per_channel_k_on_xqa &&
      ((quantized_xqa_eligible && std::is_same<T, MLFloat16>::value) || native_spec_xqa_eligible) &&
      ((parameters.head_size == 256) ||
       (parameters.head_size == 128 && parameters.is_causal &&
        kIsInt8Cache && std::is_same<T, MLFloat16>::value)) &&
      group_size == 6 &&
      max_query_len_bound > 1 && max_query_len_bound <= 8;

After:

  const bool bf16_int8_spec_xqa_eligible =
      quantized_xqa_eligible && std::is_same<T, BFloat16>::value && kIsInt8Cache &&
      parameters.head_size == 256 && (group_size == 6 || group_size == 16);
  const bool xqa_spec_dec_candidate =
      decode_eligible && has_metadata_bounds && per_channel_k_on_xqa &&
      ((quantized_xqa_eligible && std::is_same<T, MLFloat16>::value && group_size == 6) ||
       bf16_int8_spec_xqa_eligible ||
       (native_spec_xqa_eligible && group_size == 6)) &&
      ((parameters.head_size == 256) ||
       (parameters.head_size == 128 && parameters.is_causal &&
        kIsInt8Cache && std::is_same<T, MLFloat16>::value)) &&
      max_query_len_bound > 1 && max_query_len_bound <= 8;

새로운 bf16_int8_spec_xqa_eligible 변수를 도입하여 BFloat16 타입의 쿼리와 INT8 캐시, head_size = 256, group_size가 6 또는 16인 경우를 명시적으로 xqa_spec_dec_candidate에 포함시켰습니다. 또한, 기존 MLFloat16 조건에도 group_size == 6을 명시하여 조건을 더 정확하게 분리했습니다.

이와 함께, GetXQAPagedSpecDecRequiredSharedMemoryBytes 및 GetXQAPagedSpecDecWorkspaceSize 함수 호출 시 num_heads, kv_num_heads, is_bf16과 같은 추가 파라미터를 전달하도록 변경하여, 워크스페이스 및 공유 메모리 계산이 새로운 커널의 요구사항에 맞게 동적으로 이루어지도록 했습니다.

Before (GetXQAPagedSpecDecRequiredSharedMemoryBytes):

                                           ? GetXQAPagedSpecDecRequiredSharedMemoryBytes(
                                                 parameters.head_size, xqa_kv_quant_type)

After (GetXQAPagedSpecDecRequiredSharedMemoryBytes):

                                           ? GetXQAPagedSpecDecRequiredSharedMemoryBytes(
                                                 parameters.head_size, parameters.num_heads,
                                                 parameters.kv_num_heads, xqa_kv_quant_type,
                                                 std::is_same<T, BFloat16>::value)

2. `onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu`: 새로운 커널 선언 및 라우팅 로직

이 파일에서는 새로운 LaunchXQAPagedSpecDecInt8KernelBF16 커널을 선언하고, LaunchXQAPagedSpecDecKernel 함수 내에서 이 커널로 요청을 라우팅하는 로직을 추가했습니다.

Before (XQA_PAGED_SPEC_DEC_DECL):

XQA_PAGED_SPEC_DEC_DECL(LaunchXQAPagedSpecDecInt8Kernel);

After (XQA_PAGED_SPEC_DEC_DECL):

XQA_PAGED_SPEC_DEC_DECL(LaunchXQAPagedSpecDecInt8Kernel);
XQA_PAGED_SPEC_DEC_DECL(LaunchXQAPagedSpecDecInt8KernelBF16);

그리고 LaunchXQAPagedSpecDecKernel 내부에서 head_size == 256 && kv_quant_type == XqaQuantType::kInt8 && is_bf16 조건과 group_size가 6 또는 16인 경우를 검사하여 새로운 커널을 호출하도록 라우팅 로직이 업데이트되었습니다.

Before (LaunchXQAPagedSpecDecKernel 조건):

  if (device_prop.major < 8 || !IsXQAPagedSpecDecHeadSizeSupported(head_size, kv_quant_type) ||
      (head_size == 128 && !h128_int8) ||
      kv_num_heads <= 0 || num_heads != 6 * kv_num_heads) {
    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
                           "Speculative paged XQA requires SM80+, group size 6, and either head_size 256 "
                           "or FP16 query/output with I

After (LaunchXQAPagedSpecDecKernel 조건):

  const bool h256_bf16_int8 =
      head_size == 256 && kv_quant_type == XqaQuantType::kInt8 && is_bf16;
  const bool valid_group =
      kv_num_heads > 0 && (num_heads == 6 * kv_num_heads ||
                           (h256_bf16_int8 && num_heads == 16 * kv_num_heads));
  if (device_prop.major < 8 || !IsXQAPagedSpecDecHeadSizeSupported(head_size, kv_quant_type) ||
      (head_size == 128 && !h128_int8) || !valid_group) {
    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
                           "Speculative paged XQA requires SM80+, group size 6, and either head_size 256 "
                           "or FP16 query/output with I

h256_bf16_int8 변수를 추가하여 새로운 최적화 경로의 조건을 명확히 하고, valid_group 조건에 num_heads == 16 * kv_num_heads (즉, group_size = 16)인 경우를 포함시켜 BF16 쿼리/INT8 캐시 조합에 대한 지원을 확장했습니다.

3. `onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_spec_dec_bf16_int8_256.cu`: 새로운 커널 구현

이것은 새로운 Paged XQA 특수화 커널인 H256::LaunchXQAPagedSpecDecInt8KernelBF16의 실제 구현을 담고 있는 새 파일입니다. 이 파일은 HEAD_ELEMS=256, XQA_PAGED_CACHE_ELEM=1 (INT8), XQA_PAGED_INPUT_FP16=0 (BF16), XQA_PAGED_SPEC_DEC=1, XQA_PAGED_GROUP6_AND_16=1과 같은 매크로를 사용하여 컴파일됩니다. 이는 BF16 쿼리와 INT8 캐시를 직접 처리하도록 설계된 고도로 최적화된 CUDA 커널입니다.

4. `docs/contrib_ops/cuda/paged_attention.md`: 문서 업데이트

Paged Attention 문서의 §8과 §13 섹션이 업데이트되어 H256 BF16-query/INT8-cache 조합 (group_size ∈ {6, 16})에 대한 Speculative XQA 지원이 명시되었습니다. 이는 사용자가 새로운 최적화의 존재를 인지하고 활용할 수 있도록 돕습니다.

Before (일부 발췌):

>   directly at `group_size = 6`: H256 supports causal and non-causal attention with matching native
>   FP16/BF16 query and cache types, FP16 query/output with an INT8/FP8 cache, or FP16 query/output
>   with a `PER_CHANNEL` INT4 cache; causal H128 supports FP16 query/output with an INT8 cache.

After (일부 발췌):

>   directly: H256 supports causal and non-causal attention at `group_size = 6` with matching native
>   FP16/BF16 query and cache types, FP16 query/output with an INT8/FP8 cache, or FP16 query/output
>   with a `PER_CHANNEL` INT4 cache, plus BF16 query/output with an INT8 cache at
>   `group_size ∈ {6, 16}`; causal H128 supports FP16 query/output with an INT8 cache at
>   `group_size = 6`.

5. `onnxruntime/test/contrib_ops/paged_attention_op_test.cc`: 단위 테스트 추가

새로운 H256 BF16-query/INT8-cache 조합에 대한 단위 테스트가 추가되었습니다. group_size = 16 및 group_size = 6에 대해 쿼리 너비 2..8, 불규칙 배치, 단편화된 페이지 테이블, 로컬 윈도우, PER_CHANNEL 스케일, CUDA 그래프 재생 등을 포함한 다양한 시나리오를 검증합니다. 이는 새로운 최적화가 올바르게 작동하고 안정적임을 보장하는 데 필수적입니다.

왜 이게 좋은 최적화인가?

이 PR은 다음과 같은 이유로 매우 훌륭한 최적화입니다.

압도적인 성능 향상: 32K 컨텍스트에서 12개 레이어의 어텐션 검증 시간이 15.17 ms/step에서 0.72 ms/step으로 21배 빨라졌습니다. 이는 LLM 추론의 핵심 병목 구간 중 하나를 획기적으로 개선하여 전체 추론 지연 시간을 크게 줄입니다.

Metric (32K Context, 12 Full-Attention Layers)	Before (Gather + Flash varlen)	With H256 BF16/INT8 SpecDec XQA	Improvement
`GatherAndExpandPagedKVCache` (`ms/step`)	3.68 ms	0.00 ms (eliminated)	-3.68 ms
Attention kernel (`ms/step`)	11.49 ms	0.72 ms	16.0x faster
Total 12-layer attention verification (`ms/step`)	15.17 ms	0.72 ms	21.1x faster (-14.45 ms/step)

불필요한 데이터 이동 및 역양자화 제거: 기존에는 INT8 KV 캐시를 BF16으로 역양자화하여 스테이징 버퍼로 옮기는 GatherAndExpandPagedKVCache 단계가 필수적이었습니다. 새로운 커널은 INT8 캐시를 직접 처리함으로써 이 단계를 완전히 제거하여 메모리 대역폭 사용을 줄이고 지연 시간을 단축합니다.
특수화된 커널의 힘: 특정 데이터 타입(BF16, INT8)과 하드웨어(NVIDIA H200, sm_90a)에 최적화된 CUDA 커널을 개발함으로써, 일반적인 커널로는 달성하기 어려운 높은 효율성을 이끌어냈습니다. 이는 GPU의 Tensor Core와 같은 특수 하드웨어 기능을 최대한 활용하는 설계입니다.
광범위한 모델 지원: Ornith-1.5-122B와 같이 BF16 활성화, INT8 Paged KV 캐시, head_size = 256, group_size = 16 또는 6을 사용하는 모델들이 이제 효율적인 Speculative Decoding의 혜택을 받을 수 있게 됩니다. 이는 ONNX Runtime의 LLM 지원 범위를 실질적으로 확장하는 효과가 있습니다.

기술적 논의 및 교훈

PR 리뷰 과정에서 onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_spec_dec_bf16_int8_256.cu 파일에 MSVC C4459 경고 억제(pragma warning(disable : 4459))를 추가하는 논의가 있었습니다. 이는 mha_impl.cuh의 헬퍼 함수들이 mha.h의 설정 상수와 이름이 같은 headGrpSize 파라미터를 선언하여 발생하는 변수 섀도잉(shadowing) 문제 때문입니다. Windows CUDA 빌드 환경에서 경고가 오류로 처리될 경우 빌드가 실패할 수 있으므로, 기존 SPEC_DEC 번역 단위와 동일하게 이 경고를 억제하는 조치가 취해졌습니다.

이러한 논의는 고성능 CUDA 커널 개발 시 단순히 기능 구현을 넘어, 다양한 플랫폼 및 컴파일러 환경에서의 호환성과 빌드 안정성까지 고려해야 함을 보여줍니다. 사소해 보이는 경고 하나가 전체 빌드 프로세스를 멈출 수 있기 때문에, 이러한 디테일에 대한 주의는 견고한 소프트웨어 개발에 필수적입니다.

이 최적화에서 얻을 수 있는 일반적인 교훈은 다음과 같습니다.

병목 지점 정확히 파악: 프로파일링을 통해 GatherAndExpandPagedKVCache와 일반 flash_fwd_kernel이 병목임을 정확히 파악한 것이 최적화의 시작점입니다.
특수화의 가치: 특정 조건(데이터 타입, 헤드 사이즈, 그룹 사이즈)에 대해 일반적인 경로보다 훨씬 효율적인 특수화된 커널을 개발하는 것이 성능 향상에 결정적입니다.
데이터 이동 최소화: GPU 프로그래밍에서 데이터 이동은 가장 비싼 연산 중 하나입니다. INT8 캐시를 직접 처리하여 불필요한 역양자화 및 데이터 복사를 제거한 것이 핵심입니다.

마치며

이번 ONNX Runtime PR은 LLM 추론 최적화의 중요성을 다시 한번 상기시켜 줍니다. Speculative Decoding의 검증 단계에서 특정 모델 구성에 대한 비효율적인 폴백 경로를 제거하고, 고도로 최적화된 CUDA 커널을 도입함으로써 21배라는 놀라운 성능 향상을 달성했습니다. 이는 ONNX Runtime이 LLM 추론 분야에서 지속적으로 선도적인 역할을 하고 있음을 보여주는 좋은 사례이며, 앞으로도 이러한 저수준 최적화를 통해 더 많은 모델들이 고성능으로 서비스될 수 있기를 기대합니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [onnxruntime] ONNX Runtime WebGPU: Subgroup Matrix 연산을 활용한 Conv 최적화
현재글 : [onnxruntime] ONNX Runtime CUDA Paged Attention: BF16/INT8 추론 21배 가속 최적화 분석
다음글 [sglang] ROCm 환경에서의 성능 최적화: Triton을 활용한 Bit-exact QK RMSNorm 및 RoPE 융합