NVIDIA integra Universal Sparse Tensor ao nvmath-python 0.9.0 e promete até 444x de speedup em SpMV

A NVIDIA integrou o Universal Sparse Tensor (UST) ao nvmath-python v0.9.0. A ideia do UST é desacoplar a sparsity do tensor do seu layout de memória, o que dá flexibilidade e performance tanto pra computação científica quanto pra deep learning esparso.

O anúncio é continuação de um post anterior que apresentou o conceito. Aqui a NVIDIA detalha como ele funciona dentro do nvmath-python.

O que entrega

• Interop zero-copy: conversão sem movimentação de dados com PyTorch, SciPy, CuPy e NumPy.
• Formatos customizados: dá pra definir esquemas novos de sparsity via DSL.
• Operações polimórficas: funções agnósticas de sparsity (tipo matmul() e solve()) que despacham pra kernel otimizado ou geram código esparso automaticamente.
• Injeção em PyTorch: usar UST em modelos existentes sem reescrever código.
• Cache transparente: evita recompilação JIT/LTO e replanning em execuções repetidas.

Tensor format DSL

O UST descreve formatos comuns (COO, CSR, CSC, BSR) e menos comuns via uma DSL (Domain-Specific Language) simples. Um CSC, por exemplo, fica assim:

(i, j) -> (j : dense, i : compressed) # CSC

Na biblioteca C++ MatX, essa DSL é implementada via tipos e templating, o que permite otimizações em tempo de compilação. O nvmath-python pegou um caminho mais pythônico: a DSL vira objeto Python, construído dinamicamente em runtime.

i, j = (Dimension(dimension_name="i"), Dimension(dimension_name="j"))
CSC = TensorFormat([i, j], {j: LevelFormat.DENSE, i: LevelFormat.COMPRESSED})

A troca aqui é clara: perde um pouco de performance ao inspecionar objetos de formato em runtime, mas ganha flexibilidade (incluindo parsing de strings). Como essas decisões acontecem fora do caminho crítico, é troca razoável.

Interop com PyTorch, SciPy e CuPy

Converter entre formatos esparsos (COO, CSR, CSC, BSR, BSC, DIA) e UST é zero-copy: o UST só referencia os buffers originais. Exemplo de SciPy COO virando UST:

row = np.array([0, 0, 1, 1, 2, 3], dtype=np.int32)
col = np.array([0, 1, 1, 3, 2, 3], dtype=np.int32)
val = np.array([1, 2, 3, 4, 5, 6], dtype=np.float32)
coo = sps.coo_array((val, (row, col)), shape=(4, 8))
ust = Tensor.from_package(coo)

E volta com ust.to_package(), também sem cópia.

Formatos menos comuns

Dá pra definir formato próprio. O exemplo da NVIDIA converte um tensor denso do PyTorch pra um formato delta-compressed de 2 bits (parecido com CSR, mas usando diferenças entre coordenadas em runtime), que o PyTorch não suporta nativamente:

delta_format = TensorFormat([i, j], {i: LevelFormat.DENSE,
                                     j: (LevelFormat.DELTA, 2)})
S = U.convert(tensor_format=delta_format)

Aqui tá a parte interessante pra quem pesquisa sparsity: testar um formato novo não exige mais escrever kernel do zero. Define via DSL, plugga e vê o que rola.

Operações polimórficas

O desenvolvedor escreve matmul() ou solve() sem se preocupar com o formato. O sistema inspeciona os operandos e decide: se tem kernel otimizado (tipo cuSPARSE pra CSR), despacha; se não, gera código esparso automaticamente.

A operação se divide em duas fases: planning (decide dispatch ou JIT LTO, configura tudo, cacheia) e execution (roda de fato).

# c := a @ b + c
with Matmul(a, b, c, beta=1.0) as mm:
    mm.plan()
    mm.execute()

Isso amortiza o custo inicial de planejamento em execuções repetidas. Faz sentido em solvers iterativos esparsos e em linear layers podados pra deep learning.

Injeção em PyTorch

Pesquisadores de IA são resistentes a reescrever modelo só pra trocar GEMM das linear layers. O nvmath-python contorna isso embalando o UST numa subclass de torch.Tensor e fornecendo wrappers pra dot, mv, mm, linear etc.

Tem também o reformat_model(), que itera nas linear layers e chama uma função do usuário pra inspecionar, podar e esparsificar os pesos:

def reformat(weight):
    if (...):
        return TorchUST.from_torch(weight.to_sparse_csr())
    return None

Depois disso, inferência usa UST automaticamente nas camadas modificadas.

Performance

Dois experimentos. O primeiro compara SpMV (Sparse Matrix-Vector multiplication) do CuPy e PyTorch contra o matmul() do UST na matriz atmosmodl (1.489.752 × 1.489.752) do Matrix Market. Medindo só o execute() do UST (comparação justa em aplicações com multiplicação repetida):

• UST teve speedups de 1,1x até 444x.
• COO do CuPy e PyTorch foi particularmente lento.
• Em CSR, todos usam cuSPARSE por baixo. UST ganha por reusar a fase de planning.
• Em DIA, UST tem kernel especializado; CuPy converte pra CSR primeiro, PyTorch nem suporta.

O segundo experimento usa o formato MACKO (da publicação "MACKO: Sparse Matrix-Vector Multiplication for Low Sparsity"), que é basicamente o delta-compressed mencionado acima. Os autores do MACKO entregaram uma implementação CUDA que bate dense a partir de 50% de sparsity. A NVIDIA plugou esse kernel como backend do formato delta no UST.

Em matrizes 8192×8192, o MACKO via UST ficou ligeiramente melhor que o MACKO original, porque o padding para no fim de cada linha (otimização fácil de portar pro original).

A leitura prática: se você roda HPC (High-Performance Computing), cientista de dados trabalhando com matriz esparsa enorme, ou pesquisador investigando sparsity em LLM, o nvmath-python virou uma opção séria. Os 444x são em cenário específico (COO), mas mesmo os speedups menores em CSR importam em loop iterativo.

Saiba mais

Documentação completa em nvmath-python.