From 0acaa2660aa30395a4520a639808ec7a1aef293a Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 14:15:55 +0200
Subject: [PATCH 01/12] Chunked actication for intra-op parallelization.

---
 internal/gobackend/activations/activations.go | 26 ++++++++++++++-----
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/internal/gobackend/activations/activations.go b/internal/gobackend/activations/activations.go
index 8bb21e8..59d0a02 100644
--- a/internal/gobackend/activations/activations.go
+++ b/internal/gobackend/activations/activations.go
@@ -168,7 +168,7 @@ func GetSwiGLU[T gotype.Supported]() SwiGLUFn[T] {
 	return fn.(SwiGLUFn[T])
 }
 
-const minParallelizeChunk = 4096
+const minParallelizeChunk = 32768
 
 // Apply applies the given activation function in-place across the slice.
 // If the slice is large and workers are available, work is parallelized in chunks.
@@ -196,9 +196,12 @@ func Execute[T gotype.Supported](backend *gobackend.Backend, act compute.Activat
 
 	n := len(in)
 	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && n > minParallelizeChunk {
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		targetChunks := max(1, numWorkers*2)
+		chunkSize := max(minParallelizeChunk, (n+targetChunks-1)/targetChunks)
 		var wg sync.WaitGroup
-		for i := 0; i < n; i += minParallelizeChunk {
-			end := min(i+minParallelizeChunk, n)
+		for i := 0; i < n; i += chunkSize {
+			end := min(i+chunkSize, n)
 			inChunk := in[i:end]
 			outChunk := out[i:end]
 			wg.Add(1)
@@ -223,7 +226,10 @@ func ExecuteSwiGLU[T gotype.Supported](backend *gobackend.Backend, in, out []T,
 
 	totalWork := numRows * hiddenDim
 	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && totalWork > minParallelizeChunk && numRows > 1 {
-		rowsPerChunk := max(minParallelizeChunk/hiddenDim, 1)
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		targetChunks := max(1, numWorkers*2)
+		chunkSize := max(minParallelizeChunk, (totalWork+targetChunks-1)/targetChunks)
+		rowsPerChunk := max(chunkSize/hiddenDim, 1)
 		var wg sync.WaitGroup
 		for r := 0; r < numRows; r += rowsPerChunk {
 			rEnd := min(r+rowsPerChunk, numRows)
@@ -337,9 +343,12 @@ func ExecuteVJP[T gotype.Supported](backend *gobackend.Backend, act compute.Acti
 
 	n := len(dOutput)
 	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && n > minParallelizeChunk {
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		targetChunks := max(1, numWorkers*2)
+		chunkSize := max(minParallelizeChunk, (n+targetChunks-1)/targetChunks)
 		var wg sync.WaitGroup
-		for i := 0; i < n; i += minParallelizeChunk {
-			end := min(i+minParallelizeChunk, n)
+		for i := 0; i < n; i += chunkSize {
+			end := min(i+chunkSize, n)
 			var yChunk, xChunk []T
 			if len(y) > 0 {
 				yChunk = y[i:end]
@@ -394,7 +403,10 @@ func ExecuteSwiGLUVJP[T gotype.Supported](backend *gobackend.Backend, x, dOutput
 
 	totalWork := numRows * hiddenDim
 	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && totalWork > minParallelizeChunk && numRows > 1 {
-		rowsPerChunk := max(minParallelizeChunk/hiddenDim, 1)
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		targetChunks := max(1, numWorkers*2)
+		chunkSize := max(minParallelizeChunk, (totalWork+targetChunks-1)/targetChunks)
+		rowsPerChunk := max(chunkSize/hiddenDim, 1)
 		var wg sync.WaitGroup
 		for r := 0; r < numRows; r += rowsPerChunk {
 			rEnd := min(r+rowsPerChunk, numRows)

From 0107094949ffab2d0ee57180dee0ffc7c71a24fb Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 14:16:30 +0200
Subject: [PATCH 02/12] Fixed assembly pack LHS to 8 rows (used by the large
 matrix multiplication).

---
 .../gobackend/dot/matmul/avx512/avx512.go     |  14 ++
 .../dot/matmul/avx512/gen_large_bf16.go       |   4 +-
 .../dot/matmul/avx512/gen_large_f16.go        |   4 +-
 .../dot/matmul/avx512/gen_large_f64.go        |   4 +-
 .../dot/matmul/avx512/internal_test.go        |   7 +
 internal/gobackend/dot/matmul/avx512/large.go |   4 +-
 .../dot/matmul/avx512/large_amd64.go          |  11 +
 .../dot/matmul/avx512/pack_amd64_float32.s    | 201 ++++++++++++++++++
 8 files changed, 241 insertions(+), 8 deletions(-)

diff --git a/internal/gobackend/dot/matmul/avx512/avx512.go b/internal/gobackend/dot/matmul/avx512/avx512.go
index 09e42c3..6542edf 100644
--- a/internal/gobackend/dot/matmul/avx512/avx512.go
+++ b/internal/gobackend/dot/matmul/avx512/avx512.go
@@ -511,6 +511,20 @@ func avx512PackLHSKernelRows4[T gotype.ScalarNotComplex](
 	contractingColsBytes := uintptr(contractingCols) * bytesPerElement
 
 	if kernelRows == 8 {
+		if AVX512UseAsm {
+			if lhsTyped, ok := any(lhs).([]float32); ok {
+				fullRows := copyRows & ^7
+				if fullRows > 0 {
+					avx512PackLHSKernelRows8Float32Asm(lhsTyped, any(panel).([]float32), lhsRowStart, lhsColStart, lhsCols, fullRows, contractingCols)
+				}
+				if fullRows == copyRows {
+					return
+				}
+				panelOffset := (fullRows / kernelRows) * contractingCols * kernelRows
+				unsafePackLHS(lhs, panel[panelOffset:], lhsRowStart+fullRows, lhsColStart, lhsCols, copyRows-fullRows, contractingCols, kernelRows)
+				return
+			}
+		}
 		unsafePackLHS(lhs, panel, lhsRowStart, lhsColStart, lhsCols, copyRows, contractingCols, kernelRows)
 		return
 	}
diff --git a/internal/gobackend/dot/matmul/avx512/gen_large_bf16.go b/internal/gobackend/dot/matmul/avx512/gen_large_bf16.go
index f543761..28299e1 100644
--- a/internal/gobackend/dot/matmul/avx512/gen_large_bf16.go
+++ b/internal/gobackend/dot/matmul/avx512/gen_large_bf16.go
@@ -236,7 +236,7 @@ func avx512LargeMatrixSliceBFloat16( //alt:bf16
 
 				isFirstContractingPanel := contractingPanelIdx == 0
 				accumulate := !isFirstContractingPanel
-				canDirectOutput := (contractingSize <= params.PanelContractingSize) && (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
+				canDirectOutput := (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
 
 				if canDirectOutput {
 					outOffset := lhsPanelRowIdx*rhsCrossSize + rhsPanelColIdx
@@ -332,7 +332,7 @@ func avx512LargeMatrixSliceBFloat16( //alt:bf16
 			// Copy accumulated results from L2 cache to outputMatrix in a single pass (only for panels that could not direct output).
 			for mIdx, lhsPanelRowIdx := 0, rowStart; lhsPanelRowIdx < rowEnd; mIdx, lhsPanelRowIdx = mIdx+1, lhsPanelRowIdx+params.LHSPanelCrossSize {
 				lhsPanelHeight := min(params.LHSPanelCrossSize, rowEnd-lhsPanelRowIdx)
-				canDirectOutput := (contractingSize <= params.PanelContractingSize) && (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
+				canDirectOutput := (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
 				if canDirectOutput {
 					continue
 				}
diff --git a/internal/gobackend/dot/matmul/avx512/gen_large_f16.go b/internal/gobackend/dot/matmul/avx512/gen_large_f16.go
index cbb5165..9ea6afe 100644
--- a/internal/gobackend/dot/matmul/avx512/gen_large_f16.go
+++ b/internal/gobackend/dot/matmul/avx512/gen_large_f16.go
@@ -236,7 +236,7 @@ func avx512LargeMatrixSliceFloat16( //alt:f16
 
 				isFirstContractingPanel := contractingPanelIdx == 0
 				accumulate := !isFirstContractingPanel
-				canDirectOutput := (contractingSize <= params.PanelContractingSize) && (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
+				canDirectOutput := (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
 
 				if canDirectOutput {
 					outOffset := lhsPanelRowIdx*rhsCrossSize + rhsPanelColIdx
@@ -332,7 +332,7 @@ func avx512LargeMatrixSliceFloat16( //alt:f16
 			// Copy accumulated results from L2 cache to outputMatrix in a single pass (only for panels that could not direct output).
 			for mIdx, lhsPanelRowIdx := 0, rowStart; lhsPanelRowIdx < rowEnd; mIdx, lhsPanelRowIdx = mIdx+1, lhsPanelRowIdx+params.LHSPanelCrossSize {
 				lhsPanelHeight := min(params.LHSPanelCrossSize, rowEnd-lhsPanelRowIdx)
-				canDirectOutput := (contractingSize <= params.PanelContractingSize) && (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
+				canDirectOutput := (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
 				if canDirectOutput {
 					continue
 				}
diff --git a/internal/gobackend/dot/matmul/avx512/gen_large_f64.go b/internal/gobackend/dot/matmul/avx512/gen_large_f64.go
index b06018c..816d0cc 100644
--- a/internal/gobackend/dot/matmul/avx512/gen_large_f64.go
+++ b/internal/gobackend/dot/matmul/avx512/gen_large_f64.go
@@ -236,7 +236,7 @@ func avx512LargeMatrixSliceFloat64( //alt:f64
 
 				isFirstContractingPanel := contractingPanelIdx == 0
 				accumulate := !isFirstContractingPanel
-				canDirectOutput := (contractingSize <= params.PanelContractingSize) && (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
+				canDirectOutput := (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
 
 				if canDirectOutput {
 					outOffset := lhsPanelRowIdx*rhsCrossSize + rhsPanelColIdx
@@ -332,7 +332,7 @@ func avx512LargeMatrixSliceFloat64( //alt:f64
 			// Copy accumulated results from L2 cache to outputMatrix in a single pass (only for panels that could not direct output).
 			for mIdx, lhsPanelRowIdx := 0, rowStart; lhsPanelRowIdx < rowEnd; mIdx, lhsPanelRowIdx = mIdx+1, lhsPanelRowIdx+params.LHSPanelCrossSize {
 				lhsPanelHeight := min(params.LHSPanelCrossSize, rowEnd-lhsPanelRowIdx)
-				canDirectOutput := (contractingSize <= params.PanelContractingSize) && (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
+				canDirectOutput := (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
 				if canDirectOutput {
 					continue
 				}
diff --git a/internal/gobackend/dot/matmul/avx512/internal_test.go b/internal/gobackend/dot/matmul/avx512/internal_test.go
index 0c685a4..5dad9b1 100644
--- a/internal/gobackend/dot/matmul/avx512/internal_test.go
+++ b/internal/gobackend/dot/matmul/avx512/internal_test.go
@@ -29,6 +29,7 @@ func TestAVX512(t *testing.T) {
 	t.Run("Pack", func(t *testing.T) {
 		t.Run("Float32", func(t *testing.T) {
 			matmultest.RunPackLHSTests(t, avx512PackLHSKernelRows4[float32], 4)
+			matmultest.RunPackLHSTests(t, avx512PackLHSKernelRows4[float32], 8)
 			matmultest.RunPackRHSTests(t, avx512PackRHSNonTransposed[float32], 32)
 			matmultest.RunApplyPackedOutputTests(t, avx512ApplyPackedOutputFloat32)
 		})
@@ -768,6 +769,12 @@ func BenchmarkAVX512(b *testing.B) {
 			defer func() { AVX512UseAsm = orig }()
 			runBenchmarkPackLHS[float32](b, "float32", avx512PackLHSKernelRows4, s.totalRows, s.totalCols, s.panelRows, s.panelCols, 4)
 		})
+		b.Run(s.name+"/Float32/Asm8", func(b *testing.B) {
+			orig := AVX512UseAsm
+			AVX512UseAsm = true
+			defer func() { AVX512UseAsm = orig }()
+			runBenchmarkPackLHS[float32](b, "float32", avx512PackLHSKernelRows4, s.totalRows, s.totalCols, s.panelRows, s.panelCols, 8)
+		})
 
 		b.Run(s.name+"/Float64/GoSIMD", func(b *testing.B) {
 			orig := AVX512UseAsm
diff --git a/internal/gobackend/dot/matmul/avx512/large.go b/internal/gobackend/dot/matmul/avx512/large.go
index a5bfb54..77ecc73 100644
--- a/internal/gobackend/dot/matmul/avx512/large.go
+++ b/internal/gobackend/dot/matmul/avx512/large.go
@@ -230,7 +230,7 @@ func avx512LargeMatrixSliceFloat32( //alt:f32
 
 				isFirstContractingPanel := contractingPanelIdx == 0
 				accumulate := !isFirstContractingPanel
-				canDirectOutput := (contractingSize <= params.PanelContractingSize) && (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
+				canDirectOutput := (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
 
 				if canDirectOutput {
 					outOffset := lhsPanelRowIdx*rhsCrossSize + rhsPanelColIdx
@@ -326,7 +326,7 @@ func avx512LargeMatrixSliceFloat32( //alt:f32
 			// Copy accumulated results from L2 cache to outputMatrix in a single pass (only for panels that could not direct output).
 			for mIdx, lhsPanelRowIdx := 0, rowStart; lhsPanelRowIdx < rowEnd; mIdx, lhsPanelRowIdx = mIdx+1, lhsPanelRowIdx+params.LHSPanelCrossSize {
 				lhsPanelHeight := min(params.LHSPanelCrossSize, rowEnd-lhsPanelRowIdx)
-				canDirectOutput := (contractingSize <= params.PanelContractingSize) && (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
+				canDirectOutput := (lhsPanelHeight%params.LHSL1KernelRows == 0) && (rhsPanelWidth%params.RHSL1KernelCols == 0)
 				if canDirectOutput {
 					continue
 				}
diff --git a/internal/gobackend/dot/matmul/avx512/large_amd64.go b/internal/gobackend/dot/matmul/avx512/large_amd64.go
index 10ae3f6..e0759a7 100644
--- a/internal/gobackend/dot/matmul/avx512/large_amd64.go
+++ b/internal/gobackend/dot/matmul/avx512/large_amd64.go
@@ -34,6 +34,17 @@ func avx512PackLHSKernelRows4Float32Asm(
 	copyRows, contractingCols int,
 )
 
+// avx512PackLHSKernelRows8Float32Asm packs 8 rows of float32 LHS matrix in strips of 8 into panel.
+// Only full 8-row strips are packed; any remaining partial strip is handled by the caller.
+// Defined in pack_amd64_float32.s.
+//
+//go:noescape
+func avx512PackLHSKernelRows8Float32Asm(
+	lhs, panel []float32,
+	lhsRowStart, lhsColStart, lhsCols,
+	copyRows, contractingCols int,
+)
+
 // avx512LargeKernelFloat16Asm is the assembly implementation of the 8 rows x 32 cols GEMM microkernel for Float16.
 // Defined in avx512_large_amd64_float16.s.
 //
diff --git a/internal/gobackend/dot/matmul/avx512/pack_amd64_float32.s b/internal/gobackend/dot/matmul/avx512/pack_amd64_float32.s
index 05bca17..4cb03cf 100644
--- a/internal/gobackend/dot/matmul/avx512/pack_amd64_float32.s
+++ b/internal/gobackend/dot/matmul/avx512/pack_amd64_float32.s
@@ -132,3 +132,204 @@ next_strip:
 done:
 	VZEROUPPER
 	RET
+
+// func avx512PackLHSKernelRows8Float32Asm(
+//     lhs, panel []float32,
+//     lhsRowStart, lhsColStart, lhsCols,
+//     copyRows, contractingCols int)
+TEXT ·avx512PackLHSKernelRows8Float32Asm(SB), NOSPLIT, $0-88
+	MOVQ lhs_base+0(FP), R8              // R8 = lhsBasePtr
+	MOVQ panel_base+24(FP), R9           // R9 = panelBasePtr
+	MOVQ lhsRowStart+48(FP), R10         // R10 = lhsRowStart
+	MOVQ lhsColStart+56(FP), R11         // R11 = lhsColStart
+	MOVQ lhsCols+64(FP), R12             // R12 = lhsCols
+	MOVQ copyRows+72(FP), R13            // R13 = copyRows
+	MOVQ contractingCols+80(FP), R14     // R14 = contractingCols
+
+	CMPQ R13, $8
+	JL done8
+	TESTQ R14, R14
+	JLE done8
+
+	SHLQ $2, R12                         // R12 = lhsStrideBytes = lhsCols * 4
+	SHLQ $2, R11                         // R11 = lhsColStartBytes = lhsColStart * 4
+
+	// R15 = contractingCols16 = contractingCols & ~15
+	MOVQ R14, R15
+	ANDQ $~15, R15
+
+	XORQ AX, AX                          // AX = stripRowIdx = 0
+	ANDQ $~7, R13                        // R13 = fullStripLimit = copyRows & ~7
+
+loop_strip8:
+	CMPQ AX, R13
+	JGE done8
+
+	// Compute base pointer for row 0 of this strip:
+	// SI = lhsBasePtr + (lhsRowStart + stripRowIdx) * lhsStrideBytes + lhsColStartBytes
+	MOVQ R10, SI
+	ADDQ AX, SI
+	IMULQ R12, SI
+	ADDQ R8, SI
+	ADDQ R11, SI
+
+	LEAQ (SI)(R12*1), DX                 // DX = row 1
+	LEAQ (DX)(R12*1), DI                 // DI = row 2
+	LEAQ (DI)(R12*1), CX                 // CX = row 3
+
+	// BP = 4 * lhsStrideBytes (used to access rows 4, 5, 6, 7)
+	LEAQ (R12*4), BP
+
+	XORQ BX, BX                          // BX = colIdx = 0
+
+	TESTQ R15, R15
+	JZ check_tail_cols8
+
+	PCALIGN $32
+loop_cols8:
+	CMPQ BX, R15
+	JGE check_tail_cols8
+
+	// Load 16 float32s from rows 0, 1, 2, 3
+	VMOVDQU32 (SI), Z0
+	VMOVDQU32 (DX), Z1
+	VMOVDQU32 (DI), Z2
+	VMOVDQU32 (CX), Z3
+
+	// Stage 1 for rows 0..3: 32-bit unpack (intra-128-bit lane)
+	VUNPCKLPS Z1, Z0, Z8
+	VUNPCKHPS Z1, Z0, Z9
+	VUNPCKLPS Z3, Z2, Z10
+	VUNPCKHPS Z3, Z2, Z11
+
+	// Stage 2 for rows 0..3: 64-bit unpack (intra-128-bit lane)
+	VUNPCKLPD Z10, Z8, Z12
+	VUNPCKHPD Z10, Z8, Z13
+	VUNPCKLPD Z11, Z9, Z14
+	VUNPCKHPD Z11, Z9, Z15
+
+	// Stage 3 for rows 0..3: 128-bit cross-lane shuffle
+	VSHUFI32X4 $0x44, Z13, Z12, Z16
+	VSHUFI32X4 $0x44, Z15, Z14, Z17
+	VSHUFI32X4 $0xEE, Z13, Z12, Z18
+	VSHUFI32X4 $0xEE, Z15, Z14, Z19
+
+	// Stage 4 for rows 0..3: Assemble Out0..Out3
+	VSHUFI32X4 $0x88, Z17, Z16, Z24      // Z24 = cols 0..3, rows 0..3
+	VSHUFI32X4 $0xDD, Z17, Z16, Z25      // Z25 = cols 4..7, rows 0..3
+	VSHUFI32X4 $0x88, Z19, Z18, Z26      // Z26 = cols 8..11, rows 0..3
+	VSHUFI32X4 $0xDD, Z19, Z18, Z27      // Z27 = cols 12..15, rows 0..3
+
+	// Load 16 float32s from rows 4, 5, 6, 7
+	VMOVDQU32 (0)(SI)(BP*1), Z4
+	VMOVDQU32 (0)(DX)(BP*1), Z5
+	VMOVDQU32 (0)(DI)(BP*1), Z6
+	VMOVDQU32 (0)(CX)(BP*1), Z7
+
+	// Stage 1 for rows 4..7: 32-bit unpack (intra-128-bit lane)
+	VUNPCKLPS Z5, Z4, Z8
+	VUNPCKHPS Z5, Z4, Z9
+	VUNPCKLPS Z7, Z6, Z10
+	VUNPCKHPS Z7, Z6, Z11
+
+	// Stage 2 for rows 4..7: 64-bit unpack (intra-128-bit lane)
+	VUNPCKLPD Z10, Z8, Z12
+	VUNPCKHPD Z10, Z8, Z13
+	VUNPCKLPD Z11, Z9, Z14
+	VUNPCKHPD Z11, Z9, Z15
+
+	// Stage 3 for rows 4..7: 128-bit cross-lane shuffle
+	VSHUFI32X4 $0x44, Z13, Z12, Z16
+	VSHUFI32X4 $0x44, Z15, Z14, Z17
+	VSHUFI32X4 $0xEE, Z13, Z12, Z18
+	VSHUFI32X4 $0xEE, Z15, Z14, Z19
+
+	// Stage 4 for rows 4..7: Assemble Out4..Out7
+	VSHUFI32X4 $0x88, Z17, Z16, Z28      // Z28 = cols 0..3, rows 4..7
+	VSHUFI32X4 $0xDD, Z17, Z16, Z29      // Z29 = cols 4..7, rows 4..7
+	VSHUFI32X4 $0x88, Z19, Z18, Z30      // Z30 = cols 8..11, rows 4..7
+	VSHUFI32X4 $0xDD, Z19, Z18, Z31      // Z31 = cols 12..15, rows 4..7
+
+	// Interleave rows 0..3 and rows 4..7 for each column group:
+	// Cols 0..3 -> output strips cols 0, 1 (Z0) and cols 2, 3 (Z1)
+	VSHUFI32X4 $0x44, Z28, Z24, Z0
+	VSHUFI32X4 $0xD8, Z0, Z0, Z0
+	VSHUFI32X4 $0xEE, Z28, Z24, Z1
+	VSHUFI32X4 $0xD8, Z1, Z1, Z1
+
+	// Cols 4..7 -> output strips cols 4, 5 (Z2) and cols 6, 7 (Z3)
+	VSHUFI32X4 $0x44, Z29, Z25, Z2
+	VSHUFI32X4 $0xD8, Z2, Z2, Z2
+	VSHUFI32X4 $0xEE, Z29, Z25, Z3
+	VSHUFI32X4 $0xD8, Z3, Z3, Z3
+
+	// Cols 8..11 -> output strips cols 8, 9 (Z4) and cols 10, 11 (Z5)
+	VSHUFI32X4 $0x44, Z30, Z26, Z4
+	VSHUFI32X4 $0xD8, Z4, Z4, Z4
+	VSHUFI32X4 $0xEE, Z30, Z26, Z5
+	VSHUFI32X4 $0xD8, Z5, Z5, Z5
+
+	// Cols 12..15 -> output strips cols 12, 13 (Z6) and cols 14, 15 (Z7)
+	VSHUFI32X4 $0x44, Z31, Z27, Z6
+	VSHUFI32X4 $0xD8, Z6, Z6, Z6
+	VSHUFI32X4 $0xEE, Z31, Z27, Z7
+	VSHUFI32X4 $0xD8, Z7, Z7, Z7
+
+	// Store 8 output strips (512 bytes = 16 columns * 8 rows * 4 bytes)
+	VMOVDQU32 Z0, (R9)
+	VMOVDQU32 Z1, 64(R9)
+	VMOVDQU32 Z2, 128(R9)
+	VMOVDQU32 Z3, 192(R9)
+	VMOVDQU32 Z4, 256(R9)
+	VMOVDQU32 Z5, 320(R9)
+	VMOVDQU32 Z6, 384(R9)
+	VMOVDQU32 Z7, 448(R9)
+
+	ADDQ $64, SI
+	ADDQ $64, DX
+	ADDQ $64, DI
+	ADDQ $64, CX
+	ADDQ $512, R9                        // panelPtr += 512 bytes
+	ADDQ $16, BX
+	JMP loop_cols8
+
+check_tail_cols8:
+	CMPQ BX, R14
+	JGE next_strip8
+
+loop_tail_cols8:
+	VMOVSS (SI), X0
+	VMOVSS X0, (R9)
+	VMOVSS (DX), X0
+	VMOVSS X0, 4(R9)
+	VMOVSS (DI), X0
+	VMOVSS X0, 8(R9)
+	VMOVSS (CX), X0
+	VMOVSS X0, 12(R9)
+	VMOVSS (0)(SI)(BP*1), X0
+	VMOVSS X0, 16(R9)
+	VMOVSS (0)(DX)(BP*1), X0
+	VMOVSS X0, 20(R9)
+	VMOVSS (0)(DI)(BP*1), X0
+	VMOVSS X0, 24(R9)
+	VMOVSS (0)(CX)(BP*1), X0
+	VMOVSS X0, 28(R9)
+
+	ADDQ $4, SI
+	ADDQ $4, DX
+	ADDQ $4, DI
+	ADDQ $4, CX
+	ADDQ $32, R9
+	INCQ BX
+	CMPQ BX, R14
+	JL loop_tail_cols8
+
+next_strip8:
+	MOVQ lhs_base+0(FP), R8              // reload R8 = lhsBasePtr for next strip
+	ADDQ $8, AX
+	JMP loop_strip8
+
+done8:
+	VZEROUPPER
+	RET
+

From 3fc445c1e2b942de5e4cb239d29d5cd44509dba0 Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 14:17:19 +0200
Subject: [PATCH 03/12] PackLHS assembly.

---
 internal/gobackend/dot/matmul/packing.go      | 71 +++++++++++++++++++
 internal/gobackend/dot/matmul/packing_test.go |  9 ++-
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/internal/gobackend/dot/matmul/packing.go b/internal/gobackend/dot/matmul/packing.go
index 4b7acd6..cacb9a9 100644
--- a/internal/gobackend/dot/matmul/packing.go
+++ b/internal/gobackend/dot/matmul/packing.go
@@ -295,6 +295,77 @@ func unsafePackLHS[T gotype.ScalarNotComplex](
 			}
 			panelPtr += stripSizeBytes
 		}
+	case kernelRows == 8:
+		for stripRowIdx := 0; stripRowIdx < fullStripsRows; stripRowIdx += kernelRows {
+			srcIdxBase := ((lhsRowStart + stripRowIdx) * lhsCols) + lhsColStart
+			pSrcBase := lhsPtr + uintptr(srcIdxBase)*elemSize
+
+			pSrc0 := pSrcBase
+			pSrc1 := pSrc0 + lhsColsBytes
+			pSrc2 := pSrc1 + lhsColsBytes
+			pSrc3 := pSrc2 + lhsColsBytes
+			pSrc4 := pSrc3 + lhsColsBytes
+			pSrc5 := pSrc4 + lhsColsBytes
+			pSrc6 := pSrc5 + lhsColsBytes
+			pSrc7 := pSrc6 + lhsColsBytes
+
+			pDst := panelPtr
+
+			col := 0
+			for ; col+1 < contractingCols; col += 2 {
+				v0 := *(*T)(unsafe.Pointer(pSrc0))
+				v1 := *(*T)(unsafe.Pointer(pSrc1))
+				v2 := *(*T)(unsafe.Pointer(pSrc2))
+				v3 := *(*T)(unsafe.Pointer(pSrc3))
+				v4 := *(*T)(unsafe.Pointer(pSrc4))
+				v5 := *(*T)(unsafe.Pointer(pSrc5))
+				v6 := *(*T)(unsafe.Pointer(pSrc6))
+				v7 := *(*T)(unsafe.Pointer(pSrc7))
+				*(*[8]T)(unsafe.Pointer(pDst)) = [8]T{v0, v1, v2, v3, v4, v5, v6, v7}
+
+				v0_1 := *(*T)(unsafe.Pointer(pSrc0 + elemSize))
+				v1_1 := *(*T)(unsafe.Pointer(pSrc1 + elemSize))
+				v2_1 := *(*T)(unsafe.Pointer(pSrc2 + elemSize))
+				v3_1 := *(*T)(unsafe.Pointer(pSrc3 + elemSize))
+				v4_1 := *(*T)(unsafe.Pointer(pSrc4 + elemSize))
+				v5_1 := *(*T)(unsafe.Pointer(pSrc5 + elemSize))
+				v6_1 := *(*T)(unsafe.Pointer(pSrc6 + elemSize))
+				v7_1 := *(*T)(unsafe.Pointer(pSrc7 + elemSize))
+				*(*[8]T)(unsafe.Pointer(pDst + 8*elemSize)) = [8]T{v0_1, v1_1, v2_1, v3_1, v4_1, v5_1, v6_1, v7_1}
+
+				pSrc0 += 2 * elemSize
+				pSrc1 += 2 * elemSize
+				pSrc2 += 2 * elemSize
+				pSrc3 += 2 * elemSize
+				pSrc4 += 2 * elemSize
+				pSrc5 += 2 * elemSize
+				pSrc6 += 2 * elemSize
+				pSrc7 += 2 * elemSize
+				pDst += 16 * elemSize
+			}
+			for ; col < contractingCols; col++ {
+				v0 := *(*T)(unsafe.Pointer(pSrc0))
+				v1 := *(*T)(unsafe.Pointer(pSrc1))
+				v2 := *(*T)(unsafe.Pointer(pSrc2))
+				v3 := *(*T)(unsafe.Pointer(pSrc3))
+				v4 := *(*T)(unsafe.Pointer(pSrc4))
+				v5 := *(*T)(unsafe.Pointer(pSrc5))
+				v6 := *(*T)(unsafe.Pointer(pSrc6))
+				v7 := *(*T)(unsafe.Pointer(pSrc7))
+				*(*[8]T)(unsafe.Pointer(pDst)) = [8]T{v0, v1, v2, v3, v4, v5, v6, v7}
+
+				pSrc0 += elemSize
+				pSrc1 += elemSize
+				pSrc2 += elemSize
+				pSrc3 += elemSize
+				pSrc4 += elemSize
+				pSrc5 += elemSize
+				pSrc6 += elemSize
+				pSrc7 += elemSize
+				pDst += 8 * elemSize
+			}
+			panelPtr += stripSizeBytes
+		}
 
 	default:
 		// Larger values of kernelRows must be multiple of 4.
diff --git a/internal/gobackend/dot/matmul/packing_test.go b/internal/gobackend/dot/matmul/packing_test.go
index 1ef4a6f..e231ae7 100644
--- a/internal/gobackend/dot/matmul/packing_test.go
+++ b/internal/gobackend/dot/matmul/packing_test.go
@@ -234,7 +234,7 @@ func runApplyPackedOutputTests[T gotype.NumericNotComplex](t *testing.T, applyFn
 }
 
 func TestUnsafe(t *testing.T) {
-	for _, kernelRows := range []int{2, 4, 16, 32} {
+	for _, kernelRows := range []int{2, 4, 8, 16, 32} {
 		t.Run(fmt.Sprintf("PackLHS_kernelRows=%d", kernelRows), func(t *testing.T) {
 			runPackLHSTests(t, unsafePackLHS[float32], kernelRows)
 		})
@@ -322,6 +322,13 @@ func BenchmarkNoSIMD(b *testing.B) {
 		runBenchmarkPackLHS[bfloat16.BFloat16](b, "standard/bfloat16", packLHS, totalRows, totalCols, panelRows, panelCols, kernelRows)
 		runBenchmarkPackLHS[bfloat16.BFloat16](b, "unsafe/bfloat16", unsafePackLHS, totalRows, totalCols, panelRows, panelCols, kernelRows)
 	})
+	b.Run("PackLHS/kernelRows=8", func(b *testing.B) {
+		kernelRows := 8
+		runBenchmarkPackLHS[float32](b, "standard/float32", packLHS, totalRows, totalCols, panelRows, panelCols, kernelRows)
+		runBenchmarkPackLHS[float32](b, "unsafe/float32", unsafePackLHS, totalRows, totalCols, panelRows, panelCols, kernelRows)
+		runBenchmarkPackLHS[bfloat16.BFloat16](b, "standard/bfloat16", packLHS, totalRows, totalCols, panelRows, panelCols, kernelRows)
+		runBenchmarkPackLHS[bfloat16.BFloat16](b, "unsafe/bfloat16", unsafePackLHS, totalRows, totalCols, panelRows, panelCols, kernelRows)
+	})
 	b.Run("PackLHS/kernelRows=32", func(b *testing.B) {
 		kernelRows := 32
 		runBenchmarkPackLHS[float32](b, "standard/float32", packLHS, totalRows, totalCols, panelRows, panelCols, kernelRows)

From c6ed31c9f7f5ef6c15d3d4efa9c0ceb7ebd987e0 Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 14:18:23 +0200
Subject: [PATCH 04/12] Chunked and intra-op parallelized epilogue of MatMul.

---
 internal/gobackend/dot/matmul/epilogue.go | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/internal/gobackend/dot/matmul/epilogue.go b/internal/gobackend/dot/matmul/epilogue.go
index 2d2f986..df87e35 100644
--- a/internal/gobackend/dot/matmul/epilogue.go
+++ b/internal/gobackend/dot/matmul/epilogue.go
@@ -48,8 +48,10 @@ func ApplyEpilogue[T any](
 
 	// Fast path: No bias, only activation. We can apply activation on the whole buffer or in chunks.
 	if !hasBias && act != nil {
-		if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && len(output) > 8192 {
-			chunkSize := 4096
+		if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && len(output) > 32768 {
+			numWorkers := backend.Workers.AdjustedMaxParallelism()
+			targetChunks := max(1, numWorkers*2)
+			chunkSize := max(16384, (len(output)+targetChunks-1)/targetChunks)
 			var wg sync.WaitGroup
 			for i := 0; i < len(output); i += chunkSize {
 				end := min(i+chunkSize, len(output))
@@ -81,8 +83,11 @@ func ApplyEpilogue[T any](
 	}
 
 	totalElements := numRows * rowSize
-	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && numRows > 1 && totalElements > 8192 {
-		rowsPerChunk := max(1, 4096/rowSize)
+	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && numRows > 1 && totalElements > 32768 {
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		targetChunks := max(1, numWorkers*2)
+		chunkSize := max(16384, (totalElements+targetChunks-1)/targetChunks)
+		rowsPerChunk := max(1, chunkSize/rowSize)
 		var wg sync.WaitGroup
 		for r := 0; r < numRows; r += rowsPerChunk {
 			rStart := r
@@ -103,6 +108,9 @@ func addBias[T any](row, bias []T) {
 	switch r := any(row).(type) {
 	case []float32:
 		b := any(bias).([]float32)
+		if addBiasFloat32Arch(r, b) {
+			return
+		}
 		for i, val := range b {
 			r[i] += val
 		}

From 356a343cac5c3790f7e3b388bac4d10543c30743 Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 14:19:14 +0200
Subject: [PATCH 05/12] Chunked and parallelized SoftMax.

---
 internal/gobackend/fusedops/layernorm_simd.go | 74 ++++++++++++++++--
 internal/gobackend/fusedops/softmax.go        | 75 ++++++++++++-------
 2 files changed, 115 insertions(+), 34 deletions(-)

diff --git a/internal/gobackend/fusedops/layernorm_simd.go b/internal/gobackend/fusedops/layernorm_simd.go
index e5b6a42..d167613 100644
--- a/internal/gobackend/fusedops/layernorm_simd.go
+++ b/internal/gobackend/fusedops/layernorm_simd.go
@@ -7,6 +7,7 @@ package fusedops
 import (
 	"math"
 	"simd"
+	"sync"
 	"unsafe"
 
 	"github.com/gomlx/compute"
@@ -68,14 +69,71 @@ func execFusedLayerNormSIMD(backend *gobackend.Backend, node *gobackend.Node, in
 	}
 
 	outerSize := input.RawShape.Size() / normSize
-	if archFn := gobackend.GetLayerNormTrailingArchDispatcher(); archFn != nil {
-		var gPtr, bPtr unsafe.Pointer
-		if gamma != nil {
-			gPtr = gamma.UnsafePointer()
-		}
-		if beta != nil {
-			bPtr = beta.UnsafePointer()
-		}
+	archFn := gobackend.GetLayerNormTrailingArchDispatcher()
+	var gPtr, bPtr unsafe.Pointer
+	if gamma != nil {
+		gPtr = gamma.UnsafePointer()
+	}
+	if beta != nil {
+		bPtr = beta.UnsafePointer()
+	}
+
+	totalElements := input.RawShape.Size()
+	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && outerSize > 1 && totalElements > 16384 {
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		targetChunks := min(outerSize, max(1, numWorkers*2))
+		outerPerChunk := max(1, (outerSize+targetChunks-1)/targetChunks)
+		inPtr := input.UnsafePointer()
+		outPtr := output.UnsafePointer()
+		elemSize := dtype.Size()
+		var wg sync.WaitGroup
+
+		for start := 0; start < outerSize; start += outerPerChunk {
+			chunkOuter := min(start+outerPerChunk, outerSize) - start
+			rStart := start
+			wg.Add(1)
+			backend.Workers.WaitToStart(func() {
+				defer wg.Done()
+				rowOffset := rStart * normSize
+				if archFn != nil {
+					chunkIn := unsafe.Pointer(uintptr(inPtr) + uintptr(rowOffset*elemSize))
+					chunkOut := unsafe.Pointer(uintptr(outPtr) + uintptr(rowOffset*elemSize))
+					if archFn(chunkIn, chunkOut, gPtr, bPtr, chunkOuter, normSize, data.epsilon, dtype) {
+						return
+					}
+				}
+				// Fallback path
+				switch dtype {
+				case dtypes.Float32:
+					var gData, bData []float32
+					if gamma != nil {
+						gData = gamma.Flat.([]float32)
+					}
+					if beta != nil {
+						bData = beta.Flat.([]float32)
+					}
+					inSlice := input.Flat.([]float32)[rowOffset : rowOffset+chunkOuter*normSize]
+					outSlice := output.Flat.([]float32)[rowOffset : rowOffset+chunkOuter*normSize]
+					simdLayerNormTrailingAxesFloat32(inSlice, outSlice, gData, bData, normSize, data.epsilon)
+				case dtypes.Float64:
+					var gData, bData []float64
+					if gamma != nil {
+						gData = gamma.Flat.([]float64)
+					}
+					if beta != nil {
+						bData = beta.Flat.([]float64)
+					}
+					inSlice := input.Flat.([]float64)[rowOffset : rowOffset+chunkOuter*normSize]
+					outSlice := output.Flat.([]float64)[rowOffset : rowOffset+chunkOuter*normSize]
+					simdLayerNormTrailingAxesFloat64(inSlice, outSlice, gData, bData, normSize, data.epsilon)
+				}
+			})
+		}
+		wg.Wait()
+		return output, nil
+	}
+
+	if archFn != nil {
 		if archFn(input.UnsafePointer(), output.UnsafePointer(), gPtr, bPtr, outerSize, normSize, data.epsilon, dtype) {
 			return output, nil
 		}
diff --git a/internal/gobackend/fusedops/softmax.go b/internal/gobackend/fusedops/softmax.go
index 98f8c32..8b9ee99 100644
--- a/internal/gobackend/fusedops/softmax.go
+++ b/internal/gobackend/fusedops/softmax.go
@@ -2,6 +2,7 @@ package fusedops
 
 import (
 	"math"
+	"sync"
 
 	"github.com/gomlx/compute"
 	"github.com/gomlx/compute/dtypes"
@@ -59,9 +60,9 @@ func execFusedSoftmax(backend *gobackend.Backend, node *gobackend.Node, inputs [
 
 	switch input.RawShape.DType {
 	case dtypes.Float32:
-		fusedSoftmax(input.Flat.([]float32), output.Flat.([]float32), axis, node.Shape, fastmath.Exp32)
+		fusedSoftmax(backend, input.Flat.([]float32), output.Flat.([]float32), axis, node.Shape, fastmath.Exp32)
 	case dtypes.Float64:
-		fusedSoftmax(input.Flat.([]float64), output.Flat.([]float64), axis, node.Shape, math.Exp)
+		fusedSoftmax(backend, input.Flat.([]float64), output.Flat.([]float64), axis, node.Shape, math.Exp)
 	default:
 		return nil, errors.Wrapf(compute.ErrNotImplemented, "FusedSoftmax: dtype %s", input.RawShape.DType)
 	}
@@ -85,35 +86,57 @@ func fusedSoftmaxComputeAxisStrides(shape shapes.Shape, axis int) (outerSize, ax
 	return
 }
 
-func fusedSoftmax[T float32 | float64](input, output []T, axis int, shape shapes.Shape, expFn func(T) T) {
+func fusedSoftmax[T float32 | float64](backend *gobackend.Backend, input, output []T, axis int, shape shapes.Shape, expFn func(T) T) {
 	outerSize, axisSize, innerSize := fusedSoftmaxComputeAxisStrides(shape, axis)
-	for outer := range outerSize {
-		for inner := range innerSize {
-			baseIdx := outer*axisSize*innerSize + inner
-
-			// Pass 1: Find max.
-			maxVal := T(math.Inf(-1))
-			for i := range axisSize {
-				idx := baseIdx + i*innerSize
-				if input[idx] > maxVal {
-					maxVal = input[idx]
+	totalElements := outerSize * axisSize * innerSize
+
+	processOuter := func(start, end int) {
+		for outer := start; outer < end; outer++ {
+			for inner := range innerSize {
+				baseIdx := outer*axisSize*innerSize + inner
+
+				// Pass 1: Find max.
+				maxVal := T(math.Inf(-1))
+				for i := range axisSize {
+					idx := baseIdx + i*innerSize
+					if input[idx] > maxVal {
+						maxVal = input[idx]
+					}
 				}
-			}
 
-			// Pass 2: Exp and sum.
-			var sum T
-			for i := range axisSize {
-				idx := baseIdx + i*innerSize
-				output[idx] = expFn(input[idx] - maxVal)
-				sum += output[idx]
-			}
+				// Pass 2: Exp and sum.
+				var sum T
+				for i := range axisSize {
+					idx := baseIdx + i*innerSize
+					output[idx] = expFn(input[idx] - maxVal)
+					sum += output[idx]
+				}
 
-			// Pass 3: Normalize.
-			invSum := 1.0 / sum
-			for i := range axisSize {
-				idx := baseIdx + i*innerSize
-				output[idx] *= invSum
+				// Pass 3: Normalize.
+				invSum := 1.0 / sum
+				for i := range axisSize {
+					idx := baseIdx + i*innerSize
+					output[idx] *= invSum
+				}
 			}
 		}
 	}
+
+	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && outerSize > 1 && totalElements > 16384 {
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		targetChunks := min(outerSize, max(1, numWorkers*2))
+		outerPerChunk := max(1, (outerSize+targetChunks-1)/targetChunks)
+		var wg sync.WaitGroup
+		for start := 0; start < outerSize; start += outerPerChunk {
+			end := min(start+outerPerChunk, outerSize)
+			wg.Add(1)
+			backend.Workers.WaitToStart(func() {
+				processOuter(start, end)
+				wg.Done()
+			})
+		}
+		wg.Wait()
+	} else {
+		processOuter(0, outerSize)
+	}
 }

From 19d86d07a7d3bb8fc48b5ee0d84b16035e446ded Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 14:20:07 +0200
Subject: [PATCH 06/12] Fast-path for Transpose 4D.

---
 internal/gobackend/ops/ops_test.go       |   2 +-
 internal/gobackend/ops/transpose.go      | 331 ++++++++++++++++++++++-
 internal/gobackend/ops/transpose_test.go | 116 ++++++++
 3 files changed, 440 insertions(+), 9 deletions(-)

diff --git a/internal/gobackend/ops/ops_test.go b/internal/gobackend/ops/ops_test.go
index 6b6dfed..d6d69e4 100644
--- a/internal/gobackend/ops/ops_test.go
+++ b/internal/gobackend/ops/ops_test.go
@@ -39,7 +39,7 @@ func teardown() {
 	backend.Finalize()
 }
 
-func makeBuffer(t *testing.T, shape shapes.Shape, flat any) *gobackend.Buffer {
+func makeBuffer(t testing.TB, shape shapes.Shape, flat any) *gobackend.Buffer {
 	t.Helper()
 	computeBuf, err := backend.BufferFromFlatData(0, flat, shape)
 	if err != nil {
diff --git a/internal/gobackend/ops/transpose.go b/internal/gobackend/ops/transpose.go
index ce1df49..2f36041 100644
--- a/internal/gobackend/ops/transpose.go
+++ b/internal/gobackend/ops/transpose.go
@@ -4,6 +4,7 @@ package ops
 
 import (
 	"slices"
+	"sync"
 
 	"github.com/gomlx/compute"
 	"github.com/gomlx/compute/dtypes"
@@ -48,6 +49,21 @@ func init() {
 	TransposeDTypeMap.Register(dtypes.BFloat16, gobackend.PriorityTyped, execTransposeBFloat16)
 }
 
+// ExecuteTranspose transposes operand into output given permutations, using fast paths if available.
+func ExecuteTranspose(backend *gobackend.Backend, operand, output *gobackend.Buffer, permutations []int) {
+	if dispatchFastTranspose(backend, operand, output, permutations) {
+		return
+	}
+	it := NewTransposeIterator(operand.RawShape, permutations)
+	dtype := output.RawShape.DType
+	tmpAny, tmpErr := TransposeDTypeMap.Get(dtype)
+	if tmpErr != nil {
+		panic(tmpErr)
+	}
+	transposeFn := tmpAny.(func(operand, output *gobackend.Buffer, it *TransposeIterator))
+	transposeFn(operand, output, it)
+}
+
 // execTranspose implements Transpose.
 // The output will have: output.Shape.Dimension[ii] = operand.Shape.Dimension[permutations[i]].
 func execTranspose(backend *gobackend.Backend, node *gobackend.Node, inputs []*gobackend.Buffer, inputsOwned []bool) (*gobackend.Buffer, error) {
@@ -63,17 +79,316 @@ func execTranspose(backend *gobackend.Backend, node *gobackend.Node, inputs []*g
 	if backend.NoOps {
 		return output, nil
 	}
-	it := NewTransposeIterator(operand.RawShape, permutations)
-	dtype := node.Shape.DType
-	tmpAny, tmpErr := TransposeDTypeMap.Get(dtype)
-	if tmpErr != nil {
-		panic(tmpErr)
-	}
-	transposeFn := tmpAny.(func(operand, output *gobackend.Buffer, it *TransposeIterator))
-	transposeFn(operand, output, it)
+	ExecuteTranspose(backend, operand, output, permutations)
 	return output, nil
 }
 
+func isIdentityPermutation(permutations []int) bool {
+	for i, p := range permutations {
+		if p != i {
+			return false
+		}
+	}
+	return true
+}
+
+func dispatchFastTranspose(backend *gobackend.Backend, operand, output *gobackend.Buffer, permutations []int) bool {
+	dims := operand.RawShape.Dimensions
+	rank := len(dims)
+	if rank == 0 || operand.RawShape.Size() == 0 {
+		return true
+	}
+
+	// Case 0: Identity permutation
+	if isIdentityPermutation(permutations) {
+		return dispatchIdentity(operand, output)
+	}
+
+	// Case 1: 4D permutation [0, 2, 1, 3] or 3D [1, 0, 2] (multi-head attention transpose)
+	if (rank == 4 && permutations[0] == 0 && permutations[1] == 2 && permutations[2] == 1 && permutations[3] == 3) ||
+		(rank == 3 && permutations[0] == 1 && permutations[1] == 0 && permutations[2] == 2) {
+		var bDim, sDim, hDim, dDim int
+		if rank == 4 {
+			bDim, sDim, hDim, dDim = dims[0], dims[1], dims[2], dims[3]
+		} else {
+			bDim, sDim, hDim, dDim = 1, dims[0], dims[1], dims[2]
+		}
+		return dispatchTranspose4D_0213(backend, operand, output, bDim, sDim, hDim, dDim)
+	}
+
+	// Case 2: 2D transpose [1, 0]
+	if rank == 2 && permutations[0] == 1 && permutations[1] == 0 {
+		return dispatchTranspose2D(backend, operand, output, dims[0], dims[1])
+	}
+
+	// Case 3: Any rank where trailing dimension is preserved (permutations[rank-1] == rank-1)
+	if rank > 2 && permutations[rank-1] == rank-1 {
+		return dispatchTransposeTrailing(backend, operand, output, dims, permutations)
+	}
+
+	return false
+}
+
+func dispatchIdentity(operand, output *gobackend.Buffer) bool {
+	switch in := operand.Flat.(type) {
+	case []float32:
+		copy(output.Flat.([]float32), in)
+	case []float64:
+		copy(output.Flat.([]float64), in)
+	case []bfloat16.BFloat16:
+		copy(output.Flat.([]bfloat16.BFloat16), in)
+	case []float16.Float16:
+		copy(output.Flat.([]float16.Float16), in)
+	case []int32:
+		copy(output.Flat.([]int32), in)
+	case []int64:
+		copy(output.Flat.([]int64), in)
+	case []int16:
+		copy(output.Flat.([]int16), in)
+	case []int8:
+		copy(output.Flat.([]int8), in)
+	case []uint32:
+		copy(output.Flat.([]uint32), in)
+	case []uint64:
+		copy(output.Flat.([]uint64), in)
+	case []uint16:
+		copy(output.Flat.([]uint16), in)
+	case []uint8:
+		copy(output.Flat.([]uint8), in)
+	case []bool:
+		copy(output.Flat.([]bool), in)
+	default:
+		return false
+	}
+	return true
+}
+
+func dispatchTranspose4D_0213(backend *gobackend.Backend, operand, output *gobackend.Buffer, bDim, sDim, hDim, dDim int) bool {
+	switch in := operand.Flat.(type) {
+	case []float32:
+		transpose4D_0213(backend, in, output.Flat.([]float32), bDim, sDim, hDim, dDim)
+	case []float64:
+		transpose4D_0213(backend, in, output.Flat.([]float64), bDim, sDim, hDim, dDim)
+	case []bfloat16.BFloat16:
+		transpose4D_0213(backend, in, output.Flat.([]bfloat16.BFloat16), bDim, sDim, hDim, dDim)
+	case []float16.Float16:
+		transpose4D_0213(backend, in, output.Flat.([]float16.Float16), bDim, sDim, hDim, dDim)
+	case []int32:
+		transpose4D_0213(backend, in, output.Flat.([]int32), bDim, sDim, hDim, dDim)
+	case []int64:
+		transpose4D_0213(backend, in, output.Flat.([]int64), bDim, sDim, hDim, dDim)
+	case []int16:
+		transpose4D_0213(backend, in, output.Flat.([]int16), bDim, sDim, hDim, dDim)
+	case []int8:
+		transpose4D_0213(backend, in, output.Flat.([]int8), bDim, sDim, hDim, dDim)
+	case []uint32:
+		transpose4D_0213(backend, in, output.Flat.([]uint32), bDim, sDim, hDim, dDim)
+	case []uint64:
+		transpose4D_0213(backend, in, output.Flat.([]uint64), bDim, sDim, hDim, dDim)
+	case []uint16:
+		transpose4D_0213(backend, in, output.Flat.([]uint16), bDim, sDim, hDim, dDim)
+	case []uint8:
+		transpose4D_0213(backend, in, output.Flat.([]uint8), bDim, sDim, hDim, dDim)
+	case []bool:
+		transpose4D_0213(backend, in, output.Flat.([]bool), bDim, sDim, hDim, dDim)
+	default:
+		return false
+	}
+	return true
+}
+
+func transpose4D_0213[T any](backend *gobackend.Backend, in, out []T, bDim, sDim, hDim, dDim int) {
+	totalElements := bDim * sDim * hDim * dDim
+	totalOuter := bDim * hDim
+	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && totalOuter > 1 && totalElements > 16384 {
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		targetChunks := min(totalOuter, max(1, numWorkers*2))
+		outerPerChunk := max(1, (totalOuter+targetChunks-1)/targetChunks)
+		var wg sync.WaitGroup
+		for start := 0; start < totalOuter; start += outerPerChunk {
+			end := min(start+outerPerChunk, totalOuter)
+			wg.Add(1)
+			backend.Workers.WaitToStart(func() {
+				for outer := start; outer < end; outer++ {
+					b := outer / hDim
+					h := outer % hDim
+					outBase := outer * (sDim * dDim)
+					inBase := b*(sDim*hDim*dDim) + h*dDim
+					strideIn := hDim * dDim
+					for s := 0; s < sDim; s++ {
+						inOff := inBase + s*strideIn
+						outOff := outBase + s*dDim
+						copy(out[outOff:outOff+dDim], in[inOff:inOff+dDim])
+					}
+				}
+				wg.Done()
+			})
+		}
+		wg.Wait()
+	} else if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && totalOuter == 1 && sDim > 512 && totalElements > 16384 {
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		targetChunks := min(sDim, max(1, numWorkers*2))
+		sPerChunk := max(1, (sDim+targetChunks-1)/targetChunks)
+		var wg sync.WaitGroup
+		for start := 0; start < sDim; start += sPerChunk {
+			end := min(start+sPerChunk, sDim)
+			wg.Add(1)
+			backend.Workers.WaitToStart(func() {
+				strideIn := hDim * dDim
+				for s := start; s < end; s++ {
+					inOff := s * strideIn
+					outOff := s * dDim
+					copy(out[outOff:outOff+dDim], in[inOff:inOff+dDim])
+				}
+				wg.Done()
+			})
+		}
+		wg.Wait()
+	} else {
+		for outer := 0; outer < totalOuter; outer++ {
+			b := outer / hDim
+			h := outer % hDim
+			outBase := outer * (sDim * dDim)
+			inBase := b*(sDim*hDim*dDim) + h*dDim
+			strideIn := hDim * dDim
+			for s := 0; s < sDim; s++ {
+				inOff := inBase + s*strideIn
+				outOff := outBase + s*dDim
+				copy(out[outOff:outOff+dDim], in[inOff:inOff+dDim])
+			}
+		}
+	}
+}
+
+func dispatchTranspose2D(backend *gobackend.Backend, operand, output *gobackend.Buffer, H, W int) bool {
+	switch in := operand.Flat.(type) {
+	case []float32:
+		transpose2D(backend, in, output.Flat.([]float32), H, W)
+	case []float64:
+		transpose2D(backend, in, output.Flat.([]float64), H, W)
+	case []bfloat16.BFloat16:
+		transpose2D(backend, in, output.Flat.([]bfloat16.BFloat16), H, W)
+	case []float16.Float16:
+		transpose2D(backend, in, output.Flat.([]float16.Float16), H, W)
+	case []int32:
+		transpose2D(backend, in, output.Flat.([]int32), H, W)
+	case []int64:
+		transpose2D(backend, in, output.Flat.([]int64), H, W)
+	case []int16:
+		transpose2D(backend, in, output.Flat.([]int16), H, W)
+	case []int8:
+		transpose2D(backend, in, output.Flat.([]int8), H, W)
+	case []uint32:
+		transpose2D(backend, in, output.Flat.([]uint32), H, W)
+	case []uint64:
+		transpose2D(backend, in, output.Flat.([]uint64), H, W)
+	case []uint16:
+		transpose2D(backend, in, output.Flat.([]uint16), H, W)
+	case []uint8:
+		transpose2D(backend, in, output.Flat.([]uint8), H, W)
+	case []bool:
+		transpose2D(backend, in, output.Flat.([]bool), H, W)
+	default:
+		return false
+	}
+	return true
+}
+
+func transpose2D[T any](backend *gobackend.Backend, in, out []T, H, W int) {
+	tileSize := 32
+	totalElements := H * W
+	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && H > 32 && totalElements > 16384 {
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		rTiles := (H + tileSize - 1) / tileSize
+		targetChunks := min(rTiles, max(1, numWorkers*2))
+		tilesPerChunk := max(1, (rTiles+targetChunks-1)/targetChunks)
+		var wg sync.WaitGroup
+		for t := 0; t < rTiles; t += tilesPerChunk {
+			tStart := t * tileSize
+			tEnd := min((t+tilesPerChunk)*tileSize, H)
+			wg.Add(1)
+			backend.Workers.WaitToStart(func() {
+				for r0 := tStart; r0 < tEnd; r0 += tileSize {
+					rLimit := min(r0+tileSize, tEnd)
+					for c0 := 0; c0 < W; c0 += tileSize {
+						cLimit := min(c0+tileSize, W)
+						for r := r0; r < rLimit; r++ {
+							for c := c0; c < cLimit; c++ {
+								out[c*H+r] = in[r*W+c]
+							}
+						}
+					}
+				}
+				wg.Done()
+			})
+		}
+		wg.Wait()
+	} else {
+		for r0 := 0; r0 < H; r0 += tileSize {
+			rLimit := min(r0+tileSize, H)
+			for c0 := 0; c0 < W; c0 += tileSize {
+				cLimit := min(c0+tileSize, W)
+				for r := r0; r < rLimit; r++ {
+					for c := c0; c < cLimit; c++ {
+						out[c*H+r] = in[r*W+c]
+					}
+				}
+			}
+		}
+	}
+}
+
+func dispatchTransposeTrailing(backend *gobackend.Backend, operand, output *gobackend.Buffer, dims, permutations []int) bool {
+	dtype := operand.RawShape.DType
+	switch in := operand.Flat.(type) {
+	case []float32:
+		transposeTrailing(backend, in, output.Flat.([]float32), dims, permutations, dtype)
+	case []float64:
+		transposeTrailing(backend, in, output.Flat.([]float64), dims, permutations, dtype)
+	case []bfloat16.BFloat16:
+		transposeTrailing(backend, in, output.Flat.([]bfloat16.BFloat16), dims, permutations, dtype)
+	case []float16.Float16:
+		transposeTrailing(backend, in, output.Flat.([]float16.Float16), dims, permutations, dtype)
+	case []int32:
+		transposeTrailing(backend, in, output.Flat.([]int32), dims, permutations, dtype)
+	case []int64:
+		transposeTrailing(backend, in, output.Flat.([]int64), dims, permutations, dtype)
+	case []int16:
+		transposeTrailing(backend, in, output.Flat.([]int16), dims, permutations, dtype)
+	case []int8:
+		transposeTrailing(backend, in, output.Flat.([]int8), dims, permutations, dtype)
+	case []uint32:
+		transposeTrailing(backend, in, output.Flat.([]uint32), dims, permutations, dtype)
+	case []uint64:
+		transposeTrailing(backend, in, output.Flat.([]uint64), dims, permutations, dtype)
+	case []uint16:
+		transposeTrailing(backend, in, output.Flat.([]uint16), dims, permutations, dtype)
+	case []uint8:
+		transposeTrailing(backend, in, output.Flat.([]uint8), dims, permutations, dtype)
+	case []bool:
+		transposeTrailing(backend, in, output.Flat.([]bool), dims, permutations, dtype)
+	default:
+		return false
+	}
+	return true
+}
+
+func transposeTrailing[T any](backend *gobackend.Backend, in, out []T, dims, permutations []int, dtype dtypes.DType) {
+	rank := len(dims)
+	dDim := dims[rank-1]
+	outerShape := shapes.Make(dtype, dims[:rank-1]...)
+	outerPerm := permutations[:rank-1]
+	it := NewTransposeIterator(outerShape, outerPerm)
+	outerCount := outerShape.Size()
+	for m := 0; m < outerCount; m++ {
+		outIdx := it.Next()
+		inOff := m * dDim
+		outOff := outIdx * dDim
+		copy(out[outOff:outOff+dDim], in[inOff:inOff+dDim])
+	}
+}
+
+
 // TransposeIterator creates a dynamic iterator that yields output flat indices
 // for the corresponding flat index on the input operand, assuming the operand flat index is moving
 // incrementally.
diff --git a/internal/gobackend/ops/transpose_test.go b/internal/gobackend/ops/transpose_test.go
index e83d04a..39c8e5c 100644
--- a/internal/gobackend/ops/transpose_test.go
+++ b/internal/gobackend/ops/transpose_test.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/gomlx/compute/dtypes"
 	"github.com/gomlx/compute/internal/gobackend/ops"
+	"github.com/gomlx/compute/shapeinference"
 	"github.com/gomlx/compute/shapes"
 	"github.com/gomlx/compute/support/testutil"
 )
@@ -32,3 +33,118 @@ func TestTransposeIterator(t *testing.T) {
 		t.Fatalf("transposeIterator mismatch:\n%s", diff)
 	}
 }
+
+func TestTransposeFastPaths(t *testing.T) {
+	testCases := []struct {
+		name         string
+		shape        shapes.Shape
+		permutations []int
+	}{
+		{
+			name:         "4D_0213_attention",
+			shape:        shapes.Make(dtypes.Float32, 1, 64, 4, 32),
+			permutations: []int{0, 2, 1, 3},
+		},
+		{
+			name:         "4D_0213_batched",
+			shape:        shapes.Make(dtypes.Float32, 2, 8, 4, 16),
+			permutations: []int{0, 2, 1, 3},
+		},
+		{
+			name:         "3D_102",
+			shape:        shapes.Make(dtypes.Float32, 16, 8, 32),
+			permutations: []int{1, 0, 2},
+		},
+		{
+			name:         "2D_matrix_transpose",
+			shape:        shapes.Make(dtypes.Float32, 64, 128),
+			permutations: []int{1, 0},
+		},
+		{
+			name:         "trailing_preserved_5D",
+			shape:        shapes.Make(dtypes.Float32, 2, 3, 4, 5, 8),
+			permutations: []int{0, 2, 1, 3, 4},
+		},
+		{
+			name:         "identity_4D",
+			shape:        shapes.Make(dtypes.Float32, 2, 3, 4, 5),
+			permutations: []int{0, 1, 2, 3},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			size := tc.shape.Size()
+			srcData := make([]float32, size)
+			for i := range srcData {
+				srcData[i] = float32(i + 1)
+			}
+			srcBuf := makeBuffer(t, tc.shape, srcData)
+
+			outShape, err := shapeinference.Transpose(tc.shape, tc.permutations)
+			if err != nil {
+				t.Fatalf("Transpose shape failed: %+v", err)
+			}
+			dstBuf, err := backend.GetBuffer(outShape)
+			if err != nil {
+				t.Fatalf("GetBuffer failed: %+v", err)
+			}
+
+			// Run ExecuteTranspose (fast path)
+			ops.ExecuteTranspose(backend, srcBuf, dstBuf, tc.permutations)
+
+			// Compute expected via TransposeIterator
+			expected := make([]float32, size)
+			it := ops.NewTransposeIterator(tc.shape, tc.permutations)
+			for _, val := range srcData {
+				expected[it.Next()] = val
+			}
+
+			result := dstBuf.Flat.([]float32)
+			for i := range result {
+				if result[i] != expected[i] {
+					t.Fatalf("%s at idx %d: got %f, want %f", tc.name, i, result[i], expected[i])
+				}
+			}
+		})
+	}
+}
+
+func BenchmarkTranspose4D_Attention_Iterator(b *testing.B) {
+	// Shape: [1, 32768, 12, 64] -> [1, 12, 32768, 64]
+	// Permutation: [0, 2, 1, 3]
+	operandShape := shapes.Make(dtypes.Float32, 1, 32768, 12, 64)
+	permutations := []int{0, 2, 1, 3}
+	outShape := shapes.Make(dtypes.Float32, 1, 12, 32768, 64)
+
+	operandFlat := make([]float32, operandShape.Size())
+	outputFlat := make([]float32, outShape.Size())
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		it := ops.NewTransposeIterator(operandShape, permutations)
+		for _, value := range operandFlat {
+			outputFlat[it.Next()] = value
+		}
+	}
+}
+
+func BenchmarkTranspose4D_Attention_FastPath(b *testing.B) {
+	operandShape := shapes.Make(dtypes.Float32, 1, 32768, 12, 64)
+	permutations := []int{0, 2, 1, 3}
+	outShape := shapes.Make(dtypes.Float32, 1, 12, 32768, 64)
+
+	operandFlat := make([]float32, operandShape.Size())
+	srcBuf := makeBuffer(b, operandShape, operandFlat)
+	dstBuf, err := backend.GetBuffer(outShape)
+	if err != nil {
+		b.Fatalf("GetBuffer failed: %+v", err)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ops.ExecuteTranspose(backend, srcBuf, dstBuf, permutations)
+	}
+}
+
+

From 791e442468e002d43f8939bafc2be741383e57bc Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 14:20:24 +0200
Subject: [PATCH 07/12] Chunked and parallelized Where.

---
 internal/gobackend/ops/where.go | 115 ++++++++++++++++++++++++++++++++
 1 file changed, 115 insertions(+)

diff --git a/internal/gobackend/ops/where.go b/internal/gobackend/ops/where.go
index a179918..781a6d4 100644
--- a/internal/gobackend/ops/where.go
+++ b/internal/gobackend/ops/where.go
@@ -1,7 +1,12 @@
 package ops
 
 import (
+	"sync"
+
 	"github.com/gomlx/compute"
+	"github.com/gomlx/compute/dtypes"
+	"github.com/gomlx/compute/dtypes/bfloat16"
+	"github.com/gomlx/compute/dtypes/float16"
 	"github.com/gomlx/compute/internal/gobackend"
 	"github.com/gomlx/compute/shapeinference"
 )
@@ -50,6 +55,11 @@ func execWhere(backend *gobackend.Backend, node *gobackend.Node, inputs []*gobac
 	if backend.NoOps {
 		return output, nil
 	}
+
+	if dispatchWhereParallel(backend, condition, onTrue, onFalse, output) {
+		return output, nil
+	}
+
 	tmpAny, tmpErr := whereDTypeMap.Get(outputShape.DType)
 	if tmpErr != nil {
 		panic(tmpErr)
@@ -59,6 +69,111 @@ func execWhere(backend *gobackend.Backend, node *gobackend.Node, inputs []*gobac
 	return output, nil
 }
 
+func dispatchWhereParallel(backend *gobackend.Backend, conditionBuf, onTrueBuf, onFalseBuf, outputBuf *gobackend.Buffer) bool {
+	if conditionBuf.RawShape.IsScalar() {
+		return false
+	}
+	n := conditionBuf.RawShape.Size()
+	if backend == nil || backend.Workers == nil || !backend.Workers.IsEnabled() || n <= 32768 {
+		return false
+	}
+	cond := conditionBuf.Flat.([]bool)
+	switch outputBuf.RawShape.DType {
+	case dtypes.Float32:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]float32))
+	case dtypes.Float64:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]float64))
+	case dtypes.BFloat16:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]bfloat16.BFloat16))
+	case dtypes.Float16:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]float16.Float16))
+	case dtypes.Int32:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]int32))
+	case dtypes.Int64:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]int64))
+	case dtypes.Int16:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]int16))
+	case dtypes.Int8:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]int8))
+	case dtypes.Uint32:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]uint32))
+	case dtypes.Uint64:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]uint64))
+	case dtypes.Uint16:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]uint16))
+	case dtypes.Uint8:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]uint8))
+	case dtypes.Bool:
+		parallelWhere(backend, cond, onTrueBuf, onFalseBuf, outputBuf.Flat.([]bool))
+	default:
+		return false
+	}
+	return true
+}
+
+func parallelWhere[T any](backend *gobackend.Backend, cond []bool, onTrueBuf, onFalseBuf *gobackend.Buffer, out []T) {
+	n := len(cond)
+	onTrueIsScalar := onTrueBuf.RawShape.IsScalar()
+	onFalseIsScalar := onFalseBuf.RawShape.IsScalar()
+	onTrueFlat := onTrueBuf.Flat.([]T)
+	onFalseFlat := onFalseBuf.Flat.([]T)
+	var onTrueScalar, onFalseScalar T
+	if onTrueIsScalar {
+		onTrueScalar = onTrueFlat[0]
+	}
+	if onFalseIsScalar {
+		onFalseScalar = onFalseFlat[0]
+	}
+
+	numWorkers := backend.Workers.AdjustedMaxParallelism()
+	targetChunks := min(n, max(1, numWorkers*2))
+	chunkSize := max(16384, (n+targetChunks-1)/targetChunks)
+	var wg sync.WaitGroup
+
+	for start := 0; start < n; start += chunkSize {
+		end := min(start+chunkSize, n)
+		wg.Add(1)
+		backend.Workers.WaitToStart(func() {
+			switch {
+			case !onTrueIsScalar && onFalseIsScalar:
+				for i := start; i < end; i++ {
+					if cond[i] {
+						out[i] = onTrueFlat[i]
+					} else {
+						out[i] = onFalseScalar
+					}
+				}
+			case onTrueIsScalar && !onFalseIsScalar:
+				for i := start; i < end; i++ {
+					if cond[i] {
+						out[i] = onTrueScalar
+					} else {
+						out[i] = onFalseFlat[i]
+					}
+				}
+			case !onTrueIsScalar && !onFalseIsScalar:
+				for i := start; i < end; i++ {
+					if cond[i] {
+						out[i] = onTrueFlat[i]
+					} else {
+						out[i] = onFalseFlat[i]
+					}
+				}
+			default:
+				for i := start; i < end; i++ {
+					if cond[i] {
+						out[i] = onTrueScalar
+					} else {
+						out[i] = onFalseScalar
+					}
+				}
+			}
+			wg.Done()
+		})
+	}
+	wg.Wait()
+}
+
 //gobackend:dtypemap execWhereGeneric ints,uints,floats,half,bool
 var whereDTypeMap = gobackend.NewDTypeMap("Where")
 

From 0e6de03b45ecce3e7087e6d42fa2e8540948575c Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 14:28:17 +0200
Subject: [PATCH 08/12] IsAVX2Allowed and IsAVX512Allowed converted to bool
 variables.

---
 .agents/AGENTS.md                             |   8 +-
 internal/gobackend/activations/avx2/avx2.go   |   4 +-
 .../gobackend/activations/avx512/avx512.go    |   4 +-
 internal/gobackend/dot/matmul/avx2/avx2.go    |  95 ++++++++-------
 .../gobackend/dot/matmul/avx2/avx2_test.go    |   2 +-
 .../dot/matmul/avx2/internal_test.go          |   5 +-
 .../gobackend/dot/matmul/avx2/packing_test.go |   2 +-
 .../gobackend/dot/matmul/avx512/avx512.go     |   3 +-
 .../dot/matmul/avx512/avx512_test.go          |   2 +-
 .../dot/matmul/avx512/internal_test.go        |   9 +-
 .../gobackend/dot/matmul/epilogue_amd64.go    |  43 +++++++
 .../gobackend/dot/matmul/epilogue_amd64.s     | 109 ++++++++++++++++++
 .../gobackend/dot/matmul/epilogue_generic.go  |   9 ++
 .../gobackend/dot/matmul/epilogue_test.go     |  42 +++++++
 internal/gobackend/fusedops/avx2/avx2.go      |   2 +-
 .../gobackend/fusedops/avx2/layernorm_test.go |   2 +-
 internal/gobackend/fusedops/avx512/avx512.go  |   2 +-
 .../fusedops/avx512/layernorm_test.go         |   2 +-
 internal/gobackend/ops/avx2/avx2.go           |   4 +-
 .../ops/avx2/binary_trailing_test.go          |  10 +-
 .../ops/avx2/reduce_leading_sum_test.go       |   2 +-
 .../ops/avx2/reduce_trailing_sum_test.go      |   2 +-
 internal/gobackend/ops/avx512/avx512.go       |   4 +-
 .../ops/avx512/binary_trailing_test.go        |  10 +-
 .../ops/avx512/reduce_leading_sum_test.go     |   2 +-
 .../ops/avx512/reduce_trailing_sum_test.go    |   2 +-
 .../gobackend/ops/reduce_thresholds_amd64.go  |   4 +-
 internal/gobackend/ops/where_test.go          |  83 +++++++++++++
 internal/gobackend/simd_other.go              |  16 ++-
 internal/gobackend/simd_x86.go                |  22 ++--
 30 files changed, 389 insertions(+), 117 deletions(-)
 create mode 100644 internal/gobackend/dot/matmul/epilogue_amd64.go
 create mode 100644 internal/gobackend/dot/matmul/epilogue_amd64.s
 create mode 100644 internal/gobackend/dot/matmul/epilogue_generic.go
 create mode 100644 internal/gobackend/dot/matmul/epilogue_test.go
 create mode 100644 internal/gobackend/ops/where_test.go

diff --git a/.agents/AGENTS.md b/.agents/AGENTS.md
index 3a2b25e..13986fc 100644
--- a/.agents/AGENTS.md
+++ b/.agents/AGENTS.md
@@ -259,16 +259,16 @@ SIMD and assembly implementations for specific architectures (such as AVX2 or AV
 - **Always gate tests with runtime checks**: Every test targeting architecture-specific SIMD or assembly (e.g., in `avx2/` and `avx512/` subpackages) **must** verify that the instructions are supported and allowed on the current host, skipping gracefully if not:
   ```go
   // In AVX2 tests:
-  if !gobackend.IsAVX2Allowed() {
+  if !gobackend.IsAVX2Allowed {
       t.Skip("AVX2 is not supported or allowed on this host")
   }
 
   // In AVX-512 tests:
-  if !gobackend.IsAVX512Allowed() {
+  if !gobackend.IsAVX512Allowed {
       t.Skip("AVX-512 is not supported or allowed on this host")
   }
   ```
-- **Use `gobackend.IsAVX2Allowed()` / `gobackend.IsAVX512Allowed()`**: In `internal/gobackend/...`, prefer these helpers over raw `archsimd.X86.AVX2()` / `archsimd.X86.AVX512()`, as they check both hardware CPUID support and environment variables (`GOMLX_GO_SIMD_AVX2`, `GOMLX_GO_SIMD_AVX512`). Outside `internal/gobackend` (e.g. in `dtypes/float16`), check `archsimd.X86.AVX2()` / `archsimd.X86.AVX512()`.
+- **Use `gobackend.IsAVX2Allowed` / `gobackend.IsAVX512Allowed`**: In `internal/gobackend/...`, prefer these helpers over raw `archsimd.X86.AVX2()` / `archsimd.X86.AVX512()`, as they check both hardware CPUID support and environment variables (`GOMLX_GO_SIMD_AVX2`, `GOMLX_GO_SIMD_AVX512`). Outside `internal/gobackend` (e.g. in `dtypes/float16`), check `archsimd.X86.AVX2()` / `archsimd.X86.AVX512()`.
 - **Never invoke instructions directly in tests without gating**: Without this check, executing unsupported instructions triggers a `SIGILL` (illegal instruction) crash on machines that lack those CPU extensions (e.g., running AVX-512 code on a CPU without AVX-512, or in virtualized/containerized environments).
 
 ## Guidelines for Writing SIMD Assembly (AMD64 / AVX2 & AVX-512)
@@ -382,6 +382,6 @@ SIMD operations carry fixed overheads (register setup, horizontal reductions, ma
 ### 7. Gating Assembly Tests by Runtime Support
 
 Handwritten assembly functions (in `*_amd64.s`) bypass Go compiler checks and will immediately trigger a `SIGILL` crash if invoked on a CPU lacking the required instruction set:
-- Always gate unit tests in `avx2/` and `avx512/` subpackages by calling `if !gobackend.IsAVX2Allowed() { t.Skip(...) }` or `if !gobackend.IsAVX512Allowed() { t.Skip(...) }` at the top of each test function.
+- Always gate unit tests in `avx2/` and `avx512/` subpackages by calling `if !gobackend.IsAVX2Allowed { t.Skip(...) }` or `if !gobackend.IsAVX512Allowed { t.Skip(...) }` at the top of each test function.
 - Never write tests that execute raw assembly kernels without this gate.
 
diff --git a/internal/gobackend/activations/avx2/avx2.go b/internal/gobackend/activations/avx2/avx2.go
index 5d98a46..ca5675f 100644
--- a/internal/gobackend/activations/avx2/avx2.go
+++ b/internal/gobackend/activations/avx2/avx2.go
@@ -18,13 +18,13 @@ import (
 const PriorityAVX2 = gobackend.PriorityArch
 
 func init() {
-	if gobackend.IsAVX2Allowed() {
+	if gobackend.IsAVX2Allowed {
 		registerAVX2()
 	}
 }
 
 func registerAVX2() {
-// Float32
+	// Float32
 	activations.Register[float32]("avx2:relu", compute.ActivationRelu, ReluAVX2, PriorityAVX2)
 	activations.Register[float32]("avx2:hardswish", compute.ActivationHardSwish, HardSwishAVX2, PriorityAVX2)
 	activations.Register[float32]("avx2:silu", compute.ActivationSilu, SiluAVX2, PriorityAVX2)
diff --git a/internal/gobackend/activations/avx512/avx512.go b/internal/gobackend/activations/avx512/avx512.go
index 251e605..321ce81 100644
--- a/internal/gobackend/activations/avx512/avx512.go
+++ b/internal/gobackend/activations/avx512/avx512.go
@@ -18,7 +18,7 @@ import (
 const PriorityAVX512 = gobackend.PriorityArch + 1
 
 func init() {
-	if gobackend.IsAVX512Allowed() {
+	if gobackend.IsAVX512Allowed {
 		registerAVX512()
 	}
 }
@@ -406,7 +406,7 @@ func LeakyReluAVX512(data []float32) {
 }
 
 const (
-	seluScaleAVX512     = 1.0507009873554804934193349852946
+	seluScaleAVX512      = 1.0507009873554804934193349852946
 	seluScaleAlphaAVX512 = 1.0507009873554804934193349852946 * 1.6732632423543772848170429916717
 )
 
diff --git a/internal/gobackend/dot/matmul/avx2/avx2.go b/internal/gobackend/dot/matmul/avx2/avx2.go
index d34c049..892162f 100644
--- a/internal/gobackend/dot/matmul/avx2/avx2.go
+++ b/internal/gobackend/dot/matmul/avx2/avx2.go
@@ -80,7 +80,7 @@ func init() {
 		return
 	}
 
-	if gobackend.IsAVX2Allowed() {
+	if gobackend.IsAVX2Allowed {
 		registerAVX2(false)
 	}
 }
@@ -147,59 +147,59 @@ func avx2PackRHSNonTransposed[T gotype.ScalarNotComplex](
 	} else {
 		switch kernelColsBytes {
 		case 128:
-		for ; stripColIdx+kernelColsBytes <= copyColsBytesAll; stripColIdx += kernelColsBytes {
-			rhsPtr := rhsBasePtr + uintptr(rhsRowStart)*rhsStrideBytes + rhsColStartBytes + stripColIdx
-			for range contractingRows {
-				v0 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr)))
-				v1 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr + 32)))
-				v2 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr + 64)))
-				v3 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr + 96)))
-				v0.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr)))
-				v1.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr + 32)))
-				v2.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr + 64)))
-				v3.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr + 96)))
-				panelPtr += kernelColsBytes
-				rhsPtr += rhsStrideBytes
-			}
-		}
-	case 64:
-		for ; stripColIdx+kernelColsBytes <= copyColsBytesAll; stripColIdx += kernelColsBytes {
-			rhsPtr := rhsBasePtr + uintptr(rhsRowStart)*rhsStrideBytes + rhsColStartBytes + stripColIdx
-			for range contractingRows {
-				v0 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr)))
-				v1 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr + 32)))
-				v0.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr)))
-				v1.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr + 32)))
-				panelPtr += kernelColsBytes
-				rhsPtr += rhsStrideBytes
+			for ; stripColIdx+kernelColsBytes <= copyColsBytesAll; stripColIdx += kernelColsBytes {
+				rhsPtr := rhsBasePtr + uintptr(rhsRowStart)*rhsStrideBytes + rhsColStartBytes + stripColIdx
+				for range contractingRows {
+					v0 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr)))
+					v1 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr + 32)))
+					v2 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr + 64)))
+					v3 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr + 96)))
+					v0.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr)))
+					v1.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr + 32)))
+					v2.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr + 64)))
+					v3.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr + 96)))
+					panelPtr += kernelColsBytes
+					rhsPtr += rhsStrideBytes
+				}
 			}
