[sglang] SGLang, EmbeddingGemma 사전 추론 성능 최적화: CUDA 그래프와 배치 처리의 힘

PR 링크: sgl-project/sglang#32383 상태: Merged | 변경: +285 / -37

들어가며

최근 SGLang 프로젝트에서는 google/embeddinggemma-300m 모델의 사전 추론(prefill) 성능을 최적화하는 중요한 변경 사항이 있었습니다. 텍스트 임베딩 모델은 대규모 언어 모델(LLM)의 핵심 구성 요소로, 의미론적으로 유사한 데이터를 효율적으로 검색하고 분류하는 데 필수적입니다. 특히 EmbeddingGemma와 같이 양방향(bidirectional) 인코더를 사용하는 모델은 입력 시퀀스 전체를 고려해야 하므로, 기존의 순차적 처리 방식으로는 성능 병목 현상이 발생하기 쉽습니다. 이번 PR은 이러한 문제를 해결하고, SGLang의 강력한 기능인 CUDA 그래프(CUDA Graph)와 배치 처리(batching)를 활용하여 EmbeddingGemma의 사전 추론 성능을 극적으로 향상시키는 데 초점을 맞추고 있습니다.

본 글에서는 이 PR이 어떤 문제를 해결하려 했는지, 구체적으로 어떤 코드 변경이 이루어졌는지, 그리고 이러한 변경이 왜 성능 향상으로 이어지는지에 대해 심층적으로 분석해 보겠습니다.

코드 분석

이번 PR은 주로 EmbeddingGemma 모델의 특성에 맞춰 SGLang의 내부 동작 방식을 조정하고, 성능 최적화를 위한 새로운 전략을 도입하는 데 중점을 두고 있습니다. 변경 사항은 여러 파일에 걸쳐 나타나며, 각 파일의 역할과 변경 내용을 살펴보겠습니다.

1. docs_new/cookbook/autoregressive/Google/EmbeddingGemma.mdx 및 docs_new/docs.json

• 변경 내용: EmbeddingGemma 모델을 위한 새로운 문서 파일이 생성되었고, docs.json에 해당 문서가 추가되었습니다. 이 문서는 EmbeddingGemma 모델을 SGLang으로 서빙하는 방법, 필요한 사전 조건, 권장 서버 실행 명령어, 그리고 임베딩 생성 예시를 포함합니다.
• 핵심: 새로운 모델 지원을 위한 문서화는 사용자 경험과 채택률을 높이는 데 필수적입니다. 특히, EmbeddingGemma가 양방향 인코더이며, SGLang이 이를 위해 CUDA 그래프(BCG - Breakable CUDA Graph)를 기본으로 사용하고, Prefix Caching 및 Chunked Prefill을 비활성화한다는 점을 명확히 합니다.

2. python/sglang/srt/layers/attention/flashattention_backend.py

• 변경 내용: FlashAttention 백엔드의 forward_extend 함수에서 window_size 계산 로직이 수정되었습니다. 기존에는 causal 플래그가 단순하게 결정되었으나, 이제는 layer.attn_type (AttentionType.ENCODER_ONLY, AttentionType.DECODER_BIDIRECTIONAL)을 고려하여 causal 여부를 결정하고, 이에 따라 window_size의 오른쪽(right) 값을 조정합니다.

-        window_size = (layer.sliding_window_size, 0) if is_swa_layer else (-1, -1)
+        causal = not (
+            layer.is_cross_attention
+            or layer.attn_type
+            in (AttentionType.ENCODER_ONLY, AttentionType.DECODER_BIDIRECTIONAL)
+        )
+        # FlashAttention's sliding-window tuple is (left, right). Bidirectional
+        # encoder layers must see the same local context on both sides.
+        window_size = (
+            (layer.sliding_window_size, 0 if causal else layer.sliding_window_size)
+            if is_swa_layer
+            else (-1, -1)
+        )

