본문으로 건너뛰기

[flashinfer] [FlashInfer] Softmax 루프 최적화: Register Spill 방지를 위한 Tail Masking 분리 전략

PR 링크: flashinfer-ai/flashinfer#5280 상태: Merged | 변경: +133 / -13

들어가며

고성능 LLM 추론 엔진인 FlashInfer의 prims_ts 커널에서 매우 흥미로운 성능 최적화 PR이 머지되었습니다. 이번 변경사항의 핵심은 Dense Context FMHA(Fused Multi-Head Attention) 연산 시, 시퀀스 길이가 타일 크기(128)의 배수가 아닐 때 발생하는 성능 저하를 해결하는 것입니다.

기존 구조에서는 시퀀스의 마지막 부분(Tail)을 처리하기 위한 마스킹 로직이 Softmax 루프 내부에 포함되어 있었습니다. 이로 인해 컴파일러가 레지스터 할당(Register Allocation)을 최적으로 수행하지 못하고, 대량의 Register Spills가 발생하여 L2 캐시 히트율이 떨어지고 연산 사이클이 증가하는 문제가 있었습니다. 이번 포스트에서는 이 문제를 어떻게 루프 구조 변경을 통해 해결했는지 심층 분석합니다.

문제의 핵심: 루프 내부의 조건부 마스킹과 Register Spilling

FlashInfer는 어텐션 연산을 타일 단위로 처리합니다. 만약 시퀀스 길이 $S$가 128의 배수가 아니라면(예: $S=147600$, $147600 \pmod{128} = 16$), 마지막 타일은 128개 중 16개만 유효한 데이터를 가집니다.

기존 코드에서는 이 '마지막 타일'을 처리하기 위해 Softmax 루프의 매 반복(iteration)마다 해당 타일이 마지막인지 확인하고 마스킹을 적용하는 masked row-max 변형을 실행했습니다.

문제 현상

  • Register Spilling: 마지막 타일 크기가 컴파일 타임 상수로 취급되면서, 각기 다른 Tail Size마다 서로 다른 바이너리(cubin)가 생성되었습니다. 일부 크기에서는 Softmax 워프에서 심각한 레지스터 부족이 발생하여 데이터를 메모리로 빼내는 Spill이 발생했습니다.
  • 성능 저하: $S=147600$ 기준, 타일 정렬이 된 경우보다 4.7배 많은 Spill Store가 발생했으며, 이로 인해 K/V 데이터를 L2 캐시에서 밀어내어(evict) 이상적인 $S^2$ 스케일링 대비 21% 더 많은 사이클을 소모했습니다.

코드 분석: 무엇이 바뀌었는가?

1. Softmax 도메인 정책 변경 (fmha_kernel.py)

가장 먼저, Softmax가 순회하는 타일의 범위를 조정했습니다. 마지막 타일을 루프 밖에서 처리하기 위해, 루프는 $N-1$개의 '완전한' 타일만 처리하도록 변경되었습니다.

Before:

# 모든 타일을 Softmax 도메인에 포함
return FmhaDomainPolicy(
    domain_n_kwargs=domain_n_kwargs,
    domain_n_minus_1_kwargs=domain_n_minus_1_kwargs,
    softmax0_domain_kwargs=domain_n_kwargs,
    softmax1_domain_kwargs=domain_n_kwargs,
)

After:

# Dense K Tail이 있는 경우, Softmax 루프는 N-1개만 처리
softmax_uses_dense_k_tail = (
    not cfg.single_qkv_instance
    and not cfg.is_causal
    and not cfg.has_varlen
    and cfg.fixed_dense_k_tail > 0
)
softmax_domain_kwargs = (
    domain_n_minus_1_kwargs if softmax_uses_dense_k_tail else domain_n_kwargs
)
return FmhaDomainPolicy(
    # ...
    softmax0_domain_kwargs=softmax_domain_kwargs,
    softmax1_domain_kwargs=softmax_domain_kwargs,
)

2. 루프 내부 마스킹 제거 및 Tail 섹션 추가 (fmha_tasks.py)

핵심 로직인 softmax_schedule 함수에서 루프 내부의 마스킹 로직을 제거하고, 루프가 끝난 뒤 명시적인 Tail 섹션에서 마지막 타일을 처리하도록 분리했습니다.

Before (Loop Inside):

