[onnxruntime] ONNX Runtime CUDA Packed Sparse Attention 최적화: 성능 향상의 비결

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

PR 링크: microsoft/onnxruntime#33328 상태: Merged | 변경: +652 / -25

들어가며

최근 Microsoft의 ONNX Runtime 레포지토리에서는 [CUDA] Optimize packed QSA scoring and multi-token verification라는 제목의 PR이 올라왔습니다. 이 PR은 특히 CUDA 환경에서 Packed Sparse Attention (QSA)의 스코어링 및 멀티 토큰 검증 과정을 최적화하여, 기존의 비효율적인 연산을 줄이고 전반적인 성능을 향상시키는 것을 목표로 합니다. Transformer 모델, 특히 Qwen과 같이 긴 시퀀스를 처리하는 모델에서 Sparse Attention은 메모리 및 연산 효율성을 높이는 핵심 기법입니다. 하지만 Packed Sparse Attention은 종종 실제 사용되는 토큰보다 훨씬 큰 상태 용량을 예약하게 되는데, 이로 인해 불필요한 연산이 발생하고 성능 저하의 원인이 되었습니다. 본 글에서는 이 PR이 어떤 문제점을 해결하고, 어떻게 코드를 개선했으며, 이러한 변경이 왜 성능 향상으로 이어지는지에 대해 심층적으로 분석해 보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu 파일에 집중되어 있습니다. 주요 최적화 포인트는 다음과 같습니다.

1. 스코어링 로직 개선: 불필요한 연산 제거

기존 QSA 스코어링 커널(QsaScoreTileTopKKernel)은 모든 가능한 토큰 쌍에 대해 스코어를 계산했습니다. 하지만 실제로는 각 쿼리가 접근할 수 있는 키의 범위(causal length)가 제한적이며, 예약된 상태 용량 전체를 사용할 필요가 없는 경우가 많습니다. 이번 PR에서는 이 부분을 개선하여, 각 토큰이 실제로 접근 가능한 키의 범위(block_count)만을 계산하고 해당 범위 내에서만 스코어링을 수행하도록 변경했습니다.

Before:

