[sglang] SGLang, MoonViT의 최대 시퀀스 길이 메타데이터 재사용으로 성능 개선

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

PR 링크: sgl-project/sglang#30878 상태: Merged | 변경: +97 / -3

들어가며

최근 SGLang 프로젝트의 GitHub PR에서는 MoonViT 모델의 Vision Transformer (ViT) 부분에서 발생하는 성능 병목 현상을 해결하기 위한 중요한 최적화가 이루어졌습니다. 특히, Kimi 모델의 K2.5/K2.6/K2.7 버전에서 사용되는 MoonViT 아키텍처의 max_seqlen 메타데이터 계산 방식을 개선하여, GPU와 호스트 간의 불필요한 동기화 작업을 줄이는 데 초점을 맞추었습니다. 이 PR은 ViT 인코더의 각 어텐션 블록에서 max_seqlen을 개별적으로 계산하고 GPU-호스트 동기화를 반복적으로 수행하던 비효율적인 패턴을 제거함으로써, 전체적인 추론 속도를 향상시키는 것을 목표로 합니다.

코드 분석

이번 PR의 핵심 변경 사항은 python/sglang/srt/layers/attention/vision.py와 python/sglang/srt/models/kimi_k25.py 파일에서 이루어졌으며, 새로운 테스트 파일 test/registered/unit/layers/attention/test_vision_max_seqlen.py가 추가되었습니다.

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

이 파일은 Vision Transformer의 어텐션 레이어를 구현하는 핵심 로직을 담고 있습니다. 이전 코드에서는 cu_seqlens (cumulative sequence lengths)로부터 각 시퀀스의 길이를 계산하고, 그 중 최대값을 max_seqlen으로 결정했습니다. 이 과정에서 seq_lens.max().item() 호출은 GPU에서 계산된 결과를 CPU로 가져오는 동기화 작업(GPU-to-host synchronization)을 유발했습니다.

Before:

-            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
-            max_seqlen = seq_lens.max().item()
+            # Some vision encoders precompute this scalar once per encoder
+            # forward and share it across all of their attention blocks.  Use
+            # that value when available: deriving it here requires a
+            # GPU-to-host sync, so repeating it per block serializes the ViT
+            # launch stream for variable-size images.
+            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)

After:

수정된 코드에서는 kwargs.get("max_seqlen")을 통해 외부에서 max_seqlen 값이 전달되는지를 먼저 확인합니다. 만약 max_seqlen이 전달되지 않았다면, 이전과 같이 cu_seqlens로부터 계산하지만, 이 계산은 이제 max_seqlen이 이미 제공된 경우를 위한 폴백(fallback) 로직으로 간주됩니다. 중요한 점은, 외부에서 max_seqlen 값이 제공될 경우, 이 값을 그대로 사용하고 GPU-호스트 동기화 과정을 생략한다는 것입니다. 이는 특히 여러 개의 어텐션 블록을 순차적으로 호출하는 ViT 인코더에서 각 블록마다 max_seqlen을 다시 계산하고 동기화하는 오버헤드를 제거합니다.

2. `python/sglang/srt/models/kimi_k25.py`

이 파일은 Kimi 모델의 특정 구현(K2.5/K2.6/K2.7)을 담당하며, MoonViT 인코더 레이어를 포함합니다. 이전 코드에서는 MoonViTEncoderLayer의 forward 메소드 내에서 lengths.max()를 호출하여 max_seqlen을 계산했습니다. 이 역시 GPU-호스트 동기화를 유발하는 지점입니다.

Before:

-        max_seqlen = lengths.max()
+        # FlashAttention needs a host integer.  Compute it once per MoonViT
+        # forward and pass it to every encoder block instead of synchronizing
+        # once per block inside the attention backend.
+        max_seqlen = int(lengths.max().item())

After:

수정된 코드에서는 lengths.max().item()을 호출하여 max_seqlen을 계산한 후, 즉시 int()로 캐스팅하여 호스트 정수형으로 만듭니다. 그리고 이 값을 max_seqlen 변수에 저장하고, 이후 VisionFlash3Attention 호출 시 max_seqlen 인자로 전달합니다. 이 변경을 통해 MoonViTEncoderLayer의 forward 메소드 호출 시점에 max_seqlen이 한 번만 계산되고, 그 값이 모든 내부 어텐션 블록에 재사용됩니다. 이는 vision.py에서의 변경과 일관되며, 불필요한 GPU-호스트 동기화를 제거하는 효과를 가져옵니다.

