i think it makes more sense if we continue using int as the default for any hardware we are not aware of. because it's better to have the best perf on as many hardware as possible and only use float if it breaks it.

because int histograms are easier to compute than floating ones

@IlyasMoutawwakil I agree with you — using int by default makes sense for performance whenever it's reliably supported, since integer histograms are generally faster and cheaper to compute.
That said, a broader range of hardware (such as MPS on Apple Silicon, and potentially some other backends) currently doesn't support integer dtypes for torch.histc, which is what caused the NotImplementedError.
I also considered the "try int first, fallback to float" pattern. While it's more aggressive at giving good performance on unknown hardware, it has a couple of downsides:

• It can hide bugs or lead to inconsistent behavior across devices.
• The fallback path adds a small runtime cost and makes the code slightly more complex.

Because of that, I went with the explicit approach in this PR: int on CUDA (where we know it works well and gives the best perf) and float32 everywhere else for broad compatibility.
Happy to adjust if there's a cleaner way to handle this.

Given that the large majority of users (especially for MoE models) run on CUDA, this approach gives them the best performance while ensuring reliable behavior on other backends like MPS. We can always expand the int64 support list later as more backends (ROCm, XPU, etc.) stabilize.

It feels like this is a classic performance vs portability tradeoff:

• int histograms → better performance
• float histograms → broader hardware support (MPS, etc.)

One question I had: do we know how significant the performance difference is in practice for MoE workloads?

If the gap is large, it might be worth keeping int as the default and falling back only when needed. But if the difference is small, the current approach (float for non-CUDA) seems safer.

Also wondering whether a "try int → fallback to float" approach could be acceptable if the failure is rare and predictable.

It feels like this is a classic performance vs portability tradeoff:

* int histograms → better performance

* float histograms → broader hardware support (MPS, etc.)

One question I had: do we know how significant the performance difference is in practice for MoE workloads?

If the gap is large, it might be worth keeping int as the default and falling back only when needed. But if the difference is small, the current approach (float for non-CUDA) seems safer.

Also wondering whether a "try int → fallback to float" approach could be acceptable if the failure is rare and predictable.

Yes, that's correct. However, the main issue here is more about hardware reality than a classic performance vs portability tradeoff.

Integer histograms for torch.histc are still poorly supported outside of CUDA. They tend to fail on MPS, CPU, and several other backends (APU, TPU, etc.), and broader support may take a long time to mature.

My current approach uses int only on CUDA (where we get the best performance) and float32 everywhere else. This gives strong performance for the majority of MoE users while ensuring reliable behavior across hardware.

Regarding the "try int first, then fallback to float" idea — I actually considered it as well. While it could offer better performance on unknown hardware, I’m concerned it might introduce subtle issues such as problems with torch.compile, Python exception handling interfering with the memory cache, or inconsistent states on backends like Metal (MPS). That said, I might be overthinking it, so I’d really appreciate the maintainers’ conservative feedback on this.

i wouldn't wanna penalize xpu and xla/tpu for example or any other custom torch backend that supports int histc
so i'm more inclined towards expanding the float fallback backends list rather than defaulting to float

i wouldn't wanna penalize xpu and xla/tpu for example or any other custom torch backend that supports int histc so i'm more inclined towards expanding the float fallback backends list rather than defaulting to float

Thank you for the response. I now see that this repo follows a performance-first, stability-second approach, which is a valid stance. I had initially gotten the priorities the other way around.
Upon reflection, your approach makes complete sense. Going with a float-compatible hardware list is smart because PyTorch and other serious AI backends are likely to expand their integer histogram support soon. This makes the current pattern nicely future-proof.
Sorry for dragging this out. I’ll keep my lessons in mind to improve my future PRs in the repo.

basically we've had this experts impl for a couple months now, and no reports from xpu/hpu. so i'd say they probably already support int, otherwise they'd be the first to report it (we have a couple contributors from intel).

basically we've had this experts impl for a couple months now, and no reports from xpu/hpu. so i'd say they probably already support int, otherwise they'd be the first to report it (we have a couple contributors from intel).

That's reassuring.

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

I couldn't reproduce the failure using this command in my local device: uv run pytest tests/cli/test_download.py::test_cli_download_trust_remote. Perhaps local cache and automated test cache is handled differently?

Hi @IlyasMoutawwakil @ArthurZucker @Cyrilvallez, could you check the latest change? CI was failing due to a distributed race condition with pytest-xdist.
When running with -n 8, workers have isolated sys.path. If Worker A downloads a model and initializes the cache, Worker B skips initialization because the files already exist, but then fails the import because its sys.path doesn't have the cache directory.
I've updated get_class_in_module to ensure every worker is initialized. I managed to reproduce this with a script—let me know if you'd prefer this in a separate PR!

@rigen1048 thanks for investigating that's very helpful ! let's keep it in a separate PR tho since it would require reviewing from our ci/testing mantainers.