• 핵심: EmbeddingGemma와 같은 양방향 인코더 모델은 어텐션 계산 시 좌우 양쪽의 컨텍스트를 모두 봐야 합니다. 기존의 causal=True 설정은 이를 제대로 반영하지 못할 수 있습니다. 수정된 코드는 AttentionType.DECODER_BIDIRECTIONAL과 같은 경우를 명시적으로 처리하여, 양방향 어텐션이 올바르게 동작하도록 window_size를 설정합니다. 이는 모델의 정확성을 보장하는 중요한 수정입니다.

3. python/sglang/srt/layers/layernorm.py

• 변경 내용: LayerNorm 클래스의 forward_native, forward_cpu, forward_cuda 메서드에 residual 인자가 추가되었습니다. 또한, CUDA 구현에서는 gemma_fused_add_rmsnorm 함수를 사용하여 잔차 연결(residual connection)과 RMSNorm 연산을 융합(fuse)하도록 변경되었습니다.

-    def forward_native(self, x):
-        return output.type_as(x)
+    def forward_native(self, x, residual: Optional[torch.Tensor] = None):
+        if residual is not None:
+            residual = x + residual
+            x = residual
+        output = self._norm(x.float())
+        output = output * (1.0 + self.weight.float())
+        output = output.type_as(x)
+        return output if residual is None else (output, residual)
...
-    def forward_cuda(self, x):
+    def forward_cuda(self, x, residual: Optional[torch.Tensor] = None):
+        if residual is not None:
+            # The decoder residual is token-major and contiguous. The fused
+            # kernel updates both tensors in place: x becomes the normalized
+            # output and residual becomes x + residual for the next layer.
+            gemma_fused_add_rmsnorm(x, residual, self.weight.data, self.eps)
+            return x, residual
+        if x.dim() == 2:
+            return gemma_rmsnorm(x, self.weight.data, self.eps)
+        return self.forward_native(x)

• 핵심: Gemma 모델 아키텍처는 잔차 연결과 RMSNorm 연산이 밀접하게 결합되어 있습니다. 이 변경은 SGLang이 Gemma의 이러한 아키텍처 특성을 정확하게 따르도록 하여 모델의 정확성을 높입니다. 특히 CUDA에서의 융합 연산(gemma_fused_add_rmsnorm)은 연산 횟수를 줄이고 메모리 접근을 최적화하여 성능 향상에 기여합니다. 또한, forward_native 메서드에서 residual을 처리하는 방식은 x + residual 연산을 먼저 수행한 후 정규화하는 Gemma의 방식을 따릅니다.

4. python/sglang/srt/managers/tokenizer_manager.py

• 변경 내용: _tokenize_texts 함수 내에서 EmbeddingGemma 모델의 경우, 토크나이저 결과에 EOS(End-Of-Sentence) 토큰을 명시적으로 추가하는 로직이 추가되었습니다.

+        # vLLM's OpenAI embeddings endpoint includes special tokens for
+        # encoder models. EmbeddingGemma's restored Gemma tokenizer adds BOS
+        # but, by its checkpoint default, omits EOS. Add EOS explicitly here
+        # rather than mutating tokenizer-global post-processing state.
+        if (
+            self.model_config.is_embedding_gemma
+            and self.tokenizer.eos_token_id is not None
+        ):
+            input_ids = [
+                (
+                    ids
+                    if ids and ids[-1] == self.tokenizer.eos_token_id
+                    else [*ids, self.tokenizer.eos_token_id]
+                )
+                for ids in input_ids
+            ]

• 핵심: EmbeddingGemma 모델은 Hugging Face 체크포인트 기본 설정에서 EOS 토큰을 포함하지 않을 수 있습니다. 그러나 OpenAI 호환 API의 임베딩 엔드포인트는 종종 EOS 토큰을 기대하거나, 모델이 이를 통해 시퀀스의 끝을 명확히 인식하도록 설계될 수 있습니다. 이 변경은 EmbeddingGemma 모델에 대해 EOS 토큰을 일관되게 추가하여, 모델의 출력 정확성을 보장하고 다른 시스템과의 호환성을 높입니다.

5. python/sglang/srt/models/gemma3_causal.py

