Recommendation for enabling MXLinear on XPU backends

[slabhs-aws]

Hi team,

Currently MXLinear runs only when CUDA capability ≥ SM100, due to:

assert has_cuda_capability(10, 0)

This makes MXLinear inaccessible on XPUs or other non-SM100 accelerator paths.

What would be the recommended strategy for enabling MXLinear usage on XPU-based backends?

This will help ensure that frameworks relying on MXLinear can onboard non-CUDA backends without diverging from intended usage patterns.

Related question raised in TorchAO: pytorch/ao#3457

Thanks!

[tianyu-l]

Is MXLinear already supported on XPU today? If so we could start with adding another utils function similar to https://github.com/pytorch/torchtitan/blob/main/torchtitan/tools/utils.py#L20 to unblock.

[HahTK]

@tianyu-l : Trainium 3 has support for MXFP8 and MXFP4 (see recent announcement : https://aws.amazon.com/ai/machine-learning/trainium/).

Edit : Technically not an Intel XPU but backend extensions could be generic.

For internal prototypes we have already implemented a similar function. We were never blocked and are not concerned about short term progress. In fact, the question here is the result of our internal review of the prototype

IMHO, we would generally want to avoid such a list of "has_X_capability" for every XPU which is check in multiple places. That seems error prone and not scalable. The question is really : what is the long term design direction here? Is this really it or are we already planning something else?

It would seem like one possible solution is to create a registry for different backends to register ops, customizations and what not. Does the TorchAO team have something else in mind?

[HahTK]

Zooming out it seems assert has_cuda_capability(10, 0) is used as a check for mx support in cuda.One would imagine that if we go down this path we would get a set of custom has_x_capabilities(major_ver, minor_ver) everywhere. We are aiming for device agnostic implementation ala torch.accelerator which we believe will be cleaner.

Example : torch.accelerator.current_accelerator().has_mx_capability()
This would of course require accelerators/devices to implement a has_mx_capability Or perhaps a has_capability(some_capability)

@tianyu-l @vkuzo @danielvegamyhre : thoughts on doing it this way instead?

[slabhs-aws]

To be precise, the proposal is torch.accelerator.has_mx_capability()

[tianyu-l]

To be precise, the proposal is torch.accelerator.has_mx_capability()

@albanD does this sound reasonable to you?

[albanD]

Having an option to check for mx capability sounds fair yes. We have other similar features in accelerator already and we can continue adding to it for capabilities that are broadly used. We can follow up at pytorch/pytorch#143887 for this.

That being said, we obviously don't want to have too narrow of a flag there, as each new flag needs logic across all backends.
The criterion there will be based on where this capability check is needed and how many of our accelerator support that feature.