본문으로 건너뛰기

[flashinfer] FlashInfer GDN 커널: SM12x GPU에서 상삼각 행렬 MMA 타일 건너뛰기로 성능 최적화

PR 링크: flashinfer-ai/flashinfer#6371 상태: Merged | 변경: +96 / -2

들어가며

안녕하세요, 기술 블로거입니다. 오늘은 고성능 딥러닝 추론 라이브러리인 FlashInfer의 GDN(Grouped Delta Normalization) 커널에서 SM12x 아키텍처 GPU를 위한 흥미로운 성능 최적화 PR을 분석해보고자 합니다. 이 PR은 특히 NewV 및 O2 출력 계산 과정에서 발생하는 불필요한 행렬 곱셈(MMA, Multiply-Accumulate) 연산을 제거하여 상당한 성능 향상을 이끌어냈습니다.

기존 FlashInfer의 SM12x GDN non-CP(non-Context Parallel) prefill 경로는 NewV와 O2 출력 기여를 계산할 때, 우측 피연산자가 상삼각 행렬(upper triangular matrix)임에도 불구하고 밀집(dense) MMA 시퀀스를 발행했습니다. 상삼각 행렬은 대각선 아래의 모든 원소가 0인 특성을 가집니다. 이 PR은 이러한 특성을 활용하여 0으로만 구성된 MMA 타일을 건너뛰어 연산량을 줄이는 것을 목표로 합니다.

구체적으로, M16 행 타일당 64x64 우측 피연산자는 4x8 K16/N8 타일로 나뉘는데, 이 중 12개의 0 타일을 건너뛰어 M16xN8xK16 타일 MMA 연산 횟수를 32개에서 20개로 줄이는 최적화입니다. 이는 GPU 연산 시간을 크게 단축시킬 수 있는 중요한 개선입니다.

코드 분석: 불필요한 MMA 타일 건너뛰기

이번 최적화의 핵심은 flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_sm120.py 파일에 새로운 헬퍼 함수 _gemm_upper_triangular_b를 도입하고, 기존의 cute.gemm 호출을 이 함수로 대체하는 것입니다.

flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_sm120.py

이 파일은 SM120 아키텍처를 위한 GDN 델타 규칙 커널의 구현을 담고 있습니다. 주요 변경 사항은 compute_loop_body 함수 내에서 cute.gemm 호출이 _gemm_upper_triangular_b로 대체된 부분입니다.

Before:

        self._math_order_wait(wg_idx)
        cute.copy(newv_tiled_copy_B, tNewVsB, tNewVrB_cv)
        tNewVrC.fill(self.acc_dtype(0.0))
        cute.gemm(newv_tiled_mma, tNewVrC, tNewVrA, tNewVrB, tNewVrC)
        self._math_order_notify(wg_idx)
        v_pipeline.consumer_release(v_consumer_state)
        v_consumer_state.advance()

        # ── O2 output contribution ────────────────────────────────────────────
        tOrNewV = SM80.make_acc_into_op(tNewVrC, o2_tiled_mma, self.dtype)
        self._math_order_wait(wg_idx)
        cute.copy(o2_tiled_copy_B, tOsQK, tOrQK_cv)
        cute.gemm(o2_tiled_mma, tOrO, tOrNewV, tOrQK, tOrO)
        self._math_order_notify(wg_idx)

After:

        self._math_order_wait(wg_idx)
        cute.copy(newv_tiled_copy_B, tNewVsB, tNewVrB_cv)
        tNewVrC.fill(self.acc_dtype(0.0))
        self._gemm_upper_triangular_b(
            newv_tiled_mma, tNewVrA, tNewVrB, tNewVrC, sKK_opd.shape
        )
        self._math_order_notify(wg_idx)
        v_pipeline.consumer_release(v_consumer_state)
        v_consumer_state.advance()

        # ── O2 output contribution ────────────────────────────────────────────
        tOrNewV = SM80.make_acc_into_op(tNewVrC, o2_tiled_mma, self.dtype)
        self._math_order_wait(wg_idx)
        cute.copy(o2_tiled_copy_B, tOsQK, tOrQK_cv)
        self._gemm_upper_triangular_b(o2_tiled_mma, tOrNewV, tOrQK, tOrO, sQK.shape)
        self._math_order_notify(wg_idx)

