[sglang] HunyuanVideo VAE 디코딩 성능 향상: GroupNorm SiLU 커널 최적화

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

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

들어가며

최근 AI 모델, 특히 이미지 및 비디오 생성 모델의 발전 속도는 눈부십니다. 이러한 모델들은 방대한 양의 데이터를 처리하고 복잡한 연산을 수행해야 하므로, 효율적인 코드 구현이 필수적입니다. HunyuanVideo VAE 디코더는 이러한 복잡한 모델 중 하나로, 특히 GroupNorm과 SiLU 활성화 함수를 결합한 연산(GroupNorm SiLU)에서 성능 병목 현상이 발생했습니다. 본 PR은 이 병목 현상을 해결하기 위해 Triton을 사용하여 GroupNorm SiLU 연산의 특정 부분을 최적화했습니다. 이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 왜 이러한 최적화가 성능 향상에 크게 기여하는지, 그리고 그 일반적인 교훈은 무엇인지 살펴보겠습니다.

코드 분석

이번 PR의 핵심은 python/sglang/jit_kernel/diffusion/triton/group_norm_silu.py 파일에서 triton_group_norm_silu 함수의 적용(apply) 단계를 최적화하는 데 있습니다. 특히, 대규모 입력(large-shape)에 대한 청크(chunked) 적용 단계를 개선했습니다.

1. 새로운 Scalar-Affine Apply Kernel 추가 (`_group_norm_apply_scalar_affine_kernel`)

기존의 _group_norm_apply_kernel은 모든 경우에 대해 일반적인 방식으로 GroupNorm과 SiLU를 적용했습니다. 하지만 HunyuanVideo VAE 디코더와 같이 특정 크기의 입력에서는 이 방식이 비효율적일 수 있습니다. PR에서는 이러한 특정 경우를 위해 _group_norm_apply_scalar_affine_kernel이라는 새로운 커널을 도입했습니다.

기존 커널 (_group_norm_apply_kernel)의 일부:

def _group_norm_apply_kernel(
    input_ptr,
    weight_ptr,
    bias_ptr,
    output_ptr,
    stats_ptr,
    channels,
    spatial_size,
    num_groups,
    channels_per_group,
    group_size,
    chunks_per_row,
    BLOCK_SIZE: tl.constexpr,
    BLOCKS_PER_PROGRAM: tl.constexpr,
):
    # ... (중략) ...
    for block_id in range(BLOCKS_PER_PROGRAM):
        idx = chunk_start + block_id * BLOCK_SIZE + offsets
        mask = idx < group_size
        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
        # GroupNorm 적용 (mean, rstd 사용)
        y = (x - mean) * rstd
        # Affine 변환 (weight, bias)
        y = y * weight + bias
        # SiLU 활성화 함수 적용
        y = y * tl.sigmoid(y)
        tl.store(output_ptr + group_base + idx, y, mask=mask)

이 커널은 각 스레드 블록 내에서 idx // spatial_size와 같은 연산을 통해 채널 인덱스를 계산하고, 벡터 로드(vector load)를 수행합니다. 대규모 입력에서는 이 과정이 반복되면서 오버헤드가 발생할 수 있습니다.

새로운 커널 (_group_norm_apply_scalar_affine_kernel)의 일부:

