Import of nonexistent config module crashes fp8_fa3/setup.py

[tatavishnurao]

Description

Missing config module breaks fp8_fa3/setup.py

What happened

I was trying to import setup_fp8_fa3 directly and hit this:

ModuleNotFoundError: No module named 'torchao.prototype.attention.config'

Turns out torchao/prototype/attention/fp8_fa3/setup.py is importing from a module that doesn't exist:

from torchao.prototype.attention.config import LowPrecisionAttentionConfig

There is no config.py (or config/ package) anywhere under torchao/prototype/attention/. The class LowPrecisionAttentionConfig was never defined.

Impact

setup_fp8_fa3 is completely broken -- you can't even import it. The function signature references a type that doesn't exist:

def setup_fp8_fa3(
    model: nn.Module,
    config: LowPrecisionAttentionConfig,  # <-- undefined
) -> nn.Module:

However, the public API (apply_low_precision_attention) is fine. It calls shared_utils/setup.py:setup_fp8_backend() directly, passing a plain string "FA3" for the backend name, so the broken function is never reached through normal usage.

Affected code

File	Line	Problem
fp8_fa3/setup.py	11	from torchao.prototype.attention.config import LowPrecisionAttentionConfig -- module doesn't exist
fp8_fa3/setup.py	17	config: LowPrecisionAttentionConfig parameter type is undefined

How to fix

Two reasonable options:

1. Create the missing config module

Add torchao/prototype/attention/config.py defining LowPrecisionAttentionConfig and have setup_fp8_fa3 translate the config into setup_fp8_backend arguments. This was probably the original intent -- the config was planned but never implemented.

or

2. Remove setup_fp8_fa3 entirely

Since apply_low_precision_attention already handles the FA3 path through setup_fp8_backend, and setup_fp8_fa3 is the only consumer of this nonexistent config, we can just remove setup_fp8_fa3 and stop exporting it from fp8_fa3/__init__.py.

1 makes sense if there's a plan to add more backend-specific configuration later.

Steps to reproduce

## Steps to reproduce


from torchao.prototype.attention.fp8_fa3.setup import setup_fp8_fa3
# ModuleNotFoundError: No module named 'torchao.prototype.attention.config'

Environment

- torch: 2.11.0+cpu
- torchao: 0.18.0+gitde873d15
- Hardware: RTX 4060 
- OS: CachyOS Linux 7.0.0-1 (x86_64)