[onnxruntime] WebGPU 성능 최적화: Graph Capture 재사용을 위한 Session-level Buffer Pool 도입

PR 링크: microsoft/onnxruntime#28761 상태: Merged | 변경: +277 / -2

들어가며

최근 GenAI 모델을 브라우저나 클라이언트 환경에서 실행할 때 WebGPU의 Graph Capture 기능은 성능 향상의 핵심입니다. 하지만 기존 구조에서는 각 Generator(예: GenAI의 개별 요청)가 자신만의 BufferManager를 소유하고, Generator가 소멸할 때(State destructor) 해당 버퍼 캐시를 모두 폐기했습니다.

이로 인해 다음 Generator가 생성될 때마다 모든 Storage 및 Uniform 버퍼를 처음부터 다시 할당해야 하는 Cold-start Latency와 GPU 메모리 변동(Churn) 문제가 발생했습니다. 이번 PR은 세션 레벨에서 은퇴한 버퍼들을 관리하는 SessionBufferPool을 도입하여, 새로운 Generator가 이전 Generator의 버퍼를 즉시 재사용할 수 있도록 개선했습니다.

코드 분석: 핵심 변경 사항

1. IBufferCacheManager 인터페이스 확장

버퍼를 추출하고 다시 흡수하기 위해 인터페이스에 새로운 가상 함수들이 추가되었습니다.

[Before]

class IBufferCacheManager {
 public:
  virtual ~IBufferCacheManager() = default;
  virtual void OnRefresh(GraphCaptureState graph_capture_state) = 0;
};

[After]

class IBufferCacheManager {
 public:
  // ... 기존 코드 ...
  // 캐시된 버퍼를 추출하여 소유권을 이전함
  virtual std::vector<std::pair<size_t, WGPUBuffer>> ExtractCachedBuffers() {
    return {};
  }

  // 외부에서 기부받은 버퍼를 캐시에 흡수함
  virtual void AbsorbCachedBuffers(std::vector<std::pair<size_t, WGPUBuffer>>&& buffers) {
    for (auto& entry : buffers) {
      if (entry.second) {
        wgpuBufferRelease(entry.second); // 기본적으로는 해제하여 누수 방지
      }
    }
  }
};

2. GraphSimpleCacheManager의 구현

실제 그래프 모드에서 사용되는 캐시 매니저가 이 인터페이스를 구현하여, pending_buffers_와 captured_buffers_에 있는 모든 리소스를 안전하게 추출합니다.

[After]

std::vector<std::pair<size_t, WGPUBuffer>> ExtractCachedBuffers() override {
  std::vector<std::pair<size_t, WGPUBuffer>> result;
  // 1. 일반 버퍼 버킷 비우기
  for (auto& pair : buffers_) {
    for (auto& buffer : pair.second) {
      result.emplace_back(pair.first, buffer);
    }
    pair.second.clear();
  }
  // 2. Pending 및 Captured 버퍼들 추출
  for (auto& buffer : pending_buffers_) {
    result.emplace_back(static_cast<size_t>(wgpuBufferGetSize(buffer)), buffer);
  }
  pending_buffers_.clear();
  // ... captured_buffers_ 도 동일하게 처리 ...
  return result;
}

3. SessionBufferPool: 기부(Donate)와 시딩(Seed)

이 PR의 핵심 로직인 SessionBufferPool은 은퇴하는 BufferManager로부터 버퍼를 기부받고(Donate), 새로운 매니저에게 주입(SeedInto)합니다.

[New File: session_buffer_pool.cc]

void SessionBufferPool::Donate(BufferManager& retiring_mgr) {
  if (max_generations_ == 0) return;

  Slot slot;
  slot.storage = retiring_mgr.StorageCache().ExtractCachedBuffers();
  slot.uniform = retiring_mgr.UniformCache().ExtractCachedBuffers();

  // 용량 초과 시 가장 오래된(FIFO) 슬롯 제거
  while (slots_.size() >= max_generations_) {
    auto& victim = slots_.front();
    ReleaseSlotBuffers(victim.storage);
    slots_.erase(slots_.begin());
  }
  slots_.emplace_back(std::move(slot));
}

void SessionBufferPool::SeedInto(BufferManager& new_mgr) {
  if (slots_.empty()) return;
  // 가장 최근에 기부된(LIFO) 슬롯 사용
  Slot slot = std::move(slots_.back());
  slots_.pop_back();
  new_mgr.StorageCache().AbsorbCachedBuffers(std::move(slot.storage));
  new_mgr.UniformCache().AbsorbCachedBuffers(std::move(slot.uniform));
}

왜 이게 좋은 최적화인가?

1. 극적인 재할당 감소

PR 작성자의 테스트 결과에 따르면, Phi-4 모델 실행 시 두 번째 Generator부터는 storage hits=171 misses=0, uniform hits=296 misses=0을 기록했습니다. 즉, 첫 실행 이후에는 GPU 메모리 할당이 전혀 발생하지 않았음을 의미합니다.

2. 안전한 동시성 설계 (Reviewer Feedback 반영)

리뷰어 hariharans29는 GPU 작업이 진행 중일 때 버퍼를 기부하는 것의 위험성을 지적했습니다. 이에 대해 qjia7은 다음 세 가지 근거로 안전성을 증명했습니다:

• session_mutex_: WebGPU EP는 ConcurrentRunSupported() == false이므로 Run과 ReleaseCapturedGraph가 호스트 측에서 직렬화됩니다.
• Queue Flush: OnRunEnd에서 큐를 플러시하고 검증 모드에서는 펜스(Fence)를 치기 때문에 작업 제출이 보장됩니다.
• WebGPU Queue Ordering: 동일 큐에 제출된 작업은 순서대로 실행되므로, 이전 Generator의 작업이 끝나기 전에 다음 Generator가 동일 버퍼 핸들을 사용하더라도 GPU 타임라인 상에서 데이터 레이스가 발생하지 않습니다.

3. 유연한 메모리 관리

max_generations 옵션을 통해 얼마나 많은 세대의 버퍼를 유지할지 결정할 수 있습니다. 기본값은 1로 설정되어 메모리 풋프린트 증가를 최소화하면서도 즉각적인 재사용 이득을 얻도록 설계되었습니다.

일반적인 교훈

1. Lifecycle Mismatch 해결: 객체의 수명(Generator)과 리소스의 수명(Buffer)이 일치하지 않을 때, 상위 컨텍스트(Session) 레벨의 풀을 도입하는 것은 고전적이지만 매우 강력한 패턴입니다.
2. LIFO/FIFO의 조합: 최신 버퍼를 먼저 사용(LIFO)하고 오래된 것을 먼저 버리는(FIFO) 전략은 모델의 입력 크기 변화(Shape drift)에 적응하기 가장 좋은 구조입니다.
3. 엄격한 옵션 파싱: 리뷰 과정에서 std::from_chars를 사용한 엄격한 파싱이 도입되었습니다. 이는 잘못된 설정값으로 인한 런타임 오류를 방지하는 좋은 습관입니다.

나가며

이번 개선은 특히 WebGPU 기반의 GenAI 서비스에서 초기 로딩 이후의 반응 속도를 크게 개선할 것입니다. 리소스 할당 오버헤드를 줄이는 것은 단순히 속도뿐만 아니라 시스템 전체의 안정성에도 기여합니다.

참고 자료

• https://onnxruntime.ai/docs/execution-providers/WebGPU-ExecutionProvider.html
• https://www.w3.org/TR/webgpu/#dom-wgpuqueue-submit
• https://en.cppreference.com/w/cpp/utility/from_chars

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