Skip to content

[5/6] feat(attention): sparse attention and SWA interfaces - #1029

Open
maocheng23 wants to merge 1 commit into
maocheng/pr1023-4-visibilityfrom
maocheng/pr1023-5-sparse-interfaces
Open

maocheng23 wants to merge 1 commit into
maocheng/pr1023-4-visibilityfrom
maocheng/pr1023-5-sparse-interfaces

Conversation

@maocheng23

@maocheng23 maocheng23 commented Oct 27, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

Interfaces for sparse attention, carried over from #1023 (author's commit 0e361e24). Nothing is executable yet. They define how sparse selections, KV sources and auxiliary objectives flow, before any DSA/QSA implementation lands.

Modifications

  • attention/selection.py:
    • IndexedSelection holds token or block indices per named KV source.
    • SlidingWindowSelection describes a local window.
    • SparseSelection is their union, intersected with global visibility.
  • attention/backends/sparse.py: the typed SparseAttentionBackend protocol and SparseAttentionResult (values plus per-head validity).
  • AttentionOutput.auxiliary_losses / AttentionAuxiliaryLoss: named, unweighted scalars, with the graph retained.
  • AttentionConfig reserves the qsa and dsa kinds; build_attention raises NotImplementedError for them.

Note: #1023's description says QSA/DSA are out of scope, but its commits 2–3 add them. They are split into this PR and the next one.

Accuracy Test

test_sparse_contract.py (from #1023). tests/test_modeling on CPU: 196 tests, 27 skipped. The only errors are the 4 sglang/yunchang import errors described in #1025.

Stack

This is part of a split of #1023 (by @jiapingW) into reviewable layers. Each PR targets the branch of the previous one, and only its own commits show in the diff.

  1. [1/6] refactor(draft): move draft layer components under modeling/draft/layers #1025: Move draft layer components (pure move + legacy-import shims)
  2. [2/6] refactor(draft): extract GQA/MLA attention into reusable layers #1026: GQA/MLA attention extraction (bit-exact parity test)
  3. [3/6] feat(draft): residual hooks with Qwen GR and DeepSeek-V4.1 mHC #1027: Residual hooks + GR + mHC
  4. [4/6] feat(attention): unify visibility across dense attention backends #1028: Visibility unified across dense backends (behavior change)
  5. [5/6] feat(attention): sparse attention and SWA interfaces #1029: Sparse attention / SWA interfaces 👈 this PR
  6. [6/6] feat(attention): DSA and QSA draft attention training #1030: DSA and QSA training

The code is @jiapingW's from #1023 (commits are authored by them). The split adds a parity test, legacy import shims, and two fixes in the visibility PR, each noted in the PR it belongs to.

🤖 Generated with Claude Code

Split from #1023 (originally 0e361e2). Interface-only; nothing executable yet:

- attention/selection.py: IndexedSelection (token/block indices per KV source),
  SlidingWindowSelection and their union SparseSelection
- attention/backends/sparse.py: typed SparseAttentionBackend protocol and
  SparseAttentionResult (values + per-head validity)
- AttentionOutput.auxiliary_losses / AttentionAuxiliaryLoss for indexer or
  router objectives
- AttentionConfig reserves the qsa/dsa kinds; build_attention raises
  NotImplementedError for them

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@maocheng23
maocheng23 force-pushed the maocheng/pr1023-4-visibility branch from 9c5967f to 4b48be0 Compare October 27, 2026 20:45
@maocheng23
maocheng23 force-pushed the maocheng/pr1023-5-sparse-interfaces branch from a935f93 to 282a4b5 Compare October 27, 2026 20:45

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants