Skip to content

Commit 5db2349

Browse files
zackangelozackangelo
committed
use float8 crate with cudarc 16.x, fix errors
1 parent 9a5cb24 commit 5db2349

10 files changed

Lines changed: 49 additions & 23 deletions

File tree

‎Cargo.toml‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,10 @@ half = { version = "2.5.0", features = [
6262
"use-intrinsics",
6363
"rand_distr",
6464
] }
65-
float8 = { version = "0.1.3", features = ["num-traits", "rand_distr"] }
65+
float8 = { git = "https://github.com/zackangelo/float8", branch = "cudarc_0_16", features = [
66+
"num-traits",
67+
"rand_distr",
68+
] }
6669
hound = "3.5.1"
6770
image = { version = "0.25.2", default-features = false, features = [
6871
"jpeg",

‎candle-core/src/cpu_backend/mod.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2727,7 +2727,8 @@ impl BackendDevice for CpuDevice {
27272727
DType::F8E4M3 => {
27282728
let mut data = Vec::with_capacity(elem_count);
27292729
let uniform =
2730-
rand::distributions::Uniform::new(F8E4M3::from_f64(min), F8E4M3::from_f64(max));
2730+
rand::distr::Uniform::new(F8E4M3::from_f64(min), F8E4M3::from_f64(max))
2731+
.map_err(Error::wrap)?;
27312732
for _i in 0..elem_count {
27322733
data.push(rng.sample::<F8E4M3, _>(uniform))
27332734
}

‎candle-core/src/cuda_backend/device.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::backend::BackendDevice;
22
use crate::{CpuStorage, CpuStorageRef, DType, Layout, Result, Shape};
33
pub use candle_kernels as kernels;
44
pub use cudarc;
5-
use cudarc::driver::{CudaFunction, LaunchAsync, LaunchConfig};
5+
use cudarc::driver::CudaFunction;
66
use float8::F8E4M3;
77
use half::{bf16, f16};
88
use std::collections::HashMap;
@@ -328,7 +328,7 @@ impl BackendDevice for CudaDevice {
328328
CudaStorageSlice::F64(data)
329329
}
330330
DType::F8E4M3 => {
331-
let data = self.alloc_zeros::<F8E4M3>(elem_count).w()?;
331+
let data = self.alloc_zeros::<F8E4M3>(elem_count)?;
332332
CudaStorageSlice::F8E4M3(data)
333333
}
334334
};
@@ -447,7 +447,7 @@ impl BackendDevice for CudaDevice {
447447
CudaStorageSlice::F64(data)
448448
}
449449
DType::F8E4M3 => {
450-
let data = self.alloc::<F8E4M3>(elem_count).w()?;
450+
let data = self.alloc::<F8E4M3>(elem_count)?;
451451
CudaStorageSlice::F8E4M3(data)
452452
}
453453
};
@@ -488,7 +488,7 @@ impl BackendDevice for CudaDevice {
488488
CudaStorageSlice::F64(data)
489489
}
490490
CpuStorageRef::F8E4M3(storage) => {
491-
let data = self.htod_sync_copy(storage).w()?;
491+
let data = self.memcpy_stod(storage)?;
492492
CudaStorageSlice::F8E4M3(data)
493493
}
494494
};
@@ -529,7 +529,7 @@ impl BackendDevice for CudaDevice {
529529
CudaStorageSlice::F64(data)
530530
}
531531
CpuStorage::F8E4M3(storage) => {
532-
let data = self.htod_sync_copy(storage).w()?;
532+
let data = self.memcpy_stod(storage)?;
533533
CudaStorageSlice::F8E4M3(data)
534534
}
535535
};
@@ -570,7 +570,7 @@ impl BackendDevice for CudaDevice {
570570
CudaStorageSlice::F64(data)
571571
}
572572
CpuStorage::F8E4M3(storage) => {
573-
let data = self.htod_copy(storage).w()?;
573+
let data = self.memcpy_stod(&storage)?;
574574
CudaStorageSlice::F8E4M3(data)
575575
}
576576
};

‎candle-core/src/cuda_backend/mod.rs‎

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ impl crate::scalar::Scalar {
4646
Scalar::F64(v) => builder.arg(v),
4747
Scalar::F16(v) => builder.arg(v),
4848
Scalar::BF16(v) => builder.arg(v),
49+
Scalar::F8E4M3(v) => builder.arg(v),
4950
};
5051
}
5152
}
@@ -1330,6 +1331,7 @@ impl BackendStorage for CudaStorage {
13301331
S::F16(s) => (slice_ptr(s, src_o), "const_set_f16"),
13311332
S::F32(s) => (slice_ptr(s, src_o), "const_set_f32"),
13321333
S::F64(s) => (slice_ptr(s, src_o), "const_set_f64"),
1334+
S::F8E4M3(s) => (slice_ptr(s, src_o), "const_set_f8_e4m3"),
13331335
};
13341336

