Add Kimi AttentionResiduals (AttnRes) kernelFeature/add attn res kernel (#1161)

## Summary

Implements **Attention Residuals (AttnRes)** from Kimi/Moonshot AI
([arxiv.org/abs/2603.15031](https://arxiv.org/abs/2603.15031)).

  This PR addresses issue #1158.

  ### What is AttnRes?

AttnRes replaces standard residual connections with softmax attention
over depth blocks to solve the **PreNorm dilution** problem where deep
layer contributions get diluted:

  ```python
  V = stack(blocks)           # [N, B, T, D]
  K = RMSNorm(V)              # per-block normalize
  scores = einsum(w, K)       # [N, B, T] — w is [D] learned query
  alpha = softmax(scores, 0)  # over block dim
  h = einsum(alpha, V)        # [B, T, D] — weighted sum 
```
###  Implementation

  - Single fused Triton kernel: RMSNorm + dot product + softmax + weighted sum
  - Efficient memory usage: Scores stored in registers (supports N≤32 blocks)
  - Complete autograd support: Forward and backward kernels with @ensure_contiguous decorator
  - Benchmark script: Compares against PyTorch and torch.compile

###  Files Added

  - src/liger_kernel/ops/attn_res.py - Core kernel implementation (318 lines)
  - benchmark/scripts/benchmark_attn_res.py - Benchmark and correctness tests (246 lines)
  - Updated src/liger_kernel/ops/__init__.py to export LigerAttnResFunction

###  Test Plan

  - Run correctness tests: python benchmark/scripts/benchmark_attn_res.py --quick
  - Run full benchmark: python benchmark/scripts/benchmark_attn_res.py
  - Test on RTX 5090 (Blackwell architecture)
  - Verify forward pass correctness (fp16, bf16, fp32)
  - Verify backward pass correctness

  ### Benchmark Results

  Tested on: **NVIDIA GeForce RTX 5090**, CUDA 12.8

  > Note: Due to resource constraints, testing was only performed on RTX 5090. Additional testing on datacenter GPUs (A100, H100) would be valuable to validate performance across different architectures. Maintainers are welcome to run benchmarks on other hardware configurations.
  To reproduce: `python benchmark/scripts/benchmark_attn_res.py`

  #### Forward Pass Performance

  | Config | PyTorch | torch.compile | Liger AttnRes | Speedup vs PyTorch | Speedup vs compile |
  |--------|---------|---------------|---------------|--------------------|--------------------|
  | N=4, D=4096, fp16 | 5.164 ms | 0.691 ms | **0.206 ms** | **25.11x** ✨ | 3.35x |
  | N=8, D=4096, fp16 | 10.011 ms | 2.716 ms | **0.394 ms** | **25.39x** ✨ | 6.89x |
  | N=8, D=8192, fp16 | 20.076 ms | 3.183 ms | **0.780 ms** | **25.74x** ✨ | 4.08x |
  | N=16, D=4096, fp16 | 19.946 ms | 2.996 ms | **1.004 ms** | **19.86x** ✨ | 2.98x |
  | N=8, D=4096, bf16 | 10.009 ms | 1.596 ms | **0.393 ms** | **25.47x** ✨ | 4.06x |

  #### Forward + Backward Performance

  | Config | PyTorch | Liger AttnRes | Speedup |
  |--------|---------|---------------|---------|
  | N=4, D=4096, fp16 | 20.525 ms | **1.003 ms** | **20.46x** ✨ |
  | N=8, D=4096, fp16 | 39.880 ms | **1.747 ms** | **22.83x** ✨ |
  | N=8, D=8192, fp16 | 79.723 ms | **3.317 ms** | **24.03x** ✨ |
  | N=16, D=4096, fp16 | 80.480 ms | **3.555 ms** | **22.64x** ✨ |
  | N=8, D=4096, bf16 | 39.888 ms | **1.744 ms** | **22.88x** ✨ |

  #### Visualizations

  **Speed - Full (Forward + Backward):**
  <img width="1000" alt="attn_res_speed_full" src="https://github.com/user-attachments/assets/a1bb7710-1d37-4824-8d30-79fd9e14d774" />

  **Speed - Forward:**
  <img width="1000" alt="attn_res_speed_forward" src="https://github.com/user-attachments/assets/a5649c68-fbf4-4d8e-9cb6-8a119ca6f8a3" />

  **Speed - Backward:**
  <img width="1000" alt="attn_res_speed_backward" src="https://github.com/user-attachments/assets/78161703-fea0-405f-a2ec-ae5be44f72ae" />

  **Memory - Full:**
  <img width="1000" alt="attn_res_memory_full" src="https://github.com/user-attachments/assets/d902bf36-ff65-4657-a471-8e92524261e0" />

  #### Correctness Tests

  All tests passed with expected numerical precision:
  - Forward pass: max diff < 2e-3 (fp16/bf16), < 2e-6 (fp32)
  - Backward pass: max diff < 4e-3 (fp16/bf16), < 2e-6 (fp32)

###  Key Insights

  1. Exceptional speedup: 20-25x faster than PyTorch's einsum-based implementation
  2. Beats torch.compile: 3-7x faster than torch.compile in forward pass
  3. Scales well: Performance maintained across different N (number of blocks) and D (hidden dimension)
  4. Memory efficient: Single-pass fused kernel minimizes memory traffic

###  The dramatic speedup is achieved by:
  - Fusing RMSNorm + attention + weighted sum into a single kernel
  - Storing attention scores in registers (no global memory roundtrip)
  - Optimized memory access patterns for coalesced reads/writes

  ---
  Closes #1158

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: kirsten-1

benchmark/scripts/benchmark_attn_res.py

@@ -0,0 +1,129 @@
+"""
+AttnRes Benchmark: Liger (Triton) vs PyTorch
+
+Kimi Attention Residuals: softmax attention over depth blocks.
+"""
+
+import math
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+
+from benchmark_model_configs import compute_seq_len_sweep_config
+from benchmark_model_configs import estimate_kernel_peak_memory
+from benchmark_model_configs import get_benchmark_model_config
+from utils import SingleBenchmarkRunInput
+from utils import SingleBenchmarkRunOutput
+from utils import parse_benchmark_script_args
+from utils import run_benchmarks
+from utils import run_memory_benchmark
+from utils import run_speed_benchmark
+
+from liger_kernel.ops.attn_res import LigerAttnResFunction
+from liger_kernel.utils import infer_device
+
+device = infer_device()
+
+
+def _setup_attn_res(input: SingleBenchmarkRunInput):
+    """Create input tensors for AttnRes from benchmark config."""
+    cfg = input.extra_benchmark_config
+    seq_len = input.x
+
+    # V: [N, B, T, D]
+    V = torch.randn(
+        cfg["N"],
+        cfg["bsz"],
+        seq_len,
+        cfg["hidden_size"],
+        device=device,
+        dtype=cfg["dtype"],
+        requires_grad=True,
+    )
+    w_query = torch.randn(cfg["hidden_size"], device=device, dtype=cfg["dtype"]) * 0.02
+    w_norm = torch.ones(cfg["hidden_size"], device=device, dtype=cfg["dtype"])
+    eps = cfg.get("eps", 1e-6)
+
+    if input.kernel_provider == "liger":
+        fn = lambda: LigerAttnResFunction.apply(V, w_query, w_norm, eps)
+    elif input.kernel_provider == "pytorch":
+        from test.transformers.test_attn_res import pytorch_attn_res
+
+        fn = lambda: pytorch_attn_res(V, w_query, w_norm, eps)
+    else:
+        raise ValueError(f"Invalid provider: {input.kernel_provider}")
+
+    return V, fn
+
+
+def bench_speed_attn_res(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
+    V, fn = _setup_attn_res(input)
+    return run_speed_benchmark(fn, input.kernel_operation_mode, [V])
+
+
+def bench_memory_attn_res(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
+    V, fn = _setup_attn_res(input)
+    return run_memory_benchmark(fn, input.kernel_operation_mode)
+
+
+if __name__ == "__main__":
+    args = parse_benchmark_script_args()
+
+    model = get_benchmark_model_config(args.model)
+    probe_seq_len = 1024
+
+    def _probe():
+        probe_input = SingleBenchmarkRunInput(
+            x=probe_seq_len,
+            kernel_provider="pytorch",
+            extra_benchmark_config={
+                "N": 8,
+                "bsz": 1,
+                "hidden_size": model.hidden_size,
+                "dtype": model.dtype,
+                "eps": 1e-6,
+            },
+        )
+        V, fn = _setup_attn_res(probe_input)
+        return fn()
+
+    peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
+    kernel_bpt = peak_bytes // probe_seq_len
+
+    config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt)
+
+    common_configs = {
+        "kernel_name": "attn_res",
+        "x_name": "T",
+        "x_label": "sequence length",
+        "x_values": [2**i for i in range(10, int(math.log2(config.seq_len)) + 1)],
+        "kernel_providers": ["liger", "pytorch"],
+        "extra_benchmark_configs": [
+            {
+                "N": 8,
+                "bsz": config.batch_size,
+                "hidden_size": model.hidden_size,
+                "dtype": model.dtype,
+                "eps": 1e-6,
+            }
+        ],
+        "overwrite": args.overwrite,
+    }
+
+    run_benchmarks(
+        bench_test_fn=bench_speed_attn_res,
+        kernel_operation_modes=["full", "forward", "backward"],
+        metric_name="speed",
+        metric_unit="ms",
+        **common_configs,
+    )
+    run_benchmarks(
+        bench_test_fn=bench_memory_attn_res,
+        kernel_operation_modes=["full", "forward", "backward"],
+        metric_name="memory",
+        metric_unit="MB",
+        **common_configs,
+    )

src/liger_kernel/ops/__init__.py

@@ -30,6 +30,9 @@
 # All of these can be replaced by vendor-specific implementations.
 # =============================================================================
 
+from liger_kernel.ops.attn_res import LigerAttnResFunction  # noqa: F401
+from liger_kernel.ops.attn_res import attn_res_backward  # noqa: F401
+from liger_kernel.ops.attn_res import attn_res_forward  # noqa: F401
 from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction  # noqa: F401
 from liger_kernel.ops.cross_entropy import cross_entropy_backward  # noqa: F401
 from liger_kernel.ops.cross_entropy import cross_entropy_forward  # noqa: F401