핵심은 _gemm_upper_triangular_b 함수입니다. 이 함수는 t_a @ t_b.T 연산을 수행하는데, t_b가 상삼각 행렬이라는 특성을 활용합니다. t_b[n,k]에서 n < k인 경우 0이므로, t_b.T는 상삼각 행렬이 됩니다.

    @cute.jit
    def _gemm_upper_triangular_b(self, tiled_mma, t_a, t_b, t_acc, b_shape):
        """Compute C += A @ B.T; B[n,k] == 0 for n < k, so B.T is upper triangular."""
        atom = cute.make_mma_atom(tiled_mma.op)
        cB = cute.make_identity_tensor(cute.product_each(b_shape))
        # M groups share B coordinates; inspect one M group to avoid duplicate work.
        b_thread_layout = cute.slice_(tiled_mma.thr_layout_vmnk, (None, 0, None, None))
        b_coords = [
            tiled_mma.get_slice(b_thread_layout(tid)).partition_B(cB)
            for tid in range(cute.size(b_thread_layout))
        ]
        # qk_or_kk_mask zeros n < k for QK/KK in partition_B's (N,K) coordinates.
        # CollectiveInverse updates only diagonal/lower blocks; beta scaling and
        # dtype conversion also preserve these zeros.
        for k in cutlass.range_constexpr(cute.size(t_b, mode=[2])):
            for n in cutlass.range_constexpr(cute.size(t_b, mode=[1])):
                if cutlass.const_expr(
                    any(
                        thread_coords[i, n, k][0] >= thread_coords[i, n, k][1]
                        for thread_coords in b_coords
                        for i in range(cute.size(t_b, mode=[0]))
                    )
                ):
                    b = t_b[None, n, k]
                    b = cute.make_tensor(
                        b.iterator, cute.append(b.layout, cute.make_layout(1))
                    )
                    c = t_acc[None, None, n]
                    c = cute.make_tensor(
                        c.iterator, cute.append(c.layout, cute.make_layout(1))
                    )
                    cute.gemm(atom, c, t_a[None, None, k], b, c)

이 함수의 핵심은 if cutlass.const_expr(...) 조건문입니다. cutlass.const_expr은 컴파일 타임에 평가되는 조건문으로, 런타임 오버헤드 없이 특정 코드 블록의 실행 여부를 결정할 수 있게 합니다. 여기서는 b_coords를 통해 B 행렬의 각 타일이 커버하는 (N, K) 좌표를 확인하고, thread_coords[i, n, k][0] >= thread_coords[i, n, k][1] 조건을 통해 해당 타일이 완전히 대각선 아래에 있지 않은지 (즉, 0이 아닌 원소를 포함할 가능성이 있는지) 검사합니다. 만약 타일의 모든 원소가 0이라면, 해당 cute.gemm 호출을 건너뛰게 됩니다.

리뷰어 kahyunnam의 피드백에 따라, atom = cute.make_mma_atom(tiled_mma.op)과 같이 호출자의 tiled_mma에서 atom을 재구성하여, MMA 타일의 K/N 비율 변경 시에도 코드가 유연하게 대응하도록 개선되었습니다. 또한, qk_or_kk_mask와 CollectiveInverse 연산이 대각선 아래 블록만 업데이트하여 0 타일의 불변성을 보장한다는 설명이 코드 주석에 추가되어, 최적화의 근거를 명확히 했습니다.

tests/gdn/test_prefill_delta_rule.py

성능 최적화는 항상 정확성 검증이 수반되어야 합니다. 이 PR에서는 test_sm12x_triangular_tiles_match_dense라는 새로운 테스트 케이스를 추가하여 최적화의 정확성을 보장합니다.

@torch.inference_mode()
@pytest.mark.parametrize("beta", [False, True])
@pytest.mark.parametrize("alpha", [False, True])
@pytest.mark.parametrize("seq_lens", [[64], [93]])
@pytest.mark.parametrize("dtype", ["float16", "bfloat16"])
def test_sm12x_triangular_tiles_match_dense(
    qkv_factory,
    monkeypatch,
    dtype: str,
    seq_lens: list[int],
    alpha: bool,
    beta: bool,
    seed: int = int(os.environ.get("SEED", "0")),
):
    _skip_if_not_sm12x()

    import cutlass.cute as cute
    from flashinfer.gdn_kernels.delta_rule_dsl import delta_rule_sm120 as sm12x

    class DenseMmaDeltaRule(sm12x._FullyFusedDeltaRuleSm120):
        @cute.jit
        def _gemm_upper_triangular_b(self, tiled_mma, t_a, t_b, t_acc, b_shape):
            cute.gemm(tiled_mma, t_acc, t_a, t_b, t_acc)

    def run():
        return _test_prefill_kernel(
            qkv_factory,
            dtype,
            1,
            1,
            1,
            128,
            64,
            seq_lens,
            1.0,
            alpha,
            beta,
            False,
            seed,
        )

    our_o, our_state = run()
    monkeypatch.setattr(sm12x, "_FullyFusedDeltaRuleSm120", DenseMmaDeltaRule)
    monkeypatch.setattr(
        sm12x, "_get_prefill_kernel", sm12x._get_prefill_kernel.__wrapped__
    )
    ref_o, ref_state = run()

    torch.testing.assert_close(our_o, ref_o, atol=0, rtol=0)
    torch.testing.assert_close(our_state, ref_state, atol=0, rtol=0)

