[sglang] Cosmos3 T2I 가속화: QKNorm + RoPE 커널 퓨전과 BF16 정밀도 최적화

PR 링크: sgl-project/sglang#34932 상태: Merged | 변경: +135 / -2

들어가며

최신 Diffusion 모델인 Cosmos3는 비디오와 이미지(T2I) 생성을 모두 지원합니다. 기존에 Cosmos3 비디오 경로는 QKNorm, RoPE, KV-packing 연산이 하나로 합쳐진(fused) 커널을 사용하여 효율적으로 동작하고 있었습니다. 하지만 이미지 생성 경로($T=1$)는 여전히 이 연산들을 개별적으로 수행하는 'Eager split path'를 사용하고 있었습니다.

그 이유는 이미지 경로가 RoPE 연산 직전에 BF16 라운딩(Rounding) 단계를 거쳐야 했기 때문입니다. 기존의 퓨전 커널은 이 중간 라운딩 과정을 지원하지 않아, 비트 단위의 정확도(Bit-exactness)를 유지하기 위해 성능을 희생하고 개별 연산을 호출해왔습니다.

이번 PR은 퓨전 커널에 라운딩 모드를 추가하여 Cosmos3 T2I 경로에서도 가속화된 커널을 사용할 수 있게 하고, 반복적인 캐스팅 오버헤드를 제거한 최적화 사례입니다. 결과적으로 NVIDIA B300 환경에서 torch.compile 대비 약 13.4%의 성능 향상을 달성했습니다.

---

코드 분석: 핵심 변경 사항

1. 퓨전 커널 인터페이스 확장 (cosmos3video.py)

가장 먼저, 퓨전 커널을 호출하는 함수들에 round_norm_before_rope 옵션을 추가했습니다. 이를 통해 기존의 퓨전 로직을 유지하면서도 이미지 경로에 필요한 라운딩 처리를 선택적으로 활성화할 수 있게 되었습니다.

[Before]

def _apply_qwen3_qk_norm_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    q_norm: torch.nn.Module,
    k_norm: torch.nn.Module,
    head_dim: int,
    cos_sin_cache: torch.Tensor,
    rope_cache_positions: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    return apply_qk_norm_rope(
        q=q, k=k, q_norm=q_norm, k_norm=k_norm, ...
    )

[After]

def _apply_qwen3_qk_norm_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    q_norm: torch.nn.Module,
    k_norm: torch.nn.Module,
    head_dim: int,
    cos_sin_cache: torch.Tensor,
    rope_cache_positions: torch.Tensor,
    *,
    round_norm_before_rope: bool = False, # 라운딩 옵션 추가
) -> tuple[torch.Tensor, torch.Tensor]:
    return apply_qk_norm_rope(
        ...,
        round_norm_before_rope=round_norm_before_rope,
    )

2. RoPE 캐시 캐스팅 최적화

기존 코드에서는 매 레이어(GEN layer)마다 RoPE 캐시를 활성화 함수 데이터 타입(Activation dtype)으로 캐스팅하는 과정이 반복되었습니다. 이번 변경에서는 캐시를 생성할 때 한 번만 캐스팅하여 저장하도록 수정되었습니다.

[Before]

# 매 forward마다 build_rope_cache_inputs 호출 및 내부 캐스팅 발생 가능성
self.cached_gen_rope_inputs[cache_key] = (
    self.language_model.rotary_emb.build_rope_cache_inputs(
        vis_pos_ids, cache_dtype=hidden_gen.dtype
    )
)

[After]

cos_sin_gen, gen_rope_cache_positions = (
    self.language_model.rotary_emb.build_rope_cache_inputs(
        vis_pos_ids, cache_dtype=hidden_gen.dtype
    )
)
if T == 1 and not self._gen_layers_torch_compiled:
    # 퓨전 커널이 매번 캐스팅하지 않도록 미리 activation dtype으로 변환
    cos_sin_gen = cos_sin_gen.to(hidden_gen.dtype)

self.cached_gen_rope_inputs[cache_key] = (
    cos_sin_gen,
    gen_rope_cache_positions,
)

3. 퓨전 커널 활성화 조건 로직

단순히 커널을 합치는 것뿐만 아니라, torch.compile이 적용된 상태인지, 장치가 CUDA인지 등을 꼼꼼하게 체크하여 안전하게 퓨전 커널을 적용합니다.

[After]

round_norm_before_rope = T == 1
use_fused_qk_norm_rope = T > 1 or (
    hidden_gen.device.type == "cuda"
    and not torch.compiler.is_compiling() # 컴파일 중이 아닐 때만
    and not self._gen_layers_torch_compiled
    and get_sp_world_size() == 1
    and can_use_fused_inplace_qknorm_rope(
        ...,
        round_norm_before_rope=True,
        pack_kv=True,
    )
)

---

왜 이게 좋은 최적화인가?

1. 메모리 대역폭(Memory Bandwidth) 절약

QKNorm, RoPE, KV-packing은 모두 연산량 대비 메모리 접근이 많은 Memory-bound 연산들입니다. 개별 커널로 실행하면 각 단계마다 중간 결과를 VRAM에 썼다가 다시 읽어와야 합니다. 이를 하나의 커널로 퓨전하면 데이터를 GPU 레지스터나 Shared Memory에 둔 상태로 연속 처리가 가능하므로 메모리 대역폭 오버헤드를 극적으로 줄일 수 있습니다.

2. 비트 단위 정확도(Bit-exactness) 보장

딥러닝 모델 가속화에서 가장 까다로운 점은 '가속화 후 결과값이 달라지는 것'입니다. 특히 BF16은 정밀도가 낮아 연산 순서나 중간 라운딩 여부에 따라 결과가 미세하게 달라질 수 있습니다. 이 PR은 round_norm_before_rope 플래그를 통해 기존 Eager 모드의 BF16 라운딩 경계(Rounding boundary)를 완벽히 재현함으로써, 성능 향상과 결과의 동일성을 동시에 확보했습니다.

3. Redundant Cast 제거

모델의 레이어가 깊을수록 매 레이어마다 발생하는 작은 오버헤드가 누적되어 큰 성능 저하를 일으킵니다. RoPE 캐시를 미리 타겟 dtype으로 변환해두는 사소해 보이는 변경이 전체 추론 시간을 단축하는 데 기여합니다.

성능 결과

NVIDIA B300 1개 기준, cosmos3-nano-t2i 프리셋 측정 결과:

• Lossless Eager (이번 최적화 적용): 0.863 s
• torch.compile: 0.997 s

일반적으로 torch.compile이 가장 빠를 것이라 기대하지만, 특정 도메인에 최적화된 수동 퓨전 커널(Hand-written fused kernel)이 컴파일러의 자동 최적화보다 더 나은 성능을 보일 수 있음을 증명했습니다.

---

결론

이번 최적화는 단순히 "빠른 커널을 썼다"는 점을 넘어, 실제 서비스 환경에서 요구되는 수치적 정확도(Lossless)를 유지하면서도 어떻게 아키텍처 특성을 활용해 가속화할 것인가에 대한 좋은 답안을 제시합니다. 특히 torch.compile과의 상호작용을 고려하여 _gen_layers_torch_compiled 플래그로 안정성을 확보한 점은 시니어 엔지니어다운 세심한 설계라고 볼 수 있습니다.

이러한 접근 방식은 다른 Diffusion 모델이나 LLM의 커스텀 연산 최적화에도 동일하게 적용될 수 있는 중요한 패턴입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://pytorch.org/docs/stable/notes/cuda.html#tf32-on-ampere

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