diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index cd3e9d04..0c3b7c0b 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -2881,6 +2881,11 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val headDim = query.shape[3] val seqKV = key.shape[2] + // The signature default `scale = 0f` means "use the standard + // 1/sqrt(headDim)"; applying 0 literally would flatten every softmax to + // a uniform distribution. Resolve it here. + val effectiveScale = if (scale == 0f) (1.0 / kotlin.math.sqrt(headDim.toDouble())).toFloat() else scale + val qBuf = query.data.copyToFloatArray() val kBuf = key.data.copyToFloatArray() val vBuf = value.data.copyToFloatArray() @@ -2899,7 +2904,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory for (d in 0 until headDim) { dot += qBuf[qOff + d] * kBuf[kOff + d] } - scores[qi * seqKV + ki] = dot * scale + scores[qi * seqKV + ki] = dot * effectiveScale } } diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SDPAShapeValidationTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SDPAShapeValidationTest.kt index e2f07166..bb9ee9c2 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SDPAShapeValidationTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SDPAShapeValidationTest.kt @@ -3,6 +3,7 @@ package sk.ainet.exec.tensor.ops import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertFailsWith +import kotlin.test.assertTrue import sk.ainet.context.DirectCpuExecutionContext import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor @@ -77,4 +78,31 @@ class SDPAShapeValidationTest { val out = ctx.ops.scaledDotProductAttention(q, k, v, mask = null, scale = 1f, causal = true) assertEquals(Shape(1, 2, 3, 4), out.shape) } + + @Test + fun default_scale_uses_one_over_sqrt_head_dim_not_zero() { + // Regression for #860: the default scale = 0f must be resolved to + // 1/sqrt(headDim), not applied literally (which flattens the softmax + // to a uniform average and silently discards the attention pattern). + val rng = kotlin.random.Random(42) + fun rnd(shape: Shape) = ctx.fromFloatArray(shape, FP32::class, FloatArray(shape.volume) { rng.nextFloat() }) + val headDim = 4 + val q = rnd(Shape(1, 1, 3, headDim)) + val k = rnd(Shape(1, 1, 3, headDim)) + val v = rnd(Shape(1, 1, 3, headDim)) + + val defaulted = ctx.ops.scaledDotProductAttention(q, k, v, mask = null, scale = 0f, causal = true) + val explicit = ctx.ops.scaledDotProductAttention( + q, k, v, mask = null, scale = (1.0 / kotlin.math.sqrt(headDim.toDouble())).toFloat(), causal = true, + ) + // The two must agree... + val a = defaulted.data.copyToFloatArray() + val b = explicit.data.copyToFloatArray() + for (i in a.indices) assertEquals(b[i], a[i], 1e-6f) + + // ...and must differ from the degenerate scale=0 (uniform) result. + val uniform = ctx.ops.scaledDotProductAttention(q, k, v, mask = null, scale = Float.MIN_VALUE, causal = true) + val u = uniform.data.copyToFloatArray() + assertTrue(a.indices.any { kotlin.math.abs(a[it] - u[it]) > 1e-4f }, "default scale must not equal the near-zero (uniform) result") + } }