phase 5: naive multi-head attention

- Batched GEMM via cublasGemmStridedBatchedEx
- Causal mask CUDA kernel (F32 + BF16)
- Element-wise scale CUDA kernel (F32 + BF16)
- attention() composing: batched_matmul + scale + causal_mask + softmax
- Fixed to_device/contiguous infinite recursion (GPU contiguous via CPU round-trip)
- 5 attention tests passing (max_err < 3e-7 F32)
- Total: 61 tests passing across all crates

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-21 21:17:23 +08:00
parent c8e8153702
commit 6035ffdc0b
10 changed files with 550 additions and 12 deletions

View File

@@ -1,4 +1,5 @@
pub mod activation;
pub mod attention;
pub mod embedding;
pub mod gemm;
pub mod layernorm;
@@ -6,9 +7,10 @@ pub mod rmsnorm;
pub mod rope;
pub mod softmax;
pub use activation::{gelu, silu};
pub use activation::{gelu, scale, silu};
pub use attention::attention;
pub use embedding::embedding;
pub use gemm::{matmul, GemmBackend};
pub use gemm::{batched_matmul, matmul, GemmBackend};
pub use layernorm::layernorm;
pub use rmsnorm::rmsnorm;
pub use rope::{rope_inplace, RopeCache};