fused_ops: typed flash backward (FusedSDPA + VJP), not a generic CustomCall

fused_ops.go

@@ -203,8 +203,15 @@ type Quantization struct {
 	GGMLType GGMLQuantType
 }
 
-// ScaledDotProductAttentionConfig holds optional optimization hints for FusedScaledDotProductAttention.
+// ScaledDotProductAttentionConfig holds optional optimization hints and fused-attention
+// parameters for FusedScaledDotProductAttention.
 // A nil *ScaledDotProductAttentionConfig means "use defaults" (all optimizations disabled).
+// Correctness-affecting fields (Bias, QuerySeqLen, KeyValueSeqLen, dropout, etc.): backends
+// that cannot honor them MUST return ErrNotImplemented so the caller falls back to the
+// decomposed path.
+// Pure optimization hints (e.g. QuantizedMatmuls): backends that do not support them may
+// silently ignore them and compute a numerically equivalent result in float arithmetic.
+// nil/zero means "unused".
 type ScaledDotProductAttentionConfig struct {
 	// QuantizedMatmuls: if true, the backend may use dynamic per-head symmetric
 	// affine quantization (scale-only, no zero point) to convert float32 Q/K/V slices
@@ -215,6 +222,12 @@ type ScaledDotProductAttentionConfig struct {
 	// (e.g. ARM SDOT/UDOT, x86 VNNI). Backends that do not support quantized matmuls
 	// ignore this flag and use float arithmetic.
 	QuantizedMatmuls bool
+
+	// QuerySeqLen, KeyValueSeqLen are optional per-batch actual sequence lengths
+	// (int32 tensors, shape [B]). When set, the backend masks by sequence length
+	// (padding mask) instead of a materialized [S,Skv] mask. Combined with causal=true
+	// this is a padding-causal mask. nil = unused.
+	QuerySeqLen, KeyValueSeqLen Value
 }
 
 // ActivationType specifies the activation function for fused operations.
@@ -310,14 +323,37 @@ type FusedOps interface {
 	//     this method when both are needed. Backends may assume they won't both be set.
 	//   - options: optional optimization hints (nil uses defaults). See ScaledDotProductAttentionConfig.
 	//
-	// Output: same shape as query.
+	// Outputs:
+	//   - output: same shape as query.
+	//   - softmaxStats: optional (may be nil). When the backend supports a fused backward
+	//     (FusedScaledDotProductAttentionVJP), it returns the per-row softmax statistics
+	//     (log-sum-exp), shape [batch, numHeads, seqLen] f32, which the VJP needs. Backends
+	//     without a fused backward return nil here and ErrNotImplemented from the VJP, so the
+	//     caller differentiates the decomposed attention instead.
 	FusedScaledDotProductAttention(
 		query, key, value, mask Value,
 		numHeads, numKVHeads int,
 		axesLayout AxesLayout,
 		scale float64,
 		causal bool,
-		options *ScaledDotProductAttentionConfig) (Value, error)
+		options *ScaledDotProductAttentionConfig) (output, softmaxStats Value, err error)
+
+	// FusedScaledDotProductAttentionVJP computes the gradients (dQuery, dKey, dValue) of
+	// FusedScaledDotProductAttention given the forward output, the softmaxStats it returned, and
+	// the output gradient dOutput. The query/key/value/mask and numHeads..options parameters are
+	// the same values passed to the forward call.
+	//
+	// Backends that return non-nil softmaxStats from the forward and implement a fused backward
+	// (e.g. the cuDNN flash backward) implement this. Others return ErrNotImplemented so the caller
+	// falls back to differentiating the decomposed attention.
+	FusedScaledDotProductAttentionVJP(
+		query, key, value, mask Value,
+		numHeads, numKVHeads int,
+		axesLayout AxesLayout,
+		scale float64,
+		causal bool,
+		options *ScaledDotProductAttentionConfig,
+		output, softmaxStats, dOutput Value) (dQuery, dKey, dValue Value, err error)
 
 	// QuantizedEmbeddingLookup performs a quantized embedding lookup (row gather)
 	// with on-the-fly dequantization.

internal/gobackend/fusedops/sdpa.go

@@ -43,9 +43,12 @@ func FusedScaledDotProductAttention(
 	query, key, value, mask compute.Value,
 	numHeads, numKVHeads int, axesLayout compute.AxesLayout,
 	scale float64, causal bool,
-	options *compute.ScaledDotProductAttentionConfig) (compute.Value, error) {
-	return buildSDPANode(f, compute.OpTypeFusedScaledDotProductAttention, "FusedScaledDotProductAttention",
+	options *compute.ScaledDotProductAttentionConfig) (output, softmaxStats compute.Value, err error) {
+	// The Go backend has no fused flash backward, so it returns nil softmaxStats; its gradient
+	// goes through the decomposed path (FusedScaledDotProductAttentionVJP is ErrNotImplemented).
+	out, err := buildSDPANode(f, compute.OpTypeFusedScaledDotProductAttention, "FusedScaledDotProductAttention",
 		query, key, value, mask, numHeads, numKVHeads, axesLayout, scale, causal, options)
+	return out, nil, err
 }
 
 func init() {
@@ -64,13 +67,16 @@ type nodeScaledDotProductAttention struct {
 	axesLayout compute.AxesLayout
 	scale      float64
 	causal     bool
+	hasMask    bool
+	hasSeqLens 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.equalOptions(o)
 }
 
@@ -81,7 +87,11 @@ func (d *nodeScaledDotProductAttention) equalOptions(o *nodeScaledDotProductAtte
 	if d.options == nil || o.options == nil {
 		return false
 	}
-	return *d.options == *o.options
+	a, b := d.options, o.options
+	// QuerySeqLen and KeyValueSeqLen are Value (any) and may hold non-comparable types,
+	// so comparing them with == would panic. They are already passed as input operands
+	// in buildSDPANode, so node dedup via input identity covers them; skip here.
+	return a.QuantizedMatmuls == b.QuantizedMatmuls
 }
 
 // buildSDPANode builds the SDPA computation node.
@@ -95,6 +105,15 @@ func buildSDPANode(
 	if mask != nil {
 		values = append(values, mask)
 	}
+
+	if options != nil && (options.QuerySeqLen != nil) != (options.KeyValueSeqLen != nil) {
+		return nil, errors.Errorf("%s: QuerySeqLen and KeyValueSeqLen must both be set or both nil", opName)
+	}
+	hasSeqLens := options != nil && options.QuerySeqLen != nil && options.KeyValueSeqLen != nil
+	if hasSeqLens {
+		values = append(values, options.QuerySeqLen, options.KeyValueSeqLen)
+	}
+
 	inputs, err := f.VerifyAndCastValues(opName, values...)
 	if err != nil {
 		return nil, err
@@ -104,11 +123,21 @@ func buildSDPANode(
 	if qNode.Shape.Rank() != 4 {
 		return nil, errors.Errorf("%s: query must have rank 4, got %d", opName, qNode.Shape.Rank())
 	}
+	switch qNode.Shape.DType {
+	case dtypes.F8E4M3FN, dtypes.F8E5M2:
+		return nil, errors.Wrapf(compute.ErrNotImplemented,
+			"%s: float8 input dtype %s is not implemented in the go backend", opName, qNode.Shape.DType)
+	}
 	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)
 	}
 
-	data := &nodeScaledDotProductAttention{numHeads: numHeads, numKVHeads: numKVHeads, axesLayout: axesLayout, scale: scale, causal: causal, options: options}
+	data := &nodeScaledDotProductAttention{
+		numHeads: numHeads, numKVHeads: numKVHeads, axesLayout: axesLayout,
+		scale: scale, causal: causal,
+		hasMask: mask != nil, hasSeqLens: hasSeqLens,
+		options: options,
+	}
 	node, _ := f.GetOrCreateNode(opType, qNode.Shape.Clone(), inputs, data)
 	return node, nil
 }
@@ -128,10 +157,30 @@ func execFusedScaledDotProductAttention(backend *gobackend.Backend, node *goback
 	query := inputs[0]
 	key := inputs[1]
 	value := inputs[2]
+	next := 3
 	var mask *gobackend.Buffer
-	if len(inputs) > 3 {
-		mask = inputs[3]
+	if data.hasMask {
+		mask = inputs[next]
+		next++
+	}
+	var querySeqLen, keyValueSeqLen []int32
+	if data.hasSeqLens {
+		batchSize := inputs[0].RawShape.Dimensions[0]
+		var ok bool
+		querySeqLen, ok = inputs[next].Flat.([]int32)
+		if !ok || len(querySeqLen) != batchSize {
+			return nil, errors.Errorf("FusedScaledDotProductAttention: QuerySeqLen must be int32 with length batch (%d), got type %T len %d",
+				batchSize, inputs[next].Flat, len(querySeqLen))
+		}
+		next++
+		keyValueSeqLen, ok = inputs[next].Flat.([]int32)
+		if !ok || len(keyValueSeqLen) != batchSize {
+			return nil, errors.Errorf("FusedScaledDotProductAttention: KeyValueSeqLen must be int32 with length batch (%d), got type %T len %d",
+				batchSize, inputs[next].Flat, len(keyValueSeqLen))
+		}
+		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
@@ -157,9 +206,9 @@ func execFusedScaledDotProductAttention(backend *gobackend.Backend, node *goback
 
 	switch query.RawShape.DType {
 	case dtypes.Float32:
-		sdpaMultiHeadGeneric[float32](query, key, value, mask, output, data, maskBatchStride, maskHeadStride)
+		sdpaMultiHeadGeneric[float32](query, key, value, mask, output, data, maskBatchStride, maskHeadStride, querySeqLen, keyValueSeqLen)
 	case dtypes.Float64:
-		sdpaMultiHeadGeneric[float64](query, key, value, mask, output, data, maskBatchStride, maskHeadStride)
+		sdpaMultiHeadGeneric[float64](query, key, value, mask, output, data, maskBatchStride, maskHeadStride, querySeqLen, keyValueSeqLen)
 	default:
 		return nil, errors.Errorf("FusedScaledDotProductAttention: unsupported dtype %s", query.RawShape.DType)
 	}
@@ -243,24 +292,38 @@ func sdpaGeneric[T float32 | float64](
 	scores []T,
 	output []T,
 	groupSize, seqLen, kvLen, headDim int, scale T, causal bool,
+	qLimit, kvLimit int,
 ) {
 	for gIdx := range groupSize {
 		gQOff := qOff + gIdx*qGroupStride
 		gMaskOff := gIdx * maskGroupStride
 		for qIdx := range seqLen {
+			if qIdx >= qLimit {
+				// Padded query position: emit a zero output row, skip attention.
+				// Local padBase has its own scope, distinct from the outBase the
+				// accumulation loop declares later, so there is no redeclaration.
+				padBase := gQOff + qIdx*qSeqStride
+				for d := range headDim {
+					output[padBase+d] = 0
+				}
+				continue
+			}
 			rowMax := T(math.Inf(-1))
 			qBase := gQOff + qIdx*qSeqStride
 			scoreIdxBase := (gIdx*seqLen + qIdx) * kvLen
 			maskIdxBase := gMaskOff + qIdx*kvLen
 
 			kvLenUnmasked := kvLen
+			if kvLimit < kvLenUnmasked {
+				kvLenUnmasked = kvLimit
+			}
 			if causal {
-				kvLenUnmasked = min(kvLen, qIdx+1)
+				kvLenUnmasked = min(kvLenUnmasked, qIdx+1)
 			}
 
-			// Zero out scores to prevent stale data from previous iterations
-			// when boolean mask or causal mask skips positions.
-			if causal || len(booleanMask) > 0 {
+			// Zero out scores so a future loop widening past kvLenUnmasked
+			// cannot read stale scores from a prior kvHead iteration.
+			if causal || len(booleanMask) > 0 || kvLimit < kvLen {
 				for i := scoreIdxBase; i < scoreIdxBase+kvLen; i++ {
 					scores[i] = 0
 				}
@@ -362,7 +425,7 @@ func sdpaGeneric[T float32 | float64](
 	}
 }
 
-func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, output *gobackend.Buffer, data *nodeScaledDotProductAttention, maskBatchStride, maskHeadStride int) {
+func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, output *gobackend.Buffer, data *nodeScaledDotProductAttention, maskBatchStride, maskHeadStride int, querySeqLen, keyValueSeqLen []int32) {
 	q := query.Flat.([]T)
 	k := key.Flat.([]T)
 	v := value.Flat.([]T)
@@ -420,6 +483,14 @@ func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, output *
 	scores := make([]T, groupSize*seqLen*kvLen)
 	maskSliceLen := seqLen * kvLen
 	for batchIdx := range batchSize {
+		qLimit := seqLen
+		kvLimit := kvLen
+		if len(querySeqLen) > 0 {
+			qLimit = max(0, min(int(querySeqLen[batchIdx]), seqLen))
+		}
+		if len(keyValueSeqLen) > 0 {
+			kvLimit = max(0, min(int(keyValueSeqLen[batchIdx]), kvLen))
+		}
 		for kvHeadIdx := range numKVHeads {
 			qOff := batchIdx*qBatchStride + kvHeadIdx*groupSize*qHeadStride
 			kvOff := batchIdx*kvBatchStride + kvHeadIdx*kvHeadStride
@@ -447,6 +518,7 @@ func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, output *
 				scores,
 				out,
 				groupSize, seqLen, kvLen, headDim, scale, causal,
+				qLimit, kvLimit,
 			)
 		}
 	}