Add LigerAttnRes nn.Module wrapper for the AttnRes op

src/liger_kernel/transformers/__init__.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING
 
 # Always-safe imports (independent of 'transformers')
+from liger_kernel.transformers.attn_res import LigerAttnRes  # noqa: F401
 from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss  # noqa: F401
 from liger_kernel.transformers.dyt import LigerDyT  # noqa: F401
 from liger_kernel.transformers.fused_add_rms_norm import LigerFusedAddRMSNorm  # noqa: F401
@@ -172,6 +173,7 @@ def __getattr__(name: str):
 # Shared symbols in all environments
 __all__ = [
     "is_transformers_available",
+    "LigerAttnRes",
     "LigerCrossEntropyLoss",
     "LigerDyT",
     "LigerFusedLinearCrossEntropyLoss",

src/liger_kernel/transformers/attn_res.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops import LigerAttnResFunction
+
+
+class LigerAttnRes(nn.Module):
+    """Attention Residuals (AttnRes) from Kimi/Moonshot AI (arXiv:2603.15031).
+
+    Replaces the standard residual connection ``h = h_prev + f(RMSNorm(h_prev))``
+    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, and
+    the resulting per-block softmax weights produce a weighted sum. This is the
+    module wrapper around :func:`~liger_kernel.transformers.functional.liger_attn_res`.
+
+    Args:
+        hidden_size: hidden dimension ``D`` of each block output.
+        eps: epsilon for the per-block RMSNorm (default: 1e-6).
+        query_init_std: standard deviation of the normal initialization for the
+            learned pseudo-query ``w_query`` (default: 0.02). ``w_norm`` (the
+            per-block RMSNorm weight) is initialized to ones.
+
+    Shape:
+        - Input: ``[N, B, T, D]`` stacked block outputs, or a list of ``N``
+          tensors each of shape ``[B, T, D]``.
+        - Output: ``[B, T, D]``.
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6, query_init_std: float = 0.02):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.eps = eps
+        self.query_init_std = query_init_std
+        self.w_query = nn.Parameter(torch.randn(hidden_size) * query_init_std)
+        self.w_norm = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, V):
+        if isinstance(V, (list, tuple)):
+            V = torch.stack(V)
+        return LigerAttnResFunction.apply(V, self.w_query, self.w_norm, self.eps)
+
+    def extra_repr(self):
+        return f"hidden_size={self.hidden_size}, eps={self.eps}"

test/transformers/test_attn_res.py

@@ -8,6 +8,7 @@
 from test.utils import supports_bfloat16
 
 from liger_kernel.ops import LigerAttnResFunction
+from liger_kernel.transformers.attn_res import LigerAttnRes
 from liger_kernel.transformers.functional import liger_attn_res
 from liger_kernel.utils import infer_device
 
@@ -185,3 +186,99 @@ def test_correctness_functional(N, B, T, D, dtype, atol, rtol):
     y2.backward(grad)
 
     assert_verbose_allclose(V1.grad, V2.grad, atol=atol, rtol=rtol)
+
+
+@pytest.mark.flaky(reruns=3, reruns_delay=2)
+@pytest.mark.parametrize(
+    "N, B, T, D",
+    [
+        (4, 2, 64, 512),
+        (8, 2, 32, 256),
+        # weird shapes
+        (3, 5, 37, 123),
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype, atol, rtol",
+    [
+        (torch.float32, 1e-4, 1e-5),
+        (torch.float16, 1e-2, 1e-3),
+        pytest.param(
+            torch.bfloat16,
+            1e-1,
+            1e-2,
+            marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
+        ),
+    ],
+)
+def test_module_matches_reference(N, B, T, D, dtype, atol, rtol):
+    """LigerAttnRes module matches the PyTorch reference (fwd) and trains its params (bwd)."""
+    set_seed(0)
+    model = LigerAttnRes(hidden_size=D, eps=1e-6).to(device).to(dtype)
+
+    V = torch.randn(N, B, T, D, device=device, dtype=dtype)
+    V_in = V.clone().requires_grad_(True)
+    out = model(V_in)
+    ref = pytorch_attn_res(V, model.w_query.detach(), model.w_norm.detach(), eps=1e-6)
+    assert out.shape == (B, T, D)
+    assert_verbose_allclose(out, ref, atol=atol, rtol=rtol)
+
+    out.backward(torch.randn_like(out))
+    assert V_in.grad is not None and torch.isfinite(V_in.grad).all()
+    for name, p in model.named_parameters():
+        assert p.grad is not None and torch.isfinite(p.grad).all(), f"no/invalid grad for {name}"
+
+
+@pytest.mark.flaky(reruns=3, reruns_delay=2)
+@pytest.mark.parametrize(
+    "N, B, T, D",
+    [
+        (4, 2, 64, 512),
+        (3, 5, 37, 123),
+    ],
+)
+def test_module_param_gradients_match_reference(N, B, T, D):
+    """The learned params (w_query, w_norm) must receive gradients that match the
+    PyTorch reference — the whole point of the module is to train them.
+
+    Checked in fp32 only: w_query/w_norm gradients are a sum-reduction over all
+    ``N*B*T`` tokens, and the kernel accumulates that reduction in fp32 while a
+    same-dtype PyTorch reference does not, so in fp16/bf16 the reference is the
+    less-accurate baseline (verified against an fp64 ground truth) and an
+    element-wise comparison would test reduction noise, not correctness. The
+    low-precision forward and input-grad paths are covered by the tests above.
+    """
+    set_seed(0)
+    atol, rtol = 1e-3, 1e-4
+    model = LigerAttnRes(hidden_size=D, eps=1e-6).to(device)
+
+    V = torch.randn(N, B, T, D, device=device)
+    do = torch.randn(B, T, D, device=device)
+
+    # Reference: same parameter values, plain PyTorch autograd.
+    V_ref = V.clone().requires_grad_(True)
+    wq_ref = model.w_query.detach().clone().requires_grad_(True)
+    wn_ref = model.w_norm.detach().clone().requires_grad_(True)
+    pytorch_attn_res(V_ref, wq_ref, wn_ref, eps=1e-6).backward(do)
+
+    # Module (Triton kernel) path.
+    V_mod = V.clone().requires_grad_(True)
+    model(V_mod).backward(do)
+
+    assert_verbose_allclose(V_mod.grad, V_ref.grad, atol=atol, rtol=rtol)
+    assert_verbose_allclose(model.w_query.grad, wq_ref.grad, atol=atol, rtol=rtol)
+    assert_verbose_allclose(model.w_norm.grad, wn_ref.grad, atol=atol, rtol=rtol)
+
+
+def test_module_list_input_and_repr():
+    """Module accepts a list of blocks (equivalent to the stacked tensor) and reprs its config."""
+    set_seed(0)
+    N, B, T, D = 4, 2, 16, 64
+    model = LigerAttnRes(hidden_size=D).to(device)
+    blocks = [torch.randn(B, T, D, device=device) for _ in range(N)]
+
+    out_list = model(blocks)
+    out_stacked = model(torch.stack(blocks))
+    assert out_list.shape == (B, T, D)
+    assert_verbose_allclose(out_list, out_stacked, atol=1e-6, rtol=1e-6)
+    assert f"hidden_size={D}" in model.extra_repr()