[vllm] vLLM Nemotron Nano VL: Pixel Shuffle 최적화를 통한 성능 향상 분석

2026년 4월 10일수정: 2026년 4월 10일

PR 링크: vllm-project/vllm#37580 상태: Merged | 변경: +None / -None

들어가며

vLLM은 LLM 추론 속도를 혁신적으로 개선하는 라이브러리로, 특히 대규모 언어 모델의 효율적인 서빙에 중점을 두고 있습니다. 최근 vLLM 프로젝트의 Nemotron Nano VL 모델에서 Pixel Shuffle 연산을 최적화하는 PR이 올라왔습니다. 이 PR은 기존 코드의 비효율적인 텐서 연산을 개선하여, 특히 이미지 처리와 관련된 멀티모달 모델의 성능을 향상시키는 것을 목표로 합니다.

본 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 왜 이러한 변경이 성능 향상으로 이어지는지, 그리고 이 최적화가 주는 일반적인 교훈은 무엇인지 살펴보겠습니다.

코드 분석

이번 PR의 핵심은 vllm/model_executor/models/nano_nemotron_vl.py 파일 내 pixel_shuffle 함수의 최적화입니다. 기존 구현은 여러 단계의 view, permute, contiguous 연산을 통해 텐서의 형태를 변경했는데, 이는 GPU 메모리 복사 및 재배열 비용을 발생시켜 성능 저하의 원인이 되었습니다.

`pixel_shuffle` 함수의 최적화

기존 pixel_shuffle 함수는 다음과 같이 구현되어 있었습니다:

--- a/vllm/model_executor/models/nano_nemotron_vl.py
+++ b/vllm/model_executor/models/nano_nemotron_vl.py
@@ -1005,38 +1005,27 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             )
 
     def pixel_shuffle(self, x, scale_factor=0.5):
-        n, w, h, c = x.size()
-        # N, W, H, C --> N, W, H * scale, C // scale
-        x = x.view(
-            n,
-            w,
-            int(h * scale_factor),
-            int(c / scale_factor),
-        )
-        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
-        x = x.permute(0, 2, 1, 3).contiguous()
-        # N, H * scale, W, C // scale -->
-        # N, H * scale, W * scale, C // (scale ** 2)
-        x = x.view(
-            n,
-            int(h * scale_factor),
-            int(w * scale_factor),
-            int(c / (scale_factor * scale_factor)),
-        )
+        n, h, w, c = x.size()
+        r = int(1 / scale_factor)
+        new_h = h // r
+        new_w = w // r
+        new_c = c * r * r
+
+        x = x.view(n, new_h, r, new_w, r, c)
         if self.ps_version == "v1":
             warnings.warn(
                 "In ps_version 'v1', the height and width have not "
                 "been swapped back, which results in a transposed image.",
                 stacklevel=2,
             )
-            x = x.permute(0, 2, 1, 3).contiguous()
+            x = x.permute(0, 3, 1, 2, 4, 5).reshape(n, new_w, new_h, new_c)
         else:
-            x = x.permute(0, 2, 1, 3).contiguous()
+            x = x.permute(0, 1, 3, 2, 4, 5).reshape(n, new_h, new_w, new_c)
         return x
 
     def pixel_shuffle_dynamic_res(
         self, x: torch.Tensor, *, imgs_sizes: list[tuple[int, int]]
     ) -> torch.Tensor:
-        scale_factor = self.downsample_ratio
         patch_dim = self.patch_size
         seq_lens = calc_seq_lens(imgs_sizes, patch_dim)
         splits = torch.split(x, seq_lens, dim=-2)
