feat: add blackwell sm120 forward support

[minatoyukinaa]

SM120 (Blackwell) Forward Pass Support for FlashQLA

---

It implements forward-pass inference for FlashQLA on SM120 GPUs (RTX 5090 / RTX Pro 6000). issue #2 The primary challenge is that SM120 provides only ~100 KB of usable shared memory, compared to
228 KB on SM90. This PR targets a fixed configuration of chunk_size=32 and block_DV=64, and includes a new kkt_solve kernel specialized for chunk_size=32.

---

✅ Completed

• Forward pass passes accuracy checks(see below comment) .
• Benchmark results available (see below comment).
• New blackwell_sm120/ kernel directory with:
  • kkt_solve.py — 32×32 block-inversion kernel (specialized for chunk_size=32)
  • fused_fwd.py — fused GDR forward kernel(chunk_size=32 and blockDV=64)

💬 Discussion

1. Enabling Auto-CP in Forward Pass

The bottleneck is prepare_h. We observed that m_shared_L and m_shared_R can potentially reuse the memory of **h_shared**: h_shared finishes computation at barrier 1, while m_shared_L/R are only needed after
barrier 2. On SM90 this was infeasible because T.gemm(A, B[0:DK]) with sliced indexing triggered compilation errors. However, on SM120 this pattern is supported. Once this reuse is implemented, shared memory
usage should drop below 100 KB.

2. Reducing Shared Memory in Backward Pass by Reusing Buffers

q_shared ↔ tmp_shared_2_1 reuse (validated on H20 SM90):

By moving T.copy(q_shared, odot_fragment_2) (line 662 of fused_bwd.py) to before the write of dV' into tmp_shared_2_1, we can let tmp_shared_2_1 take over q_shared's role for both prefetching and
producer-side reads. This eliminates the standalone q_shared allocation. At fp16 + block_S=64, this saves ~16 KB of shared memory (block_S × DK × 2B / 1024).

k_shared / do_shared reuse — difficult:

Unlike q_shared, there is no intermediate fragment (odot_fragment_2) that can recover the values of k and do after their shared memory is overwritten. The conflict arises from the usage of W and U, making
straightforward reuse infeasible.

Remaining memory budget (with block_S=32):

Even with successful q_shared elimination, the total shared memory footprint still exceeds 100 KB:

┌───────────────────────────────────────────┬────────┐
│ Variable │ Size │
├───────────────────────────────────────────┼────────┤
│ h_shared + tmp_shared_4_1 (128×128, fp32) │ 64 KB │
├───────────────────────────────────────────┼────────┤
│ tmp_shared_2_x (3 × block_S × DK × 2B) │ 24 KB │
├───────────────────────────────────────────┼────────┤
│ tmp_shared_1_x (3 × block_S × block_S) │ 6 KB │
├───────────────────────────────────────────┼────────┤
│ q, k, v, do (4 × block_S × DK × 2B) │ 32 KB │
├───────────────────────────────────────────┼────────┤
│ dqkv_shared │ 8 KB │
├───────────────────────────────────────────┼────────┤
│ a_shared │ 2 KB │
├───────────────────────────────────────────┼────────┤
│ Total │ 136 KB │
└───────────────────────────────────────────┴────────┘

Reducing block_S to 16 would bring it under 100 KB, but breaks correctness in kkt_solve (16×16 matrix is too small for the block-inversion algorithm). The remaining options are further precision adjustments
or deeper buffer reuse.

---

[Starmys]

Thank you for your contribution! Using chunk_size=32 to address the shared memory limitation is a truly brilliant design.

[minatoyukinaa]

---

Summary

This PR rebases onto the latest main (previously built on library v0.1.10), resolves merge conflicts, and introduces a shared memory optimization for prepare_h on the SM120 (Blackwell consumer) architecture.

SMEM Optimization for prepare_h

To reduce shared memory usage, h_shared is reused to replace both m_shared_L and m_shared_R.

Motivation

