From 369cb6c6fee0481ce31cf0575a57527d0b07dfb9 Mon Sep 17 00:00:00 2001 From: Pedro Augusto <16762743+pedronaugusto@users.noreply.github.com> Date: Tue, 21 Apr 2026 16:27:02 +0100 Subject: [PATCH] sparse/norm: CPU-staged zeros_like workaround for MPS SparseGroupNorm / SparseLayerNorm used torch.zeros_like which fails with "DispatchStub: missing kernel for mps" on PyTorch builds compiled with both CUDA and MPS backends (the user's local build hit this). Added a _zeros_like_safe helper that builds zeros on CPU and transfers when the reference tensor is on MPS; Apple Silicon unified memory makes the transfer metadata-only, so the overhead vs a working MPS zeros kernel is negligible. On CPU, behaves identically to torch.zeros_like. Co-Authored-By: Claude Opus 4.7 (1M context) --- trellis2/modules/sparse/norm.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/trellis2/modules/sparse/norm.py b/trellis2/modules/sparse/norm.py index 9571120..654c0c2 100644 --- a/trellis2/modules/sparse/norm.py +++ b/trellis2/modules/sparse/norm.py @@ -12,12 +12,24 @@ __all__ = [ ] +def _zeros_like_safe(t: torch.Tensor) -> torch.Tensor: + """zeros_like equivalent that works around PyTorch MPS builds missing + fp16/fp32 zeros kernels (DispatchStub missing). Constructs zeros on CPU + and transfers — Apple Silicon unified memory makes the transfer a + metadata-only operation, so the overhead vs a true MPS zeros kernel is + negligible when it does exist.""" + if t.device.type == 'mps': + cpu_zeros = torch.zeros(t.shape, dtype=t.dtype) + return cpu_zeros.to(t.device) + return torch.zeros_like(t) + + class SparseGroupNorm(nn.GroupNorm): def __init__(self, num_groups, num_channels, eps=1e-5, affine=True): super(SparseGroupNorm, self).__init__(num_groups, num_channels, eps, affine) def forward(self, input: VarLenTensor) -> VarLenTensor: - nfeats = torch.zeros_like(input.feats) + nfeats = _zeros_like_safe(input.feats) for k in range(input.shape[0]): bfeats = input.feats[input.layout[k]] bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1) @@ -32,7 +44,7 @@ class SparseLayerNorm(nn.LayerNorm): super(SparseLayerNorm, self).__init__(normalized_shape, eps, elementwise_affine) def forward(self, input: VarLenTensor) -> VarLenTensor: - nfeats = torch.zeros_like(input.feats) + nfeats = _zeros_like_safe(input.feats) for k in range(input.shape[0]): bfeats = input.feats[input.layout[k]] bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1)