[Feature] support reusable final_state buffer in fused recurrent GDN

benchmarks/ops/benchmark_fused_recurrent_final_state.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import argparse
+import statistics
+
+import torch
+import torch.nn.functional as F
+
+from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Benchmark reusable final_state for fused recurrent GDN.')
+    parser.add_argument('--B', type=int, default=1)
+    parser.add_argument('--T', type=int, default=1)
+    parser.add_argument('--H', type=int, default=32)
+    parser.add_argument('--HV', type=int, default=32)
+    parser.add_argument('--K', type=int, default=128)
+    parser.add_argument('--V', type=int, default=128)
+    parser.add_argument('--dtype', choices=('float16', 'bfloat16'), default='bfloat16')
+    parser.add_argument('--warmup', type=int, default=50)
+    parser.add_argument('--iters', type=int, default=1000)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError('CUDA is required for this benchmark.')
+    if args.HV % args.H != 0:
+        raise ValueError(f'HV ({args.HV}) must be divisible by H ({args.H}).')
+
+    device = torch.device('cuda')
+    dtype = getattr(torch, args.dtype)
+    torch.manual_seed(42)
+    q = torch.randn(args.B, args.T, args.H, args.K, dtype=dtype, device=device)
+    k = torch.randn_like(q)
+    v = torch.randn(args.B, args.T, args.HV, args.V, dtype=dtype, device=device)
+    beta = torch.sigmoid(torch.randn(args.B, args.T, args.HV, dtype=dtype, device=device))
+    g = F.logsigmoid(torch.randn(args.B, args.T, args.HV, dtype=torch.float32, device=device))
+    initial_state = torch.randn(args.B, args.HV, args.K, args.V, dtype=torch.float32, device=device)
+    final_state = torch.empty_like(initial_state)
+
+    def run(use_buffer):
+        kwargs = dict(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            initial_state=initial_state,
+            output_final_state=True,
+            use_qk_l2norm_in_kernel=True,
+        )
+        if use_buffer:
+            kwargs['final_state'] = final_state
+        return fused_recurrent_gated_delta_rule(**kwargs)
+
+    ref_o, ref_ht = run(False)
+    buf_o, buf_ht = run(True)
+    torch.testing.assert_close(ref_o, buf_o, rtol=0.002, atol=0.002)
+    torch.testing.assert_close(ref_ht, buf_ht, rtol=0.002, atol=0.002)
+    if buf_ht.data_ptr() != final_state.data_ptr():
+        raise AssertionError('final_state buffer was not reused.')
+
+    def measure(use_buffer):
+        for _ in range(args.warmup):
+            run(use_buffer)
+        torch.cuda.synchronize()
+        samples = []
+        for _ in range(5):
+            start = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
+            start.record()
+            for _ in range(args.iters):
+                run(use_buffer)
+            end.record()
+            end.synchronize()
+            samples.append(start.elapsed_time(end) / args.iters)
+        return samples
+
+    baseline = measure(False)
+    buffered = measure(True)
+    baseline_median = statistics.median(baseline)
+    buffered_median = statistics.median(buffered)
+    print(f'GPU: {torch.cuda.get_device_name()}')
+    print(f'shape: B={args.B}, T={args.T}, H={args.H}, HV={args.HV}, K={args.K}, V={args.V}, dtype={args.dtype}')
+    print(f'baseline ms/call: {baseline}')
+    print(f'buffered ms/call: {buffered}')
+    print(f'baseline median: {baseline_median:.6f} ms')
+    print(f'buffered median: {buffered_median:.6f} ms')
+    print(f'speedup: {baseline_median / buffered_median:.4f}x')
+    print(f'final_state data_ptr: {final_state.data_ptr()}')
+
+
+if __name__ == '__main__':
+    main()

benchmarks/ops/benchmark_fused_recurrent_final_state_autoregressive.py

