[sglang] SGLang DFlash 최적화: 호스트-디바이스 동기화 제거를 통한 추론 성능 향상

PR 링크: sgl-project/sglang#31468 상태: Merged | 변경: +528 / -50

들어가며

LLM 추론 엔진에서 Speculative Decoding은 성능 최적화의 핵심입니다. SGLang의 DFlash는 효율적인 추론을 제공하지만, 기존 구현에서는 매 디코드 단계마다 cudaStreamSynchronize를 호출하여 호스트와 디바이스 간의 동기화가 발생했습니다. 이로 인해 CPU가 다음 단계의 작업을 준비하지 못하고 대기하는 병목 현상이 발생했습니다. 본 PR은 이러한 per-step 동기화를 제거하여 CPU가 한 단계 앞서 작업을 큐에 넣을 수 있도록(Run-ahead) 개선한 사례입니다.

코드 분석

1. Triton 커널을 통한 Compact Draft-Cache 재구축

기존에는 lengths.max().item() 호출과 같은 D2H(Device-to-Host) 작업이 동기화를 강제했습니다. 이를 rebuild_compact_draft_req_to_token이라는 단일 Triton 커널로 통합했습니다.

@triton.jit
def rebuild_compact_draft_req_to_token(...):
    # ... (생략) ...
    for i in range(num_loop):
        col = offs + i * BLOCK
        in_prefix = col < prefix_len
        in_block = (col >= prefix_len) & (col < total)
        src = tl.load(src_row + start + col, mask=in_prefix, other=0)
        blk = tl.load(verify_row + (col - prefix_len), mask=in_block, other=0)
        val = tl.where(in_prefix, src, blk)
        tl.store(dst_row + col, val, mask=in_prefix | in_block)

이 방식은 데이터 의존적인 형상을 커널 내부에서 처리함으로써 호스트가 GPU의 상태를 읽기 위해 대기할 필요를 없앴습니다.

2. HybridAttnBackend의 CPU 시퀀스 길이 위임

HybridAttnBackend에서 needs_cpu_seq_lens 속성을 하위 백엔드에 위임하도록 수정하여, 불필요한 동기화를 방지했습니다.

# Before: 기본값 True로 인해 무조건 동기화 발생
# After:
self.needs_cpu_seq_lens = (
    prefill_backend.needs_cpu_seq_lens or decode_backend.needs_cpu_seq_lens
)

3. TP(Tensor Parallelism) 환경에서의 샘플러 최적화

TP > 1 환경에서 샘플링 로직을 CUDA 그래프 내부에 통합하여, 호스트로 데이터를 가져오는 과정을 제거했습니다.

왜 이게 좋은가

이번 최적화의 핵심은 'Launch Lead'의 확보입니다. 기존에는 cudaStreamSynchronize로 인해 CPU가 GPU의 작업 완료를 기다려야 했으나, 개선 후에는 launch lead가 약 -51 µs에서 +19.1 ms로 비약적으로 상승했습니다. 결과적으로:

• 성능 향상: 낮은 동시성(cc1)에서 약 1213%, 높은 동시성에서 35%의 토큰 처리량(tok/s/user) 향상을 확인했습니다.
• 교훈: 추론 엔진 설계 시, 호스트와 디바이스 사이의 데이터 동기화 지점을 최소화하고, 가능한 한 모든 연산을 GPU 커널 내에서 닫힌 루프(Closed-loop)로 처리하는 것이 성능 최적화의 핵심임을 보여줍니다.

리뷰 과정에서 Event() 할당과 같은 부수적인 최적화는 별도의 PR로 분리하는 등, 기술적 엄밀함을 유지하며 점진적으로 개선을 진행한 점이 인상적입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.cuda.Stream.synchronize.html
• https://triton-lang.org/main/index.html
• https://docs.sglang.ai/developer_guide/contribution_guide.html

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