[maocheng23]

Motivation

This PR makes DFlash-family attention reusable: GQA/MLA become components that any draft can build. Behavior must not change.

Modifications

• layers/attention/base.py:
  • The Attention[State] contract, AttentionContext (position embeddings, mask, KV prefix, cache position) and AttentionOutput.
  • DenseAttention, which holds the shared execution formerly in Qwen3DFlashAttentionBase.forward: KV-cache update, backend dispatch, zeroing of empty queries, and o_proj.
• layers/attention/{gqa,mla,rotary}.py: the GQA/MHA and expanded-MLA projections and RoPE, moved out of dflash.py.
• layers/attention/backends/dense.py: eager/SDPA/Flash/Flex dispatch, unchanged from main. Masks are used as supplied, and eager's boolean→additive conversion is kept.
• layers/config.py (LayerKernels, AttentionConfig), layers/factory.py (build_attention, build_ffn) and layers/decoder.py (a pre-norm DecoderLayer).
• DFlash, Domino and DSpark use DecoderLayer. DFlash2 keeps its convolution placement and calls attention through the new context.
• Compatibility kept:
  • Qwen3DFlashAttention and Qwen3DFlashMLAAttention remain as compatibility constructors.
  • DFlashKernels is an alias of LayerKernels.
  • Checkpoint keys and parameter order are unchanged.

Differences from the corresponding part of #1023: the mask-handling changes (causal/window applied by logical position, SDPA/Flash zeroing of empty queries, Flex LSE) are moved to the visibility PR. This PR's numerics are therefore exactly main's.

Accuracy Test

tests/test_modeling/test_layers/test_refactor_parity.py embeds the pre-refactor Qwen3DFlashAttentionBase, GQA/MLA attention and DFlash/DFlash2 decoder forwards verbatim from e3ab28d. It swaps them into a copy of each model and requires zero-tolerance equality on:

• outputs, gradients and two AdamW steps, for:
  • DFlash/DFlash2/DSpark/Domino
  • GQA and MLA
  • eager and SDPA
  • full, sliding+full and is_causal=True layouts
  • attention_bias=True, with a dropped draft block
• cached block decoding, with two consecutive blocks on a DynamicCache

A mutation check confirms the test isn't vacuous. Each of these fails it:

• scaling GQA V by (1+1e-6)
• scaling MLA Q by (1+1e-6)
• dropping eager's empty-row zeroing

tests/test_modeling on CPU: 169 tests, 23 skipped. The only errors are the 4 sglang/yunchang import errors described in the move PR.

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) 👈 this PR
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
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