• Mar. 16th, 2026: We released the Mixture-of-Depths Attention paper on arXiv. Code is coming soon.

• Release Mixture-of-Depths Attention (MoDA) paper on arXiv.
• Release MoDA Triton kernel and corresponding test units.
• Release Chunk-Visible MoDA Triton kernel and corresponding test units.
• Release Non-Causal MoDA Triton kernel and corresponding test units.
• Release full LLM training recipe and reproducible configs.
• Release full vision tasks training recipe, i.e., Classification on ImageNet.

Scaling depth is a key driver for large language models (LLMs). Yet, as LLMs become deeper, they often suffer from signal degradation: informative features formed in shallow layers are gradually diluted by repeated residual updates, making them harder to recover in deeper layers. We introduce mixture-of-depths attention (MoDA), a mechanism that allows each attention head to attend to sequence KV pairs at the current layer and depth KV pairs from preceding layers. We further describe a hardware-efficient algorithm for MoDA that resolves non-contiguous memory-access patterns, achieving 97.3% of FlashAttention-2's efficiency at a sequence length of 64K. Experiments on 1.5B-parameter models demonstrate that MoDA consistently outperforms strong baselines. Notably, it improves average perplexity by 0.2 across 10 validation benchmarks and increases average performance by 2.11% on 10 downstream tasks, with a negligible 3.7% FLOPs computational overhead. We also find that combining MoDA with post-norm yields better performance than using it with pre-norm. These results suggest that MoDA is a promising primitive for depth scaling.

Conceptual comparison of mechanisms that utilize the depth stream. (a) Depth Residual reads the current representation and writes back by addition. (b) Depth Dense reads a set of historical representations and linearly projects them back; it writes back by concatenation along depth. (c) Depth Attention uses attention to read historical depth KV pairs in a data-dependent way. (d) Mixture-of-Depths Attention (MoDA) combines depth attention with standard sequence attention and writes both the current layer's output and its KV pairs to depth streams for subsequent layers.

Left: Flash-compatible hardware-efficient MoDA achieves higher efficiency than torch-implemented MoDA. However, it keeps a depth KV cache of length T×L for each sequence, so each query potentially scans a long concatenated depth KV. Right: Chunk/Group-aware MoDA groups queries by chunk size C and reorganizes depth KV by chunk, reducing the effective depth span from T×L to (C×L)/G per chunk, where G is the GQA group number. This layout improves depth KV calculation efficiency and reduces memory access overhead.

Scaling Sequence Length T (G=8, Hq=64, Hk=8, L=64)

Scaling GQA Group Size G (T=16384, Hk=8, L=64)

Scaling Model Depth L (T=16384, G=8, Hq=64, Hk=8)

MoDA reaches 97.3% of FlashAttention-2 efficiency at a sequence length of 64K. Extra time consistently decreases as sequence length or group size increases.

MoDA attention heatmaps with the combined-softmax formulation. Columns correspond to uniformly sampled layers {0, 11, 23, 35}, and rows correspond to randomly selected heads in each layer. The first column shows attention over Sequence KV only, while the remaining columns show the concatenated Sequence KV | Depth KV; the red dashed line marks the boundary between the two KV blocks. Across layers and heads, substantial attention mass is consistently assigned to the Depth KV block, indicating that MoDA effectively leverages depth information in addition to standard sequence attention.

The following requirements should be satisfied:

• PyTorch >= 2.5
• Triton >= 3.0
• einops
• transformers >= 4.53.0
• datasets >= 3.3.0
• causal-conv1d >= 1.4.0

Install the local MoDA-enabled fla package from this repository:

cd libs/moda_triton
pip install -e .
cd ../..

Note: Please install from libs/moda_triton instead of PyPI, since the MoDA Triton kernels are maintained in this local directory.

python3 libs/moda_triton/fla/ops/moda/moda_v14.py

This project is based on OLMo2 (paper) and Flash Linear Attention (paper). Thanks for their wonderful works.

If you find MoDA useful in your research or applications, please consider giving us a star ⭐ and citing it by the following BibTeX entry.

@article{zhu2026moda,
  title   = {Mixture-of-Depths Attention},
  author  = {Zhu, Lianghui and Fang, Yuxin and Liao, Bencheng and Wang, Shijie and Cheng, Tianheng and Huang, Zilong and Chen, Chen and Wei, Lai and Zeng, Yutao and Wang, Ya and Lin, Yi and Li, Yu and Wang, Xinggang},
  journal = {arXiv preprint arXiv:2603.15619},
  year    = {2026}
}