Skip to content

Add LigerAttnRes nn.Module wrapper for the AttnRes op#1293

Open
abtonmoy wants to merge 1 commit into
linkedin:mainfrom
abtonmoy:attn-res-module
Open

Add LigerAttnRes nn.Module wrapper for the AttnRes op#1293
abtonmoy wants to merge 1 commit into
linkedin:mainfrom
abtonmoy:attn-res-module

Conversation

@abtonmoy

@abtonmoy abtonmoy commented Jul 7, 2026

Copy link
Copy Markdown

Summary

Adds the missing LigerAttnRes nn.Module wrapper for the existing Attention Residuals
(AttnRes) op, completing its artifact set.

The attn_res kernel already ships a torch.autograd.Function (LigerAttnResFunction), a
functional entry point (liger_attn_res), a test, and a benchmark — but unlike every other op
(LigerDyT, LigerPolyNorm, LigerRMSNorm, …) it has no nn.Module and is not exported from
liger_kernel.transformers. This PR adds that module so AttnRes can be dropped into a model like
any other Liger layer, holding its two learnable parameters.

Details

AttnRes (Kimi/Moonshot AI, arXiv:2603.15031) replaces the
standard residual connection with a softmax attention over the depth (block) dimension: the stacked
outputs of N blocks are each RMSNorm'd, scored against a learned pseudo-query w_query, and the
per-block softmax weights produce a weighted sum.

  • src/liger_kernel/transformers/attn_res.py — new LigerAttnRes(hidden_size, eps=1e-6, query_init_std=0.02). It owns the two learnable parameters (w_query, initialized randn * 0.02;
    w_norm, initialized to ones) and calls LigerAttnResFunction. Follows the LigerDyT module
    pattern.
  • src/liger_kernel/transformers/__init__.py — registers the import (always-safe section) and
    adds LigerAttnRes to __all__.
  • test/transformers/test_attn_res.py — adds module-level tests (see below).

Two small notes:

  • List inputs. The op's docstrings advertise passing V as a list of N block tensors, but
    LigerAttnResFunction.forward reads V_stacked.shape before the stacking path in
    attn_res_forward runs, so a list reaches .apply() and raises AttributeError. The module
    therefore stacks a list/tuple itself before calling the Function, so LigerAttnRes honors the
    documented list-or-[N,B,T,D] contract. (Fixing the raw Function/functional to accept lists is a
    separate change and is left out of this PR to keep it focused.)
  • Parameter-gradient accuracy. w_query/w_norm gradients are a sum-reduction over all
    N*B*T tokens, which the backward kernel accumulates in fp32. Verified against an fp64 ground
    truth, this makes the kernel's low-precision parameter gradients more accurate than a naive
    same-dtype PyTorch reference, so the parameter-gradient value test is asserted in fp32 (where the
    comparison is meaningful); the fp16/bf16 forward and input-gradient paths are covered separately.

No README change: the kernel table does not list composable building blocks (LigerDyT and the
other norm-style modules are likewise absent).

Testing Done

Added to test/transformers/test_attn_res.py (the existing Function/functional tests are unchanged):

  • test_module_matches_reference — module forward matches the PyTorch reference and every parameter
    receives a finite gradient, across fp32/fp16/bf16 and several shapes (incl. a non-power-of-two,
    odd-shape case).
  • test_module_param_gradients_match_referencew_query/w_norm/input gradients match the
    PyTorch reference in fp32 (rationale above).
  • test_module_list_input_and_repr — list input matches the stacked tensor; extra_repr reports
    the config.

All 39 tests in the file pass; ruff check and ruff format are clean.

  • Hardware Type: RTX 4060 Laptop GPU (8 GB), CUDA 12.6, Triton 3.x
  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence — not applicable; this adds a composable module wrapper around an
    existing, already-tested kernel, with no change to any modeled layer's numerics.

The attn_res kernel already ships a Function, a functional entry point,
a test, and a benchmark, but had no nn.Module wrapper and was not exported
from liger_kernel.transformers (unlike every other op). This adds the module
so Attention Residuals can be used like any other Liger layer.

- LigerAttnRes holds the learnable w_query/w_norm parameters and calls
  LigerAttnResFunction; follows the LigerDyT module pattern
- export it from liger_kernel.transformers (import + __all__)
- module accepts a list of blocks (stacked internally) or a stacked
  [N, B, T, D] tensor
- add module tests: forward parity vs the PyTorch reference, parameter and
  input gradient correctness, list-input equivalence, and extra_repr
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.

1 participant