[sglang] SANA-Video 2.0 anchor attention의 FP32 upcast 제거로 720p 생성 52% 단축하기

PR 링크: sgl-project/sglang#49219 상태: Merged | 변경: +211 / -9

들어가며

SANA-Video 2.0은 release config의 fp32_attention: true 설정을 따라 softmax anchor layer의 Q/K/V를 scaled_dot_product_attention(SDPA) 호출 전에 FP32로 upcast합니다. 겉보기엔 "정밀도를 높이는 안전한 선택"이지만, 실제로는 두 가지 문제가 있었습니다.

1. 정보 이득이 없다. Q/K/V는 RMSNorm과 RoPE를 거치면서 이미 activation dtype(BF16)으로 cast된 값입니다. BF16 값을 FP32로 바꿔도 새로운 정보는 생기지 않습니다.
2. 커널 선택을 망가뜨린다. PyTorch의 flash attention backend는 FP32 입력을 받지 않습니다. 결과적으로 SDPA는 FP32 memory-efficient 커널(fmha_cutlassF_f32)로 fallback합니다.

그 대가는 컸습니다. 720p/193-frame 요청에서 anchor layer 하나가 forward당 fmha_cutlassF_f32에서 529 ms를 쓰고, 8개의 anchor layer가 denoising 시간의 약 절반을 차지했습니다.

이 PR은 quality=high 요청에서만 BF16 Q/K/V를 그대로 SDPA에 넘기는 request-gated site를 추가합니다. 기본값인 quality=exact는 main과 bit-identical하게 유지됩니다.

코드 분석

1. sites/sana_video2_softmax_attention_site.py (신규)

핵심은 QualityGatedFusion 인스턴스와, 조건을 만족할 때만 결과를 반환하는 try_* 함수입니다.

SANA_VIDEO2_SOFTMAX_ATTENTION = QualityGatedFusion(
    name="SANA-Video 2.0 BF16-input softmax attention",
    marker_attr="_sgl_sana_video2_softmax_attention_site",
    enabled_attr="_sgl_sana_video2_softmax_attention_enabled",
    reject_reason=_site_reject_reason,
)


def try_sana_video2_softmax_attention(site, query, key, value):
    """Return the quality-gated attention result, or ``None`` to fall back."""
    if not (
        SANA_VIDEO2_SOFTMAX_ATTENTION.is_enabled(site)
        and query.is_cuda
        and query.dtype == torch.bfloat16
        and key.dtype == query.dtype
        and value.dtype == query.dtype
        and query.dim() == 4
        and query.shape == key.shape == value.shape
    ):
        return None
    return F.scaled_dot_product_attention(query, key, value, dropout_p=0.0)

주목할 점:

• 새 커널이 없습니다. 반환 경로는 그냥 F.scaled_dot_product_attention입니다. 바뀐 것은 입력 dtype뿐이고, 그것만으로 PyTorch dispatcher가 flash attention을 선택할 수 있게 됩니다.
• 방어적 guard. CUDA, BF16, 4D, 동일 shape일 때만 진입하고, 나머지는 None을 반환해 reference 경로로 fallback합니다. 이 Optional 반환 프로토콜 덕분에 호출부는 "시도 후 실패하면 원래대로"라는 단순한 구조를 유지합니다.
• 이 패턴은 기존 SANA-Video 2B의 SANA_VIDEO_LINEAR_ATTENTION과 동일해, 코드베이스 내 일관성을 지킵니다.

2. models/dits/sana_video2.py — GatedSoftmaxAttention

Before:

if self.fp32_attention:
    q, k, v = q.float(), k.float(), v.float()
out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)

After:

out = try_sana_video2_softmax_attention(self, q, k, v)
if out is None:
    if self.fp32_attention:
        q, k, v = q.float(), k.float(), v.float()
    out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)

그리고 생성자에서 자신을 site로 표시합니다.

self.output_gate = nn.Linear(dim, dim)
SANA_VIDEO2_SOFTMAX_ATTENTION.mark(self)

gated 경로를 FP32 upcast 이전에 시도한다는 순서가 핵심입니다. upcast 이후에 시도하면 이미 dtype이 FP32라 guard에서 탈락합니다.

3. Stage와 handler 등록

