monster/trellis-2-mrp-mlx
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

sparse/norm: CPU-staged zeros_like workaround for MPS

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
Pedro Augusto 2026-04-21 16:27:02 +01:00 committed by Jourloy
parent c869eb995a
commit 369cb6c6fe
1 changed files with 14 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
trellis2/modules/sparse/norm.py
View File

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