MSLK 커널 레퍼런스(웹사이트)

MSLK 커널 레퍼런스는 Transformer attention, 선형 레이어, 희소 모델 등 7개 도메인별로 GPU 커널 API와 사용 경로를 정리한다. 자동 디스패치를 기본으로 하되, 특정 아키텍처·paged KV·split-K·결정적 실행이 필요할 때만 백엔드를 명시하도록 안내한다.

설치 예제는 CUDA 13.0 및 ROCm 7.1용 wheel을 제공한다. import mslk로 연산자를 등록한 뒤 mslk.attention, mslk.gemm, mslk.moe 같은 도메인 모듈을 불러와 torch.ops.mslk 엔트리를 사용할 수 있다.

주요 예제는 다음과 같다.

• memory_efficient_attention를 이용한 causal attention
• FP8 row-wise 양자화 후 BF16 GEMM 실행
• 여러 K/V 청크의 partial attention과 정확한 softmax 병합
• index_shuffling 및 gather/scatter 기반 MoE 라우팅

GEMM은 일반적으로 활성화 [M,K]와 가중치 [N,K]를 받아 X @ W.T를 계산하며, 양자화 스케일 레이아웃은 생성한 양자화 방식과 맞춰야 한다. 문서 전반에서는 마지막 차원의 stride가 1이어야 한다는 공통 규칙과 함께 attention·GEMM·MoE 텐서 shape를 설명하고, 약 400개 API 항목으로 확장 가능한 참조 정보를 제공한다.