[sglang] SGLang: ROCm 환경에서 RMSNorm의 배치 불변성 확보를 통한 결정론적 추론 개선

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

PR 링크: sgl-project/sglang#28787 상태: Merged | 변경: +29 / -0

들어가며

최근 대규모 언어 모델(LLM)의 발전과 함께, 이들을 효율적으로 서빙하고 추론하는 고성능 프레임워크의 중요성이 커지고 있습니다. SGLang은 이러한 요구사항을 충족하기 위해 설계된 강력한 프레임워크 중 하나입니다. LLM 추론에서 성능 최적화는 필수적이지만, 때로는 최적화된 커널이 특정 조건에서 예측 불가능한 동작을 유발할 수 있습니다. 특히, '결정론적 추론(deterministic inference)'은 동일한 입력에 대해 항상 동일한 출력을 보장해야 하므로, 모델의 신뢰성과 재현성을 위해 매우 중요합니다.

이번에 분석할 SGLang의 GitHub PR ([AMD] Fix RMSNorm batch-invariance on ROCm under deterministic inference)은 ROCm (AMD GPU) 환경에서 RMSNorm 레이어의 배치 불변성(batch-invariance) 문제를 해결하여 결정론적 추론을 가능하게 한 중요한 개선 사항입니다. 기존에는 최적화된 rms_norm 또는 fused_add_rms_norm 커널이 배치 형태에 따라 동일한 입력에 대해 미묘하게 다른 출력을 생성하는 문제가 있었습니다. 이는 특히 --enable-deterministic-inference 모드에서 logprob "wobble" 현상으로 나타나, 추론 결과의 일관성을 해쳤습니다. 이 PR은 forward_cuda 경로에서 이미 적용되던 보호 로직을 forward_hip (ROCm) 경로에도 미러링하여 이 문제를 해결합니다.

코드 분석

이 PR의 핵심 변경 사항은 python/sglang/srt/layers/layernorm.py 파일 내의 RMSNorm 클래스에 있습니다. 주로 forward_hip 메서드와 리뷰어의 피드백을 반영하여 forward_aiter 메서드에 조건부 로직을 추가하여 배치 불변성을 확보합니다.

`forward_hip` 메서드 변경

forward_hip 메서드는 ROCm 환경에서 RMSNorm 연산을 처리합니다. 기존에는 VLLM의 최적화된 rms_norm 커널을 사용하여 성능을 높였지만, 이 커널이 배치 불변성을 보장하지 않는다는 문제가 있었습니다. PR은 is_batch_invariant_mode_enabled() 함수를 통해 결정론적 추론 모드인지 확인하고, 이 모드일 경우 배치 불변성이 보장되는 경로로 폴백하도록 합니다.

Before:

        if not _has_vllm_rms_norm:
            return self.forward_native(x, residual, post_residual_addition)

        if not x.is_contiguous():
            # NOTE: Remove this if aiter kernel supports discontinuous input
            x = x.contiguous()

After:

        if not _has_vllm_rms_norm:
            return self.forward_native(x, residual, post_residual_addition)

        if is_batch_invariant_mode_enabled():
            if (
                residual is not None
                or self.cast_x_before_out_mul
                or get_global_server_args().rl_on_policy_target == "fsdp"
            ):
                return self.forward_native(x, residual, post_residual_addition)
            return rms_norm_batch_invariant(
                x,
                self.weight.data,
                self.variance_epsilon,
            )

        if not x.is_contiguous():
            # NOTE: Remove this if aiter kernel supports discontinuous input
            x = x.contiguous()

변경 후 코드를 보면, is_batch_invariant_mode_enabled()가 True일 경우, residual이 존재하거나, cast_x_before_out_mul 플래그가 설정되었거나, rl_on_policy_target이 fsdp인 특정 조건에서는 self.forward_native (PyTorch의 기본 구현)로 폴백합니다. 그 외의 경우에는 rms_norm_batch_invariant라는 배치 불변성이 보장되는 함수를 호출합니다. 이는 최적화된 퓨즈드 커널 대신, 정확성을 우선하는 경로를 선택함으로써 결정론적 추론을 가능하게 합니다.

`forward_aiter` 메서드 변경

초기 PR에서는 forward_hip에만 변경이 적용되었으나, 리뷰어 @1am9trash의 피드백에 따라 forward_aiter 메서드에도 동일한 보호 로직이 추가되었습니다. forward_aiter는 RMSNorm의 또 다른 추론 경로로, 이 경로에서도 배치 불변성을 확보하는 것이 중요합니다.

Before:

        if needs_reshape:
            x = x.contiguous().reshape(-1, original_shape[-1])
        elif not x.is_contiguous():
            x = x.contiguous()
        # Fused (add +) rmsnorm + zero-pad path. Triggered when caller
        # constructed RMSNorm with x_pad_to_multiple > 0. Output last
        # dim is padded up; residual_out stays at original width. Used

After:

        if needs_reshape:
            x = x.contiguous().reshape(-1, original_shape[-1])
        elif not x.is_contiguous():
            x = x.contiguous()
        if is_batch_invariant_mode_enabled():
            if (
                residual is not None
                or self.cast_x_before_out_mul
                or get_global_server_args().rl_on_policy_target == "fsdp"
                or (self._fused_pad_kernel is not None and self.x_pad_to_multiple > 0)
            ):
                return self.forward_native(x, residual, post_residual_addition)
            out = rms_norm_batch_invariant(
                x,
                self.weight.data,
                self.variance_epsilon,
            )
            if needs_reshape:
                out = out.reshape(original_shape)
            return out
        # Fused (add +) rmsnorm + zero-pad path. Triggered when caller
        # constructed RMSNorm with x_pad_to_multiple > 0. Output last
        # dim is padded up; residual_out stays at original width. Used

