Fix missing cutlass.utils.ampere_helpers on arm64

[JannikSt]

Summary

Fix ModuleNotFoundError: No module named 'cutlass.utils.ampere_helpers' crash on arm64 (GB200) trainer startup.

Root Cause

• nvidia-cutlass-dsl 4.4.1 dropped ampere_helpers.py from cutlass.utils (now only ships hopper_helpers and blackwell_helpers)
• flash-attn 2.8.3's flash_attn.cute unconditionally imports it at module level for the Sm80 forward kernel (flash_attn.cute.flash_fwd -> import cutlass.utils.ampere_helpers)
• This import is NOT gated by GPU architecture, so it triggers even on GB200 (sm_100)
• The file still exists in flashinfer's vendored copy of cutlass

Fix

Copy ampere_helpers.py from flashinfer's vendored cutlass into the nvidia-cutlass-dsl package during the Docker build (arm64 only).

Can be removed once flash-attn gates the import by architecture or nvidia-cutlass-dsl re-adds the file.

---

Note

Low Risk
A build-time, arm64-gated Dockerfile workaround that only affects container image contents; main risk is brittleness if upstream package paths change.

Overview
Adds an arm64-only Docker build workaround that copies cutlass.utils.ampere_helpers.py from flashinfer’s vendored CUTLASS into the installed nvidia-cutlass-dsl package inside the venv.

This prevents flash-attn/flash_attn.cute from crashing at import time with ModuleNotFoundError on arm64/GB200 when nvidia-cutlass-dsl 4.4.1 omits that module.

^{Written by Cursor Bugbot for commit e4f5326. This will update automatically on new commits. Configure here.}