flash.go

// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0

package xla

import (
	"encoding/json"
	"maps"
	"strconv"

	"github.com/gomlx/compute"
	"github.com/gomlx/compute/dtypes"
	"github.com/gomlx/compute/shapes"
	"github.com/pkg/errors"
)

// cuDNN fused-attention custom-call targets. S1 wires the standard softmax pair only.
// [S2] (Task 2b) adds the fmhaScaleBias* / *Dropout / fmhaScaleBias*Dropout target rows.
// FP8 targets (__cudnn$fmhaSoftmaxF8 / …BackwardF8) are intentionally NOT defined:
// fp8 fmha is paused (no local sm_8.9+ hardware to test). fp8 input dtype falls to
// ErrNotImplemented in selectFMHAVariant. Add that row when wiring fp8 on Hopper/Ada.
//
// These custom-calls, their backend_config, and the operand/result layouts below are not part of
// the public XLA op set. They are the same calls JAX emits for
// jax.nn.dot_product_attention(..., implementation="cudnn"). Provenance:
//   - target names + backend_config field mapping (CudnnfMHABackendConfig, _custom_name_maps):
//     https://github.com/jax-ml/jax/blob/main/jax/_src/cudnn/fused_attention_stablehlo.py
//   - the exact backend_config JSON and layouts here were captured by lowering that JAX call to
//     StableHLO (jax.jit(fn).lower(...).compiler_ir("stablehlo")) and reading the emitted
//     stablehlo.custom_call. The stats (log-sum-exp) output appears only in the jax.grad lowering.
//   - XLA custom_call contract (documents api_version=4; these calls use api_version=2):
//     https://openxla.org/xla/custom_call
//   - cuDNN fMHA performance context: https://github.com/jax-ml/jax/issues/24934
const (
	fmhaSoftmaxFwd          = "__cudnn$fmhaSoftmax"
	fmhaSoftmaxBwd          = "__cudnn$fmhaSoftmaxBackward"
	fmhaScaleBiasSoftmaxFwd = "__cudnn$fmhaScaleBiasSoftmax"
	fmhaScaleBiasSoftmaxBwd = "__cudnn$fmhaScaleBiasSoftmaxBackward"
)

// fmhaVariant captures the config-derived custom-call selection: the fwd/bwd targets, the
// backend_config mask_type, and the operand-set flags. Built by selectFMHAVariant.
type fmhaVariant struct {
	fwdTarget, bwdTarget string
	maskType             string // "CAUSAL" | "PADDING" | "PADDING_CAUSAL" | "NO_MASK"
	elementType          string // XLA element type of q/k/v: "BF16" | "F16". Drives the intermediate
	// tensor's element_type in the backend_config; a mismatch fails the backward custom-call.
	hasSeqLens bool
	hasBias    bool
}

