transformers: compute CPU FlashSdpa P·V as a tile GEMM (~3.1×)

[czoli1976]

What

The CPU FlashSdpaOp forward pass (transformers/src/ops/flash_sdpa.rs) computed the P·V step as head_dim separate dot products against strided columns of V:

for j in 0..head_dim {
    let sv = s.row(i).dot(&vblock.column(j));
    ...
}

vblock.column(j) is a strided view (stride = head_dim), so this is neither a GEMM nor vectorization-friendly. This replaces it with a single contiguous tile GEMM plus the online-softmax rescale:

let sv_tile = s.dot(&vblock);
for i in 0..q_range.len() {
    ...
    for j in 0..head_dim {
        oblock[(i, j)] = oblock[(i, j)] * mul_o + sv_tile[(i, j)] * mul_sv;
    }
}

The QK^T step above already used a 2-D .dot(); this just brings P·V to parity.

Why

• Correctness: bit-exact — max_abs = 0.0 vs a naive softmax(QKᵀ·scale)·V reference.
• Performance: ~3.1× on the whole flash_attention_gqa — H=8, Sq=Sk=512, D=64, release build: 76.2 → 24.6 ms/run. The strided P·V was ~68% of the runtime.

Tests

• flash_sdpa_pv_matches_naive — correctness vs naive reference (non-square, multi-head).
• bench_flash_sdpa_pv (#[ignore]) — the before/after A/B used for the numbers above.
• Full tract-transformers suite green; fmt + clippy clean.

🤖 Generated with Claude Code