WIDE: 토큰 레벨의 동적 폭(width) 가지치기를 통한 적응형 LLM 추론 성능 향상

WIDE: Boosting Adaptive LLM Inference via Token-level Dynamic Width Pruning

가지치기는 LLM의 효율성을 향상시키는 유망한 방법입니다. 기존의 정적인 구조 기반 가지치기 방법은 하드웨어 친화적이며 실질적인 처리량 증가를 제공하지만, 입력 데이터에 독립적인 계산 할당 방식으로 인해 공격적인 희소성(sparsity) 수준에서 상당한 정확도 저하가 발생할 수 있습니다. 최근 동적 희소성 방법들은 개별 입력에 따라 계산을 조정하여 품질 유지율을 개선하지만, 여전히 대부분 거친 수준의 구조적 결정에 국한되며 실제 추론 시나리오에서의 실질적인 가속화는 어려운 과제입니다. 이러한 문제점을 해결하기 위해, 본 논문에서는 프리필(prefill) 및 디코딩(decode) 시나리오 모두에 적용 가능한 최초의 엔드-투-엔드 미분 가능 토큰 레벨 동적 폭 가지치기 프레임워크인 WIDE를 제안합니다. WIDE는 각 토큰이 어텐션 헤드 그룹과 FFN 채널 그룹을 동적으로 선택할 수 있도록 함으로써, 레이어 수준의 결정에서 벗어나 뉴런-블록 수준의 세분화된 계산 할당을 가능하게 합니다. 두 단계의 학습 파이프라인을 통해 WIDE는 효과적인 토큰 단위 희소 실행 패턴을 학습하고 기존 방법보다 훨씬 우수한 품질 유지율을 달성합니다. 이러한 세분화된 동적 가지치기를 실용적으로 구현하기 위해, 우리는 가지치기와 커널 설계를 통합하는 프레임워크를 추가로 제안합니다. 이 프레임워크는 동적 희소성 가속화를 마스크 재정렬, 하드웨어에 독립적인 블록 수준 건너뛰기 및 하드웨어 의존적인 내부 블록 건너뛰기로 분해하여 다양한 세분화 수준에서 효율적인 실행을 가능하게 합니다. 50%의 희소성 수준에서 WIDE는 캘리브레이션(calibration)만 사용했을 때, 최첨단 동적 깊이 가지치기 방법에 비해 55.1%의 성능 향상을 제공합니다. 프리필 및 디코딩 추론 워크로드에서 WIDE는 이론적인 커널 레벨 가속도의 거의 최대인, 각각 1.98배(프리필) 및 4.95배(디코딩)에 가까운 속도 향상과 더불어 전체 시스템적으로 1.68배 및 1.55배의 가속 효과를 달성합니다. 저희 코드들은 https://github.com/EIT-NLP/LLM-Pruning/tree/main/WIDE 에서 확인하실 수 있습니다.

Pruning is a promising approach for improving the efficiency of LLMs. Existing static structured pruning methods are hardware-friendly and can deliver practical throughput gains, but their input-agnostic computation allocation often causes substantial accuracy degradation under aggressive sparsity. Recent dynamic sparsity methods improve quality retention by adapting computation to individual inputs, yet they remain largely limited to coarse-grained structural decisions and their practical acceleration under real-world inference scenarios remains challenging. To address these challenges, we present WIDE, the first end-to-end differentiable token-level dynamic width pruning framework designed for both prefill and decode scenarios. WIDE enables fine-grained computation allocation by allowing each token to dynamically select attention-head groups and FFN-channel groups, extending dynamic pruning beyond layer-level decisions to neuron-block-level granularity. Through a two-stage training pipeline, WIDE learns effective token-wise sparse execution patterns and achieves substantially better quality retention than existing approaches. To make such fine-grained dynamic pruning practical, we further propose a pruning--kernel co-design framework that decomposes dynamic sparsity acceleration into mask reordering, hardware-agnostic block-level skipping, and hardware-dependent intra-block skipping, enabling efficient execution across different granularities. At 50% sparsity, WIDE provides 55.1% performance boost when compared to the state-of-the-art dynamic depth pruning under calibration-only settings. Under prefill and decoding inference workloads, WIDE achieves close-to-theoretical kernel-level speedups of up to 1.98x for prefill and 4.95x for decoding, as well as 1.68x and 1.55x end-to-end acceleration. Our code is available at https://github.com/EIT-NLP/LLM-Pruning/tree/main/WIDE.