Direction for extending MXGemmKernelChoice/MXLinearRecipeName for XPU support

[slabhs-aws]

Hi team,

TorchAO today provides 3 MXGemm kernel paths:

EMULATED
CUTLASS
CUBLAS

For users running MX workloads on XPUs, it’s unclear how future MX kernel selection is intended to evolve.

Questions:

1. Should XPU-based execution map into MXGemmKernelChoice, or is there a different path to onboard new XPU architectures?
2. Similarly, how do we scale MXLinearRecipeName

This context affects how upstream frameworks (e.g., torchtitan) gate MX execution.

Cross-referencing with the related torchtitan discussion: pytorch/torchtitan#2120

Thank you!

[danielvegamyhre]

Hi, thanks for your interest in this - a few questions for you:

For users running MX workloads on XPUs, it’s unclear how future MX kernel selection is intended to evolve.

1. Are you the user, or are you interested in developing MX + XPU support on behalf of users? If so, do you have a sense for how many users are doing this?
2. What XPU hardware specifically natively supports / accelerates operations with MX dtypes?

cc @vkuzo @drisspg who may be interested in this topic as well

[vkuzo]

@slabhs-aws we actually just deprecated MXGemmKernelChoice as we move to graduate the MX workflows out of prototype, it was a debug enum for early days. We are now using KernelPreference (

ao/torchao/quantization/quantize_/common/kernel_preference.py

Line 14 in 51fd90e

class KernelPreference(str, Enum):

) which is standard among torchao workflows. KernelPreference should be device agnostic, so if this level of control is enough for you then you could make XPU work for KernelPreference.AUTO|TORCH|EMULATED.

Let us know if this works for what you have in mind?

[HahTK]

@danielvegamyhre :

1. We are developing MX+XPU **
2. 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.

@vkuzo : Yes we are trying to avoid polluting the namespace with XPU specific names. Do you have an example for how mapping an XPU into KernelPreference.AUTO|TORCH|EMULATED ?

[vkuzo]

@HahTK , would tensors on your hardware have something like tensor.device == 'xpu:0'? Do you plan to have your mxfp8 gemms landed in PyTorch Core's torch.nn.functional.scaled_mm op? What kind of dependencies do you expect to require to be able to run your mx kernels on XPU? All of that will help answer ^.