[flashinfer] FlashInfer: sm_107a 아키텍처를 위한 Fused Decode List Scorer 프리페칭 최적화

PR 링크: flashinfer-ai/flashinfer#8210 상태: Merged | 변경: +144 / -99

들어가며

FlashInfer의 cake_msa_icp_indexer는 대규모 언어 모델의 인덱스-컨텍스트 병렬(ICP) 희소 어텐션을 처리하는 핵심 컴포넌트입니다. 특히 sm_107a(NVIDIA VR200)와 같은 최신 아키텍처에서 디코딩 성능은 전체 추론 처리량에 큰 영향을 미칩니다. 이번 PR은 fused 디코딩 체인(List Scorer + List Publisher)에서 List Scorer가 PDL(Programmatic Dependent Launch) 대기 상태에 들어가기 전, 필요한 History Index-K 페이지를 미리 로드(Prefetch)하도록 최적화하여 GPU 유휴 시간을 줄이고 성능을 개선했습니다.

코드 분석

1. cake_jit.py 및 build.json의 커널 교체

이번 변경은 sm_107a 아키텍처를 타겟으로 하는 커널 심볼을 업데이트하여 프리페칭 로직이 포함된 새로운 커널로 교체하는 작업을 포함합니다.

# Before
"kernel_symbol": "kernel_cake_msa_icp_indexer_fd54418e72fdf464db9d",
"library": "cake_msa_icp_indexer_fd54418e72fdf464db9d",

# After
"kernel_symbol": "kernel_cake_msa_icp_indexer_4f9fe98ff0678c29ec27",
"library": "cake_msa_icp_indexer_4f9fe98ff0678c29ec27",

이러한 심볼 변경은 JIT 컴파일된 커널이 더 효율적인 메모리 액세스 패턴을 갖도록 수정되었음을 의미합니다. build.json에서도 기존 sm_107a 라이브러리 항목이 제거되고, 최적화된 로직을 포함한 새로운 라이브러리 경로가 등록되었습니다.

2. 문서화 및 동작 명세 업데이트 (docs/api/sparse.rst)

코드의 동작 변경에 맞춰 API 문서 또한 업데이트되었습니다. 이제 sm_107a 환경의 32 또는 64 로컬 K 행 페이지 설정에서 fused 체인도 score plane과 동일하게 프리페칭을 수행함을 명시합니다.

- the fused chain has no such read.
+ both the score plane and the fused chain's list scorer, read `seq_lens` and 
+ `block_table` and start loading the history index-K pages before their wait.

왜 이게 좋은가

이번 최적화의 핵심은 메모리 레이턴시 은폐(Memory Latency Hiding)입니다. 기존에는 PDL 대기(이전 커널의 완료를 기다리는 작업)가 끝난 후에야 History Index-K 데이터를 읽기 시작했으나, 이제는 대기 시간 동안 미리 데이터를 로드함으로써 실제 연산이 시작될 때 데이터가 이미 캐시나 레지스터에 준비되어 있게 됩니다.

성능 개선 수치

• TP4 P128 Steady 모드: 9.409ms → 9.357ms (약 1.006x 속도 향상)
• TP2 P128 Steady 모드: 14.442ms → 14.340ms (약 1.007x 속도 향상)

이러한 미세한 최적화는 대규모 추론 서비스에서 누적될 때 전체적인 토큰 생성 처리량(Throughput)을 높이는 데 기여합니다. 특히 sm_107a와 같은 고성능 아키텍처에서는 메모리 대역폭을 얼마나 효율적으로 활용하느냐가 성능의 병목을 결정짓기 때문에, 이러한 프리페칭 전략은 매우 효과적인 일반적 교훈을 제공합니다.

결론

이번 PR은 특정 아키텍처(sm_107a)의 특성을 깊이 이해하고, 커널 간의 의존성을 관리하는 PDL 메커니즘을 활용하여 메모리 액세스를 최적화한 사례입니다. 하드웨어의 특성에 맞춘 세밀한 커널 튜닝이 어떻게 실질적인 성능 향상으로 이어지는지 잘 보여줍니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-data-copies

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