// selectFMHAVariant maps the q/k/v dtype and (causal, seqlens, bias) to a cuDNN variant.
// Dtype gate: f16/bf16 only; anything else (incl. fp8 e4m3fn/e5m2 -- paused, no local hardware) ->
// ErrNotImplemented, and the caller falls back to the decomposed path.
// mask_type derives from causal + seqlens: PADDING_CAUSAL (both), PADDING (seqlens only),
// CAUSAL (causal only), NO_MASK (neither).
//
// Bias routing: cfg.Bias non-nil selects the fmhaScaleBias targets. cuDNN's ScaleBias kernel does
// not accept seqlen operands, so bias+seqlens returns ErrNotImplemented and the caller falls back
// to the decomposed path.
func selectFMHAVariant(op string, qkvDType dtypes.DType, causal bool,
	cfg *compute.ScaledDotProductAttentionConfig) (fmhaVariant, error) {
	var v fmhaVariant
	hasSeqLens := cfg != nil && cfg.QuerySeqLen != nil && cfg.KeyValueSeqLen != nil
	hasBias := cfg != nil && cfg.Bias != nil

	switch qkvDType {
	case dtypes.BFloat16:
		v.elementType = "BF16"
	case dtypes.Float16:
		v.elementType = "F16"
	default:
		// fp8 (e4m3fn/e5m2) lands here too: paused, not wired. NotImplemented -> decomposed.
		return v, errors.Wrapf(compute.ErrNotImplemented,
			"%s: cuDNN fmha needs f16/bf16, got %s", op, qkvDType)
	}

	if hasBias && hasSeqLens {
		// cuDNN ScaleBias kernel does not accept seqlen operands; fall back to decomposed.
		return v, errors.Wrapf(compute.ErrNotImplemented,
			"%s: cuDNN fmha bias+seqlens not supported; use decomposed path", op)
	}
	if hasBias {
		v.fwdTarget, v.bwdTarget = fmhaScaleBiasSoftmaxFwd, fmhaScaleBiasSoftmaxBwd
		v.hasBias = true
	} else {
		v.fwdTarget, v.bwdTarget = fmhaSoftmaxFwd, fmhaSoftmaxBwd
	}

	switch {
	case causal && hasSeqLens:
		v.maskType = "PADDING_CAUSAL"
	case hasSeqLens:
		v.maskType = "PADDING"
	case causal:
		v.maskType = "CAUSAL"
	default:
		v.maskType = "NO_MASK"
	}
	v.hasSeqLens = hasSeqLens
	return v, nil
}

// flashBackendConfigV builds a cudnn_fmha_backend_config for the given variant: mask_type comes
// from v, the score-matrix dims [B,H,S,S] and fmha_scale from the shape, dotDimNumbers carries the
// bmm dot_dimension_numbers JSON (the fwd/bwd-specific part). S1 has no dropout, so dropout_rate is
// the literal 0; Task 2b [S2] switches it to formatScale(v.dropoutRate).
func flashBackendConfigV(b, h, s int, scale float64, dotDimNumbers map[string]any, v fmhaVariant) string {
	config := map[string]any{
		"operation_queue_id": "0",
		"cudnn_fmha_backend_config": map[string]any{
			"algorithm": map[string]any{
				"algo_id":           "0",
				"math_type":         "TENSOR_OP_MATH",
				"tuning_knobs":      map[string]string{"17": "1", "24": "0"},
				"is_cudnn_frontend": true,
				"workspace_size":    "0",
			},
			"fmha_scale": scale,
			"intermediate_tensor_shape": map[string]any{
				"element_type": v.elementType,
				"dimensions":   []string{strconv.Itoa(b), strconv.Itoa(h), strconv.Itoa(s), strconv.Itoa(s)},
				"tuple_shapes": []any{},
				"layout": map[string]any{
					"dim_level_types":                     []any{},
					"dim_unique":                          []any{},
					"dim_ordered":                         []any{},
					"minor_to_major":                      []string{"3", "2", "1", "0"},
					"tiles":                               []any{},
					"element_size_in_bits":                "0",
					"memory_space":                        "0",
					"index_primitive_type":                "PRIMITIVE_TYPE_INVALID",
					"pointer_primitive_type":              "PRIMITIVE_TYPE_INVALID",
					"dynamic_shape_metadata_prefix_bytes": "0",
				},
				"is_dynamic_dimension": []bool{false, false, false, false},
			},
			"is_flash_attention":    true,
			"mask_type":             v.maskType,
			"dropout_rate":          0,
			"seed":                  42,
			"sliding_window_length": 0,
			"max_seg_per_batch":     1,
			"is_paged_attention":    false,
		},
	}

	fmhaConfig := config["cudnn_fmha_backend_config"].(map[string]any)
	maps.Copy(fmhaConfig, dotDimNumbers)

	bytes, err := json.Marshal(config)
	if err != nil {
		panic(err)
	}
	return string(bytes)
}

