Merge branch 'main' into mf/cpu-varlen-flash-attention

michaelfeil · web-flow · commit 5293865e6ccc · 2026-01-02T08:37:27.000+01:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -6,6 +6,7 @@ members = [
     "candle-nn",
     "candle-pyo3",
     "candle-transformers",
+    "candle-ug",
     "candle-wasm-examples/*",
     "candle-wasm-tests",
     "tensor-tools",
@@ -43,6 +44,7 @@ candle-metal-kernels = { path = "./candle-metal-kernels", version = "0.9.2-alpha
 candle-nn = { path = "./candle-nn", version = "0.9.2-alpha.2" }
 candle-onnx = { path = "./candle-onnx", version = "0.9.2-alpha.2" }
 candle-transformers = { path = "./candle-transformers", version = "0.9.2-alpha.2" }
+candle-ug = { path = "./candle-ug", version = "0.9.2-alpha.2" }
 clap = { version = "4.2.4", features = ["derive"] }
 criterion = { version = "0.7.0", default-features = false }
 cudarc = { version = "0.18.1", features = [
@@ -65,10 +67,7 @@ half = { version = "2.5.0", features = [
     "use-intrinsics",
     "rand_distr",
 ] }
-float8 = { version = "0.5.0", features = [
-    "num-traits",
-    "rand_distr",
-] }
+float8 = { version = "0.5.0", features = ["num-traits", "rand_distr"] }
 hound = "3.5.1"
 image = { version = "0.25.2", default-features = false, features = [
     "jpeg",
diff --git a/candle-core/Cargo.toml b/candle-core/Cargo.toml
@@ -35,9 +35,7 @@ yoke = { workspace = true }
 zip = { workspace = true }
 
 [target.'cfg(all(not(target_arch = "wasm32"), not(target_os = "ios")))'.dependencies]
-ug = { workspace = true }
-ug-cuda = { workspace = true, optional = true }
-ug-metal = { workspace = true, optional = true }
+candle-ug = { workspace = true, optional = true }
 
 [dev-dependencies]
 anyhow = { workspace = true }
@@ -46,7 +44,7 @@ criterion = { workspace = true }
 
 [features]
 default = []
-cuda = ["cudarc", "dep:candle-kernels", "dep:ug-cuda", "float8/cuda"]
+cuda = ["cudarc", "dep:candle-kernels", "candle-ug?/cuda", "float8/cuda"]
 cudnn = ["cuda", "cudarc/cudnn"]
 nccl = ["cuda", "cudarc/nccl"]
 mkl = ["dep:libc", "dep:intel-mkl-src"]
@@ -55,8 +53,9 @@ metal = [
     "dep:objc2-metal",
     "dep:objc2-foundation",
     "dep:candle-metal-kernels",
-    "dep:ug-metal",
+    "candle-ug?/metal",
 ]
+ug = ["dep:candle-ug"]
 
 [[bench]]
 name = "bench_main"
diff --git a/candle-core/src/cuda_backend/device.rs b/candle-core/src/cuda_backend/device.rs
@@ -173,14 +173,14 @@ impl CudaDevice {
         self.context.is_event_tracking()
     }
 
-    #[cfg(not(target_arch = "wasm32"))]
+    #[cfg(all(feature = "ug", not(target_arch = "wasm32")))]
     pub fn compile(
         &self,
         func_name: &'static str,
-        kernel: ug::lang::ssa::Kernel,
+        kernel: candle_ug::lang::ssa::Kernel,
     ) -> Result<CudaFunc> {
         let mut buf = vec![];
-        ug_cuda::code_gen::gen(&mut buf, func_name, &kernel)?;
+        candle_ug::cuda::code_gen::gen(&mut buf, func_name, &kernel)?;
         let cuda_code = String::from_utf8(buf)?;
         let opts = cudarc::nvrtc::CompileOptions {
             use_fast_math: Some(true),
diff --git a/candle-core/src/cuda_backend/mod.rs b/candle-core/src/cuda_backend/mod.rs
@@ -1012,7 +1012,7 @@ impl Map1 for UpsampleBilinear2D {
 
         // SAFETY: Set later by running the kernel.
         let out = unsafe { dev.alloc::<T>(dst_el)? };
-        let ds = dev.memcpy_stod(&ds)?;
+        let ds = dev.clone_htod(&ds)?;
 
         let mut builder = func.builder();
         barg!(builder, out_w);
diff --git a/candle-core/src/custom_op.rs b/candle-core/src/custom_op.rs
@@ -376,6 +376,7 @@ impl Tensor {
     }
 }
 
+#[cfg(feature = "ug")]
 pub struct UgIOp1 {
     name: &'static str,
     #[cfg(feature = "cuda")]
@@ -384,12 +385,13 @@ pub struct UgIOp1 {
     func: candle_metal_kernels::metal::ComputePipeline,
 }
 
+#[cfg(feature = "ug")]
 impl UgIOp1 {
     #[allow(unused)]
     #[cfg(all(not(target_arch = "wasm32"), not(target_os = "ios")))]
     pub fn new(
         name: &'static str,
-        kernel: ug::lang::ssa::Kernel,
+        kernel: candle_ug::lang::ssa::Kernel,
         device: &crate::Device,
     ) -> Result<Self> {
         #[cfg(feature = "cuda")]
@@ -414,6 +416,7 @@ impl UgIOp1 {
     }
 }
 
+#[cfg(feature = "ug")]
 impl InplaceOp1 for UgIOp1 {
     fn name(&self) -> &'static str {
         self.name
diff --git a/candle-core/src/error.rs b/candle-core/src/error.rs
@@ -174,9 +174,9 @@ pub enum Error {
     #[error("Metal error {0}")]
     Metal(#[from] MetalError),
 
-    #[cfg(all(not(target_arch = "wasm32"), not(target_os = "ios")))]
+    #[cfg(all(not(target_arch = "wasm32"), not(target_os = "ios"), feature = "ug"))]
     #[error(transparent)]
-    Ug(#[from] ug::Error),
+    Ug(#[from] candle_ug::Error),
 
     #[error(transparent)]
     TryFromIntError(#[from] core::num::TryFromIntError),
diff --git a/candle-core/src/lib.rs b/candle-core/src/lib.rs
@@ -92,7 +92,9 @@ mod variable;
 pub use cuda_backend::cudnn;
 
 pub use cpu_backend::{CpuStorage, CpuStorageRef};
-pub use custom_op::{CustomOp1, CustomOp2, CustomOp3, InplaceOp1, InplaceOp2, InplaceOp3, UgIOp1};
+#[cfg(feature = "ug")]
+pub use custom_op::UgIOp1;
+pub use custom_op::{CustomOp1, CustomOp2, CustomOp3, InplaceOp1, InplaceOp2, InplaceOp3};
 pub use device::{Device, DeviceLocation, NdArray};
 pub use dtype::{DType, DTypeParseError, FloatDType, IntDType, WithDType};
 pub use dummy_dtype::{F4, F6E2M3, F6E3M2, F8E8M0};
diff --git a/candle-core/src/metal_backend/device.rs b/candle-core/src/metal_backend/device.rs
@@ -1,13 +1,17 @@
 use crate::{DType, Result};
+
+#[cfg(feature = "ug")]
+use candle_metal_kernels::metal::ComputePipeline;
 use candle_metal_kernels::{
     metal::{
-        BlitCommandEncoder, Buffer, BufferMap, Commands, ComputeCommandEncoder, ComputePipeline,
-        Device, MTLResourceOptions,
+        BlitCommandEncoder, Buffer, BufferMap, Commands, ComputeCommandEncoder, Device,
+        MTLResourceOptions,
     },
     Kernels,
 };
 use objc2_foundation::NSURL;
 use objc2_metal::{MTLCaptureDescriptor, MTLCaptureDestination, MTLCaptureManager};
+
 use std::path::Path;
 use std::sync::{Arc, Mutex, RwLock};
 
@@ -88,14 +92,14 @@ impl std::ops::Deref for MetalDevice {
 }
 
 impl MetalDevice {
-    #[cfg(all(not(target_arch = "wasm32"), not(target_os = "ios")))]
+    #[cfg(all(feature = "ug", not(target_arch = "wasm32"), not(target_os = "ios")))]
     pub fn compile(
         &self,
         func_name: &'static str,
-        kernel: ug::lang::ssa::Kernel,
+        kernel: candle_ug::lang::ssa::Kernel,
     ) -> Result<ComputePipeline> {
         let mut buf = vec![];
-        ug_metal::code_gen::gen(&mut buf, func_name, &kernel)?;
+        candle_ug::metal::code_gen::gen(&mut buf, func_name, &kernel)?;
         let metal_code = String::from_utf8(buf)?;
         let lib = self
             .device
diff --git a/candle-core/src/op.rs b/candle-core/src/op.rs
@@ -247,31 +247,31 @@ pub trait BinaryOpT {
     fn i64_vec(_xs1: &[i64], _xs2: &[i64], _ys: &mut [i64]) {}
 }
 
-pub(crate) struct Add;
-pub(crate) struct Div;
-pub(crate) struct Mul;
-pub(crate) struct Sub;
-pub(crate) struct Maximum;
-pub(crate) struct Minimum;
-pub(crate) struct Exp;
-pub(crate) struct Log;
-pub(crate) struct Sin;
-pub(crate) struct Cos;
-pub(crate) struct Abs;
-pub(crate) struct Neg;
-pub(crate) struct Recip;
-pub(crate) struct Sqr;
-pub(crate) struct Sqrt;
-pub(crate) struct Gelu;
-pub(crate) struct GeluErf;
-pub(crate) struct Erf;
-pub(crate) struct Relu;
-pub(crate) struct Silu;
-pub(crate) struct Tanh;
-pub(crate) struct Floor;
-pub(crate) struct Ceil;
-pub(crate) struct Round;
-pub(crate) struct Sign;
+pub struct Add;
+pub struct Div;
+pub struct Mul;
+pub struct Sub;
+pub struct Maximum;
+pub struct Minimum;
+pub struct Exp;
+pub struct Log;
+pub struct Sin;
+pub struct Cos;
+pub struct Abs;
+pub struct Neg;
+pub struct Recip;
+pub struct Sqr;
+pub struct Sqrt;
+pub struct Gelu;
+pub struct GeluErf;
+pub struct Erf;
+pub struct Relu;
+pub struct Silu;
+pub struct Tanh;
+pub struct Floor;
+pub struct Ceil;
+pub struct Round;
+pub struct Sign;
 
 macro_rules! bin_op {
     ($op:ident, $name: literal, $e: expr, $f32_vec: ident, $f64_vec: ident) => {
diff --git a/candle-core/src/sort.rs b/candle-core/src/sort.rs
@@ -85,9 +85,11 @@ mod cuda {
             let ncols = self.last_dim;
             let nrows = elem_count / ncols;
             let ncols_pad = next_power_of_2(ncols);
+            // Limit block dim to 1024 threads, which is the maximum on modern CUDA gpus.
+            let block_dim = ncols_pad.min(1024);
             let cfg = LaunchConfig {
                 grid_dim: (nrows as u32, 1, 1),
-                block_dim: (ncols_pad as u32, 1, 1),
+                block_dim: (block_dim as u32, 1, 1),
                 shared_mem_bytes: (ncols_pad * std::mem::size_of::<u32>()) as u32,
             };
             let stream = dev.cuda_stream();
diff --git a/candle-core/tests/custom_op_tests.rs b/candle-core/tests/custom_op_tests.rs
@@ -145,20 +145,20 @@ fn inplace_op1() -> Result<()> {
     Ok(())
 }
 
-#[cfg(any(feature = "cuda", feature = "metal"))]
+#[cfg(all(feature = "ug", any(feature = "cuda", feature = "metal")))]
 #[allow(clippy::approx_constant)]
 #[test]
 fn ug_op() -> Result<()> {
     let kernel = {
-        use ug::lang::op;
+        use candle_ug::lang::op;
 
-        let layout = ug::Layout::from_shape(&[12]);
-        let ptr = op::Arg::ptr(ug::DType::F32);
-        let src = op::load(ptr.id(), layout.clone(), ug::DType::F32)?;
+        let layout = candle_ug::Layout::from_shape(&[12]);
+        let ptr = op::Arg::ptr(candle_ug::DType::F32);
+        let src = op::load(ptr.id(), layout.clone(), candle_ug::DType::F32)?;
         let src = op::unary(op::UnaryOp::Exp, src)?;
         let st = op::store(ptr.id(), layout, src)?;
         let kernel = op::Kernel::new("exp".to_string(), vec![ptr], vec![st]);
-        let opts: ug::lower_op::Opts = Default::default();
+        let opts: candle_ug::lower_op::Opts = Default::default();
         kernel.lower(&opts)?
     };
     let device = if candle_core::utils::cuda_is_available() {
diff --git a/candle-core/tests/tensor_tests.rs b/candle-core/tests/tensor_tests.rs
@@ -219,6 +219,26 @@ fn asort(device: &Device) -> Result<()> {
     Ok(())
 }
 
+/// Test sorting a large tensor that exceeds 1024 elements.
+fn asort_big(device: &Device) -> Result<()> {
+    // Skip on metal for now
+    if device.is_metal() {
+        return Ok(());
+    }
+    const SIZE: usize = 2000;
+    let data: Vec<f32> = (0..SIZE).map(|x| (SIZE - x) as f32).collect();
+    let tensor = Tensor::new(data.as_slice(), device)?;
+
+    let indexes = tensor.arg_sort_last_dim(true)?;
+    let expected_indexes: Vec<u32> = (0..SIZE).rev().map(|x| x as u32).collect();
+    assert_eq!(indexes.to_vec1::<u32>()?, expected_indexes);
+
+    let indexes = tensor.arg_sort_last_dim(false)?;
+    let expected_indexes: Vec<u32> = (0..SIZE).map(|x| x as u32).collect();
+    assert_eq!(indexes.to_vec1::<u32>()?, expected_indexes);
+    Ok(())
+}
+
 fn unary_op(device: &Device) -> Result<()> {
     let data = &[[-3f32, 1., 4., -0.1, 0.5], [2.7, -1.8, -0.28, 1.8, 2.8]];
     let tensor = Tensor::new(data, device)?;
@@ -1707,6 +1727,7 @@ test_device!(
 test_device!(randn, randn_cpu, randn_gpu, randn_metal);
 test_device!(clamp, clamp_cpu, clamp_gpu, clamp_metal);
 test_device!(asort, asort_cpu, asort_gpu, asort_metal);
+test_device!(asort_big, asort_big_cpu, asort_big_gpu, asort_big_metal);
 test_device!(var, var_cpu, var_gpu, var_metal);
 test_device!(zero_dim, zero_dim_cpu, zero_dim_gpu, zero_dim_metal);
 
diff --git a/candle-examples/examples/mnist-training/main.rs b/candle-examples/examples/mnist-training/main.rs
@@ -137,7 +137,7 @@ fn training_loop_cnn(
     let test_labels = m.test_labels.to_dtype(DType::U32)?.to_device(&dev)?;
     let n_batches = train_images.dim(0)? / BSIZE;
     let mut batch_idxs = (0..n_batches).collect::<Vec<usize>>();
-    for epoch in 1..args.epochs {
+    for epoch in 1..=args.epochs {
         let mut sum_loss = 0f32;
         batch_idxs.shuffle(&mut rng());
         for batch_idx in batch_idxs.iter() {
@@ -194,7 +194,7 @@ fn training_loop<M: Model>(
     let mut sgd = candle_nn::SGD::new(varmap.all_vars(), args.learning_rate)?;
     let test_images = m.test_images.to_device(&dev)?;
     let test_labels = m.test_labels.to_dtype(DType::U32)?.to_device(&dev)?;
-    for epoch in 1..args.epochs {
+    for epoch in 1..=args.epochs {
         let logits = model.forward(&train_images)?;
         let log_sm = ops::log_softmax(&logits, D::Minus1)?;
         let loss = loss::nll(&log_sm, &train_labels)?;
diff --git a/candle-kernels/src/sort.cu b/candle-kernels/src/sort.cu
diff --git a/candle-nn/src/loss.rs b/candle-nn/src/loss.rs
diff --git a/candle-nn/tests/loss.rs b/candle-nn/tests/loss.rs
diff --git a/candle-transformers/src/models/quantized_gemma3.rs b/candle-transformers/src/models/quantized_gemma3.rs
diff --git a/candle-ug/Cargo.toml b/candle-ug/Cargo.toml
diff --git a/candle-ug/src/lib.rs b/candle-ug/src/lib.rs

Original file line number	Diff line number	Diff line change
`@@ -376,6 +376,7 @@ impl Tensor {`
`376`	`376`	`}`
`377`	`377`	`}`
`378`	`378`
	`379`	`+#[cfg(feature = "ug")]`
`379`	`380`	`pub struct UgIOp1 {`
`380`	`381`	`name: &'static str,`
`381`	`382`	`#[cfg(feature = "cuda")]`
`@@ -384,12 +385,13 @@ pub struct UgIOp1 {`
`384`	`385`	`func: candle_metal_kernels::metal::ComputePipeline,`
`385`	`386`	`}`
`386`	`387`
	`388`	`+#[cfg(feature = "ug")]`
`387`	`389`	`impl UgIOp1 {`
`388`	`390`	`#[allow(unused)]`
`389`	`391`	`#[cfg(all(not(target_arch = "wasm32"), not(target_os = "ios")))]`
`390`	`392`	`pub fn new(`
`391`	`393`	`name: &'static str,`
`392`		`- kernel: ug::lang::ssa::Kernel,`
	`394`	`+ kernel: candle_ug::lang::ssa::Kernel,`
`393`	`395`	`device: &crate::Device,`
`394`	`396`	`) -> Result<Self> {`
`395`	`397`	`#[cfg(feature = "cuda")]`
`@@ -414,6 +416,7 @@ impl UgIOp1 {`
`414`	`416`	`}`
`415`	`417`	`}`
`416`	`418`
	`419`	`+#[cfg(feature = "ug")]`
`417`	`420`	`impl InplaceOp1 for UgIOp1 {`
`418`	`421`	`fn name(&self) -> &'static str {`
`419`	`422`	`self.name`