MISA: Смесь разреженного внимания на основе индексатора для инференции LLM с длинным контекстом

MISA: Mixture of Indexer Sparse Attention for Long-Context LLM Inference

May 8, 2026

Авторы: Ruijie Zhou, Fanxu Meng, Yufei Xu, Tongxuan Liu, Guangming Lu, Muhan Zhang, Wenjie Pei

cs.AI

Аннотация

Разреженное внимание DeepSeek (DSA) устанавливает новый уровень качества для мелкозернистого разреженного внимания на этапе инференса, вводя обучаемый потокеновый индексатор, который оценивает каждый префиксный токен и выбирает наиболее релевантные для основного внимания. Для сохранения выразительности индексатор использует множество голов запросов (например, 64 в DeepSeek-V3.2), которые разделяют один и тот же выбранный набор токенов; именно эта многоголовочная конструкция делает индексатор доминирующей вычислительной нагрузкой на длинных контекстах. Мы предлагаем MISA (Смесь индексаторов разреженного внимания) — прямую замену индексатора DSA, которая рассматривает его головы индексатора как пул смеси экспертов. Легковесный маршрутизатор использует дешёвую статистику на уровне блоков, чтобы выбрать зависимое от запроса подмножество лишь из нескольких активных голов, и только эти головы выполняют тяжелое выставление оценок на уровне токенов. Это сохраняет разнообразие исходного пула индексаторов, одновременно снижая стоимость на запрос: вместо оценки каждого префиксного токена каждой головой она производится лишь несколькими маршрутизированными головами, плюс пренебрежимо малый член маршрутизатора, вычисляемый на небольшом наборе пулированных ключей. Мы также представляем иерархический вариант MISA, который использует маршрутизированный проход для сохранения расширенного набора кандидатов, а затем переранжирует его с помощью исходного индексатора DSA, чтобы восстановить итоговые выбранные токены почти в точности. Используя всего восемь активных голов и без дополнительного обучения, MISA соответствует плотному индексатору DSA на наборе LongBench как для DeepSeek-V3.2, так и для GLM-5, работая при этом с восемью и четырьмя раза меньшим количеством голов индексатора соответственно, и в среднем превосходит HISA. Он также сохраняет полностью зелёные тепловые карты «Иголка в стоге сена» вплоть до контекста длиной 128K токенов и восстанавливает более 92% токенов, выбранных индексатором DSA на каждый слой. Наше ядро TileLang обеспечивает ускорение примерно в 3,82 раза по сравнению с исходным ядром индексатора DSA на одном графическом процессоре NVIDIA H200.

English

DeepSeek Sparse Attention (DSA) sets the state of the art for fine-grained inference-time sparse attention by introducing a learned token-wise indexer that scores every prefix token and selects the most relevant ones for the main attention. To remain expressive, the indexer uses many query heads (for example, 64 on DeepSeek-V3.2) that share the same selected token set; this multi-head design is precisely what makes the indexer the dominant cost on long contexts. We propose MISA (Mixture of Indexer Sparse Attention), a drop-in replacement for the DSA indexer that treats its indexer heads as a pool of mixture-of-experts. A lightweight router uses cheap block-level statistics to pick a query-dependent subset of only a few active heads, and only those heads run the heavy token-level scoring. This preserves the diversity of the original indexer pool while reducing the per-query cost from scoring every prefix token with every head to scoring it with only a handful of routed heads, plus a negligible router term computed on a small set of pooled keys. We further introduce a hierarchical variant of MISA that uses the routed pass to keep an enlarged candidate set and then re-ranks it with the original DSA indexer to recover the final selected tokens almost exactly. With only eight active heads and no additional training, MISA matches the dense DSA indexer on LongBench across DeepSeek-V3.2 and GLM-5 while running with eight and four times fewer indexer heads respectively, and outperforms HISA on average. It also preserves fully green Needle-in-a-Haystack heatmaps up to a 128K-token context and recovers more than 92% of the tokens selected by the DSA indexer per layer. Our TileLang kernel delivers roughly a 3.82 times speedup over DSA's original indexer kernel on a single NVIDIA H200 GPU.