Add semi-structured sparse + dynamic int8 subclasses

[jcaip]

This PR adds in int8 dynamic quantization + semi-structured sparsity support into torchao.

This is implemented by extending the existing quantization subclasses to use sparse ops.
Ideally we would be able to compose subclasses, and call to_sparse_semi_structured from inside the quantization subclass, but ATM nested subclass tracing does not work with torch.compile and for stuff like fusing scales into the sparse multiply you would probably want to implement it like this anyways.

In particular, this PR adds in two new subclasses:

Int8DynamicallyQuantizedSemiStructuredSparseLinearWeight
Int8DynamicallyQuantized24CusparseltLinearFuseMulWeight

For the cuSPARSELt subclasse, I can extend Int8DynamicallyQuantizedWeightBase by storing the compressed representation in W_int_repr.

FuseMulWeight will fuse one of the multiplies for the dequant into the cuSPARSELt matmul op. However cuSPARSELt expects this in a float32 format, so this eats into our previous speedups since we're now passing this as a bfloat16 tensor.

However for the general subclass, I need to extend QuantizeWeightBase, because I need to pass two tensors (packed and meta) for the CUTLASS sparse mm op. This relies on to_sparse_semi_structured to decide between CUTLASS and cuSPARSELt, which is the right choice for UI but makes benchmarking between them kind of difficult, since it's a class var that decides which backend gets used. Maybe we should add a flag to to_sparse_semi_structured because you might mix between cutlass and cusparselt.

I've also added a benchmarking script for SAM. I don't know how we plan on handling dependencies in torchao, but let me know if there's a better place for that.

On batch size 32, I see a 1.16x speedup over bfloat16 torch.compile baseline, from 21.96 -> 25.54 img/s.

Other benchmarks (BS=16)

[HDCharles]

can you move the benchmark_sam and other .py files to one of the other directories? Maybe make a torch/benchmarks dir?

[cpuhrsch]

Is it currently possible to replace this with

from torch.sparse import to_sparse_semi_structured
int_data = to_sparse_semi_structured(int_data)

[jcaip]

Leaving this one here b/c it's the cuSPARSELt fuse mul special one, but I have changed the subclass to be backend agnostic (Int8DynamicallyQuantizedSemiStructuredSparseLinearWeight)

side note - do you care much about naming convention? This name is so long I kind of want to change it to something simpler like QuantizedSemiSparseLinearWeight

[jerryzh168]

is sparsity also implemented with tensor subclass? I thought we should be able to compose them in some way?

[jcaip]

We don't have nested subclassing support currently for tracing, so we can't compose them currently :( hence why we're landing in prototype.

I can tag you in the issue i'll make to raise this for core.

[cpuhrsch]

Approving for prototype. Thanks for sending this :D

[msaroufim]

@jcaip just saw the CI failure issue is our CI GPU is too old so updating it now to A10G which should work for your code so after I merge this #176 make sure to rebase to main