[sglang] SGLang Kimi K2.5/K2.7 멀티모달 인코더-DP 성능 최적화: 이미지 피처 전송 샤딩 및 배치 처리

2026년 7월 16일수정: 2026년 7월 16일

PR 링크: sgl-project/sglang#31227 상태: Merged | 변경: +972 / -99

들어가며

대규모 멀티모달 모델(LMM)은 텍스트와 이미지 등 다양한 양식의 데이터를 처리하며 복잡한 추론을 수행합니다. 특히 Kimi K2.5/K2.7과 같은 모델은 고해상도 이미지를 처리하기 위해 Vision Transformer (ViT) 인코더를 사용하며, 이 과정에서 이미지 피처의 효율적인 전송과 전처리가 전체 시스템 성능에 critical한 영향을 미칩니다. 본 PR은 SGLang 프레임워크 내 Kimi K2.5/K2.7 모델의 인코더-DP(Data Parallelism) 모드에서 발생하는 비효율적인 이미지 피처 전송 및 GPU 전처리 문제를 해결하여, 상당한 성능 향상을 달성했습니다.

해결하려는 문제:

중복된 이미지 피처 복사: 기존에는 --mm-enable-dp-encoder 옵션이 활성화된 상태에서도, 일반적인 멀티모달 경로가 모든 TP(Tensor Parallelism) 랭크에서 모든 원본 이미지 피처를 재구성하거나 이동시켰습니다. 예를 들어, TP8 환경에서 4개의 이미지를 처리할 때, 비전 인코더는 총 4개의 이미지만 소비함에도 불구하고 32번의 원본 피처 복사가 발생했습니다. 이는 불필요한 GPU 간 데이터 이동과 메모리 대역폭 낭비로 이어졌습니다.
비효율적인 GPU 이미지 전처리: GPU 전처리 단계에서 NaViT 타겟 사이즈로만 그룹화되어, 동일한 레이아웃을 가진 이미지라도 개별적으로 bicubic resize 커널이 실행되었습니다. 이는 배치 처리의 이점을 충분히 활용하지 못했습니다.
max_seqlen 계산 오버헤드: 어텐션 블록마다 cu_seqlens로부터 max_seqlen을 동적으로 계산하는 과정에서 GPU-to-Host 동기화가 발생하여, 가변 크기 이미지 처리 시 ViT 실행 스트림을 직렬화하는 병목이 있었습니다.

이 PR은 이러한 문제들을 해결하여 Kimi 모델의 멀티모달 처리 파이프라인을 최적화하는 것을 목표로 합니다.

코드 분석

1. `python/sglang/srt/layers/attention/vision.py`

이 파일에서는 어텐션 레이어에서 max_seqlen을 결정하는 방식이 개선되었습니다. 기존에는 각 어텐션 블록마다 cu_seqlens에서 max_seqlen을 계산하여 GPU-to-Host 동기화를 유발했습니다. 새로운 resolve_precomputed_max_seqlen 함수를 도입하여, 인코더가 미리 계산한 max_seqlen 값을 활용할 수 있도록 했습니다.

Before:

            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
            max_seqlen = seq_lens.max().item()

            # ... (context_attention_fwd call)

            max_seqlen = kwargs.get("max_seqlen")
            if max_seqlen is None:
                seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
                max_seqlen = int(seq_lens.max().item())
            elif isinstance(max_seqlen, torch.Tensor):
                max_seqlen = int(max_seqlen.item())
            else:
                max_seqlen = int(max_seqlen)

            # ... (flash_attn_varlen_func call)

            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
            max_seqlen = seq_lens.max().item()

After:

def resolve_precomputed_max_seqlen(
    cu_seqlens: torch.Tensor, max_seqlen: int | torch.Tensor | None
) -> int:
    """Use an encoder-provided max sequence length when one is available.

    Packed vision encoders execute many attention blocks for one image batch.
    Deriving the max from GPU ``cu_seqlens`` in every block synchronizes the
    launch stream, whereas the encoder can materialize this host scalar once.
    """
    if max_seqlen is None:
        seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
        return int(seq_lens.max().item())
    if isinstance(max_seqlen, torch.Tensor):
        return int(max_seqlen.item())
    return int(max_seqlen)

            # ... (in forward method)

            max_seqlen = resolve_precomputed_max_seqlen(
                cu_seqlens, kwargs.get("max_seqlen")
            )
            context_attention_fwd(
                q,
                k,
                v,
                output,
                cu_seqlens.to(q.device),
                seq_lens,
                max_seqlen,
                is_causal=False,
                sm_scale=softmax_scale,
                # ...
            )

            # ... (similar change for flash_attn_varlen_func call)

