Couple of minor bugs.

[githubsgi]

Bug description

1, Does not appear to be bracketed by compile flag.

torchtitan/torchtitan/experiments/kernels/moe/indices.py

Line 16 in ad06609

@triton.jit

2. Calls to cuda instead of accelerator.
   

   torchtitan/torchtitan/distributed/expert_parallel.py

   Line 197 in ad06609

   torch.cuda.current_stream().synchronize()

Versions

TorchTitan and PyTorch master.

[tianyu-l]

cc @xmfan on issue 1 (honestly I don't know what it means)

[githubsgi]

1 should be enabled only when torch.compile is enabled.

[xmfan]

@tianyu-l I think this is about hardware heterogeneity, we have a triton kernel in the main path of the MoE, so it wouldn't work for a device triton doesn't support.

Gating it only under compile wouldn't be a good approach because using compile doesn't mean you need to use triton. You probably just want to check if the device is supported.