기술 요약: BASTION

문제 정의

대규모 언어 모델 (LLM) 은 자기회귀 (AR) 디코딩의 순차적 특성으로 인해 높은 서비스 비용을 겪습니다. 이는 생성된 토큰마다 별도의 타겟 모델 순전파 (forward pass) 가 필요하기 때문입니다. 추측적 디코딩 (Speculative decoding) 은 경량 드래프터 (drafter) 를 사용하여 여러 후보 토큰을 제안하고, 타겟 모델이 이를 병렬로 검증함으로써 이를 해결합니다.

최근 블록 확산 드래프터 (block-diffusion drafters) (예: DFlash, TiDAR) 의 발전으로 단일 단계에서 여러 미래 토큰을 병렬 예측할 수 있게 되었으나, 이들은 중요한 한계를 도입합니다: 이러한 모델들은 완전히 조건부인 시퀀스 대신 위치별 주변분포 (position-wise marginals) 를 생성합니다. 결과적으로, 단일 탐욕적 (top-1) 경로에 의존하는 것은 타겟 모델의 분포 하에서 비일관된 시퀀스를 초래하는 경우가 많습니다. 기존 트리 기반 추측적 디코딩 방법들 (예: EAGLE-3) 은 일반적으로 특정 쿼리나 하드웨어 제약에 적응하지 않는 정적, 휴리스틱으로 사전 정의된 트리 위상 (static, heuristically predefined tree topologies) 에 의존하여, 드래프트 품질과 검증 오버헤드 간의 비최적 균형을 초래합니다.

방법론: BASTION

BASTION(Budget-Aware Speculative Decoding with Tree-structured Block Diffusion Drafting, 예산 인지 트리 구조 블록 확산 드래프팅을 통한 추측적 디코딩) 은 디코딩 속도 향상을 극대화하기 위해 쿼리 의존적 트리 위상을 동적으로 구성하는 학습 없는 (training-free) 프레임워크입니다. 이는 세 가지 시너지 구성 요소를 통합합니다:

1. 트리 기반 블록 확산 드래프팅

토큰을 순차적으로 생성하는 자기회귀 드래프터와 달리, BASTION 은 단일 순전파에서 위치별 주변분포 $\{q_k(\cdot)\}$ 를 생성하기 위해 블록 확산 드래프터를 사용합니다.

• 위상 분리 (Topology Decoupling): 블록이 병렬로 처리되므로 드래프팅 비용 ($T_{draft}$) 은 트리 구조 (체인, 넓은 트리, 또는 깊은 트리) 와 무관합니다.
• 주변분포 독립성: 검증된 컨텍스트가 주어지면 서로 다른 위치의 후보들은 서로 독립적이므로, 이를 자유롭게 접두사 트리 $\mathcal{T}$ 로 결합할 수 있습니다.
• 경로 점수 매기기: 깊이 $d(i)$ 에 있는 트리의 각 노드 $i$ 는 드래프터 분포 하에서 경로의 결합 확률을 나타내는 경로 점수 $\rho(i) = \prod_{k=1}^{d(i)} q_k(x_{ik})$ 를 할당받습니다.

2. 수용 대용량 및 최우선 확장

타겟 모델 검증의 비용을 들이지 않고 검증 트리에 어떤 노드를 포함할지 결정하기 위해, BASTION 은 경로 신뢰도 대용량 (path-confidence surrogate) 을 사용합니다.

• 대용량 추정: 기대 수용 길이 $\mathbb{E}[A(\mathcal{T})]$ 는 경로 점수의 합으로 근사됩니다: $\hat{A}(\mathcal{T}) = \sum_{i \in \mathcal{T}} \rho(i)$. 이는 높은 확률의 드래프터 경로가 타겟 모델의 선호 궤적과 일치한다는 전제에 기반합니다.
• 최적 구성: 논문은 고정된 예산 하에서 $\hat{A}(\mathcal{T})$ 를 극대화하기 위해 최우선 확장 (best-first expansion) 전략이 증명적으로 최적임을 입증합니다. 이 알고리즘은 가장 높은 $\rho(i)$ 를 가진 유효 노드를 반복적으로 추가하여, 트리가 가장 유망한 궤적을 따라 먼저 성장하도록 보장합니다. 이는 트리 크기와 기대 수용 간의 오목 관계를 초래하며, 체감 수익을 보입니다.

3. 하드웨어 인지 온라인 예산 제어기