-		}
-	case 32:
-		for ; stripColIdx+kernelColsBytes <= copyColsBytesAll; stripColIdx += kernelColsBytes {
-			rhsPtr := rhsBasePtr + uintptr(rhsRowStart)*rhsStrideBytes + rhsColStartBytes + stripColIdx
-			for range contractingRows {
-				v0 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr)))
-				v0.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr)))
-				panelPtr += kernelColsBytes
-				rhsPtr += rhsStrideBytes
+		case 64:
+			for ; stripColIdx+kernelColsBytes <= copyColsBytesAll; stripColIdx += kernelColsBytes {
+				rhsPtr := rhsBasePtr + uintptr(rhsRowStart)*rhsStrideBytes + rhsColStartBytes + stripColIdx
+				for range contractingRows {
+					v0 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr)))
+					v1 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr + 32)))
+					v0.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr)))
+					v1.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr + 32)))
+					panelPtr += kernelColsBytes
+					rhsPtr += rhsStrideBytes
+				}
 			}
-		}
-	default:
-		for ; stripColIdx+kernelColsBytes <= copyColsBytesAll; stripColIdx += kernelColsBytes {
-			rhsPtr := rhsBasePtr + uintptr(rhsRowStart)*rhsStrideBytes + rhsColStartBytes + stripColIdx
-			for range contractingRows {
-				rowRhsPtr := rhsPtr
-				for kernelColIdx := uintptr(0); kernelColIdx < kernelColsBytes; kernelColIdx += 32 {
-					v0 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rowRhsPtr)))
+		case 32:
+			for ; stripColIdx+kernelColsBytes <= copyColsBytesAll; stripColIdx += kernelColsBytes {
+				rhsPtr := rhsBasePtr + uintptr(rhsRowStart)*rhsStrideBytes + rhsColStartBytes + stripColIdx
+				for range contractingRows {
+					v0 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rhsPtr)))
 					v0.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr)))