@@ -0,0 +1,104 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import argparse
+import statistics
+
+import torch
+import torch.nn.functional as F
+
+from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Benchmark autoregressive final_state reuse for fused recurrent GDN.')
+    parser.add_argument('--B', type=int, default=8)
+    parser.add_argument('--H', type=int, default=32)
+    parser.add_argument('--HV', type=int, default=32)
+    parser.add_argument('--K', type=int, default=128)
+    parser.add_argument('--V', type=int, default=256)
+    parser.add_argument('--steps', type=int, default=1000)
+    parser.add_argument('--warmup', type=int, default=20)
+    parser.add_argument('--rounds', type=int, default=5)
+    parser.add_argument('--dtype', choices=('float16', 'bfloat16'), default='bfloat16')
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError('CUDA is required for this benchmark.')
+    if args.HV % args.H != 0:
+        raise ValueError(f'HV ({args.HV}) must be divisible by H ({args.H}).')
+
+    device = torch.device('cuda')
+    dtype = getattr(torch, args.dtype)
+    torch.manual_seed(42)
+    q_steps = [torch.randn(args.B, 1, args.H, args.K, dtype=dtype, device=device) for _ in range(args.steps)]
+    k_steps = [torch.randn_like(q) for q in q_steps]
+    v_steps = [torch.randn(args.B, 1, args.HV, args.V, dtype=dtype, device=device) for _ in range(args.steps)]
+    beta_steps = [torch.sigmoid(torch.randn(args.B, 1, args.HV, dtype=dtype, device=device)) for _ in range(args.steps)]
+    g_steps = [
+        F.logsigmoid(torch.randn(args.B, 1, args.HV, dtype=torch.float32, device=device))
+        for _ in range(args.steps)
+    ]
+    initial_state = torch.randn(args.B, args.HV, args.K, args.V, dtype=torch.float32, device=device)
+
+    def run_chain(use_buffer):
+        state = initial_state.clone()
+        for q, k, v, beta, g in zip(q_steps, k_steps, v_steps, beta_steps, g_steps):
+            kwargs = dict(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=state,
+                output_final_state=True,
+                use_qk_l2norm_in_kernel=True,
+            )
+            if use_buffer:
+                kwargs['final_state'] = state
+            _, state = fused_recurrent_gated_delta_rule(**kwargs)
+        return state
+
+    ref_state = run_chain(False)
+    buf_state = run_chain(True)
+    torch.testing.assert_close(ref_state, buf_state, rtol=0.002, atol=0.002)
+
+    def measure_interleaved():
+        samples = {False: [], True: []}
+        for round_idx in range(args.rounds):
+            order = (False, True) if round_idx % 2 == 0 else (True, False)
+            for use_buffer in order:
+                for _ in range(args.warmup):
+                    run_chain(use_buffer)
+                torch.cuda.synchronize()
+                start = torch.cuda.Event(enable_timing=True)
+                end = torch.cuda.Event(enable_timing=True)
+                start.record()
+                run_chain(use_buffer)
+                end.record()
+                end.synchronize()
+                samples[use_buffer].append(start.elapsed_time(end))
+        return samples[False], samples[True]
+
+    baseline, buffered = measure_interleaved()
+    baseline_median = statistics.median(baseline)
+    buffered_median = statistics.median(buffered)
+    print(f'GPU: {torch.cuda.get_device_name()}')
+    print(f'shape per step: B={args.B}, T=1, H={args.H}, HV={args.HV}, K={args.K}, V={args.V}, dtype={args.dtype}')
+    print(f'autoregressive steps: {args.steps}')
+    print(f'baseline total ms: {baseline}')
+    print(f'buffered total ms: {buffered}')
+    print(f'baseline median: {baseline_median:.6f} ms ({baseline_median / args.steps:.6f} ms/step)')
+    print(f'buffered median: {buffered_median:.6f} ms ({buffered_median / args.steps:.6f} ms/step)')
+    print(f'speedup: {baseline_median / buffered_median:.4f}x')
+
+
+if __name__ == '__main__':
+    main()

benchmarks/ops/issue_872_final_state_buffer.md