무엇이 왜 좋은가: resolve_precomputed_max_seqlen 함수는 max_seqlen이 이미 계산되어 kwargs로 전달된 경우 이를 재활용합니다. 이는 seq_lens.max().item() 호출을 반복적으로 수행하여 발생하는 GPU-to-Host 동기화를 방지합니다. 비전 인코더는 여러 어텐션 블록에 걸쳐 동일한 이미지 배치에 대해 max_seqlen을 한 번만 계산하여 호스트 스칼라로 저장할 수 있으므로, 이 최적화는 가변 크기 이미지 처리 시 ViT 실행 스트림의 직렬화를 줄여 성능을 향상시킵니다.

2. `python/sglang/srt/managers/mm_utils.py`

이 파일에서는 멀티모달 피처의 디바이스 이동 및 캐시 처리 방식이 개선되었습니다. 특히 Kimi 모델의 인코더-DP 모드에서 불필요한 데이터 복사를 제거하고, CUDA IPC 캐시 히트 시 리소스 관리를 강화했습니다.

Before:

def _can_skip_pre_embed_feature_move(data_embedding_func: DataEmbeddingFunc) -> bool:
    # ... (Qwen models listed)

def _move_items_to_device(
    items: List[MultimodalDataItem], device: torch.device
) -> None:
    for item in items:
        if item.feature is not None:
            item.feature = item.feature.to(device, non_blocking=True)

    # ... (in _get_chunked_embedding_by_item)
    if miss_items:
        miss_item_list = [item for _, item, _, _ in miss_items]
        _move_items_to_device(miss_item_list, device)

After:

def _can_skip_pre_embed_feature_move(data_embedding_func: DataEmbeddingFunc) -> bool:
    # ... (existing Qwen models)
    "KimiK25ForConditionalGeneration", # Kimi 모델 추가
}

def _acknowledge_deferred_cuda_ipc_cache_hits(
    items: List[MultimodalDataItem],
) -> None:
    # ... (logic to acknowledge consumption of lazy IPC slices)

    # ... (in _get_chunked_embedding_full)
    else:
        _acknowledge_deferred_cuda_ipc_cache_hits(embedding_items_per_req)

    # ... (in _get_chunked_embedding_by_item)
    if miss_items:
        miss_item_list = [item for _, item, _, _ in miss_items]
        if not _can_skip_pre_embed_feature_move(data_embedding_func):
            _move_items_to_device(miss_item_list, device)

무엇이 왜 좋은가:

_can_skip_pre_embed_feature_move 함수에 KimiK25ForConditionalGeneration이 추가되어, Kimi 모델이 내부적으로 피처를 디바이스로 옮기는 경우 _move_items_to_device 호출을 건너뛸 수 있게 되었습니다. 이는 인코더-DP 모드에서 각 TP 랭크가 자신이 소유한 이미지 피처만 구체화(materialize)하도록 하여, 불필요한 GPU 간 복사를 대폭 줄입니다.
_acknowledge_deferred_cuda_ipc_cache_hits 함수는 캐시 히트가 발생하여 ViT 실행이 건너뛰어질 때, 지연된 CUDA IPC 슬라이스를 해제합니다. 이는 고정 풀(fixed-pool) 리소스의 라이프사이클을 올바르게 관리하며, 불필요한 GPU-to-GPU 복사를 다시 도입하지 않으면서도 리소스 누수를 방지합니다.

3. `python/sglang/srt/managers/schedule_batch.py`

이 파일에서는 MultimodalDataItem 클래스가 CUDA IPC 프록시 텐서의 재구성(reconstruction) 및 소비(consumption)를 더 유연하게 처리하도록 개선되었습니다.

Before:

    def reconstruct(self, target_device: int):
        if isinstance(self.feature, CudaIpcTensorTransportProxy):
            self.feature = self.feature.reconstruct_on_target_device(target_device)
        # ... (other proxy reconstructions)

    # ... (in from_processor_output)
    for mm_item in mm_items:
        if mm_item.has_cuda_ipc_proxy():
            if reconstruct_device is None:
                reconstruct_device = torch.cuda.current_device()
            mm_item.reconstruct(reconstruct_device)

