[Test] Add Pytest tests for cumsum_gdn and the fused sigmoid gating rule

cumsum_gdn parametrised over the forward and reverse scan, both layouts and the
fragment buffer. Sequence lengths that are not a multiple of the chunk (250
against 32) and odd head counts stay in: those are what exercise the tail.

The fused sigmoid gating rule takes its shapes as parameters and asserts the
output and the final state against golden(), so the test drives it directly on
the hundred alternating short sequences its guard uses. That guard wraps the
call in a thread pool, which only matters when several cases run at once and
there is one.

Author: dlaerbinjwy

examples/fused_sigmoid_gating_delta_rule/test_fused_sigmoid_gating_delta_rule_varlen.py

@@ -0,0 +1,33 @@
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _load_example() -> ModuleType:
+    source = Path(__file__).with_name("fused_sigmoid_gating_delta_rule_varlen.py")
+    spec = importlib.util.spec_from_file_location("_fused_sigmoid_gating_example_for_test", source)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Cannot load example module: {source}")
+
+    module = importlib.util.module_from_spec(spec)
+    original_argv = sys.argv
+    try:
+        sys.argv = [str(source)]
+        spec.loader.exec_module(module)
+    finally:
+        sys.argv = original_argv
+
+    return module
+
+
+# main() takes its shapes as parameters and asserts the output and the final
+# state against golden(), so this drives it directly. The example's own guard
+# wraps the same call in a thread pool, which only matters when several cases
+# run at once and there is one.
+def test_fused_sigmoid_gating_delta_rule_varlen() -> None:
+    example = _load_example()
+
+    # A hundred short sequences of alternating length, which is what makes this
+    # the varlen case rather than a padded batch.
+    example.main(seqlens=[4, 8] * 50, nk=16, nv=32, dk=128, dv=128)

examples_experiment/cumsum_gdn/test_example_cumsum.py

@@ -0,0 +1,77 @@
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+import pytest
+
+
+def _load_example() -> ModuleType:
+    source = Path(__file__).with_name("example_cumsum.py")
+    spec = importlib.util.spec_from_file_location("_cumsum_gdn_example_for_test", source)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Cannot load example module: {source}")
+
+    module = importlib.util.module_from_spec(spec)
+    original_argv = sys.argv
+    try:
+        sys.argv = [str(source)]
+        spec.loader.exec_module(module)
+    finally:
+        sys.argv = original_argv
+
+    return module
+
+
+# Shapes come from the example. Sequence lengths that are not a multiple of the
+# chunk (250 against 32) and head counts that are odd are what exercise the tail
+# handling, so both stay in.
+@pytest.mark.parametrize(
+    "batch, heads, seq_len, chunk, reverse, head_first",
+    [
+        (2, 32, 256, 32, False, True),
+        (2, 32, 256, 32, True, True),
+        (2, 7, 250, 32, False, False),
+        (2, 7, 250, 32, True, False),
+        (4, 8, 512, 64, True, True),
+    ],
+)
+def test_chunk_cumsum(batch, heads, seq_len, chunk, reverse, head_first) -> None:
+    import torch
+
+    example = _load_example()
+
+    shape = (batch, heads, seq_len) if head_first else (batch, seq_len, heads)
+    torch.manual_seed(0)
+    g = torch.randn(shape).npu().to(torch.float)
+
+    actual = example.chunk_cumsum(g, chunk, reverse=reverse, head_first=head_first)
+    expected = example.ref_chunk_cumsum(g, chunk, reverse=reverse, head_first=head_first)
+
+    torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=1e-5, atol=1e-5)
+
+
+# use_fragment selects a different buffer for the running sum; the reference is
+# the same either way.
+@pytest.mark.parametrize(
+    "batch, heads, seq_len, chunk, reverse, head_first",
+    [
+        (2, 32, 256, 32, False, True),
+        (2, 32, 256, 32, True, True),
+        (2, 7, 250, 32, False, False),
+        (1, 16, 128, 64, True, True),
+    ],
+)
+def test_chunk_cumsum_fragment(batch, heads, seq_len, chunk, reverse, head_first) -> None:
+    import torch
+
+    example = _load_example()
+
+    shape = (batch, heads, seq_len) if head_first else (batch, seq_len, heads)
+    torch.manual_seed(0)
+    g = torch.randn(shape).npu().to(torch.float)
+
+    actual = example.chunk_cumsum(g, chunk, reverse=reverse, head_first=head_first, use_fragment=True)
+    expected = example.ref_chunk_cumsum(g, chunk, reverse=reverse, head_first=head_first)
+
+    torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=1e-5, atol=1e-5)