• The shape of h_shared is (DK, DV). In the current implementation, DV = DK = 128 is asserted. The shape of m_shared_L / m_shared_R is (DK, DK // 2), and both use the qkva_dtype.
• By analyzing lifetimes: h_shared finishes its usage at bar_1, while m_shared_L / m_shared_R are not needed until after bar_2. This makes memory reuse possible.
• After applying this optimization, shared memory usage drops below 100 KB, allowing compilation to succeed.

Why This Optimization Does Not Apply to Hopper

When attempting the same approach on H20 with tilelang v0.1.8, using indexed slicing in T.gemm(A, B[0:DK//2]) fails — the wgmma instruction does not support indexed addressing. On SM120 (consumer-grade Blackwell), the
mma instruction is used instead. Notably, SM120 lacks tmem and does not use SM100's tcgen05.mma instruction.

Code Changes

h_shared[:, 0:DK//2] replaces m_shared_L, and h_shared[:, DK//2:] replaces m_shared_R.

Hopper path (unchanged):
T.copy(m_fragment_L, m_shared_L)

SM120 path (new):
T.barrier_wait(bar_2, i_s % 2)
T.copy(m_fragment_L, h_shared[:, :DK // 2])

A T.barrier_wait(bar_2, ...) is added because the last producer warp group performs store_h — we must wait for store_h to complete before reusing h_shared. Correspondingly, T.barrier_arrive(bar_2) is emitted after
store_h. Since the last producer warp group has 32 threads, the barrier count increases: 384 + 32 = 416.

• bar_2 = T.alloc_barrier(arrive_count=384) → bar_2 = T.alloc_barrier(arrive_count=416)
• T.copy(h_shared, h) → T.copy(h_shared, h) + T.barrier_arrive(bar_2)

The same pattern applies when replacing m_shared_R.

Testing

Benchmark

python benchmark/bench_gated_delta_rule.py

Forward-pass benchmark results on NVIDIA GeForce RTX 5090 are available in benchmark/benchmark_results_5090.txt. FlashQLA achieves up to ~2x speedup over FLA and up to ~2x over FlashInfer across Qwen3.5 family model
configurations.

Unit Tests

python -m pytest tests/test_gdr_unit.py -k "fwd and not mixed_cp" -v

All 56 forward-only tests pass. Mixed CP tests are skipped because they include backward pass, which is not yet implemented for SM120.

Notes

• FlashInfer compatibility: profiling FI performance requires flashinfer-python >= 0.6.14 and CUDA > 12.8, which may break PyTorch version compatibility.
• Backward pass: cp_bwd.py and fused_bwd.py are not yet implemented for SM120. They are currently copied from the Hopper implementation as placeholders to keep function signatures correct. The SM120 implementation uses
  chunk_size=32; if the backward pass is accidentally triggered, the Hopper chunk_size=64 assertion will catch it.
• kkt_solve performance will be further optimized in a future PR.

---

cc @Starmys

[Starmys]

Can we remove the chunk_size here and apply the default chunk_size for different devices?

[minatoyukinaa]

I’ve removed the hard‑coded chunk_size values (64 and 32) from the forward passes, and added a helper _get_default_chunk_size() that returns a device‑dependent default:

For SM120 (compute version "12.0") → 32

For all other supported devices (SM90, SM100) → 64

[Starmys]

To protect users from calling an unsupported backward path, we’d suggest dropping the SM120 backward kernels and mark backward pass as not implemented for SM120

[minatoyukinaa]

Summary

Drop unsupported SM120 backward kernels and mark backward pass as not implemented
for SM120. This prevents users from silently running incorrect backward
computations on Blackwell SM120 (compute capability 12.0) GPUs.

Changes

Backward disabled for SM120

File	Change
chunk/__init__.py	SM120 branch: removed fused_gdr_bwd / fused_gdr_dh imports, set both to None. Added NotImplementedError guard in chunk_gated_delta_rule_bwd.
chunk/cp_context.py	SM120 branch: removed cp_bwd import, fused_gdr_dh = None. Added guard in intra_card_cp_preprocess_bwd.
blackwell_sm120/__init__.py	Removed fused_gdr_bwd from imports and __all__.
blackwell_sm120/fused_bwd.py	Deleted — unsupported backward kernel.
blackwell_sm120/cp_bwd.py	Deleted — unsupported CP backward kernel.

Supporting changes

• tests/conftest.py: Added SM120 arch detection and requires_sm120 marker.
• pytest.ini: Registered the sm120 marker.

Note

test_bwd_works_on_non_sm120 was not run on real SM90 or SM100 hardware due to
environment constraints. Before merging, please confirm that the existing
backward tests in test_gdr_unit.py still pass on both architectures.