Skip to content

Enable softmax-sink FMHA kernels for TransformerEngine - #7

Open
shurale-nkn wants to merge 6 commits into
ROCm:mainfrom
shurale-nkn:knasyrov/fixed_sink
Open

Enable softmax-sink FMHA kernels for TransformerEngine#7
shurale-nkn wants to merge 6 commits into
ROCm:mainfrom
shurale-nkn:knasyrov/fixed_sink

Conversation

@shurale-nkn

Copy link
Copy Markdown

Motivation

Required by ROCm/TransformerEngine#678[ROCm] Jax Add softmax sink (learnable
off-by-one) support for the ROCm/CK fused attention backend
.

Receipt 700's forward rule carried cond &= kernel_ctx.pipeline.F_sink == "f", so
codegen emitted zero sink instances and a forward call with has_sink=true found no
matching trait — CK's generated fmha_fwd() returned -1, surfacing in TE as
fused attn configs not supported in ck_fused_attn fwd pass. That invariant was written
when TE had no sink support, so it is now stale. This PR drops the F_sink condition
from 0003-ck-te-receipt-700.patch and updates the two spots that document the
invariant (patch preamble and the generate.py help text).

Also adds 0008-fmha-bwd-sink-fix.patch: CK's fmha_bwd_kernel.hpp indexed sink_ptr
as i_batch * nhead + i_nhead, but the buffer holds one value per head (shape
[nhead]), matching how the forward kernels already read it, so i_batch must not
factor into the offset.

Technical Details

Test Plan

Test Result

Submission Checklist

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants