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 ^.

[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 ?

For example, if you integrate gemm on your hardware into torch._scaled_mm, then your logic could go into KernelPreference.TORCH (for native pytorch supported ops) and be gated by something like x.device == XPU.

If you don't plan to integrate gemm on your hardware into torch._scaled_mm, we should align on a different extension point, see my q's above.

[HahTK]

Yes we are looking for device agnostic implementations. Wouldn't gating KernelPreference with x.device == my_device defeat that purpose? We'd like to abstract device out of torchao as well.

Here is an example of what I mean on Titan : pytorch/torchtitan#2120 (reply in thread)
What would be the equivalent here?

Context : We are evaluating an out of tree backend implementation envisioned by PyTorch's PrivateUse1 and would implement lowerings for torch._scaled_mm. We are PR-ing a new quantize aten to PyTorch.

What would be the equivalent to that in this case?
If we just did kernel registration for new backend (https://docs.pytorch.org/tutorials/advanced/privateuseone.html#register-kernels-for-the-new-backend)

[vkuzo]

Yes we are looking for device agnostic implementations. Wouldn't gating KernelPreference with x.device == my_device defeat that purpose? We'd like to abstract device out of torchao as well.

Sorry, I was not clear. I meant "gated by device XPU internally in scaled_mm". The need of gating by device in torchao would depend on how the XPU kernels are integrated with PyTorch core - if the CUDA and XPU functions are aligned, we might not need to gate.

We are evaluating an out of tree backend implementation envisioned by PyTorch's PrivateUse1 and would implement lowerings for torch._scaled_mm.

Makes sense! My view of how that would work with torchao's KernelPreference:

1. You can use KernelPreference.TORCH for the main path, and we should make sure that KernelPreference.AUTO aliases to KernelPreference.TORCH in your case.
2. If you register all the aten op lowerings correctly for PrivateUse1, it should "just work", as KernelPreference.TORCH path only calls aten ops.
3. If we land a new quantize op to PyTorch Core, torchao can onboard to it as needed.

Lmk your thoughts?

[slabhs-aws]

@vkuzo makes sense!
For #3, does it make sense to implement it like:

def to_mx(
    data_hp: torch.Tensor,
    elem_dtype: Union[torch.dtype, str],
    block_size: int,
    scaling_mode: ScaleCalculationMode = ScaleCalculationMode.FLOOR,
    is_swizzled_scales: bool = False,
    kernel_preference: KernelPreference = KernelPreference.EMULATED,
):
    """
    Takes a high precision tensor and converts to MX scale and raw data, in
    naive layout (scale and raw data are separate tensors).
    """
    assert data_hp.dtype in (
        torch.bfloat16,
        torch.float,
    ), f"{data_hp.dtype} is not supported yet"
    # TODO(future PR): consider supporting padding
    assert data_hp.shape[-1] % block_size == 0, (
        f"the last dimension of shape {data_hp.shape} must be divisible by block_size {block_size}"
    )
    assert data_hp.is_contiguous(), "unsupported"
    assert elem_dtype in SUPPORTED_ELEM_DTYPES, "unsupported"

    if kernel_preference == KernelPreference.TORCH :
        return torch.quantize_mx(data_hp, block_size, .....)

    orig_shape = data_hp.shape
    data_hp = data_hp.reshape(
        *orig_shape[:-1], orig_shape[-1] // block_size, block_size
    )

    # find max value of the data
    # Note: this only implements the `minimally supported` version of
    # https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
    # section 6.3.
    max_abs = torch.amax(torch.abs(data_hp), -1).unsqueeze(-1)
........

[vkuzo]

there is good discussion in pytorch/rfcs#86 about this

[slabhs-aws]

Yup! Once our RFC goes through and the aten op is added, can we consume it in TorchAO as above (see torch.quantize_mx(data_hp, block_size, .....)