Allow quantized linear registration in a different file

[jerryzh168]

Summary:

Previously there was some ordering that we need to maintain for quantized linear dispatch table in AffineQuantizedTensor, the reason is there is a fallback entry that dequantizes the input:

ao/torchao/dtypes/affine_quantized_tensor.py

Line 1195 in ba2d3b1

(_linear_quantized_act_fallback_check, _linear_quantized_act_fallback_impl),

so the dispatches with two inputs quantized (static or dynamic quantization) must come before this entry and dispatches with weight only quantization, however the fallback is not really used/needed in practice, since people typically just want to call into a very specific kernel.

From offline discussions with @drisspg and @HDCharles, it might be useful to have a "quantized_linear_impl" for LayoutType, this allows people to specify and check which quantized_linear_impl they want to use to make sure they can call into the specific kernel, when this field is set, we'll not run the fallback path for quantized linear either (dequantize all activation and weight tensors and run the floating point linear op)
I think this can be added for a specific layout type if people want to and we don't have to enforce this in the base LayoutType, otherwise we'd have to specify this for all LayoutType which seems a bit cumbersome. The contract here is when people want to skip the fallback path, they can do:

class MyLayoutType:
    quantized_linear_impl = "some_impl"
    ...

and have the dispatch_condition function check for this:

def dispatch_condition(...):
    return weight_tensor.layout_type.quantized_linear_impl == "some_impl" and ...

if there are some issues with the implementation, it will not call the fallback in this case

Test Plan:
python test/dtypes/test_affine_quantized.py -k test_register_new_dispatch

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/783

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit b7c1512 with merge base e2dad4a ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vayuda]

When should this be used?

[jerryzh168]

this is used to remove some registration, most of the time probably don't need to use it, I'm using this just to make sure test pass since the registration is global