Skip to content

Move hip_fmha op schemas next to their implementations#245

Closed
bottler wants to merge 1 commit intometa-pytorch:mainfrom
bottler:export-D97933992
Closed

Move hip_fmha op schemas next to their implementations#245
bottler wants to merge 1 commit intometa-pytorch:mainfrom
bottler:export-D97933992

Conversation

@bottler
Copy link
Copy Markdown
Contributor

@bottler bottler commented Mar 24, 2026

Summary:
When MSLK_BUILD_HIP_FMHA=0, attention.cpp was still compiled into the main library and registered schemas for efficient_attention_forward_ck, efficient_attention_backward_ck, and _ck_rand_uniform via TORCH_LIBRARY_FRAGMENT, even though the TORCH_LIBRARY_IMPL bindings and kernel implementations (in mslk_hip_fmha) were absent. This left unimplemented operators registered in the library — the op appears in the dispatcher but calling it fails.

Fix by moving each m.def into the same file as its TORCH_LIBRARY_IMPL, inside hip_fmha/. Since those files are only compiled as part of the mslk_hip_fmha static library, schema and implementation now come and go together. The decoder ops remain in attention.cpp since their situation differs.

Differential Revision: D97933992

Summary:
When `MSLK_BUILD_HIP_FMHA=0`, `attention.cpp` was still compiled into the main library and registered schemas for `efficient_attention_forward_ck`, `efficient_attention_backward_ck`, and `_ck_rand_uniform` via `TORCH_LIBRARY_FRAGMENT`, even though the `TORCH_LIBRARY_IMPL` bindings and kernel implementations (in `mslk_hip_fmha`) were absent. This left unimplemented operators registered in the library — the op appears in the dispatcher but calling it fails.

Fix by moving each `m.def` into the same file as its `TORCH_LIBRARY_IMPL`, inside `hip_fmha/`. Since those files are only compiled as part of the `mslk_hip_fmha` static library, schema and implementation now come and go together. The decoder ops remain in `attention.cpp` since their situation differs.

Differential Revision: D97933992
@meta-cla meta-cla bot added the cla signed label Mar 24, 2026
@meta-codesync
Copy link
Copy Markdown

meta-codesync bot commented Mar 24, 2026

@bottler has exported this pull request. If you are a Meta employee, you can view the originating Diff in D97933992.

@meta-codesync
Copy link
Copy Markdown

meta-codesync bot commented Mar 25, 2026

This pull request has been merged in a271a6d.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant