SGLang에서 멀티모달 인코더용 CUDA Graph

왜 필요한가 (Motivation)

멀티모달 추론 서비스에서 시각 인코더(ViT / Vision Transformer)는 대체로 이런 특징을 가져요.

레이어가 많고 연산자가 쪼개져 있어요: 각 레이어가 LN, QKV projection, attention, MLP, residual connection 등을 포함해 커널 런칭이 극도로 잦아요.

서버 쪽 "소규모 배치 / 낮은 지연"이 흔해요: 배치 크기가 아주 작아서("플래트닝"하면 1개처럼 보일 때도 있죠), 커널 런칭 오버헤드가 엔드투엔드 지연의 큰 부분을 차지해요.

입력 토큰 수(패치 수)가 자주 변해요: 이미지·비디오 해상도와 배치 구성이 달라지면 시퀀스 길이 S가 달라지는데, 이게 바로 CUDA Graph의 가장 큰 장애물(불안정한 shape)이에요.

CUDA Graph의 가치는 이래요. 고정된 shape과 고정된 메모리 주소를 가진 긴 GPU 커널 시퀀스를 하나의 그래프로 캡처해 두고, 나중에 같은 shape이 오면 그래프를 그대로 재생(replay)할 수 있어요. 그 결과 런칭 오버헤드가 크게 줄고 GPU 스케줄링이 더 촘촘해져요.

그래서 SGLang은 ViT 성능을 높이기 위해 CUDA Graph 지원 기능을 만들었어요.

출처: 공식문서

설계와 제약

CUDA Graph를 지원하는 새 ViT 로직은 ViTCudaGraphRunner 위에 구축돼요. 이 러너는 비전 트랜스포머의 "blocks + merger + deepstack merger(선택)" 부분을 CUDA 그래프로 캡처하고, 동일한 shape에 대해 재생해요. 세부 설계 고려와 제약은 아래를 확인하세요.

CUDA Graph의 정적 제약에 맞추기 위한 동적 입력

ViT에선 가변 시퀀스 길이 S가 아주 흔해요. 하지만 CUDA Graph는 고정된 shape만 요구하죠. 해법은 S를 기준으로 그래프 캐시를 만드는 거예요(예: graph_key = S). 처음 새 S가 나타나면 그래프를 캡처하고, 이후에는 재생하면 돼요.

서로 다른 S 값이 많다면, 여러 그래프의 개인 메모리 풀 때문에 VRAM 사용량이 늘어나요.

안정적인 주소 (Stable addresses)

"파라미터 같은" 것들은 모두 정적 버퍼가 돼요:

block_input / block_ws / block_output
cu_full_len / cu_window_len 및 그 kk 변형
sin_cos_ws

이렇게 해서 재생 중 텐서를 바꿔치기할 수 없고 내용만 수정할 수 있다는 기본 요건을 해결해요.

어텐션 백엔드 인자

어텐션 백엔드 인자는 그래프 안에서 고정돼요.

TritonAttn은 [cu_seqlens, cu_seqlens_kk, max_len]을 기대하고, FA3는 [cu_seqlens, max_len]을 기대해요.

max_len은 int 상수로 고정돼요. cu_seqlens는 create_graph() 동안 dict에 캐시되고, 이후 재생 중엔 내용이 갱신되지 않아요.

같은 graph_key = S에 대해 입력 shape뿐 아니라 cu_seqlens(와 window seqlens)의 세그먼테이션 패턴도 동일해야 해요. 그렇지 않으면 attention이 시퀀스를 잘못 세그먼트해요.

Rotary 버퍼 관리

seq_len이 늘어나면 더 큰 sin_cos_ws를 재할당해요. max_content_len은 할당된 rotary 버퍼의 최대 크기를 보장하는 데 쓰여요.

명령 예시

ViT에 CUDA Graph를 켜려면 환경변수 SGLANG_VIT_ENABLE_CUDA_GRAPH=1을 설정하면 돼요.

SGLANG_VIT_ENABLE_CUDA_GRAPH=1 \
python3 -m sglang.launch_server \
  --model Qwen/Qwen3-VL-8B-Instruct

ViT 그래프를 디코더의 프리필 위에 Piecewise CUDA Graph와 함께 돌리려면 프리필 백엔드를 명시적으로 요청해야 해요. Qwen3-VL은 멀티모달 허용 목록(allowlist)에 없어서 기본 백엔드는 프리필을 disabled로 해석하기 때문이에요. 백엔드를 명명하면 그 선택이 고정되고 해당 규칙을 건너뛰게 돼요.

SGLANG_VIT_ENABLE_CUDA_GRAPH=1 \
python3 -m sglang.launch_server \
  --model Qwen/Qwen3-VL-8B-Instruct \
  --cuda-graph-backend-prefill tc_piecewise \
  --cuda-graph-max-bs-prefill 4096 \
  --cuda-graph-tc-compiler eager

지원이 확인된 모델

Qwen2.5-VL (PR #14422)
Qwen3-VL (PR #15320)

SGLang에서 멀티모달 인코더용 CUDA Graph