From 2a0f64375c96a1653a2225595c424406d07110f3 Mon Sep 17 00:00:00 2001 From: John Date: Sun, 2 Aug 2026 10:09:57 +1000 Subject: [PATCH] Cache the prepared device gather index (5.9x total on the hot op) Caching only the raw indice map left two thirds of the runtime on the table. Every layer still rebuilt the missing->N sentinel substitution over 5.6M elements in numpy and re-uploaded a 22.6MB index to the GPU. Isolating the kernel showed ~13ms of actual GPU work behind ~26ms of CPU bookkeeping. The transposed, sentinel-substituted index depends only on the coordinate set - the same invariant that justifies caching the indice map - so it is now cached on-device whole. m3ultra 128^3/128ch, cumulative: per-offset loop 81.6ms 2.57 Mvox/s fused gather+matmul 39.3ms 5.34 cached device index 13.8ms 15.25 <- 5.9x total A chunk-size sweep confirmed 256MB (11 dispatches) is at the optimum; unchunked is marginally slower (14.2ms), so the chunking is free insurance for small-memory boxes. --- lato_mlx/sparse/conv.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/lato_mlx/sparse/conv.py b/lato_mlx/sparse/conv.py index 38daee7..224404a 100644 --- a/lato_mlx/sparse/conv.py +++ b/lato_mlx/sparse/conv.py @@ -107,14 +107,24 @@ class SubMConv3d(nn.Module): if bias: self.bias = mx.zeros((out_channels,)) - def _indice_map(self, x: SparseTensor) -> np.ndarray: - key = f"imap_k{self.kernel_size}_{self.indice_key}" + def _gather_index(self, x: SparseTensor, n: int) -> mx.array: + """Cached [N, K^3] gather index, already on-device. + + Caching only the raw indice map is not enough: rebuilding the "missing -> N" + substitution and re-uploading the index cost more than the convolution itself. + At 128^3/128ch that overhead was ~26ms against ~13ms of actual GPU work, i.e. + two thirds of the measured time was CPU-side bookkeeping repeated every layer. + The transposed, sentinel-substituted device array depends only on the + coordinate set, so it is cached whole. + """ + key = f"gidx_k{self.kernel_size}_{self.indice_key}" cached = x.cache_get(key) if cached is not None: return cached imap = build_indice_map(x.coords, self.kernel_size) - x.cache_put(key, imap) - return imap + idx_t = mx.array(np.where(imap == _MISSING, n, imap).T) # [N, K^3] + x.cache_put(key, idx_t) + return idx_t def __call__(self, x: SparseTensor) -> SparseTensor: n = x.feats.shape[0] @@ -127,15 +137,13 @@ class SubMConv3d(nn.Module): out = out + self.bias return x.replace(out) - imap = self._indice_map(x) - # One appended zero row: absent neighbours index it and contribute nothing, # which avoids a per-offset boolean mask. feats_pad = mx.concatenate( [x.feats, mx.zeros((1, self.in_channels), dtype=x.feats.dtype)], axis=0 ) - idx = np.where(imap == _MISSING, n, imap) # [K^3, N] - k3 = imap.shape[0] + idx_t = self._gather_index(x, n) # [N, K^3], cached on device + k3 = self.kernel_size**3 # Fuse the K^3 taps into ONE gather + ONE matmul. # @@ -157,8 +165,6 @@ class SubMConv3d(nn.Module): w_flat = self.weight.reshape(k3 * self.in_channels, self.out_channels).astype( x.feats.dtype ) - idx_t = mx.array(idx.T) # [N, K^3] - bytes_per_row = k3 * self.in_channels * 4 chunk = max(1, min(n, (256 << 20) // max(bytes_per_row, 1)))