viable/strict/1793100530: [MPS] Migrate repeat to metal (#202121)

viable/strict/1793100530: [MPS] Migrate repeat to metal (#202121)
Migrates repeat from MPS Graph to metal. 2D input with two repeat factors use the repeat_2d gather kernel. This was needed because otherwise hitting same perf as MPS Graph was not possible with only copy. All other inputs unfold the contiguous output into a view and run one copy with input broadcast over the repeated dims. Perf: Model usage Input shape Layout dtype Repeats MPSGraph (us) Metal (us) Speedup FLUX prompt embeds (1, 512, 4096) contiguous bf16 (1, 2, 1) 72.2 26.2 2.76x SDXL prompt embeds (1, 77, 2048) contiguous fp16 (1, 1, 1) 17.9 3.5 5.11x FLUX pooled embeds (1, 768) contiguous bf16 (1, 2) 17.2 2.3 7.48x Z-Image CFG latents (1, 16, 128, 128) contiguous bf16 (2, 1, 1, 1) 19.2 4.8 4.00x LTX VAE causal pad (1, 128, 1, 256, 352) frame slice bf16 (1, 1, 2, 1, 1) 320.8 264.3 1.21x Whisper position ids (1, 1) contiguous int64 (1, 1) 18.2 6.6 2.76x RT-DETR top-k index (1, 300, 1) contiguous int64 (1, 1, 256) 17.1 4.0 4.28x Mask2Former attention mask (1, 1, 100, 16384) contiguous fp32 (1, 8, 1, 1) 211.3 183.9 1.15x Llama 4 MoE (512, 5120) contiguous bf16 (16, 1) 559.5 519.7 1.08x Transposed matrix (2048, 2048) transposed fp32 (2, 2) 426.0 415.0 1.03x SAM image pos embed (1, 256, 64, 64) permuted fp32 (1, 1, 1, 1) 37.0 56.7 0.65x HunyuanVideo i2v frame broadcast (1, 16, 1, 90, 160) contiguous fp32 (1, 1, 33, 1, 1) 88.9 124.3 0.72x Pull Request resolved: #202121 Approved by: https://github.com/Skylion007, https://github.com/malfet

Take Your Experience to the Next Level

New

Download our mobile app for a faster and better experience.

Comments

0
U

Join the discussion

Sign in to leave a comment

0:000:00