// flashFwdBackendConfig builds the forward cudnn_fmha_backend_config for q,k,v [B,S,H,D] or [B,H,S,D]
// (score matrix [B,H,S,S]). Only the scale and score-matrix dims vary with shape.
func flashFwdBackendConfig(b, h, s int, scale float64, layout compute.AxesLayout, v fmhaVariant) string {
	var dotDimNumbers map[string]any
	if layout == compute.AxesLayoutBHSD {
		dotDimNumbers = map[string]any{
			"bmm1_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"3"},
				"rhs_contracting_dimensions": []string{"3"},
				"lhs_batch_dimensions":       []string{"0", "1"},
				"rhs_batch_dimensions":       []string{"0", "1"},
			},
			"bmm2_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"3"},
				"rhs_contracting_dimensions": []string{"2"},
				"lhs_batch_dimensions":       []string{"0", "1"},
				"rhs_batch_dimensions":       []string{"0", "1"},
			},
		}
	} else { // compute.AxesLayoutBSHD
		dotDimNumbers = map[string]any{
			"bmm1_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"3"},
				"rhs_contracting_dimensions": []string{"3"},
				"lhs_batch_dimensions":       []string{"0", "2"},
				"rhs_batch_dimensions":       []string{"0", "2"},
			},
			"bmm2_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"3"},
				"rhs_contracting_dimensions": []string{"1"},
				"lhs_batch_dimensions":       []string{"0", "1"},
				"rhs_batch_dimensions":       []string{"0", "2"},
			},
		}
	}
	return flashBackendConfigV(b, h, s, scale, dotDimNumbers, v)
}

// flashBwdBackendConfig is the backward counterpart: the four backward-gemm dot_dimension_numbers.
func flashBwdBackendConfig(b, h, s int, scale float64, layout compute.AxesLayout, v fmhaVariant) string {
	var dotDimNumbers map[string]any
	if layout == compute.AxesLayoutBHSD {
		dotDimNumbers = map[string]any{
			"bmm1_grad_gemm1_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"2"},
				"rhs_contracting_dimensions": []string{"2"},
				"lhs_batch_dimensions":       []string{"0", "1"},
				"rhs_batch_dimensions":       []string{"0", "1"},
			},
			"bmm1_grad_gemm2_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"3"},
				"rhs_contracting_dimensions": []string{"2"},
				"lhs_batch_dimensions":       []string{"0", "1"},
				"rhs_batch_dimensions":       []string{"0", "1"},
			},
			"bmm2_grad_gemm1_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"2"},
				"rhs_contracting_dimensions": []string{"2"},
				"lhs_batch_dimensions":       []string{"0", "1"},
				"rhs_batch_dimensions":       []string{"0", "1"},
			},
			"bmm2_grad_gemm2_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"3"},
				"rhs_contracting_dimensions": []string{"3"},
				"lhs_batch_dimensions":       []string{"0", "1"},
				"rhs_batch_dimensions":       []string{"0", "1"},
			},
		}
	} else { // compute.AxesLayoutBSHD
		dotDimNumbers = map[string]any{
			"bmm1_grad_gemm1_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"2"},
				"rhs_contracting_dimensions": []string{"1"},
				"lhs_batch_dimensions":       []string{"0", "1"},
				"rhs_batch_dimensions":       []string{"0", "2"},
			},
			"bmm1_grad_gemm2_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"3"},
				"rhs_contracting_dimensions": []string{"1"},
				"lhs_batch_dimensions":       []string{"0", "1"},
				"rhs_batch_dimensions":       []string{"0", "2"},
			},
			"bmm2_grad_gemm1_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"2"},
				"rhs_contracting_dimensions": []string{"1"},
				"lhs_batch_dimensions":       []string{"0", "1"},
				"rhs_batch_dimensions":       []string{"0", "2"},
			},
			"bmm2_grad_gemm2_dot_dimension_numbers": map[string]any{
				"lhs_contracting_dimensions": []string{"3"},
				"rhs_contracting_dimensions": []string{"3"},
				"lhs_batch_dimensions":       []string{"0", "2"},
				"rhs_batch_dimensions":       []string{"0", "2"},
			},
		}
	}
	return flashBackendConfigV(b, h, s, scale, dotDimNumbers, v)
}

