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:
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user