본문으로 건너뛰기

[sglang] [SGLang] VLM 추론 성능의 비약적 향상: Cross-request ViT Batching과 Metadata 재사용 기법

PR 링크: sgl-project/sglang#24013 상태: Merged | 변경: +686 / -213

들어가며

VLM(Vision-Language Model) 서비스에서 가장 비용이 많이 드는 작업 중 하나는 이미지를 텍스트 모델이 이해할 수 있는 임베딩으로 변환하는 ViT(Vision Transformer) Encoding 과정입니다. 기존 SGLang의 chunked prefill 로직에서는 여러 개의 VLM 요청이 들어오더라도, 각 요청에 포함된 이미지를 개별적으로(request-by-request) 처리했습니다. 이는 특히 다수의 요청이 동시에 들어오는 서빙 환경에서 ViT 인코더를 불필요하게 여러 번 호출하게 만들고, GPU 커널 런칭 오버헤드를 가중시키는 원인이 되었습니다.

이번 PR([VLM] Batch cross-request ViT encoding and reuse attention metadata)은 이러한 병목을 해결하기 위해 Cross-request ViT Batching과 Attention Metadata 재사용이라는 두 가지 핵심 최적화를 도입했습니다. 이를 통해 대형 VLM 모델에서 TTFT(Time to First Token)를 획기적으로 줄이고 처리량(Throughput)을 높였습니다.

코드 분석: 무엇이 어떻게 바뀌었나?

1. Vision Attention Metadata 도입 및 재사용

기존 코드에서는 ViT의 각 레이어를 통과할 때마다 cu_seqlens를 계산하거나 GPU에서 CPU로 데이터를 가져오는 동기화(Synchronization) 작업이 발생할 수 있었습니다. 특히 .item()이나 .max()와 같은 호출은 CPU가 GPU의 작업 완료를 기다리게 만들어 심각한 성능 저하를 유발합니다.

Before: 레이어마다 cu_seqlens를 다시 계산하거나 인자로 전달받아 처리했습니다.

# 기존 방식: forward 내에서 매번 계산 또는 resolve
if cu_seqlens is None:
    cu_seqlens = _get_cu_seqlens_for_shape(bsz, seq_len, device=q.device)
# ... 중략 ...
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
max_seqlen = resolve_precomputed_max_seqlen(cu_seqlens, kwargs.get("max_seqlen"))

After: VisionAttentionMetadata 데이터 클래스를 도입하여, 인코딩 시작 전 단 한 번만 메타데이터를 계산하고 이를 모든 레이어에서 공유합니다.

@dataclasses.dataclass
class VisionAttentionMetadata:
    cu_seqlens: torch.Tensor
    seq_lens: torch.Tensor
    max_seqlen: int
    # ... 기타 필드

def prepare_vision_attention_metadata(...) -> VisionAttentionMetadata:
    # 모든 레이어 루프 이전에 단 한 번만 수행
    cu_seqlens = cu_seqlens.to(device=device, dtype=torch.int32, non_blocking=True)
    seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
    max_seqlen = int(seq_lens.max().item()) # 여기서 발생하는 sync를 최소화
    return VisionAttentionMetadata(...)

이 변경을 통해 mickqian 리뷰어가 지적한 대로 Host-Sync 관련 이벤트가 85개에서 4개로 95% 감소하는 놀라운 결과를 얻었습니다.

2. Cross-request ViT Batching

이 PR의 가장 큰 변화는 _get_chunked_prefill_embedding()의 구조적 개선입니다. 이전에는 루프를 돌며 각 요청의 이미지를 하나씩 인코딩했다면, 이제는 현재 배치에 포함된 모든 요청의 'Cache-miss' 이미지들을 모아서 한 번의 ViT 호출로 처리합니다.

Before (Conceptual):

for request in requests:
    if image_not_in_cache:
        embedding = vit_encode(image) # 요청마다 개별 호출

After (Conceptual):

# 1. 모든 요청에서 누락된 이미지를 수집 및 해시 기반 중복 제거
missed_images = collect_and_dedup_misses(requests)

# 2. 단 한 번의 배칭된 ViT 호출
if missed_images:
    batched_embeddings = vit_encode(missed_images)

# 3. 결과를 각 요청의 청크에 다시 할당 (Assemble)
for request in requests:
    request.embeddings = assemble_from_batched(batched_embeddings)

이 방식은 특히 작은 이미지(360p 등)가 여러 개 들어올 때 GPU의 병렬 처리 능력을 극대화합니다.

3. CUDA Graph Warmup 및 Dynamo 컴파일 최적화

CUDA Graph는 커널 런칭 오버헤드를 줄여주지만, 그래프 캡처 도중 torch.compile에 의한 Dynamo 컴파일이 발생하면 에러가 발생하거나 성능이 저하됩니다. 이를 방지하기 위해 캡처 전 'Warmup' 단계를 추가했습니다.

# sglang/srt/layers/attention/vision.py
def flash_attn_func(*args, ver: int = 3, **kwargs):
    if ver == 4:
        from sglang.jit_kernel.flash_attention_v4 import flash_attn_varlen_func as flash_attn_varlen_func_fa4
        return flash_attn_varlen_func_fa4(*args, **kwargs)
    return flash_attn_varlen_func(*args, **kwargs)

또한, FlashAttention-4(FA4) 지원을 위한 래퍼 함수를 추가하여 최신 커널 활용 가능성을 열어두었습니다.

왜 이게 좋은 최적화인가?

1. 성능 수치로 증명된 효과

Qwen3-VL-8B 모델 기준, H200 GPU에서 다음과 같은 성능 향상이 관찰되었습니다.

  • 360p (Small images): CUDA Graph와 결합 시 Throughput이 약 25.91 req/s로 향상 (Launch-bound 영역 해결).
  • 1080p (Large images): Cross-request batching 덕분에 TTFT가 823ms에서 673ms로 약 18% 개선 (Compute-bound 영역 효율화).
  • Qwen3.6 MoE VLM: mickqian의 벤치마크에 따르면 Median TTFT가 17.2% 감소하고 Throughput이 19.8% 증가했습니다.

2. CPU-GPU 동기화 제거의 중요성

딥러닝 추론 엔진에서 가장 흔한 실수 중 하나는 루프 내부에서 .item()을 호출하는 것입니다. 이번 PR은 VisionAttentionMetadata를 통해 이러한 동기화 지점을 레이어 밖으로 밀어냄으로써, GPU가 멈추지 않고 계속해서 커널을 실행할 수 있는 환경을 만들었습니다.

3. 유연한 설계와 안정성

AgainstEntropy 리뷰어의 피드백을 수용하여, 기존의 ViTCudaGraphRunner 경로를 유지하면서도 새로운 배칭 로직을 통합했습니다. 이는 작은 이미지에서는 CUDA Graph의 이점을, 큰 이미지나 복잡한 배치 상황에서는 Batching의 이점을 모두 누릴 수 있게 합니다.

마치며: 시니어 엔지니어의 관점

이번 PR은 단순히 "코드를 빠르게" 만드는 것을 넘어, 데이터 흐름의 구조(Data Flow Architecture)를 개선한 훌륭한 사례입니다.

  1. Batching at the right level: 스케줄링 레벨에서 이미지를 모으는 결정은 시스템 전체의 효율을 결정합니다.
  2. Metadata reuse: 중복 계산을 피하는 것은 기본이지만, '동기화 비용'을 고려한 설계는 시니어급 엔지니어링의 정수입니다.
  3. Deduplication: 동일한 이미지가 여러 요청에 포함될 경우(예: 동일한 비디오 프레임에 대한 다른 질문)를 해시로 처리한 점도 실무적인 통찰이 돋보입니다.

VLM 서빙 성능을 고민하고 있다면, 이 PR에서 보여준 '동기화 최소화'와 'Cross-request 리소스 집약' 전략을 반드시 참고하시기 바랍니다.

참고 자료

  • torch.compile — 이 PR에서 최적화한 Dynamo 컴파일 및 CUDA Graph 관련 공식 문서
  • FlashAttention — PR에서 사용된 고성능 어텐션 커널 라이브러리
  • SGLang GitHub PR #22132 — 이 PR의 기반이 된 이전 최적화 논의

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

댓글

관련 포스트

PR Analysis 의 다른글