MPS Support

[ludsvick]

🔥 Summary

Adds Metal Performance Shaders (MPS) for PyTorch to boost performance on MacOS with Apple Silicon/AMD GPUs. Also updates the Cache class with a device-aware update method.

📖 📷 Additional context

• Updates to kv_cache now handled with update method
  • Uses static kernel ops for CUDA Graph Mode
  • Falls back to slicing on unsupported devices
• generate and vq_vae_encode_decode updated to include mps support

🛠 Testing instructions

Mesh Generation, CPU

import torch
import trimesh
from cube3d.inference.engine import Engine

# load ckpt
config_path = "cube3d/configs/open_model.yaml"
gpt_ckpt_path = "model_weights/shape_gpt.safetensors"
shape_ckpt_path = "model_weights/shape_tokenizer.safetensors"
engine = Engine(
    config_path, 
    gpt_ckpt_path, 
    shape_ckpt_path, 
    device=torch.device("cpu"),
)

# inference
input_prompt = "A scaly gecko with a long winding tail"
mesh_v_f = engine.t2s([input_prompt], use_kv_cache=True, resolution_base=6.5, top_p=0.0)

# save output
vertices, faces = mesh_v_f[0][0], mesh_v_f[0][1]
_ = trimesh.Trimesh(vertices=vertices, faces=faces).export("outputs/cpu_output.obj")

Output

generating: 100%|█████████████████████████████| 512/512 [00:56<00:00,  8.98it/s]
extracting geometry: 100%|█████████████████████| 8/8 [05:28<00:00, 41.03s/chunk]

Mesh Generation, MPS

import torch
import trimesh
from cube3d.inference.engine import Engine

# load ckpt
config_path = "cube3d/configs/open_model.yaml"
gpt_ckpt_path = "model_weights/shape_gpt.safetensors"
shape_ckpt_path = "model_weights/shape_tokenizer.safetensors"
engine = Engine(
    config_path, 
    gpt_ckpt_path, 
    shape_ckpt_path, 
    device=torch.device("mps"),
)

# inference
input_prompt = "A scaly gecko with a long winding tail"
mesh_v_f = engine.t2s([input_prompt], use_kv_cache=True, resolution_base=6.5, top_p=0.0)

# save output
vertices, faces = mesh_v_f[0][0], mesh_v_f[0][1]
_ = trimesh.Trimesh(vertices=vertices, faces=faces).export("outputs/mps_output.obj")

Output

generating: 100%|█████████████████████████████| 512/512 [01:12<00:00,  7.08it/s]
extracting geometry: 100%|█████████████████████| 8/8 [00:17<00:00,  2.15s/chunk]

✅ Checklist

• Provide testing instructions
• Update relevant documentation

[animan42]

Thank you for getting this out quick! I will test it on some cuda devices later today

[animan42]

Lets move device selection to a util function in https://github.com/Roblox/cube/blob/main/cube3d/inference/utils.py since its being used at-least twice

[ludsvick]

select_device() added to utils 👍