@triton.jit
def _group_norm_apply_scalar_affine_kernel(
    input_ptr,
    weight_ptr,
    bias_ptr,
    output_ptr,
    stats_ptr,
    channels,
    spatial_size,
    num_groups,
    channels_per_group,
    group_size,
    chunks_per_row,
    BLOCK_SIZE: tl.constexpr,
    BLOCKS_PER_PROGRAM: tl.constexpr,
):
    row = tl.program_id(0).to(tl.int64)
    chunk_id = tl.program_id(1).to(tl.int64)

    batch_id = row // num_groups
    group_id = row - batch_id * num_groups
    chunk_start = chunk_id * BLOCK_SIZE * BLOCKS_PER_PROGRAM
    group_base = batch_id * channels * spatial_size + group_id * group_size

    # 각 프로그램이 단일 채널을 커버하도록 변경
    channel_id = chunk_start // spatial_size
    affine_offset = group_id * channels_per_group + channel_id
    weight = tl.load(weight_ptr + affine_offset).to(tl.float32)
    bias = tl.load(bias_ptr + affine_offset).to(tl.float32)

    mean = tl.load(stats_ptr + row * 2)
    rstd = tl.load(stats_ptr + row * 2 + 1)
    offsets = tl.arange(0, BLOCK_SIZE)

    for block_id in range(BLOCKS_PER_PROGRAM):
        idx = chunk_start + block_id * BLOCK_SIZE + offsets
        mask = idx < group_size
        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
        y = (x - mean) * rstd
        y = y * weight + bias
        y = y * tl.sigmoid(y)
        tl.store(output_ptr + group_base + idx, y, mask=mask)

이 새로운 커널의 핵심 아이디어는 다음과 같습니다:

Scalar-Affine Apply: 각 프로그램(program)이 단일 채널(channel)을 담당하도록 하여, idx // spatial_size와 같은 불필요한 연산을 제거합니다. 이는 각 채널에 대한 weight와 bias 로드를 한 번만 수행하게 하여 효율성을 높입니다.
Vector Affine Loads 회피: 기존의 벡터 로드 방식 대신, 각 채널별로 필요한 weight와 bias를 직접 로드하고 적용합니다. 이는 메모리 접근 패턴을 더 예측 가능하게 만들고, 특정 하드웨어에서 더 효율적일 수 있습니다.
조건부 적용: 이 최적화된 커널은 입력 텐서의 spatial_size가 _CHUNK_SIZE로 나누어 떨어지고, chunks_per_row가 충분히 클 때만 사용됩니다. 그렇지 않은 경우에는 기존의 범용 커널(_group_norm_apply_kernel)이 fallback으로 사용됩니다. 이는 최적화가 모든 경우에 적용되는 것이 아니라, 가장 효과적인 상황에만 적용되도록 보장합니다.

2. 기존 커널 유지 및 Fallback

PR은 새로운 커널만을 사용하는 것이 아니라, 기존의 _group_norm_apply_kernel을 그대로 유지했습니다. 이는 새로운 커널이 적용되지 않는 작은 크기의 입력이나 정렬되지 않은(unaligned) 청크의 경우에도 올바르게 동작하도록 보장합니다. 이러한 fallback 메커니즘은 코드의 견고성을 높여줍니다.

3. 벤치마킹 및 프로파일링 도구 추가

성능 개선을 측정하고 검증하기 위해, python/sglang/jit_kernel/benchmark/diffusion/bench_group_norm_silu.py 파일이 추가되었습니다. 이 스크립트는 다양한 케이스와 데이터 타입에 대해 기존 구현과 최적화된 구현의 성능을 비교하고, 결과를 CSV 파일로 저장하거나 콘솔에 출력할 수 있습니다. 또한, 프로파일링 모드를 지원하여 특정 연산의 성능을 더 깊이 분석할 수 있도록 돕습니다.

왜 이게 좋은가?

H200에서의 극적인 성능 향상

H200 GPU에서 HunyuanVideo VAE의 대규모 디코딩 시나리오에 대한 벤치마크 결과는 매우 인상적입니다.

Case	Shape	Branch	Native us	Fused us	Speedup
HunyuanVideo VAE large	`1x128x20x256x256`	`main`	4716.98	295.01	15.989x
HunyuanVideo VAE large	`1x128x20x256x256`	`PR`	4714.42	259.66	18.156x

PR이 적용된 버전(PR)은 기존 main 브랜치 대비 약 18배의 속도 향상을 보여주었습니다. 이는 단순히 GroupNorm SiLU 연산의 apply 단계만 최적화했음에도 불구하고 전체 VAE 디코딩 성능에 상당한 영향을 미칠 수 있음을 시사합니다.