• _QUALITY_FUSION_HANDLERS에 이 site를 tier high로 추가했습니다.
• SanaVideo2DenoisingStage가 quality 변경 시 request-gated fusion을 mount/unmount하도록 수정했습니다. PR 설명에 따르면, 기존에는 이 model-specific stage가 어떤 site도 mount하지 않고 있었습니다. 공용 DenoisingStage나 Helios stage는 하던 일을 빠뜨렸던 것으로, 이 PR이 아니었다면 기존의 다른 fusion도 이 stage에선 적용되지 않았을 것입니다.

4. 문서

fused_kernels.mdx는 fusion을 "Same math, different rounding(lossless)"과 "Approximate(high)"로 구분합니다. 이 변경은 probability를 PV GEMM 전에 BF16으로 반올림하므로 operand 정밀도를 낮추는 Approximate 범주에 들어가며, 그래서 lossless가 아닌 high에서만 mount됩니다. 분류 기준을 문서에 명시해둔 것이 리뷰와 유지보수를 쉽게 만듭니다.

왜 이게 좋은가

성능 수치

profile	exact e2e	high e2e	변화
832x480, 81f T2V	35.91 s	30.89 s	-14.0%
1280x736, 193f T2V	364.07 s	172.37 s	-52.7%

layer 단위로는 720p에서 554.65 → 78.06 ms이며, attention 커널 자체는 fmha_cutlassF_f32 529.30 ms → flash_fwd_kernel 54.07 ms로 약 10배 빨라졌습니다. 8개 layer 절감분이 관측된 step당 절감(480p 100 ms, 720p 3.81 s)과 정확히 맞아떨어집니다. Peak reserved memory는 동일합니다.

시퀀스가 길수록 이득이 커지는 이유는 attention이 O(N²)이라 커널 효율 차이가 그대로 증폭되기 때문입니다.

정확도에 대한 정직한 보고

저자는 단점도 숨기지 않았습니다. 480p TI2V만 레포의 기본 video gate(SSIM ≥ 0.92, PSNR ≥ 24 dB)를 통과했고, 720p T2V는 최소 SSIM 0.8293까지 떨어집니다. 편차는 frame index와 clip 길이에 따라 커지며, 50 step denoising이 미세한 rounding 차이를 증폭하기 때문입니다. 다만 장면 구성, 카메라 움직임, 조명은 동일하고 차이는 edge·텍스처에 집중됩니다. 그래서 기본값이 아닌 opt-in(--quality high)으로 제공됩니다.

FP16 입력도 시도했지만 layer 단위 rel-L2는 2.2e-3 → 1.45e-3으로 줄었어도 end-to-end SSIM 개선은 미미했습니다(0.898 → 0.903). "drift는 probability 정밀도로 통제되지 않는다"는 결론으로 BF16을 유지한 것은 근거 기반의 좋은 판단입니다.

리뷰 과정

CI 실패가 반복되자 리뷰어(BBuf)가 원인 확인을 요청했고, 저자는 각 실패 job이 sdpa.py cuDNN priority 테스트, mem_cache 등 이 PR이 건드리지 않은 영역이며 다른 PR에서도 동일하게 실패함을 하나씩 짚었습니다. 리뷰어도 coverage 우려를 철회했습니다.

일반적 교훈

1. "안전한 upcast"가 커널 선택을 막을 수 있다. SDPA처럼 dtype에 따라 backend를 dispatch하는 API에서는 입력 dtype 하나가 10배 성능 차이를 만듭니다. upcast 전에 값이 이미 저정밀인지 확인하세요.
2. 커널 내부 accumulation은 이미 FP32다. flash attention은 BF16 입력에서도 QK^T, softmax, PV를 FP32로 누적합니다. 입력 upcast는 중복 비용일 수 있습니다.
3. 근사 최적화는 gate 뒤에 두라. 기본 경로를 bit-identical하게 유지하고 품질 tier로 opt-in시키면, 성능과 재현성을 동시에 지킬 수 있습니다.
4. profiler로 커널 이름을 확인하라. fmha_cutlassF_f32라는 이름 하나가 문제의 전부를 말해줍니다.

참고 자료

• torch.nn.functional.scaled_dot_product_attention — 이 PR에서 BF16 입력으로 호출해 flash attention backend를 활성화한 SDPA 공식 문서
• torch.nn.attention.sdpa_kernel — SDPA backend(flash / memory-efficient) 선택 관련 공식 문서
• torch.Tensor.float — reference 경로에서 Q/K/V를 FP32로 upcast하는 메서드

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