AI 모델이 점점 더 정교한 문제를 해결하기 위해 기능을 확장하면서, 테스트 시간 확장 또는 추론 시간 확장이라는 개념이 등장하고 있습니다. 이 기술은 AI가 추론 중에 추가적인 컴퓨팅 리소스를 할당하여 여러 가지 가능한 결과를 평가한 후, 가장 적합한 신경망을 선택함으로써 모델의 성능을 향상시키는 방식입니다. 이를 통해 AI는 인간이 복잡한 문제를 분석하고 해결하는 방식과 유사하게, 전략적이고 체계적인 접근법을 취할 수 있습니다.
이 글에서는 NVIDIA 엔지니어들이 최신 오픈 소스 모델 중 하나인 DeepSeek-R1 모델과 추론 과정에서 추가 컴퓨팅 성능을 함께 사용해 복잡한 문제를 해결한 실험에 대해 이야기합니다. 이 실험의 목표는 명시적인 프로그래밍 없이도 수치적으로 정확하고 다양한 어텐션 구조에 최적화된 GPU 어텐션 커널을 자동으로 생성하는 것이었습니다.
그 결과, 경우에 따라 숙련된 엔지니어가 개발한 최적화된 커널보다 더 나은 것으로 나타났습니다.
최적화된 어텐션 커널의 필요성 및 관련 과제
어텐션은 거대 언어 모델(LLM) 개발에 혁신을 가져온 핵심 개념입니다. AI 모델이 작업을 수행할 때 입력 중 가장 관련성이 높은 부분을 선택적으로 집중하도록 도와주는 강력한 메커니즘으로, 중요한 정보에 집중함으로써 더 정확한 예측을 하고 데이터 속 숨은 패턴을 효과적으로 찾아낼 수 있습니다.
어텐션 연산의 계산 복잡도는 입력 시퀀스 길이에 따라 4제곱으로 증가합니다. 따라서 단순한 구현 방식으로는 런타임 오류(예: 메모리 부족)가 발생할 수 있어, 계산 효율성을 높이기 위해 최적화된 저수준 구현(GPU 커널 등)이 필요합니다.
어텐션에는 인과 관계, 상대적 위치 임베딩, 알리바이 등 다양한 변형이 있으며, 엔지니어는 주어진 작업에 따라 이러한 변형을 적절히 조합해 활용해야 합니다.
다중 모달 모델(예: 비전 트랜스포머)은 컴퓨터 비전이나 비디오 생성 모델처럼 시공간 정보를 유지해야 하는 경우가 많아, 공간 이웃 어텐션(Spatial Neighborhood Attention) 같은 특수한 어텐션 메커니즘이 필요합니다. 이로 인해 추가적인 기술적 과제가 뒤따르게 됩니다.
어텐션에 최적화된 GPU 커널을 생성하려면 숙련된 소프트웨어 엔지니어라도 많은 기술과 시간이 필요합니다.
DeepSeek-R1과 같은 최신 LLM은 코드 생성에서 높은 가능성을 보여주었지만, 첫 시도만으로 최적화된 코드를 만들기는 아직 어렵습니다. 따라서 추론 과정에서 다른 전략을 활용해 더 최적화된 코드를 생성하는 접근이 필요합니다.
다음 프롬프트는 상대 위치 임베딩 어텐션 커널에 대한 샘플 사용자 입력입니다.
Please write a GPU attention kernel to support relative position encodings. Implement the relative positional encoding on the fly within the kernel. The complete code should be returned, including the necessary modifications.
Use the following function to compute the relative positional encoding:
def relative_positional(score, b, h, q_idx, kv_idx):
return score + (q_idx - kv_idx)
When implementing the kernel, keep in mind that a constant scaling factor 1.44269504 should be applied to the relative positional encoding due to qk_scale = sm_scale * 1.44269504. The PyTorch reference does not need to scale the relative positional encoding, but in the GPU kernel, use:
qk = qk * qk_scale + rel_pos * 1.44269504
Please provide the complete updated kernel code that incorporates these changes, ensuring that the relative positional encoding is applied efficiently within the kernel operations.LLM은 때때로 환각 코드를 생성하거나 다른 언어 또는 프레임워크의 구문을 혼합하여 즉각적인 코드 오류 또는 비효율성을 유발할 수 있습니다. 최적의 GPU 스레드 매핑을 계산하는 것도 간단하지 않고 까다로운 작업으로, 정확하고 효율적인 커널을 얻기 위해 반복적인 개선이 필요한 경우가 많습니다
최적화된 GPU 커널 생성을 위한 추론 시간 스케일링
최적화된 어텐션 커널에서 최상의 결과를 얻기 위해, NVIDIA 엔지니어들은 일정 기간 동안 폐루프 방식으로 추론하면서 DeepSeek-R1 모델과 함께 특수 검증기를 포함하는 새로운 워크플로우를 개발했습니다.
워크플로우가 먼저 수동 프롬프트에 의해 초기화되고 DeepSeek-R1 모델이 첫 번째 패스에서 GPU 코드(즉, 커널)를 생성합니다. 검증기는 NVIDIA H100 GPU에서 실행됩니다. 이 검증기는 생성된 커널을 분석하고 새로운 프롬프트를 생성하여 DeepSeek-R1 모델에 입력으로 제공합니다.
이 폐쇄 루프 접근 방식은 매번 다른 방식으로 코드를 생성하도록 유도하여 코드 생성 프로세스를 개선합니다. 연구팀은 이 프로세스를 15분 동안 계속 진행하면 어텐션 커널이 개선된다는 사실을 발견했습니다.
이 워크플로우에서는 스탠퍼드의 KernelBench 벤치마크에서 테스트한 결과 레벨 1 문제의 100%, 레벨 2 문제의 96%에 대해 수치적으로 정확한 커널을 생성했습니다.
KernelBench의 레벨 1 해결률은 특정 연산 작업에 대해 효율적인 GPU 커널을 생성하는 LLM의 능력을 평가하는 데 사용되는 수치상의 정확한 지표를 의미합니다. 이 테스트는 GPU 프로그래밍에서 최신 LLM의 능력을 테스트하기 위한 일련의 과제 중 하나입니다.
그림 4는 추론 시간 예산이 에이전트의 해결 속도에 어떤 영향을 미치는지 보여줍니다. 레벨 1 카테고리에서 문제당 10분 이상을 할당하면 워크플로우에서 100개의 문제 대부분에 대해 수치적으로 올바른 코드를 생성할 수 있습니다.
DeepSeek-R1에서 최적화된 GPU 커널
이 결과는 최신 DeepSeek-R1 모델을 활용해 추론 시간 동안 더 많은 컴퓨팅 성능을 투입함으로써 최적의 GPU 커널을 생성하는 방법을 보여줍니다. 아직 새로운 연구 분야이지만, 효과적인 어텐션 커널을 자동으로 생성하는 유망한 접근 방식에 대한 초기 결과가 나오고 있습니다.
좋은 출발을 보였지만, 더 다양한 문제에서 일관되게 우수한 결과를 얻으려면 추가적인 연구와 개선이 필요합니다. 최근 DeepSeek-R1의 발전과 그 잠재력에 대한 기대가 큽니다.
자세한 내용을 알아보거나 시작하려면 현재 build.nvidia.com에서 제공되는 DeepSeek-R1 NIM 마이크로서비스를 확인해 보세요.
관련 리소스
- DLI 과정: 프로덕션 규모에서 추론을 위한 모델 배포하기
- GTC 세션: NVIDIA GPU에서 추론 가속화하기
- GTC 세션: AI와 그 이상을 위한 커널 최적화: Nsight 컴퓨팅의 성능 활용하기
- GTC 세션: 다차원 동형을 통한 딥 러닝을 위한 GPU 코드 생성 및 최적화
- NGC 컨테이너: NVIDIA MLPerf 추론
- NGC 컨테이너: NVIDIA MLPerf 추론
