[Feature]: Remove the FA4 install workaround after the PyTorch 26.05 upgrade

🚀 The feature, motivation and pitch

With the move to the NVIDIA PyTorch 26.05 container, revisit the workaround in docker/common/install_fa4.sh and install FlashAttention 4 through the normal requirements flow.

Current main already uses nvcr.io/nvidia/pytorch:26.05-py3 and includes flash-attn-4==4.0.0b19 in requirements.txt. However, docker/Dockerfile.multi still invokes install_fa4.sh, which independently pins and installs the same package.

The remaining cleanup is to rely on the requirements entry and remove the dedicated script and its Dockerfile invocation, once the original installation issue is confirmed resolved on 26.05. This would keep the FA4 version in one place and simplify container dependency management.

Alternatives

Keep the dedicated installation step if the original import issue still occurs on 26.05, and document the remaining blocker.

Additional context

The workaround was introduced in #14706 because FA4 installed by a non-root user landed in ~/.local, separate from the base image's system-installed flash_attn package. The FA4 wheel lacked a top-level flash_attn/__init__.py, so flash_attn.cute could not be imported. Installing FA4 as root in the devel image co-located the packages.

Before removing the workaround, validate that:

• The normal requirements installation on 26.05 makes flash_attn.cute importable, including the non-root installation path that originally failed.
• A VisualGen FA4 smoke test passes without the dedicated script.
• Supported devel and release container workflows still provide FA4 where needed.

This is a cleanup proposal based on code inspection; the installation behavior on 26.05 has not been tested here.

References:

• PyTorch base image
• Existing FA4 requirement
• Dedicated FA4 install script
• Dockerfile invocation