[sglang] sglang 성능 최적화: torch.compile 퓨전 복원을 통한 TopK 후처리 개선

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

PR 링크: sgl-project/sglang#21771 상태: Merged | 변경: +None / -None

들어가며

최근 sglang 프로젝트에서는 모델 추론 성능 향상을 위한 다양한 최적화 작업이 진행되고 있습니다. 특히 Expert Parallelism (EPLB)과 같은 고급 병렬 처리 기법을 사용할 때, 효율적인 TopK 연산 후처리는 전체 성능에 지대한 영향을 미칩니다. 이번 PR(#16945)은 이전 버전에서 torch.compile을 통한 커널 퓨전(kernel fusion)으로 최적화되었던 TopK 후처리 로직이 리팩토링 과정에서 분리되면서 발생한 성능 회귀(regression)를 해결하는 데 중점을 둡니다.

이전에는 _biased_grouped_topk_postprocess라는 @torch.compile 데코레이터가 적용된 함수를 통해 여러 연산이 하나의 CUDA 커널로 융합되어 실행되었습니다. 하지만 PR #16945에서는 이 로직이 _post_process_topk_ids 함수 내로 인라인(inline)되면서, topk_ids_logical_to_physical과 _mask_topk_ids_padded_region이 별도의 CUDA 커널로 분리되어 실행되는 문제가 발생했습니다. 이는 특히 Expert Parallelism 환경에서 불필요한 커널 실행 오버헤드를 유발하여 성능 저하의 원인이 되었습니다.

본 글에서는 이 PR이 어떻게 이전의 최적화된 커널 퓨전 상태를 복원하고, 그로 인해 어떤 성능 개선 효과를 얻을 수 있는지 코드 변경 사항을 중심으로 상세히 분석해보겠습니다.

코드 분석: 파일별 변경점

`python/sglang/srt/layers/moe/topk.py`

이 PR의 핵심 변경은 topk.py 파일 내 _post_process_topk_ids 함수의 구현 방식에 있습니다. 이전 리팩토링 이전에는 @torch.compile로 최적화된 _biased_grouped_topk_postprocess 함수가 호출되었으나, PR #16945 이후에는 해당 함수가 인라인되어 두 개의 별도 함수 호출로 대체되었습니다.

변경 전 (Before - PR #16945 이후, 성능 회귀 발생 시점):

-        topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
-        _mask_topk_ids_padded_region(topk_ids, num_token_non_padded)

위 코드는 topk_ids_logical_to_physical과 _mask_topk_ids_padded_region 함수를 각각 호출합니다. 이 함수들은 @torch.compile 데코레이터가 적용되어 있었음에도 불구하고, 호출되는 방식이 변경되면서 torch.compile의 퓨전 최적화 대상에서 제외되었습니다. 결과적으로 이 두 연산은 CUDA 상에서 별도의 커널로 실행되어 오버헤드가 발생했습니다.

변경 후 (After - 본 PR):

+        topk_ids = _biased_grouped_topk_postprocess(
+            topk_ids, expert_location_dispatch_info, num_token_non_padded
+        )

본 PR에서는 이 두 개의 개별 함수 호출을 다시 기존의 @torch.compile(dynamic=True, backend=get_compiler_backend()) 데코레이터가 적용된 _biased_grouped_topk_postprocess 함수 호출로 대체했습니다. 이 변경을 통해 topk_ids_logical_to_physical과 _mask_topk_ids_padded_region에서 수행되던 연산들이 다시 _biased_grouped_topk_postprocess 함수 내부로 포함되어, torch.compile에 의해 하나의 CUDA 커널로 퓨전될 수 있게 됩니다. 이는 불필요한 커널 실행을 줄여 성능을 향상시킵니다.

왜 이게 좋은가: 성능 개선과 일반적인 교훈

성능 개선

이 PR의 주된 목표는 성능 회귀를 복원하는 것입니다. PR 설명에 따르면, 이전 버전에서는 TopK 후처리 로직이 torch.compile을 통해 효과적으로 퓨전되어 단일 커널로 실행되었습니다. 하지만 리팩토링 이후 두 개의 별도 연산으로 분리되면서 CUDA 상에서 각각 별도의 커널로 실행되었고, 이는 특히 Expert Parallelism (EPLB)과 같이 많은 전문가(expert)를 활용하는 복잡한 추론 경로에서 상당한 오버헤드를 발생시켰습니다.

PR 설명에 첨부된 이미지들은 이러한 성능 변화를 시각적으로 보여줍니다:

Current ToT (Top of Tree - 회귀 발생 후): 두 개의 별도 연산이 실행됨을 나타냅니다.
This PR (복원 후): 단일 연산으로 퓨전되어 실행됨을 나타냅니다.

이러한 퓨전 복원을 통해 불필요한 커널 호출 및 데이터 전송 오버헤드가 제거되어, 전체적인 추론 속도가 향상될 것으로 기대됩니다. 구체적인 성능 수치(예: ms 단위의 속도 향상 또는 처리량 증가)가 PR에 명시되어 있다면 더 좋았겠지만, torch.compile의 퓨전 복원이라는 기술적 개선 자체만으로도 상당한 성능 향상을 기대할 수 있습니다.

일반적인 교훈

torch.compile의 중요성: PyTorch에서 torch.compile은 연산 퓨전, 커널 최적화 등을 통해 상당한 성능 향상을 제공하는 강력한 도구입니다. 복잡한 모델이나 연산 그래프에서는 torch.compile의 효과를 극대화하는 것이 중요합니다.
리팩토링 시 성능 회귀 주의: 코드의 가독성이나 유지보수성을 높이기 위한 리팩토링 과정에서 의도치 않은 성능 회귀가 발생할 수 있습니다. 특히 @torch.compile과 같이 성능에 직접적인 영향을 미치는 최적화 기법이 적용된 코드를 수정할 때는, 해당 최적화가 유지되는지 반드시 검증해야 합니다.
커널 퓨전의 이점: 여러 개의 작은 연산을 하나의 큰 커널로 융합하는 것은 GPU 컴퓨팅에서 매우 중요합니다. 이는 커널 실행 오버헤드를 줄이고, 메모리 접근 패턴을 개선하며, 데이터 재사용률을 높여 전반적인 성능을 향상시킵니다.
테스트 및 모니터링: PR 설명에 언급된 것처럼, 리뷰어(@fzyzcjy)의 지적을 통해 문제가 발견되었습니다. 이는 코드 변경 사항에 대한 철저한 리뷰와 함께, 성능 지표를 지속적으로 모니터링하는 자동화된 테스트의 중요성을 강조합니다.

리뷰 댓글 분석

주요 리뷰 댓글은 @fzyzcjy의 지적으로, 이 PR이 해결하고자 하는 문제의 핵심을 정확히 짚었습니다:

qq: does this mean this will launch a kernel while this should be fused in many cases

이 댓글은 _post_process_topk_ids 함수 내에서 topk_ids_logical_to_physical과 _mask_topk_ids_padded_region이 별도의 함수로 분리되어 호출될 때, 이 연산들이 원래는 @torch.compile에 의해 하나의 커널로 퓨전될 수 있었음에도 불구하고 분리되어 각각 커널을 실행하게 된다는 점을 지적합니다. 이는 명백한 성능 저하 요인이었으며, 본 PR은 이 지적을 바탕으로 torch.compile의 퓨전 기능을 복원하는 방향으로 수정되었습니다.