3. `test/registered/unit/layers/attention/test_vision_max_seqlen.py`

새로 추가된 이 테스트 파일은 이번 PR의 변경 사항을 검증하기 위한 단위 테스트를 포함합니다. test_vision_flash3_uses_precomputed_max_seqlen 함수는 VisionFlash3Attention이 외부에서 제공된 max_seqlen 값을 올바르게 사용하는지, 그리고 test_kimi_moonvit_forwards_one_precomputed_max_seqlen 함수는 MoonViTEncoderLayer가 계산된 max_seqlen 스칼라 값을 내부 어텐션 블록에 올바르게 전달하는지를 테스트합니다. 또한, CPU 환경에서의 회귀 테스트도 포함하여 다양한 환경에서의 안정성을 보장합니다.

왜 이게 좋은가?

이번 PR의 가장 큰 장점은 GPU-호스트 동기화 횟수를 줄여 성능을 개선했다는 점입니다. 특히, ViT와 같이 여러 어텐션 블록을 순차적으로 사용하는 모델에서 각 블록마다 max_seqlen을 계산하고 그 결과를 CPU로 가져오는 작업은 상당한 오버헤드를 발생시킬 수 있습니다. GPU 연산은 병렬 처리에 최적화되어 있지만, GPU와 CPU 간의 데이터 전송 및 동기화는 직렬적인 작업이 될 수밖에 없기 때문입니다.

PR 설명에 따르면, H100 GPU에서 ViT 마이크로벤치마크를 수행한 결과 다음과 같은 성능 향상이 관찰되었습니다:

ViT median: 32.13 ms → 31.08 ms (-3.27%)
ViT mean: 32.14 ms → 31.07 ms (-3.35%)

이 수치는 특정 모델과 하드웨어 환경에서의 결과이지만, GPU-호스트 동기화 오버헤드가 성능에 미치는 영향을 명확히 보여줍니다. 비록 이 최적화가 엔드-투-엔드 서빙 성능 전체를 극적으로 향상시키는 것은 아니지만, 모델의 핵심 연산 부분에서 불필요한 지연을 제거함으로써 전체적인 효율성을 높이는 중요한 단계입니다.

일반적인 교훈:

GPU-호스트 동기화 최소화: GPU 커널 실행 전후의 동기화는 성능 저하의 주요 원인이 될 수 있습니다. 가능한 한 GPU 내에서 연산을 완료하거나, 필요한 데이터만 효율적으로 전송해야 합니다.
메타데이터 재사용: 모델의 여러 부분에서 동일한 계산 결과를 필요로 할 때, 해당 계산을 한 번만 수행하고 그 결과를 재사용하는 것이 효율적입니다. 특히 max_seqlen과 같이 전체 시퀀스 길이에 의존하는 값은 모델의 상위 레벨에서 한 번 계산하여 하위 레벨의 여러 컴포넌트에 전달하는 것이 좋습니다.
테스트 커버리지 확보: 새로운 최적화 로직이 도입될 때는, 해당 로직이 의도한 대로 동작하는지, 그리고 다양한 환경(CPU, GPU)에서 안정적으로 작동하는지를 검증하는 단위 테스트를 반드시 추가해야 합니다.

리뷰 댓글 분석

제공된 리뷰 댓글은 주로 CI 상태 확인 및 재실행 요청에 관한 것이었습니다. 특별한 기술적 논의나 피드백은 포함되어 있지 않았습니다. 이는 해당 PR의 변경 사항이 명확하고, 코드 리뷰어들이 큰 이견 없이 받아들였음을 시사할 수 있습니다. 다만, /tag-and-rerun-ci와 같은 댓글은 CI 파이프라인의 정상적인 작동을 확인하고, 변경 사항이 통합되기 전에 모든 테스트가 통과했는지 확인하는 과정의 일부로 볼 수 있습니다.