[WIP] SM120 fused backward kernel: exploratory shared-memory reduction

[minatoyukinaa]

⚠️ Draft — do NOT merge before tilelang is upgraded on main

This PR is developed against tilelang 0.1.13 (latest). Earlier tilelang versions have suboptimal shared-memory planning, tilelang==0.1.9 will have more shared memory use. Please keep this as a draft until the main branch upgrades tilelang.This is also published as an explored solution rather than a finished one: it is not yet fully implemented, and no speedup has been measured yet.。 And h_dtype is assert in fp16...

Summary

Adds the SM120 (sm_120) fused backward kernel (fused_gdr_bwd) for gated delta rule chunked computation, ported/adapted from the Hopper fused_bwd. The main focus of this PR is an aggressive reduction of shared-memory usage through three techniques.

Shared-memory reduction strategies

1.Reuse of q_shared / tmp_shared — buffer aliasing as previously discussed in #30 and #21.
2. Merge the S and K consumers — done to make strategy 3 feasible. (It is still unclear whether strategy 3 could be achieved without merging S and K; feedback welcome.)
3. Halve the DK dimension of h_shared and tmp_shared_4_1 — each original gemm() over DK is split into two half-DK computations:
tmp_shared_4_1 is the easy case: it simply caches dk, which is straightforward to handle once the S and K consumers are merged.

h_shared is much more involved: since only half of h fits at a time, a dedicated producer streams DK-halves from global memory while the consumer overlaps compute on the other half. This is coordinated with a new barrier set bar_h0..bar_h7: even-indexed barriers signal "producer has finished loading this window", odd-indexed barriers signal "consumer has finished, producer may reload".
Side effect of layout inference: some transfers could no longer be expressed as a simple T.copy() and had to be rewritten as explicitly unrolled loops.
Current status
✅ Compiles successfully; shared-memory usage reduced as intended
❌ Deadlocks at runtime — suspected to be caused by the complex producer/consumer synchronization scheme (bar_h* barriers)
❌ No correctness validation (pytest) passed yet
❌ No benchmark numbers yet
TODO
1.Debug the deadlock。
Pass tests/test_gdr_unit.py
2.Benchmark.
3.Evaluate whether strategy 3 can be applied without merging S and K consumers
4.Wait for tilelang ≥ 0.1.13 on main before merging
I will add and organize the comparison between the pipeline after merging S and K and the previous Hopper version later.

Since the bwd operator is quite complex, and debugging on TileLang is also difficult, we need to proceed carefully with the development and design. Any suggestions would be very helpful to me.

[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

[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.