[NVIDIA][CuTe,Fwd,sm120] Fix use_tma_O crash on SM120 (issue #2649)

[johnnynunez]

Summary

Fixes #2649. On SM120 (Blackwell GeForce / RTX PRO 6000 / DGX Spark) the CuTe-DSL forward kernel is currently unusable — every call crashes with:

AttributeError: 'NoneType' object has no attribute '_trait'
  in cpasync.tma_partition (flash_fwd.py epilogue O-store)

Root cause

FlashAttentionForwardSm120 inherits the forward __call__, which sets:

self.use_tma_O = self.arch >= Arch.sm_90

This enables the TMA-based O-store epilogue for SM120. But SM120 does not build the TMA store atom (tma_atom_O is None), because it lacks the WGMMA-era TMA-store epilogue path. The first tma_partition on the null atom raises the AttributeError.

Fix

Restrict the TMA O-store to sm_90 .. sm_119:

self.use_tma_O = Arch.sm_90 <= self.arch < Arch.sm_120

SM120 then falls back to the register→gmem O store already used by the SM80 path, which is what the CpAsync SM120 kernel expects.

Testing

Verified on RTX PRO 6000 Blackwell (sm_120, cc 12.0), torch 2.12.0+cu130, nvidia-cutlass-dsl 4.5.2:

• Before: every SM120 forward call raises the AttributeError above.
• After: forward runs and matches PyTorch SDPA reference (bf16), max abs error:
  • hdim=64 non-causal: 1.3e-3
  • hdim=96 non-causal: 1.1e-3
  • hdim=128 causal: 8.1e-3

One-line behavioral change; no effect on sm80/sm90/sm100 paths.

[thad0ctor]

#2634 disables the broken SM80-base TMA-O path for SM120, though with a broader self.use_tma_O = False approach in the SM80-base path.