Add LigerAttnRes nn.Module wrapper for the AttnRes op#1293
Open
abtonmoy wants to merge 1 commit into
Open
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds the missing
LigerAttnResnn.Modulewrapper for the existing Attention Residuals(AttnRes) op, completing its artifact set.
The
attn_reskernel already ships atorch.autograd.Function(LigerAttnResFunction), afunctional entry point (
liger_attn_res), a test, and a benchmark — but unlike every other op(
LigerDyT,LigerPolyNorm,LigerRMSNorm, …) it has nonn.Moduleand is not exported fromliger_kernel.transformers. This PR adds that module so AttnRes can be dropped into a model likeany 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
Nblocks are each RMSNorm'd, scored against a learned pseudo-queryw_query, and theper-block softmax weights produce a weighted sum.
src/liger_kernel/transformers/attn_res.py— newLigerAttnRes(hidden_size, eps=1e-6, query_init_std=0.02). It owns the two learnable parameters (w_query, initializedrandn * 0.02;w_norm, initialized to ones) and callsLigerAttnResFunction. Follows theLigerDyTmodulepattern.
src/liger_kernel/transformers/__init__.py— registers the import (always-safe section) andadds
LigerAttnResto__all__.test/transformers/test_attn_res.py— adds module-level tests (see below).Two small notes:
Vas a list ofNblock tensors, butLigerAttnResFunction.forwardreadsV_stacked.shapebefore the stacking path inattn_res_forwardruns, so a list reaches.apply()and raisesAttributeError. The moduletherefore stacks a list/tuple itself before calling the Function, so
LigerAttnReshonors thedocumented list-or-
[N,B,T,D]contract. (Fixing the raw Function/functional to accept lists is aseparate change and is left out of this PR to keep it focused.)
w_query/w_normgradients are a sum-reduction over allN*B*Ttokens, which the backward kernel accumulates in fp32. Verified against an fp64 groundtruth, 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 (
LigerDyTand theother 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 parameterreceives a finite gradient, across fp32/fp16/bf16 and several shapes (incl. a non-power-of-two,
odd-shape case).
test_module_param_gradients_match_reference—w_query/w_norm/input gradients match thePyTorch reference in fp32 (rationale above).
test_module_list_input_and_repr— list input matches the stacked tensor;extra_reprreportsthe config.
All 39 tests in the file pass;
ruff checkandruff formatare clean.make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergence— not applicable; this adds a composable module wrapper around anexisting, already-tested kernel, with no change to any modeled layer's numerics.