Stage 2: additive attention bias (config field + CPU reference)#4
Closed
guygrigsby wants to merge 3 commits into
Closed
Stage 2: additive attention bias (config field + CPU reference)#4guygrigsby wants to merge 3 commits into
guygrigsby wants to merge 3 commits into
Conversation
Stage 1: seqlen config fields + CPU reference + fp8 gate
Add Bias Value field to ScaledDotProductAttentionConfig. Thread it as an operand through buildSDPANode (deduped via input identity, not equality check). Extract it in the executor, transpose BSHD rank-4 bias to BHSD, compute broadcast strides via sdpaComputeMaskStrides, and add to scores before softmax in sdpaGeneric. No capabilities change; no dropout; no fp8. Tests: TestSDPADirect_WithBias (bias shifts weights to strongly-favored key) and TestSDPADirect_BiasWithCausal (causal masks future positions while bias shifts scores within valid positions). Both confirm RED before implementation and GREEN after.
…QA bias offset Bias operand now checked: float dtype, rank 2-4, each dim broadcastable (1 or equal) to the corresponding score dim [B,H,S,Skv]. Returns a wrapped error on mismatch instead of panicking deep in the kernel. Add one-line comment on the *groupSize term in the bias offset computation (per-Q-head bias under GQA, mirroring the mask comment). TestSDPADirect_BiasValidation covers rank-1, bad-kvdim rank-2, and bad-seqdim rank-4 cases; existing WithBias/BiasWithCausal tests stay green.
Owner
Author
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
1 participant
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.
Reopens the Stage 2 bias work (compute#2 was accidentally merged into stage1-seqlen-config and reverted; stage1-seqlen-config has been reset to its clean Stage 1 head). Stacked on the Stage 1 PR branch. Bias only (dropout, bias+dropout, fp8 all cut as NotImplemented seams).
ScaledDotProductAttentionConfig.Bias Value(additive attention-score bias [B,H,Sq,Skv]; NOT the Q/K/V projection bias). Deduped via operand identity (equality compares the hasBias flag, not the Value).validateBias(rank-4, exact shape) returns a clear error, not a panic. Capability staysfalse; direct unit tests.Validated on RTX 3070 Ti via the downstream go-xla/gomlx bias PRs: cuDNN ScaleBias applies the bias forward + backward.