Hold the global vector cumsum to the tolerance its example allows

CI reached 1.4e-5 on five of 262144 values against the 1e-5 this asserted. The
example holds three of its four variants to 1e-5 and this one to 1e-4, which the
test flattened to 1e-5 for all four.

The looser bound is the right one here: the other three accumulate within a
chunk while this one runs the length of the sequence, and the kernel and the
reference reach the same sum in a different order. The order depends on how the
work divides across cores, which is why 256 elements passed on one machine and
not on the other.

Author: dlaerbinjwy

examples_experiment/cumsum_kda/test_example_cumsum_kda.py

@@ -119,4 +119,9 @@ def test_chunk_global_cumsum_vector(batch, heads, seq_len, s_dim, reverse, head_
     actual = example.chunk_global_cumsum_vector(s, reverse=reverse, head_first=head_first)
     expected = example.ref_chunk_global_cumsum_vector(s, reverse=reverse, head_first=head_first)
 
-    torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=1e-5, atol=1e-5)
+    # 1e-4 rather than the 1e-5 the other three hold, which is what the example
+    # itself allows here: this variant accumulates across the whole sequence
+    # instead of within a chunk, and the kernel and the reference reach the same
+    # sum in a different order. At 256 elements that showed up as five values in
+    # 262144 landing 1.4e-5 apart.
+    torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=1e-4, atol=1e-4)