Stage 2: additive attention bias (config field + CPU reference)

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.

internal/gobackend/fusedops/sdpa.go

@@ -69,14 +69,15 @@ type nodeScaledDotProductAttention struct {
 	causal     bool
 	hasMask    bool
 	hasSeqLens bool
+	hasBias    bool
 	options    *compute.ScaledDotProductAttentionConfig
 }
 
 func (d *nodeScaledDotProductAttention) EqualNodeData(other gobackend.NodeDataComparable) bool {
 	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 {
@@ -131,11 +136,52 @@ 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,
 		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 +226,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 +262,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 +347,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 +364,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 +380,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 +412,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 +497,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 +511,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 +588,23 @@ 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 {
+				// *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 {
+					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,

internal/gobackend/fusedops/sdpa_direct_test.go

@@ -198,6 +198,105 @@ 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)
+	}
+}
+
+// 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