@@ -1045,22 +1034,8 @@ def pixel_shuffle_dynamic_res(
             h = imgs_sizes[i][0] // patch_dim
             w = imgs_sizes[i][1] // patch_dim
             sv = sv.reshape(sv.shape[0], h, w, -1)
-
-            n, h, w, c = sv.size()
-
-            sv = sv.view(n, h, int(w * scale_factor), int(c / scale_factor))
-            sv = sv.permute(0, 2, 1, 3).contiguous()
-            sv = sv.view(
-                n,
-                int(w * scale_factor),
-                int(h * scale_factor),
-                int(c / (scale_factor * scale_factor)),
-            )
-
-            if self.ps_version == "v2":
-                sv = sv.permute(0, 2, 1, 3).contiguous()
-
-            sv = sv.reshape(sv.shape[0], -1, sv.shape[-1])
+            sv = self.pixel_shuffle(sv, scale_factor=self.downsample_ratio)
+            sv = sv.flatten(1, 2)
             out.append(sv)
 
         x = torch.cat(out, dim=-2)

기존 코드에서는 pixel_shuffle 함수 내에서 다음과 같은 연산이 수행되었습니다:

x.view(...): 텐서의 형태를 변경합니다. 이 과정에서 데이터의 실제 메모리 레이아웃은 변경되지 않습니다.
x.permute(...): 텐서의 차원 순서를 변경합니다. 이 연산은 종종 텐서를 비연속적인(non-contiguous) 상태로 만듭니다.
.contiguous(): permute 등으로 인해 비연속적이 된 텐서를 연속적인 메모리 레이아웃으로 만듭니다. 이 과정에서 데이터 복사가 발생할 수 있습니다.
다시 x.view(...): 변경된 형태에 맞춰 텐서를 재구성합니다.

이러한 view -> permute -> contiguous -> view 패턴은 여러 번의 중간 연산과 잠재적인 메모리 복사를 유발하여 비효율적이었습니다.

개선된 코드는 이 과정을 훨씬 간결하게 처리합니다:

n, h, w, c = x.size()
r = int(1 / scale_factor)
new_h = h // r
new_w = w // r
new_c = c * r * r

x = x.view(n, new_h, r, new_w, r, c)

if self.ps_version == "v1":
    x = x.permute(0, 3, 1, 2, 4, 5).reshape(n, new_w, new_h, new_c)
else:
    x = x.permute(0, 1, 3, 2, 4, 5).reshape(n, new_h, new_w, new_c)

핵심 변경 사항은 다음과 같습니다:

단일 view 연산으로 차원 병합: x.view(n, new_h, r, new_w, r, c)를 통해 원래 h와 w 차원에 포함된 scale_factor 관련 정보를 한 번에 분해합니다. 이는 h * scale_factor와 w * scale_factor를 계산하고 이를 다시 view하는 이전 방식보다 훨씬 직접적입니다.
contiguous() 제거: 개선된 코드는 permute 후에도 불필요한 contiguous() 호출을 제거했습니다. reshape 연산은 종종 permute와 함께 사용될 때 내부적으로 연속적인 메모리 레이아웃을 처리할 수 있어, 명시적인 contiguous() 호출이 필요 없을 수 있습니다. 이를 통해 불필요한 메모리 복사를 방지합니다.
변수명 변경: h, w 변수의 순서가 x.size()에서 가져온 실제 텐서의 차원 순서와 일치하도록 수정되었습니다 (n, h, w, c -> n, w, h, c 에서 n, h, w, c로). 이는 코드의 가독성을 높이고 잠재적인 버그를 예방합니다.

`pixel_shuffle_dynamic_res` 함수의 변경

pixel_shuffle_dynamic_res 함수에서도 유사한 최적화가 적용되었습니다. 기존에는 각 이미지 크기에 대해 개별적으로 pixel_shuffle과 유사한 로직을 반복 수행하고, 마지막에 reshape하는 방식이었습니다.

개선된 코드에서는 이 부분을 self.pixel_shuffle 함수 호출로 대체하고, 마지막에 flatten(1, 2)를 사용하여 시퀀스 길이를 재조정합니다.

@@ -1045,22 +1034,8 @@ def pixel_shuffle_dynamic_res(
             h = imgs_sizes[i][0] // patch_dim
             w = imgs_sizes[i][1] // patch_dim
             sv = sv.reshape(sv.shape[0], h, w, -1)
-
-            n, h, w, c = sv.size()
-
-            sv = sv.view(n, h, int(w * scale_factor), int(c / scale_factor))
-            sv = sv.permute(0, 2, 1, 3).contiguous()
-            sv = sv.view(
-                n,
-                int(w * scale_factor),
-                int(h * scale_factor),
-                int(c / (scale_factor * scale_factor)),
-            )
-
-            if self.ps_version == "v2":
-                sv = sv.permute(0, 2, 1, 3).contiguous()
-
-            sv = sv.reshape(sv.shape[0], -1, sv.shape[-1])
+            sv = self.pixel_shuffle(sv, scale_factor=self.downsample_ratio)
+            sv = sv.flatten(1, 2)
             out.append(sv)
 
         x = torch.cat(out, dim=-2)

이 변경은 코드 중복을 제거하고, pixel_shuffle 함수의 최적화된 로직을 재사용함으로써 전체적인 코드의 유지보수성과 효율성을 높입니다.

왜 이게 좋은가?

이번 PR의 가장 큰 장점은 성능 향상입니다. PR 설명에 포함된 벤치마크 결과는 이를 명확히 보여줍니다:

TTFT (Time To First Token) 감소: 특히 높은 동시성(Concurrency 32) 환경에서 TTFT가 최대 **-41.67%**까지 감소했습니다. 이는 첫 토큰을 생성하기까지 걸리는 시간이 크게 단축되었음을 의미하며, 사용자 경험에 직접적인 영향을 미칩니다.
처리량 (Throughput) 증가: 초당 처리하는 토큰 수가 최대 +16.34% 증가했습니다. 이는 동일 시간 동안 더 많은 요청을 처리할 수 있음을 의미하며, 서버의 효율성을 크게 높입니다.
정확도 유지: OCRBenchV2를 사용한 테스트에서 영어 및 중국어 평균 정확도에 미미한 변화만 있었으며, 이는 성능 향상이 모델의 정확성을 해치지 않았음을 보여줍니다.

이러한 성능 향상은 다음과 같은 이유로 가능했습니다:

불필요한 메모리 복사 제거: contiguous() 호출 제거는 GPU 메모리 복사 횟수를 줄여 직접적인 성능 향상을 가져옵니다. PyTorch에서 contiguous()는 텐서가 메모리 상에서 연속적인 블록을 차지하도록 보장하는 연산인데, permute와 같은 연산은 종종 텐서를 비연속적으로 만들고, 이를 다시 연속적으로 만들기 위해 데이터 복사가 필요합니다. 이 PR은 이러한 복사를 피하도록 연산 순서를 재구성했습니다.
연산 횟수 감소: 여러 단계의 view, permute, contiguous 연산을 단일 view와 permute/reshape 조합으로 줄임으로써, GPU 커널 실행 횟수를 줄이고 오버헤드를 감소시켰습니다.
코드 간결화 및 재사용: pixel_shuffle_dynamic_res에서 pixel_shuffle 함수를 재사용하고, flatten 연산을 추가하여 동적 해상도 처리 로직을 더 효율적으로 만들었습니다.

일반적인 교훈:

텐서 연산의 효율성: PyTorch와 같은 딥러닝 프레임워크에서 텐서 연산의 순서와 방식은 성능에 지대한 영향을 미칩니다. view, permute, contiguous, reshape와 같은 연산의 특성을 정확히 이해하고, 불필요한 메모리 복사를 최소화하는 것이 중요합니다.
GPU 최적화: GPU는 병렬 처리에 강하지만, 메모리 접근 패턴과 복사 비용에 민감합니다. 가능한 한 GPU 커널 호출 횟수를 줄이고, 메모리 접근을 효율적으로 만드는 것이 성능 향상의 핵심입니다.
코드 가독성과 성능의 균형: 이 PR은 복잡했던 텐서 연산을 더 간결하고 효율적인 방식으로 재구성하여 가독성과 성능을 동시에 개선했습니다. 때로는 더 직접적인 연산이 더 나은 성능을 보장합니다.

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

PR Analysis 의 다른글

이전글 [vllm] AMD ROCm을 위한 Triton 기반 W4A16 커널 도입: MI300X 성능 최적화 분석
현재글 : [vllm] vLLM Nemotron Nano VL: Pixel Shuffle 최적화를 통한 성능 향상 분석
다음글 [vllm] vLLM 성능 최적화: H2D 메모리 복사 병목 해결을 통한 추론 처리량 개선

[vllm] vLLM Nemotron Nano VL: Pixel Shuffle 최적화를 통한 성능 향상 분석

들어가며

코드 분석

`pixel_shuffle` 함수의 최적화

`pixel_shuffle_dynamic_res` 함수의 변경

왜 이게 좋은가?

리뷰 피드백

결론

References

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글