13351337
let func = dev.get_or_load_func(kernel_name, &kernels::FILL)?;
@@ -1448,9 +1450,16 @@ impl BackendStorage for CudaStorage {
14481450
CudaStorageSlice::F64(out)
14491451
}
14501452
DType::F8E4M3 => {
1451-
let out = unsafe { dev.alloc::<F8E4M3>(el) }.w()?;
1452-
let params = (el, dims.len(), &ds, *inp, &out);
1453-
unsafe { func.launch(cfg, params) }.w()?;
1453+
let out: CudaSlice<F8E4M3> = unsafe { dev.alloc::<F8E4M3>(el) }?;
1454+
1455+
let mut builder = func.builder();
1456+
barg!(builder, el);
1457+
barg!(builder, dims.len());
1458+
ds.builder_arg(&mut builder);
1459+
barg!(builder, *inp);
1460+
builder.arg(&out);
1461+
unsafe { builder.launch(cfg) }.w()?;
1462+
14541463
CudaStorageSlice::F8E4M3(out)
14551464
}
14561465
};
@@ -1538,8 +1547,7 @@ impl BackendStorage for CudaStorage {
15381547
Ok(CpuStorage::F64(cpu_storage))
15391548
}
15401549
CudaStorageSlice::F8E4M3(slice) => {
1541-
let dev = slice.device();
1542-
let cpu_storage = dev.dtoh_sync_copy(slice).w()?;
1550+
let cpu_storage = slice.stream().memcpy_dtov(slice).w()?;
15431551
Ok(CpuStorage::F8E4M3(cpu_storage))
15441552
}
15451553
}
@@ -2119,13 +2127,17 @@ impl BackendStorage for CudaStorage {
21192127
(CudaStorageSlice::F8E4M3(src), CudaStorageSlice::F8E4M3(dst)) => {
21202128
let (src, mut dst) = slice_src_and_dst(src, src_l, dst, dst_offset);
21212129
if src_l.is_contiguous() {
2122-
dev.dtod_copy(&src, &mut dst).w()?
2130+
dev.memcpy_dtod(&src, &mut dst)?
21232131
} else {
2124-
let func = dev.get_or_load_func("ucopy_f8_e4m3", kernels::UNARY)?;
2125-
// SAFETY: Set later by running the kernel.
2126-
let params = (el_count, dims.len(), &ds, &src, &mut dst);
2132+
let func = dev.get_or_load_func("ucopy_f8_e4m3", &kernels::UNARY)?;
2133+
let mut builder = func.builder();
2134+
barg!(builder, el_count);
2135+
barg!(builder, dims.len());
2136+
ds.builder_arg(&mut builder);
2137+
builder.arg(&src);
2138+
builder.arg(&mut dst);
21272139
// SAFETY: ffi.
2128-
unsafe { func.launch(cfg, params) }.w()?
2140+
unsafe { builder.launch(cfg) }.w()?;
21292141
}
21302142
}
21312143
(CudaStorageSlice::U8(src), CudaStorageSlice::U8(dst)) => {

‎candle-core/src/cuda_backend/utils.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ pub trait Map2InPlace {
123123
(S::F16(dst), S::F16(src)) => self.f(dst, dst_l, src, src_l, d),
124124
(S::F32(dst), S::F32(src)) => self.f(dst, dst_l, src, src_l, d),
125125
(S::F64(dst), S::F64(src)) => self.f(dst, dst_l, src, src_l, d),
126-
(S::F8E4M3(dst), S::F8E4M3(src)) => self.f(dst, dst_s, src, src_l, d),
126+
(S::F8E4M3(dst), S::F8E4M3(src)) => self.f(dst, dst_l, src, src_l, d),
127127
_ => Err(CudaError::InternalError("dtype mismatch in binary op"))?,
128128
}
129129
}

‎candle-core/src/scalar.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
//! TensorScalar Enum and Trait
22
//!
33
use crate::{DType, Result, Tensor, WithDType};
4+
use float8::F8E4M3;
45
use half::{bf16, f16};
56

67
#[derive(Debug, Clone, Copy, PartialEq)]
@@ -12,6 +13,7 @@ pub enum Scalar {
1213
F16(f16),
1314
F32(f32),
1415
F64(f64),
16+
F8E4M3(F8E4M3),
1517
}
1618

1719
impl<T: WithDType> From<T> for Scalar {
@@ -30,6 +32,7 @@ impl Scalar {
3032
DType::F16 => Scalar::F16(f16::ZERO),
3133
DType::F32 => Scalar::F32(0.0),
3234
DType::F64 => Scalar::F64(0.0),
35+
DType::F8E4M3 => Scalar::F8E4M3(F8E4M3::ZERO),
3336
}
3437
}
3538

@@ -42,6 +45,7 @@ impl Scalar {
4245
DType::F16 => Scalar::F16(f16::ONE),
4346
DType::F32 => Scalar::F32(1.0),
4447
DType::F64 => Scalar::F64(1.0),
48+
DType::F8E4M3 => Scalar::F8E4M3(F8E4M3::ONE),
4549
}
4650
}
4751

