viable/strict/1788331551: [ROCm][CK] Remove host-side FMHA sequence padding workaround (#195498)
What changed: The external Q/K/V sequence padding, O/LSE scratch allocation, and output copies from the PyTorch CK FMHA wrapper were removed. The wrapper now forwards the original logical tensors to CK, which now rejects unpadded, unsafe instances via its generated dispatcher. Operator coverage was added for unaligned sequence lengths, dense bias, dropout replay, and grouped-query attention.
Why it matters: This simplifies the FMHA path and reduces host-side padding complexity, potentially improving performance and reliability on ROCm GPUs. The release also broadens test coverage to ensure correct behavior with unaligned sequences and related features.
Why it matters: Removes padding-related overhead and clarifies data flow to CK, with broader validation across unaligned lengths and related features.
AI-assisted brief
source-only
Add a comment
No account required