-  const int64_t total_tiles = static_cast<int64_t>(params.total_tokens) * tile_count;
-  for (int64_t work = blockIdx.x; work < total_tiles; work += gridDim.x) {
-    const int tile = static_cast<int>(work % tile_count);
-    const int token = static_cast<int>(work / tile_count);
-    const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token);
-    const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength];
-    const int64_t abs_position =
-        params.has_position_ids
-            ? position_ids[token]
-            : static_cast<int64_t>(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
-    const int64_t causal_count = SaiCausalThreshold(abs_position, kQwenCompressRatio);
-    const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast<int64_t>(key_len_after);
-    const int block_count = overflow_flags[batch] == 0 ? static_cast<int>(visible_64 < 0 ? 0 : visible_64) : 0;
+  const int token = static_cast<int>(blockIdx.y);
+  const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token);
+  const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength];
+  const int64_t abs_position =
+      params.has_position_ids
+          ? position_ids[token]
+          : static_cast<int64_t>(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
+  const int64_t causal_count = SaiCausalThreshold(abs_position, kQwenCompressRatio);
+  const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast<int64_t>(key_len_after);
+  const int block_count = overflow_flags[batch] == 0 ? static_cast<int>(visible_64 < 0 ? 0 : visible_64) : 0;
+  const int active_tiles = (block_count + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks;
+  const int key_stride = tile_count * kHierarchicalTileBlocks;
+  const int64_t row_base = static_cast<int64_t>(token) * key_stride;
+  for (int index = active_tiles * kHierarchicalTileBlocks + blockIdx.x * blockDim.x + threadIdx.x;
+       index < key_stride; index += gridDim.x * blockDim.x) {
+    tile_keys[row_base + index] = topk::kPaddingSortKey;
+  }
+
+  for (int tile = static_cast<int>(blockIdx.x); tile < active_tiles; tile += static_cast<int>(gridDim.x)) {
     const int block_index = tile * kHierarchicalTileBlocks + block_slot;
     float score = 0.0f;
     if (block_index < block_count) {
@@ -538,7 +543,7 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float*
       } 
     }
     __syncthreads();
-    const int64_t output_base = work * kHierarchicalTileBlocks;
+    const int64_t output_base = row_base + tile * kHierarchicalTileBlocks;
     if (threadIdx.x < kHierarchicalTileBlocks) {
       tile_keys[output_base + threadIdx.x] = keys[threadIdx.x];
     }

기존 코드에서는 total_tiles를 계산하고 work 변수를 통해 모든 가능한 token과 tile 조합을 순회했습니다. 하지만 변경된 코드에서는 blockIdx.y를 통해 token을 직접 가져오고, 계산된 block_count를 기반으로 active_tiles를 결정하여 실제 유효한 스코어링 작업만 수행합니다. 또한, tile_keys 초기화 로직이 추가되어 불필요한 패딩 키를 미리 설정함으로써 후속 작업의 효율성을 높입니다.

2. 멀티 토큰 검증 최적화: 쿼리 재사용 및 상태 압축

멀티 토큰 검증(Multi-Token Verification, MTP) 시나리오에서는 여러 쿼리가 키(Key) 상태를 공유할 수 있습니다. 이 PR은 이러한 상황에서 쿼리별로 키 벡터를 재사용하고, 후보 리스트를 압축하는 로직을 추가했습니다. QsaPairedScoreTileTopKKernel 커널은 두 개의 쿼리(first_token, first_token + 1)를 동시에 처리하며, 이들의 키 상태 접근 범위를 계산하고(block_counts), 이를 기반으로 스코어링 및 키 재사용을 수행합니다.

Before (개념적, 실제 diff는 더 복잡함): 각 쿼리가 독립적으로 키 상태에 접근하고 스코어를 계산했을 가능성이 높음.

After:

+template <typename T>
+__global__ void QsaPairedScoreTileTopKKernel(const T* present_key_state, const float* query_rotated, 
+                                             const int32_t* cumulative_sequence_lengths, 
+                                             const int32_t* past_sequence_lengths, const int64_t* position_ids, 
+                                             const int32_t* present_state_lengths, const int32_t* overflow_flags, 
+                                             uint64_t* tile_keys, int tile_count, 
+                                             PackedSparseAttentionIndexerParams params) {
+  __shared__ uint64_t keys[2][kHierarchicalTileBlocks];
+  __shared__ int block_counts[2];
+  const int first_token = static_cast<int>(blockIdx.y) * 2;
+  // ... (중략) ...
+  const int visible = max(block_counts[0], block_counts[1]);
+  const int active_tiles = (visible + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks;
+  for (int tile = static_cast<int>(blockIdx.x); tile < active_tiles; tile += static_cast<int>(gridDim.x)) {
+    // ... (중략) ...
+    for (int query_index = 0; query_index < 2; ++query_index) {
+      float total = 0.0f;
+      // ... (중략) ...
+      uint64_t key = topk::kPaddingSortKey;
+      if (block < block_counts[query_index]) { // 각 쿼리의 유효 범위 내에서만 키 생성
+        const float scaled = total * params.scale;
+        key = topk::PackStableSortKey(scaled == 0.0f ? 0.0f : scaled, block);
+      }
+      if (lane == 0) {
+        keys[query_index][warp] = key;
+      }
+    }
+    // ... (중략) ...
+    if (threadIdx.x < 2 * kHierarchicalTileBlocks) {
+      const int query_index = static_cast<int>(threadIdx.x) / kHierarchicalTileBlocks;
+      const int slot = static_cast<int>(threadIdx.x) % kHierarchicalTileBlocks;
+      if (first_token + query_index < params.total_tokens &&
+          tile * kHierarchicalTileBlocks < block_counts[query_index]) { // 각 쿼리의 유효 범위 내에서만 결과 저장
+        const int64_t output_base = static_cast<int64_t>(first_token + query_index) * key_stride +
+                                    tile * kHierarchicalTileBlocks;
+        tile_keys[output_base + slot] = keys[query_index][slot];
+      }
+    }
+    __syncthreads();
+  }
+}

QsaPairedScoreTileTopKKernel에서는 두 쿼리의 block_counts를 계산하고, 이를 바탕으로 visible 범위를 결정합니다. 이후 각 쿼리의 block_counts[query_index]를 사용하여 해당 쿼리의 유효 범위 내에서만 keys를 생성하고 tile_keys에 저장합니다. 이는 불필요한 키 계산 및 저장을 방지하여 효율성을 높입니다. 또한, QsaCompactTileTopKKernel 커널은 이전 단계에서 생성된 input_keys를 받아, kBoundedTopKMax 크기로 압축된 output_keys를 생성합니다. 이는 후속 병합 단계에서 처리해야 할 데이터 양을 줄여줍니다.

3. 테스트 커버리지 확장

PR 설명에 따르면, 이번 변경 사항은 다음과 같은 시나리오에 대한 회귀 테스트 커버리지를 확장했습니다:

Ragged batches: 길이가 다른 배치 처리
Empty requests: 빈 요청 처리
Multi-query verification: 멀티 쿼리 검증

이는 다양한 엣지 케이스에서도 최적화된 코드가 올바르게 동작함을 보장합니다.

왜 이게 좋은가?

이 PR의 최적화는 다음과 같은 이유로 매우 훌륭합니다.

불필요한 연산 제거: 스코어링 단계에서 실제 접근 가능한 키의 범위만을 계산하고 처리함으로써, GPU 연산량을 획기적으로 줄였습니다. 이는 특히 긴 시퀀스나 큰 배치 크기에서 상당한 성능 향상을 가져올 수 있습니다.
메모리 대역폭 절약: 쿼리별 키 재사용 및 후보 리스트 압축은 GPU 메모리 접근 횟수를 줄이고, 필요한 데이터만 처리하도록 하여 메모리 대역폭 사용을 최적화합니다. 이는 GPU 성능의 병목 현상을 완화하는 데 중요합니다.
호스트 오버헤드 감소: 이전에는 호스트로 메타데이터를 읽어와야 했던 부분을 디바이스 단에서 처리하도록 변경하여, 호스트-디바이스 간 통신 오버헤드를 줄였습니다.
일반화 가능성: 이러한 최적화 기법들은 Packed Sparse Attention을 사용하는 다른 모델이나 프레임워크에도 적용될 수 있는 일반적인 교훈을 제공합니다. 즉, 실제 사용 범위를 정확히 파악하고, 디바이스 단에서 최대한 많은 연산을 처리하며, 데이터 압축을 통해 후속 단계를 효율화하는 것이 중요합니다.

리뷰 댓글에서 언급된 것처럼, 테스트 커버리지 확장은 이러한 최적화가 다양한 실제 사용 사례에서도 안정적으로 작동함을 보장하는 데 필수적입니다. 특히 packed_sparse_attention_indexer_op_test.cc의 테스트 케이스 수정은 past length와 visible blocks 간의 경계 조건을 조절하여, 최적화된 로직이 올바르게 동작하는지 검증하는 좋은 예시입니다.

결론

이번 ONNX Runtime PR은 CUDA 환경에서의 Packed Sparse Attention 성능을 크게 향상시키는 중요한 개선을 이루었습니다. 불필요한 연산 제거, 쿼리 재사용, 데이터 압축 등 다각적인 최적화를 통해 GPU 활용률을 높이고 응답 속도를 단축했습니다. 이는 Transformer 기반 모델, 특히 긴 시퀀스를 처리하는 모델의 효율성을 높이는 데 크게 기여할 것입니다. 이러한 최적화는 단순히 특정 연산의 속도를 높이는 것을 넘어, GPU 컴퓨팅의 근본적인 효율성을 개선하는 방향을 제시합니다.

참고 자료

https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer, NVIDIA 최신 GPU를 위한 AlphaMoE NVFP4 커널 최적화로 성능 대폭 향상
현재글 : [onnxruntime] ONNX Runtime CUDA Packed Sparse Attention 최적화: 성능 향상의 비결
다음글 [onnxruntime] ONNX Runtime CUDA QMoE 및 FP8 최적화: 성능 향상 및 메모리 효율성 증대

댓글

관련 포스트

PR Analysis 의 다른글