SM120/121 : Add support for fused backward kernel

[minatoyukinaa]

Summary

Adds the SM120/SM121 (sm_120) fused backward kernel (fused_gdr_bwd) for the chunked gated delta rule.

The implementation largely preserves the Hopper pipeline. The main difference is a shared-memory reduction strategy designed for SM120's ~99 KB shared-memory limit. Only the minimum required barriers are added to
synchronise the new half-state HBM reload path, avoiding the deadlocks and data races encountered in earlier versions.

---

Design and explored approaches

A direct Hopper port exceeds the SM120 shared-memory limit by approximately 37 KB.

The following approaches were explored:

• Buffer aliasing: reuse q_shared and temporary buffers as discussed in [RFC][SM120] Strategy for Implementing Backward Pass for sm120 #30 and feat: add blackwell sm120 forward support #21. This reduces shared memory but is insufficient on its own.
• Merge the S/K consumers: intended to keep the state entirely in registers and eliminate tmp_shared_4_1, but was rejected because of incompatible TileLang layouts, high register pressure, and fragile
  synchronisation.
• Selected approach: preserve the Hopper-style producer/consumer pipeline while halving the DK dimension of h_shared and tmp_shared_4_1. One half remains in shared memory, while the other half is loaded from HBM
  on demand and processed using split GEMMs.
• Remove v_shared: read v directly from HBM to further reduce shared-memory usage.

This design introduces some additional HBM traffic, but fits within the SM120 shared-memory limit while retaining the original pipeline structure and stable numerical behaviour.

---

Current status

• ✅ Full forward and backward support on SM120/SM121.
• ✅ Fixed-length and variable-length inputs.
• ✅ KV/VK state layouts and initial/final states.
• ✅ Autograd and intra-card CP, including mixed forward/backward CP settings.
• ✅ Tested with tilelang==0.1.13.
• ✅ tests/test_gdr_unit.py: 120 passed, 0 skipped.
• ✅ Backward reaches up to 2.89× speedup over FLA in the current RTX 5090 benchmarks.

Backward performance has not yet been fully optimised. The current implementation prioritises correctness, stability, and fitting within the SM120 shared-memory limit. Further pipeline scheduling, HBM overlap, and
configuration-specific tuning will be explored in future work.

[minatoyukinaa]

I found serious problem in producer and prefretch logic in pipeline, I will fix them soon. I realize, once we split h_shared into half,we only prefretch half of the window. We do some T.gemm(xxx,
h_share[:,0:dk//2 ] ). And We can rearange the half of gemm(T.gemm[dk//2:] in other operartion to overlap compute and memory access.
The reason of deadlock is casued by bar_h1, for example. h1 has cosumer S+K+A threads, this will immediately block the right part of gemm([dk//2:])
And I found merge S ＆ K is Non-essential，the key is how to duel the layout and half of h_shared and tmp_share41
update 8.28
finally，we found dh_tmp （half）can be reused in tmp-4-1。

[Starmys]

Found numeric errors on Hopper with tilelang==0.1.13. Will fix when I have some spare time.

[minatoyukinaa]

For now, I've implemented a solution that avoids merging the S and K paths. However, please note that I have not yet resolved the layout conflict problem—the current progress is blocked because the compiler fails to find a suitable layout inference for the generated code.

The core idea is as follows:
fused_bwd (2).py

Split the DK dimension of h_shared and tmp_shared_4_1 into two halves.

The remaining half is reloaded from HBM on the fly.

Since we always keep one half resident in shared memory, for every GEMM operation that involves both tmp_shared_4_1 and h_shared, an extra HBM copy and an additional GEMM are performed.

To reduce overhead, I've merged some adjacent h_shared accesses (around the bar_08 synchronization points) and added new barrier states to coordinate the extra data movement.
Update 8.21
i nearly finish fragment confilct，bwd is coming soon。