[2/N] Rs/vllm quantization - Refactor refactor to support non-uniform via config

[robertgshaw2-redhat]

Refactored to support nonuniform quantization by adding a new layer of Abstraction.

Now, SmoothQuantLinearMethod can hold a SmoothQuantFormat, which implements the details of how to do quant and dequant operations. There are two SmoothQuantFormat classes:

• SmoothQuantDynamicPerToken
• SmoothQuantStaticPerTensor

We have the following lifecycle:

• LinearMethod is created during get_model, has access to QuantizationConfig
• Layer is initialized and passed a LinearMethod
• Layer calls LinearMethod.create_weights, which creates a dictionary of weights and metadata
• Layer calls LinearMethod.apply_weights during inference, passing the dictionary created during create_weights

This PR modifies the LinearMethod.create_weights API to receive a layer_name as argument. The LinearMethod then looks in the config to determine which SmoothQuantFormat to use for the layer with layer_name

• As a result, the LinearMethod is responsible for parsing the config from disk and making decisions about what the inference format should look like. In this specific case, since the SmoothQuantConfig is not very good, we just match on the suffix qkv to determine what each layer should use --> but for SparseMLConfig, we could use a similar structure

In this PR, the SmoothQuantFormat is passed in the dictionary returned by create_weights and then is used by apply_weights

In Summary

I think this is a good overall structure because it:

• (a) allows us to make minimal changes to the existing models
• (b) allows us to make no changes to the model loading lifecycle (i.e. config / constructor / linear method) ** critically requires having one LinearMethod that propagates through the whole model
• (c) encapsulates the nonuniform logic into the LinearMethod, allowing us to have a clean interface into

For SparseML Models

We could imagine the following architecture:

Config

Config is responsible for:

• loading config from disk
• mapping layer_names --> SparseMLFormat

class SparseMLConfig
    def from_dict()
    
    def get_layer_format(layer_name):
         return SparseMLFormat

LinearMethod

Config is responsible for:

• interface between layers and kernels (so LinearMethod is what is used by the model)

class SparseMLLinearMethod:
    def __init__(self, sparseml_config)
          self.sparseml_config = sparseml_config
          
    def create_weights(layer_name, ...):
          # this, e.g. is where nonuniform might be supported
          format = self.sparseml_config.get_layer_format(layer_name)
          
          weights = format.get_weights()
          weights["format"] = format
          
          return weights
     
     # wrapper around the SparseML format
     def apply_weights(x, weights, ...)
           format = weights["format"]
           weights = weights["weights"]
           
           return format.apply_weights(x, weights)

SparseMLFormat

Format is responsible for:

• actual weight creation and forward

class SparseMLLinearMethod:
    def __init__(self, sparseml_config)
          self.sparseml_config = sparseml_config
          
    def get_weights(sizes):
         # returns dictionary , e.g.
         return {
             "weights": x
             "scales": y
         }
     
     def apply_weights(weights, x):
         # calls cuda kernel 
         return output

Sample Formats:
- W8A8DynamicPerToken
- SparseW8A8StaticPerTensorAsymmetric
- W4A8DynamicPerToken
- ...

[robertgshaw2-redhat]

This should be moved to SmoothQuantConfig I think

[robertgshaw2-redhat]

Or at least the core logic should. This api probably still needs to exist, but this logic fits better conceptually with the spirit of the QuantizationConfig

[varun-sundar-rabindranath]

LGTM 👍