monster/lato.2_mrp_mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Cache the prepared device gather index (5.9x total on the hot op)

Browse Source 
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.
This commit is contained in:
John 2026-08-02 10:09:57 +10:00
parent 6fcf677cef
commit 2a0f64375c
1 changed files with 16 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
lato_mlx/sparse/conv.py
View File

@ -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)))