Alternative approach to support torch.compile

[Giuseppe5]

This works by assuming that most of the quantization process has already taken place, and it is no longer needed to propagate QuantTensors.

Typical usage:

with inference_mode(model):
     model(input)
     compile_model = torch.compile(model, fullgraph=True)
     # Rest of the computation under compile goes here. Once out of the context manager, the original model must be used.

[nickfraser]

This needs tests!

[nickfraser]

It looks like these won't work with Groupwise quantization, correct? So inference_mode + MX won't work?

[Giuseppe5]

You're right, I forgot to add the export handler for MX INT and MX Float

[Giuseppe5]

Postponed to another update

[nickfraser]

LGTM!