[sglang] SGLang의 FLUX.2 추론 최적화: Eager 모드에서의 커널 융합 기법

PR 링크: sgl-project/sglang#34616 상태: Merged | 변경: +293 / -20

들어가며

최신 대규모 확산 모델인 FLUX.2는 뛰어난 생성 성능을 자랑하지만, 추론 과정에서 발생하는 빈번한 메모리 접근과 연산 오버헤드가 병목으로 작용합니다. 특히 Eager 모드에서 실행될 때, 개별적인 LayerNorm, Modulation, 그리고 SwiGLU 연산들은 GPU 메모리 대역폭을 낭비하게 됩니다. 본 PR은 이러한 연산들을 Triton 커널로 융합하여 메모리 접근을 최소화하고, 비트 단위로 정확한(bit-exact) 결과를 보장하는 최적화를 수행했습니다.

코드 분석

1. SwiGLU 연산의 융합 (silu_mul_bitexact.py)

기존에는 SwiGLU 연산을 위해 입력을 두 개로 슬라이싱하고 각각 연산한 뒤 결합하는 과정이 필요했습니다. 이를 _packed_silu_mul_kernel을 통해 하나의 커널에서 처리하도록 변경했습니다.

@triton.jit
def _packed_silu_mul_kernel(out_ptr, x_ptr, num_rows, row_stride, D: tl.constexpr, BLOCK: tl.constexpr):
    # ... (생략) ...
    a = tl.load(x_ptr + row_base + cols, mask=mask, other=0.0).to(tl.float32)
    b = tl.load(x_ptr + row_base + D + cols, mask=mask, other=0.0).to(tl.float32)
    s = round_bf16_to_fp32(a * tl.sigmoid(a))
    tl.store(out_ptr + row * D + cols, s * b, mask=mask)

이 방식은 데이터를 메모리에서 두 번 읽지 않고, 하나의 커널 내에서 a와 b를 로드하여 즉시 연산함으로써 캐시 효율을 극대화합니다.

2. AdaLN 및 Modulation 융합 (flux_2.py)

FLUX.2의 핵심인 AdaLN(Adaptive LayerNorm)과 Modulation 연산을 _flux2_norm_modulate 함수로 통합했습니다. 이 함수는 torch.compile 환경에서는 원본 표현을 유지하되, Eager 모드에서는 최적화된 Triton 커널을 호출합니다.

def _flux2_norm_modulate(norm, x, scale, shift):
    # ... (생략) ...
    if can_use_fused_layernorm_modulate(x, scale_row, shift_row):
        return fused_layernorm_modulate_raw(x, scale_row, shift_row, norm.eps)
    return norm(x) * (1 + scale) + shift

왜 이게 좋은가

이번 최적화의 핵심은 '비트 정확성(Bit-exactness)을 유지하면서 메모리 대역폭을 절약'하는 것입니다.

1. 성능 향상: B300 환경에서 FLUX.2 Klein Base 모델 기준, Eager 모드에서만 약 12.81%의 Denoise 속도 향상을 기록했습니다. torch.compile과 결합 시 성능 격차는 거의 사라집니다.
2. 안전한 폴백(Fallback): BitExactFusionGate를 사용하여 새로운 레이아웃이나 지원되지 않는 데이터 타입이 들어올 경우 자동으로 원본 연산으로 돌아가도록 설계되어 있어, 운영 환경에서의 안정성을 보장합니다.
3. CUDA Graph 호환성: 새로운 SwiGLU 시그니처가 감지될 경우, CUDA 그래프 캡처 중에 호스트 동기화가 발생하지 않도록 하여 런타임 성능 저하를 방지했습니다.

이러한 접근은 대규모 모델 추론 시 연산 효율을 높이는 표준적인 패턴으로, 메모리 바운드(memory-bound) 연산이 많은 모델에 매우 효과적입니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://pytorch.org/docs/stable/generated/torch.compile.html

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