Skip to content

Commit bca3912

Browse files
committed
Revert "perf(cuda): use scalar tensor for attention neg_large instead of full allocation"
This reverts commit 0077696a57b9cf6723aa38cf53d4b933de7910c7.
1 parent 3d2073e commit bca3912

1 file changed

Lines changed: 2 additions & 2 deletions

File tree

  • cake-core/src/backends/cuda

cake-core/src/backends/cuda/mod.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,8 @@ impl ComputeBackend for CudaBackend {
6363
.reshape((1, 1, 1, kv_len))?;
6464
let mask = (rows + offset as f64)?.broadcast_ge(&cols)?;
6565
let mask = mask.broadcast_as(attn.shape())?;
66-
let neg_large = Tensor::new(-1e9f32, q.device())?;
67-
mask.where_cond(&attn, &neg_large.broadcast_as(attn.shape())?)?
66+
let neg_large = Tensor::full(-1e9f32, attn.shape(), q.device())?;
67+
mask.where_cond(&attn, &neg_large)?
6868
}
6969
} else {
7070
attn

0 commit comments

Comments
 (0)