• 변경 내용: Gemma3CausalForSglang 클래스의 forward 메서드에서 잔차 연결(residual connection) 처리가 수정되었습니다. input_layernorm과 pre_feedforward_layernorm 통과 후 residual을 업데이트하고, 최종 출력을 (hidden_states, residual) 튜플로 반환하도록 변경되었습니다.

-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
+        # Keep the residual live across layers so the add preceding the next
+        # RMSNorm is fused by Gemma3RMSNorm. This matches the upstream Gemma3
+        # residual layout and is safe to capture in a breakable CUDA graph.
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
...
-        hidden_states = residual + hidden_states
-
-        residual = hidden_states
-        hidden_states = self.pre_feedforward_layernorm(hidden_states)
+        hidden_states, residual = self.pre_feedforward_layernorm(
+            hidden_states, residual
+        )
+        # ... (MLP and post_feedforward_layernorm)
+        # hidden_states = residual + hidden_states # This line is removed
+
+        outputs = (hidden_states, residual)

• 핵심: 이 변경은 layernorm.py에서의 수정과 연계됩니다. Gemma3CausalForSglang 모델은 각 레이어의 출력을 다음 레이어의 입력 잔차로 전달해야 합니다. 이전에는 hidden_states가 업데이트된 후 residual에 다시 할당되었지만, 이제는 input_layernorm과 pre_feedforward_layernorm에서 residual을 직접 업데이트하고, 최종적으로 (hidden_states, residual) 쌍을 반환합니다. 이는 CUDA 그래프 캡처 시 잔차 연결의 상태를 정확하게 유지하는 데 중요하며, gemma_fused_add_rmsnorm과 같은 융합 연산이 올바르게 작동하도록 보장합니다.

왜 이게 좋은가: 성능 향상과 일반화된 교훈

이번 PR의 핵심은 EmbeddingGemma 모델의 특성(양방향 인코더)을 이해하고, SGLang의 고성능 기능을 이에 맞게 적용하여 사전 추론 성능을 극대화하는 것입니다. 특히 다음과 같은 점들이 성능 향상에 기여합니다.

1. CUDA 그래프 (BCG - Breakable CUDA Graph) 활용:

   • PR 설명: "batch embedding requests before scheduling so BCG captures the complete prefill" (스케줄링 전에 임베딩 요청을 배치하여 BCG가 전체 사전 추론을 캡처하도록 함)
   • 이유: CUDA 그래프는 GPU 커널 실행을 캡처하여 반복 실행 시 오버헤드를 크게 줄입니다. EmbeddingGemma와 같이 입력 시퀀스 전체를 처리하는 경우, 각 요청마다 그래프를 새로 생성하는 것은 비효율적입니다. 이 PR에서는 여러 임베딩 요청을 하나의 큰 배치로 묶어, 해당 배치의 전체 사전 추론 과정을 단일 CUDA 그래프로 캡처합니다. 이렇게 하면 그래프 캡처 및 실행 오버헤드가 줄어들고, GPU가 더 효율적으로 활용됩니다.
   • 성능 수치: PR 설명에 따르면, H200 GPU에서 SGLang BCG 모드가 vLLM의 기본 모드보다 4.37배 더 빠릅니다 (33.01ms vs 144.29ms median p50 latency).

2. 배치 처리 (Batching) 및 FA3/FA4 활용:

   • PR 설명: "batch embedding requests before scheduling so BCG captures the complete prefill" (스케줄링 전에 임베딩 요청을 배치하여 BCG가 전체 사전 추론을 캡처하도록 함)
   • 이유: 여러 개의 짧은 임베딩 요청을 하나의 긴 요청으로 묶어 처리하면, GPU의 병렬 처리 능력을 최대한 활용할 수 있습니다. 특히, H100/H200과 같은 최신 GPU에서는 FlashAttention 3 (FA3) 또는 4 (FA4)와 같은 최적화된 어텐션 구현을 활용하여 처리량을 높입니다. SGLang은 EmbeddingGemma 모델에 대해 자동으로 FA3/FA4를 선택하고, 16,384 토큰까지의 BCG 티어를 캡처하도록 설정하여 이러한 이점을 극대화합니다.

