본문으로 건너뛰기

[sglang] [Kimi K3] CPU 전송 이미지의 지연 전처리(Deferred Preprocessing)를 통한 VLM 성능 최적화

PR 링크: sgl-project/sglang#33921 상태: Merged | 변경: +470 / -82

들어가며

VLM(Vision Language Model) 서빙에서 이미지 데이터는 텍스트 토큰에 비해 압도적으로 큰 크기를 차지합니다. 특히 분산 환경(Data Parallelism, Tensor Parallelism)에서 이미지를 처리할 때, 어떤 GPU가 해당 이미지를 처리할지 결정하기 전에 미리 전처리를 수행하고 이를 모든 노드에 브로드캐스트하는 방식은 심각한 네트워크 및 메모리 병목을 초래합니다.

이번 PR은 Kimi-K3 모델에서 CPU 멀티모달 전송(CPU transport) 시 발생하는 비효율을 해결하기 위해, 이미지 전처리를 실제 연산을 담당할 GPU(Vision Owner)가 결정될 때까지 미루는 지연 전처리(Deferred Preprocessing) 기법을 도입했습니다. 이를 통해 불필요한 FP32 텐서 전송을 제거하고 전체적인 서빙 성능을 비약적으로 향상시켰습니다.

문제 상황: 전처리 후 전송의 비효율성

기존 Kimi-K3 구현에서는 다음과 같은 비효율적인 워크플로우가 존재했습니다:

  1. Tokenizer Worker: GPU에서 이미지를 전처리(Resize, Normalization, Patchification)하여 FP32 텐서로 변환합니다.
  2. CPU Copy: 변환된 거대한 FP32 텐서를 CPU로 복사합니다.
  3. Broadcast: 모든 스케줄러 GPU로 이 데이터를 전송합니다.
  4. Vision Owner 선택: DP(Data Parallel) 로드 밸런싱을 통해 실제 연산을 수행할 GPU를 정한 뒤, 다시 해당 GPU로 데이터를 복사합니다.

이미지는 uint8 상태일 때보다 전처리 후 fp32 상태일 때 크기가 약 4배 이상 커지며, 패치화(Patchification) 과정에서 차원이 확장되면 그 차이는 더 벌어집니다. 이를 모든 GPU에 뿌리는 것은 엄청난 낭비입니다.

코드 분석: 지연 전처리의 구현

1. kimi_k3.py: Vision Owner에서의 실행 시점 전처리

가장 핵심적인 변화는 materialize_item_features 함수입니다. 기존에는 이미 전처리된 item.feature를 가져오기만 했다면, 이제는 DEFERRED_PREPROCESSING_KEY가 있는지 확인하여 현 시점에서 GPU 전처리를 수행합니다.

Before:

# 전처리된 피처를 단순히 리스트에 담아 반환
features = []
for image_index in image_indices:
    item = items[image_index]
    feature = item.feature
    features.append(feature)
return materialize_multimodal_features(features, device=device, dtype=target_dtype)

After:

# 지연 전처리 설정이 있는지 확인
deferred = [
    item.model_specific_data.get(DEFERRED_PREPROCESSING_KEY)
    for item in selected_items
]

if any(config is not None for config in deferred):
    # ... 생략 ...
    # 실제 Vision Owner GPU에서 전처리를 직접 수행
    pixel_values, _ = _gpu_preprocess_images(
        [item.feature for item in selected_items], # 원본 uint8 데이터
        [config["resize_config"] for config in deferred],
        image_scale, image_bias,
        self.vision_tower.patch_size,
        to_chw=lambda image: to_chw_uint8(image, device=device),
        post_resize=lambda x: fill_transparent_bg(x, first_config["transparent_bg_config"]),
    )
    return pixel_values.to(dtype=target_dtype)

이 변경을 통해 데이터 전송 단계에서는 가벼운 uint8 상태를 유지하고, 실제 계산이 필요한 GPU에서만 fp32로 확장하며 전처리를 수행하게 됩니다.

2. kimi_k3_image_processing.py: 공통 로직의 분리 및 최적화

전처리를 지연시키기 위해 uint8 데이터를 안전하게 다루고, GPU에서 배경 채우기(Background filling) 등을 수행할 유틸리티가 추가되었습니다.

def to_chw_uint8(image: Union[torch.Tensor, Image.Image], device: torch.device | str | None = None) -> torch.Tensor:
    if isinstance(image, Image.Image):
        # PIL 이미지를 uint8 텐서로 변환 (메모리 절약)
        array = np.array(image.convert("RGBA" if has_alpha else "RGB"), copy=True)
        image = torch.from_numpy(array).permute(2, 0, 1)
    
    if image.dtype != torch.uint8:
        raise ValueError(f"Kimi-K3 preprocessing expects raw uint8 pixels, got {image.dtype}")
    # ...
    if device is not None:
        image = image.to(device)
    return image

이 코드는 이미지를 fp32로 미리 바꾸지 않고 uint8 상태로 GPU에 올린 뒤, _gpu_preprocess_images 내부에서 연산 직전에 정규화를 수행하도록 돕습니다.

3. processors/kimi_k3.py: 조건부 지연 결정

모든 경우에 지연 전처리가 유리한 것은 아닙니다. 원본 이미지(uint8)가 너무 커서 전처리된 결과물보다 데이터량이 많을 경우(예: 매우 큰 무손실 이미지)에는 오히려 미리 전처리하는 것이 나을 수 있습니다. PR에서는 이를 판단하는 로직이 포함되었습니다.

# (PR 설명 요약) raw payload가 processed FP32 텐서보다 클 경우 eager preprocessing으로 폴백
if raw_payload_size > processed_tensor_size:
    return self._eager_preprocess(input_data)

왜 이게 좋은가?

1. 전송 오버헤드의 극적인 감소

마이크로 벤치마크 결과에 따르면, 4장의 1536x1024 이미지를 처리할 때 CPU 페이로드가 73.0 MiB에서 18.0 MiB로 약 75% 감소했습니다. 이는 uint8과 fp32의 데이터 크기 차이 및 패치화 전후의 데이터 구조 차이에서 기인합니다.

2. TTFT(Time To First Token) 개선

서버 벤치마크에서 부하가 높은 상황(unlimited request rate)일 때, Median TTFT가 4742.8ms에서 1008.6ms로 무려 78.7%나 단축되었습니다. 전처리 연산 자체를 분산된 GPU 중 하나에서만 수행하고, 무거운 데이터 이동을 최소화했기 때문에 동시 요청이 많을수록 그 효과가 극대화됩니다.

3. 유연한 폴백 전략

모든 이미지를 강제로 지연 처리하는 것이 아니라, uint8 원본이 더 큰 특이 케이스나 이미 float 텐서로 들어온 입력에 대해서는 기존의 eager 경로를 유지함으로써 안정성을 확보했습니다.

결론

이번 최적화의 핵심 교훈은 "데이터는 가능한 한 작게 유지하고, 확장은 소비 직전에 수행하라"는 것입니다. 분산 시스템에서 노드 간 통신 비용은 연산 비용보다 비싼 경우가 많습니다. 특히 VLM처럼 입력 데이터의 부피가 큰 모델일수록, 전처리 시점을 런타임 스케줄링 이후로 미루는 전략이 시스템 전체의 처리량(Throughput)과 응답 속도(Latency)를 동시에 잡는 열쇠가 됩니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글