// formatScale renders a float as a JSON number (no quotes, shortest round-trip form).
func formatScale(scale float64) string {
	return strconv.FormatFloat(scale, 'g', -1, 64)
}

// flashSupported reports whether the cuDNN flash path can serve this call. cuDNN fMHA here is
// f16/bf16 (fp8 paused), BSHD-layout, equal-head, on a cuda plugin. Causality and
// per-batch seqlen padding are supported (mask_type derives from them in selectFMHAVariant);
// an explicit materialized mask is not (use seqlens instead). Anything else -> ErrNotImplemented.
func (f *Function) flashSupported(op string, qkvDType dtypes.DType, mask compute.Value, numHeads, numKVHeads, featureDim int, axesLayout compute.AxesLayout, causal bool, options *compute.ScaledDotProductAttentionConfig) error {
	_ = causal
	if !f.builder.backend.plugin.IsCUDA() {
		return errors.Wrapf(compute.ErrNotImplemented, "%s: cuDNN flash needs the cuda plugin, have %q", op, f.builder.backend.pluginName)
	}
	if !qkvDType.IsHalfPrecision() {
		return errors.Wrapf(compute.ErrNotImplemented,
			"%s: cuDNN flash needs half-precision (float16/bfloat16), got %s", op, qkvDType)
	}
	if mask != nil {
		return errors.Wrapf(compute.ErrNotImplemented,
			"%s: cuDNN flash path takes seqlens, not a materialized mask", op)
	}
	if featureDim%8 != 0 {
		return errors.Wrapf(compute.ErrNotImplemented,
			"%s: cuDNN flash path requires query/key/value last (feature) dim to be multiple of 8, got %d", op, featureDim)
	}
	if axesLayout != compute.AxesLayoutBSHD && axesLayout != compute.AxesLayoutBHSD {
		return errors.Wrapf(compute.ErrNotImplemented,
			"%s: cuDNN flash path supports only BSHD or BHSD layouts (got layout %v)", op, axesLayout)
	}
	if numKVHeads != numHeads {
		return errors.Wrapf(compute.ErrNotImplemented,
			"%s: cuDNN flash path requires equal number of q/kv heads (got layout=%v heads=%d/%d)",
			op, axesLayout, numHeads, numKVHeads)
	}
	// One of QuerySeqLen/KeyValueSeqLen set without the other is ambiguous.
	if options != nil && (options.QuerySeqLen != nil) != (options.KeyValueSeqLen != nil) {
		return errors.Wrapf(compute.ErrNotImplemented,
			"%s: cuDNN flash padding mask needs both QuerySeqLen and KeyValueSeqLen", op)
	}
	return nil
}

// dtypeOf returns the element dtype of a value (used to pick the cuDNN fmha variant before
// any bf16 cast).
func (f *Function) dtypeOf(op string, v compute.Value) (dtypes.DType, error) {
	nodes, err := f.verifyAndCastValues(op, v)
	if err != nil {
		return 0, err
	}
	return nodes[0].shape.DType, nil
}

// bshdDims returns the [B,S,H,D] (or [B,H,S,D]) dimensions of a rank-4 value depending on layout.
func (f *Function) bshdDims(op string, v compute.Value, layout compute.AxesLayout) (b, s, h, d int, err error) {
	nodes, err := f.verifyAndCastValues(op, v)
	if err != nil {
		return 0, 0, 0, 0, err
	}
	dims := nodes[0].shape.Dimensions
	if len(dims) != 4 {
		return 0, 0, 0, 0, errors.Errorf("%s: expected rank-4 tensor, got shape %v", op, dims)
	}
	if layout == compute.AxesLayoutBSHD {
		return dims[0], dims[1], dims[2], dims[3], nil
	}
	return dims[0], dims[2], dims[1], dims[3], nil
}

