deepseek-ai/DeepSelect

해결하는 문제

DeepSelect는 특정 AI 워크로드에서 torch.topk를 대체하기 위해 설계된 고성능 TopK 커널 구현체입니다. 표준 PyTorch 구현체의 성능 저하 문제를 해결하며, 텐서에서 상위 k 요소를 선택할 때 메모리 대역폭 효율을 2x에서 20x까지 향상시킵니다.

작동 방식

DeepSelect는 두 가지 주요 시나리오에 최적화된 CUDA 커널을 제공합니다: bfloat16 입력용 'Lightning Indexer' (DeepSeek 스파스 어텐션에서 사용)와 float32 입력용 'Sampling Scenario' (LLM 샘플링에서 사용). 메모리 대역폭을 최적화하고, 정렬을 비활성화하거나 값 출력을 건너뛰어 성능을 추가로 향상시킬 수 있습니다. 또한 가변 길이 행을 지원하고 내장된 NaN 체크 기능도 제공합니다.

대상 사용자

DeepSeek 모델(V3.2, V4, V4.1 등)을 사용하는 개발자 및 연구자, 또는 대규모 LLM 추론 및 학습에 매우 최적화된 TopK 연산이 필요한 사용자에게 적합합니다.

주요 특징

• 매우 빠른 성능: torch.topk보다 2x에서 20x까지 빠릅니다.
• 특화된 시나리오: bfloat16 (스파스 어텐션) 및 float32 (샘플링) 데이터 유형에 최적화되어 있습니다.
• 유연한 출력: 값 출력을 건너뛰거나 정렬된 인덱스를 비활성화하여 성능을 극대화할 수 있습니다.
• 가변 길이 지원: end 텐서를 사용해 길이가 다른 행을 처리할 수 있습니다.
• 강건성: 계산 중 오류를 방지하기 위해 내장된 NaN 체크 기능이 있습니다.