3. 양방향 어텐션 및 잔차 연결 최적화:

   • 변경 내용: flashattention_backend.py와 gemma3_causal.py에서의 수정은 EmbeddingGemma의 양방향 어텐션 메커니즘과 Gemma 아키텍처의 잔차 연결 방식을 정확하게 구현합니다. layernorm.py에서의 융합 연산(gemma_fused_add_rmsnorm)은 연산 효율성을 높입니다.
   • 이유: 모델의 아키텍처를 정확하게 이해하고 내부 연산을 최적화하는 것은 정확성뿐만 아니라 성능에도 직결됩니다. 특히, 잔차 연결과 RMSNorm을 융합하는 것은 GPU 커널 수준에서의 최적화를 통해 메모리 접근 및 연산 시간을 단축시킵니다.

4. EOS 토큰 처리:

   • 변경 내용: tokenizer_manager.py에서 EmbeddingGemma에 EOS 토큰을 명시적으로 추가합니다.
   • 이유: 이는 모델의 정확성을 보장하고, API 호환성을 높이는 데 기여합니다. 비록 직접적인 성능 향상보다는 정확성 및 안정성 측면이 강하지만, 잘못된 토큰 처리로 인한 재처리나 오류를 방지함으로써 간접적으로 효율성을 높일 수 있습니다.

일반화된 교훈:

• 모델 특성 이해의 중요성: 특정 모델 아키텍처(예: 양방향 인코더, 잔차 연결 방식)의 고유한 특성을 이해하고 SGLang의 기능을 이에 맞게 조정하는 것이 성능 최적화의 핵심입니다.
• CUDA 그래프의 전략적 활용: 반복적인 연산 패턴(특히 사전 추론)이 있는 경우, CUDA 그래프를 사용하여 오버헤드를 줄이고 GPU 활용률을 극대화할 수 있습니다. 여러 요청을 묶어 하나의 그래프로 캡처하는 전략은 매우 효과적입니다.
• 배치 처리의 힘: 가능한 경우 요청을 배치하여 처리하면 GPU의 병렬 처리 능력을 최대한 활용할 수 있습니다.
• 융합 연산의 이점: 여러 연산을 GPU 커널 수준에서 융합하면 메모리 접근 및 연산 시간을 크게 단축할 수 있습니다.

리뷰 댓글 분석

이번 PR에 대한 리뷰 댓글은 주로 CI 상태 확인에 집중되어 있었습니다. [mickqian]님의 "urgent PR but blocked by CI. will bypassing while keeping an eye on CI"라는 댓글은 이 PR이 긴급하게 처리되어야 했음을 시사합니다. CI가 차단되었음에도 불구하고, 리뷰어는 CI를 우회하면서도 PR을 주시하겠다는 의지를 보였습니다. 이는 코드 변경 자체의 중요성을 인정하고, 신속한 병합을 통해 문제를 해결하려는 의도로 해석됩니다. 코드 변경 내용 자체에 대한 깊은 기술적 논의는 없었지만, 이는 변경 사항이 기존 SGLang 프레임워크 내에서 잘 통합될 것으로 예상되었거나, 또는 변경의 필요성이 명확했기 때문일 수 있습니다.

결론

이번 SGLang PR은 EmbeddingGemma 모델의 사전 추론 성능을 획기적으로 개선했습니다. CUDA 그래프를 활용하여 여러 임베딩 요청을 효율적으로 처리하고, 모델의 양방향 어텐션 및 잔차 연결 특성을 정확하게 반영하며, 최신 하드웨어 기능을 활용함으로써 이전 대비 상당한 성능 향상을 달성했습니다. 이는 LLM 서빙 엔진이 특정 모델 아키텍처에 맞춰 최적화될 때 얼마나 강력한 성능을 발휘할 수 있는지를 보여주는 좋은 사례입니다.

참고 자료

• https://huggingface.co/google/embeddinggemma-300m
• https://pytorch.org/docs/stable/generated/torch.compile.html

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