-					panelPtr += 32
-					rowRhsPtr += 32
+					panelPtr += kernelColsBytes
+					rhsPtr += rhsStrideBytes
+				}
+			}
+		default:
+			for ; stripColIdx+kernelColsBytes <= copyColsBytesAll; stripColIdx += kernelColsBytes {
+				rhsPtr := rhsBasePtr + uintptr(rhsRowStart)*rhsStrideBytes + rhsColStartBytes + stripColIdx
+				for range contractingRows {
+					rowRhsPtr := rhsPtr
+					for kernelColIdx := uintptr(0); kernelColIdx < kernelColsBytes; kernelColIdx += 32 {
+						v0 := archsimd.LoadUint8x32Array((*[32]uint8)(unsafe.Pointer(rowRhsPtr)))
+						v0.StoreArray((*[32]uint8)(unsafe.Pointer(panelPtr)))
+						panelPtr += 32
+						rowRhsPtr += 32
+					}
+					rhsPtr += rhsStrideBytes
 				}
-				rhsPtr += rhsStrideBytes
 			}
 		}
 	}
-	}
 
 	copyColsBytes := copyColsBytesAll - stripColIdx
 	if copyColsBytes == 0 {
@@ -525,4 +525,3 @@ func unsafePackLHS[T gotype.ScalarNotComplex](
 	lhsRowStart, lhsColStart, lhsCols, copyRows, contractingCols, kernelRows int) {
 	matmul.UnsafePackLHS(lhs, panel, lhsRowStart, lhsColStart, lhsCols, copyRows, contractingCols, kernelRows)
 }
-
diff --git a/internal/gobackend/dot/matmul/avx2/avx2_test.go b/internal/gobackend/dot/matmul/avx2/avx2_test.go
index 4614160..e78a6a2 100644
--- a/internal/gobackend/dot/matmul/avx2/avx2_test.go
+++ b/internal/gobackend/dot/matmul/avx2/avx2_test.go
@@ -17,7 +17,7 @@ import (
 )
 
 func TestAVX2(t *testing.T) {
-	if !gobackend.IsAVX2Allowed() {
+	if !gobackend.IsAVX2Allowed {
 		t.Skip("AVX2 is not supported on this architecture")
 	}
 
diff --git a/internal/gobackend/dot/matmul/avx2/internal_test.go b/internal/gobackend/dot/matmul/avx2/internal_test.go
index 7dcfcd3..fda7980 100644
--- a/internal/gobackend/dot/matmul/avx2/internal_test.go
+++ b/internal/gobackend/dot/matmul/avx2/internal_test.go
@@ -4,7 +4,6 @@
 
 package avx2
 
-
 import (
 	"simd/archsimd"
 	"testing"
@@ -18,7 +17,7 @@ import (
 )
 
 func TestAVX2(t *testing.T) {
-	if !gobackend.IsAVX2Allowed() {
+	if !gobackend.IsAVX2Allowed {
 		t.Skip("AVX2 is not supported on this architecture")
 	}
 
@@ -436,5 +435,3 @@ func BenchmarkAVX2SmallMatMul(b *testing.B) {
 		}
 	}
 }
-
-
diff --git a/internal/gobackend/dot/matmul/avx2/packing_test.go b/internal/gobackend/dot/matmul/avx2/packing_test.go
index bee0c8d..6cfabfc 100644
--- a/internal/gobackend/dot/matmul/avx2/packing_test.go
+++ b/internal/gobackend/dot/matmul/avx2/packing_test.go
@@ -14,7 +14,7 @@ import (
 )
 
 func TestAVX2Packing(t *testing.T) {
-	if !gobackend.IsAVX2Allowed() {
+	if !gobackend.IsAVX2Allowed {
 		t.Skip("AVX2 is not supported on this architecture")
 	}
 
diff --git a/internal/gobackend/dot/matmul/avx512/avx512.go b/internal/gobackend/dot/matmul/avx512/avx512.go
index 6542edf..e7d8305 100644
--- a/internal/gobackend/dot/matmul/avx512/avx512.go
+++ b/internal/gobackend/dot/matmul/avx512/avx512.go
@@ -97,7 +97,7 @@ func init() {
 		return
 	}
 
-	if gobackend.IsAVX512Allowed() {
+	if gobackend.IsAVX512Allowed {
 		registerAVX512(false)
 	}
 }
@@ -720,4 +720,3 @@ func unsafePackLHS[T gotype.ScalarNotComplex](
 	lhsRowStart, lhsColStart, lhsCols, copyRows, contractingCols, kernelRows int) {
 	matmul.UnsafePackLHS(lhs, panel, lhsRowStart, lhsColStart, lhsCols, copyRows, contractingCols, kernelRows)
 }
-
diff --git a/internal/gobackend/dot/matmul/avx512/avx512_test.go b/internal/gobackend/dot/matmul/avx512/avx512_test.go
index 763c3a5..ab542b7 100644
--- a/internal/gobackend/dot/matmul/avx512/avx512_test.go
+++ b/internal/gobackend/dot/matmul/avx512/avx512_test.go
@@ -17,7 +17,7 @@ import (
 )
 
 func TestAVX512(t *testing.T) {
-	if !gobackend.IsAVX512Allowed() {
+	if !gobackend.IsAVX512Allowed {
 		t.Skip("AVX512 is not supported on this architecture")
 	}
 
diff --git a/internal/gobackend/dot/matmul/avx512/internal_test.go b/internal/gobackend/dot/matmul/avx512/internal_test.go
index 5dad9b1..4899feb 100644
--- a/internal/gobackend/dot/matmul/avx512/internal_test.go
+++ b/internal/gobackend/dot/matmul/avx512/internal_test.go
@@ -22,7 +22,7 @@ import (
 )
 
 func TestAVX512(t *testing.T) {
-	if !gobackend.IsAVX512Allowed() {
+	if !gobackend.IsAVX512Allowed {
 		t.Skip("AVX512 is not supported on this architecture")
 	}
 
@@ -817,10 +817,10 @@ func BenchmarkAVX512(b *testing.B) {
 	}
 
 	rhsSizes := []struct {
-		name                 string
-		contractingRows, rhsCols int
+		name                        string
+		contractingRows, rhsCols    int
 		panelContracting, panelCols int
-		kernelCols           int
+		kernelCols                  int
 	}{
 		{"Large-1_1920x1024", 1920, 1024, 192, 384, 64},
 		{"Large-2_1920x1536", 1920, 1536, 192, 384, 64},
@@ -1014,4 +1014,3 @@ func BenchmarkSmallMatMul(b *testing.B) {
 		}
 	}
 }
-
diff --git a/internal/gobackend/dot/matmul/epilogue_amd64.go b/internal/gobackend/dot/matmul/epilogue_amd64.go
new file mode 100644
index 0000000..dc3da60
--- /dev/null
+++ b/internal/gobackend/dot/matmul/epilogue_amd64.go
@@ -0,0 +1,43 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+//go:build amd64
+
+package matmul
+
+import (
+	"unsafe"
+
+	"github.com/gomlx/compute/internal/gobackend"
+)
+
+//go:noescape
+func addBiasFloat32AVX512Asm(row, bias unsafe.Pointer, n int)
+
+//go:noescape
+func addBiasFloat32AVX2Asm(row, bias unsafe.Pointer, n int)
+
+var (
+	hasAVX512 bool
+	hasAVX2   bool
+)
+
+func init() {
+	hasAVX512 = gobackend.IsAVX512Allowed
+	hasAVX2 = gobackend.IsAVX2Allowed
+}
+
+func addBiasFloat32Arch(row, bias []float32) bool {
+	n := min(len(row), len(bias))
+	if n == 0 {
+		return true
+	}
+	if hasAVX512 {
+		addBiasFloat32AVX512Asm(unsafe.Pointer(&row[0]), unsafe.Pointer(&bias[0]), n)
+		return true
+	}
+	if hasAVX2 {
+		addBiasFloat32AVX2Asm(unsafe.Pointer(&row[0]), unsafe.Pointer(&bias[0]), n)
+		return true
+	}
+	return false
+}
diff --git a/internal/gobackend/dot/matmul/epilogue_amd64.s b/internal/gobackend/dot/matmul/epilogue_amd64.s
new file mode 100644
index 0000000..2a1ba0d
--- /dev/null
+++ b/internal/gobackend/dot/matmul/epilogue_amd64.s
@@ -0,0 +1,109 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+//go:build amd64
+
+#include "textflag.h"
+
+// func addBiasFloat32AVX512Asm(row, bias unsafe.Pointer, n int)
+TEXT ·addBiasFloat32AVX512Asm(SB), NOSPLIT, $0-24
+	MOVQ row+0(FP), DI
+	MOVQ bias+8(FP), SI
+	MOVQ n+16(FP), DX
+
+loop64:
+	CMPQ DX, $64
+	JL loop16
+	VMOVUPS 0(SI), Z0
+	VMOVUPS 64(SI), Z1
+	VMOVUPS 128(SI), Z2
+	VMOVUPS 192(SI), Z3
+	VADDPS 0(DI), Z0, Z0
+	VADDPS 64(DI), Z1, Z1
+	VADDPS 128(DI), Z2, Z2
+	VADDPS 192(DI), Z3, Z3
+	VMOVUPS Z0, 0(DI)
+	VMOVUPS Z1, 64(DI)
+	VMOVUPS Z2, 128(DI)
+	VMOVUPS Z3, 192(DI)
+	ADDQ $256, DI
+	ADDQ $256, SI
+	SUBQ $64, DX
+	JMP loop64
+
+loop16:
+	CMPQ DX, $16
+	JL loop1
+	VMOVUPS 0(SI), Z0
+	VADDPS 0(DI), Z0, Z0
+	VMOVUPS Z0, 0(DI)
+	ADDQ $64, DI
+	ADDQ $64, SI
+	SUBQ $16, DX
+	JMP loop16
+
+loop1:
+	CMPQ DX, $0
+	JLE done
+	MOVSS 0(SI), X0
+	ADDSS 0(DI), X0
+	MOVSS X0, 0(DI)
+	ADDQ $4, DI
+	ADDQ $4, SI
+	DECQ DX
+	JMP loop1
+
+done:
+	VZEROUPPER
+	RET
+
+// func addBiasFloat32AVX2Asm(row, bias unsafe.Pointer, n int)
+TEXT ·addBiasFloat32AVX2Asm(SB), NOSPLIT, $0-24
+	MOVQ row+0(FP), DI
+	MOVQ bias+8(FP), SI
+	MOVQ n+16(FP), DX
+
+loop32:
+	CMPQ DX, $32
+	JL loop8
+	VMOVUPS 0(SI), Y0
+	VMOVUPS 32(SI), Y1
+	VMOVUPS 64(SI), Y2
+	VMOVUPS 96(SI), Y3
+	VADDPS 0(DI), Y0, Y0
+	VADDPS 32(DI), Y1, Y1
+	VADDPS 64(DI), Y2, Y2
+	VADDPS 96(DI), Y3, Y3
+	VMOVUPS Y0, 0(DI)
+	VMOVUPS Y1, 32(DI)
+	VMOVUPS Y2, 64(DI)
+	VMOVUPS Y3, 96(DI)
+	ADDQ $128, DI
+	ADDQ $128, SI
+	SUBQ $32, DX
+	JMP loop32
+
+loop8:
+	CMPQ DX, $8
+	JL loop1_avx2
+	VMOVUPS 0(SI), Y0
+	VADDPS 0(DI), Y0, Y0
+	VMOVUPS Y0, 0(DI)
+	ADDQ $32, DI
+	ADDQ $32, SI
+	SUBQ $8, DX
+	JMP loop8
+
+loop1_avx2:
+	CMPQ DX, $0
+	JLE done_avx2
+	MOVSS 0(SI), X0
+	ADDSS 0(DI), X0
+	MOVSS X0, 0(DI)
+	ADDQ $4, DI
+	ADDQ $4, SI
+	DECQ DX
+	JMP loop1_avx2
+
+done_avx2:
+	VZEROUPPER
+	RET
diff --git a/internal/gobackend/dot/matmul/epilogue_generic.go b/internal/gobackend/dot/matmul/epilogue_generic.go
new file mode 100644
index 0000000..cf0881c
--- /dev/null
+++ b/internal/gobackend/dot/matmul/epilogue_generic.go
@@ -0,0 +1,9 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+//go:build !amd64
+
+package matmul
+
+func addBiasFloat32Arch(row, bias []float32) bool {
+	return false
+}
diff --git a/internal/gobackend/dot/matmul/epilogue_test.go b/internal/gobackend/dot/matmul/epilogue_test.go
new file mode 100644
index 0000000..0ebe555
--- /dev/null
+++ b/internal/gobackend/dot/matmul/epilogue_test.go
@@ -0,0 +1,42 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+package matmul
+
+import (
+	"testing"
+)
+
+func TestAddBiasFloat32(t *testing.T) {
+	for _, size := range []int{0, 1, 3, 7, 8, 15, 16, 31, 32, 63, 64, 65, 127, 128, 1536} {
+		row := make([]float32, size)
+		bias := make([]float32, size)
+		expected := make([]float32, size)
+		for i := 0; i < size; i++ {
+			row[i] = float32(i) * 1.5
+			bias[i] = float32(i) * 2.5
+			expected[i] = row[i] + bias[i]
+		}
+		addBias(row, bias)
+		for i := 0; i < size; i++ {
+			if row[i] != expected[i] {
+				t.Fatalf("size=%d, idx=%d: got %f, want %f", size, i, row[i], expected[i])
+			}
+		}
+	}
+}
+
+func BenchmarkAddBiasFloat32(b *testing.B) {
+	rowSize := 1536
+	row := make([]float32, rowSize)
+	bias := make([]float32, rowSize)
+	for i := range bias {
+		bias[i] = float32(i)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		addBias(row, bias)
+	}
+}
+
+
diff --git a/internal/gobackend/fusedops/avx2/avx2.go b/internal/gobackend/fusedops/avx2/avx2.go
index 1f2bfff..c7777f0 100644
--- a/internal/gobackend/fusedops/avx2/avx2.go
+++ b/internal/gobackend/fusedops/avx2/avx2.go
@@ -12,7 +12,7 @@ import (
 )
 
 func init() {
-	if gobackend.IsAVX2Allowed() {
+	if gobackend.IsAVX2Allowed {
 		registerAVX2()
 	}
 }
diff --git a/internal/gobackend/fusedops/avx2/layernorm_test.go b/internal/gobackend/fusedops/avx2/layernorm_test.go
index 512722b..490c63b 100644
--- a/internal/gobackend/fusedops/avx2/layernorm_test.go
+++ b/internal/gobackend/fusedops/avx2/layernorm_test.go
@@ -15,7 +15,7 @@ import (
 )
 
 func TestAVX2LayerNormCorrectness(t *testing.T) {
-	if !gobackend.IsAVX2Allowed() {
+	if !gobackend.IsAVX2Allowed {
 		t.Skip("AVX-2 is not supported or allowed on this machine")
 	}
 	rand.Seed(42)
diff --git a/internal/gobackend/fusedops/avx512/avx512.go b/internal/gobackend/fusedops/avx512/avx512.go
index 957ad1e..4bc5eb5 100644
--- a/internal/gobackend/fusedops/avx512/avx512.go
+++ b/internal/gobackend/fusedops/avx512/avx512.go
@@ -12,7 +12,7 @@ import (
 )
 
 func init() {
-	if gobackend.IsAVX512Allowed() {
+	if gobackend.IsAVX512Allowed {
 		registerAVX512()
 	}
 }
diff --git a/internal/gobackend/fusedops/avx512/layernorm_test.go b/internal/gobackend/fusedops/avx512/layernorm_test.go
index ea24a17..71872ae 100644
--- a/internal/gobackend/fusedops/avx512/layernorm_test.go
+++ b/internal/gobackend/fusedops/avx512/layernorm_test.go
@@ -15,7 +15,7 @@ import (
 )
 
 func TestAVX512LayerNormCorrectness(t *testing.T) {
-	if !gobackend.IsAVX512Allowed() {
+	if !gobackend.IsAVX512Allowed {
 		t.Skip("AVX-512 is not supported or allowed on this machine")
 	}
 	rand.Seed(42)
diff --git a/internal/gobackend/ops/avx2/avx2.go b/internal/gobackend/ops/avx2/avx2.go
index c1a2fdc..1e0bb38 100644
--- a/internal/gobackend/ops/avx2/avx2.go
+++ b/internal/gobackend/ops/avx2/avx2.go
@@ -13,7 +13,7 @@ import (
 )
 
 func init() {
-	if gobackend.IsAVX2Allowed() {
+	if gobackend.IsAVX2Allowed {
 		registerAVX2()
 	}
 }
@@ -394,5 +394,3 @@ func DispatchBinaryTrailingAVX2(op compute.OpType, isLHS bool, lhs, rhs, out uns
 	}
 	return false
 }
-
-
diff --git a/internal/gobackend/ops/avx2/binary_trailing_test.go b/internal/gobackend/ops/avx2/binary_trailing_test.go
index db22c76..579a044 100644
--- a/internal/gobackend/ops/avx2/binary_trailing_test.go
+++ b/internal/gobackend/ops/avx2/binary_trailing_test.go
@@ -15,7 +15,7 @@ import (
 )
 
 func TestAVX2BinaryTrailingCorrectness(t *testing.T) {
-	if !gobackend.IsAVX2Allowed() {
+	if !gobackend.IsAVX2Allowed {
 		t.Skip("AVX2 not allowed or not supported on this host")
 	}
 
@@ -211,10 +211,10 @@ func TestAVX2BinaryTrailingCorrectness(t *testing.T) {
 					expected := make([]int32, A*B)
 
 					for i := range lhs {
-						lhs[i] = int32((i%100) - 50)
+						lhs[i] = int32((i % 100) - 50)
 					}
 					for i := range rhs {
-						rhs[i] = int32((i%10) - 5)
+						rhs[i] = int32((i % 10) - 5)
 					}
 
 					for a := 0; a < A; a++ {
@@ -354,10 +354,10 @@ func TestAVX2BinaryTrailingCorrectness(t *testing.T) {
 					expected := make([]int64, A*B)
 
 					for i := range lhs {
-						lhs[i] = int64((i%100) - 50)
+						lhs[i] = int64((i % 100) - 50)
 					}
 					for i := range rhs {
-						rhs[i] = int64((i%10) - 5)
+						rhs[i] = int64((i % 10) - 5)
 					}
 
 					for a := 0; a < A; a++ {
diff --git a/internal/gobackend/ops/avx2/reduce_leading_sum_test.go b/internal/gobackend/ops/avx2/reduce_leading_sum_test.go
index 9570b1c..a8dee8b 100644
--- a/internal/gobackend/ops/avx2/reduce_leading_sum_test.go
+++ b/internal/gobackend/ops/avx2/reduce_leading_sum_test.go
@@ -39,7 +39,7 @@ func assertEqual[T comparable](t *testing.T, expected, actual T, msg string) {
 }
 
 func TestAVX2LeadingSumCorrectness(t *testing.T) {
-	if !gobackend.IsAVX2Allowed() {
+	if !gobackend.IsAVX2Allowed {
 		t.Skip("AVX2 not allowed or not supported on this host")
 	}
 	bValues := []int{1, 2, 3, 4, 7, 8, 9, 15, 16, 23, 24, 31, 32, 33, 48, 63, 64, 65, 100, 127, 128, 129, 255, 256, 512, 1024}
diff --git a/internal/gobackend/ops/avx2/reduce_trailing_sum_test.go b/internal/gobackend/ops/avx2/reduce_trailing_sum_test.go
index eabd330..d565e4f 100644
--- a/internal/gobackend/ops/avx2/reduce_trailing_sum_test.go
+++ b/internal/gobackend/ops/avx2/reduce_trailing_sum_test.go
@@ -17,7 +17,7 @@ import (
 )
 
 func TestAVX2TrailingSumCorrectness(t *testing.T) {
-	if !gobackend.IsAVX2Allowed() {
+	if !gobackend.IsAVX2Allowed {
 		t.Skip("AVX2 not allowed or not supported on this host")
 	}
 	bValues := []int{1, 2, 3, 4, 7, 8, 9, 15, 16, 23, 24, 31, 32, 33, 48, 63, 64, 65, 100, 127, 128, 129, 255, 256, 512, 1024}
diff --git a/internal/gobackend/ops/avx512/avx512.go b/internal/gobackend/ops/avx512/avx512.go
index 876c490..e81973b 100644
--- a/internal/gobackend/ops/avx512/avx512.go
+++ b/internal/gobackend/ops/avx512/avx512.go
@@ -13,7 +13,7 @@ import (
 )
 
 func init() {
-	if gobackend.IsAVX512Allowed() {
+	if gobackend.IsAVX512Allowed {
 		registerAVX512()
 	}
 }
@@ -394,5 +394,3 @@ func DispatchBinaryTrailingAVX512(op compute.OpType, isLHS bool, lhs, rhs, out u
 	}
 	return false
 }
-
-
diff --git a/internal/gobackend/ops/avx512/binary_trailing_test.go b/internal/gobackend/ops/avx512/binary_trailing_test.go
index 2d30b3f..8e723fd 100644
--- a/internal/gobackend/ops/avx512/binary_trailing_test.go
+++ b/internal/gobackend/ops/avx512/binary_trailing_test.go
@@ -15,7 +15,7 @@ import (
 )
 
 func TestAVX512BinaryTrailingCorrectness(t *testing.T) {
-	if !gobackend.IsAVX512Allowed() {
+	if !gobackend.IsAVX512Allowed {
 		t.Skip("AVX-512 not allowed or not supported on this host")
 	}
 
@@ -210,10 +210,10 @@ func TestAVX512BinaryTrailingCorrectness(t *testing.T) {
 					expected := make([]int32, A*B)
 
 					for i := range lhs {
-						lhs[i] = int32((i%100) - 50)
+						lhs[i] = int32((i % 100) - 50)
 					}
 					for i := range rhs {
-						rhs[i] = int32((i%10) - 5)
+						rhs[i] = int32((i % 10) - 5)
 					}
 
 					for a := 0; a < A; a++ {
@@ -353,10 +353,10 @@ func TestAVX512BinaryTrailingCorrectness(t *testing.T) {
 					expected := make([]int64, A*B)
 
 					for i := range lhs {
-						lhs[i] = int64((i%100) - 50)
+						lhs[i] = int64((i % 100) - 50)
 					}
 					for i := range rhs {
-						rhs[i] = int64((i%10) - 5)
+						rhs[i] = int64((i % 10) - 5)
 					}
 
 					for a := 0; a < A; a++ {
diff --git a/internal/gobackend/ops/avx512/reduce_leading_sum_test.go b/internal/gobackend/ops/avx512/reduce_leading_sum_test.go
index 6ae2a91..f646508 100644
--- a/internal/gobackend/ops/avx512/reduce_leading_sum_test.go
+++ b/internal/gobackend/ops/avx512/reduce_leading_sum_test.go
@@ -39,7 +39,7 @@ func assertEqual[T comparable](t *testing.T, expected, actual T, msg string) {
 }
 
 func TestAVX512LeadingSumCorrectness(t *testing.T) {
-	if !gobackend.IsAVX512Allowed() {
+	if !gobackend.IsAVX512Allowed {
 		t.Skip("AVX-512 not allowed or not supported on this host")
 	}
 	bValues := []int{1, 2, 3, 4, 7, 8, 9, 15, 16, 23, 24, 31, 32, 33, 48, 63, 64, 65, 100, 127, 128, 129, 255, 256, 512, 1024}
diff --git a/internal/gobackend/ops/avx512/reduce_trailing_sum_test.go b/internal/gobackend/ops/avx512/reduce_trailing_sum_test.go
index cb1bc5d..9e66bcf 100644
--- a/internal/gobackend/ops/avx512/reduce_trailing_sum_test.go
+++ b/internal/gobackend/ops/avx512/reduce_trailing_sum_test.go
@@ -17,7 +17,7 @@ import (
 )
 
 func TestAVX512TrailingSumCorrectness(t *testing.T) {
-	if !gobackend.IsAVX512Allowed() {
+	if !gobackend.IsAVX512Allowed {
 		t.Skip("AVX-512 not allowed or not supported on this host")
 	}
 	bValues := []int{1, 2, 3, 4, 7, 8, 9, 15, 16, 23, 24, 31, 32, 33, 48, 63, 64, 65, 100, 127, 128, 129, 255, 256, 512, 1024}
diff --git a/internal/gobackend/ops/reduce_thresholds_amd64.go b/internal/gobackend/ops/reduce_thresholds_amd64.go
index 8560e5c..735019b 100644
--- a/internal/gobackend/ops/reduce_thresholds_amd64.go
+++ b/internal/gobackend/ops/reduce_thresholds_amd64.go
@@ -106,9 +106,9 @@ var avx512ReduceThresholds = ReduceThresholdsConfig{
 }
 
 func init() {
-	if gobackend.IsAVX512Allowed() {
+	if gobackend.IsAVX512Allowed {
 		reduceThresholds = avx512ReduceThresholds
-	} else if gobackend.IsAVX2Allowed() {
+	} else if gobackend.IsAVX2Allowed {
 		reduceThresholds = avx2ReduceThresholds
 	}
 }
diff --git a/internal/gobackend/ops/where_test.go b/internal/gobackend/ops/where_test.go
new file mode 100644
index 0000000..5ee121a
--- /dev/null
+++ b/internal/gobackend/ops/where_test.go
@@ -0,0 +1,83 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+package ops_test
+
+import (
+	"testing"
+
+	"github.com/gomlx/compute"
+	"github.com/gomlx/compute/dtypes"
+	"github.com/gomlx/compute/internal/gobackend"
+	"github.com/gomlx/compute/internal/gobackend/ops"
+	"github.com/gomlx/compute/shapes"
+)
+
+func TestWhere(t *testing.T) {
+	for _, size := range []int{10, 1000, 40000} {
+		shape := shapes.Make(dtypes.Float32, size)
+		condShape := shapes.Make(dtypes.Bool, size)
+		scalarShape := shapes.Make(dtypes.Float32)
+
+		condData := make([]bool, size)
+		onTrueData := make([]float32, size)
+		onFalseData := []float32{-10000.0} // scalar onFalse
+
+		for i := 0; i < size; i++ {
+			condData[i] = (i % 2 == 0)
+			onTrueData[i] = float32(i + 1)
+		}
+
+		builder := backend.Builder("test_where").(*gobackend.Builder)
+		main := builder.Main().(*gobackend.Function)
+
+		condNode, err := main.Parameter("cond", condShape, nil)
+		if err != nil {
+			t.Fatalf("Parameter cond failed: %+v", err)
+		}
+		onTrueNode, err := main.Parameter("onTrue", shape, nil)
+		if err != nil {
+			t.Fatalf("Parameter onTrue failed: %+v", err)
+		}
+		onFalseNode, err := main.Parameter("onFalse", scalarShape, nil)
+		if err != nil {
+			t.Fatalf("Parameter onFalse failed: %+v", err)
+		}
+
+		outNode, err := ops.Where(main, condNode, onTrueNode, onFalseNode)
+		if err != nil {
+			t.Fatalf("Where failed: %+v", err)
+		}
+		err = main.Return([]compute.Value{outNode}, nil)
+		if err != nil {
+			t.Fatalf("Return failed: %+v", err)
+		}
+
+		exec, err := builder.Compile()
+		if err != nil {
+			t.Fatalf("Compile failed: %+v", err)
+		}
+		defer exec.Finalize()
+
+		condBuf := makeBuffer(t, condShape, condData)
+		onTrueBuf := makeBuffer(t, shape, onTrueData)
+		onFalseBuf := makeBuffer(t, scalarShape, onFalseData)
+
+		outputs, err := exec.Execute([]compute.Buffer{condBuf, onTrueBuf, onFalseBuf}, nil, 0)
+		if err != nil {
+			t.Fatalf("Execute failed: %+v", err)
+		}
+
+		outFlat := outputs[0].(*gobackend.Buffer).Flat.([]float32)
+		for i := 0; i < size; i++ {
+			var expected float32
+			if condData[i] {
+				expected = onTrueData[i]
+			} else {
+				expected = onFalseData[0]
+			}
+			if outFlat[i] != expected {
+				t.Fatalf("size=%d, idx=%d: got %f, want %f", size, i, outFlat[i], expected)
+			}
+		}
+	}
+}
diff --git a/internal/gobackend/simd_other.go b/internal/gobackend/simd_other.go
index c7566d9..48dc7d6 100644
--- a/internal/gobackend/simd_other.go
+++ b/internal/gobackend/simd_other.go
@@ -1,15 +1,13 @@
 // Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
 
-//go:build (!amd64) || (!goexperiment.simd)
+//go:build !amd64 || !goexperiment.simd
 
 package gobackend
 
-// IsAVX512Allowed returns false on non-amd64 architectures or when SIMD is disabled.
-func IsAVX512Allowed() bool {
-	return false
-}
+var (
+	// IsAVX512Allowed returns false on non-amd64 architectures or when SIMD is disabled.
+	IsAVX512Allowed = false
 
-// IsAVX2Allowed returns false on non-amd64 architectures or when SIMD is disabled.
-func IsAVX2Allowed() bool {
-	return false
-}
+	// IsAVX2Allowed returns false on non-amd64 architectures or when SIMD is disabled.
+	IsAVX2Allowed = false
+)
diff --git a/internal/gobackend/simd_x86.go b/internal/gobackend/simd_x86.go
index 3fd85dd..69e886f 100644
--- a/internal/gobackend/simd_x86.go
+++ b/internal/gobackend/simd_x86.go
@@ -10,16 +10,14 @@ import (
 	"github.com/gomlx/compute/support/envutil"
 )
 
-// IsAVX512Allowed returns true if AVX-512 is supported by the CPU and not disabled by GOMLX_GO_SIMD_AVX512.
-// Note that this only controls AVX512-specific implementations (e.g. specialized matmul and activations);
-// generic portable SIMD operations will still use hardware vector features if available.
-func IsAVX512Allowed() bool {
-	return envutil.MustReadBool(envutil.GoBackendSIMD_AVX512, true) && archsimd.X86.AVX512()
-}
+var (
+	// IsAVX2Allowed returns true if AVX2 is supported by the CPU and not disabled by GOMLX_GO_SIMD_AVX2.
+	// Note that this only controls AVX2-specific implementations (e.g. specialized matmul and activations);
+	// generic portable SIMD operations will still use hardware vector features if available.
+	IsAVX2Allowed = envutil.MustReadBool(envutil.GoBackendSIMD_AVX2, true) && archsimd.X86.AVX2()
 
-// IsAVX2Allowed returns true if AVX2 is supported by the CPU and not disabled by GOMLX_GO_SIMD_AVX2.
-// Note that this only controls AVX2-specific implementations (e.g. specialized matmul and activations);
-// generic portable SIMD operations will still use hardware vector features if available.
-func IsAVX2Allowed() bool {
-	return envutil.MustReadBool(envutil.GoBackendSIMD_AVX2, true) && archsimd.X86.AVX2()
-}
+	// IsAVX512Allowed returns true if AVX-512 is supported by the CPU and not disabled by GOMLX_GO_SIMD_AVX512.
+	// Note that this only controls AVX512-specific implementations (e.g. specialized matmul and activations);
+	// generic portable SIMD operations will still use hardware vector features if available.
+	IsAVX512Allowed = envutil.MustReadBool(envutil.GoBackendSIMD_AVX512, true) && archsimd.X86.AVX512()
+)

From 9a5833bc9d19d3e6e67b62f645080b8a75a9cac1 Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 15:18:19 +0200
Subject: [PATCH 09/12] Updated `Where` operator to support implicit broadcast
 of any of the 3 operands.

---
 ops.go                                |   6 +-
 shapeinference/shapeinference.go      | 100 +++++++++++++++++++++-----
 shapeinference/shapeinference_test.go |  92 ++++++++++++++++++++++++
 support/backendtest/special.go        |  53 ++++++++++++++
 4 files changed, 234 insertions(+), 17 deletions(-)

diff --git a/ops.go b/ops.go
index cf5580e..428ca5b 100644
--- a/ops.go
+++ b/ops.go
@@ -830,7 +830,11 @@ type StandardOps interface {
 	//
 	// The condition must be boolean, and onTrue and onFalse must have the same dtype.
 	//
-	// If either condition, onTrue or onFalse is a scalar, it will be broadcasted to the shape of the other operands.
+	// Standard implicit broadcasting rules apply across all three operands:
+	// - Scalar operands are implicitly broadcast to the shape of the other operands.
+	// - Non-scalar operands must have the same rank, and for each axis, their dimensions must either match or be 1.
+	//   Dimension 1 is broadcast to the larger dimension.
+	// - The resulting shape has dimension max(dim_condition, dim_onTrue, dim_onFalse) on each axis, and the dtype of onTrue/onFalse.
 	Where(condition, onTrue, onFalse Value) (Value, error)
 
 	// OptimizationBarrier introduces an optimization barrier.
diff --git a/shapeinference/shapeinference.go b/shapeinference/shapeinference.go
index a954436..bcb23a4 100644
--- a/shapeinference/shapeinference.go
+++ b/shapeinference/shapeinference.go
@@ -418,7 +418,16 @@ func UnaryOp(opType compute.OpType, operand shapes.Shape) (output shapes.Shape,
 //     must be Bool.
 //
 // Note: If you need to select between values of different dtypes, use ConvertDType to convert them
-// to a common dtype before calling Where.
+// Where returns the shape resulting from the Where operation.
+//
+// The condition must be of boolean dtype.
+// The onTrue and onFalse must have the same dtype.
+//
+// Standard implicit broadcasting rules apply across all three operands:
+// - Scalar operands are implicitly broadcast to the shape of the other operands.
+// - Non-scalar operands must have the same rank, and for each axis, their dimensions must either match or be 1.
+//   Dimension 1 is broadcast to the larger dimension.
+// - The resulting shape has dimension max(dim_condition, dim_onTrue, dim_onFalse) on each axis, and the dtype of onTrue/onFalse.
 func Where(condition, onTrue, onFalse shapes.Shape) (output shapes.Shape, err error) {
 	if condition.DType != dtypes.Bool {
 		err = errors.Errorf("condition for Where() must be a boolean, got %s instead", condition)
@@ -429,28 +438,87 @@ func Where(condition, onTrue, onFalse shapes.Shape) (output shapes.Shape, err er
 			onTrue.DType, onFalse.DType)
 		return
 	}
-	if !onTrue.IsScalar() && !onFalse.IsScalar() && !onTrue.Equal(onFalse) {
-		err = errors.Errorf("onTrue (%s) and onFalse (%s) values for Where() must either be scalar or match each other's shape",
-			onTrue, onFalse)
-		return
+
+	// Trivial case: all operands are scalar.
+	if condition.IsScalar() && onTrue.IsScalar() && onFalse.IsScalar() {
+		return shapes.Make(onTrue.DType), nil
 	}
 
-	output = onTrue
-	if output.IsScalar() {
-		output = onFalse
-		if output.IsScalar() && !condition.IsScalar() {
-			output = condition.Clone()
-			output.DType = onTrue.DType
+	// Collect non-scalar operands.
+	var nonScalars []shapes.Shape
+	for _, s := range []shapes.Shape{condition, onTrue, onFalse} {
+		if !s.IsScalar() {
+			nonScalars = append(nonScalars, s)
 		}
 	}
 
-	if !condition.IsScalar() && slices.Compare(condition.Dimensions, output.Dimensions) != 0 {
-		err = errors.Errorf("condition for Where() must either be a scalar or match the output shape (not the DType), instead got shapes condition=%s, onTrue=%s and onFalse=%s",
-			condition, onTrue, onFalse)
-		return
+	// If only one operand is non-scalar, it determines the shape.
+	if len(nonScalars) == 1 {
+		output = nonScalars[0].Clone()
+		output.DType = onTrue.DType
+		return output, nil
 	}
 
-	return
+	// Check that all non-scalar operands have the same rank.
+	rank := nonScalars[0].Rank()
+	for _, s := range nonScalars[1:] {
+		if s.Rank() != rank {
+			err = errors.Errorf("if operands are not scalars, their rank must match for Where(); got shapes condition=%s, onTrue=%s and onFalse=%s",
+				condition, onTrue, onFalse)
+			return shapes.Invalid(), err
+		}
+	}
+
+	output = shapes.Make(onTrue.DType, make([]int, rank)...)
+	for axis := range rank {
+		targetDim := 1
+		hasDynamic := false
+		var dynamicName string
+		for _, s := range nonScalars {
+			d := s.Dimensions[axis]
+			if d == shapes.DynamicDim {
+				hasDynamic = true
+				if dynamicName == "" {
+					dynamicName = s.AxisName(axis)
+				} else if !shapes.AxisNameEqual(dynamicName, s.AxisName(axis)) {
+					err = errors.Errorf("axis #%d is dynamic for multiple operands, but have different axis names for Where(), "+
+						"they cannot be implicitly broadcast; got shapes condition=%s, onTrue=%s and onFalse=%s",
+						axis, condition, onTrue, onFalse)
+					return shapes.Invalid(), err
+				}
+			} else if d != 1 {
+				if targetDim != 1 && targetDim != d {
+					err = errors.Errorf("dimension of axis #%d doesn't match and cannot be broadcast for Where(), got shapes condition=%s, onTrue=%s and onFalse=%s",
+						axis, condition, onTrue, onFalse)
+					return shapes.Invalid(), err
+				}
+				targetDim = d
+			}
+		}
+
+		if hasDynamic {
+			if targetDim != 1 {
+				err = errors.Errorf("axis #%d is dynamic for one operand and non-1 (%d) for another for Where(), "+
+					"they cannot be implicitly broadcast; got shapes condition=%s, onTrue=%s and onFalse=%s",
+					axis, targetDim, condition, onTrue, onFalse)
+				return shapes.Invalid(), err
+			}
+			output.Dimensions[axis] = shapes.DynamicDim
+		} else {
+			output.Dimensions[axis] = targetDim
+		}
+	}
+
+	// Unify axis names across all non-scalar operands.
+	axisNames := nonScalars[0].AxisNames
+	for _, s := range nonScalars[1:] {
+		axisNames, err = shapes.UnifyAxisNames(shapes.Shape{Dimensions: output.Dimensions, AxisNames: axisNames}, s)
+		if err != nil {
+			return shapes.Invalid(), errors.Wrapf(err, "axis name conflict in Where()")
+		}
+	}
+	output.AxisNames = axisNames
+	return output, nil
 }
 
 // Reshape to the given dimensions: trivial output shape, but this function also checks
diff --git a/shapeinference/shapeinference_test.go b/shapeinference/shapeinference_test.go
index 3195e9e..a9b29ad 100644
--- a/shapeinference/shapeinference_test.go
+++ b/shapeinference/shapeinference_test.go
@@ -1584,3 +1584,95 @@ func TestScatterOp_Dynamic(t *testing.T) {
 		}
 	})
 }
+
+func TestWhere(t *testing.T) {
+	// All scalars.
+	out, err := Where(S(Bool), S(F32), S(F32))
+	if err != nil {
+		t.Fatalf("Where failed: %v", err)
+	}
+	if !out.Equal(S(F32)) {
+		t.Fatalf("Expected scalar F32, got %s", out)
+	}
+
+	// Two scalars, one tensor.
+	out, err = Where(S(Bool, 2, 3), S(F32), S(F32))
+	if err != nil {
+		t.Fatalf("Where failed: %v", err)
+	}
+	if !out.Equal(S(F32, 2, 3)) {
+		t.Fatalf("Expected F32[2, 3], got %s", out)
+	}
+
+	// Scalar onFalse.
+	out, err = Where(S(Bool, 2, 3), S(F32, 2, 3), S(F32))
+	if err != nil {
+		t.Fatalf("Where failed: %v", err)
+	}
+	if !out.Equal(S(F32, 2, 3)) {
+		t.Fatalf("Expected F32[2, 3], got %s", out)
+	}
+
+	// Scalar onTrue.
+	out, err = Where(S(Bool, 2, 3), S(F32), S(F32, 2, 3))
+	if err != nil {
+		t.Fatalf("Where failed: %v", err)
+	}
+	if !out.Equal(S(F32, 2, 3)) {
+		t.Fatalf("Expected F32[2, 3], got %s", out)
+	}
+
+	// 2D multidirectional broadcasting: [2, 1], [1, 3], [2, 3].
+	out, err = Where(S(Bool, 2, 1), S(F32, 1, 3), S(F32, 2, 3))
+	if err != nil {
+		t.Fatalf("Where failed: %v", err)
+	}
+	if !out.Equal(S(F32, 2, 3)) {
+		t.Fatalf("Expected F32[2, 3], got %s", out)
+	}
+
+	// 3D broadcasting with trailing dimensions.
+	out, err = Where(S(Bool, 2, 1, 4), S(F32, 1, 3, 4), S(F32, 2, 3, 1))
+	if err != nil {
+		t.Fatalf("Where failed: %v", err)
+	}
+	if !out.Equal(S(F32, 2, 3, 4)) {
+		t.Fatalf("Expected F32[2, 3, 4], got %s", out)
+	}
+
+	// Axis name unification.
+	sCond := S(Bool, 2, 1).WithAxisNames("batch", "")
+	sTrue := S(F32, 1, 3).WithAxisNames("", "seq")
+	sFalse := S(F32, 2, 3).WithAxisNames("batch", "seq")
+	out, err = Where(sCond, sTrue, sFalse)
+	if err != nil {
+		t.Fatalf("Where failed: %v", err)
+	}
+	if !out.Equal(S(F32, 2, 3).WithAxisNames("batch", "seq")) {
+		t.Fatalf("Expected F32[2, 3] with axis names [batch, seq], got %s", out)
+	}
+
+	// Error: condition not bool.
+	_, err = Where(S(I32, 2), S(F32, 2), S(F32, 2))
+	if err == nil {
+		t.Fatalf("Expected error for non-boolean condition, got nil")
+	}
+
+	// Error: onTrue vs onFalse dtype mismatch.
+	_, err = Where(S(Bool, 2), S(F32, 2), S(I32, 2))
+	if err == nil {
+		t.Fatalf("Expected error for dtype mismatch, got nil")
+	}
+
+	// Error: rank mismatch among non-scalars.
+	_, err = Where(S(Bool, 2), S(F32, 2, 3), S(F32, 2, 3))
+	if err == nil {
+		t.Fatalf("Expected error for rank mismatch, got nil")
+	}
+
+	// Error: dimension mismatch.
+	_, err = Where(S(Bool, 2, 4), S(F32, 2, 3), S(F32, 2, 3))
+	if err == nil {
+		t.Fatalf("Expected error for dimension mismatch, got nil")
+	}
+}
diff --git a/support/backendtest/special.go b/support/backendtest/special.go
index 59e9e93..7443ba3 100644
--- a/support/backendtest/special.go
+++ b/support/backendtest/special.go
@@ -79,6 +79,59 @@ func TestSpecialOps(t *testing.T, b compute.Backend) {
 		if ok, diff := testutil.IsEqual([]float32{101, 2, 3}, y3); !ok {
 			t.Errorf("Where (non-scalar cond and values) mismatch:\n%s", diff)
 		}
+
+		// Vector cond, vector onTrue, scalar onFalse.
+		y4, err := testutil.Exec1(b, []any{[]bool{true, false, true}, []float32{1, 2, 3}, float32(10)}, buildWhere)
+		if err != nil {
+			t.Fatalf("Where (vector cond, vector onTrue, scalar onFalse) failed: %v", err)
+		}
+		if ok, diff := testutil.IsEqual([]float32{1, 10, 3}, y4); !ok {
+			t.Errorf("Where (vector cond, vector onTrue, scalar onFalse) mismatch:\n%s", diff)
+		}
+
+		// Vector cond, scalar onTrue, vector onFalse.
+		y5, err := testutil.Exec1(b, []any{[]bool{true, false, true}, float32(10), []float32{1, 2, 3}}, buildWhere)
+		if err != nil {
+			t.Fatalf("Where (vector cond, scalar onTrue, vector onFalse) failed: %v", err)
+		}
+		if ok, diff := testutil.IsEqual([]float32{10, 2, 10}, y5); !ok {
+			t.Errorf("Where (vector cond, scalar onTrue, vector onFalse) mismatch:\n%s", diff)
+		}
+
+		// Vector cond, scalar onTrue, scalar onFalse.
+		y6, err := testutil.Exec1(b, []any{[]bool{true, false, true}, float32(10), float32(20)}, buildWhere)
+		if err != nil {
+			t.Fatalf("Where (vector cond, scalar onTrue, scalar onFalse) failed: %v", err)
+		}
+		if ok, diff := testutil.IsEqual([]float32{10, 20, 10}, y6); !ok {
+			t.Errorf("Where (vector cond, scalar onTrue, scalar onFalse) mismatch:\n%s", diff)
+		}
+
+		// 2D multidirectional broadcasting: cond [2, 1], onTrue [2, 3], onFalse [1, 3] -> output [2, 3].
+		y7, err := testutil.Exec1(b, []any{
+			[][]bool{{true}, {false}},
+			[][]float32{{1, 2, 3}, {4, 5, 6}},
+			[][]float32{{10, 20, 30}},
+		}, buildWhere)
+		if err != nil {
+			t.Fatalf("Where (2D multidirectional broadcast) failed: %v", err)
+		}
+		if ok, diff := testutil.IsEqual([][]float32{{1, 2, 3}, {10, 20, 30}}, y7); !ok {
+			t.Errorf("Where (2D multidirectional broadcast) mismatch:\n%s", diff)
+		}
+
+		// 3D broadcasting with trailing dimensions: cond [1, 2, 1], onTrue [1, 2, 2], onFalse scalar.
+		y8, err := testutil.Exec1(b, []any{
+			[][][]bool{{{true}, {false}}},
+			[][][]float32{{{1, 2}, {3, 4}}},
+			float32(99),
+		}, buildWhere)
+		if err != nil {
+			t.Fatalf("Where (3D broadcast with trailing dimension) failed: %v", err)
+		}
+		if ok, diff := testutil.IsEqual([][][]float32{{{1, 2}, {99, 99}}}, y8); !ok {
+			t.Errorf("Where (3D broadcast with trailing dimension) mismatch:\n%s", diff)
+		}
 	})
 
 	t.Run("Reshape", func(t *testing.T) {

From b98a9b8d0e573c10e9e41b8537db7bc607678e30 Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 15:31:14 +0200
Subject: [PATCH 10/12] Implemented support for implicit broadcast for `Where`.
 AVX2 and AVX512 implementations for `Where`.

---
 internal/gobackend/ops/where.go               | 282 ++++++++++++++++--
 internal/gobackend/ops/where_amd64.go         |  54 ++++
 internal/gobackend/ops/where_amd64.s          | 233 +++++++++++++++
 internal/gobackend/ops/where_other.go         |   9 +
 internal/gobackend/ops/where_test.go          | 160 ++++++----
 .../gobackend/passes/implicitbroadcast.go     |  14 +-
 6 files changed, 666 insertions(+), 86 deletions(-)
 create mode 100644 internal/gobackend/ops/where_amd64.go
 create mode 100644 internal/gobackend/ops/where_amd64.s
 create mode 100644 internal/gobackend/ops/where_other.go

diff --git a/internal/gobackend/ops/where.go b/internal/gobackend/ops/where.go
index 781a6d4..35494da 100644
--- a/internal/gobackend/ops/where.go
+++ b/internal/gobackend/ops/where.go
@@ -1,6 +1,7 @@
 package ops
 
 import (
+	"slices"
 	"sync"
 
 	"github.com/gomlx/compute"
@@ -9,6 +10,7 @@ import (
 	"github.com/gomlx/compute/dtypes/float16"
 	"github.com/gomlx/compute/internal/gobackend"
 	"github.com/gomlx/compute/shapeinference"
+	"github.com/gomlx/compute/shapes"
 )
 
 func init() {
@@ -73,8 +75,18 @@ func dispatchWhereParallel(backend *gobackend.Backend, conditionBuf, onTrueBuf,
 	if conditionBuf.RawShape.IsScalar() {
 		return false
 	}
-	n := conditionBuf.RawShape.Size()
-	if backend == nil || backend.Workers == nil || !backend.Workers.IsEnabled() || n <= 32768 {
+	outputShape := outputBuf.RawShape
+	n := outputShape.Size()
+	if backend == nil || backend.Workers == nil || !backend.Workers.IsEnabled() || n < 8192 {
+		return false
+	}
+	if !conditionBuf.RawShape.Equal(outputShape) {
+		return false
+	}
+	if !onTrueBuf.RawShape.Equal(outputShape) && !onTrueBuf.RawShape.IsScalar() {
+		return false
+	}
+	if !onFalseBuf.RawShape.Equal(outputShape) && !onFalseBuf.RawShape.IsScalar() {
 		return false
 	}
 	cond := conditionBuf.Flat.([]bool)
@@ -130,10 +142,32 @@ func parallelWhere[T any](backend *gobackend.Backend, cond []bool, onTrueBuf, on
 	chunkSize := max(16384, (n+targetChunks-1)/targetChunks)
 	var wg sync.WaitGroup
 
+	_, isFloat32 := any(onTrueScalar).(float32)
+
 	for start := 0; start < n; start += chunkSize {
 		end := min(start+chunkSize, n)
 		wg.Add(1)
 		backend.Workers.WaitToStart(func() {
+			if isFloat32 {
+				var onTrueChunk, onFalseChunk []float32
+				condChunk := cond[start:end]
+				outChunk := any(out[start:end]).([]float32)
+				if onTrueIsScalar {
+					onTrueChunk = any(onTrueFlat[:1]).([]float32)
+				} else {
+					onTrueChunk = any(onTrueFlat[start:end]).([]float32)
+				}
+				if onFalseIsScalar {
+					onFalseChunk = any(onFalseFlat[:1]).([]float32)
+				} else {
+					onFalseChunk = any(onFalseFlat[start:end]).([]float32)
+				}
+				if whereFloat32Arch(condChunk, onTrueChunk, onFalseChunk, outChunk, onTrueIsScalar, onFalseIsScalar) {
+					wg.Done()
+					return
+				}
+			}
+
 			switch {
 			case !onTrueIsScalar && onFalseIsScalar:
 				for i := start; i < end; i++ {
@@ -174,6 +208,54 @@ func parallelWhere[T any](backend *gobackend.Backend, cond []bool, onTrueBuf, on
 	wg.Wait()
 }
 
+// ternaryZipIterator iterates over the flat indices of three broadcast operands and the target buffer.
+type ternaryZipIterator struct {
+	tgtSize          int
+	tgtDims          []int
+	condStrides      []int
+	trueStrides      []int
+	falseStrides     []int
+	condIsBroadcast  []bool
+	trueIsBroadcast  []bool
+	falseIsBroadcast []bool
+	condIsScalar     bool
+	trueIsScalar     bool
+	falseIsScalar    bool
+}
+
+func newTernaryZipIterator(condShape, trueShape, falseShape, tgtShape shapes.Shape) *ternaryZipIterator {
+	rank := tgtShape.Rank()
+	zi := &ternaryZipIterator{
+		tgtSize:          tgtShape.Size(),
+		tgtDims:          slices.Clone(tgtShape.Dimensions),
+		condIsScalar:     condShape.IsScalar(),
+		trueIsScalar:     trueShape.IsScalar(),
+		falseIsScalar:    falseShape.IsScalar(),
+		condIsBroadcast:  make([]bool, rank),
+		trueIsBroadcast:  make([]bool, rank),
+		falseIsBroadcast: make([]bool, rank),
+	}
+	if !zi.condIsScalar {
+		zi.condStrides = condShape.Strides()
+		for axis := range rank {
+			zi.condIsBroadcast[axis] = condShape.Dimensions[axis] != tgtShape.Dimensions[axis]
+		}
+	}
+	if !zi.trueIsScalar {
+		zi.trueStrides = trueShape.Strides()
+		for axis := range rank {
+			zi.trueIsBroadcast[axis] = trueShape.Dimensions[axis] != tgtShape.Dimensions[axis]
+		}
+	}
+	if !zi.falseIsScalar {
+		zi.falseStrides = falseShape.Strides()
+		for axis := range rank {
+			zi.falseIsBroadcast[axis] = falseShape.Dimensions[axis] != tgtShape.Dimensions[axis]
+		}
+	}
+	return zi
+}
+
 //gobackend:dtypemap execWhereGeneric ints,uints,floats,half,bool
 var whereDTypeMap = gobackend.NewDTypeMap("Where")
 
@@ -188,25 +270,167 @@ func execWhereGeneric[T gobackend.SupportedTypesConstraints](conditionBuf, onTru
 		return
 	}
 
-	conditionFlat := conditionBuf.Flat.([]bool)
-	onTrueFlat := onTrueBuf.Flat.([]T)
-	onFalseFlat := onFalseBuf.Flat.([]T)
-	outputFlat := outputBuf.Flat.([]T)
-	onTrueIsScalar := onTrueBuf.RawShape.IsScalar()
-	onFalseIsScalar := onFalseBuf.RawShape.IsScalar()
-	onTrue := onTrueFlat[0]
-	onFalse := onFalseFlat[0]
-	for outputIdx, condition := range conditionFlat {
-		if condition {
-			if !onTrueIsScalar {
-				onTrue = onTrueFlat[outputIdx]
+	condShape := conditionBuf.RawShape
+	trueShape := onTrueBuf.RawShape
+	falseShape := onFalseBuf.RawShape
+	tgtShape := outputBuf.RawShape
+
+	// Fast path: condition matches target shape and value operands match or are scalar.
+	if condShape.Equal(tgtShape) &&
+		(trueShape.Equal(tgtShape) || trueShape.IsScalar()) &&
+		(falseShape.Equal(tgtShape) || falseShape.IsScalar()) {
+		condFlat := conditionBuf.Flat.([]bool)
+		trueFlat := onTrueBuf.Flat.([]T)
+		falseFlat := onFalseBuf.Flat.([]T)
+		outputFlat := outputBuf.Flat.([]T)
+		trueIsScalar := trueShape.IsScalar()
+		falseIsScalar := falseShape.IsScalar()
+
+		if _, ok := any(trueFlat[0]).(float32); ok {
+			var trueChunk, falseChunk []float32
+			if trueIsScalar {
+				trueChunk = any(trueFlat[:1]).([]float32)
+			} else {
+				trueChunk = any(trueFlat).([]float32)
 			}
-			outputFlat[outputIdx] = onTrue
-		} else {
-			if !onFalseIsScalar {
-				onFalse = onFalseFlat[outputIdx]
+			if falseIsScalar {
+				falseChunk = any(falseFlat[:1]).([]float32)
+			} else {
+				falseChunk = any(falseFlat).([]float32)
+			}
+			if whereFloat32Arch(condFlat, trueChunk, falseChunk, any(outputFlat).([]float32), trueIsScalar, falseIsScalar) {
+				return
+			}
+		}
+
+		var tVal, fVal T
+		if trueIsScalar {
+			tVal = trueFlat[0]
+		}
+		if falseIsScalar {
+			fVal = falseFlat[0]
+		}
+		for i, c := range condFlat {
+			if c {
+				if trueIsScalar {
+					outputFlat[i] = tVal
+				} else {
+					outputFlat[i] = trueFlat[i]
+				}
+			} else {
+				if falseIsScalar {
+					outputFlat[i] = fVal
+				} else {
+					outputFlat[i] = falseFlat[i]
+				}
+			}
+		}
+		return
+	}
+
+	// General multi-dimensional broadcasting:
+	zi := newTernaryZipIterator(condShape, trueShape, falseShape, tgtShape)
+	condFlat := conditionBuf.Flat.([]bool)
+	trueFlat := onTrueBuf.Flat.([]T)
+	falseFlat := onFalseBuf.Flat.([]T)
+	outFlat := outputBuf.Flat.([]T)
+
+	rank := len(zi.tgtDims)
+	trailingLen := 1
+	splitAxis := -1
+	for axis := rank - 1; axis >= 0; axis-- {
+		condBroadcast := !zi.condIsScalar && zi.condIsBroadcast[axis]
+		trueBroadcast := !zi.trueIsScalar && zi.trueIsBroadcast[axis]
+		falseBroadcast := !zi.falseIsScalar && zi.falseIsBroadcast[axis]
+		if condBroadcast || trueBroadcast || falseBroadcast {
+			splitAxis = axis
+			break
+		}
+		trailingLen *= zi.tgtDims[axis]
+	}
+
+	trueIsScalar := zi.trueIsScalar
+	falseIsScalar := zi.falseIsScalar
+	_, isFloat32 := any(trueFlat[0]).(float32)
+
+	perAxesIdx := make([]int, rank)
+	for dstIdx := 0; dstIdx < zi.tgtSize; dstIdx += trailingLen {
+		condOffset, trueOffset, falseOffset := 0, 0, 0
+		if !zi.condIsScalar {
+			for a := 0; a <= splitAxis; a++ {
+				idx := perAxesIdx[a]
+				if zi.condIsBroadcast[a] {
+					idx = 0
+				}
+				condOffset += idx * zi.condStrides[a]
+			}
+		}
+		if !trueIsScalar {
+			for a := 0; a <= splitAxis; a++ {
+				idx := perAxesIdx[a]
+				if zi.trueIsBroadcast[a] {
+					idx = 0
+				}
+				trueOffset += idx * zi.trueStrides[a]
+			}
+		}
+		if !falseIsScalar {
+			for a := 0; a <= splitAxis; a++ {
+				idx := perAxesIdx[a]
+				if zi.falseIsBroadcast[a] {
+					idx = 0
+				}
+				falseOffset += idx * zi.falseStrides[a]
+			}
+		}
+
+		handled := false
+		if isFloat32 {
+			var tChunk, fChunk []float32
+			if trueIsScalar {
+				tChunk = any(trueFlat[:1]).([]float32)
+			} else {
+				tChunk = any(trueFlat[trueOffset : trueOffset+trailingLen]).([]float32)
+			}
+			if falseIsScalar {
+				fChunk = any(falseFlat[:1]).([]float32)
+			} else {
+				fChunk = any(falseFlat[falseOffset : falseOffset+trailingLen]).([]float32)
+			}
+			cChunk := condFlat[condOffset : condOffset+trailingLen]
+			oChunk := any(outFlat[dstIdx : dstIdx+trailingLen]).([]float32)
+			handled = whereFloat32Arch(cChunk, tChunk, fChunk, oChunk, trueIsScalar, falseIsScalar)
+		}
+
+		if !handled {
+			for j := 0; j < trailingLen; j++ {
+				var tVal, fVal T
+				if trueIsScalar {
+					tVal = trueFlat[0]
+				} else {
+					tVal = trueFlat[trueOffset+j]
+				}
+				if falseIsScalar {
+					fVal = falseFlat[0]
+				} else {
+					fVal = falseFlat[falseOffset+j]
+				}
+				if condFlat[condOffset+j] {
+					outFlat[dstIdx+j] = tVal
+				} else {
+					outFlat[dstIdx+j] = fVal
+				}
+			}
+		}
+
+		if splitAxis >= 0 {
+			for axis := splitAxis; axis >= 0; axis-- {
+				perAxesIdx[axis]++
+				if perAxesIdx[axis] < zi.tgtDims[axis] {
+					break
+				}
+				perAxesIdx[axis] = 0
 			}
-			outputFlat[outputIdx] = onFalse
 		}
 	}
 }
@@ -216,15 +440,23 @@ func execWhereSetOutputWithValue[T gobackend.SupportedTypesConstraints](outputBu
 		// The output is reusing the value buffer, nothing to do.
 		return
 	}
+	outputSlice := outputBuf.Flat.([]T)
+	valSlice := valueBuf.Flat.([]T)
 	if valueBuf.RawShape.Equal(outputBuf.RawShape) {
 		// Copy over values.
-		copy(outputBuf.Flat.([]T), valueBuf.Flat.([]T))
+		copy(outputSlice, valSlice)
 		return
 	}
-	// Value must then be a scalar:
-	c := valueBuf.Flat.([]T)[0]
-	outputSlice := outputBuf.Flat.([]T)
-	for outputIdx := range outputSlice {
-		outputSlice[outputIdx] = c
+	if valueBuf.RawShape.IsScalar() {
+		c := valSlice[0]
+		for outputIdx := range outputSlice {
+			outputSlice[outputIdx] = c
+		}
+		return
+	}
+	// General broadcast from valueBuf to outputBuf.
+	zi := gobackend.NewZippedBroadcastIterator(valueBuf.RawShape, valueBuf.RawShape, outputBuf.RawShape)
+	for idxs := range zi.IterFlatIndices() {
+		outputSlice[idxs.TgtFlatIdx] = valSlice[idxs.LHSFlatIdx]
 	}
 }
diff --git a/internal/gobackend/ops/where_amd64.go b/internal/gobackend/ops/where_amd64.go
new file mode 100644
index 0000000..9d3dbb7
--- /dev/null
+++ b/internal/gobackend/ops/where_amd64.go
@@ -0,0 +1,54 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+//go:build amd64
+
+package ops
+
+import (
+	"unsafe"
+
+	"github.com/gomlx/compute/internal/gobackend"
+)
+
+//go:noescape
+func whereVVFloat32AVX2Asm(cond, onTrue, onFalse, out unsafe.Pointer, count int)
+
+//go:noescape
+func whereVSFloat32AVX2Asm(cond, onTrue, onFalseScalar, out unsafe.Pointer, count int)
+
+//go:noescape
+func whereSVFloat32AVX2Asm(cond, onTrueScalar, onFalse, out unsafe.Pointer, count int)
+
+//go:noescape
+func whereSSFloat32AVX2Asm(cond, onTrueScalar, onFalseScalar, out unsafe.Pointer, count int)
+
+func whereFloat32Arch(cond []bool, onTrue, onFalse, out []float32, onTrueIsScalar, onFalseIsScalar bool) bool {
+	if len(out) == 0 {
+		return true
+	}
+	condPtr := unsafe.Pointer(&cond[0])
+	outPtr := unsafe.Pointer(&out[0])
+	var truePtr, falsePtr unsafe.Pointer
+	if len(onTrue) > 0 {
+		truePtr = unsafe.Pointer(&onTrue[0])
+	}
+	if len(onFalse) > 0 {
+		falsePtr = unsafe.Pointer(&onFalse[0])
+	}
+
+	if gobackend.IsAVX2Allowed || gobackend.IsAVX512Allowed {
+		switch {
+		case !onTrueIsScalar && !onFalseIsScalar:
+			whereVVFloat32AVX2Asm(condPtr, truePtr, falsePtr, outPtr, len(out))
+		case !onTrueIsScalar && onFalseIsScalar:
+			whereVSFloat32AVX2Asm(condPtr, truePtr, falsePtr, outPtr, len(out))
+		case onTrueIsScalar && !onFalseIsScalar:
+			whereSVFloat32AVX2Asm(condPtr, truePtr, falsePtr, outPtr, len(out))
+		case onTrueIsScalar && onFalseIsScalar:
+			whereSSFloat32AVX2Asm(condPtr, truePtr, falsePtr, outPtr, len(out))
+		}
+		return true
+	}
+
+	return false
+}
diff --git a/internal/gobackend/ops/where_amd64.s b/internal/gobackend/ops/where_amd64.s
new file mode 100644
index 0000000..1356c0d
--- /dev/null
+++ b/internal/gobackend/ops/where_amd64.s
@@ -0,0 +1,233 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+//go:build amd64
+
+#include "textflag.h"
+
+// ----------------------------------------------------------------------------
+// AVX2 FLOAT32
+// ----------------------------------------------------------------------------
+
+// func whereVVFloat32AVX2Asm(cond, onTrue, onFalse, out unsafe.Pointer, count int)
+TEXT ·whereVVFloat32AVX2Asm(SB), NOSPLIT, $0-40
+	MOVQ cond+0(FP), AX
+	MOVQ onTrue+8(FP), BX
+	MOVQ onFalse+16(FP), CX
+	MOVQ out+24(FP), DX
+	MOVQ count+32(FP), SI
+
+	CMPQ SI, $8
+	JL tail_vv2
+
+	VPXOR Y1, Y1, Y1
+
+loop8_vv2:
+	VMOVQ (AX), X0
+	VPMOVSXBD X0, Y0
+	VPCMPGTD Y1, Y0, Y0
+	VMOVUPS (BX), Y2
+	VMOVUPS (CX), Y3
+	VANDPS Y0, Y2, Y4
+	VANDNPS Y3, Y0, Y5
+	VORPS Y4, Y5, Y6
+	VMOVUPS Y6, (DX)
+
+	ADDQ $8, AX
+	ADDQ $32, BX
+	ADDQ $32, CX
+	ADDQ $32, DX
+	SUBQ $8, SI
+	CMPQ SI, $8
+	JGE loop8_vv2
+
+tail_vv2:
+	TESTQ SI, SI
+	JLE done_vv2
+
+loop_scalar_vv2:
+	MOVB (AX), R8
+	TESTB R8, R8
+	JZ false_scalar_vv2
+	MOVL (BX), R9
+	MOVL R9, (DX)
+	JMP next_scalar_vv2
+false_scalar_vv2:
+	MOVL (CX), R9
+	MOVL R9, (DX)
+next_scalar_vv2:
+	INCQ AX
+	ADDQ $4, BX
+	ADDQ $4, CX
+	ADDQ $4, DX
+	DECQ SI
+	JNZ loop_scalar_vv2
+
+done_vv2:
+	VZEROUPPER
+	RET
+
+// func whereVSFloat32AVX2Asm(cond, onTrue, onFalseScalar, out unsafe.Pointer, count int)
+TEXT ·whereVSFloat32AVX2Asm(SB), NOSPLIT, $0-40
+	MOVQ cond+0(FP), AX
+	MOVQ onTrue+8(FP), BX
+	MOVQ onFalseScalar+16(FP), CX
+	MOVQ out+24(FP), DX
+	MOVQ count+32(FP), SI
+
+	VBROADCASTSS (CX), Y3
+	VPXOR Y1, Y1, Y1
+
+	CMPQ SI, $8
+	JL tail_vs2
+
+loop8_vs2:
+	VMOVQ (AX), X0
+	VPMOVSXBD X0, Y0
+	VPCMPGTD Y1, Y0, Y0
+	VMOVUPS (BX), Y2
+	VANDPS Y0, Y2, Y4
+	VANDNPS Y3, Y0, Y5
+	VORPS Y4, Y5, Y6
+	VMOVUPS Y6, (DX)
+
+	ADDQ $8, AX
+	ADDQ $32, BX
+	ADDQ $32, DX
+	SUBQ $8, SI
+	CMPQ SI, $8
+	JGE loop8_vs2
+
+tail_vs2:
+	TESTQ SI, SI
+	JLE done_vs2
+	MOVL (CX), R10
+
+loop_scalar_vs2:
+	MOVB (AX), R8
+	TESTB R8, R8
+	JZ false_scalar_vs2
+	MOVL (BX), R9
+	MOVL R9, (DX)
+	JMP next_scalar_vs2
+false_scalar_vs2:
+	MOVL R10, (DX)
+next_scalar_vs2:
+	INCQ AX
+	ADDQ $4, BX
+	ADDQ $4, DX
+	DECQ SI
+	JNZ loop_scalar_vs2
+
+done_vs2:
+	VZEROUPPER
+	RET
+
+// func whereSVFloat32AVX2Asm(cond, onTrueScalar, onFalse, out unsafe.Pointer, count int)
+TEXT ·whereSVFloat32AVX2Asm(SB), NOSPLIT, $0-40
+	MOVQ cond+0(FP), AX
+	MOVQ onTrueScalar+8(FP), BX
+	MOVQ onFalse+16(FP), CX
+	MOVQ out+24(FP), DX
+	MOVQ count+32(FP), SI
+
+	VBROADCASTSS (BX), Y2
+	VPXOR Y1, Y1, Y1
+
+	CMPQ SI, $8
+	JL tail_sv2
+
+loop8_sv2:
+	VMOVQ (AX), X0
+	VPMOVSXBD X0, Y0
+	VPCMPGTD Y1, Y0, Y0
+	VMOVUPS (CX), Y3
+	VANDPS Y0, Y2, Y4
+	VANDNPS Y3, Y0, Y5
+	VORPS Y4, Y5, Y6
+	VMOVUPS Y6, (DX)
+
+	ADDQ $8, AX
+	ADDQ $32, CX
+	ADDQ $32, DX
+	SUBQ $8, SI
+	CMPQ SI, $8
+	JGE loop8_sv2
+
+tail_sv2:
+	TESTQ SI, SI
+	JLE done_sv2
+	MOVL (BX), R10
+
+loop_scalar_sv2:
+	MOVB (AX), R8
+	TESTB R8, R8
+	JZ false_scalar_sv2
+	MOVL R10, (DX)
+	JMP next_scalar_sv2
+false_scalar_sv2:
+	MOVL (CX), R9
+	MOVL R9, (DX)
+next_scalar_sv2:
+	INCQ AX
+	ADDQ $4, CX
+	ADDQ $4, DX
+	DECQ SI
+	JNZ loop_scalar_sv2
+
+done_sv2:
+	VZEROUPPER
+	RET
+
+// func whereSSFloat32AVX2Asm(cond, onTrueScalar, onFalseScalar, out unsafe.Pointer, count int)
+TEXT ·whereSSFloat32AVX2Asm(SB), NOSPLIT, $0-40
+	MOVQ cond+0(FP), AX
+	MOVQ onTrueScalar+8(FP), BX
+	MOVQ onFalseScalar+16(FP), CX
+	MOVQ out+24(FP), DX
+	MOVQ count+32(FP), SI
+
+	VBROADCASTSS (BX), Y2
+	VBROADCASTSS (CX), Y3
+	VPXOR Y1, Y1, Y1
+
+	CMPQ SI, $8
+	JL tail_ss2
+
+loop8_ss2:
+	VMOVQ (AX), X0
+	VPMOVSXBD X0, Y0
+	VPCMPGTD Y1, Y0, Y0
+	VANDPS Y0, Y2, Y4
+	VANDNPS Y3, Y0, Y5
+	VORPS Y4, Y5, Y6
+	VMOVUPS Y6, (DX)
+
+	ADDQ $8, AX
+	ADDQ $32, DX
+	SUBQ $8, SI
+	CMPQ SI, $8
+	JGE loop8_ss2
+
+tail_ss2:
+	TESTQ SI, SI
+	JLE done_ss2
+	MOVL (BX), R10
+	MOVL (CX), R11
+
+loop_scalar_ss2:
+	MOVB (AX), R8
+	TESTB R8, R8
+	JZ false_scalar_ss2
+	MOVL R10, (DX)
+	JMP next_scalar_ss2
+false_scalar_ss2:
+	MOVL R11, (DX)
+next_scalar_ss2:
+	INCQ AX
+	ADDQ $4, DX
+	DECQ SI
+	JNZ loop_scalar_ss2
+
+done_ss2:
+	VZEROUPPER
+	RET
diff --git a/internal/gobackend/ops/where_other.go b/internal/gobackend/ops/where_other.go
new file mode 100644
index 0000000..28ea309
--- /dev/null
+++ b/internal/gobackend/ops/where_other.go
@@ -0,0 +1,9 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+//go:build !amd64
+
+package ops
+
+func whereFloat32Arch(cond []bool, onTrue, onFalse, out []float32, onTrueIsScalar, onFalseIsScalar bool) bool {
+	return false
+}
diff --git a/internal/gobackend/ops/where_test.go b/internal/gobackend/ops/where_test.go
index 5ee121a..52a0bd2 100644
--- a/internal/gobackend/ops/where_test.go
+++ b/internal/gobackend/ops/where_test.go
@@ -13,71 +13,119 @@ import (
 )
 
 func TestWhere(t *testing.T) {
-	for _, size := range []int{10, 1000, 40000} {
-		shape := shapes.Make(dtypes.Float32, size)
-		condShape := shapes.Make(dtypes.Bool, size)
-		scalarShape := shapes.Make(dtypes.Float32)
+	sizes := []int{1, 3, 7, 8, 9, 15, 16, 23, 32, 65, 1000, 50000}
+	cases := []struct {
+		name          string
+		onTrueScalar  bool
+		onFalseScalar bool
+	}{
+		{"VectorVector", false, false},
+		{"VectorScalar", false, true},
+		{"ScalarVector", true, false},
+		{"ScalarScalar", true, true},
+	}
 
-		condData := make([]bool, size)
-		onTrueData := make([]float32, size)
-		onFalseData := []float32{-10000.0} // scalar onFalse
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			for _, size := range sizes {
+				condShape := shapes.Make(dtypes.Bool, size)
+				var trueShape, falseShape shapes.Shape
+				if c.onTrueScalar {
+					trueShape = shapes.Make(dtypes.Float32)
+				} else {
+					trueShape = shapes.Make(dtypes.Float32, size)
+				}
+				if c.onFalseScalar {
+					falseShape = shapes.Make(dtypes.Float32)
+				} else {
+					falseShape = shapes.Make(dtypes.Float32, size)
+				}
 
-		for i := 0; i < size; i++ {
-			condData[i] = (i % 2 == 0)
-			onTrueData[i] = float32(i + 1)
-		}
+				condData := make([]bool, size)
+				for i := 0; i < size; i++ {
+					condData[i] = (i%3 != 0)
+				}
 
-		builder := backend.Builder("test_where").(*gobackend.Builder)
-		main := builder.Main().(*gobackend.Function)
+				var onTrueData, onFalseData []float32
+				if c.onTrueScalar {
+					onTrueData = []float32{42.0}
+				} else {
+					onTrueData = make([]float32, size)
+					for i := 0; i < size; i++ {
+						onTrueData[i] = float32(i + 1)
+					}
+				}
 
-		condNode, err := main.Parameter("cond", condShape, nil)
-		if err != nil {
-			t.Fatalf("Parameter cond failed: %+v", err)
-		}
-		onTrueNode, err := main.Parameter("onTrue", shape, nil)
-		if err != nil {
-			t.Fatalf("Parameter onTrue failed: %+v", err)
-		}
-		onFalseNode, err := main.Parameter("onFalse", scalarShape, nil)
-		if err != nil {
-			t.Fatalf("Parameter onFalse failed: %+v", err)
-		}
+				if c.onFalseScalar {
+					onFalseData = []float32{-999.0}
+				} else {
+					onFalseData = make([]float32, size)
+					for i := 0; i < size; i++ {
+						onFalseData[i] = float32(-(i + 1))
+					}
+				}
 
-		outNode, err := ops.Where(main, condNode, onTrueNode, onFalseNode)
-		if err != nil {
-			t.Fatalf("Where failed: %+v", err)
-		}
-		err = main.Return([]compute.Value{outNode}, nil)
-		if err != nil {
-			t.Fatalf("Return failed: %+v", err)
-		}
+				builder := backend.Builder("test_where").(*gobackend.Builder)
+				main := builder.Main().(*gobackend.Function)
 
-		exec, err := builder.Compile()
-		if err != nil {
-			t.Fatalf("Compile failed: %+v", err)
-		}
-		defer exec.Finalize()
+				condNode, err := main.Parameter("cond", condShape, nil)
+				if err != nil {
+					t.Fatalf("Parameter cond failed: %+v", err)
+				}
+				onTrueNode, err := main.Parameter("onTrue", trueShape, nil)
+				if err != nil {
+					t.Fatalf("Parameter onTrue failed: %+v", err)
+				}
+				onFalseNode, err := main.Parameter("onFalse", falseShape, nil)
+				if err != nil {
+					t.Fatalf("Parameter onFalse failed: %+v", err)
+				}
 
-		condBuf := makeBuffer(t, condShape, condData)
-		onTrueBuf := makeBuffer(t, shape, onTrueData)
-		onFalseBuf := makeBuffer(t, scalarShape, onFalseData)
+				outNode, err := ops.Where(main, condNode, onTrueNode, onFalseNode)
+				if err != nil {
+					t.Fatalf("Where failed: %+v", err)
+				}
+				err = main.Return([]compute.Value{outNode}, nil)
+				if err != nil {
+					t.Fatalf("Return failed: %+v", err)
+				}
 
-		outputs, err := exec.Execute([]compute.Buffer{condBuf, onTrueBuf, onFalseBuf}, nil, 0)
-		if err != nil {
-			t.Fatalf("Execute failed: %+v", err)
-		}
+				exec, err := builder.Compile()
+				if err != nil {
+					t.Fatalf("Compile failed: %+v", err)
+				}
 
-		outFlat := outputs[0].(*gobackend.Buffer).Flat.([]float32)
-		for i := 0; i < size; i++ {
-			var expected float32
-			if condData[i] {
-				expected = onTrueData[i]
-			} else {
-				expected = onFalseData[0]
-			}
-			if outFlat[i] != expected {
-				t.Fatalf("size=%d, idx=%d: got %f, want %f", size, i, outFlat[i], expected)
+				condBuf := makeBuffer(t, condShape, condData)
+				onTrueBuf := makeBuffer(t, trueShape, onTrueData)
+				onFalseBuf := makeBuffer(t, falseShape, onFalseData)
+
+				outputs, err := exec.Execute([]compute.Buffer{condBuf, onTrueBuf, onFalseBuf}, nil, 0)
+				if err != nil {
+					t.Fatalf("Execute failed: %+v", err)
+				}
+
+				outFlat := outputs[0].(*gobackend.Buffer).Flat.([]float32)
+				for i := 0; i < size; i++ {
+					var expected float32
+					if condData[i] {
+						if c.onTrueScalar {
+							expected = onTrueData[0]
+						} else {
+							expected = onTrueData[i]
+						}
+					} else {
+						if c.onFalseScalar {
+							expected = onFalseData[0]
+						} else {
+							expected = onFalseData[i]
+						}
+					}
+					if outFlat[i] != expected {
+						t.Fatalf("%s size=%d, idx=%d: got %f, want %f", c.name, size, i, outFlat[i], expected)
+					}
+				}
+				exec.Finalize()
 			}
-		}
+		})
 	}
 }
diff --git a/internal/gobackend/passes/implicitbroadcast.go b/internal/gobackend/passes/implicitbroadcast.go
index fb23f5d..7fa5565 100644
--- a/internal/gobackend/passes/implicitbroadcast.go
+++ b/internal/gobackend/passes/implicitbroadcast.go
@@ -39,19 +39,21 @@ func (p *ImplicitBroadcastFusion) Apply(b *gobackend.Builder) (bool, error) {
 		for _, node := range f.Nodes {
 			isBinary := shapeinference.AllBinaryOperations.Has(node.OpType)
 			isUnary := shapeinference.StandardUnaryOperations.Has(node.OpType)
-			if !isBinary && !isUnary {
+			isWhere := node.OpType == compute.OpTypeWhere
+			if !isBinary && !isUnary && !isWhere {
 				continue
 			}
 
 			changed := false
-			// Binary and Unary operations for which implicit broadcasting applies.
+			// Binary, Unary, and Where operations for which implicit broadcasting applies.
 			for i, input := range node.Inputs {
 				if input.OpType != compute.OpTypeBroadcastInDim {
 					continue
 				}
 
-				// Only fuse if rank doesn't change: implicit broadcasting only works if ranks are equal.
-				if input.Inputs[0].Shape.Rank() == node.Shape.Rank() {
+				// Only fuse if rank doesn't change or if the broadcast input is scalar:
+				// implicit broadcasting supports equal ranks and scalar operands.
+				if input.Inputs[0].Shape.Rank() == node.Shape.Rank() || input.Inputs[0].Shape.IsScalar() {
 					// Fuse: use the input of the broadcast directly.
 					node.Inputs[i] = input.Inputs[0]
 					changed = true
@@ -61,7 +63,9 @@ func (p *ImplicitBroadcastFusion) Apply(b *gobackend.Builder) (bool, error) {
 			if changed {
 				var newShape shapes.Shape
 				var err error
-				if isBinary {
+				if isWhere {
+					newShape, err = shapeinference.Where(node.Inputs[0].Shape, node.Inputs[1].Shape, node.Inputs[2].Shape)
+				} else if isBinary {
 					if shapeinference.StandardBinaryOperations.Has(node.OpType) {
 						newShape, err = shapeinference.BinaryOp(node.OpType, node.Inputs[0].Shape, node.Inputs[1].Shape)
 					} else {

From 60c75432ba84cd7f3ceca34712107e0b26fcaccf Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 16:35:44 +0200
Subject: [PATCH 11/12] SIMD FusedSPDA implementation.

---
 internal/gobackend/capabilities.go         |  17 +-
 internal/gobackend/fusedops/avx2/sdpa.go   | 669 +++++++++++++++++
 internal/gobackend/fusedops/avx512/sdpa.go | 804 +++++++++++++++++++++
 internal/gobackend/fusedops/sdpa.go        | 174 +++--
 internal/gobackend/sdpa_arch.go            |  39 +
 5 files changed, 1640 insertions(+), 63 deletions(-)
 create mode 100644 internal/gobackend/fusedops/avx2/sdpa.go
 create mode 100644 internal/gobackend/fusedops/avx512/sdpa.go
 create mode 100644 internal/gobackend/sdpa_arch.go

diff --git a/internal/gobackend/capabilities.go b/internal/gobackend/capabilities.go
index d49b9bd..5164b58 100644
--- a/internal/gobackend/capabilities.go
+++ b/internal/gobackend/capabilities.go
@@ -143,16 +143,13 @@ var Capabilities = compute.Capabilities{
 		compute.OpTypeSort:  true,
 
 		// Fused operations:
-		compute.OpTypeFusedSoftmax:    true,
-		compute.OpTypeFusedLayerNorm:  true,
-		compute.OpTypeFusedActivation:    true,
-		compute.OpTypeFusedActivationVJP: true,
-		compute.OpTypeFusedDense:         true,
-		compute.OpTypeFusedDenseVJP:      true,
-		// - Fused SPDA: Temporarily DISABLED, the new matmul with SIMD support is much faster (+3x faster),
-		//   so this fused op ends up being slower. TODO: add a SIMD version of the fused SPDA -- or split it
-		//   into a normal matmul (and use the SIMD matmul) + a fused softmax.
-		compute.OpTypeFusedScaledDotProductAttention: false,
+		compute.OpTypeFusedSoftmax:                   true,
+		compute.OpTypeFusedLayerNorm:                 true,
+		compute.OpTypeFusedActivation:                true,
+		compute.OpTypeFusedActivationVJP:             true,
+		compute.OpTypeFusedDense:                     true,
+		compute.OpTypeFusedDenseVJP:                  true,
+		compute.OpTypeFusedScaledDotProductAttention: true,
 		compute.OpTypeFusedAttentionQKVProjection:    true,
 		compute.OpTypeFusedQuantizedDense:            true,
 		compute.OpTypeQuantizedEmbeddingLookup:       true,
diff --git a/internal/gobackend/fusedops/avx2/sdpa.go b/internal/gobackend/fusedops/avx2/sdpa.go
new file mode 100644
index 0000000..1ab8fab
--- /dev/null
+++ b/internal/gobackend/fusedops/avx2/sdpa.go
@@ -0,0 +1,669 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+//go:build amd64 && goexperiment.simd
+
+package avx2
+
+import (
+	"math"
+	"simd/archsimd"
+
+	"github.com/gomlx/compute/internal/fastmath"
+	"github.com/gomlx/compute/internal/gobackend"
+)
+
+func init() {
+	if gobackend.IsAVX2Allowed {
+		gobackend.SetSDPAArchDispatcher(gobackend.PriorityArch, DispatchSDPAAVX2)
+	}
+}
+
+// reduceSum8 reduces a Float32x8 vector to a single float32 sum.
+func reduceSum8(v archsimd.Float32x8) float32 {
+	v4 := v.GetLo().Add(v.GetHi())
+	v2 := v4.ConcatAddPairs(v4)
+	return v2.GetElem(0) + v2.GetElem(1)
+}
+
+// reduceMax8 reduces a Float32x8 vector to its maximum element.
+func reduceMax8(v archsimd.Float32x8) float32 {
+	v4 := v.GetLo().Max(v.GetHi())
+	var arr [4]float32
+	v4.StoreArray(&arr)
+	m := max(arr[0], arr[1])
+	return max(m, max(arr[2], arr[3]))
+}
+
+// exp256 approximates e^x for 8 float32s using Cephes degree-7 Horner polynomial.
+func exp256(x archsimd.Float32x8) archsimd.Float32x8 {
+	const (
+		maxLogF = 88.02969187150841
+		minLogF = -88.02969187150841
+		log2E   = 1.44269504088896341
+		ln2Hi   = 0.693359375
+		ln2Lo   = -2.12194440e-4
+
+		p7 = 1.9875691500e-4
+		p6 = 1.3981999507e-3
+		p5 = 8.3334519073e-3
+		p4 = 4.1665795894e-2
+		p3 = 1.6666665459e-1
+		p2 = 5.0000001201e-1
+	)
+	vMaxLog := archsimd.BroadcastFloat32x8(maxLogF)
+	vMinLog := archsimd.BroadcastFloat32x8(minLogF)
+	vLog2E := archsimd.BroadcastFloat32x8(log2E)
+	vHalf := archsimd.BroadcastFloat32x8(0.5)
+	vLn2Hi := archsimd.BroadcastFloat32x8(ln2Hi)
+	vLn2Lo := archsimd.BroadcastFloat32x8(ln2Lo)
+
+	vP7 := archsimd.BroadcastFloat32x8(p7)
+	vP6 := archsimd.BroadcastFloat32x8(p6)
+	vP5 := archsimd.BroadcastFloat32x8(p5)
+	vP4 := archsimd.BroadcastFloat32x8(p4)
+	vP3 := archsimd.BroadcastFloat32x8(p3)
+	vP2 := archsimd.BroadcastFloat32x8(p2)
+	vOne := archsimd.BroadcastFloat32x8(1.0)
+	v127 := archsimd.BroadcastUint32x8(127)
+
+	xClamped := x.Max(vMinLog).Min(vMaxLog)
+	z := xClamped.MulAdd(vLog2E, vHalf).Floor()
+	g := xClamped.Sub(z.Mul(vLn2Hi)).Sub(z.Mul(vLn2Lo))
+
+	n := z.ConvertToInt32().AsUint32x8().Add(v127).ShiftAllLeft(23).BitsToFloat32()
+
+	poly := vP7.MulAdd(g, vP6)
+	poly = poly.MulAdd(g, vP5)
+	poly = poly.MulAdd(g, vP4)
+	poly = poly.MulAdd(g, vP3)
+	poly = poly.MulAdd(g, vP2)
+	poly = poly.Mul(g).Mul(g).Add(g).Add(vOne)
+
+	return n.Mul(poly)
+}
+
+// DispatchSDPAAVX2 executes scaled dot-product attention using AVX2.
+func DispatchSDPAAVX2(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen, headDim int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) bool {
+	switch headDim {
+	case 32:
+		sdpaFloat32AVX2Dim32(
+			q, k, v, output,
+			qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride,
+			additiveMask, booleanMask, maskGroupStride,
+			additiveBias, biasGroupStride,
+			scoresScratch,
+			groupSize, seqLen, kvLen,
+			scale, causal, qLimit, kvLimit,
+		)
+		return true
+	case 64:
+		sdpaFloat32AVX2Dim64(
+			q, k, v, output,
+			qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride,
+			additiveMask, booleanMask, maskGroupStride,
+			additiveBias, biasGroupStride,
+			scoresScratch,
+			groupSize, seqLen, kvLen,
+			scale, causal, qLimit, kvLimit,
+		)
+		return true
+	case 128:
+		sdpaFloat32AVX2Dim128(
+			q, k, v, output,
+			qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride,
+			additiveMask, booleanMask, maskGroupStride,
+			additiveBias, biasGroupStride,
+			scoresScratch,
+			groupSize, seqLen, kvLen,
+			scale, causal, qLimit, kvLimit,
+		)
+		return true
+	default:
+		sdpaFloat32AVX2General(
+			q, k, v, output,
+			qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride,
+			additiveMask, booleanMask, maskGroupStride,
+			additiveBias, biasGroupStride,
+			scoresScratch,
+			groupSize, seqLen, kvLen, headDim,
+			scale, causal, qLimit, kvLimit,
+		)
+		return true
+	}
+}
+
+func softmaxAndValueAccumDim32(
+	v, output []float32,
+	outBase, kvOff, kvSeqStride int,
+	scoresScratch []float32,
+	kvLenUnmasked int,
+	rowMax float32,
+) {
+	var sum float32
+	k := 0
+	vRowMax := archsimd.BroadcastFloat32x8(rowMax)
+	vSum := archsimd.BroadcastFloat32x8(0)
+	for ; k+8 <= kvLenUnmasked; k += 8 {
+		vScores := archsimd.LoadFloat32x8(scoresScratch[k : k+8])
+		expV := exp256(vScores.Sub(vRowMax))
+		expV.Store(scoresScratch[k : k+8])
+		vSum = vSum.Add(expV)
+	}
+	if k > 0 {
+		sum += reduceSum8(vSum)
+	}
+	for ; k < kvLenUnmasked; k++ {
+		s := scoresScratch[k]
+		if s == float32(math.Inf(-1)) {
+			scoresScratch[k] = 0
+		} else {
+			e := fastmath.Exp32(s - rowMax)
+			scoresScratch[k] = e
+			sum += e
+		}
+	}
+
+	var invSum float32
+	if sum > 0 {
+		invSum = 1.0 / sum
+	}
+
+	k = 0
+	vInvSum := archsimd.BroadcastFloat32x8(invSum)
+	for ; k+8 <= kvLenUnmasked; k += 8 {
+		vScores := archsimd.LoadFloat32x8(scoresScratch[k : k+8])
+		vScores.Mul(vInvSum).Store(scoresScratch[k : k+8])
+	}
+	for ; k < kvLenUnmasked; k++ {
+		scoresScratch[k] *= invSum
+	}
+
+	out0 := archsimd.BroadcastFloat32x8(0)
+	out1 := archsimd.BroadcastFloat32x8(0)
+	out2 := archsimd.BroadcastFloat32x8(0)
+	out3 := archsimd.BroadcastFloat32x8(0)
+
+	for ki := 0; ki < kvLenUnmasked; ki++ {
+		w := scoresScratch[ki]
+		if w == 0 {
+			continue
+		}
+		vW := archsimd.BroadcastFloat32x8(w)
+		vBase := kvOff + ki*kvSeqStride
+		v0 := archsimd.LoadFloat32x8(v[vBase : vBase+8])
+		v1 := archsimd.LoadFloat32x8(v[vBase+8 : vBase+16])
+		v2 := archsimd.LoadFloat32x8(v[vBase+16 : vBase+24])
+		v3 := archsimd.LoadFloat32x8(v[vBase+24 : vBase+32])
+		out0 = out0.MulAdd(vW, v0)
+		out1 = out1.MulAdd(vW, v1)
+		out2 = out2.MulAdd(vW, v2)
+		out3 = out3.MulAdd(vW, v3)
+	}
+
+	out0.Store(output[outBase : outBase+8])
+	out1.Store(output[outBase+8 : outBase+16])
+	out2.Store(output[outBase+16 : outBase+24])
+	out3.Store(output[outBase+24 : outBase+32])
+}
+
+func sdpaFloat32AVX2Dim32(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) {
+	for gIdx := range groupSize {
+		gQOff := qOff + gIdx*qGroupStride
+		gMaskOff := gIdx * maskGroupStride
+		gBiasOff := gIdx * biasGroupStride
+
+		for qIdx := range seqLen {
+			outBase := gQOff + qIdx*qSeqStride
+			if qIdx >= qLimit {
+				for d := range 32 {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			kvLenUnmasked := kvLen
+			if kvLimit < kvLenUnmasked {
+				kvLenUnmasked = kvLimit
+			}
+			if causal && qIdx+1 < kvLenUnmasked {
+				kvLenUnmasked = qIdx + 1
+			}
+
+			if kvLenUnmasked <= 0 {
+				for d := range 32 {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			qBase := gQOff + qIdx*qSeqStride
+			maskIdxBase := gMaskOff + qIdx*kvLen
+			biasIdxBase := gBiasOff + qIdx*kvLen
+
+			q0 := archsimd.LoadFloat32x8(q[qBase : qBase+8])
+			q1 := archsimd.LoadFloat32x8(q[qBase+8 : qBase+16])
+			q2 := archsimd.LoadFloat32x8(q[qBase+16 : qBase+24])
+			q3 := archsimd.LoadFloat32x8(q[qBase+24 : qBase+32])
+
+			rowMax := float32(math.Inf(-1))
+
+			for ki := range kvLenUnmasked {
+				maskIdx := maskIdxBase + ki
+				if len(booleanMask) > 0 && !booleanMask[maskIdx] {
+					scoresScratch[ki] = float32(math.Inf(-1))
+					continue
+				}
+
+				kBase := kvOff + ki*kvSeqStride
+				k0 := archsimd.LoadFloat32x8(k[kBase : kBase+8])
+				k1 := archsimd.LoadFloat32x8(k[kBase+8 : kBase+16])
+				k2 := archsimd.LoadFloat32x8(k[kBase+16 : kBase+24])
+				k3 := archsimd.LoadFloat32x8(k[kBase+24 : kBase+32])
+
+				sumVec := q0.Mul(k0).Add(q1.Mul(k1)).Add(q2.Mul(k2)).Add(q3.Mul(k3))
+				dot := reduceSum8(sumVec)
+				s := dot * scale
+				if len(additiveBias) > 0 {
+					s += additiveBias[biasIdxBase+ki]
+				}
+				if len(additiveMask) > 0 {
+					s += additiveMask[maskIdx]
+				}
+				scoresScratch[ki] = s
+				if s > rowMax {
+					rowMax = s
+				}
+			}
+
+			softmaxAndValueAccumDim32(v, output, outBase, kvOff, kvSeqStride, scoresScratch, kvLenUnmasked, rowMax)
+		}
+	}
+}
+
+func softmaxAndValueAccumDim64(
+	v, output []float32,
+	outBase, kvOff, kvSeqStride int,
+	scoresScratch []float32,
+	kvLenUnmasked int,
+	rowMax float32,
+) {
+	var sum float32
+	k := 0
+	vRowMax := archsimd.BroadcastFloat32x8(rowMax)
+	vSum := archsimd.BroadcastFloat32x8(0)
+	for ; k+8 <= kvLenUnmasked; k += 8 {
+		vScores := archsimd.LoadFloat32x8(scoresScratch[k : k+8])
+		expV := exp256(vScores.Sub(vRowMax))
+		expV.Store(scoresScratch[k : k+8])
+		vSum = vSum.Add(expV)
+	}
+	if k > 0 {
+		sum += reduceSum8(vSum)
+	}
+	for ; k < kvLenUnmasked; k++ {
+		s := scoresScratch[k]
+		if s == float32(math.Inf(-1)) {
+			scoresScratch[k] = 0
+		} else {
+			e := fastmath.Exp32(s - rowMax)
+			scoresScratch[k] = e
+			sum += e
+		}
+	}
+
+	var invSum float32
+	if sum > 0 {
+		invSum = 1.0 / sum
+	}
+
+	k = 0
+	vInvSum := archsimd.BroadcastFloat32x8(invSum)
+	for ; k+8 <= kvLenUnmasked; k += 8 {
+		vScores := archsimd.LoadFloat32x8(scoresScratch[k : k+8])
+		vScores.Mul(vInvSum).Store(scoresScratch[k : k+8])
+	}
+	for ; k < kvLenUnmasked; k++ {
+		scoresScratch[k] *= invSum
+	}
+
+	out0 := archsimd.BroadcastFloat32x8(0)
+	out1 := archsimd.BroadcastFloat32x8(0)
+	out2 := archsimd.BroadcastFloat32x8(0)
+	out3 := archsimd.BroadcastFloat32x8(0)
+	out4 := archsimd.BroadcastFloat32x8(0)
+	out5 := archsimd.BroadcastFloat32x8(0)
+	out6 := archsimd.BroadcastFloat32x8(0)
+	out7 := archsimd.BroadcastFloat32x8(0)
+
+	for ki := 0; ki < kvLenUnmasked; ki++ {
+		w := scoresScratch[ki]
+		if w == 0 {
+			continue
+		}
+		vW := archsimd.BroadcastFloat32x8(w)
+		vBase := kvOff + ki*kvSeqStride
+		v0 := archsimd.LoadFloat32x8(v[vBase : vBase+8])
+		v1 := archsimd.LoadFloat32x8(v[vBase+8 : vBase+16])
+		v2 := archsimd.LoadFloat32x8(v[vBase+16 : vBase+24])
+		v3 := archsimd.LoadFloat32x8(v[vBase+24 : vBase+32])
+		v4 := archsimd.LoadFloat32x8(v[vBase+32 : vBase+40])
+		v5 := archsimd.LoadFloat32x8(v[vBase+40 : vBase+48])
+		v6 := archsimd.LoadFloat32x8(v[vBase+48 : vBase+56])
+		v7 := archsimd.LoadFloat32x8(v[vBase+56 : vBase+64])
+		out0 = out0.MulAdd(vW, v0)
+		out1 = out1.MulAdd(vW, v1)
+		out2 = out2.MulAdd(vW, v2)
+		out3 = out3.MulAdd(vW, v3)
+		out4 = out4.MulAdd(vW, v4)
+		out5 = out5.MulAdd(vW, v5)
+		out6 = out6.MulAdd(vW, v6)
+		out7 = out7.MulAdd(vW, v7)
+	}
+
+	out0.Store(output[outBase : outBase+8])
+	out1.Store(output[outBase+8 : outBase+16])
+	out2.Store(output[outBase+16 : outBase+24])
+	out3.Store(output[outBase+24 : outBase+32])
+	out4.Store(output[outBase+32 : outBase+40])
+	out5.Store(output[outBase+40 : outBase+48])
+	out6.Store(output[outBase+48 : outBase+56])
+	out7.Store(output[outBase+56 : outBase+64])
+}
+
+func sdpaFloat32AVX2Dim64(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) {
+	for gIdx := range groupSize {
+		gQOff := qOff + gIdx*qGroupStride
+		gMaskOff := gIdx * maskGroupStride
+		gBiasOff := gIdx * biasGroupStride
+
+		for qIdx := range seqLen {
+			outBase := gQOff + qIdx*qSeqStride
+			if qIdx >= qLimit {
+				for d := range 64 {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			kvLenUnmasked := kvLen
+			if kvLimit < kvLenUnmasked {
+				kvLenUnmasked = kvLimit
+			}
+			if causal && qIdx+1 < kvLenUnmasked {
+				kvLenUnmasked = qIdx + 1
+			}
+
+			if kvLenUnmasked <= 0 {
+				for d := range 64 {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			qBase := gQOff + qIdx*qSeqStride
+			maskIdxBase := gMaskOff + qIdx*kvLen
+			biasIdxBase := gBiasOff + qIdx*kvLen
+
+			q0 := archsimd.LoadFloat32x8(q[qBase : qBase+8])
+			q1 := archsimd.LoadFloat32x8(q[qBase+8 : qBase+16])
+			q2 := archsimd.LoadFloat32x8(q[qBase+16 : qBase+24])
+			q3 := archsimd.LoadFloat32x8(q[qBase+24 : qBase+32])
+			q4 := archsimd.LoadFloat32x8(q[qBase+32 : qBase+40])
+			q5 := archsimd.LoadFloat32x8(q[qBase+40 : qBase+48])
+			q6 := archsimd.LoadFloat32x8(q[qBase+48 : qBase+56])
+			q7 := archsimd.LoadFloat32x8(q[qBase+56 : qBase+64])
+
+			rowMax := float32(math.Inf(-1))
+
+			for ki := range kvLenUnmasked {
+				maskIdx := maskIdxBase + ki
+				if len(booleanMask) > 0 && !booleanMask[maskIdx] {
+					scoresScratch[ki] = float32(math.Inf(-1))
+					continue
+				}
+
+				kBase := kvOff + ki*kvSeqStride
+				k0 := archsimd.LoadFloat32x8(k[kBase : kBase+8])
+				k1 := archsimd.LoadFloat32x8(k[kBase+8 : kBase+16])
+				k2 := archsimd.LoadFloat32x8(k[kBase+16 : kBase+24])
+				k3 := archsimd.LoadFloat32x8(k[kBase+24 : kBase+32])
+				k4 := archsimd.LoadFloat32x8(k[kBase+32 : kBase+40])
+				k5 := archsimd.LoadFloat32x8(k[kBase+40 : kBase+48])
+				k6 := archsimd.LoadFloat32x8(k[kBase+48 : kBase+56])
+				k7 := archsimd.LoadFloat32x8(k[kBase+56 : kBase+64])
+
+				sumVec0 := q0.Mul(k0).Add(q1.Mul(k1)).Add(q2.Mul(k2)).Add(q3.Mul(k3))
+				sumVec1 := q4.Mul(k4).Add(q5.Mul(k5)).Add(q6.Mul(k6)).Add(q7.Mul(k7))
+				dot := reduceSum8(sumVec0.Add(sumVec1))
+				s := dot * scale
+				if len(additiveBias) > 0 {
+					s += additiveBias[biasIdxBase+ki]
+				}
+				if len(additiveMask) > 0 {
+					s += additiveMask[maskIdx]
+				}
+				scoresScratch[ki] = s
+				if s > rowMax {
+					rowMax = s
+				}
+			}
+
+			softmaxAndValueAccumDim64(v, output, outBase, kvOff, kvSeqStride, scoresScratch, kvLenUnmasked, rowMax)
+		}
+	}
+}
+
+func sdpaFloat32AVX2Dim128(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) {
+	// For Dim128 on AVX2, delegate to General (step 8).
+	sdpaFloat32AVX2General(
+		q, k, v, output,
+		qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride,
+		additiveMask, booleanMask, maskGroupStride,
+		additiveBias, biasGroupStride,
+		scoresScratch,
+		groupSize, seqLen, kvLen, 128,
+		scale, causal, qLimit, kvLimit,
+	)
+}
+
+func sdpaFloat32AVX2General(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen, headDim int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) {
+	for gIdx := range groupSize {
+		gQOff := qOff + gIdx*qGroupStride
+		gMaskOff := gIdx * maskGroupStride
+		gBiasOff := gIdx * biasGroupStride
+
+		for qIdx := range seqLen {
+			outBase := gQOff + qIdx*qSeqStride
+			if qIdx >= qLimit {
+				for d := range headDim {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			kvLenUnmasked := kvLen
+			if kvLimit < kvLenUnmasked {
+				kvLenUnmasked = kvLimit
+			}
+			if causal && qIdx+1 < kvLenUnmasked {
+				kvLenUnmasked = qIdx + 1
+			}
+
+			if kvLenUnmasked <= 0 {
+				for d := range headDim {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			qBase := gQOff + qIdx*qSeqStride
+			maskIdxBase := gMaskOff + qIdx*kvLen
+			biasIdxBase := gBiasOff + qIdx*kvLen
+
+			rowMax := float32(math.Inf(-1))
+
+			for ki := range kvLenUnmasked {
+				maskIdx := maskIdxBase + ki
+				if len(booleanMask) > 0 && !booleanMask[maskIdx] {
+					scoresScratch[ki] = float32(math.Inf(-1))
+					continue
+				}
+
+				kBase := kvOff + ki*kvSeqStride
+				d := 0
+				sumVec := archsimd.BroadcastFloat32x8(0)
+				for ; d+8 <= headDim; d += 8 {
+					qV := archsimd.LoadFloat32x8(q[qBase+d : qBase+d+8])
+					kV := archsimd.LoadFloat32x8(k[kBase+d : kBase+d+8])
+					sumVec = sumVec.Add(qV.Mul(kV))
+				}
+				var dot float32
+				if d > 0 {
+					dot = reduceSum8(sumVec)
+				}
+				for ; d < headDim; d++ {
+					dot += q[qBase+d] * k[kBase+d]
+				}
+
+				s := dot * scale
+				if len(additiveBias) > 0 {
+					s += additiveBias[biasIdxBase+ki]
+				}
+				if len(additiveMask) > 0 {
+					s += additiveMask[maskIdx]
+				}
+				scoresScratch[ki] = s
+				if s > rowMax {
+					rowMax = s
+				}
+			}
+
+			// Softmax
+			var sum float32
+			ki := 0
+			vRowMax := archsimd.BroadcastFloat32x8(rowMax)
+			vSum := archsimd.BroadcastFloat32x8(0)
+			for ; ki+8 <= kvLenUnmasked; ki += 8 {
+				vScores := archsimd.LoadFloat32x8(scoresScratch[ki : ki+8])
+				expV := exp256(vScores.Sub(vRowMax))
+				expV.Store(scoresScratch[ki : ki+8])
+				vSum = vSum.Add(expV)
+			}
+			if ki > 0 {
+				sum += reduceSum8(vSum)
+			}
+			for ; ki < kvLenUnmasked; ki++ {
+				s := scoresScratch[ki]
+				if s == float32(math.Inf(-1)) {
+					scoresScratch[ki] = 0
+				} else {
+					e := fastmath.Exp32(s - rowMax)
+					scoresScratch[ki] = e
+					sum += e
+				}
+			}
+
+			var invSum float32
+			if sum > 0 {
+				invSum = 1.0 / sum
+			}
+
+			ki = 0
+			vInvSum := archsimd.BroadcastFloat32x8(invSum)
+			for ; ki+8 <= kvLenUnmasked; ki += 8 {
+				vScores := archsimd.LoadFloat32x8(scoresScratch[ki : ki+8])
+				vScores.Mul(vInvSum).Store(scoresScratch[ki : ki+8])
+			}
+			for ; ki < kvLenUnmasked; ki++ {
+				scoresScratch[ki] *= invSum
+			}
+
+			// Value accumulation along d
+			d := 0
+			for ; d+8 <= headDim; d += 8 {
+				outV := archsimd.BroadcastFloat32x8(0)
+				for kIdx := range kvLenUnmasked {
+					w := scoresScratch[kIdx]
+					if w == 0 {
+						continue
+					}
+					vW := archsimd.BroadcastFloat32x8(w)
+					vBase := kvOff + kIdx*kvSeqStride + d
+					vVal := archsimd.LoadFloat32x8(v[vBase : vBase+8])
+					outV = outV.MulAdd(vW, vVal)
+				}
+				outV.Store(output[outBase+d : outBase+d+8])
+			}
+			for ; d < headDim; d++ {
+				var acc float32
+				for kIdx := range kvLenUnmasked {
+					w := scoresScratch[kIdx]
+					if w != 0 {
+						acc += w * v[kvOff+kIdx*kvSeqStride+d]
+					}
+				}
+				output[outBase+d] = acc
+			}
+		}
+	}
+}
diff --git a/internal/gobackend/fusedops/avx512/sdpa.go b/internal/gobackend/fusedops/avx512/sdpa.go
new file mode 100644
index 0000000..e8ae2ee
--- /dev/null
+++ b/internal/gobackend/fusedops/avx512/sdpa.go
@@ -0,0 +1,804 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+//go:build amd64 && goexperiment.simd
+
+package avx512
+
+import (
+	"math"
+	"simd/archsimd"
+
+	"github.com/gomlx/compute/internal/fastmath"
+	"github.com/gomlx/compute/internal/gobackend"
+)
+
+func init() {
+	if gobackend.IsAVX512Allowed {
+		gobackend.SetSDPAArchDispatcher(gobackend.PriorityArch+1, DispatchSDPAAVX512)
+	}
+}
+
+// reduceSum8 reduces a Float32x8 vector to a single float32 sum.
+func reduceSum8(v archsimd.Float32x8) float32 {
+	v4 := v.GetLo().Add(v.GetHi())
+	v2 := v4.ConcatAddPairs(v4)
+	return v2.GetElem(0) + v2.GetElem(1)
+}
+
+// reduceSum16 reduces a Float32x16 vector to a single float32 sum.
+func reduceSum16(v archsimd.Float32x16) float32 {
+	return reduceSum8(v.GetLo().Add(v.GetHi()))
+}
+
+// reduceMax8 reduces a Float32x8 vector to its maximum element.
+func reduceMax8(v archsimd.Float32x8) float32 {
+	v4 := v.GetLo().Max(v.GetHi())
+	var arr [4]float32
+	v4.StoreArray(&arr)
+	m := max(arr[0], arr[1])
+	return max(m, max(arr[2], arr[3]))
+}
+
+// reduceMax16 reduces a Float32x16 vector to its maximum element.
+func reduceMax16(v archsimd.Float32x16) float32 {
+	return reduceMax8(v.GetLo().Max(v.GetHi()))
+}
+
+// exp512 approximates e^x for 16 float32s using Cephes degree-7 Horner polynomial.
+func exp512(x archsimd.Float32x16) archsimd.Float32x16 {
+	const (
+		maxLogF = 88.02969187150841
+		minLogF = -88.02969187150841
+		log2E   = 1.44269504088896341
+		ln2Hi   = 0.693359375
+		ln2Lo   = -2.12194440e-4
+
+		p7 = 1.9875691500e-4
+		p6 = 1.3981999507e-3
+		p5 = 8.3334519073e-3
+		p4 = 4.1665795894e-2
+		p3 = 1.6666665459e-1
+		p2 = 5.0000001201e-1
+	)
+	vMaxLog := archsimd.BroadcastFloat32x16(maxLogF)
+	vMinLog := archsimd.BroadcastFloat32x16(minLogF)
+	vLog2E := archsimd.BroadcastFloat32x16(log2E)
+	vHalf := archsimd.BroadcastFloat32x16(0.5)
+	vLn2Hi := archsimd.BroadcastFloat32x16(ln2Hi)
+	vLn2Lo := archsimd.BroadcastFloat32x16(ln2Lo)
+
+	vP7 := archsimd.BroadcastFloat32x16(p7)
+	vP6 := archsimd.BroadcastFloat32x16(p6)
+	vP5 := archsimd.BroadcastFloat32x16(p5)
+	vP4 := archsimd.BroadcastFloat32x16(p4)
+	vP3 := archsimd.BroadcastFloat32x16(p3)
+	vP2 := archsimd.BroadcastFloat32x16(p2)
+	vOne := archsimd.BroadcastFloat32x16(1.0)
+	v127 := archsimd.BroadcastUint32x16(127)
+
+	xClamped := x.Max(vMinLog).Min(vMaxLog)
+	z := xClamped.MulAdd(vLog2E, vHalf).RoundScaled(0)
+	g := xClamped.Sub(z.Mul(vLn2Hi)).Sub(z.Mul(vLn2Lo))
+
+	n := z.ConvertToInt32().AsUint32x16().Add(v127).ShiftAllLeft(23).BitsToFloat32()
+
+	poly := vP7.MulAdd(g, vP6)
+	poly = poly.MulAdd(g, vP5)
+	poly = poly.MulAdd(g, vP4)
+	poly = poly.MulAdd(g, vP3)
+	poly = poly.MulAdd(g, vP2)
+	poly = poly.Mul(g).Mul(g).Add(g).Add(vOne)
+
+	return n.Mul(poly)
+}
+
+// DispatchSDPAAVX512 executes scaled dot-product attention using AVX-512.
+func DispatchSDPAAVX512(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen, headDim int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) bool {
+	switch headDim {
+	case 32:
+		sdpaFloat32AVX512Dim32(
+			q, k, v, output,
+			qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride,
+			additiveMask, booleanMask, maskGroupStride,
+			additiveBias, biasGroupStride,
+			scoresScratch,
+			groupSize, seqLen, kvLen,
+			scale, causal, qLimit, kvLimit,
+		)
+		return true
+	case 64:
+		sdpaFloat32AVX512Dim64(
+			q, k, v, output,
+			qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride,
+			additiveMask, booleanMask, maskGroupStride,
+			additiveBias, biasGroupStride,
+			scoresScratch,
+			groupSize, seqLen, kvLen,
+			scale, causal, qLimit, kvLimit,
+		)
+		return true
+	case 128:
+		sdpaFloat32AVX512Dim128(
+			q, k, v, output,
+			qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride,
+			additiveMask, booleanMask, maskGroupStride,
+			additiveBias, biasGroupStride,
+			scoresScratch,
+			groupSize, seqLen, kvLen,
+			scale, causal, qLimit, kvLimit,
+		)
+		return true
+	default:
+		sdpaFloat32AVX512General(
+			q, k, v, output,
+			qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride,
+			additiveMask, booleanMask, maskGroupStride,
+			additiveBias, biasGroupStride,
+			scoresScratch,
+			groupSize, seqLen, kvLen, headDim,
+			scale, causal, qLimit, kvLimit,
+		)
+		return true
+	}
+}
+
+func softmaxAndValueAccumDim32(
+	v, output []float32,
+	outBase, kvOff, kvSeqStride int,
+	scoresScratch []float32,
+	kvLenUnmasked int,
+	rowMax float32,
+) {
+	// Softmax: exp(scores - rowMax) and sum
+	var sum float32
+	k := 0
+	vRowMax := archsimd.BroadcastFloat32x16(rowMax)
+	vSum := archsimd.BroadcastFloat32x16(0)
+	for ; k+16 <= kvLenUnmasked; k += 16 {
+		vScores := archsimd.LoadFloat32x16(scoresScratch[k : k+16])
+		expV := exp512(vScores.Sub(vRowMax))
+		expV.Store(scoresScratch[k : k+16])
+		vSum = vSum.Add(expV)
+	}
+	if k > 0 {
+		sum += reduceSum16(vSum)
+	}
+	for ; k < kvLenUnmasked; k++ {
+		s := scoresScratch[k]
+		if s == float32(math.Inf(-1)) {
+			scoresScratch[k] = 0
+		} else {
+			e := fastmath.Exp32(s - rowMax)
+			scoresScratch[k] = e
+			sum += e
+		}
+	}
+
+	var invSum float32
+	if sum > 0 {
+		invSum = 1.0 / sum
+	}
+
+	// Normalize scores
+	k = 0
+	vInvSum := archsimd.BroadcastFloat32x16(invSum)
+	for ; k+16 <= kvLenUnmasked; k += 16 {
+		vScores := archsimd.LoadFloat32x16(scoresScratch[k : k+16])
+		vScores.Mul(vInvSum).Store(scoresScratch[k : k+16])
+	}
+	for ; k < kvLenUnmasked; k++ {
+		scoresScratch[k] *= invSum
+	}
+
+	// Value accumulation into vector registers
+	out0 := archsimd.BroadcastFloat32x16(0)
+	out1 := archsimd.BroadcastFloat32x16(0)
+
+	for ki := 0; ki < kvLenUnmasked; ki++ {
+		w := scoresScratch[ki]
+		if w == 0 {
+			continue
+		}
+		vW := archsimd.BroadcastFloat32x16(w)
+		vBase := kvOff + ki*kvSeqStride
+		v0 := archsimd.LoadFloat32x16(v[vBase : vBase+16])
+		v1 := archsimd.LoadFloat32x16(v[vBase+16 : vBase+32])
+		out0 = out0.MulAdd(vW, v0)
+		out1 = out1.MulAdd(vW, v1)
+	}
+
+	out0.Store(output[outBase : outBase+16])
+	out1.Store(output[outBase+16 : outBase+32])
+}
+
+func sdpaFloat32AVX512Dim32(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) {
+	for gIdx := range groupSize {
+		gQOff := qOff + gIdx*qGroupStride
+		gMaskOff := gIdx * maskGroupStride
+		gBiasOff := gIdx * biasGroupStride
+
+		for qIdx := range seqLen {
+			outBase := gQOff + qIdx*qSeqStride
+			if qIdx >= qLimit {
+				for d := range 32 {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			kvLenUnmasked := kvLen
+			if kvLimit < kvLenUnmasked {
+				kvLenUnmasked = kvLimit
+			}
+			if causal && qIdx+1 < kvLenUnmasked {
+				kvLenUnmasked = qIdx + 1
+			}
+
+			if kvLenUnmasked <= 0 {
+				for d := range 32 {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			qBase := gQOff + qIdx*qSeqStride
+			maskIdxBase := gMaskOff + qIdx*kvLen
+			biasIdxBase := gBiasOff + qIdx*kvLen
+
+			q0 := archsimd.LoadFloat32x16(q[qBase : qBase+16])
+			q1 := archsimd.LoadFloat32x16(q[qBase+16 : qBase+32])
+
+			rowMax := float32(math.Inf(-1))
+
+			for ki := range kvLenUnmasked {
+				maskIdx := maskIdxBase + ki
+				if len(booleanMask) > 0 && !booleanMask[maskIdx] {
+					scoresScratch[ki] = float32(math.Inf(-1))
+					continue
+				}
+
+				kBase := kvOff + ki*kvSeqStride
+				k0 := archsimd.LoadFloat32x16(k[kBase : kBase+16])
+				k1 := archsimd.LoadFloat32x16(k[kBase+16 : kBase+32])
+
+				dot := reduceSum16(q0.Mul(k0).Add(q1.Mul(k1)))
+				s := dot * scale
+				if len(additiveBias) > 0 {
+					s += additiveBias[biasIdxBase+ki]
+				}
+				if len(additiveMask) > 0 {
+					s += additiveMask[maskIdx]
+				}
+				scoresScratch[ki] = s
+				if s > rowMax {
+					rowMax = s
+				}
+			}
+
+			softmaxAndValueAccumDim32(v, output, outBase, kvOff, kvSeqStride, scoresScratch, kvLenUnmasked, rowMax)
+		}
+	}
+}
+
+func softmaxAndValueAccumDim64(
+	v, output []float32,
+	outBase, kvOff, kvSeqStride int,
+	scoresScratch []float32,
+	kvLenUnmasked int,
+	rowMax float32,
+) {
+	var sum float32
+	k := 0
+	vRowMax := archsimd.BroadcastFloat32x16(rowMax)
+	vSum := archsimd.BroadcastFloat32x16(0)
+	for ; k+16 <= kvLenUnmasked; k += 16 {
+		vScores := archsimd.LoadFloat32x16(scoresScratch[k : k+16])
+		expV := exp512(vScores.Sub(vRowMax))
+		expV.Store(scoresScratch[k : k+16])
+		vSum = vSum.Add(expV)
+	}
+	if k > 0 {
+		sum += reduceSum16(vSum)
+	}
+	for ; k < kvLenUnmasked; k++ {
+		s := scoresScratch[k]
+		if s == float32(math.Inf(-1)) {
+			scoresScratch[k] = 0
+		} else {
+			e := fastmath.Exp32(s - rowMax)
+			scoresScratch[k] = e
+			sum += e
+		}
+	}
+
+	var invSum float32
+	if sum > 0 {
+		invSum = 1.0 / sum
+	}
+
+	k = 0
+	vInvSum := archsimd.BroadcastFloat32x16(invSum)
+	for ; k+16 <= kvLenUnmasked; k += 16 {
+		vScores := archsimd.LoadFloat32x16(scoresScratch[k : k+16])
+		vScores.Mul(vInvSum).Store(scoresScratch[k : k+16])
+	}
+	for ; k < kvLenUnmasked; k++ {
+		scoresScratch[k] *= invSum
+	}
+
+	out0 := archsimd.BroadcastFloat32x16(0)
+	out1 := archsimd.BroadcastFloat32x16(0)
+	out2 := archsimd.BroadcastFloat32x16(0)
+	out3 := archsimd.BroadcastFloat32x16(0)
+
+	for ki := 0; ki < kvLenUnmasked; ki++ {
+		w := scoresScratch[ki]
+		if w == 0 {
+			continue
+		}
+		vW := archsimd.BroadcastFloat32x16(w)
+		vBase := kvOff + ki*kvSeqStride
+		v0 := archsimd.LoadFloat32x16(v[vBase : vBase+16])
+		v1 := archsimd.LoadFloat32x16(v[vBase+16 : vBase+32])
+		v2 := archsimd.LoadFloat32x16(v[vBase+32 : vBase+48])
+		v3 := archsimd.LoadFloat32x16(v[vBase+48 : vBase+64])
+		out0 = out0.MulAdd(vW, v0)
+		out1 = out1.MulAdd(vW, v1)
+		out2 = out2.MulAdd(vW, v2)
+		out3 = out3.MulAdd(vW, v3)
+	}
+
+	out0.Store(output[outBase : outBase+16])
+	out1.Store(output[outBase+16 : outBase+32])
+	out2.Store(output[outBase+32 : outBase+48])
+	out3.Store(output[outBase+48 : outBase+64])
+}
+
+func sdpaFloat32AVX512Dim64(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) {
+	for gIdx := range groupSize {
+		gQOff := qOff + gIdx*qGroupStride
+		gMaskOff := gIdx * maskGroupStride
+		gBiasOff := gIdx * biasGroupStride
+
+		for qIdx := range seqLen {
+			outBase := gQOff + qIdx*qSeqStride
+			if qIdx >= qLimit {
+				for d := range 64 {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			kvLenUnmasked := kvLen
+			if kvLimit < kvLenUnmasked {
+				kvLenUnmasked = kvLimit
+			}
+			if causal && qIdx+1 < kvLenUnmasked {
+				kvLenUnmasked = qIdx + 1
+			}
+
+			if kvLenUnmasked <= 0 {
+				for d := range 64 {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			qBase := gQOff + qIdx*qSeqStride
+			maskIdxBase := gMaskOff + qIdx*kvLen
+			biasIdxBase := gBiasOff + qIdx*kvLen
+
+			q0 := archsimd.LoadFloat32x16(q[qBase : qBase+16])
+			q1 := archsimd.LoadFloat32x16(q[qBase+16 : qBase+32])
+			q2 := archsimd.LoadFloat32x16(q[qBase+32 : qBase+48])
+			q3 := archsimd.LoadFloat32x16(q[qBase+48 : qBase+64])
+
+			rowMax := float32(math.Inf(-1))
+
+			for ki := range kvLenUnmasked {
+				maskIdx := maskIdxBase + ki
+				if len(booleanMask) > 0 && !booleanMask[maskIdx] {
+					scoresScratch[ki] = float32(math.Inf(-1))
+					continue
+				}
+
+				kBase := kvOff + ki*kvSeqStride
+				k0 := archsimd.LoadFloat32x16(k[kBase : kBase+16])
+				k1 := archsimd.LoadFloat32x16(k[kBase+16 : kBase+32])
+				k2 := archsimd.LoadFloat32x16(k[kBase+32 : kBase+48])
+				k3 := archsimd.LoadFloat32x16(k[kBase+48 : kBase+64])
+
+				sumVec := q0.Mul(k0).Add(q1.Mul(k1)).Add(q2.Mul(k2)).Add(q3.Mul(k3))
+				dot := reduceSum16(sumVec)
+				s := dot * scale
+				if len(additiveBias) > 0 {
+					s += additiveBias[biasIdxBase+ki]
+				}
+				if len(additiveMask) > 0 {
+					s += additiveMask[maskIdx]
+				}
+				scoresScratch[ki] = s
+				if s > rowMax {
+					rowMax = s
+				}
+			}
+
+			softmaxAndValueAccumDim64(v, output, outBase, kvOff, kvSeqStride, scoresScratch, kvLenUnmasked, rowMax)
+		}
+	}
+}
+
+func softmaxAndValueAccumDim128(
+	v, output []float32,
+	outBase, kvOff, kvSeqStride int,
+	scoresScratch []float32,
+	kvLenUnmasked int,
+	rowMax float32,
+) {
+	var sum float32
+	k := 0
+	vRowMax := archsimd.BroadcastFloat32x16(rowMax)
+	vSum := archsimd.BroadcastFloat32x16(0)
+	for ; k+16 <= kvLenUnmasked; k += 16 {
+		vScores := archsimd.LoadFloat32x16(scoresScratch[k : k+16])
+		expV := exp512(vScores.Sub(vRowMax))
+		expV.Store(scoresScratch[k : k+16])
+		vSum = vSum.Add(expV)
+	}
+	if k > 0 {
+		sum += reduceSum16(vSum)
+	}
+	for ; k < kvLenUnmasked; k++ {
+		s := scoresScratch[k]
+		if s == float32(math.Inf(-1)) {
+			scoresScratch[k] = 0
+		} else {
+			e := fastmath.Exp32(s - rowMax)
+			scoresScratch[k] = e
+			sum += e
+		}
+	}
+
+	var invSum float32
+	if sum > 0 {
+		invSum = 1.0 / sum
+	}
+
+	k = 0
+	vInvSum := archsimd.BroadcastFloat32x16(invSum)
+	for ; k+16 <= kvLenUnmasked; k += 16 {
+		vScores := archsimd.LoadFloat32x16(scoresScratch[k : k+16])
+		vScores.Mul(vInvSum).Store(scoresScratch[k : k+16])
+	}
+	for ; k < kvLenUnmasked; k++ {
+		scoresScratch[k] *= invSum
+	}
+
+	out0 := archsimd.BroadcastFloat32x16(0)
+	out1 := archsimd.BroadcastFloat32x16(0)
+	out2 := archsimd.BroadcastFloat32x16(0)
+	out3 := archsimd.BroadcastFloat32x16(0)
+	out4 := archsimd.BroadcastFloat32x16(0)
+	out5 := archsimd.BroadcastFloat32x16(0)
+	out6 := archsimd.BroadcastFloat32x16(0)
+	out7 := archsimd.BroadcastFloat32x16(0)
+
+	for ki := 0; ki < kvLenUnmasked; ki++ {
+		w := scoresScratch[ki]
+		if w == 0 {
+			continue
+		}
+		vW := archsimd.BroadcastFloat32x16(w)
+		vBase := kvOff + ki*kvSeqStride
+		v0 := archsimd.LoadFloat32x16(v[vBase : vBase+16])
+		v1 := archsimd.LoadFloat32x16(v[vBase+16 : vBase+32])
+		v2 := archsimd.LoadFloat32x16(v[vBase+32 : vBase+48])
+		v3 := archsimd.LoadFloat32x16(v[vBase+48 : vBase+64])
+		v4 := archsimd.LoadFloat32x16(v[vBase+64 : vBase+80])
+		v5 := archsimd.LoadFloat32x16(v[vBase+80 : vBase+96])
+		v6 := archsimd.LoadFloat32x16(v[vBase+96 : vBase+112])
+		v7 := archsimd.LoadFloat32x16(v[vBase+112 : vBase+128])
+		out0 = out0.MulAdd(vW, v0)
+		out1 = out1.MulAdd(vW, v1)
+		out2 = out2.MulAdd(vW, v2)
+		out3 = out3.MulAdd(vW, v3)
+		out4 = out4.MulAdd(vW, v4)
+		out5 = out5.MulAdd(vW, v5)
+		out6 = out6.MulAdd(vW, v6)
+		out7 = out7.MulAdd(vW, v7)
+	}
+
+	out0.Store(output[outBase : outBase+16])
+	out1.Store(output[outBase+16 : outBase+32])
+	out2.Store(output[outBase+32 : outBase+48])
+	out3.Store(output[outBase+48 : outBase+64])
+	out4.Store(output[outBase+64 : outBase+80])
+	out5.Store(output[outBase+80 : outBase+96])
+	out6.Store(output[outBase+96 : outBase+112])
+	out7.Store(output[outBase+112 : outBase+128])
+}
+
+func sdpaFloat32AVX512Dim128(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) {
+	for gIdx := range groupSize {
+		gQOff := qOff + gIdx*qGroupStride
+		gMaskOff := gIdx * maskGroupStride
+		gBiasOff := gIdx * biasGroupStride
+
+		for qIdx := range seqLen {
+			outBase := gQOff + qIdx*qSeqStride
+			if qIdx >= qLimit {
+				for d := range 128 {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			kvLenUnmasked := kvLen
+			if kvLimit < kvLenUnmasked {
+				kvLenUnmasked = kvLimit
+			}
+			if causal && qIdx+1 < kvLenUnmasked {
+				kvLenUnmasked = qIdx + 1
+			}
+
+			if kvLenUnmasked <= 0 {
+				for d := range 128 {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			qBase := gQOff + qIdx*qSeqStride
+			maskIdxBase := gMaskOff + qIdx*kvLen
+			biasIdxBase := gBiasOff + qIdx*kvLen
+
+			q0 := archsimd.LoadFloat32x16(q[qBase : qBase+16])
+			q1 := archsimd.LoadFloat32x16(q[qBase+16 : qBase+32])
+			q2 := archsimd.LoadFloat32x16(q[qBase+32 : qBase+48])
+			q3 := archsimd.LoadFloat32x16(q[qBase+48 : qBase+64])
+			q4 := archsimd.LoadFloat32x16(q[qBase+64 : qBase+80])
+			q5 := archsimd.LoadFloat32x16(q[qBase+80 : qBase+96])
+			q6 := archsimd.LoadFloat32x16(q[qBase+96 : qBase+112])
+			q7 := archsimd.LoadFloat32x16(q[qBase+112 : qBase+128])
+
+			rowMax := float32(math.Inf(-1))
+
+			for ki := range kvLenUnmasked {
+				maskIdx := maskIdxBase + ki
+				if len(booleanMask) > 0 && !booleanMask[maskIdx] {
+					scoresScratch[ki] = float32(math.Inf(-1))
+					continue
+				}
+
+				kBase := kvOff + ki*kvSeqStride
+				k0 := archsimd.LoadFloat32x16(k[kBase : kBase+16])
+				k1 := archsimd.LoadFloat32x16(k[kBase+16 : kBase+32])
+				k2 := archsimd.LoadFloat32x16(k[kBase+32 : kBase+48])
+				k3 := archsimd.LoadFloat32x16(k[kBase+48 : kBase+64])
+				k4 := archsimd.LoadFloat32x16(k[kBase+64 : kBase+80])
+				k5 := archsimd.LoadFloat32x16(k[kBase+80 : kBase+96])
+				k6 := archsimd.LoadFloat32x16(k[kBase+96 : kBase+112])
+				k7 := archsimd.LoadFloat32x16(k[kBase+112 : kBase+128])
+
+				sumVec0 := q0.Mul(k0).Add(q1.Mul(k1)).Add(q2.Mul(k2)).Add(q3.Mul(k3))
+				sumVec1 := q4.Mul(k4).Add(q5.Mul(k5)).Add(q6.Mul(k6)).Add(q7.Mul(k7))
+				dot := reduceSum16(sumVec0.Add(sumVec1))
+				s := dot * scale
+				if len(additiveBias) > 0 {
+					s += additiveBias[biasIdxBase+ki]
+				}
+				if len(additiveMask) > 0 {
+					s += additiveMask[maskIdx]
+				}
+				scoresScratch[ki] = s
+				if s > rowMax {
+					rowMax = s
+				}
+			}
+
+			softmaxAndValueAccumDim128(v, output, outBase, kvOff, kvSeqStride, scoresScratch, kvLenUnmasked, rowMax)
+		}
+	}
+}
+
+func sdpaFloat32AVX512General(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen, headDim int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) {
+	for gIdx := range groupSize {
+		gQOff := qOff + gIdx*qGroupStride
+		gMaskOff := gIdx * maskGroupStride
+		gBiasOff := gIdx * biasGroupStride
+
+		for qIdx := range seqLen {
+			outBase := gQOff + qIdx*qSeqStride
+			if qIdx >= qLimit {
+				for d := range headDim {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			kvLenUnmasked := kvLen
+			if kvLimit < kvLenUnmasked {
+				kvLenUnmasked = kvLimit
+			}
+			if causal && qIdx+1 < kvLenUnmasked {
+				kvLenUnmasked = qIdx + 1
+			}
+
+			if kvLenUnmasked <= 0 {
+				for d := range headDim {
+					output[outBase+d] = 0
+				}
+				continue
+			}
+
+			qBase := gQOff + qIdx*qSeqStride
+			maskIdxBase := gMaskOff + qIdx*kvLen
+			biasIdxBase := gBiasOff + qIdx*kvLen
+
+			rowMax := float32(math.Inf(-1))
+
+			for ki := range kvLenUnmasked {
+				maskIdx := maskIdxBase + ki
+				if len(booleanMask) > 0 && !booleanMask[maskIdx] {
+					scoresScratch[ki] = float32(math.Inf(-1))
+					continue
+				}
+
+				kBase := kvOff + ki*kvSeqStride
+				d := 0
+				sumVec := archsimd.BroadcastFloat32x16(0)
+				for ; d+16 <= headDim; d += 16 {
+					qV := archsimd.LoadFloat32x16(q[qBase+d : qBase+d+16])
+					kV := archsimd.LoadFloat32x16(k[kBase+d : kBase+d+16])
+					sumVec = sumVec.Add(qV.Mul(kV))
+				}
+				var dot float32
+				if d > 0 {
+					dot = reduceSum16(sumVec)
+				}
+				for ; d < headDim; d++ {
+					dot += q[qBase+d] * k[kBase+d]
+				}
+
+				s := dot * scale
+				if len(additiveBias) > 0 {
+					s += additiveBias[biasIdxBase+ki]
+				}
+				if len(additiveMask) > 0 {
+					s += additiveMask[maskIdx]
+				}
+				scoresScratch[ki] = s
+				if s > rowMax {
+					rowMax = s
+				}
+			}
+
+			// Softmax
+			var sum float32
+			ki := 0
+			vRowMax := archsimd.BroadcastFloat32x16(rowMax)
+			vSum := archsimd.BroadcastFloat32x16(0)
+			for ; ki+16 <= kvLenUnmasked; ki += 16 {
+				vScores := archsimd.LoadFloat32x16(scoresScratch[ki : ki+16])
+				expV := exp512(vScores.Sub(vRowMax))
+				expV.Store(scoresScratch[ki : ki+16])
+				vSum = vSum.Add(expV)
+			}
+			if ki > 0 {
+				sum += reduceSum16(vSum)
+			}
+			for ; ki < kvLenUnmasked; ki++ {
+				s := scoresScratch[ki]
+				if s == float32(math.Inf(-1)) {
+					scoresScratch[ki] = 0
+				} else {
+					e := fastmath.Exp32(s - rowMax)
+					scoresScratch[ki] = e
+					sum += e
+				}
+			}
+
+			var invSum float32
+			if sum > 0 {
+				invSum = 1.0 / sum
+			}
+
+			ki = 0
+			vInvSum := archsimd.BroadcastFloat32x16(invSum)
+			for ; ki+16 <= kvLenUnmasked; ki += 16 {
+				vScores := archsimd.LoadFloat32x16(scoresScratch[ki : ki+16])
+				vScores.Mul(vInvSum).Store(scoresScratch[ki : ki+16])
+			}
+			for ; ki < kvLenUnmasked; ki++ {
+				scoresScratch[ki] *= invSum
+			}
+
+			// Value accumulation along d
+			d := 0
+			for ; d+16 <= headDim; d += 16 {
+				outV := archsimd.BroadcastFloat32x16(0)
+				for kIdx := range kvLenUnmasked {
+					w := scoresScratch[kIdx]
+					if w == 0 {
+						continue
+					}
+					vW := archsimd.BroadcastFloat32x16(w)
+					vBase := kvOff + kIdx*kvSeqStride + d
+					vVal := archsimd.LoadFloat32x16(v[vBase : vBase+16])
+					outV = outV.MulAdd(vW, vVal)
+				}
+				outV.Store(output[outBase+d : outBase+d+16])
+			}
+			for ; d < headDim; d++ {
+				var acc float32
+				for kIdx := range kvLenUnmasked {
+					w := scoresScratch[kIdx]
+					if w != 0 {
+						acc += w * v[kvOff+kIdx*kvSeqStride+d]
+					}
+				}
+				output[outBase+d] = acc
+			}
+		}
+	}
+}
diff --git a/internal/gobackend/fusedops/sdpa.go b/internal/gobackend/fusedops/sdpa.go
index 34155fc..dbb9f01 100644
--- a/internal/gobackend/fusedops/sdpa.go
+++ b/internal/gobackend/fusedops/sdpa.go
@@ -2,6 +2,8 @@ package fusedops
 
 import (
 	"math"
+	"sync"
+	"sync/atomic"
 
 	"github.com/gomlx/compute"
 	"github.com/gomlx/compute/dtypes"
@@ -51,13 +53,8 @@ func FusedScaledDotProductAttention(
 }
 
 func init() {
-	// DISABLED: the new matmul with SIMD support is much faster (+3x faster), so thi fused op ends up being slower,
-	// at least in a small sentence embedding model.
-	// TODO: add a SIMD version and re-evaluate.
-	if false {
-		gobackend.RegisterFusedScaledDotProductAttention.Register(FusedScaledDotProductAttention, gobackend.PriorityGeneric)
-		gobackend.SetNodeExecutor(compute.OpTypeFusedScaledDotProductAttention, gobackend.PriorityTyped, execFusedScaledDotProductAttention)
-	}
+	gobackend.RegisterFusedScaledDotProductAttention.Register(FusedScaledDotProductAttention, gobackend.PriorityGeneric)
+	gobackend.SetNodeExecutor(compute.OpTypeFusedScaledDotProductAttention, gobackend.PriorityTyped, execFusedScaledDotProductAttention)
 }
 
 type nodeScaledDotProductAttention struct {
@@ -135,9 +132,11 @@ func buildSDPANode(
 		return nil, errors.Errorf("%s: key must have rank 4, got %d", opName, kNode.Shape.Rank())
 	}
 	switch qNode.Shape.DType {
-	case dtypes.F8E4M3FN, dtypes.F8E5M2:
+	case dtypes.Float32, dtypes.Float64:
+		// Supported in go backend.
+	default:
 		return nil, errors.Wrapf(compute.ErrNotImplemented,
-			"%s: float8 input dtype %s is not implemented in the go backend", opName, qNode.Shape.DType)
+			"%s: dtype %s is not implemented in the go backend", opName, qNode.Shape.DType)
 	}
 
 	numHeads := qNode.Shape.Dimensions[axesLayout.HeadsAxis()]
@@ -290,11 +289,11 @@ func execFusedScaledDotProductAttention(backend *gobackend.Backend, node *goback
 
 	switch query.RawShape.DType {
 	case dtypes.Float32:
-		sdpaMultiHeadGeneric[float32](query, key, value, mask, bias, output, data, maskBatchStride, maskHeadStride, biasBatchStride, biasHeadStride, querySeqLen, keyValueSeqLen)
+		sdpaMultiHeadGeneric[float32](backend, query, key, value, mask, bias, output, data, maskBatchStride, maskHeadStride, biasBatchStride, biasHeadStride, querySeqLen, keyValueSeqLen)
 	case dtypes.Float64:
-		sdpaMultiHeadGeneric[float64](query, key, value, mask, bias, output, data, maskBatchStride, maskHeadStride, biasBatchStride, biasHeadStride, querySeqLen, keyValueSeqLen)
+		sdpaMultiHeadGeneric[float64](backend, 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)
+		return nil, errors.Wrapf(compute.ErrNotImplemented, "FusedScaledDotProductAttention: unsupported dtype %s", query.RawShape.DType)
 	}
 
 	return output, nil
@@ -518,7 +517,13 @@ func sdpaGeneric[T float32 | float64](
 	}
 }
 
-func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, bias, output *gobackend.Buffer, data *nodeScaledDotProductAttention, maskBatchStride, maskHeadStride, biasBatchStride, biasHeadStride int, querySeqLen, keyValueSeqLen []int32) {
+func sdpaMultiHeadGeneric[T float32 | float64](
+	backend *gobackend.Backend,
+	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)
@@ -577,9 +582,10 @@ func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, bias, ou
 		kvBatchStride = numKVHeads * kvLen * headDim
 	}
 
-	scores := make([]T, groupSize*seqLen*kvLen)
 	maskSliceLen := seqLen * kvLen
-	for batchIdx := range batchSize {
+	archFn := gobackend.GetSDPAArchDispatcher()
+
+	processHead := func(batchIdx, kvHeadIdx int, scores []T) {
 		qLimit := seqLen
 		kvLimit := kvLen
 		if len(querySeqLen) > 0 {
@@ -588,49 +594,111 @@ func sdpaMultiHeadGeneric[T float32 | float64](query, key, value, mask, bias, ou
 		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
-
-			// Compute mask slice and group stride for this KV head group.
-			var additiveMaskSlice []T
-			var booleanMaskSlice []bool
-			maskGroupStride := 0
-			if len(additiveMask) > 0 || len(booleanMask) > 0 {
-				maskOffset := batchIdx*maskBatchStride + kvHeadIdx*groupSize*maskHeadStride
-				maskEnd := maskOffset + maskSliceLen
-				if maskHeadStride > 0 && groupSize > 1 {
-					maskEnd = maskOffset + (groupSize-1)*maskHeadStride + maskSliceLen
-					maskGroupStride = maskHeadStride
+		qOff := batchIdx*qBatchStride + kvHeadIdx*groupSize*qHeadStride
+		kvOff := batchIdx*kvBatchStride + kvHeadIdx*kvHeadStride
+
+		// Compute mask slice and group stride for this KV head group.
+		var additiveMaskSlice []T
+		var booleanMaskSlice []bool
+		maskGroupStride := 0
+		if len(additiveMask) > 0 || len(booleanMask) > 0 {
+			maskOffset := batchIdx*maskBatchStride + kvHeadIdx*groupSize*maskHeadStride
+			maskEnd := maskOffset + maskSliceLen
+			if maskHeadStride > 0 && groupSize > 1 {
+				maskEnd = maskOffset + (groupSize-1)*maskHeadStride + maskSliceLen
+				maskGroupStride = maskHeadStride
+			}
+			if len(additiveMask) > 0 {
+				additiveMaskSlice = additiveMask[maskOffset:maskEnd]
+			} else {
+				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]
+		}
+
+		if archFn != nil {
+			if qF32, ok := any(q).([]float32); ok {
+				kF32 := any(k).([]float32)
+				vF32 := any(v).([]float32)
+				outF32 := any(out).([]float32)
+				var addMaskF32 []float32
+				if len(additiveMaskSlice) > 0 {
+					addMaskF32 = any(additiveMaskSlice).([]float32)
 				}
-				if len(additiveMask) > 0 {
-					additiveMaskSlice = additiveMask[maskOffset:maskEnd]
-				} else {
-					booleanMaskSlice = booleanMask[maskOffset:maskEnd]
+				var addBiasF32 []float32
+				if len(additiveBiasSlice) > 0 {
+					addBiasF32 = any(additiveBiasSlice).([]float32)
+				}
+				scratchF32 := any(scores).([]float32)
+				if archFn(
+					qF32, kF32, vF32, outF32,
+					qOff, kvOff, qSeqStride, kvSeqStride, qHeadStride,
+					addMaskF32, booleanMaskSlice, maskGroupStride,
+					addBiasF32, biasGroupStride,
+					scratchF32,
+					groupSize, seqLen, kvLen, headDim,
+					float32(scale), causal,
+					qLimit, kvLimit,
+				) {
+					return
 				}
 			}
-			// 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
+		}
+
+		sdpaGeneric(
+			q, k, v, qOff, kvOff, qSeqStride, kvSeqStride, qHeadStride,
+			additiveMaskSlice, booleanMaskSlice, maskGroupStride,
+			additiveBiasSlice, biasGroupStride,
+			scores,
+			out,
+			groupSize, seqLen, kvLen, headDim, scale, causal,
+			qLimit, kvLimit,
+		)
+	}
+
+	totalTasks := batchSize * numKVHeads
+	if backend != nil && backend.Workers != nil && backend.Workers.IsEnabled() && totalTasks > 1 {
+		numWorkers := backend.Workers.AdjustedMaxParallelism()
+		workersToStart := min(totalTasks, numWorkers)
+		var taskCounter atomic.Int64
+		var wg sync.WaitGroup
+		scratchSize := max(groupSize*seqLen*kvLen, kvLen)
+
+		for range workersToStart {
+			wg.Add(1)
+			backend.Workers.WaitToStart(func() {
+				defer wg.Done()
+				workerScores := make([]T, scratchSize)
+				for {
+					taskIdx := int(taskCounter.Add(1) - 1)
+					if taskIdx >= totalTasks {
+						break
+					}
+					batchIdx := taskIdx / numKVHeads
+					kvHeadIdx := taskIdx % numKVHeads
+					processHead(batchIdx, kvHeadIdx, workerScores)
 				}
-				additiveBiasSlice = additiveBias[biasOffset:biasEnd]
+			})
+		}
+		wg.Wait()
+	} else {
+		scores := make([]T, max(groupSize*seqLen*kvLen, kvLen))
+		for batchIdx := range batchSize {
+			for kvHeadIdx := range numKVHeads {
+				processHead(batchIdx, kvHeadIdx, scores)
 			}
-			sdpaGeneric(
-				q, k, v, qOff, kvOff, qSeqStride, kvSeqStride, qHeadStride,
-				additiveMaskSlice, booleanMaskSlice, maskGroupStride,
-				additiveBiasSlice, biasGroupStride,
-				scores,
-				out,
-				groupSize, seqLen, kvLen, headDim, scale, causal,
-				qLimit, kvLimit,
-			)
 		}
 	}
 }
diff --git a/internal/gobackend/sdpa_arch.go b/internal/gobackend/sdpa_arch.go
new file mode 100644
index 0000000..65a0917
--- /dev/null
+++ b/internal/gobackend/sdpa_arch.go
@@ -0,0 +1,39 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+package gobackend
+
+// SDPAFloat32ArchFn is a function hook to execute Scaled Dot-Product Attention for float32
+// using architecture-specific kernels (e.g. AVX2, AVX-512).
+// It processes one KV-head group (with groupSize query heads) with zero transpositions.
+type SDPAFloat32ArchFn func(
+	q, k, v, output []float32,
+	qOff, kvOff, qSeqStride, kvSeqStride, qGroupStride int,
+	additiveMask []float32,
+	booleanMask []bool,
+	maskGroupStride int,
+	additiveBias []float32,
+	biasGroupStride int,
+	scoresScratch []float32,
+	groupSize, seqLen, kvLen, headDim int,
+	scale float32, causal bool,
+	qLimit, kvLimit int,
+) bool
+
+var (
+	sdpaFloat32ArchFn       SDPAFloat32ArchFn
+	sdpaFloat32ArchPriority RegisterPriority
+)
+
+// SetSDPAArchDispatcher registers the architecture-specific SDPA dispatcher with a priority.
+// Higher priority replaces lower priority dispatcher.
+func SetSDPAArchDispatcher(priority RegisterPriority, fn SDPAFloat32ArchFn) {
+	if priority >= sdpaFloat32ArchPriority {
+		sdpaFloat32ArchPriority = priority
+		sdpaFloat32ArchFn = fn
+	}
+}
+
+// GetSDPAArchDispatcher returns the registered architecture-specific SDPA dispatcher, if any.
+func GetSDPAArchDispatcher() SDPAFloat32ArchFn {
+	return sdpaFloat32ArchFn
+}

From dffb83f383191b2a2b892519b3a1a3d8282a6065 Mon Sep 17 00:00:00 2001
From: Jan Pfeifer <pfeifer@gmail.com>
Date: Thu, 10 Sep 2026 16:38:58 +0200
Subject: [PATCH 12/12] Update notes on optimization of the FusedSPDA.

---
 .agents/AGENTS.md            | 53 ++++++++++++++++++++++++++++++++++++
 internal/gobackend/README.md | 52 +++++++++++++++++++++++++++++++++--
 2 files changed, 103 insertions(+), 2 deletions(-)

diff --git a/.agents/AGENTS.md b/.agents/AGENTS.md
index 13986fc..8871c07 100644
--- a/.agents/AGENTS.md
+++ b/.agents/AGENTS.md
@@ -385,3 +385,56 @@ Handwritten assembly functions (in `*_amd64.s`) bypass Go compiler checks and wi
 - Always gate unit tests in `avx2/` and `avx512/` subpackages by calling `if !gobackend.IsAVX2Allowed { t.Skip(...) }` or `if !gobackend.IsAVX512Allowed { t.Skip(...) }` at the top of each test function.
 - Never write tests that execute raw assembly kernels without this gate.
 
+## Guidelines for Fused Kernels and Intra-Op Parallelization
+
+Lessons learned from optimizing high-performance fused operations (such as `FusedScaledDotProductAttention`, `FusedLayerNorm`, and `FusedDense`):
+
+### 1. Zero-Transpose Strided Memory Access
+Complex multi-dimensional operations frequently encounter multiple tensor layout conventions across frameworks (e.g. `[B, S, H, D]` vs `[B, H, S, D]` in attention):
+- **The Pitfall**: Decomposing operations or permuting tensors using explicit transposition (`Transpose`) introduces massive memory copying overhead. In transformer benchmarks, 4D transpositions accounted for **19.8% of total CPU time** and dominated `runtime.memmove` (26.7%).
+- **The Principle**: In row-major tensors, the inner head dimension $D$ remains contiguous (unit stride) regardless of whether heads $H$ or sequence $S$ comes first.
+- **The Implementation**:
+  - Parametrize the inner kernel with sequence strides (`qSeqStride`, `kvSeqStride`) and head/group strides (`qGroupStride`).
+  - Read operands and write outputs directly using strided indexing over the native slices without creating intermediate tensors.
+  - Completely eliminates intermediate transposition passes and tensor allocations, cutting memory movement by >70%.
+
+### 2. Cache-Aware Task Decomposition (KV-Head Grouping)
+In multi-head attention with Grouped Query Attention (GQA) or Multi-Query Attention (MQA), $N_{query}$ query heads share $N_{kv}$ key/value heads ($N_{query} \ge N_{kv}$):
+- **Avoid Independent Query Head Parallelization**: If tasks are split per query head, multiple worker threads concurrently compete for memory bandwidth fetching identical $K$ and $V$ tensors from L3 cache or DRAM.
+- **Group by KV-Head**: Partition the workload into $N_{tasks} = Batch \times N_{kv}$ chunks.
+- Each worker processes an entire $(batch, kv\_head)$ unit, iterating through all $G = N_{query} / N_{kv}$ query heads sequentially. This guarantees that the key and value sequences ($S_{kv} \times D$) are loaded once and stay resident in the core's private L1/L2 cache across all $G$ query heads.
+
+### 3. Dynamic Work-Stealing with Atomic Counters
+- **Channel vs Atomic Counter**: Avoid distributing tasks via Go channels. In tight compute loops, channel send/receive operations introduce channel lock contention, context switching, and runtime scheduler latency.
+- **Static Slicing Pitfall**: Evenly dividing $N_{tasks}$ across $W$ workers upfront leads to thread imbalance if sequence lengths vary (due to padding) or if task counts do not divide evenly by worker count.
+- **Optimal Pattern**:
+  1. Spawn $W = \min(N_{tasks}, \text{backend.Workers})$ goroutines using `backend.Workers.Saturate` or the worker pool.
+  2. Use an `atomic.Int64` task counter shared across workers:
+     ```go
+     for {
+         taskIdx := int(taskCounter.Add(1) - 1)
+         if taskIdx >= numTasks {
+             break
+         }
+         // execute taskIdx
+     }
+     ```
+  3. This provides lock-free, zero-allocation dynamic load balancing where faster cores naturally process more tasks.
+
+### 4. In-Cache Reductions & Register Accumulation
+- **Never Materialize Intermediate Full Matrices**: Materializing $S_q \times S_{kv}$ attention logit matrices in DRAM creates severe memory-bandwidth bottlenecks for sequence lengths $S \ge 512$.
+- **Per-Worker Scratch Buffers**: Allocate a small scratch slice ($S_{kv}$ floats) per worker (or reuse a buffer across rows).
+- **Fused Math in Vector Registers**:
+  - Immediately fold scale factors, additive masks, and attention bias into vector registers during the dot product ($Q_q \cdot K_k$).
+  - For causal masking, immediately short-circuit keys where $k > q$.
+  - Evaluate softmax normalizations using vectorized polynomial approximations (e.g. degree-7 Cephes `exp512` / `exp256`).
+  - Accumulate the weighted value vectors ($P \cdot V$) directly into vector accumulators and store directly into the destination buffer.
+
+### 5. Head Dimension Specialization ($D \in \{32, 64, 128\}$)
+- When the inner reduction loop over dimension $D$ has a dynamic or variable upper bound, compilers and SIMD abstractions cannot unroll completely and must maintain loop counters and boundary checks.
+- Generating specialized fast paths for standard power-of-two head dimensions ($D=32, 64, 128$):
+  - On AVX-512 (`Float32x16`): $D=32$ unrolls into exactly 2 vector registers, $D=64$ into 4 vectors, $D=128$ into 8 vectors.
+  - On AVX2 (`Float32x8`): $D=32$ unrolls into 4 vector registers, $D=64$ into 8 vectors, $D=128$ into 16 vectors.
+  - Completely unrolling the inner dimension eliminates loop overhead and allows optimal compiler/assembler instruction scheduling across FMA execution ports.
+
+
diff --git a/internal/gobackend/README.md b/internal/gobackend/README.md
index edfd426..f4d2c14 100644
--- a/internal/gobackend/README.md
+++ b/internal/gobackend/README.md
@@ -30,8 +30,8 @@ But there are many relatively "low-hanging fruits" for optimization, a few obvio
   loop over the data many times, each time applying the unary function.
 * Fuse binary/unary ops: perform unary functions while traversing the data for binary functions. Again to save
   memory accesses.
-* Further in-operation parallelization: only DotGeneral has been parallelized so far: it is usually the one that consumes most of the time.
-* Use intrinsics/SIMD on platforms that allow it. It was announced as experimental in Go 1.25.
+* In-operation parallelization: `DotGeneral` and `FusedScaledDotProductAttention` are parallelized across worker pools using dynamic lock-free work-stealing.
+* Use intrinsics/SIMD on platforms that allow it: AVX2 and AVX-512 SIMD kernels are implemented for `DotGeneral` (GEMM), `FusedScaledDotProductAttention`, `Where`, `FusedLayerNorm`, and activations (`Gelu`).
 * ~~Eliminate common sub-expressions.~~
 
 ## SIMD Thresholding
@@ -123,3 +123,51 @@ In tight loops or row-by-row reductions (such as LayerNorm, RMSNorm, or Softmax)
    Standard Go compiler code generation uses legacy SSE instructions for scalar float math. If an assembly routine leaves the upper halves of YMM/ZMM registers in a "dirty" state upon returning, the very next Go-compiled SSE instruction will trigger an AVX $\to$ SSE penalty.
    Always emit `VZEROUPPER` immediately before every `RET` instruction in functions using 256-bit or 512-bit registers.
 
+## Fused Operations & Intra-Op Parallelization Architecture
+
+When implementing compute-intensive fused operations (such as `FusedScaledDotProductAttention`), the Go backend employs several core design principles to maximize hardware efficiency on modern multi-core x86_64 processors:
+
+### 1. Dynamic Work-Stealing with Atomic Counters
+
+Rather than using Go channels or statically partitioning work across goroutines:
+- **Avoid Go Channels in Inner Loops**: Distributing work units through Go channels introduces channel mutex lock contention, channel buffer overhead, and goroutine parking/unparking latency.
+- **Avoid Static Chunking**: Dividing tasks evenly across $W$ workers upfront leads to thread imbalance when individual tasks have variable workloads (e.g. varying sequence lengths due to padding) or when the number of tasks does not divide evenly by worker count.
+- **Atomic Work-Stealing**:
+  1. Workers are spawned up to $\min(N_{\text{tasks}}, \text{backend.Workers})$.
+  2. Workers dynamically steal the next available task index using an atomic counter:
+     ```go
+     for {
+         taskIdx := int(taskCounter.Add(1) - 1)
+         if taskIdx >= numTasks {
+             break
+         }
+         // Process taskIdx
+     }
+     ```
+  3. This provides zero-allocation, lock-free dynamic load balancing where faster cores naturally process more tasks.
+
+### 2. Cache-Aware Task Decomposition (KV-Head Grouping in GQA/MQA)
+
+In Grouped Query Attention (GQA) and Multi-Query Attention (MQA), multiple query heads share a single key/value (KV) head ($N_{\text{query}} \ge N_{\text{kv}}$):
+- If tasks are decomposed per individual query head, multiple worker goroutines concurrently compete for memory bandwidth fetching identical $K$ and $V$ tensors from L3 cache or DRAM.
+- By defining each parallel task as a $(batch, kv\_head)$ chunk, a single worker iterates over all $G = N_{\text{query}} / N_{\text{kv}}$ query heads sequentially.
+- This guarantees that the key and value sequences ($S_{\text{kv}} \times D$) are loaded once and remain resident in the core's private L1/L2 cache while computing all $G$ query heads.
+
+### 3. Zero-Transpose Direct Strided Memory Access
+
+Deep learning models frequently switch between layout conventions (e.g. `LayoutBSHD` $[B, S, H, D]$ and `LayoutBHSD` $[B, H, S, D]$):
+- Decomposing attention into explicit 4D tensor permutations (`Transpose`) incurs severe memory copying penalties (accounting for ~20% of CPU time and dominating `runtime.memmove`).
+- In row-major tensors, the inner head dimension $D$ is contiguous (stride 1) in both layouts.
+- By parametrizing inner loops with sequence strides (`qSeqStride`, `kvSeqStride`) and query group strides (`qGroupStride`), compute kernels can read operands and write outputs directly in their native layouts with zero tensor copying.
+
+### 4. In-Cache Softmax & Register-Accumulated Values
+
+- **Avoid Materializing $S \times S$ Matrices in Memory**: Writing intermediate attention logits or probabilities to DRAM creates severe bandwidth bottlenecks.
+- **Per-Worker Scratch Buffers**: Each worker allocates a small temporary slice ($S_{\text{kv}}$ floats) reused across tokens.
+- **In-Register Fused Math**: Scale factors, causal masks, additive masks, and attention biases are folded directly into vector registers during dot-product accumulation. Softmax exponentiation uses fast vectorized polynomial approximations (`exp512` / `exp256`), and weighted values ($P \cdot V$) accumulate directly into vector registers before writing directly to the final destination buffer.
+
+### 5. Head Dimension Specialization ($D \in \{32, 64, 128\}$)
+
+When the inner loop over dimension $D$ has a dynamic upper bound, the compiler cannot unroll vector loops and must emit loop counter branches. Specializing kernels for standard head dimensions allows completely unrolling into a fixed set of vector registers (e.g. 2 `Float32x16` registers for $D=32$ on AVX-512, 4 for $D=64$, 8 for $D=128$), sustaining near-peak FMA throughput.
+
+