Stage 2: cuDNN fmhaScaleBias variant (additive attention bias, cumulative on #37)#39
Merged
Merged
Conversation
This was referenced Jul 1, 2026
…bias selectFMHAVariant routes cfg.Bias to the fmhaScaleBias targets (bias+seqlens -> ErrNotImplemented, cuDNN limitation). Bias operand inserted before seqlens (fwd) and before O (bwd); backward declares the dBias result slot (not propagated). validateBias requires bias dtype == q/k/v dtype (no auto-conversion, matching the backend's explicit-dtype policy). Adapts onto the refactored dispatch/config.
…was hardcoded BF16)
The intermediate_tensor_shape element_type was hardcoded BF16, so the Float16
fused BACKWARD custom-call failed to compile ('Input datatypes do not match')
while the forward tolerated it. Derive it from the q/k/v dtype (BF16/F16) in
selectFMHAVariant, enabling f16 fused training (fwd+bwd).
Contributor
Author
|
Beyond the This repairs f16 fused training generally, not just the bias path. Validated on cuda (RTX 3070 Ti): f16 fused forward+backward now compiles and matches the decomposed reference; full fusion suite green with no bf16 regression. |
janpfeifer
approved these changes
Jul 7, 2026
janpfeifer
left a comment
janpfeifer
left a comment
Contributor
There was a problem hiding this comment.
Thank you!! Looks good to me, I'll wait till the changes in compute are merged to merge this one.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stage 2 of the fused-attention work: the cuDNN ScaleBias variant for additive attention bias. Cumulative on top of Stage 1 (#37), the prerequisite; the Stage 1 commits show here too until #37 merges to main, then this shrinks to the bias delta.
selectFMHAVariant routes a non-nil bias to __cudnn$fmhaScaleBiasSoftmax and its backward; bias combines with causal, and bias+seqlens returns ErrNotImplemented (the ScaleBias kernel takes no seqlen operands) so it decomposes. Forward has the bias at operand index 3; backward at index 5 with an added dBias result at index 3, matching the JAX-lowered dot_product_attention(bias=..., implementation=cudnn) reference. Bias is cast to bf16 and shape-validated.
Validated on an RTX 3070 Ti: the non-uniform-V test shows the bias changes the output (it is genuinely applied), and forward and backward both match the decomposed reference.