[flashinfer] Blackwell 시대를 위한 최적화: FlashInfer의 SM120 Block-Sparse Attention 백엔드 도입기

PR 링크: flashinfer-ai/flashinfer#4951 상태: Merged | 변경: +18238 / -0

들어가며

LLM(Large Language Model) 추론의 효율성을 극대화하기 위해서는 하드웨어 아키텍처에 최적화된 커널 설계가 필수적입니다. 특히 NVIDIA의 최신 아키텍처인 Blackwell(SM120)은 기존 아키텍처보다 강력한 연산 능력과 새로운 메모리 접근 기술을 제공합니다.

이번 글에서는 FlashInfer 라이브러리에 새롭게 추가된 cake_sage 백엔드를 분석합니다. 이 PR은 SM120 아키텍처에서 Block-Sparse Attention(BSA)의 성능을 끌어올리기 위해 'Cake'라는 코드 생성 엔진을 활용한 JIT(Just-In-Time) 로더와 최적화된 커널을 도입했습니다. 기존 CuTe DSL 기반 구현체 대비 기하평균 1.11배(약 11%)의 성능 향상을 달성한 비결이 무엇인지 실제 코드를 통해 살펴보겠습니다.

---

코드 분석: 핵심 변경 사항

1. Manifest 기반의 JIT 로더와 커널 특수화(Specialization)

가장 큰 변화는 cake_sage_block_sparse_attention_manifest.json 파일의 도입입니다. 이 파일은 다양한 연산 형태(Shape)와 설정에 따라 최적화된 커널을 매핑하는 역할을 합니다.

[Before: 기존 방식 (추정)] 기존에는 범용적인 커널 하나가 런타임에 조건문을 통해 다양한 상황을 처리하거나, 수동으로 정의된 몇 가지 특수화된 커널을 사용했습니다.

[After: Manifest 기반 JIT 로딩]

// csrc/cake_sage_block_sparse_attention/cake_sage_block_sparse_attention_manifest.json
{
  "routes": [
    {
      "arch": "sm_120a",
      "metadata": {
        "operator": "prequantized_direct_block_sparse_attention",
        "specializations": {
          "BLOCK_SIZES_MODE": 0,
          "CONTIGUOUS_BLOCK_INDICES": 1,
          "FULL_K64_TILES": 1,
          "HAS_BLOCK_NUMS": 0,
          "UNIFORM_NONEMPTY": 1
        }
      },
      "route": "sage_q64_nums0_sizes0_full1_uniform1_contiguous1",
      "stages": [{
        "module": {"name": "cake_sage_block_sparse_attention_4532f6a2faed9609fb8e"},
        "source_build": {
          "compile_options": ["--use_fast_math"],
          "source_tma_abi": "pointer"
        },
        "tma_workspace_bytes": 512
      }]
    }
  ]
}

이 방식의 장점은 컴파일 타임 최적화를 극대화할 수 있다는 것입니다. BLOCK_SIZES_MODE, FULL_K64_TILES와 같은 파라미터를 커널의 템플릿 인자로 고정함으로써, 런타임 오버헤드를 줄이고 컴파일러가 루프 언롤링(Loop Unrolling)이나 레지스터 할당을 더 효율적으로 수행하게 합니다.

2. TMA(Tensor Memory Accelerator) 활용 및 메모리 관리 전략

SM120(Blackwell) 아키텍처의 핵심 기능 중 하나는 TMA(Tensor Memory Accelerator)입니다. 이번 PR에서는 TMA 디스크립터 워크스페이스를 호출자(Caller)가 소유하도록 설계하여 불필요한 메모리 할당 오버헤드를 제거했습니다.

[After: TMA 및 Caller-owned Workspace 설정]

"metadata": {
  "caller_owned_output": true,
  "caller_owned_tma_workspace": true
},
"tma_abi": "pointer",
"tma_workspace_bytes": 512

커널 내부에서 매번 메모리를 할당하거나 관리하는 대신, 외부(Python/C++ API)에서 미리 할당된 tma_workspace를 전달받습니다. 이는 특히 CUDA Graph를 사용하는 환경에서 매우 중요한데, 커널 실행 시 내부 할당이 발생하면 그래프 캡처가 복잡해지거나 성능 저하가 발생할 수 있기 때문입니다.

3. Prequantized Sage Attention 지원

이번 업데이트는 INT8 또는 FP8로 양자화된 입력을 처리하는 prequantized 모드를 지원합니다. 이는 최신 LLM 서빙 엔진들이 메모리 대역폭 한계를 극복하기 위해 양자화를 적극적으로 사용하는 트렌드를 반영한 것입니다.

---

왜 이게 좋은가: 성능과 설계의 교훈

1. 11.3%의 성능 향상 (Geometric Mean)

PR 설명에 포함된 벤치마크 결과는 놀랍습니다. 기존의 CuTe DSL 기반 구현체(Baseline)와 비교했을 때, 거의 모든 성능 테스트 케이스(perf_*)에서 우위를 점했습니다.

Performance shape	Baseline / Export (Speedup)
perf_b8_h32_s1024_density90	1.155x
perf_b8_h32_s2048_density50	1.139x
perf_b8_h32_s32768_density10	1.099x

특히 데이터 밀도(Density)가 높거나 시퀀스 길이가 긴 경우에도 안정적으로 10% 이상의 성능 향상을 보여줍니다. 이는 Cake 백엔드가 생성하는 커널이 Blackwell의 하드웨어 파이프라인을 더 촘촘하게 활용하고 있음을 시사합니다.

2. 유연성과 엄격함의 조화

리뷰어 yyihuang은 다양한 테스트 스크립트(test_cake_sage_block_sparse_attention.py)를 실행하며 검증을 진행했습니다. Manifest 파일에 정의된 21개의 Shape(정확성 테스트 13개, 성능 테스트 8개)를 모두 통과해야만 배포가 가능하도록 설계된 프로세스는, 고성능 커널이 자칫 놓치기 쉬운 엣지 케이스(Ragged Q/K, Empty rows 등)에 대한 신뢰성을 보장합니다.

3. 일반적인 교훈: 하드웨어 특화 JIT의 힘

모든 상황에 맞는 하나의 거대한 커널(Monolithic Kernel)을 만드는 시대는 지났습니다. 하드웨어 아키텍처(SM120)와 연산의 특성(Sparse pattern, Block size)에 맞춰 최적화된 커널 조각들을 생성하고, 이를 Manifest로 관리하며 런타임에 적재하는 방식이 현대 고성능 컴퓨팅의 표준이 되고 있습니다.

---

마치며

이번 FlashInfer의 PR은 최신 하드웨어인 Blackwell을 어떻게 소프트웨어적으로 정복해야 하는지를 잘 보여주는 사례입니다. TMA와 같은 하드웨어 가속 기능을 적극 활용하고, JIT 컴파일을 통해 커널을 특수화하며, 메모리 관리의 주도권을 사용자에게 넘겨 오버헤드를 최소화했습니다.

Blackwell GPU를 도입할 예정이거나 고성능 어텐션 커널 구현에 관심이 있는 엔지니어라면, FlashInfer의 cake_sage 구현 방식을 깊이 있게 연구해 보시길 권장합니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#tensor-memory-accelerator-tma
• https://github.com/flashinfer-ai/flashinfer
• https://pytorch.org/docs/stable/generated/torch.compile.html

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