forward_aiter에서도 forward_hip과 유사하게 is_batch_invariant_mode_enabled() 체크가 추가되었으며, 특정 조건에서 forward_native로 폴백하거나 rms_norm_batch_invariant를 호출하도록 변경되었습니다. 이는 SGLang의 다양한 추론 경로에서 일관된 결정론적 동작을 보장하는 데 기여합니다.

왜 이게 좋은 최적화/개선인가

이 PR은 단순한 버그 수정 이상의 의미를 가집니다. 이는 고성능 LLM 추론 프레임워크 개발에서 중요한 여러 원칙을 보여주는 좋은 예시입니다.

정확성 및 신뢰성 확보: LLM의 추론 결과는 재현 가능해야 합니다. 특히 연구, 개발, 배포 단계에서 동일한 입력에 대해 항상 동일한 출력을 얻는 것은 디버깅, 모델 검증, 그리고 사용자 신뢰도 측면에서 매우 중요합니다. 이 PR은 logprob "wobble"과 같은 미묘한 비결정론적 동작을 제거하여 SGLang의 신뢰성을 크게 향상시킵니다.
안전하고 조건부적인 최적화: 이 개선은 무조건적으로 최적화된 커널을 비활성화하는 것이 아닙니다. is_batch_invariant_mode_enabled()라는 명확한 조건부 플래그를 통해, 결정론적 추론이 필요한 경우에만 배치 불변성이 보장되는 경로를 사용하도록 합니다. 이는 결정론적 동작이 필수가 아닌 일반적인 고성능 추론 시나리오에서는 기존의 최적화된 퓨즈드 커널을 계속 활용할 수 있도록 하여, 성능 저하 없이 정확성을 개선하는 '안전한 최적화' 접근 방식입니다.
ROCm 환경 지원 강화: AMD GPU를 사용하는 ROCm 환경은 LLM 추론에서 점점 더 중요해지고 있습니다. 이 PR은 ROCm 환경에서 SGLang의 안정성과 기능적 완성도를 CUDA 환경과 동등한 수준으로 끌어올리는 데 기여합니다. 이는 더 넓은 하드웨어 생태계에서 SGLang의 활용성을 높이는 중요한 단계입니다.
철저한 테스트 및 검증: PR 설명에 따르면, qwen-3-30B 모델을 사용한 결정론적 테스트 (--test-mode single, --test-mode prefix)를 통과했습니다. 또한, CI 상태 보고서 (amd-bot의 댓글)는 변경된 코드가 test/registered/attention/test_gemma4_swa_triton_oob_regression.py와 같은 실제 시나리오 테스트에서 실행되었고 통과했음을 명확히 밝히고 있습니다. 이는 이 개선이 실제 환경에서 효과적으로 작동함을 입증합니다.

일반적 교훈

이 PR은 다음과 같은 중요한 교훈을 제공합니다:

성능과 정확성의 균형: 고성능 시스템을 개발할 때는 항상 성능 최적화와 정확성(특히 결정론적 동작) 사이의 균형을 신중하게 고려해야 합니다. 때로는 미묘한 정확성 문제가 최적화된 커널에서 발생할 수 있으며, 이를 해결하기 위한 명확한 전략이 필요합니다.
하드웨어별 특성 이해: CUDA와 ROCm과 같은 다른 하드웨어 플랫폼은 동일한 연산에 대해서도 미묘하게 다른 동작 특성을 가질 수 있습니다. 이러한 플랫폼별 특성을 이해하고 적절히 대응하는 것이 크로스 플랫폼 호환성을 보장하는 데 중요합니다.
기능 플래그의 활용: is_batch_invariant_mode_enabled()와 같은 기능 플래그를 사용하여 특정 동작 모드(예: 결정론적 모드)에서만 다른 코드 경로를 활성화하는 것은 코드의 유연성과 유지보수성을 높이는 효과적인 방법입니다.

결론

SGLang의 이번 PR은 ROCm 환경에서 RMSNorm의 배치 불변성 문제를 해결하여 결정론적 추론의 신뢰성을 크게 향상시켰습니다. 이는 LLM 추론 프레임워크가 단순히 빠르기만 한 것이 아니라, 예측 가능하고 신뢰할 수 있어야 한다는 중요한 원칙을 다시 한번 상기시켜 줍니다. 이러한 세심한 개선들이 모여 SGLang과 같은 프레임워크가 더욱 견고하고 신뢰할 수 있는 LLM 추론 솔루션으로 발전하는 기반이 됩니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [vllm] vLLM XPU 백엔드 최적화: W8A8 및 W8A16 FP8 Linear 커널 도입
현재글 : [sglang] SGLang: ROCm 환경에서 RMSNorm의 배치 불변성 확보를 통한 결정론적 추론 개선
다음글 [uv] uv의 패키지 설치 성능 최적화: 불필요한 site-packages 스캔 제거

댓글

관련 포스트

PR Analysis 의 다른글