// validateSeqLen checks that v is a rank-1 int32 vector of length == batch.
// name is used in the error message (e.g. "QuerySeqLen"). Returns a wrapped
// error on mismatch so callers can distinguish validation failures.
func validateSeqLen(name string, v compute.Value, batch int) error {
	n, ok := v.(*Node)
	if !ok {
		return errors.Errorf("seqlen %s: expected *Node, got %T", name, v)
	}
	sh := n.shape
	if sh.DType != dtypes.Int32 {
		return errors.Errorf("seqlen %s: must be int32, got %s", name, sh.DType)
	}
	if sh.Rank() != 1 {
		return errors.Errorf("seqlen %s: must be rank-1 [B], got rank %d (shape %v)", name, sh.Rank(), sh.Dimensions)
	}
	if sh.Dimensions[0] != batch {
		return errors.Errorf("seqlen %s: length %d != batch size %d", name, sh.Dimensions[0], batch)
	}
	return nil
}

// validateBias checks that v is a rank-4 tensor of exactly [B,H,S,Skv] with the same dtype as the
// q/k/v operands (wantDType). The backend does no automatic conversion, so a mismatched dtype is a
// caller error, not silently cast. name is used in error messages.
func validateBias(name string, v compute.Value, wantDType dtypes.DType, b, h, s, skv int) error {
	n, ok := v.(*Node)
	if !ok {
		return errors.Errorf("bias %s: expected *Node, got %T", name, v)
	}
	sh := n.shape
	if sh.DType != wantDType {
		return errors.Errorf("bias %s: dtype %s must match q/k/v dtype %s (no automatic conversion)", name, sh.DType, wantDType)
	}
	if sh.Rank() != 4 {
		return errors.Errorf("bias %s: must be rank-4 [B,H,S,Skv], got rank %d (shape %v)", name, sh.Rank(), sh.Dimensions)
	}
	if d := sh.Dimensions; d[0] != b || d[1] != h || d[2] != s || d[3] != skv {
		return errors.Errorf("bias %s: shape %v != [%d,%d,%d,%d]", name, d, b, h, s, skv)
	}
	return nil
}

// fwdResultLayouts: output BHSD [3,1,2,0], stats [2,1,0], scratch u8 [0].
var fwdResultLayouts = [][]int{{3, 1, 2, 0}, {2, 1, 0}, {0}}

// bwdResultLayouts: dQ, dK, dV BHSD, scratch u8.
var bwdResultLayouts = [][]int{{3, 1, 2, 0}, {3, 1, 2, 0}, {3, 1, 2, 0}, {0}}