이 테스트는 monkeypatch를 사용하여 최적화된 _gemm_upper_triangular_b 함수를 항상 cute.gemm을 호출하는 DenseMmaDeltaRule 클래스로 임시로 교체합니다. 이렇게 하면 최적화가 적용되지 않은 '밀집(dense)' MMA 연산 결과와 최적화가 적용된 현재 코드의 결과를 비교할 수 있습니다. torch.testing.assert_close(..., atol=0, rtol=0)를 통해 두 결과가 비트 단위로 동일한지 확인하여, 최적화가 성능을 향상시키면서도 기존의 정확성을 완벽하게 유지함을 입증합니다.

왜 이게 좋은 최적화인가?

이 PR은 여러 면에서 좋은 최적화 사례를 보여줍니다.

1. 명확한 성능 향상

PR 설명에 따르면, 이 최적화는 다양한 구성에서 GPU 시간을 유의미하게 단축시킵니다.

  • RTX 5060 (SM120): 멀티 청크 구성에서 5.39%에서 최대 6.77%까지 GPU 시간이 감소했습니다.
  • RTX PRO 6000 Blackwell (SM121): waynehacking8 리뷰어의 벤치마크 결과에 따르면, 모든 30개 구성에서 3.1%에서 최대 7.9%까지 성능이 향상되었습니다. 특히, RTX 5060에서는 미미했던 L=64, B=32와 같은 짧은 시퀀스에서도 FP16 5.3%, BF16 7.9%의 명확한 개선을 보였습니다.

이러한 성능 향상은 불필요한 MMA 연산 타일(32개 중 12개)을 건너뛰어 직접적으로 연산량을 줄였기 때문에 발생합니다. 이는 GPU의 연산 유닛이 0으로 채워진 데이터를 처리하는 데 낭비되는 시간을 줄여줍니다.

2. 정확성 보장

새롭게 추가된 test_sm12x_triangular_tiles_match_dense 테스트는 최적화된 코드와 최적화되지 않은(밀집) 코드의 출력이 비트 단위로 동일함을 검증합니다. 이는 성능 최적화가 기능적 회귀(functional regression)를 일으키지 않았음을 강력하게 보장합니다. 성능 개선만큼이나 중요한 것이 기존 동작의 정확성 유지인데, 이 PR은 이를 철저히 지켰습니다.

3. 도메인 지식의 효과적인 활용

이 최적화는 GDN 커널의 우측 피연산자가 상삼각 행렬이라는 도메인 특성을 깊이 이해하고 이를 활용한 결과입니다. 단순히 일반적인 최적화 기법을 적용하는 것을 넘어, 특정 문제의 수학적/알고리즘적 특성을 파고들어 최적의 해결책을 찾았다는 점에서 높은 가치를 가집니다.

4. 저수준 GPU 프로그래밍 기법의 숙련도

CUTLASS DSL(cute.gemm, cutlass.const_expr)을 사용하여 GPU의 MMA 유닛 동작을 세밀하게 제어하고 최적화하는 것은 상당한 저수준 프로그래밍 지식과 경험을 요구합니다. cutlass.const_expr을 통해 컴파일 타임에 불필요한 연산을 제거하는 방식은 런타임 오버헤드를 최소화하면서 성능을 극대화하는 모범적인 사례입니다.

5. 리뷰어 피드백 반영을 통한 코드 품질 향상

리뷰 과정에서 kahyunnam 리뷰어의 날카로운 지적(예: tiled_mma 파라미터 사용, 0 타일 불변성 설명 추가)이 있었고, PR 작성자는 이를 적극적으로 반영하여 코드의 유연성, 가독성, 그리고 견고함을 더욱 높였습니다. 이는 협업을 통한 코드 품질 향상의 좋은 예시입니다.

결론

FlashInfer의 GDN 커널 최적화 PR은 SM12x GPU에서 상삼각 행렬의 특성을 활용하여 불필요한 MMA 연산을 제거함으로써, 최대 7.9%에 달하는 GPU 시간 단축을 달성했습니다. 이는 도메인 지식, 저수준 GPU 프로그래밍 기술, 그리고 철저한 테스트를 통한 정확성 검증이 결합된 훌륭한 최적화 사례입니다. 이러한 최적화는 딥러닝 모델의 추론 성능을 한 단계 끌어올리는 데 기여하며, 고성능 컴퓨팅 분야에서 미세한 개선이 얼마나 큰 차이를 만들 수 있는지 다시 한번 보여줍니다.

References

"CUTLASS: CUDA Templates for Linear Algebra Subroutines" GitHub Repository PyTorch torch.testing.assert_close Documentation

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글