BASTION 은 수용 길이의 한계적 이득과 검증 지연 시간의 한계적 증가를 균형 있게 고려하여 각 디코딩 단계에서 최적의 트리 크기 $N$ 을 동적으로 결정합니다.

• 지연 시간 모델링: 이 프레임워크는 메모리 대역폭과 연산 제약을 고려하며, 커널 퓨전 및 스케줄링 오버헤드를 보정하기 위해 관찰된 실행 시간을 통해 보정된 하드웨어 인지 로프라인 모델 (hardware-aware roofline model) 을 사용하여 검증 지연 시간 $T_{verify}(N)$ 을 추정합니다.
• 적응형 중단: 제어기는 추정된 속도 향상 $S(N) = \frac{\hat{A}(\mathcal{T}_N) \cdot L_{AR}}{C(N)}$ 이 감소하기 시작할 때까지 트리를 확장합니다. 이는 추가 노드를 검증하는 비용이 더 긴 수용 시퀀스의 이득을 초과할 때 시스템이 확장을 정확히 중단하도록 보장합니다.

주요 기여

1. 트리 기반 블록 확산 드래프팅: 단일 경로 탐욕적 디코딩을 넘어, 블록 확산 로그잇에서 고품질 접두사 폐쇄형 후보 트리를 구축하기 위해 드래프터 로그잇 경로 점수를 활용하는 최우선 접두사 트리 구성 알고리즘을 도입했습니다.
2. 하드웨어 인지 온라인 예산 제어기: 기대 수용 이득과 하드웨어 보정 검증 지연 시간을 균형 있게 고려하여 각 단계에서 최적의 트리 크기를 동적으로 결정하는 새로운 제어기로, 정적 고정 크기 트리의 필요성을 제거합니다.
3. 광범위한 경험적 검증: 다양한 타겟 모델 (Qwen3, Llama-3.1), 벤치마크 (수학, 코드, 채팅), 그리고 GPU 아키텍처 (A100, H100, A6000, RTX PRO 6000) 에 대한 포괄적 평가를 통해 설정별 튜닝 없이 상당한 속도 향상을 입증했습니다.

실험 결과

NVIDIA A100 GPU 에서 탐욕적 디코딩 (temperature 0) 으로 평가:

• 속도 향상: BASTION 은 Qwen3-8B 에서 표준 자기회귀 디코딩 대비 평균 6.61 배의 속도 향상을 달성했습니다.
• 베이스라인 비교:
  • EAGLE-3(최첨단 AR 기반 트리 디코딩) 보다 2.45 배 우월합니다 (평균 6.61 배 vs 2.70 배).
  • DFlash(최첨단 블록 확산 단일 경로) 보다 1.39 배 우월합니다 (평균 6.61 배 vs 4.76 배).
• 수용 길이: BASTION 은 단일 경로 베이스라인에 비해 훨씬 더 큰 평균 수용 길이 (AAL) 를 달성합니다 (예: Qwen3-8B 에서 8.48 대 6.24). 이는 트리 구조 검증이 더 긴 유효 연속을 성공적으로 포착함을 확인시켜 줍니다.
• 일반화: 성능 향상은 다양한 GPU 세대 (H100 및 Blackwell 아키텍처 포함) 와 모델 계열에 걸쳐 일관되게 나타나며, 이 방법이 특정 하드웨어나 모델 구성에 종속되지 않음을 나타냅니다.

중요성 및 주장

이 논문은 BASTION 이 블록 확산 드래프터의 병렬 효율성과 타겟 모델이 요구하는 순차적 일관성 간의 긴장을 해결함으로써 추측적 디코딩에서 중요한 발전을 이루었다고 주장합니다. 정적 휴리스틱 트리 형태를 동적, 예산 인지 구성으로 대체함으로써, BASTION 은 사용 가능한 병렬 연산 자원을 완전히 활용하면서도 저확률 경로의 과도한 검증으로 인한 지연 시간 패널티를 피합니다.

중요하게도, 이 방법은 학습이 필요 없으며 설정별 튜닝이 필요하지 않아 다양한 하드웨어 및 모델 구성에 즉시 배포 가능합니다. 저자들은 그들의 접근 방식이 블록 확산 드래프팅의 이론적 잠재력과 실제 고처리량 추론 간의 간극을 효과적으로 메우며, 실시간 시스템 제약에 맞춰 드래프트 구조를 적응적으로 정렬함으로써 거의 최적의 속도 향상을 달성한다고 주장합니다.