huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 02:38:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Adding bfloat16 support for the cast kernels. (#1520)

Browse Source
This commit is contained in:
Nicolas Patry
2024-01-04 12:12:56 +01:00
committed by GitHub
parent 135ae5f3eb
commit fa3ea98ba9
2 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
candle-core/src/metal_backend.rs
View File

@ -596,6 +596,8 @@ impl BackendStorage for MetalStorage {
(DType::F32, DType::F16) => "cast_f32_f16",
(DType::F16, DType::F32) => "cast_f16_f32",
(DType::I64, DType::F32) => "cast_i64_f32",
(DType::F32, DType::BF16) => "cast_f32_bf16",
(DType::BF16, DType::F32) => "cast_bf16_f32",
(left, right) => {
crate::bail!("Metal contiguous to_dtype {left:?} {right:?} not implemented")
}
@ -622,6 +624,8 @@ impl BackendStorage for MetalStorage {
(DType::F32, DType::F16) => "cast_f32_f16_strided",
(DType::F16, DType::F32) => "cast_f16_f32_strided",
(DType::I64, DType::F32) => "cast_i64_f32_strided",
(DType::F32, DType::BF16) => "cast_f32_bf16_strided",
(DType::BF16, DType::F32) => "cast_bf16_f32_strided",
(left, right) => {
crate::bail!("Metal strided to_dtype {left:?} {right:?} not implemented")
}

2
candle-metal-kernels/src/cast.metal
View File

@ -59,4 +59,6 @@ CAST(cast_i64_f32, cast_i64_f32_strided, int64_t, float)
#endif
#if __METAL_VERSION__ >= 310
CAST(cast_bf16_f32, cast_bf16_f32_strided, bfloat, float)
CAST(cast_f32_bf16, cast_f32_bf16_strided, float, bfloat)
#endif