본문으로 건너뛰기

[sglang] ERNIE-Image의 RoPE와 GELU-mul 융합 및 RoPE cos/sin 호이스팅을 통한 성능 최적화

PR 링크: sgl-project/sglang#34306 상태: Merged | 변경: +371 / -13

들어가며

ERNIE-Image 모델은 이미지 생성 분야에서 뛰어난 성능을 보여주지만, 기존 구현에서는 특히 H100과 같은 최신 GPU에서 torch.compile 대비 느린 속도를 보였습니다. H200에서는 오히려 torch.compile이 성능 저하를 일으키는 문제도 있었습니다. 이러한 성능 병목 현상의 주된 원인은 RoPE(Rotary Positional Embedding)와 GELU(Gaussian Error Linear Unit) 연산이 포함된 복잡한 연쇄적인 연산(elementwise soup)에 있었습니다. 이 PR은 이러한 연산들을 효과적으로 융합하고 불필요한 재계산을 제거함으로써 ERNIE-Image의 추론 속도를 크게 향상시키는 것을 목표로 합니다.

코드 분석

이번 PR의 핵심은 세 가지 주요 최적화 포인트에 집중되어 있습니다. 각 변경 사항은 sglang 라이브러리의 커널 레벨에서 이루어졌으며, 기존의 연산과 동일한 결과를 보장하면서도 성능을 개선합니다.

1. RoPE cos/sin 재계산 제거 및 융합

기존 ERNIE-Image 구현에서는 RoPE의 cos와 sin 값이 각 레이어의 각 프로젝션마다 불필요하게 재계산되었습니다. 하지만 이 값들은 모델의 순방향(forward) 계산 동안 변하지 않는 상수입니다. 이 PR은 cos와 sin 값을 전체 순방향 계산 동안 단 한 번만 계산하도록 변경했습니다.

Before:

# 기존 구현 (개념적 예시, 실제 diff와는 다를 수 있음)
freqs = ... # (S, B, 1, rot_dim)
cos_ = torch.cos(freqs).to(dtype)
sin_ = torch.sin(freqs).to(dtype)
# ... 각 레이어, 각 프로젝션마다 반복

After:

_precompute_rope_cos_sin 함수가 도입되어 freqs로부터 cos와 sin을 한 번만 계산합니다. 이 값들은 (B * S, rot_dim) 형태로 미리 계산되어 재사용됩니다.

# sglang/multimodal_gen/runtime/models/dits/ernie_image.py
def _precompute_rope_cos_sin(
    freqs: torch.Tensor, dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]:
    """cos/sin of the rotary embedding, computed once per forward.

    ``freqs`` is the ``(S, B, 1, rot_dim)`` output of :class:`EmbedND3`; the
    eager chain recomputed ``torch.cos(freqs).to(dtype)`` per layer per
    projection.  Returns bit-identical ``(B * S, rot_dim)`` rows.
    """
    freqs = freqs.permute(1, 0, 2, 3)
    cos_ = torch.cos(freqs).to(dtype)
    sin_ = torch.sin(freqs).to(dtype)
    rot_dim = freqs.shape[-1]
    return cos_.reshape(-1, rot_dim), sin_.reshape(-1, rot_dim)

# ... 이후 forward 함수에서 이 함수를 호출하여 cos/sin을 미리 계산

2. rotate-half 연산의 Triton 커널 융합

기존의 rotate-half 연산은 여러 개의 커널(chunk, neg, cat, 두 번의 mul, add, tail cat 등)로 구성되어 있었습니다. 이는 각 프로젝션마다 약 7개의 커널을 필요로 했습니다. 이 PR은 이 복잡한 연쇄를 하나의 Triton 커널(rope_rotate_half_bitexact.py)로 통합했습니다.

이 Triton 커널은 bf16 연산의 반올림(rounding) 경계까지 정확하게 재현하도록 설계되었습니다. 즉, round(round(x1·cos1) + round(-x2·sin1))와 같은 연산을 각 요소별로 정확히 수행합니다. 이를 통해 기존의 여러 커널 호출과 임시 텐서 생성 오버헤드를 제거했습니다.

Before (개념적):

# ... cos/sin 계산 후
x1, x2 = x_rot.chunk(2, dim=-1)
x_rotated = torch.cat((-x2, x1), dim=-1) # neg + cat
x_rot = x_rot * cos_ + x_rotated * sin_ # two muls + add
# ... tail 처리 및 최종 cat

After:

fused_rope_rotate_half_bitexact 함수가 새로운 Triton 커널을 호출합니다.

# sglang/kernels/ops/diffusion/triton/rope_rotate_half_bitexact.py
# ... Triton 커널 정의 (_rope_rotate_half_kernel)