@@ -0,0 +1,128 @@
+# Issue #872: Reusable `final_state` Buffer
+
+## Summary
+
+`fused_recurrent_gated_delta_rule` now accepts an optional preallocated `final_state` buffer. When provided, the recurrent kernel writes the final state in-place and returns the same Tensor. Existing calls without a buffer keep the original allocation behavior.
+
+## Correctness
+
+Hardware: NVIDIA B200
+
+Command:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m pytest -q tests/ops/test_gdn.py
+```
+
+Result:
+
+```text
+80 passed, 20 skipped
+```
+
+The added coverage includes dense inputs, V-first state layout, variable-length inputs, repeated buffer reuse, initial/final state aliasing, NaN write coverage, and invalid buffer validation.
+
+Dependent layer/model test command:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m pytest -q \
+  tests/ops/test_gdn.py \
+  tests/layers/test_gated_deltanet.py \
+  tests/layers/test_layer_cache_layer_idx.py \
+  tests/models/test_modeling_gated_deltanet.py \
+  tests/models/test_modeling_gated_deltaproduct.py \
+  tests/models/test_modeling_mom.py \
+  tests/models/test_modeling_yoco.py
+```
+
+Result:
+
+```text
+130 passed, 29 skipped, 22 warnings
+```
+
+The warnings are existing TorchScript deprecations and layer configuration warnings; they are unrelated to the reusable buffer change.
+
+## Performance methodology
+
+- GPU: NVIDIA B200
+- dtype: `bfloat16`
+- mode: fused recurrent GDN
+- timing: CUDA events, median of repeated rounds
+- baseline: allocate `final_state` inside each call
+- buffered: reuse a preallocated `final_state`
+
+The single-call benchmark is implemented in `benchmark_fused_recurrent_final_state.py`. The autoregressive benchmark feeds each step's final state into the next step and is implemented in `benchmark_fused_recurrent_final_state_autoregressive.py`.
+
+## Single-call decode results
+
+Workload: `T=1, H=32, HV=32, K=128, V=256, dtype=bfloat16`, NVIDIA B200.
+
+| B | Baseline (ms) | Buffered (ms) | Speedup |
+|---:|---:|---:|---:|
+| 1 | 0.074139 | 0.071357 | 1.0390x |
+| 2 | 0.073129 | 0.071429 | 1.0238x |
+| 4 | 0.073106 | 0.071662 | 1.0201x |
+| 8 | 0.073223 | 0.070268 | 1.0421x |
+| 16 | 0.073739 | 0.071307 | 1.0341x |
+| 32 | 0.073152 | 0.071405 | 1.0245x |
+
+## Autoregressive results
+
+Workload: `B=8, T=1, H=32, HV=32, K=128, V=256, dtype=bfloat16`, NVIDIA B200.
+
+| Steps | Baseline total (ms) | Buffered total (ms) | Speedup |
+|---:|---:|---:|---:|
+| 1,000, run 1 | 77.253777 | 74.325905 | 1.0394x |
+| 1,000, run 2 | 75.828991 | 73.272305 | 1.0349x |
+
+## Final B-sweep
+
+Command:
+
+```bash
+for B in 1 2 4 8 16 32; do
+  CUDA_VISIBLE_DEVICES=0 \
+  python benchmarks/ops/benchmark_fused_recurrent_final_state_autoregressive.py \
+    --B $B --H 32 --HV 32 --K 128 --V 256 \
+    --dtype bfloat16 --steps 4096 --warmup 20 --rounds 5
+done
+```
+
+Result: 4096 autoregressive steps, five timing rounds per configuration.
+
+| B | Baseline ms/step | Buffered ms/step | Speedup |
+|---:|---:|---:|---:|
+| 1 | 0.075271 | 0.072985 | 1.0313x |
+| 2 | 0.073729 | 0.070771 | 1.0418x |
+| 4 | 0.074615 | 0.072983 | 1.0224x |
+| 8 | 0.073829 | 0.072487 | 1.0185x |
+| 16 | 0.074288 | 0.072566 | 1.0237x |
+| 32 | 0.079126 | 0.077348 | 1.0230x |
+
+The speedup range is `1.0185x` to `1.0418x`, with an equal-weight geometric mean of approximately `1.0268x` across the six batch sizes.
+
+## Conclusion
+
+The correctness suite passes, and autoregressive decode shows a consistent approximately 2–4% latency improvement across batch sizes from 1 to 32 on NVIDIA B200. Long prefill workloads are expected to show little benefit because the recurrent computation dominates the one-time output-buffer allocation.
+
+## PR summary
+
+### Summary
+
+Add an optional reusable `final_state` output buffer to fused recurrent GDN. The buffer is validated for shape, dtype, device, contiguity, and gradient safety, then written in-place by the existing recurrent kernel. Calls that do not provide a buffer remain backward compatible.
+
+### Test plan
+
+- `pytest -q tests/ops/test_gdn.py`: 80 passed, 20 skipped on NVIDIA B200.
+- Added dense, V-first, varlen, repeated reuse, aliasing, NaN write, and validation coverage.
+- Dependent layer/model tests should be run before opening the PR.
+
+### Benchmark
+
+- Hardware: NVIDIA B200.
+- Workload: autoregressive `T=1`, `B=1/2/4/8/16/32`, `H=HV=32`, `K=128`, `V=256`, `bfloat16`.
+- 4096-step autoregressive speedup: `1.0185x`–`1.0418x`, geometric mean approximately `1.0268x`.
+- Conclusion: consistent 2–4% decode improvement; no material benefit is expected for long prefill.
+
+No NCU profile was collected because the Triton kernel computation and launch configuration are unchanged; this change targets output-buffer allocation and reuse.