[sglang] SGLang에서 DP Attention, TBO, Shared Experts Fusion 동시 최적화 구현

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

PR 링크: sgl-project/sglang#27510 상태: Merged | 변경: +79 / -8

들어가며

최근 대규모 언어 모델(LLM) 추론 최적화에서 DP(Data Parallel) Attention, TBO(Two-Batch Overlap), 그리고 Shared Experts Fusion은 성능 향상을 위한 핵심 기술입니다. 하지만 이 기능들을 동시에 활성화할 때 발생하는 리소스 충돌과 동기화 문제는 시스템 안정성을 저해하는 요소였습니다. 본 PR은 SGLang에서 이 세 가지 기능을 동시에 사용할 수 있도록 아키텍처를 개선하고, 그 과정에서 발생한 런타임 에러를 해결하여 전체적인 추론 처리량을 약 2.5% 향상시켰습니다.

코드 분석

1. `python/sglang/srt/server_args.py`: 제약 조건 제거

기존에는 TBO와 Shared Experts Fusion을 동시에 사용하는 것을 명시적으로 금지하고 있었습니다. 이를 제거하여 사용자가 두 기능을 동시에 활성화할 수 있도록 허용했습니다.

# Before
if self.enable_two_batch_overlap and self.enforce_shared_experts_fusion:
    raise ValueError("--enable-two-batch-overlap and --enforce-shared-experts-fusion cannot be used together.")

# After: 위 블록 삭제

2. `python/sglang/srt/batch_overlap/two_batch_overlap.py`: 배치 필터링 로직 수정

DP Attention 환경에서 각 랭크마다 토큰 수가 다를 때 발생하는 패딩 문제로 인해 rids 개수가 일치하지 않는 버그가 있었습니다. 이를 해결하기 위해 rids 슬라이싱 로직을 추가했습니다.

# After: rids 불일치 시 슬라이싱 처리
elif key == "rids" and len(old_value) != num_seqs:
    output_dict[key] = old_value[
        start_seq_index : min(end_seq_index, len(old_value))
    ]
    continue

3. `python/sglang/srt/models/deepseek_v2.py`: IDLE 배치 처리 개선

IDLE 배치에서도 패딩으로 인해 더미 행이 생성되는데, 기존 로직은 이를 제대로 처리하지 못해 DeepEP 디스패치 시 차원 불일치 에러를 유발했습니다.

# Before
if is_non_idle_and_non_empty(state.forward_batch.forward_mode, state.hidden_states_mlp_input):

# After
if state.hidden_states_mlp_input.shape[0] > 0:

왜 이게 좋은가

이번 최적화는 단순히 기능을 여는 것에 그치지 않고, 복잡한 분산 환경에서의 데이터 정합성을 보장했습니다. 벤치마크 결과, Output throughput이 179.66 tok/s에서 184.07 tok/s로 약 2.5% 향상되었으며, Mean TTFT(Time To First Token) 또한 961.83ms에서 936.71ms로 단축되는 성과를 거두었습니다.

핵심 교훈

분산 환경의 비대칭성: DP Attention 사용 시 각 랭크의 토큰 수가 다를 수 있음을 항상 고려해야 합니다.
패딩과 메타데이터의 동기화: 패딩을 적용할 때 단순히 텐서만 늘리는 것이 아니라, rids와 같은 메타데이터도 함께 슬라이싱하거나 패딩해야 런타임 에러를 방지할 수 있습니다.
조건부 로직의 유연성: is_non_idle과 같은 추상화된 체크보다는 실제 데이터의 shape를 기반으로 로직을 분기하는 것이 엣지 케이스 처리에 더 견고합니다.