// FusedScaledDotProductAttention runs the cuDNN flash forward. query/key/value are [B,S,H,D]
// (BSHD), bf16. It returns the [B,S,H,D] bf16 output and the [B,H,S] f32 softmax statistics
// (log-sum-exp) the flash backward needs. The [B,H,S,S] scores never materialize. On non-cuda
// plugins or unsupported option combinations it returns ErrNotImplemented.
func (f *Function) FusedScaledDotProductAttention(query, key, value, mask compute.Value, _, _ int, axesLayout compute.AxesLayout, scale float64, causal bool, options *compute.ScaledDotProductAttentionConfig) (output compute.Value, statesForVJP []compute.Value, err error) {
	op := compute.OpTypeFusedScaledDotProductAttention.String()
	batchSize, seqLen, numHeads, featureDim, err := f.bshdDims(op, query, axesLayout)
	if err != nil {
		return nil, nil, err
	}
	_, _, numKVHeads, _, err := f.bshdDims(op, key, axesLayout)
	if err != nil {
		return nil, nil, err
	}
	qDType, err := f.dtypeOf(op, query)
	if err != nil {
		return nil, nil, err
	}
	if err = f.flashSupported(op, qDType, mask, numHeads, numKVHeads, featureDim, axesLayout, causal, options); err != nil {
		return nil, nil, err
	}
	variant, err := selectFMHAVariant(op, qDType, causal, options)
	if err != nil {
		return nil, nil, err
	}
	// Operand order cuDNN expects: q, k, v, [bias], [seqQ, seqKV]. Bias goes before seqlens
	// (bias and seqlens are mutually exclusive here; see selectFMHAVariant).
	operands := []compute.Value{query, key, value}
	operandLayouts := [][]int{{3, 2, 1, 0}, {3, 2, 1, 0}, {3, 2, 1, 0}}
	if variant.hasBias {
		if err = validateBias("Bias", options.Bias, qDType, batchSize, numHeads, seqLen, seqLen); err != nil {
			return nil, nil, err
		}
		operands = append(operands, options.Bias)
		operandLayouts = append(operandLayouts, []int{3, 2, 1, 0}) // [B,H,S,Skv] row-major
	}
	if variant.hasSeqLens {
		if err = validateSeqLen("QuerySeqLen", options.QuerySeqLen, batchSize); err != nil {
			return nil, nil, err
		}
		if err = validateSeqLen("KeyValueSeqLen", options.KeyValueSeqLen, batchSize); err != nil {
			return nil, nil, err
		}
		operands = append(operands, options.QuerySeqLen, options.KeyValueSeqLen)
		operandLayouts = append(operandLayouts, nil, nil) // int32 [B], row-major
	}
	bhsd := shapes.Make(qDType, batchSize, numHeads, seqLen, featureDim)
	stats := shapes.Make(dtypes.Float32, batchSize, numHeads, seqLen)
	scratch := shapes.Make(dtypes.Uint8, 0)
	fwdResultLayouts := [][]int{{3, 1, 2, 0}, {2, 1, 0}, {0}}
	if axesLayout == compute.AxesLayoutBHSD {
		fwdResultLayouts[0] = []int{3, 2, 1, 0}
	}
	outs, err := f.customCall(variant.fwdTarget, flashFwdBackendConfig(batchSize, numHeads, seqLen, scale, axesLayout, variant),
		operandLayouts, []shapes.Shape{bhsd, stats, scratch}, fwdResultLayouts, operands...)
	if err != nil {
		return nil, nil, err
	}
	// If axesLayout == BSHD: outs[0] is BHSD; transpose to BSHD to match the query layout.
	// If axesLayout == BHSD: outs[0] is already BHSD; no transpose needed.
	if axesLayout == compute.AxesLayoutBSHD {
		output, err = f.Transpose(outs[0], 0, 2, 1, 3)
		if err != nil {
			return nil, nil, err
		}
	} else {
		output = outs[0]
	}
	// statesForVJP carries only the f32 softmax stats; outs[2] (fwd-only scratch) is not passed to the backward.
	return output, []compute.Value{outs[1]}, nil
}