After:

    def reconstruct(self, target_device: int, ipc_consumer_count: int = 1):
        if isinstance(self.feature, CudaIpcTensorTransportProxy):
            if ipc_consumer_count == 1:
                self.feature = self.feature.reconstruct_on_target_device(target_device)
            else:
                self.feature = self.feature.reconstruct_on_target_device(
                    target_device, consumer_count=ipc_consumer_count
                )
        # ... (other proxy reconstructions)

    def can_defer_cuda_ipc_feature_reconstruction(self) -> bool:
        # ... (logic to check if reconstruction can be deferred)

    def acknowledge_deferred_cuda_ipc_feature(self, consumer_count: int = 1):
        if isinstance(self.feature, CudaIpcTensorTransportProxy):
            self.feature.acknowledge_consumption(consumer_count)

    # ... (in from_processor_output)
    for mm_item in mm_items:
        if (
            mm_item.has_cuda_ipc_proxy()
            and not mm_item.can_defer_cuda_ipc_feature_reconstruction()
        ):
            if reconstruct_device is None:
                reconstruct_device = torch.cuda.current_device()
            mm_item.reconstruct(reconstruct_device)

무엇이 왜 좋은가:

reconstruct 메서드에 ipc_consumer_count 파라미터가 추가되어, CUDA IPC 프록시 텐서를 재구성할 때 해당 텐서를 소비하는 랭크의 수를 명시할 수 있게 되었습니다. 이는 분산 환경에서 공유 메모리 리소스의 정확한 관리를 가능하게 합니다.
can_defer_cuda_ipc_feature_reconstruction 메서드는 특정 조건(예: Kimi 인코더-DP 모드)에서 이미지 피처의 CUDA IPC 재구성을 지연할 수 있는지 여부를 판단합니다. 이는 각 TP 랭크가 자신이 소유한 이미지 피처만 재구성하도록 하여, 불필요한 데이터 복사를 방지합니다.
acknowledge_deferred_cuda_ipc_feature 메서드는 지연된 IPC 피처가 캐시 히트로 인해 ViT 실행을 건너뛸 때, 해당 피처의 소비를 명시적으로 알립니다. 이는 CUDA IPC 풀의 리소스가 올바르게 해제되도록 보장합니다.
from_processor_output 함수는 이제 can_defer_cuda_ipc_feature_reconstruction의 결과에 따라 reconstruct 호출을 조건부로 수행합니다. 이는 Kimi와 같은 DP-aware 모델에 대해 지연된 재구성 경로를 활성화합니다.

4. `python/sglang/srt/models/kimi_k25.py` (PR 설명 기반)

PR 설명에 따르면, Kimi K2.5/K2.7 모델은 이제 DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY 플래그를 사용하여 이미지 피처 전송을 샤딩하고, 동일한 소스 형상 및 데이터 타입을 가진 이미지에 대해서만 GPU 이미지 리사이즈를 배치 처리합니다. 이는 materialize_multimodal_featur 및 run_dp_sharded_mrope_vision_model과 같은 내부 함수들을 통해 구현됩니다.

무엇이 왜 좋은가: Kimi 모델 자체에서 이러한 최적화 메커니즘을 활용하도록 변경함으로써, 각 TP 랭크가 필요한 이미지 피처만 처리하고, GPU 리사이즈 작업을 효율적으로 배치 처리하여 전반적인 멀티모달 인코더의 성능을 크게 향상시킵니다.

왜 이게 좋은가

이 PR은 SGLang의 Kimi K2.5/K2.7 멀티모달 모델의 성능을 여러 측면에서 획기적으로 개선했습니다. 주요 성능 지표는 다음과 같습니다.

인코더-DP 피처 전송: NVIDIA H100, TP8 환경에서 4개의 이미지에 대해 원본 피처 복사 횟수를 32회에서 4회로 줄여, 5.58배에서 7.52배의 속도 향상을 달성했습니다. 이는 불필요한 데이터 복사로 인한 메모리 대역폭 낭비를 극적으로 줄였습니다.
- [patches, 3, 14, 14] 형태의 FP32 원본 패치 텐서 4개 이미지 입력 기준:
  - 1,024 패치/이미지: 0.405 ms -> 0.073 ms (5.58x)
  - 4,096 패치/이미지: 0.962 ms -> 0.128 ms (7.52x)
