Add Optional Attention Residuals (AttnRes) Integration for Unsloth Fast Paths

[danielhanchen]

Staging mirror of unslothai#4863

Original PR: unslothai#4863
Author: @kleeedolinux

This is a staging copy for review and editing. Once finalized, changes will be pushed back to the original PR.

---

Original description

1. Why this PR exists

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

$$x_{l+1} = x_l + f_l(x_l)$$

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:

$$\tilde{a}_l = a_l + \alpha \cdot \sum_i w_{l,i} v_i$$

with:

$$w_{l,i} = \mathrm{softmax}\left(\frac{q_l \cdot k_i}{\sqrt{d}}\right)$$

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):

$$\mathcal{C}_l = \{S_1, \dots, S_{m-1}\} \cup \{h_{m,1}, \dots, h_{m,t-1}\}$$

• Residual attention mix:

$$r_l = \sum_{c \in \mathcal{C}_l} \mathrm{softmax}\left(\frac{q_l \cdot c}{\sqrt{d}}\right) c$$

• Output transform:

$$\tilde{a}_l = a_l + \alpha r_l$$

• Block summary update at boundary:

$$S_m = \sum_{j \in B_m} h_j$$

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

[danielhanchen]

/gemini review

[danielhanchen]

@codex review

[gemini-code-assist]

Code Review

This pull request introduces 'Attention Residuals' (AttnRes), a feature designed to accumulate and mix residual states across transformer layers to potentially improve model performance or stability. The implementation includes a core state management module in unsloth/models/attnres.py, integration into several model architectures (Llama, Gemma2, Cohere, Falcon, and Granite), and hooks within the central attention dispatch utility to support various backends. The review feedback suggests minor logic simplifications in the new modules, specifically regarding how the number of layers is determined, how model boundary conditions are checked, and the redundancy of state initialization logic.

[gemini-code-assist]

The logic for determining num_layers is redundant. getattr(config, "num_hidden_layers", 0) already returns 0 if the attribute is missing, and the subsequent check if num_layers <= 0 and hasattr(model, "layers") is sufficient. The explicit int() cast is also unnecessary as num_hidden_layers is typically an integer in Hugging Face configs.

Suggested change

	num_layers = int(getattr(config, "num_hidden_layers", 0))
	if num_layers <= 0 and hasattr(model, "layers"):
	num_layers = len(model.layers)
	num_layers = getattr(config, "num_hidden_layers", 0)
	if num_layers <= 0 and hasattr(model, "layers"):

[gemini-code-assist]

The condition at_model_end can be simplified. Since layer_idx is 0-indexed, layer_idx + 1 is the 1-based layer count. Comparing layer_idx + 1 == attnres_state.num_layers is sufficient.

Suggested change

	at_block_end = ((layer_idx + 1) % attnres_state.block_size) == 0
	at_model_end = (attnres_state.num_layers > 0) and (
	(layer_idx + 1) >= attnres_state.num_layers
	)
	at_model_end = (attnres_state.num_layers > 0) and ((layer_idx + 1) == attnres_state.num_layers)

[gemini-code-assist]

The enabled property logic in __post_init__ is redundant because begin() performs the same check. Consider simplifying the state initialization.

Suggested change

	self.enabled = bool(
	self.enabled or self.config.enabled or self.kernel_hook is not None
	)
	def __post_init__(self):
	if self.kernel_hook is None:
	self.kernel_hook = self.config.kernel_hook
	self.enabled = self.enabled or self.config.enabled or self.kernel_hook is not None

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 513888151c

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Normalize candidate devices before stacking AttnRes states

AttnRes currently stacks completed_block_summaries and current_block_states as-is, but multi-GPU pipeline inference moves activations per layer (move_to_device in the inference loops) while reusing one shared attnres_state. As soon as layers run on different _per_layer_device_index values, this list contains tensors from different GPUs and torch.stack fails with a cross-device error. The candidates should be moved to query.device (or tracked per device) before mixing.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Preserve attention output dtype when applying AttnRes mix

attention_output can be fp16/bf16, but the mix is computed from query and may be fp32 (for example, fast inference paths pass a float32 residual buffer as query). The expression attention_output + (alpha * mixed) then promotes the result to fp32, which pushes subsequent layers off the intended low-precision fast path and increases memory/latency whenever AttnRes is enabled. Cast mixed (or query) to attention_output.dtype before the add.

Useful? React with 👍 / 👎.