ToSA: Trainable Online Sparse Attention

Processing long sequences is expensive because dense self-attention has quadratic complexity. Existing sparse attention methods, e.g., Native Sparse Attention, can skip score computation, but prediction and indexing introduce additional overhead. Online block skipping avoids a separate selection stage during inference, but high sparsity often reduces model accuracy. To address this issue, we propose ToSA, a trainable online sparse attention method for long-sequence modeling. ToSA reuses online-softmax statistics in a fused Triton kernel to identify low-contribution KV blocks after the attention scores are computed. For each skipped block, it bypasses value loading, softmax computation, and the subsequent P V accumulation. Its sparse backward kernel replays the block decisions recorded in the forward pass, enabling end-to-end training. We further propose a three-stage schedule that begins with dense training, gradually increases spar- sity, and ends with training at the desired sparsity. We evaluate ToSA with Qwen3.5-2B on long-context and long-video understanding. Compared with the dense fine-tuning baseline, ToSA improves the LongBench score from 38.18 to 39.62 and the LongVideoBench score from 58.19 to 60.02. In terms of efficiency, ToSA achieves up to a 1.77× attention-kernel speedup at a 128K sequence length. Together, these results demonstrate that ToSA substantially accelerates long-sequence modeling while preserving long-sequence understanding.