공유 피처 구체화(materialization): NVIDIA RTX 5090 환경에서 가변 길이 FP32 피처 텐서의 공유 패킹 작업에서 1.03배에서 1.26배의 속도 향상을 보였습니다. 이는 불필요한 임시 버퍼 생성을 줄여 메모리 효율성을 높였습니다.
GPU 이미지 전처리: NVIDIA H100에서 4개의 GPU 상주 uint8 이미지를 448x448로 리사이즈하는 작업에서 1.57배의 속도 향상을 달성했습니다. 이는 동일한 소스 크기의 이미지에 대해 개별 커널 실행 대신 단일 배치 리사이즈 커널을 사용함으로써 가능했습니다.
End-to-end Kimi K2.7-Code 서빙: NVIDIA GB300, TP4/EP4 환경에서 --mm-enable-dp-encoder를 사용한 워크로드에서 요청 처리율(req/s) 및 입력 토큰 처리율(input tok/s) 10.0% 증가, 평균 TTFT(Time To First Token) 5.9% 감소, P99 TTFT 7.7% 감소를 기록했습니다.

일반적인 교훈:

데이터 이동 최소화: 분산 환경에서 GPU 간 또는 CPU-GPU 간 불필요한 데이터 복사를 최소화하는 것이 성능 최적화의 핵심입니다. 각 워커가 필요한 데이터만 처리하도록 샤딩하는 전략은 매우 효과적입니다.
배치 처리의 극대화: GPU는 병렬 처리에 최적화되어 있으므로, 가능한 한 많은 작업을 배치로 묶어 처리하는 것이 중요합니다. 이미지 리사이즈와 같은 전처리 작업도 예외는 아닙니다.
불필요한 동기화 회피: torch.Tensor.item()과 같은 GPU-to-Host 동기화는 성능 병목을 유발할 수 있습니다. 가능한 경우 호스트에서 미리 계산하거나, 텐서 형태로 전달하여 동기화를 피해야 합니다.
정교한 리소스 관리: CUDA IPC와 같은 공유 메모리 리소스를 사용할 때는 라이프사이클 관리가 매우 중요합니다. 지연된 재구성 및 명시적인 소비 확인 메커니즘은 리소스 누수를 방지하고 안정성을 높입니다.
모델 아키텍처에 특화된 최적화: 모든 모델에 적용되는 일반적인 최적화 외에도, Kimi의 인코더-DP와 같은 특정 모델 아키텍처의 특성을 고려한 맞춤형 최적화가 큰 성능 향상을 가져올 수 있습니다.

결론

이 PR은 SGLang의 Kimi K2.5/K2.7 멀티모달 모델에서 이미지 피처 전송 및 전처리 파이프라인의 주요 병목 지점을 성공적으로 해결했습니다. 샤딩된 피처 전송, 효율적인 GPU 이미지 리사이즈 배치 처리, 그리고 max_seqlen 계산 오버헤드 감소를 통해 상당한 성능 향상을 달성했으며, 이는 대규모 멀티모달 모델의 서빙 효율성을 높이는 데 크게 기여할 것입니다. 이러한 최적화는 분산 환경에서 고성능 컴퓨팅 시스템을 구축하는 데 있어 중요한 모범 사례를 제공합니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [cutlass] NVIDIA CUTLASS SM100 GEMM 커널 최적화 및 분산 환경 개선 분석
현재글 : [sglang] SGLang Kimi K2.5/K2.7 멀티모달 인코더-DP 성능 최적화: 이미지 피처 전송 샤딩 및 배치 처리
다음글 [onnxruntime] ONNX Runtime: fpA_intB GEMM 최적화 및 CUDA 그래프 호환성 강화

[sglang] SGLang Kimi K2.5/K2.7 멀티모달 인코더-DP 성능 최적화: 이미지 피처 전송 샤딩 및 배치 처리

들어가며

코드 분석

1. `python/sglang/srt/layers/attention/vision.py`

2. `python/sglang/srt/managers/mm_utils.py`

3. `python/sglang/srt/managers/schedule_batch.py`

4. `python/sglang/srt/models/kimi_k25.py` (PR 설명 기반)

왜 이게 좋은가

결론

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글