[vllm] [vLLM 성능 최적화] Kimi-K2.5/K2.6 이미지 전처리 10배 가속화: Numba와 퓨전 기법 활용
PR 링크: vllm-project/vllm#47416 상태: Merged | 변경: +402 / -2
들어가며
vLLM은 대규모 언어 모델(LLM)의 서빙 성능을 극대화하는 데 초점을 맞춘 오픈소스 추론 엔진입니다. 최근 멀티모달 LLM의 중요성이 커지면서, 이미지와 같은 비전 데이터를 효율적으로 처리하는 것이 전체 시스템 성능에 큰 영향을 미치게 되었습니다. 특히 Kimi-K2.5/K2.6과 같은 비전 청크(vision chunk)를 사용하는 모델의 경우, 입력 이미지의 전처리 과정이 CPU 바운드 병목 지점이 될 수 있습니다.
기존의 이미지 전처리 파이프라인은 여러 단계(리사이즈, 패딩, 정규화, 패치화)를 개별적으로 수행하며, 각 단계마다 Python 오버헤드와 메모리 접근 비용이 발생하여 비효율적일 수 있습니다. 이러한 문제를 해결하기 위해, 이 PR은 Kimi-K2.5/K2.6 비전 청크를 위한 퓨전(fused) CPU 이미지 전처리 경로를 추가하여 성능을 대폭 개선했습니다. numba 라이브러리를 활용하여 여러 전처리 단계를 하나의 최적화된 커널로 통합하고, 정규화 과정에 룩업 테이블(Lookup Table, LUT)을 적용하여 계산 비용을 줄인 것이 핵심입니다.
이 글에서는 해당 PR의 코드 변경사항을 분석하고, 이러한 최적화가 왜 효과적인지, 그리고 이를 통해 얻을 수 있는 일반적인 교훈은 무엇인지 살펴보겠습니다.
코드 분석
1. vllm/model_executor/models/kimi_k25.py: 퓨전 프로세서 조건부 로딩
이 파일은 Kimi-K2.5 모델의 초기화 로직을 담당합니다. PR의 핵심 변경사항 중 하나는 numba 라이브러리 사용 가능 여부에 따라 최적화된 KimiK25FusedVisionProcessor를 조건부로 로드하도록 한 것입니다. 이는 시스템 환경에 따라 유연하게 대응할 수 있도록 합니다.
Before:
image_processor = cached_get_image_processor(
self.ctx.model_config.model,
revision=self.ctx.model_config.revision,
trust_remote_code=self.ctx.model_config.trust_remote_code,
)
After:
tokenizer = self.get_tokenizer()
processor_cls = KimiK25FusedVisionProcessor if is_numba_available() else None
logger.info_once(
"Using %s image preprocessing for Kimi-K2.5/K2.6 vision chunks.",
"fused CPU" if processor_cls is not None else "remote HF",
)
image_processor = cached_get_image_processor(
self.ctx.model_config.model,
revision=self.ctx.model_config.revision,
trust_remote_code=self.ctx.model_config.trust_remote_code,
processor_cls_overrides=processor_cls,
)
설명:
is_numba_available() 함수를 통해 numba가 설치되어 있는지 확인하고, 만약 가능하다면 KimiK25FusedVisionProcessor를 사용하도록 processor_cls 변수에 할당합니다. 그렇지 않으면 None으로 설정하여 기본 Hugging Face 프로세서를 사용하도록 합니다. 이 로직은 cached_get_image_processor 함수의 processor_cls_overrides 인자를 통해 주입됩니다. 이는 numba가 없는 환경에서도 vLLM이 정상적으로 동작하도록 하는 중요한 폴백(fallback) 메커니즘입니다.
2. vllm/transformers_utils/processor.py: 커스텀 프로세서 주입 기능 추가
이 파일은 Hugging Face AutoImageProcessor를 로드하는 유틸리티 함수를 포함합니다. 위에서 언급한 processor_cls_overrides 인자를 실제로 받아들이고 적용할 수 있도록 수정되었습니다.
Before:
def get_image_processor(
processor_name: str,
*args: Any,
revision: str | None = None,
trust_remote_code: bool = False,
**kwargs: Any,
):
"""Load an image processor for the given model name via HuggingFace."""
try:
processor_name = convert_model_repo_to_path(processor_name)
processor = AutoImageProcessor.from_pretrained(
processor_name,
*args,
revision=revision,
trust_remote_code=trust_remote_code,
**kwargs,
)
After:
def get_image_processor(
*args: Any,
revision: str | None = None,
trust_remote_code: bool = False,
processor_cls_overrides: type[_I] | None = None,
**kwargs: Any,
):
"""Load an image processor for the given model name via HuggingFace."""
try:
processor_name = convert_model_repo_to_path(processor_name)
processor_cls = processor_cls_overrides or AutoImageProcessor
processor = processor_cls.from_pretrained(
processor_name,
*args,
revision=revision,
trust_remote_code=trust_remote_code,
**kwargs,
)
설명:
get_image_processor 함수에 processor_cls_overrides라는 새로운 인자가 추가되었습니다. 이 인자가 제공되면 AutoImageProcessor 대신 해당 클래스를 사용하여 프로세서를 인스턴스화합니다. 이는 vLLM이 Hugging Face의 기본 로딩 메커니즘을 유지하면서도, 특정 모델에 대해 커스텀 최적화된 프로세서를 유연하게 삽입할 수 있는 확장성을 제공합니다. 리뷰어 @Isotr0py의 제안으로 깔끔하게 구현된 부분입니다.
3. vllm/transformers_utils/processors/kimi_k25_vision_fused.py: 퓨전 프로세서 구현 (핵심)
이 파일은 PR의 핵심인 KimiK25FusedVisionProcessor와 numba 기반의 최적화된 함수들을 정의합니다.
_write_fused_patches 함수
이 함수는 numba의 njit 데코레이터를 사용하여 JIT 컴파일되며, parallel=True 옵션으로 병렬 실행을 지원합니다. 이미지의 패딩, 정규화, 패치화 과정을 하나의 루프에서 처리하여 Python 오버헤드를 최소화하고 캐시 효율성을 높입니다.
if is_numba_available():
from numba import njit, prange
@njit(parallel=True, cache=True)
def _write_fused_patches(
frames: np.ndarray,
out: np.ndarray,
out_offset: int,
new_h: int,
new_w: int,
padded_h: int,
padded_w: int,
patch_size: int,
normalize_lut: np.ndarray,
) -> None:
# frames: [T, new_h, new_w, 3] uint8, without padding.
# out: [total_patches, 3, patch_size, patch_size] float32.
t_size = frames.shape[0]
patch_h = padded_h // patch_size
patch_w = padded_w // patch_size
total = t_size * padded_h * padded_w * 3
hwc = padded_h * padded_w * 3
wc = padded_w * 3
for linear in prange(total):
t = linear // hwc
rem = linear - t * hwc
y = rem // wc
rem = rem - y * wc
x = rem // 3
c = rem - x * 3
value = frames[t, y, x, c] if y < new_h and x < new_w else 0
patch_idx = (
out_offset
+ t * patch_h * patch_w
+ (y // patch_size) * patch_w
+ (x // patch_size)
)
out[patch_idx, c, y % patch_size, x % patch_size] = normalize_lut[value, c]
else:
def _write_fused_patches(*args: Any, **kwargs: Any) -> None:
raise RuntimeError("numba is required for fused Kimi image preprocessing")
설명:
@njit(parallel=True, cache=True):numba를 사용하여 Python 코드를 기계어 코드로 컴파일하고, 병렬 실행을 활성화하며, 컴파일된 코드를 캐싱하여 재사용합니다.prange(total):numba의 병렬 루프를 사용하여 이미지의 모든 픽셀을 효율적으로 처리합니다.- 퓨전: 이 단일 루프 내에서 다음 작업들이 동시에 이루어집니다.
- 패딩(Padding):
if y < new_h and x < new_w else 0조건문을 통해 이미지 경계를 벗어나는 픽셀은 0으로 채워집니다. - 정규화(Normalization):
normalize_lut[value, c]를 사용하여uint8픽셀 값을float32로 정규화합니다. 이는 픽셀마다 부동소수점 연산을 반복하는 대신 미리 계산된 룩업 테이블에서 값을 가져오므로 매우 효율적입니다. 이 아이디어는 리뷰어 @Isotr0py의 제안으로 추가되었으며, 추가적인 성능 향상을 가져왔습니다. - 패치화(Patchification):
patch_idx,y % patch_size,x % patch_size계산을 통해 픽셀이 올바른 패치 위치에 배치됩니다.
- 패딩(Padding):
KimiK25FusedVisionProcessor 클래스
이 클래스는 BaseImageProcessor를 상속받아 Kimi-K2.5/K2.6 모델의 이미지 전처리 로직을 구현합니다.
class KimiK25FusedVisionProcessor(BaseImageProcessor):
model_type = "kimi_k25"
def __init__(self, media_proc_cfg: dict[str, Any], **kwargs: Any) -> None:
super().__init__(**kwargs)
media_proc_cfg = dict(media_proc_cfg)
merge_kernel_size = media_proc_cfg["merge_kernel_size"]
if isinstance(merge_kernel_size, (list, tuple)):
media_proc_cfg["merge_kernel_size"] = int(merge_kernel_size[0])
self.media_proc_cfg = media_proc_cfg
self.num_frames_per_chunk = media_proc_cfg["temporal_merge_kernel_size"]
values = np.arange(256, dtype=np.float32)[:, None]
image_mean = np.asarray(media_proc_cfg["image_mean"], dtype=np.float32)
image_std_inv = 1.0 / np.asarray(media_proc_cfg["image_std"], dtype=np.float32)
self.normalize_lut = (values / 255.0 - image_mean[None, :]) * image_std_inv[
None, :
]
# ... (생략: media_tokens_calculator, get_resize_config, resize_image)
def preprocess(
self, medias: list[dict[str, Any]], return_tensors: str | TensorType | None = None,
) -> BatchFeature:
if not isinstance(medias, list):
medias = [medias]
if not medias:
return BatchFeature(data={}, tensor_type=return_tensors)
if njit is None:
raise RuntimeError("numba is required for fused Kimi image preprocessing")
patch_size = int(self.media_proc_cfg["patch_size"])
# ... (중략: 이미지 리사이즈 및 패치 계산 로직)
# _write_fused_patches 호출
_write_fused_patches(
frames_np,
output_patches_np,
output_offset,
new_h,
new_w,
padded_h,
padded_w,
patch_size,
self.normalize_lut,
)
# ... (중략: 결과 텐서 반환 로직)
설명:
__init__메서드에서normalize_lut를 한 번만 계산하여 저장합니다. 이 룩업 테이블은 0부터 255까지의 각 픽셀 값에 대해 미리 정규화된float32값을 담고 있습니다. 이는 런타임에 반복적인 부동소수점 연산을 피하는 핵심 전략입니다.preprocess메서드는 입력 이미지를 PIL을 사용하여 Kimi-K2.5/K2.6의 NaViT 리사이즈 정책에 따라 리사이즈한 후,_write_fused_patches함수를 호출하여 나머지 전처리 단계를numba로 가속화합니다. 이로써 리사이즈는 기존 PIL의 동작을 유지하면서, 이후의 CPU 바운드 작업들은 최적화된 경로를 따르게 됩니다.
왜 이게 좋은가
1. 압도적인 성능 향상
PR 설명에 포함된 벤치마크 결과는 이 최적화의 효과를 명확히 보여줍니다. 8개의 2048x2048 이미지에 대한 벤치마크에서:
- Original processor: 평균 1315.969 ms
- PIL resize + LUT fused: 평균 189.631 ms
- OpenCV resize + LUT fused: 평균 128.490 ms (현재 PR에서는 PIL resize를 사용)
관찰된 속도 향상은 다음과 같습니다:
- PIL LUT fused vs origin: 6.940배
- OpenCV LUT fused vs origin: 10.242배
이러한 성능 향상은 다음과 같은 이유로 가능했습니다:
- 연산 퓨전 (Operation Fusion): 패딩, 정규화, 패치화와 같은 여러 개별 단계를
numba를 통해 하나의 JIT 컴파일된 루프에서 처리함으로써, 각 단계별 Python 인터프리터 오버헤드와 데이터 복사 비용을 제거했습니다. - 룩업 테이블 (Lookup Table): 픽셀 정규화는 0-255 범위의
uint8값을float32로 변환하는 반복적인 연산입니다. 이를__init__단계에서 미리 계산된 룩업 테이블로 대체함으로써, 런타임에 수백만 번의 부동소수점 연산을 피하고 메모리 참조로 대체하여 속도를 크게 높였습니다. - 병렬 처리 (Parallel Processing):
numba.prange를 사용하여 CPU 코어를 최대한 활용, 병렬로 픽셀을 처리함으로써 대규모 이미지 처리 시간을 단축했습니다. - 캐시 효율성 (Cache Efficiency): 데이터를 한 번에 읽고 여러 작업을 수행함으로써 CPU 캐시 활용도를 높여 메모리 접근 지연 시간을 줄였습니다.
2. 견고한 설계와 유연성
- 조건부 활성화:
numba의 가용성에 따라 최적화된 경로와 기존 경로를 자동으로 전환함으로써, 사용자 환경에 대한 유연성과 안정성을 확보했습니다. - 확장성:
processor_cls_overrides인자를 통해cached_get_image_processor함수가 커스텀 프로세서 클래스를 주입할 수 있도록 하여, 향후 다른 모델이나 추가적인 최적화가 필요할 때 vLLM 코어 로직의 변경 없이 쉽게 확장할 수 있는 기반을 마련했습니다. - 수치적 정합성 유지: OpenCV resize가 더 빠르다는 벤치마크 결과가 있었음에도 불구하고, 기존 PIL resize의 수치적 출력과 일치하지 않는다는 이유로 PIL resize를 유지했습니다. 이는 성능 최적화와 더불어 모델의 예측 결과에 영향을 미치지 않도록 수치적 정합성을 중요하게 고려했음을 보여줍니다.
3. 일반적인 교훈
이 PR은 성능 최적화에 대한 몇 가지 중요한 교훈을 제공합니다.
- 병목 지점 식별: 멀티모달 LLM에서 이미지 전처리가 CPU 바운드 병목이 될 수 있음을 인지하는 것이 첫 단계입니다.
- 퓨전의 힘: 여러 작은 연산을 하나의 큰 연산으로 퓨전하는 것은 Python의 오버헤드를 줄이고 캐시 효율성을 높여 성능을 크게 향상시킬 수 있습니다.
- 특수 라이브러리 활용:
numba와 같이 특정 유형의 연산(CPU 바운드 수치 계산)에 최적화된 라이브러리를 적극적으로 활용하는 것이 중요합니다. - 룩업 테이블의 가치: 반복적인 계산이 필요한 경우, 미리 계산된 룩업 테이블을 사용하는 것은 CPU 사이클을 크게 절약할 수 있는 강력한 기법입니다.
- 폴백 메커니즘: 최적화된 경로가 특정 환경에서 작동하지 않을 경우를 대비하여 안정적인 폴백 메커니즘을 제공하는 것이 프로덕션 시스템에서는 필수적입니다.
결론
vLLM의 Kimi-K2.5/K2.6 이미지 전처리 최적화 PR은 numba 기반의 퓨전 기법과 룩업 테이블을 성공적으로 적용하여 최대 10배에 달하는 성능 향상을 달성했습니다. 이는 멀티모달 LLM의 추론 지연 시간을 크게 줄이고, vLLM의 전반적인 처리량을 향상시키는 데 기여할 것입니다. 이 사례는 복잡한 데이터 전처리 파이프라인에서 성능 병목을 해결하기 위한 효과적인 전략들을 잘 보여줍니다.
References
- numba.njit
- numba.prange
- PIL.Image.Image.resize
- transformers.image_processing_utils.BaseImageProcessor
참고 자료
- https://numba.pydata.org/numba-doc/latest/user/jit.html
- https://numba.pydata.org/numba-doc/latest/user/parallel.html#prange
- https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.resize
- https://huggingface.co/docs/transformers/main/en/main_classes/image_processor#transformers.image_processing_utils.BaseImageProcessor
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
PR Analysis 의 다른글
- 이전글 [sglang] sglang, 멀티모달 모델 인코더 병렬 처리 최적화: 전체 복제본 활용으로 성능 향상
- 현재글 : [vllm] [vLLM 성능 최적화] Kimi-K2.5/K2.6 이미지 전처리 10배 가속화: Numba와 퓨전 기법 활용
- 다음글 [triton] Triton: Blackwell 아키텍처를 위한 TMEM Load-Reduce 연산 퓨전 최적화
댓글