[remote-controlled-man]

What does this PR do?

attach_layerwise_casting_hooks() breaks inference through torch.nn.MultiheadAttention: it casts out_proj to the storage dtype, but attention reads that submodule's weights directly without invoking its forward hook. A float32 attention input with FP16 or FP8 storage then raises a dtype mismatch in both the native and functional attention paths.

Treat the attention module as one casting unit, so all projection weights use the compute dtype during forward and return to the storage dtype afterwards. If a skip pattern or class selects a child such as out_proj, skip the whole attention unit to preserve the exclusion without mixing dtypes. The function docstring describes this behavior.

Regression coverage compares outputs and attention weights with an independently cast reference, repeats inference with/without attention weights, and covers FP16/FP8 storage, both input layouts, combined/separate projection weights, and parent/child skip patterns and classes.

Validation

All 8 core regression cases fail on unchanged upstream with a dtype mismatch. On final commit d164abd:

• python -m pytest -c pyproject.toml tests/test_hooks.py -q -rs: 24 passed, 3 skipped. The skips require multiple hardware accelerators or CUDA/XPU offload.
• Ruff 0.13.1: ruff check . and ruff format --check . pass (183 files formatted).
• The original minimal public-API reproducer also passes after the fix.

Validation ran on CPU with Python 3.12.14 and PyTorch 2.15.0. No GPU, distributed, or offload integration validation is claimed.

Before submitting

• Read the contribution guide.
• Added regression tests.
• Updated the public function docstring.
• Prior issue/forum discussion: none; directly reproduced bug, no overlapping issue or PR found in targeted searches.

AI assistance was used for investigation, implementation, tests, and this description.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.