gobackend: SIMD FusedSDPA, Where implicit broadcast & AVX, and attention/matmul optimizations

.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,59 @@ 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.
 
+## 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.
+
+

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.
+
+

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)

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)

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
 )
 

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,