@@ -54,6 +58,7 @@ impl Scalar {
5458
Scalar::F16(_) => DType::F16,
5559
Scalar::F32(_) => DType::F32,
5660
Scalar::F64(_) => DType::F64,
61+
Scalar::F8E4M3(_) => DType::F8E4M3,
5762
}
5863
}
5964

@@ -66,6 +71,7 @@ impl Scalar {
6671
Scalar::F16(v) => v.to_f64(),
6772
Scalar::F32(v) => *v as f64,
6873
Scalar::F64(v) => *v,
74+
Scalar::F8E4M3(v) => v.to_f64(),
6975
}
7076
}
7177
}

‎candle-kernels/src/fill.cu‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,4 +83,5 @@ CONST_SET_OP(__nv_bfloat16, const_set_bf16)
8383

8484
extern "C" __global__ void fill_f8_e4m3(__nv_fp8_e4m3 *buf, __nv_fp8_e4m3 value, const size_t numel) { fill_with(buf, value, numel); }
8585
COPY2D_OP(__nv_fp8_e4m3, copy2d_f8_e4m3)
86+
CONST_SET_OP(__nv_fp8_e4m3, const_set_f8_e4m3)
8687
#endif

‎candle-pyo3/Cargo.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ candle-onnx = { workspace = true, optional = true }
2121
half = { workspace = true }
2222
intel-mkl-src = { workspace = true, optional = true }
2323
pyo3 = { version = "0.22.0", features = ["extension-module", "abi3-py311"] }
24+
float8 = { workspace = true }
2425

2526
[build-dependencies]
2627
pyo3-build-config = "0.22"
@@ -29,6 +30,5 @@ pyo3-build-config = "0.22"
2930
default = []
3031
accelerate = ["dep:accelerate-src", "candle/accelerate"]
3132
cuda = ["candle/cuda"]
32-
mkl = ["dep:intel-mkl-src","candle/mkl"]
33+
mkl = ["dep:intel-mkl-src", "candle/mkl"]
3334
onnx = ["dep:candle-onnx"]
34-

‎candle-pyo3/src/lib.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#![allow(clippy::redundant_closure_call)]
22
#![allow(clippy::useless_conversion)]
3+
use float8::F8E4M3;
4+
use half::{bf16, f16};
35
use pyo3::exceptions::{PyTypeError, PyValueError};
46
use pyo3::prelude::*;
57
use pyo3::pyclass::CompareOp;
@@ -9,8 +11,6 @@ use std::collections::hash_map::DefaultHasher;
911
use std::hash::{Hash, Hasher};
1012
use std::sync::Arc;
1113

12-
use half::{bf16, f16};
13-
1414
#[cfg(feature = "mkl")]
1515
extern crate intel_mkl_src;
1616

@@ -158,6 +158,7 @@ pydtype!(f16, f32::from);
158158
pydtype!(bf16, f32::from);
159159
pydtype!(f32, |v| v);
160160
pydtype!(f64, |v| v);
161+
pydtype!(F8E4M3, f32::from);
161162

162163
fn actual_index(t: &Tensor, dim: usize, index: i64) -> ::candle::Result<usize> {
163164
let dim = t.dim(dim)?;
@@ -205,6 +206,7 @@ trait MapDType {
205206
DType::F16 => self.f::<f16>(t),
206207
DType::F32 => self.f::<f32>(t),
207208
DType::F64 => self.f::<f64>(t),
209+
DType::F8E4M3 => self.f::<F8E4M3>(t),
208210
}
209211
}
210212
}

‎candle-transformers/src/models/deepseek2.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ impl CustomOp1 for NonZero {
5050
candle::CpuStorage::F16(vs) => self.nonzero(vs, layout),
5151
candle::CpuStorage::F32(vs) => self.nonzero(vs, layout),
5252
candle::CpuStorage::F64(vs) => self.nonzero(vs, layout),
53+
candle::CpuStorage::F8E4M3(vs) => self.nonzero(vs, layout),
5354
};
5455
let index_len = layout.dims().len();
5556
let result_len = result.len() / index_len;

0 commit comments

Comments
 (0)