MSLK: biblioteca de kernels GPU fusionados para cargas transformer

O que é a MSLK

A MSLK é uma biblioteca de kernels GPU fusionados voltada para cargas de trabalho transformer. Ela cobre sete domínios: atenção, GEMM de baixa precisão, quantização, roteamento de Mixture of Experts (MoE), convolução, runtime/dispositivo e utilitários C++. A maioria das operações é acessada via torch.ops.mslk.* após import mslk. A versão atual é a 1.3.0, compatível com PyTorch 2.13, Python 3.10–3.14, CUDA 13.0/13.2 (arquiteturas SM80, SM90a, SM100a, SM120a) e ROCm 7.1/7.2 (gfx942).

Domínios principais

Atenção: O ponto de entrada é memory_efficient_attention, com dispatch automático baseado em dtype, dimensão de head, máscara, dropout e hardware. Há backends específicos como CUTLASS, Flash, Flash3, CuTe Hopper, CuTe Blackwell, CK (AMD) e Triton split-K, cada um com suporte distinto a dtypes, backward, dropout, varlen e paged KV.

GEMM de baixa precisão: As operações seguem uma convenção de nomes que codifica o contrato - por exemplo, f8f8bf16_rowwise significa FP8 × FP8 → BF16 com escala por linha. Há variantes rowwise, blockwise, groupwise, INT4, FP4 (NVFP4, MXFP4), MXFP8×MXFP4, e até emulação de FP32 via nove produtos BF16 (bf16x9_gemm). O usuário é responsável por fornecer as escalas na granularidade exata que o kernel espera.

MoE: A biblioteca expõe as peças de roteamento separadamente (index_shuffling, gather, scatter), permitindo composição flexível. Também oferece módulos compostos como BaselineMoE e MetaShufflingMoE.

Convenções de forma

As convenções são consistentes: atenção usa [B, M, H, K], GEMM usa ativações [M, K] e pesos [N, K] (computando X @ W.T), e MoE usa scores [T, E]. A última dimensão deve sempre ter stride 1.

Armadilhas relevantes

• O registro de operadores é baseado em imports: import mslk carrega a biblioteca nativa, mas módulos específicos (como mslk.gemm.triton.int4_gemm no ROCm) precisam ser importados separadamente.
• Dois schemas de MoE (gather_scale_quant_dense_tokens e silu_mul_quant) têm histórico de divergência entre registro Python e schema declarado - vale validar antes de exportar.
• GQA broadcasting é explícito: a API aceita [B,M,G,H,K] mas não expande K/V automaticamente.
• MSLK_PYTHON_ONLY=1 não torna kernels nativos disponíveis em CPU.

Licença e instalação

A MSLK é licenciada sob BSD. A instalação via pip usa wheels específicas por plataforma (CUDA ou ROCm) a partir do índice do PyTorch.