From c3de179e1ff9a5d5b8336bf1b734e9f66ffbdb2e Mon Sep 17 00:00:00 2001 From: Guy J Grigsby Date: Tue, 30 Jun 2026 15:11:01 -0600 Subject: [PATCH 1/2] fused_ops: add additive attention bias (ALiBi/relative-position) to SDPA 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. --- fused_ops.go | 4 + internal/gobackend/fusedops/sdpa.go | 74 +++++++++++++++---- .../gobackend/fusedops/sdpa_direct_test.go | 65 ++++++++++++++++ 3 files changed, 130 insertions(+), 13 deletions(-) diff --git a/fused_ops.go b/fused_ops.go index db62fec..fd297ef 100644 --- a/fused_ops.go +++ b/fused_ops.go @@ -228,6 +228,10 @@ type ScaledDotProductAttentionConfig struct { // (padding mask) instead of a materialized [S,Skv] mask. Combined with causal=true // this is a padding-causal mask. nil = unused. QuerySeqLen, KeyValueSeqLen Value + + // Bias is an optional additive attention-score bias broadcast to [B,H,S,Skv] + // (ALiBi / relative-position). NOT the Q/K/V projection bias. nil = unused. + Bias Value } // ActivationType specifies the activation function for fused operations. diff --git a/internal/gobackend/fusedops/sdpa.go b/internal/gobackend/fusedops/sdpa.go index 587ac76..84db65e 100644 --- a/internal/gobackend/fusedops/sdpa.go +++ b/internal/gobackend/fusedops/sdpa.go @@ -69,6 +69,7 @@ type nodeScaledDotProductAttention struct { causal bool hasMask bool hasSeqLens bool + hasBias bool options *compute.ScaledDotProductAttentionConfig } @@ -76,7 +77,7 @@ func (d *nodeScaledDotProductAttention) EqualNodeData(other gobackend.NodeDataCo o := other.(*nodeScaledDotProductAttention) return d.numHeads == o.numHeads && d.numKVHeads == o.numKVHeads && d.axesLayout == o.axesLayout && d.scale == o.scale && d.causal == o.causal && - d.hasMask == o.hasMask && d.hasSeqLens == o.hasSeqLens && + d.hasMask == o.hasMask && d.hasSeqLens == o.hasSeqLens && d.hasBias == o.hasBias && d.equalOptions(o) } @@ -113,6 +114,10 @@ func buildSDPANode( if hasSeqLens { values = append(values, options.QuerySeqLen, options.KeyValueSeqLen) } + hasBias := options != nil && options.Bias != nil + if hasBias { + values = append(values, options.Bias) + } inputs, err := f.VerifyAndCastValues(opName, values...) if err != nil { @@ -135,7 +140,7 @@ func buildSDPANode( data := &nodeScaledDotProductAttention{ numHeads: numHeads, numKVHeads: numKVHeads, axesLayout: axesLayout, scale: scale, causal: causal, - hasMask: mask != nil, hasSeqLens: hasSeqLens, + hasMask: mask != nil, hasSeqLens: hasSeqLens, hasBias: hasBias, options: options, } node, _ := f.GetOrCreateNode(opType, qNode.Shape.Clone(), inputs, data) @@ -180,16 +185,29 @@ func execFusedScaledDotProductAttention(backend *gobackend.Backend, node *goback } next++ } + var bias *gobackend.Buffer + if data.hasBias { + bias = inputs[next] + next++ + } _ = next - // For rank-4 BSHD masks [batch, seq, heads, kvLen], transpose to BHSD so that - // per-head mask data is contiguous [seqLen, kvLen]. The mask is small (no headDim - // axis), so this is cheap. Rank ≤ 3 masks have no head dimension and work as-is. - if data.axesLayout == compute.AxesLayoutBSHD && mask != nil && mask.RawShape.Rank() == 4 { - var err error - mask, err = transposeBuffer(backend, mask, []int{0, 2, 1, 3}) - if err != nil { - return nil, err + // For rank-4 BSHD masks/bias [batch, seq, heads, kvLen], transpose to BHSD so that + // per-head data is contiguous [seqLen, kvLen]. Rank ≤ 3 have no head axis and work as-is. + if data.axesLayout == compute.AxesLayoutBSHD { + if mask != nil && mask.RawShape.Rank() == 4 { + var err error + mask, err = transposeBuffer(backend, mask, []int{0, 2, 1, 3}) + if err != nil { + return nil, err + } + } + if bias != nil && bias.RawShape.Rank() == 4 { + var err error + bias, err = transposeBuffer(backend, bias, []int{0, 2, 1, 3}) + if err != nil { + return nil, err + } } } @@ -203,12 +221,16 @@ func execFusedScaledDotProductAttention(backend *gobackend.Backend, node *goback if mask != nil { maskBatchStride, maskHeadStride = sdpaComputeMaskStrides(mask.RawShape.Dimensions) } + var biasBatchStride, biasHeadStride int + if bias != nil { + biasBatchStride, biasHeadStride = sdpaComputeMaskStrides(bias.RawShape.Dimensions) + } switch query.RawShape.DType { case dtypes.Float32: - sdpaMultiHeadGeneric[float32](query, key, value, mask, output, data, maskBatchStride, maskHeadStride, querySeqLen, keyValueSeqLen) + sdpaMultiHeadGeneric[float32](query, key, value, mask, bias, output, data, maskBatchStride, maskHeadStride, biasBatchStride, biasHeadStride, querySeqLen, keyValueSeqLen) case dtypes.Float64: - sdpaMultiHeadGeneric[float64](query, key, value, mask, output, data, maskBatchStride, maskHeadStride, querySeqLen, keyValueSeqLen) + sdpaMultiHeadGeneric[float64](query, key, value, mask, bias, output, data, maskBatchStride, maskHeadStride, biasBatchStride, biasHeadStride, querySeqLen, keyValueSeqLen) default: return nil, errors.Errorf("FusedScaledDotProductAttention: unsupported dtype %s", query.RawShape.DType) } @@ -284,11 +306,15 @@ func transposeBuffer(backend *gobackend.Backend, buf *gobackend.Buffer, permutat // scores is a dense [groupSize, seqLen, kvLen] scratch buffer. // Masks are dense per-head [seqLen, kvLen] buffers, shared across the group when // maskGroupStride is 0, or offset by maskGroupStride per group member for per-head masks. +// additiveBias (optional) is added to scores before additiveMask and softmax; indexed +// by biasGroupStride per group member (same stride convention as maskGroupStride). func sdpaGeneric[T float32 | float64]( q, k, v []T, qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int, additiveMask []T, booleanMask []bool, maskGroupStride int, + additiveBias []T, + biasGroupStride int, scores []T, output []T, groupSize, seqLen, kvLen, headDim int, scale T, causal bool, @@ -297,6 +323,7 @@ func sdpaGeneric[T float32 | float64]( for gIdx := range groupSize { gQOff := qOff + gIdx*qGroupStride gMaskOff := gIdx * maskGroupStride + gBiasOff := gIdx * biasGroupStride for qIdx := range seqLen { if qIdx >= qLimit { // Padded query position: emit a zero output row, skip attention. @@ -312,6 +339,7 @@ func sdpaGeneric[T float32 | float64]( qBase := gQOff + qIdx*qSeqStride scoreIdxBase := (gIdx*seqLen + qIdx) * kvLen maskIdxBase := gMaskOff + qIdx*kvLen + biasIdxBase := gBiasOff + qIdx*kvLen kvLenUnmasked := kvLen if kvLimit < kvLenUnmasked { @@ -343,6 +371,9 @@ func sdpaGeneric[T float32 | float64]( dot += q[qBase+d] * k[kBase+d] } s := dot * scale + if len(additiveBias) > 0 { + s += additiveBias[biasIdxBase+kvIdx] + } if len(additiveMask) > 0 { s += additiveMask[maskIdx] } @@ -425,7 +456,7 @@ func sdpaGeneric[T float32 | float64]( } } -func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, output *gobackend.Buffer, data *nodeScaledDotProductAttention, maskBatchStride, maskHeadStride int, querySeqLen, keyValueSeqLen []int32) { +func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, bias, output *gobackend.Buffer, data *nodeScaledDotProductAttention, maskBatchStride, maskHeadStride, biasBatchStride, biasHeadStride int, querySeqLen, keyValueSeqLen []int32) { q := query.Flat.([]T) k := key.Flat.([]T) v := value.Flat.([]T) @@ -439,6 +470,10 @@ func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, output * additiveMask = mask.Flat.([]T) } } + var additiveBias []T + if bias != nil { + additiveBias = bias.Flat.([]T) + } dims := query.RawShape.Dimensions batchSize := dims[0] @@ -512,9 +547,22 @@ func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, output * booleanMaskSlice = booleanMask[maskOffset:maskEnd] } } + // Compute bias slice and group stride for this KV head group. + var additiveBiasSlice []T + biasGroupStride := 0 + if len(additiveBias) > 0 { + biasOffset := batchIdx*biasBatchStride + kvHeadIdx*groupSize*biasHeadStride + biasEnd := biasOffset + maskSliceLen + if biasHeadStride > 0 && groupSize > 1 { + biasEnd = biasOffset + (groupSize-1)*biasHeadStride + maskSliceLen + biasGroupStride = biasHeadStride + } + additiveBiasSlice = additiveBias[biasOffset:biasEnd] + } sdpaGeneric( q, k, v, qOff, kvOff, qSeqStride, kvSeqStride, qHeadStride, additiveMaskSlice, booleanMaskSlice, maskGroupStride, + additiveBiasSlice, biasGroupStride, scores, out, groupSize, seqLen, kvLen, headDim, scale, causal, diff --git a/internal/gobackend/fusedops/sdpa_direct_test.go b/internal/gobackend/fusedops/sdpa_direct_test.go index 518adf5..35d582e 100644 --- a/internal/gobackend/fusedops/sdpa_direct_test.go +++ b/internal/gobackend/fusedops/sdpa_direct_test.go @@ -198,6 +198,71 @@ func TestSDPADirect_SeqLenClamped(t *testing.T) { } } +// TestSDPADirect_WithBias verifies that an additive attention-score bias shifts +// softmax weights and produces an output that matches a decomposed reference. +// Bias strongly favors key 1 so the test is sensitive to whether bias is applied. +func TestSDPADirect_WithBias(t *testing.T) { + b := newGoBackend(t) + // batch=1, 1 head, seqLen=2 queries, kvLen=2 keys, headDim=1. + // Raw scores (scale=1): q@k^T = [[1,1],[1,1]]. + // Bias [1,1,2,2]: [[[[-10, 10], [-10, 10]]]] strongly favors key 1. + // After bias: scores = [[-9, 11], [-9, 11]]. + // softmax([-9,11]) ≈ [~0, ~1] -> output ≈ v[1] = 20 for both queries. + q := [][][][]float32{{{{1}, {1}}}} // [1,1,2,1] + k := [][][][]float32{{{{1}, {1}}}} // [1,1,2,1] + v := [][][][]float32{{{{10}, {20}}}} // [1,1,2,1] + bias := [][][][]float32{{{{-10, 10}, {-10, 10}}}} // [1,1,2,2] batch,head,seqLen,kvLen + got, err := testutil.Exec1(b, []any{q, k, v, bias}, func(f compute.Function, params []compute.Value) (compute.Value, error) { + cfg := &compute.ScaledDotProductAttentionConfig{Bias: params[3]} + out, _, err := f.FusedScaledDotProductAttention(params[0], params[1], params[2], nil, 1, 1, compute.AxesLayoutBHSD, 1.0, false, cfg) + return out, err + }) + if err != nil { + t.Fatalf("SDPA with bias failed: %+v", err) + } + // softmax([-9,11]) ≈ [sigma_neg, sigma_pos]; sigma_neg ≈ 1/(1+e^20) ≈ 2e-9. + // output ≈ 20 - tiny correction; tolerance 0.01 covers the residual. + softmaxPos := float32(1.0 / (1.0 + math.Exp(-20))) + softmaxNeg := 1 - softmaxPos + wantVal := float32(softmaxNeg*10 + softmaxPos*20) + want := [][][][]float32{{{{wantVal}, {wantVal}}}} + if ok, diff := testutil.IsInDelta(want, got, 0.01); !ok { + t.Errorf("SDPA with bias mismatch:\n%s", diff) + } +} + +// TestSDPADirect_BiasWithCausal verifies that additive bias and causal mask combine +// correctly: causal zeroes future positions, bias shifts scores in valid positions. +func TestSDPADirect_BiasWithCausal(t *testing.T) { + b := newGoBackend(t) + // batch=1, 1 head, seqLen=2, kvLen=2, headDim=1, scale=1. + // Bias [1,1,2,2]: [[[[ 0, -100], [0, 10]]]] — large negative on (q0,k1) and large positive on (q1,k1). + // Causal: q0 sees only k0 (k1 is future); q1 sees k0 and k1. + // q0: causal masks k1 -> attends only k0 -> output = v[0] = 10 (bias on k1 irrelevant). + // q1: scores = [1+0, 1+10] = [1, 11]; softmax([1,11]) -> large weight on k1 -> output ≈ 20. + q := [][][][]float32{{{{1}, {1}}}} // [1,1,2,1] + k := [][][][]float32{{{{1}, {1}}}} // [1,1,2,1] + v := [][][][]float32{{{{10}, {20}}}} // [1,1,2,1] + bias := [][][][]float32{{{{0, -100}, {0, 10}}}} // [1,1,2,2] + got, err := testutil.Exec1(b, []any{q, k, v, bias}, func(f compute.Function, params []compute.Value) (compute.Value, error) { + cfg := &compute.ScaledDotProductAttentionConfig{Bias: params[3]} + out, _, err := f.FusedScaledDotProductAttention(params[0], params[1], params[2], nil, 1, 1, compute.AxesLayoutBHSD, 1.0, true, cfg) + return out, err + }) + if err != nil { + t.Fatalf("SDPA bias+causal failed: %+v", err) + } + // q0: causal -> only k0 -> output=10. + // q1: softmax([1+0, 1+10]) = softmax([1,11]) -> weight on k1 ≈ 1/(1+e^-10) ≈ ~1. + softmaxQ1K1 := float32(1.0 / (1.0 + math.Exp(-10))) + softmaxQ1K0 := 1 - softmaxQ1K1 + wantQ1 := float32(softmaxQ1K0*10 + softmaxQ1K1*20) + want := [][][][]float32{{{{10}, {wantQ1}}}} + if ok, diff := testutil.IsInDelta(want, got, 0.01); !ok { + t.Errorf("SDPA bias+causal mismatch:\n%s", diff) + } +} + func TestSDPADirect_FP8NotImplemented(t *testing.T) { b := newGoBackend(t) // The go backend rejects F8 parameters at creation, so feed F8 by converting From 0de8b0d3d0fc8bcb1bc803e7789d384fdd0f86fe Mon Sep 17 00:00:00 2001 From: Guy J Grigsby Date: Tue, 30 Jun 2026 15:16:08 -0600 Subject: [PATCH 2/2] fused_ops: validate bias rank/shape/dtype in buildSDPANode; comment GQA 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. --- internal/gobackend/fusedops/sdpa.go | 42 +++++++++++++++++++ .../gobackend/fusedops/sdpa_direct_test.go | 34 +++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/internal/gobackend/fusedops/sdpa.go b/internal/gobackend/fusedops/sdpa.go index 84db65e..da7c2d6 100644 --- a/internal/gobackend/fusedops/sdpa.go +++ b/internal/gobackend/fusedops/sdpa.go @@ -136,6 +136,47 @@ func buildSDPANode( if numHeads <= 0 || numKVHeads <= 0 || numHeads%numKVHeads != 0 { return nil, errors.Errorf("%s: numHeads (%d) must be positive and divisible by numKVHeads (%d)", opName, numHeads, numKVHeads) } + if hasBias { + // Bias input is the last appended value; resolve its node from inputs. + biasIdx := len(inputs) - 1 + biasNode := inputs[biasIdx] + if !biasNode.Shape.DType.IsFloat() { + return nil, errors.Errorf("%s: bias must be a float dtype matching the compute dtype, got %s", + opName, biasNode.Shape.DType) + } + biasRank := biasNode.Shape.Rank() + if biasRank < 2 || biasRank > 4 { + return nil, errors.Errorf("%s: bias must have rank 2, 3, or 4, got rank %d", opName, biasRank) + } + // Score shape is [B, H, S, Skv]. Bias dims are right-aligned and each must be 1 + // or equal to the corresponding score dim (standard broadcasting contract). + kNode := inputs[1] + var batchDim, numHeadsDim, seqDim, kvDim int + if axesLayout == compute.AxesLayoutBSHD { + batchDim = qNode.Shape.Dimensions[0] + numHeadsDim = numHeads + seqDim = qNode.Shape.Dimensions[1] + kvDim = kNode.Shape.Dimensions[1] + } else { + // BHSD: [batch, heads, seq, headDim] + batchDim = qNode.Shape.Dimensions[0] + numHeadsDim = numHeads + seqDim = qNode.Shape.Dimensions[2] + kvDim = kNode.Shape.Dimensions[2] + } + scoreDims := [4]int{batchDim, numHeadsDim, seqDim, kvDim} + biasDims := biasNode.Shape.Dimensions + // Right-align bias dims against score dims. + offset := 4 - biasRank + for i, bd := range biasDims { + sd := scoreDims[offset+i] + if bd != 1 && bd != sd { + return nil, errors.Errorf( + "%s: bias dim %d is %d, not broadcastable to score dim %d (score shape [%d,%d,%d,%d])", + opName, i, bd, sd, scoreDims[0], scoreDims[1], scoreDims[2], scoreDims[3]) + } + } + } data := &nodeScaledDotProductAttention{ numHeads: numHeads, numKVHeads: numKVHeads, axesLayout: axesLayout, @@ -551,6 +592,7 @@ func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, bias, ou var additiveBiasSlice []T biasGroupStride := 0 if len(additiveBias) > 0 { + // *groupSize: bias is per-Q-head under GQA; first Q-head of this KV group starts at kvHeadIdx*groupSize. biasOffset := batchIdx*biasBatchStride + kvHeadIdx*groupSize*biasHeadStride biasEnd := biasOffset + maskSliceLen if biasHeadStride > 0 && groupSize > 1 { diff --git a/internal/gobackend/fusedops/sdpa_direct_test.go b/internal/gobackend/fusedops/sdpa_direct_test.go index 35d582e..6a2ff44 100644 --- a/internal/gobackend/fusedops/sdpa_direct_test.go +++ b/internal/gobackend/fusedops/sdpa_direct_test.go @@ -263,6 +263,40 @@ func TestSDPADirect_BiasWithCausal(t *testing.T) { } } +// TestSDPADirect_BiasValidation verifies that a malformed bias operand returns a clear +// validation error from buildSDPANode rather than silently panicking in the kernel. +func TestSDPADirect_BiasValidation(t *testing.T) { + b := newGoBackend(t) + q := [][][][]float32{{{{1}, {1}}}} // [1,1,2,1] BHSD: batch=1 head=1 seq=2 dim=1 + k := [][][][]float32{{{{1}, {1}}}} // [1,1,2,1] kv=2 + v := [][][][]float32{{{{10}, {20}}}} + + cases := []struct { + name string + bias any + }{ + // rank-1 bias: invalid (must be rank 2-4). + {"rank1", []float32{1, 2}}, + // rank-2 bias [2,3]: second dim (3) != kvLen (2) and != 1, not broadcastable. + {"rank2_bad_kvdim", [][]float32{{1, 2, 3}, {1, 2, 3}}}, + // rank-4 bias [1,1,3,2]: seqLen dim 3 != score seqLen 2 and != 1. + {"rank4_bad_seqdim", [][][][]float32{{{{1, 2}, {1, 2}, {1, 2}}}}}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + _, err := testutil.Exec1(b, []any{q, k, v, tc.bias}, func(f compute.Function, params []compute.Value) (compute.Value, error) { + cfg := &compute.ScaledDotProductAttentionConfig{Bias: params[3]} + out, _, err := f.FusedScaledDotProductAttention(params[0], params[1], params[2], nil, 1, 1, compute.AxesLayoutBHSD, 1.0, false, cfg) + return out, err + }) + if err == nil { + t.Fatalf("case %s: expected bias validation error, got nil", tc.name) + } + }) + } +} + func TestSDPADirect_FP8NotImplemented(t *testing.T) { b := newGoBackend(t) // The go backend rejects F8 parameters at creation, so feed F8 by converting