[maocheng23]

Motivation

Make visibility mean the same thing on every dense backend. Each attention layer now enforces its own AttentionConfig (causality, sliding window) intersected with the algorithm's mask, using logical positions for packed layouts. Queries with no visible key produce zero output everywhere, not only on eager.

This is the behavior-changing part of #1023. It was split out so the move, refactor and residual PRs stay bit-exact. It overlaps #970 (draft causality decoupled from SWA layout), and the two should be reconciled.

Modifications (from #1023)

• backends/visibility.py: resolve_visibility intersects causal/window with an explicit mask for eager, SDPA (dense) and Flex (composed mask_mod). prepare_tensor_mask stabilizes rows that are fully masked.
• AttentionContext.query_positions/key_positions. DFlash passes its anchor-based positions during packed training.
• SDPA and Flash now zero empty queries after o_proj, as eager already did. Empty heads are zeroed before projection.

Changes relative to #1023

1. No per-layer mask rebuild for DFlash. With Refactor draft layers and add GR/mHC residuals #1023 as written, every sliding layer (and every layer when is_causal=True) rebuilt a mask identical to the one DFlash had already built: an extra create_block_mask per layer under Flex, or a dense [B,1,Q,K] tensor under SDPA. Measured on CPU SDPA with 2 sliding + 2 full layers, 2 of 4 layers rebuilt, and all rebuilt masks were bit-identical. Now AttentionContext.mask_encodes_visibility lets an algorithm declare that its mask already encodes the layer's rules, and DFlash sets it.
2. Flex requests the LSE only when o_proj has a bias. Flex already returns zeros for empty rows and heads; the LSE only exists to stop a bias from reviving them. The default bias-free path therefore keeps its previous kernel call.

Accuracy Test

• test_visibility.py (from Refactor draft layers and add GR/mHC residuals #1023), plus:
  • DFlash with sliding+full and is_causal=True layers calls resolve_visibility 0 times. Forcing the rebuild gives a bit-identical output.
  • Flex sets return_lse only with a bias.
  • A CUDA test checks that bias-free Flex still yields zero output and gradients for empty rows and heads.
• The parity test stays bit-exact except for the one intended change. Under SDPA, a dropped draft block no longer emits the o_proj bias, so the test now excludes those rows from the loss and comparison, the same way training excludes them.
• tests/test_modeling on CPU: 192 tests, 27 skipped. The only errors are the 4 sglang/yunchang import errors described in [1/6] refactor(draft): move draft layer components under modeling/draft/layers #1025. CUDA-only tests (Flex/Flash) were not run locally and will run in CI.

Not measured

No GPU throughput A/B yet. The default DFlash path changes as follows, and it is worth checking on the Qwen3.8 SWA recipe:

• SDPA computes per-head validity (mask.any(-1)) per layer.
• Flex with a bias computes the LSE.

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