본문으로 건너뛰기

[sglang] SGLang 성능 최적화: DeepSeek-v4 SWA 인덱스 변환 Hoisting 및 백엔드 통합

PR 링크: sgl-project/sglang#34890 상태: Merged | 변경: +68 / -30

들어가며

LLM 추론 엔진인 SGLang에서 Sliding Window Attention (SWA)은 메모리 효율성을 높이는 핵심 기술입니다. 하지만 SWA를 사용할 때 전체 KV 캐시 인덱스를 SWA 전용 인덱스로 변환하는 과정이 수반되는데, 이 과정이 모델의 모든 레이어에서 반복적으로 발생하면 무시할 수 없는 오버헤드가 됩니다.

이번 PR([Perf] Hoist DSv4 draft-extend SWA write locs; unify SWA graph buffer naming)은 특히 DeepSeek-v4 (DSv4)의 draft-extend 모드에서 발생하는 비효율을 해결합니다. 핵심 아이디어는 "CUDA Graph 내부에서 매 레이어마다 수행되던 인덱스 변환 로직을 그래프 외부(Host/Step 단위)로 끌어올리는(Hoist)" 것입니다. 또한, 여러 하드웨어 백엔드(NVIDIA, Ascend, XPU 등) 간에 파편화되어 있던 SWA 관련 버퍼 명칭을 통일하여 코드 응집력을 높였습니다.

코드 분석: DeepSeek-v4 백엔드 최적화

1. SWA 쓰기 위치 계산의 Hoisting

기존 구조에서는 out_cache_loc을 SWA 인덱스로 변환하는 작업이 CUDA Graph 내부에 캡처되어, 그래프가 리플레이될 때마다 모든 레이어의 store_cache 단계에서 개별적으로 수행되었습니다. 이를 개선하기 위해 그래프 상태 초기화 시점에 전용 버퍼를 미리 할당하고, 매 스텝 시작 시점에 한 번만 변환하여 버퍼에 써두도록 변경했습니다.

Before: 기존에는 swa_out_cache_loc이 명시적으로 관리되지 않아, 그래프 내부에서 레이어별로 변환이 일어날 가능성이 있었습니다.

After:

# python/sglang/srt/layers/attention/deepseek_v4_backend.py

def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int) -> None:
    # ... 기존 로직 ...
    if self.is_draft_runner:
        # Draft-extend SWA write-target buffer 할당
        self.cuda_graph_swa_out_cache_loc = torch.zeros(
            max_num_tokens, dtype=torch.int32, device=self.device
        )

그리고 매 스텝의 메타데이터 초기화 시점에 이 버퍼를 채워줍니다.

# python/sglang/srt/layers/attention/deepseek_v4_backend.py

def _fill_cuda_graph_swa_out_cache_loc(
    self, out_cache_loc: Optional[torch.Tensor]
) -> Optional[torch.Tensor]:
    buf = self.cuda_graph_swa_out_cache_loc
    if buf is None or out_cache_loc is None or out_cache_loc.shape[0] > buf.shape[0]:
        return None
    
    n = out_cache_loc.shape[0]
    buf[n:].zero_()
    # 그래프 외부에서 미리 변환 수행 (Hoist)
    buf[:n].copy_(
        self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to(torch.int32)
    )
    return buf[:n]

2. Eager Mode에서의 최적화

CUDA Graph를 사용하지 않는 Eager 모드나 배치 사이즈가 너무 커서 그래프 범위를 벗어나는 경우에도, 레이어마다 변환하는 대신 스텝당 한 번만 변환하도록 로직을 개선했습니다.

Before/After 비교:

# python/sglang/srt/layers/attention/deepseek_v4_backend.py

# After 로직
swa_out_cache_loc = self._fill_cuda_graph_swa_out_cache_loc(out_cache_loc)
if swa_out_cache_loc is None and out_cache_loc is not None:
    # Eager-only miss: 레이어별 store_cache 시점이 아닌, 여기서 한 번만 변환
    swa_out_cache_loc = self.token_to_kv_pool.translate_loc_from_full_to_swa(
        out_cache_loc
    ).to(torch.int32)

코드 분석: 백엔드 간 명명 규칙 통일 (Cross-backend alignment)

기존에는 하드웨어 백엔드마다 SWA 버퍼를 부르는 이름이 달랐습니다 (swa_out_cache_loc_buf vs cuda_graph_swa_out_cache_loc). 이를 cuda_graph_swa_out_cache_loc으로 통일하여 유지보수성을 강화했습니다.

FlashAttention 백엔드 사례

Before:

self.swa_out_cache_loc_buf = torch.zeros(max_num_tokens, ...)
# ...
metadata.swa_out_cache_loc = self.swa_out_cache_loc_buf[:num_tokens]

After:

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

self.cuda_graph_swa_out_cache_loc = torch.zeros(max_num_tokens, ...)
# ...
metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[:num_tokens]

이 변경사항은 ascend_backend.py 등 다른 백엔드 파일들에도 동일하게 적용되어, 전체 프로젝트 수준에서 일관된 인터페이스를 제공하게 되었습니다.

왜 이게 좋은 최적화인가?

  1. CUDA Graph 리플레이 오버헤드 감소: CUDA Graph는 정적인 연산을 빠르게 실행하는 데 최적화되어 있습니다. 그래프 내부에 복잡한 인덱싱 연산(translate_loc_from_full_to_swa)이 포함되면, 비록 GPU 연산이라 할지라도 매 레이어(DeepSeek-V4의 경우 수십 개의 레이어)마다 반복 실행되면서 지연 시간을 발생시킵니다. 이를 그래프 외부로 Hoisting 함으로써, 그래프 내부 연산은 단순한 메모리 참조로 대체됩니다.

  2. 중복 연산 제거 (Redundancy Elimination): draft-extend 과정에서 모든 레이어는 동일한 out_cache_loc을 공유합니다. 이전에는 레이어 수만큼 반복되던 변환 연산을 스텝당 단 1회로 줄였습니다.

  3. 메모리 안정성: init_cuda_graph_state에서 버퍼를 미리 할당(pre-allocate)하고 관리함으로써, 런타임 중에 임시 텐서가 생성되어 발생할 수 있는 메모리 파편화나 주소 캡처 오류를 방지합니다.

  4. 추상화 수준의 통일: 다양한 하드웨어(NPU, GPU, XPU)를 지원하는 SGLang의 특성상, 동일한 역할을 하는 버퍼의 이름을 통일하는 것은 코드 가독성과 향후 공통 기능 구현 시 매우 중요한 기반이 됩니다.

결론

이번 PR은 대규모 모델 추론 시 발생하는 미세한 오버헤드를 잡기 위한 정교한 엔지니어링의 결과물입니다. 특히 DeepSeek-v4와 같이 복잡한 구조를 가진 모델에서 CUDA Graph의 이점을 극대화하기 위해 "무엇을 그래프 안에 남기고, 무엇을 밖으로 뺄 것인가"에 대한 좋은 기준을 제시합니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글