# 루프 내부에서 매번 체크
elif tmem_sp.uses_fixed_dense_k_tail_mask:
    old_row_max, row_max = sp.fixed_dense_k_tail_masked_row_max(
        row_max=row_max,
    )

After (Loop Outside - Tail Section):

# 루프가 끝난 후 Tail 섹션에서 한 번만 실행
elif tmem_sp.uses_fixed_dense_k_tail_mask:
    # Dense tail: zero-filled lanes을 마스킹하며 한 단계 더 수행
    sp.wait()
    old_row_max, row_max = sp.fixed_dense_k_tail_masked_row_max(
        row_max=row_max,
        section=FmhaStage.Tail,
    )
    vec.store_vec(old_row_max=old_row_max, row_max=row_max, row_sum=row_sum)
    vec.commit()
    # ... 이후 Softmax 통계치 업데이트 로직

이 변경을 통해 루프 본체는 마스킹 조건문이 없는 'Pure'한 상태가 되어, 시퀀스 길이 $S$에 상관없이 동일한 레지스터 할당을 보장받게 되었습니다.

3. 리소스 관리 로직 업데이트 (fmha_resources.py)

fixed_dense_k_tail_masked_row_max 함수는 이제 루프 내부(Loop)와 루프 외부(Tail) 모두에서 호출될 수 있도록 확장되었습니다.

After:

def fixed_dense_k_tail_masked_row_max(
    self, stage_info: StageInfo, *, row_max: SoftmaxScalar,
    section: cutlass.Constexpr[FmhaStage] = FmhaStage.Loop,
) -> tuple[SoftmaxScalar, SoftmaxScalar]:
    # ...
    if cutlass.const_expr(section == FmhaStage.Tail):
        is_tail_tile = True
    else:
        is_tail_tile = stage_info.loop_offset == stage_info.loop_end - Int32(1)
    
    if is_tail_tile:
        # 유효하지 않은 영역을 -inf로 마스킹
        # ...

왜 이게 좋은 최적화인가?

1. Register Pressure 완화

CUDA 커널에서 루프 내부에 복잡한 조건문이나 컴파일 타임 상수에 의존하는 로직이 있으면, 컴파일러는 최악의 경우를 대비해 레지스터를 보수적으로 할당하거나 Spill을 발생시킵니다. 마스킹 로직을 루프 밖으로 빼냄으로써, 가장 많이 실행되는 'Hot Loop'의 레지스터 요구량을 최소화하고 최적의 파이프라이닝을 가능하게 했습니다.

2. 성능 수치 (B200 GPU 기준)

  • S=147600: 실행 사이클이 1283M에서 1037M으로 약 19.2% 감소했습니다.
  • Spill Stores: 764 GB에서 55 GB로 무려 92.8% 감소했습니다. 이는 메모리 대역폭 낭비를 획기적으로 줄였음을 의미합니다.
  • 실제 워크로드: Wan2.2-T2V 모델의 E2E 성능에서도 유의미한 개선이 확인되었습니다.

3. 유연성과 안정성 (Reviewer Feedback 반영)

리뷰 과정에서 saltyminty는 variable-window 모드와의 충돌 가능성을 지적했습니다. 이를 반영하여 uses_fixed_dense_k_tail_mask 조건에서 has_variable_window를 제외함으로써, 특정 기능 사용 시 발생할 수 있는 행(hang)이나 오동작을 방지했습니다.

# fmha_resources.py 수정 사항
def uses_fixed_dense_k_tail_mask(self) -> bool:
    return (
        not self.cfg.is_causal
        and not self.cfg.has_varlen
        and not self.cfg.has_variable_window # 리뷰 피드백 반영
        and self.cfg.fixed_dense_k_tail > 0
    )

결론 및 교훈

이번 최적화는 "루프 내부를 최대한 가볍고 균일하게 유지하라"는 GPU 프로그래밍의 고전적인 원칙을 다시 한번 증명했습니다. 특히 컴파일 타임 상수를 활용한 마스킹이 오히려 레지스터 할당에 독이 될 수 있다는 점은 시니어 엔지니어들에게 시사하는 바가 큽니다.

FlashInfer는 이처럼 하드웨어의 특성을 극한으로 활용하여, 정렬되지 않은 시퀀스 길이에 대해서도 정렬된 경우와 거의 동일한 성능을 낼 수 있도록 진화하고 있습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글