[sglang] [성능 최적화] SGLang `prepare_for_decode`에서 `latest_output_ids` H2D 복사 비동기화로 디코딩 처리량 30% 향상

2026년 6월 17일수정: 2026년 6월 17일

PR 링크: sgl-project/sglang#28491 상태: Merged | 변경: +7 / -11

들어가며

안녕하세요, 시니어 소프트웨어 엔지니어이자 기술 블로거입니다. 오늘은 대규모 언어 모델(LLM) 추론을 가속화하는 프레임워크인 SGLang에서 발생했던 흥미로운 성능 병목 현상과 이를 해결하기 위한 최적화 PR([Perf] Make latest_output_ids H2D non-blocking in prepare_for_decode)에 대해 심층적으로 분석해보려 합니다. 이 PR은 단 하나의 코드 변경으로 디코딩 처리량을 30% 이상 향상시키는 놀라운 결과를 가져왔습니다.

SGLang과 같은 LLM 추론 시스템에서 prepare_for_decode 함수는 매 디코딩 스텝마다 호출되는 핵심 경로에 위치합니다. 이 함수는 다음 토큰 생성을 위해 필요한 다양한 정보를 준비하는데, 그 중 latest_output_ids는 페널티(penalty) 계산 등 중요한 로직에 사용됩니다. 기존 구현에서는 이 latest_output_ids를 호스트(CPU)에서 생성한 후 디바이스(GPU)로 복사하는 과정에서 동기적인(blocking) torch.tensor(..., device=...) 호출을 사용하고 있었습니다. 이는 SGLang의 overlap scheduler가 GPU 연산과 데이터 전송을 오버랩시켜 효율을 극대화하려는 의도와 상충하며, 매 디코딩 스텝마다 불필요한 cudaStreamSynchronize를 유발하여 CPU가 GPU 작업 완료를 기다리게 만들었습니다. 특히 배치 사이즈가 작거나 낮은 지연 시간(low latency)이 요구되는 시나리오에서 이러한 CPU stall은 전체 시스템의 성능을 저해하는 주된 원인이 되었습니다.

이번 최적화는 이 동기적인 H2D(Host-to-Device) 복사 작업을 비동기적으로 변경하여, CPU와 GPU가 서로를 기다리지 않고 동시에 작업을 수행할 수 있도록 함으로써 전체 파이프라인의 효율성을 극대화하는 데 초점을 맞추고 있습니다.

코드 변경사항 분석

핵심 변경은 python/sglang/srt/managers/schedule_batch.py 파일의 prepare_for_decode 함수 내에서 이루어졌습니다.

`python/sglang/srt/managers/schedule_batch.py`

latest_output_ids 텐서를 생성하고 GPU로 복사하는 방식이 변경되었습니다.

Before (Blocking H2D):

            latest_output_ids = torch.tensor(
                [
                    (
                        req.output_ids[-1]
                        if len(req.output_ids)
                        else req.origin_input_ids[-1]
                    )
                    for req in self.reqs
                ],
                dtype=torch.int64,
                device=self.device,
            )

After (Non-blocking H2D):

            last_tokens = [
                req.output_ids[-1] if len(req.output_ids) else req.origin_input_ids[-1]
                for req in self.reqs
            ]
            # Non-blocking H2D so this per-step copy doesn't sync behind the forward.
            latest_output_ids = torch.tensor(last_tokens, dtype=torch.int64).to(
                self.device, non_blocking=True
            )

무엇이 변경되었고 왜 중요한가?

기존 코드에서는 torch.tensor() 함수를 호출할 때 device=self.device 인자를 직접 전달했습니다. PyTorch에서 torch.tensor()에 device 인자를 지정하면, 기본적으로 호스트 메모리에서 텐서를 생성한 후 지정된 디바이스(GPU)로 데이터를 동기적으로 복사합니다. 이 동기적 복사 작업은 CPU가 GPU로의 데이터 전송이 완료될 때까지 기다리게 만들며, 이는 cudaStreamSynchronize와 같은 내부적인 동기화 호출로 이어질 수 있습니다.

변경 후 코드에서는 두 단계로 나뉘어 처리됩니다.

먼저 last_tokens 리스트를 Python(CPU)에서 생성합니다.
그 다음, torch.tensor(last_tokens, dtype=torch.int64)를 호출하여 CPU 메모리에 텐서를 생성합니다.
마지막으로, .to(self.device, non_blocking=True)를 사용하여 이 텐서를 GPU로 비동기적으로 복사합니다.

여기서 핵심은 non_blocking=True 인자입니다. 이 인자는 PyTorch에게 GPU로의 데이터 전송을 비동기적으로 수행하도록 지시합니다. 즉, CPU는 데이터 전송이 완료될 때까지 기다리지 않고 즉시 다음 작업을 시작할 수 있습니다. GPU는 백그라운드에서 데이터를 수신하며, 이 과정은 다른 GPU 연산과 오버랩될 수 있습니다. SGLang의 forward_stream이 이미 schedule_stream을 기다리도록 설계되어 있기 때문에, 비동기 복사가 forward 연산이 데이터를 사용하기 전에 완료되는 순서 일관성(ordering consistency)은 유지됩니다.

이러한 비동기 H2D 복사는 특히 min_new_tokens와 같이 latest_output_ids 텐서의 실제 값을 읽지 않고 단순히 존재 여부만 확인하는 페널티 로직의 경우 순수한 오버헤드를 제거하는 효과를 가져옵니다.

왜 이 최적화가 좋은가?

이 최적화는 SGLang의 LLM 추론 성능에 상당한 긍정적인 영향을 미쳤습니다. 벤치마크 결과는 다음과 같습니다.

벤치마크 환경:

모델: openai/gpt-oss-120b
TP(Tensor Parallelism): 4 (GB300)
스케줄러: overlap scheduler
워크로드: aiperf 1k-in / 1k-out, 동시성 1 (bs=1 디코딩), min_tokens 설정으로 페널티 로직 활성화.

	fwd occupancy	decode throughput	inter-token latency
before (blocking H2D)	71.4%	255 tok/s	3.91 ms
after (non-blocking H2D)	86.9%	330 tok/s	2.96 ms

주요 개선점:

GPU 활용도(fwd occupancy) 증가: 71.4%에서 86.9%로 크게 증가했습니다. 이는 GPU가 데이터 전송을 기다리지 않고 더 많은 시간 동안 실제 연산(forward pass)을 수행할 수 있게 되었음을 의미합니다. 비동기 복사를 통해 데이터 전송과 연산이 효과적으로 오버랩되었기 때문입니다.
디코딩 처리량(decode throughput) 30% 향상: 255 tok/s에서 330 tok/s로 증가했습니다. 이는 초당 처리할 수 있는 토큰 수가 크게 늘어났음을 보여주며, 시스템의 전반적인 처리 능력이 향상되었음을 의미합니다.
인터-토큰 지연 시간(inter-token latency) 약 24% 감소: 3.91 ms에서 2.96 ms로 줄어들었습니다. 이는 각 토큰을 생성하는 데 걸리는 시간이 단축되어, 사용자 경험 측면에서 더 빠른 응답 속도를 제공할 수 있게 되었음을 의미합니다.

이러한 성능 향상은 특히 작은 배치 사이즈나 낮은 지연 시간이 중요한 실시간 LLM 서비스 환경에서 매우 중요합니다. 불필요한 동기화 지점을 제거함으로써 CPU stall을 줄이고, GPU의 컴퓨팅 자원을 최대한 활용할 수 있게 된 것이 핵심입니다.

일반적인 교훈:

H2D/D2H 전송의 중요성: GPU 가속 애플리케이션에서 호스트-디바이스 간 데이터 전송은 종종 간과되는 성능 병목입니다. 세심한 프로파일링을 통해 이러한 전송이 동기적으로 이루어지고 있는지 확인하고, 가능한 경우 비동기 전송으로 전환하는 것이 중요합니다.
non_blocking=True 활용: PyTorch에서 to() 메서드나 cuda() 메서드 사용 시 non_blocking=True 인자를 적극적으로 활용하여 데이터 전송과 연산 간의 오버랩을 최대화해야 합니다. 이는 CUDA Stream과 함께 사용하여 GPU 파이프라인의 효율성을 높이는 핵심 기법입니다.
숨겨진 동기화 지점 찾기: 코드의 한 줄이 전체 시스템의 성능을 저해할 수 있습니다. 특히 device 인자를 직접 지정하는 텐서 생성과 같은 작업은 예상치 못한 동기화를 유발할 수 있으므로 주의 깊게 검토해야 합니다.
프로파일링의 힘: 이 최적화는 sglang:fwd_occupancy와 같은 메트릭을 통해 병목을 정확히 식별하고, 개선 효과를 수치로 검증했기 때문에 가능했습니다. 정기적인 프로파일링은 성능 최적화의 필수 요소입니다.

리뷰 댓글 분석

PR에 대한 리뷰 댓글은 주로 CI/테스트 재실행 요청(hnyls2002] /tag-and-rerun-ci, [hnyls2002] /rerun-test registered/sampling/test_penalty.py registered/core/test_srt_endpoint.py)이었습니다. 이는 변경 사항이 기존의 샘플링 로직, 특히 페널티 계산과 관련된 부분에 영향을 미 미치지 않는지 확인하려는 의도였음을 보여줍니다. 코드 자체에 대한 심도 있는 기술적 논의는 없었지만, 이러한 테스트 요청은 성능 최적화가 기존 기능의 정확성을 훼손하지 않도록 보장하는 중요한 과정임을 시사합니다. latest_output_ids가 penalizer_orchestrator에 사용되므로, 관련 테스트가 통과하는 것이 중요합니다.

결론

이번 SGLang PR은 작은 코드 변경 하나가 시스템 전체 성능에 얼마나 큰 영향을 미 미칠 수 있는지 보여주는 좋은 사례입니다. latest_output_ids의 H2D 복사를 비동기화함으로써, SGLang은 GPU 활용도를 높이고, 디코딩 처리량을 30% 향상시키며, 인터-토큰 지연 시간을 크게 줄였습니다. 이는 LLM 추론 시스템을 개발할 때 데이터 전송 방식과 동기화 지점에 대한 깊은 이해가 필수적임을 다시 한번 일깨워줍니다. 앞으로도 이러한 세밀한 최적화를 통해 더욱 빠르고 효율적인 LLM 서비스를 제공할 수 있기를 기대합니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [vllm] vLLM에서 Flashinfer 기반 Non-gated MoE bf16 지원 최적화 분석
현재글 : [sglang] [성능 최적화] SGLang `prepare_for_decode`에서 `latest_output_ids` H2D 복사 비동기화로 디코딩 처리량 30% 향상
다음글 [sglang] SGLang 성능 최적화: Speculative Decoding의 H2D 병목 해결 및 코드 중복 제거