1. Why this PR exists

Unsloth fast paths currently use standard residual updates around attention blocks, e.g.:

This is stable and efficient, but all prior layer information is mixed with fixed weights (implicit weight = 1 for the direct residual stream at each step). The AttnRes idea is to make residual mixing content-adaptive while preserving the same forward structure and default behavior.

Reference paper:

• Attention Residuals: https://arxiv.org/pdf/2603.15031

This PR introduces an opt-in, low-risk integration that keeps Unsloth architecture and coding style intact.

2. High-level design

The integration goal is:

• Do not break existing kernels / fast paths.
• Keep default behavior identical when AttnRes is disabled.
• Touch only attention-branch outputs before existing residual merges.

Implemented approach:

1. Build lightweight AttnRes state per forward pass.
2. Track block-local states and block summaries.
3. Before residual add, compute an AttnRes contribution from prior states.
4. Add that contribution to the attention output, then continue existing residual code.

So if current attention output is (a_l), the transformed output is:

with:

where in this practical implementation:

• query (q_l): derived from current residual/query tensor,
• keys/values ((k_i, v_i)): prior block summaries + in-block previous states,
• (\alpha): configurable residual-mixing scale.

3. Mathematical mapping to implementation

The paper discusses full and block residual attention. This PR implements a practical block-style variant:

• Maintain two sets:
  • completed block summaries (S_1, \dots, S_{m-1})
  • in-progress block states in current block (h_{m,1}, \dots, h_{m,t-1})
• Candidate set for layer (l):

• Residual attention mix:

• Output transform:

• Block summary update at boundary:

4. What was changed

New files

• unsloth/models/attnres.py
  • forward-state init
  • block summary accumulation
  • attention-output transform
• unsloth/utils/attnres.py
  • utility-level AttnRes state/config/hook helpers

Updated files

• unsloth/utils/attention_dispatch.py
  • optional AttnRes-aware post-attention finalization hook path
• unsloth/utils/__init__.py
• unsloth/kernels/__init__.py
  • export AttnRes helpers

Model integrations (attention branch only)

• unsloth/models/llama.py
• unsloth/models/gemma2.py
• unsloth/models/cohere.py
• unsloth/models/granite.py
• unsloth/models/falcon_h1.py

These sites now call the AttnRes transform immediately after attention output and before existing residual merges.

5. Why this is maintainer-friendly

• Backward-compatible by default: if AttnRes is off, behavior remains unchanged.
• Minimal surface area: logic is centralized in models/attnres.py and reused.
• No kernel rewrites required: existing fast attention paths remain intact.
• Per-architecture safe integration: only attention branch is modified; non-attention math is unchanged.
• Reviewable diff: most edits are one-step hooks at known residual merge points.

6. Safety and risk controls

Implemented controls:

• Opt-in behavior only.
• Guard against risky stateful contexts (e.g., gradient-checkpointing replay desync).
• Preserve dtype/device behavior from existing paths.
• Keep all legacy residual-add code paths (e.g., +=, torch.add) untouched except attention pre-transform.

Known scope limits in this PR:

• This is not a full reproduction of paper-level distributed systems features (e.g., cross-stage caching infrastructure, full two-phase inference scheduler stack).
• This PR focuses on architecture-level integration inside current Unsloth fast-path constraints.

7. How to review quickly

1. Start with unsloth/models/attnres.py for core logic.
2. Check one model integration (llama.py) for pattern.
3. Confirm same pattern replicated in gemma2/cohere/granite/falcon_h1.
4. Verify default-off behavior and unchanged residual merge code.
5. Validate utility exports and dispatch hooks.

8. Validation performed

• Static compile check on changed Python files:
  • python -m py_compile ... (pass)

Suggested follow-up validation (maintainer side):

• A/B forward parity when AttnRes disabled.
• Small training smoke run with AttnRes enabled.
• Throughput/memory check on at least one Llama-family and one non-Llama-family model.

9. Expected impact

When enabled, this should allow attention outputs to incorporate dynamically weighted prior layer information (via block summaries + in-block history) while preserving Unsloth’s fast-path structure and compatibility defaults.