Context
demo_14 and demo_15 (all-sparse path) proved that sparse attention projections (qkv and o at 70% sparsity) work end-to-end and are additive with FFN sparsity. But the current implementation requires the user to manually construct SparseLinear layers inside a custom attention module.
A first-class API would let users write something like:
attn = sparselab.SparseAttention(d_model=384, n_heads=6, attn_sparsity=0.7)
What needs to happen
- Design the
SparseAttention API — should it wrap PyTorch's F.scaled_dot_product_attention or replace it?
- Decide: do we sparsify the attention weights (qkv/o projections) or the attention pattern (the softmax output)? v0.1 does the former. The latter is a different research direction (sparse attention masks like BigBird, Longformer).
- Implement as an
nn.Module in sparselab/nn.py.
- Tests + a demo showing the API.
Files to read
examples/demo_14_sparse_attention.py — current manual implementationexamples/demo_15_mini_gpt.py — CausalSelfAttention class that switches between dense and sparse projectionsdocs/demos/milestone_10.md — results showing sparse attention works at 10k-step scale
Context
demo_14 and demo_15 (all-sparse path) proved that sparse attention projections (qkv and o at 70% sparsity) work end-to-end and are additive with FFN sparsity. But the current implementation requires the user to manually construct
SparseLinearlayers inside a custom attention module.A first-class API would let users write something like:
What needs to happen
SparseAttentionAPI — should it wrap PyTorch'sF.scaled_dot_product_attentionor replace it?nn.Moduleinsparselab/nn.py.Files to read
examples/demo_14_sparse_attention.py— current manual implementationexamples/demo_15_mini_gpt.py—CausalSelfAttentionclass that switches between dense and sparse projectionsdocs/demos/milestone_10.md— results showing sparse attention works at 10k-step scale