// FusedScaledDotProductAttentionVJP runs the cuDNN flash backward. It threads the softmax stats
// (statesForVJP[0], from the forward) plus the forward output and the output gradient dOutput into
// the cuDNN backward custom-call, so the [B,H,S,S] scores never materialize in the backward either.
// Returns dQuery, dKey, dValue as [B,S,H,D] bf16.
func (f *Function) FusedScaledDotProductAttentionVJP(query, key, value, mask compute.Value, _, _ int, axesLayout compute.AxesLayout, scale float64, causal bool, options *compute.ScaledDotProductAttentionConfig, output compute.Value, statesForVJP []compute.Value, dOutput compute.Value) (dQuery, dKey, dValue compute.Value, err error) {
	op := compute.OpTypeFusedScaledDotProductAttentionVJP.String()
	batchSize, seqLen, numHeads, featureDim, err := f.bshdDims(op, query, axesLayout)
	if err != nil {
		return nil, nil, nil, err
	}
	_, _, numKVHeads, _, err := f.bshdDims(op, key, axesLayout)
	if err != nil {
		return nil, nil, nil, err
	}
	if len(statesForVJP) == 0 {
		return nil, nil, nil, errors.Wrapf(compute.ErrNotImplemented, "%s: statesForVJP is empty (forward produced no fused backward state)", op)
	}
	softmaxStats := statesForVJP[0]
	qDType, err := f.dtypeOf(op, query)
	if err != nil {
		return nil, nil, nil, err
	}
	if err = f.flashSupported(op, qDType, mask, numHeads, numKVHeads, featureDim, axesLayout, causal, options); err != nil {
		return nil, nil, nil, err
	}
	variant, err := selectFMHAVariant(op, qDType, causal, options)
	if err != nil {
		return nil, nil, nil, err
	}
	// Operand order cuDNN expects: q, k, v, softmax_sum, dO, [bias], O. Bias goes at index 5,
	// before O (bias and seqlens are mutually exclusive; see selectFMHAVariant).
	operands := []compute.Value{query, key, value, softmaxStats, dOutput}
	operandLayouts := [][]int{{3, 2, 1, 0}, {3, 2, 1, 0}, {3, 2, 1, 0}, {2, 1, 0}, {3, 2, 1, 0}}
	if variant.hasBias {
		if err = validateBias("Bias", options.Bias, qDType, batchSize, numHeads, seqLen, seqLen); err != nil {
			return nil, nil, nil, err
		}
		operands = append(operands, options.Bias)
		operandLayouts = append(operandLayouts, []int{3, 2, 1, 0})
	}
	operands = append(operands, output)
	operandLayouts = append(operandLayouts, []int{3, 2, 1, 0})
	if variant.hasSeqLens {
		if err = validateSeqLen("QuerySeqLen", options.QuerySeqLen, batchSize); err != nil {
			return nil, nil, nil, err
		}
		if err = validateSeqLen("KeyValueSeqLen", options.KeyValueSeqLen, batchSize); err != nil {
			return nil, nil, nil, err
		}
		operands = append(operands, options.QuerySeqLen, options.KeyValueSeqLen)
		operandLayouts = append(operandLayouts, nil, nil)
	}
	bhsd := shapes.Make(qDType, batchSize, numHeads, seqLen, featureDim)
	scratch := shapes.Make(dtypes.Uint8, 0)
	bwdResultLayouts := [][]int{{3, 1, 2, 0}, {3, 1, 2, 0}, {3, 1, 2, 0}, {0}}
	if axesLayout == compute.AxesLayoutBHSD {
		bwdResultLayouts[0] = []int{3, 2, 1, 0}
		bwdResultLayouts[1] = []int{3, 2, 1, 0}
		bwdResultLayouts[2] = []int{3, 2, 1, 0}
	}
	resultShapes := []shapes.Shape{bhsd, bhsd, bhsd, scratch}
	if variant.hasBias {
		// The ScaleBias backward emits a dBias [B,H,S,Skv] result between dV and scratch. We do not
		// propagate it (bias is not differentiated through this op), but the slot must be declared.
		biasShape := shapes.Make(qDType, batchSize, numHeads, seqLen, seqLen)
		resultShapes = []shapes.Shape{bhsd, bhsd, bhsd, biasShape, scratch}
		bwdResultLayouts = append(bwdResultLayouts[:3:3], []int{3, 2, 1, 0}, []int{0})
	}
	grads, err := f.customCall(variant.bwdTarget, flashBwdBackendConfig(batchSize, numHeads, seqLen, scale, axesLayout, variant),
		operandLayouts, resultShapes, bwdResultLayouts, operands...)
	if err != nil {
		return nil, nil, nil, err
	}
	// If axesLayout == BSHD: grads[0..2] = dQ, dK, dV (BHSD); transpose to BSHD to match q,k,v.
	// If axesLayout == BHSD: grads[0..2] are already BHSD; no transpose needed.
	if axesLayout == compute.AxesLayoutBSHD {
		if dQuery, err = f.Transpose(grads[0], 0, 2, 1, 3); err != nil {
			return nil, nil, nil, err
		}
		if dKey, err = f.Transpose(grads[1], 0, 2, 1, 3); err != nil {
			return nil, nil, nil, err
		}
		if dValue, err = f.Transpose(grads[2], 0, 2, 1, 3); err != nil {
			return nil, nil, nil, err
		}
	} else {
		dQuery = grads[0]
		dKey = grads[1]
		dValue = grads[2]
	}
	return dQuery, dKey, dValue, nil
}