[Fix] Fix shared memory race in tilelang chunk_bwd dg_last accumulation

[Erix025]

Summary

• Fix shared memory race condition at fla/ops/common/backends/tilelang/chunk_bwd.py line 112 in [Bug] Non-atomic shared-memory accumulation causes nondeterministic dg gradients in tilelang chunk_bwd #889
• Add two regression tests:
  • test_chunk_bwd_dg_shared_mem_race: runs backward 10 times with g=0, asserts dg is bit-exact across runs
  • test_chunk_bwd_dg_accuracy_zero_gate: runs backward 100 times with g=0, checks worst-case dg error vs naive recurrent reference stays within 2%

Root Cause

In the V-loop (line 85), after two T.reduce_sum operations produce f_hdh_scalar[0], every thread in the thread block executes:

s_dg_last_acc[0] = s_dg_last_acc[0] + f_hdh_scalar[0]

With multiple warps executing concurrently, inter-warp timing causes non-deterministic read-modify-write results on shared memory.

Fix

Guard the shared memory write with T.Parallel(1), consistent with the initialization pattern at lines 80-81:

# Before (all threads write — race condition):
s_dg_last_acc[0] = s_dg_last_acc[0] + f_hdh_scalar[0]

# After (single-thread write):
for _i in T.Parallel(1):
    s_dg_last_acc[0] = s_dg_last_acc[0] + f_hdh_scalar[0]

[gemini-code-assist]

Code Review

This pull request addresses a shared memory race condition in the chunk_bwd kernel by wrapping the s_dg_last_acc accumulation in a T.Parallel(1) block and adds regression tests for gradient determinism. Feedback points out a potential visibility issue where a T.sync_threads() call is needed after the accumulation to ensure the updated value is correctly synchronized across all threads before subsequent reads.

[yzhangcs]

@Erix025 hey could you fix some conflicts

[zhiyuan1i]

These codes should not be changed and require further investigation into the underlying causes。 CC @yzhangcs

[tzj-fxz]

As I understand it, this kernel will contribute f_hdh_scalar to the final dg when USE_G is set. Since f_hdh_scalar has only one element and is summed by all elements in the buffer f_hdh, it may be helpful to change the f_hdh_scalar into local variable of each thread in the register rather than shared memory to avoid data race.

s_dg_last_acc = T.alloc_shared((1,), T.float32)
# can be changed to
s_dg_last_acc = T.alloc_var((1,), T.float32)

Similar to https://github.com/tile-ai/tilelang/blob/bcb2da336a28cf641f88e7f1ae53a95c475b6202/examples/gdn/example_chunk_o_bwd.py#L203, all references to s_dg_last_acc should remove [0] since it is no longer a shared buffer.

CC @Erix025