Design Space 03: Compositional Sparsity and Structured Tensors

Series: AI-Native Programming Language Research — Design Space Memos Document 3 of 6 (planned: shapes, AD-through-effects, sparsity, distribution, verification, hardware abstraction)

1. Problem statement

A "tensor" in modern ML is rarely just a dense block of floats. It is, increasingly:

• 50% structured-sparse (NVIDIA Ampere 2:4) for inference acceleration.
• Block-sparse (BigBird, Longformer) for long-context attention.
• Low-rank (LoRA, IA³, DoRA) for parameter-efficient fine-tuning.
• Dynamically sparse (Mixtral, Switch Transformer) via MoE routing where only a few experts activate per token.
• Quantized (INT4, FP8, NF4) with associated scale/zero-point metadata.
• Compressed in the KV-cache (PagedAttention, KV quantization, eviction policies).

Each of these is a structure on the tensor that the compiler should know about and exploit. None of them are language-level concepts in any current system. They are library types, kernel implementations, or runtime metadata, glued together by hand-written conversion code at every boundary.

This is the third large class of "bolted-on" features in the Python ML stack. The first two (shape, AD over effects) at least have a clean theory waiting to be implemented; sparsity has a clean theory for static patterns (TACO, MLIR's sparse_tensor dialect, Finch) and almost no theory for the dynamic and composed cases that dominate modern workloads.

This memo: maps the taxonomy, surveys what exists, identifies where compositionality breaks, and proposes that sparsity-as-type is the right framing — with sparsity patterns composing under tensor operations the way effects compose under sequencing in memo 2.

2. Taxonomy of sparsity

Two orthogonal axes: when the pattern is known (static vs. dynamic) and what kind of structure it has (point, block, low-rank, conditional). The cross product produces eight or so cells, of which six matter:

Real workloads compose multiple cells. Mixtral does dynamic conditional routing into experts whose weights are dense but whose backward gradient flow is sparse. LoRA is static low-rank applied as a delta on top of dense pretrained weights. Long-context attention with KV cache uses block sparsity (sliding window) plus dynamic eviction.

The compositionality problem is exactly that the cells don't compose cleanly under tensor operations. Multiplying a block-sparse matrix by a low-rank matrix does not produce a clean type — the result depends on which dimensions overlap with which structure. Current systems handle this by giving up: convert one side to dense, do the operation, hope the result has discoverable structure.

3. The static end: TACO and its descendants

The static-sparsity story is, comparatively, in good shape.

TACO (Kjolstad et al., OOPSLA 2017) is the foundational work. It generates code for sparse tensor operations from a tensor index notation, given format specifications for each operand. The user writes A(i,j) = B(i,k) * C(k,j) and declares whether each tensor is dense, CSR, CSC, COO, or one of many other formats. TACO's compiler generates the appropriate iteration code, including the awkward inner loops where you intersect the nonzero patterns of two operands.

The key contribution is format abstraction (Chou, Kjolstad & Amarasinghe, OOPSLA 2018): a sparse format is a sequence of level formats (dense, compressed, singleton, etc.) along each dimension. This factors the format design space and makes it possible to compose formats — block-CSR is just a CSR with a dense inner level. This is the first piece of compositional theory that actually works.

MLIR's sparse_tensor dialect (Bik, Lewis, et al.) is the production realization of TACO's ideas inside an industrial compiler. It accepts annotated tensor types, generates iteration code, and integrates with the rest of MLIR's optimization infrastructure. Mojo's sparsity story will lower through this.

SparseTIR (Ye et al., ASPLOS 2023) brings format composability into TVM, with a focus on deep-learning workloads. It supports format hierarchies and integrates with auto-scheduling.

Finch (Ahrens et al.) is a Julia-based sparse array language that extends the TACO model with control flow — it can express algorithms that branch, conditionally compute, and combine sparse and dense computations in one program. This is the most mature work on integrating sparse arrays with general programming language constructs.

What these get right:

• Sparsity formats compose via level abstractions.
• Code generation is automatic from index notation.
• The result type of a sparse operation can be inferred from operand types.
• AD-friendly variants exist (sparse Jacobians fall out of the same machinery).

What they don't address:

• Dynamic sparsity (data-dependent patterns like MoE routing).
• Hardware-imposed sparsity (Ampere 2:4 isn't expressible as a level format).
• Composition with other structure (low-rank, quantization).
• Distribution (sharding a sparse tensor is an open question).

4. The dynamic end: MoE and routed computation

Mixture of Experts is the production case of dynamic sparsity, and it is currently a mess.

The standard implementation: a gating network produces routing weights for each token; the top-K experts are selected per token; tokens are dispatched to expert weight matrices via scatter operations; results are gathered back via a sum. This is conceptually a sparse matrix multiplication where the sparsity pattern is determined at runtime per batch.

The implementation reality (Switch Transformer, GShard, Mixtral) is that this is implemented as a series of dense kernels with explicit scatter/gather, plus all-to-all communication if experts are sharded across devices. The kernels don't know about sparsity; they see dense ragged inputs.

MegaBlocks (Gale, Narayanan, Young & Zaharia, MLSys 2023) is the breakthrough: it reformulates MoE as a single block-sparse matrix multiplication where each token-to-expert assignment is a block, and uses dense GPUs to compute it efficiently via specialized block-sparse kernels. This is the closest existing work has come to treating MoE as sparsity rather than as routing.

The language-design observation: MoE routing is a conditional operation — for each token, multiply by the expert it routes to. Conditional computation is exactly what the if statement and the match expression are for in normal languages. The reason MoE doesn't look like an if is that the if is over batched data, and current frameworks have no batched-conditional construct. JAX has vmap over lax.cond, but it has to materialize both branches; what you want is vmap over a router that picks one branch per element without materializing the unused ones.

This is a research-language opportunity. A first-class route or mux construct, where the type captures "this is a sparse selection," would let the compiler generate MegaBlocks-style code automatically and would compose with AD via the effect machinery from memo 2 (the routing decision is the stochastic effect; the gradient through it is the estimator).

5. Hardware-imposed sparsity

A separate category, because the constraints are different.

NVIDIA Ampere structured sparsity (2:4 pattern): 2 of every 4 contiguous values must be zero. Hardware tensor cores accelerate matmul with this pattern by 2× via dedicated metadata paths. The pattern is rigid; you can't choose 3:8 or 4:8.

AMD CDNA, Intel AMX, and various dataflow architectures (Cerebras, Graphcore, Tenstorrent) each have their own structured-sparsity stories with their own constraints. None of them are expressible as TACO level formats — the constraint is which positions can be nonzero, not how many.

The language-design problem: hardware structured sparsity is a predicate on the tensor (positions satisfying a constraint), not a format. This is closer to a refinement type than to a sparse format. Tensor f16 [M, K] satisfying TwoOfFourSparse is a refinement on the dense type, not a different storage scheme.

The cleanest framing I know: separate the logical structure (what positions are zero) from the physical representation (how the nonzeros are stored). TACO conflates these for good engineering reasons, but the hardware-sparsity case wants them separated so that the compiler can match a logical refinement to a physical layout.

This is also where quantization fits. INT4 + scale + zero-point is a compressed representation of a tensor whose logical type is still f32 (with quantization error). The compositionality question — can a tensor be both 2:4 sparse and INT4 quantized? — is yes in principle, used in production, and currently implemented as bespoke kernel code with no type-system support whatsoever.

6. Comparison across the design space

The composition row is the key observation. Every system handles dense + sparse, and most handle their own internal composition. None handles cross-cutting composition (sparse + low-rank + quantized + sharded), which is exactly what production LLM workloads need.

7. Where compositionality actually breaks

The mathematical core: the result type of a tensor operation depends on the input types in non-obvious ways.

For dense × dense → dense, this is trivial. For sparse-format compositions:

• CSR × CSR → ? The result has a sparsity pattern determined by the intersection-then-union of nonzero positions, plus arithmetic cancellation. TACO handles this by generating per-format code; the result format is user-specified.
• Block-sparse × low-rank → ? No clean type. The result is dense in general, but has structure if the block structure aligns with the low-rank factorization, which it usually doesn't.
• 2:4 structured × 2:4 structured → ? No longer 2:4 in general. The hardware accelerates the operation but produces a dense result.
• MoE-routed × dense → MoE-routed. Composes cleanly if you treat routing as a type. Currently nobody does.
• Quantized × quantized → ? Depends on whether you accumulate in higher precision (yes, in practice). The result type carries error bounds that affect downstream operations.

The pattern: sparse-style operations don't form a closed algebra. Most of them lose structure under composition, and the lost structure is exactly what makes the next operation slow.

This has a precise analogy in algebra. Sparse matrices form a category but not a group — composition is defined but doesn't preserve the special structure. The right type-system response is to allow the result type to widen (to a less-structured supertype) when composition doesn't preserve structure, and to narrow (to a more-structured subtype) when the user can prove it does.

This is exactly the territory of subtyping with refinement. A 2:4Sparse Tensor is a refinement of Tensor. Multiplication of two 2:4Sparse Tensors produces a Tensor (widened). The user can opt back into 2:4Sparse Tensor via an explicit cast (which checks the refinement at runtime, or the compiler proves it statically in special cases like masked operations).

This is, again, not a current feature of any production system.

8. AD interactions

Sparsity and AD interact in three distinct ways:

8.1 Sparse gradients

If the loss is sparse in its dependence on a parameter (e.g., embedding lookups), the gradient is naturally sparse. PyTorch handles this via torch.sparse.SparseTensor for embeddings; nobody handles it generally. A compiler that knew about sparsity in the type could propagate this through the backward pass automatically.

8.2 Differentiating through routing

MoE routing is a hard-decision (top-K) that breaks differentiability. Workarounds:

• Straight-through estimator: pretend the routing is identity in the backward pass. Biased.
• Sparsemax (Martins & Astudillo, 2016) and entmax (Peters et al., 2019): differentiable sparse activations that produce exact zeros without breaking the gradient.
• Switch Transformer's router loss: an auxiliary loss that pushes routing toward balance, sidestepping the gradient issue.
• Reparameterization-style relaxations for MoE: not yet standard.

This is where memo 2's effect framework pays direct dividends. Routing is a stochastic effect; the gradient through it is determined by the handler. Sparsemax is one handler, straight-through is another. The user picks (or the compiler picks) based on numerical and statistical considerations.

8.3 Sparse Jacobians for second-order methods

Hessians are usually sparse for structured problems (banded, block-diagonal). Newton's method, K-FAC, and other second-order methods need the sparsity exposed. Current frameworks compute dense Hessians and lose the structure. Sparse-aware AD (TACO has parts of this) is a research area in itself.

9. Open research problems

1. A unified type for sparsity. Static formats, dynamic routing, hardware-imposed patterns, and low-rank factorization are currently four separate concepts. Is there a single type-level abstraction (perhaps a structure predicate on the tensor's value space) that covers all four? My guess: yes, via refinement types parameterized by a predicate language, with TACO-style level formats as one specialization.

2. Composition rules under widening. When sparse × sparse loses structure, the type system should track that — and ideally let the user prove preservation in special cases (e.g., elementwise operations preserve all sparsity patterns, masked operations preserve structured sparsity). This is a theorem-proving problem dressed as a type-checking problem.

3. Compiler-chosen physical layout. Given a logical sparsity refinement, the compiler should pick the storage format. Currently the user picks (TACO requires format specifications). With enough information, this is autoschedulable.

4. MoE-as-sparse-matmul as a language feature. MegaBlocks shows the implementation works; nobody has packaged it as a language construct (route, mux, or dispatch). Doing so would let the compiler generate the kernel automatically and compose with AD via the routing-as-effect framing.

5. Sparse + distributed. Sharding a sparse tensor across devices interacts non-trivially with the sparsity pattern. If the nonzeros aren't distributed evenly, you get load imbalance. If you partition by nonzero count, you lose the spatial structure that made the sparsity useful. This is open even for static patterns and wide open for dynamic ones.

6. Sparse + quantized. As above, but with different metadata (scales, zero-points). Production systems implement specific combinations (e.g., GPTQ + 2:4) as bespoke kernels. A compositional theory is missing.

7. Verification of sparse code. Sparse kernels are notoriously bug-prone (off-by-one in level traversal, missed cancellations, format conversion errors). Verified sparse code generation would be a real contribution. TACO has some correctness arguments; full verification is open. (This connects directly to memo 5.)

8. Sparsity-aware AD as a first-class transformation. Currently AD systems either drop sparsity (Zygote) or handle it with special cases (PyTorch's sparse autograd). A from-scratch design where AD propagates sparsity types through the backward pass is missing.

10. Recommendation for an AI-native language design

Concrete positions:

• Sparsity as a type refinement, not a separate type. A Tensor f16 [M, N] with refinement [Sparse Block 16, RowDistributed] is more flexible than a BlockSparseRowDistributedTensor f16 [M, N]. The refinement is a predicate; the compiler picks the storage layout.

• TACO-style level formats as the lower layer. The refinement layer is for users; the level-format layer is for the compiler to generate code. This is roughly what MLIR's sparse_tensor dialect does and is the right architecture.

• Routing as a first-class construct. route or dispatch operators with types that capture the sparse selection pattern. MegaBlocks-style code generation falls out automatically.

• Composition produces widened types unless proven otherwise. Sparse × sparse → dense by default; let the user prove (via refinement type checking) that special cases preserve structure.

• Sparsity-aware AD by construction. The AD transformation propagates sparsity refinements through the backward pass. A sparse loss has a sparse gradient; a routed forward has a routed backward with a chosen estimator.

• Hardware structured sparsity as a refinement, separate from format. 2:4Sparse is a property of a dense tensor's values, not a storage format. The compiler matches the property to the hardware capability.

• Quantization in the same framework. Quantization is compression with error bounds; the type carries the bounds, the compiler picks INT4 vs. FP8 vs. etc. based on the bounds.

The thesis: sparsity is currently fragmented because we lack a type-level abstraction over "structure on a tensor." TACO solved this for one specific kind of structure (storage format). Everything since has been point solutions for other kinds (hardware, routing, low-rank). A from-scratch language can unify them via refinement-typed tensors with a predicate-based structure language.

11. References

Sparse tensor compilation foundations

• Kjolstad, F. et al. (2017). The Tensor Algebra Compiler. OOPSLA 2017. — TACO. The starting point for compositional sparse code generation.
• Chou, S., Kjolstad, F. & Amarasinghe, S. (2018). Format Abstraction for Sparse Tensor Algebra Compilers. OOPSLA 2018. — Level-format theory.
• Senanayake, R. et al. (2020). A Sparse Iteration Space Transformation Framework for Sparse Tensor Algebra. OOPSLA 2020. — Loop transformations.
• Henry, R. et al. (2021). Compilation of sparse array programming models. OOPSLA 2021.
• Bik, A. et al. (2022). Compiler Support for Sparse Tensor Computations in MLIR. ACM TACO. — MLIR's sparse_tensor dialect.
• Ye, Z. et al. (2023). SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning. ASPLOS 2023.
• Ahrens, W. et al. (2023). Finch: Sparse and Structured Array Programming with Control Flow. — Sparse arrays integrated with general control flow.

MoE and dynamic sparsity

• Shazeer, N. et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017.
• Lepikhin, D. et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ICLR 2021.
• Fedus, W., Zoph, B. & Shazeer, N. (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR.
• Gale, T. et al. (2023). MegaBlocks: Efficient Sparse Training with Mixture-of-Experts. MLSys 2023. — MoE-as-block-sparse.
• Jiang, A. et al. (2024). Mixtral of Experts. arXiv:2401.04088.

Differentiable sparsity

• Martins, A. & Astudillo, R. (2016). From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification. ICML 2016.
• Peters, B., Niculae, V. & Martins, A. (2019). Sparse Sequence-to-Sequence Models. ACL 2019. (Entmax.)
• Correia, G. et al. (2019). Adaptively Sparse Transformers. EMNLP 2019.
• Bengio, Y. et al. (2013). Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. arXiv:1308.3432. (Straight-through.)

Sparse attention

• Child, R. et al. (2019). Generating Long Sequences with Sparse Transformers. arXiv:1904.10509.
• Beltagy, I., Peters, M. & Cohan, A. (2020). Longformer: The Long-Document Transformer. arXiv:2004.05150.
• Zaheer, M. et al. (2020). Big Bird: Transformers for Longer Sequences. NeurIPS 2020.
• Kitaev, N., Kaiser, Ł. & Levskaya, A. (2020). Reformer: The Efficient Transformer. ICLR 2020.

Low-rank methods

• Hu, E. et al. (2021). LoRA: Low-Rank Adaptation of Large Language Models. ICLR 2022.
• Liu, S. et al. (2024). DoRA: Weight-Decomposed Low-Rank Adaptation. ICML 2024.

Hardware structured sparsity

• Mishra, A. et al. (2021). Accelerating Sparse Deep Neural Networks. NVIDIA whitepaper. — Ampere 2:4 design.
• Zhou, A. et al. (2021). Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch. ICLR 2021.

Quantization (composes with sparsity)

• Frantar, E. et al. (2023). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers. ICLR 2023.
• Dettmers, T. et al. (2024). QLoRA: Efficient Finetuning of Quantized LLMs. NeurIPS 2023.

Adjacent — refinement types and structure

• Rondon, P., Kawaguchi, M. & Jhala, R. (2008). Liquid Types. PLDI 2008. — Refinement types in a mainstream-language setting.
• Vazou, N. et al. (2014). Refinement Types for Haskell. ICFP 2014.

12. Connections to the rest of the series

• Memo 1 (shapes): dynamic-shape MoE batches and ragged sparsity sit in the intersection of these two memos. Solving them requires both shape and sparsity to be expressible in the same type system. KV-cache compression is the canonical example.
• Memo 2 (AD-through-effects): gradient through routing is a stochastic-estimator handler. Sparsemax/entmax/straight-through are competing handlers. The framework from memo 2 directly applies.
• Memo 4 (distribution): sharding interacts with sparsity in ways nobody has worked out. Load imbalance in sharded MoE is the production pain point; expert parallelism is the current mitigation, and it is not principled.
• Memo 5 (verification): sparse code generation is bug-prone in ways that justify formal verification. TACO's correctness arguments are a starting point; full verification is open.

13. Next memo

Design Space 04: Distributed Programming Model. Modern training and inference run across hundreds to thousands of devices. The current model (MPI primitives, all-reduce, pipeline/tensor/data parallelism, ZeRO) is cargo-culted from HPC and bolted onto framework code. The question for an AI-native language: should sharding be in the type system, what does a mesh look like as a language construct, and how does this compose with shapes (memo 1), effects (memo 2), and sparsity (this memo)?