@register_custom_op(
    op_name="triton_fused_rope_rotate_half_bitexact",
    mutates_args=[],
    fake_impl=_fake_rope_rotate_half,
)
def fused_rope_rotate_half_bitexact(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> torch.Tensor:
    """Rotate-half RoPE over the leading ``cos.shape[-1]`` columns of ``x``.

    ``x`` is ``(B, S, H, D)``; ``cos``/``sin`` are ``(B * S, rot_dim)`` rows.
    Bit-exact vs the eager chunk/neg/cat/mul/add chain.
    """
    # ... Triton 커널 호출 로직
    with torch.cuda.device(x.device):
        _rope_rotate_half_kernel[...] 
    return out

# sglang/multimodal_gen/runtime/models/dits/ernie_image.py
def _ernie_rope(
    x: torch.Tensor, cos_: torch.Tensor, sin_: torch.Tensor
) -> torch.Tensor:
    # ... can_use_fused_rope_rotate_half 체크 후
    out = fused_rope_rotate_half_bitexact(x, cos_, sin_)
    # ...

3. GEGLU 연산 융합

ERNIE-Image의 Feed-Forward Network(FFN)에서 사용되는 GEGLU 활성화 함수는 up * F.gelu(gate) 형태로 구현되어 있었습니다. 이 연산은 FFN 중간 결과 텐서에 대해 두 번의 전체 통과(pass)를 수행했습니다. 이 PR은 이 연산을 gelu_and_mul_with_activation_rounding 함수를 통해 단일 연산으로 융합했습니다.

이 새로운 함수는 기존의 두 단계 연산(GELU 계산 후 곱셈)과 동일한 bf16 반올림 동작을 정확히 모방하여 비트 단위 정확성을 유지합니다.

Before:

# sglang/multimodal_gen/runtime/models/dits/ernie_image.py (기존)
gate, up = self.gate_up_proj(x)
x = up * F.gelu(gate)

After:

gelu_and_mul_with_activation_rounding 함수를 사용합니다.

# sglang/kernels/ops/activation/activation.py (신규)
def gelu_and_mul_with_activation_rounding(
    input: torch.Tensor,
    out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    hidden_size = input.shape[-1] // 2
    if out is None:
        out = input.new_empty(*input.shape[:-1], hidden_size)
    _run_activation_with_rounding_inplace("gelu", input, out) # 내부적으로 GELU와 곱셈 수행
    return out

# sglang/multimodal_gen/runtime/models/dits/ernie_image.py (수정)
def forward(self, x: torch.Tensor) -> torch.Tensor:
    gate_up, _ = self.gate_up_proj(x)
    x = _ERNIE_GEGLU(gate_up) # 융합된 연산 사용
    x, _ = self.linear_fc2(x)
    return x

_ERNIE_GEGLU는 BitExactFusionGate를 사용하여 융합된 연산의 정확성을 검증하고, 문제가 발생하면 기존 연산으로 fallback합니다.

왜 이게 좋은가?

이 PR의 최적화는 여러 측면에서 뛰어납니다.

  1. 성능 향상:

    • H100 카드에서 ERNIE-Image의 denoise 단계 추론 시간이 16.2% 감소했습니다 (15.789s -> 13.226s).
    • H200 카드에서는 12.7% 감소했습니다 (15.184s -> 13.249s).
    • 이는 기존의 느린 eager 모드보다 빨라졌으며, H200에서는 torch.compile의 회귀 문제까지 해결했습니다.

  2. 비트 단위 정확성 보장:

    • 모든 최적화는 기존 연산의 bf16 반올림 경계까지 정확하게 재현하도록 설계되었습니다. 이는 BitExactFusionGate를 통해 첫 실행 시 검증되며, 일치하지 않으면 자동으로 기존의 eager 연산으로 fallback합니다.
    • 실제 실행 결과, H100과 H200 모두에서 0번의 fallback이 발생했습니다. 이는 최적화된 경로가 기존 경로와 비트 단위로 동일함을 의미합니다.

  3. 코드 간결성 및 효율성:

    • 불필요한 cos/sin 재계산을 제거하여 연산량을 줄였습니다.
    • 여러 개의 작은 커널 호출을 하나의 Triton 커널로 통합하여 커널 실행 오버헤드를 줄이고 GPU 활용도를 높였습니다.
    • GELU와 곱셈을 융합하여 연산 단계를 줄였습니다.

  4. 일반적인 교훈:

    • 연산 융합(Operator Fusion): 특히 elementwise 연산이 많은 경우, 여러 연산을 하나의 커널로 융합하는 것은 성능 향상의 큰 기회가 됩니다. 이는 커널 호출 오버헤드를 줄이고 데이터 재사용성을 높입니다.
    • 불필요한 재계산 제거: 모델의 순방향 계산에서 변하지 않는 값(예: RoPE의 cos/sin)은 미리 계산하여 재사용해야 합니다. 이는 단순하지만 효과적인 최적화 기법입니다.
    • 정확성 검증의 중요성: 성능 최적화 시, 특히 부동 소수점 연산에서는 비트 단위 정확성을 보장하는 것이 중요합니다. BitExactFusionGate와 같은 메커니즘은 이러한 검증을 자동화하여 안전한 최적화를 가능하게 합니다.
    • Triton 활용: 복잡한 연산이나 여러 연산의 조합을 최적화할 때 Triton과 같은 DSL은 GPU 커널을 효율적으로 작성하고 최적화하는 강력한 도구가 될 수 있습니다.

리뷰 피드백 반영

PR 설명에 따르면, 이 PR은 BitExactFusionGate를 사용하여 첫 실행 시 기존 eager 연산과 융합된 연산의 결과를 torch.equal로 비교합니다. 만약 모양(shape), 데이터 타입(dtype), 또는 플랫폼(platform)이 지원되지 않거나 결과가 일치하지 않으면, 해당 융합은 영구적으로 비활성화되고 eager fallback 경로를 사용하게 됩니다. 이는 리뷰어들이 우려할 수 있는 정확성 문제를 효과적으로 해결하며, 최적화된 경로가 항상 안전하게 사용될 수 있도록 보장합니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글