hanzo_ml/
hc_sinkhorn.rs

1//! DeepSeek-V4 mHC fused Sinkhorn.
2//!
3//! The hyper-connection combine matrix is tiny (`[.., nh, nh]`, `nh` = n_hc = 4)
4//! but the 20-iteration row/col alternation, expressed as elementwise tensor ops,
5//! was ~120 tiny kernel launches per call — a swarm of 16-thread kernels that
6//! dominated V4 decode (nsys: ~20% of GPU time, ~10.7k tiny kernels/token). This
7//! runs the whole iteration in ONE launch, on-device, in f32 (the tensor-op path
8//! ran it in bf16 — low precision for a 4×4 normalization). Semantics match the
9//! reference exactly: `softmax(src) + eps`, `col_norm` (denom = colsum+eps), then
10//! `(iters-1)× (row_norm, col_norm)`.
11
12use crate::{CpuStorage, Layout, Result, Shape, Tensor};
13
14#[derive(Debug, Clone, Copy)]
15struct HcSinkhorn {
16    iters: usize,
17    eps: f64,
18}
19
20fn row_norm(mat: &mut [f32], nh: usize, eps: f32) {
21    for r in 0..nh {
22        let mut s = eps;
23        for c in 0..nh {
24            s += mat[r * nh + c];
25        }
26        for c in 0..nh {
27            mat[r * nh + c] /= s;
28        }
29    }
30}
31
32fn col_norm(mat: &mut [f32], nh: usize, eps: f32) {
33    for c in 0..nh {
34        let mut s = eps;
35        for r in 0..nh {
36            s += mat[r * nh + c];
37        }
38        for r in 0..nh {
39            mat[r * nh + c] /= s;
40        }
41    }
42}
43
44/// CPU reference (also the fallback for non-CUDA devices). Bit-for-bit the same
45/// arithmetic as the CUDA kernel (both f32).
46fn sinkhorn_cpu_f32(data: &[f32], nmat: usize, nh: usize, iters: usize, eps: f32) -> Vec<f32> {
47    let mut out = vec![0f32; data.len()];
48    for m in 0..nmat {
49        let base = m * nh * nh;
50        let src = &data[base..base + nh * nh];
51        let mat = &mut out[base..base + nh * nh];
52        for r in 0..nh {
53            let row = &src[r * nh..r * nh + nh];
54            let mx = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
55            let mut sum = 0f32;
56            for c in 0..nh {
57                let e = (row[c] - mx).exp();
58                mat[r * nh + c] = e;
59                sum += e;
60            }
61            for c in 0..nh {
62                mat[r * nh + c] = mat[r * nh + c] / sum + eps;
63            }
64        }
65        col_norm(mat, nh, eps);
66        for _ in 1..iters {
67            row_norm(mat, nh, eps);
68            col_norm(mat, nh, eps);
69        }
70    }
71    out
72}
73
74impl crate::CustomOp1 for HcSinkhorn {
75    fn name(&self) -> &'static str {
76        "hc_sinkhorn"
77    }
78
79    fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> Result<(CpuStorage, Shape)> {
80        let (o1, o2) = match layout.contiguous_offsets() {
81            Some(x) => x,
82            None => crate::bail!("hc_sinkhorn requires a contiguous input"),
83        };
84        let vs = match storage {
85            CpuStorage::F32(v) => &v[o1..o2],
86            _ => crate::bail!("hc_sinkhorn: f32 only"),
87        };
88        let dims = layout.shape().dims();
89        let nh = dims[dims.len() - 1];
90        let nmat = vs.len() / (nh * nh);
91        let out = sinkhorn_cpu_f32(vs, nmat, nh, self.iters, self.eps as f32);
92        Ok((CpuStorage::F32(out), layout.shape().clone()))
93    }
94
95    #[cfg(feature = "cuda")]
96    fn cuda_fwd(
97        &self,
98        storage: &crate::CudaStorage,
99        layout: &Layout,
100    ) -> Result<(crate::CudaStorage, Shape)> {
101        use crate::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg};
102        use crate::cuda_backend::{kernels, CudaStorageSlice as S, WrapErr};
103
104        let dev = &storage.device;
105        let (o1, o2) = match layout.contiguous_offsets() {
106            Some(x) => x,
107            None => crate::bail!("hc_sinkhorn requires a contiguous input"),
108        };
109        let src = storage.as_cuda_slice::<f32>()?.slice(o1..o2);
110        let dims = layout.shape().dims();
111        let nh = dims[dims.len() - 1];
112        let elem = layout.shape().elem_count();
113        let nmat = elem / (nh * nh);
114
115        let dst = unsafe { dev.alloc::<f32>(elem)? };
116        let func = dev.get_or_load_func("hc_sinkhorn_f32", &kernels::REDUCE)?;
117        let cfg = LaunchConfig {
118            grid_dim: (nmat as u32, 1, 1),
119            block_dim: (32, 1, 1),
120            shared_mem_bytes: (nh * nh * std::mem::size_of::<f32>()) as u32,
121        };
122        let stream = dev.cuda_stream();
123        let mut builder = stream.launch_builder(&func);
124        let nh_i = nh as i32;
125        let iters_i = self.iters as i32;
126        let eps = self.eps as f32;
127        builder
128            .arg(&src)
129            .arg(&dst)
130            .arg(&nh_i)
131            .arg(&iters_i)
132            .arg(&eps);
133        unsafe { builder.launch(cfg) }.w()?;
134
135        let out = crate::cuda_backend::CudaStorage {
136            slice: S::F32(dst),
137            device: dev.clone(),
138        };
139        Ok((out, layout.shape().clone()))
140    }
141}
142
143impl Tensor {
144    /// Fused mHC Sinkhorn over the trailing `[nh, nh]` matrices: `softmax(src)+eps`,
145    /// `col_norm`, then `(iters-1)× (row_norm, col_norm)`, all in f32. Runs the whole
146    /// iteration in one CUDA launch (falls back to a CPU reference off-device).
147    /// Returns the caller's dtype.
148    pub fn hc_sinkhorn(&self, iters: usize, eps: f64) -> Result<Tensor> {
149        let orig = self.dtype();
150        let x = self.to_dtype(crate::DType::F32)?.contiguous()?;
151        let r = x.apply_op1_no_bwd(&HcSinkhorn { iters, eps })?;
152        r.to_dtype(orig)
153    }
154}
hanzo_ml/hc_sinkhorn.rs

hanzo_ml/
hc_sinkhorn.rs
