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

@@ -137,8 +137,13 @@ impl Tensor {
if self.is_contiguous() {
return self.clone();
}
// Copy to contiguous layout on CPU
assert_eq!(self.device(), Device::Cpu, "contiguous() on GPU not yet supported");
// For GPU tensors: round-trip through CPU (correct but slow).
// TODO: write a GPU contiguous-copy kernel for performance.
if matches!(self.device(), Device::Cuda(_)) {
let cpu = self.to_device(Device::Cpu);
let contig = cpu.contiguous();
return contig.to_device(self.device());
}
let numel = self.numel();
let elem_size = self.dtype.size_bytes();
let src_bytes = self.storage.as_cpu_bytes();
@@ -173,17 +178,18 @@ impl Tensor {
// --- Device transfer ---
pub fn to_device(&self, device: Device) -> Self {
let t = if self.is_contiguous() { self.clone() } else { self.contiguous() };
if t.device() == device {
return t;
if self.device() == device {
return self.clone();
}
let new_storage = t.storage.to_device(device).expect("device transfer failed");
// Transfer the raw storage (preserving strides/offset).
// Non-contiguous layout is preserved — the user can call contiguous() after.
let new_storage = self.storage.to_device(device).expect("device transfer failed");
Self {
storage: new_storage,
shape: t.shape,
strides: t.strides,
offset: 0,
dtype: t.dtype,
shape: self.shape.clone(),
strides: self.strides.clone(),
offset: self.offset,
dtype: self.dtype,
}
}