NCU(NVIDIA Nsight Compute) 프로파일링 결과 또한 이러한 성능 향상의 원인을 명확히 보여줍니다:

Apply Stage 시간 감소: 255.17 us에서 178.08 us로 감소했습니다.
레지스터 사용량 감소: 48개에서 43개로 감소했습니다.
DRAM 처리량 증가: 51.86%에서 74.13%로 증가했습니다.

이러한 지표들은 새로운 _group_norm_apply_scalar_affine_kernel이 메모리 접근을 최적화하고 연산 효율성을 높여 GPU 자원을 더 효과적으로 사용하고 있음을 나타냅니다.

일반적인 교훈

특정 워크로드에 대한 커널 최적화의 중요성: 범용 커널은 많은 경우에 잘 작동하지만, 특정 하드웨어 및 입력 크기 조합에서는 비효율적일 수 있습니다. HunyuanVideo VAE와 같이 성능이 중요한 애플리케이션에서는 특정 패턴을 식별하고 이를 위한 맞춤형 커널을 개발하는 것이 큰 성능 향상을 가져올 수 있습니다.
조건부 실행 및 Fallback 전략: 최적화된 커널을 모든 경우에 적용하려 하기보다는, 성능 향상이 확실한 경우에만 사용하고 그렇지 않은 경우에는 기존의 안정적인 코드를 사용하는 fallback 전략은 코드의 견고성과 유지보수성을 높입니다.
정확한 프로파일링 및 벤치마킹: 성능 개선의 효과를 정량적으로 입증하고 병목 지점을 정확히 파악하기 위해서는 체계적인 벤치마킹 및 프로파일링 도구가 필수적입니다. 이 PR에 추가된 벤치마크 스크립트가 좋은 예시입니다.
메모리 접근 패턴 최적화: GPU 성능은 종종 메모리 대역폭과 접근 패턴에 의해 좌우됩니다. 새로운 커널이 DRAM 처리량을 높인 것은 메모리 접근을 더 효율적으로 설계했기 때문이며, 이는 많은 GPU 최적화의 핵심입니다.

리뷰 댓글 분석

제공된 리뷰 댓글은 주로 CI(Continuous Integration) 관련 설정 및 실행 결과에 대한 내용이었습니다. 예를 들어, [BBuf] /tag-and-rerun-ci와 같은 댓글은 CI 파이프라인을 재실행하거나 특정 태그를 지정하는 등의 운영적인 측면에 초점을 맞추고 있습니다. 코드 자체의 기술적인 깊이나 논쟁보다는, PR이 CI 환경에서 올바르게 빌드되고 테스트되는지를 확인하는 데 중점을 둔 것으로 보입니다. 이는 대규모 프로젝트에서 코드 변경 사항이 안정적으로 통합되기 위한 필수적인 절차입니다.

결론

이번 PR은 HunyuanVideo VAE 디코더의 성능을 크게 향상시키는 중요한 최적화를 수행했습니다. Triton을 활용하여 GroupNorm SiLU 연산의 apply 단계를 위한 새로운, 더 효율적인 커널을 도입함으로써, 특히 대규모 입력에서 극적인 속도 향상을 달성했습니다. 이는 AI 모델의 성능 최적화가 단순히 알고리즘 개선뿐만 아니라, 저수준 커널 구현의 최적화를 통해서도 이루어질 수 있음을 보여주는 훌륭한 사례입니다. 이러한 최적화 기법은 다른 딥러닝 모델 및 프레임워크에서도 널리 적용될 수 있을 것입니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [sglang] SGLang NPU 성능 최적화: INT8 TP 통신 압축 도입
현재글 : [sglang] HunyuanVideo VAE 디코딩 성능 향상: GroupNorm SiLU 커널 최적화
다음글 [openclaw] OpenClaw: 런타임 플러그인 레지스트리 재사용을 통한 성능 최적화

댓글

관련 포스트

PR Analysis 의 다른글