viva_tensor/sparsity

2:4 Structured Sparsity

NVIDIA Tensor Cores Structured Sparsity Ampere+ Architecture (RTX 3000/4000, A100, H100)

CONCEITO: Em cada grupo de 4 elementos, apenas 2 são não-zero = 50% dos elementos são zero, mas em padrão ESTRUTURADO

POR QUE ESTRUTURADO > ALEATÓRIO:

Sparsity aleatória: difícil de acelerar em hardware
Sparsity estruturada: hardware pode pular zeros eficientemente

FORMATO DE ARMAZENAMENTO:

2 valores FP16 (32 bits)
2-bit máscara indicando posições (4 bits para 4 posições)
Total: 36 bits para 4 elementos = 9 bits/elemento vs 16 bits/elemento
Compressão: ~1.8x

PERFORMANCE:

2x throughput em Tensor Cores (pula multiplicações por zero)
Combinado com INT8: 4x speedup total!

Types

PruneMetrics

</>

Métricas de poda

pub type PruneMetrics {
  PruneMetrics(
    pruned_count: Int,
    total_count: Int,
    approximation_error: Float,
    kept_magnitude_mean: Float,
    pruned_magnitude_mean: Float,
  )
}

Constructors

```
PruneMetrics(
  pruned_count: Int,
  total_count: Int,
  approximation_error: Float,
  kept_magnitude_mean: Float,
  pruned_magnitude_mean: Float,
)
```
Arguments

pruned_count

Elementos podados (zerados)

total_count

Total de elementos

approximation_error

Erro de aproximação

kept_magnitude_mean

Magnitude média mantida

pruned_magnitude_mean

Magnitude média podada

Sparse24Block

</>

Bloco 2:4 (4 elementos com 2 não-zeros)

pub type Sparse24Block {
  Sparse24Block(values: #(Float, Float), positions: #(Int, Int))
}

Constructors

```
Sparse24Block(values: #(Float, Float), positions: #(Int, Int))
```
Arguments

values

Os 2 valores não-zero

positions

Máscara 2-bit indicando posições (0-3 para cada valor)

Sparse24Tensor

</>

Tensor com esparsidade 2:4

pub type Sparse24Tensor {
  Sparse24Tensor(
    blocks: List(Sparse24Block),
    shape: List(Int),
    num_elements: Int,
    memory_bytes: Int,
    sparsity_percent: Float,
  )
}

Constructors

Sparse24Tensor(
  blocks: List(Sparse24Block),
  shape: List(Int),
  num_elements: Int,
  memory_bytes: Int,
  sparsity_percent: Float,
)

Arguments

blocks: Blocos 2:4
shape: Shape original
num_elements: Número de elementos originais
memory_bytes: Memória em bytes
sparsity_percent: Sparsity real (%)

Values

benchmark_sparsity

</>

pub fn benchmark_sparsity() -> Nil

compute_metrics

</>

pub fn compute_metrics(
  original: tensor.Tensor,
  sparse: Sparse24Tensor,
) -> PruneMetrics

Calcula métricas de poda

decompress

</>

pub fn decompress(sparse: Sparse24Tensor) -> tensor.Tensor

Reconstrói tensor denso a partir de 2:4 sparse

main

</>

pub fn main() -> Nil

prune_24_gradient

</>

pub fn prune_24_gradient(
  weights: tensor.Tensor,
  gradients: tensor.Tensor,
) -> Sparse24Tensor

Poda baseada em gradiente (para treinamento) Mantém elementos com maior |valor × gradiente|

prune_24_magnitude

</>

pub fn prune_24_magnitude(t: tensor.Tensor) -> Sparse24Tensor

Aplica poda 2:4: mantém os 2 maiores em cada grupo de 4 Estratégia: magnitude (abs) - padrão da NVIDIA

sparse_matmul

</>

pub fn sparse_matmul(
  sparse_a: Sparse24Tensor,
  dense_b: tensor.Tensor,
) -> #(tensor.Tensor, Float)

Matmul com matriz esparsa 2:4 Em hardware real (Tensor Cores), isso é 2x mais rápido!