acext User Guide (v1.5.1)

1.1.1 fpAintB GEMM

The fpA_intB_gemm.h header file implements fpAintB mixed-precision quantization. In this method, A is a floating-point type and B is an integer type. Three template parameters are used to define the specific types of A and B, and the quantization algorithm.

template <typename T, typename WeightType, cutlass::WeightOnlyQuantOp QuantOp>
class CutlassFpAIntBGemmRunner : public virtual CutlassFpAIntBGemmRunnerInterface
{
public:
    CutlassFpAIntBGemmRunner();
    ~CutlassFpAIntBGemmRunner();

    void gemm(const void* A, const void* B, const void* weight_scales, void* C, int m, int n, int k,
        tkc::CutlassGemmConfig gemmConfig, char* workspace_ptr, const size_t workspace_bytes,
        cudaStream_t stream) override;

    void gemm(const void* A, const void* B, const void* weight_scales, const void* weight_zero_points,
        const void* biases, void* C, int m, int n, int k, const int group_size, tkc::CutlassGemmConfig gemmConfig,
        char* workspace_ptr, const size_t workspace_bytes, cudaStream_t stream) override;

    // Returns desired workspace size in bytes.
    size_t getWorkspaceSize(const int m, const int n, const int k) override;

    tkc::CutlassGemmConfig getChosenConfig(const void*          A,
                                           const void* B,
                                           const void*          weight_scales,
                                           const void*          weight_zero_points,
                                           const void*          biases,
                                           void*                C,
                                           int               m,
                                           int               n,
                                           int               k,
                                           const int         group_size,
                                           char*             workspace_ptr,
                                           const size_t      workspace_bytes,
                                           cudaStream_t      stream) override;
};

template <typename ActType, WeightOnlyQuantType QType, cutlass::WeightOnlyQuantOp QuantOp>
std::shared_ptr<CutlassFpAIntBGemmRunnerInterface> createFpAIntBGemmRunner();

1.1.2 w8a8 GEMM

The int8_gemm.h header file provides a quantization algorithm where both A and B are integers. The template parameter T specifies the output data type.

template <typename T>
class CutlassInt8GemmRunner : public virtual CutlassInt8GemmRunnerInterface
{
public:
    CutlassInt8GemmRunner();
    ~CutlassInt8GemmRunner();

    void gemm(const void* A, const void* B, tk::QuantMode quantOption, const float* alphaCol, const float* alphaRow,
        void* C, int m, int n, int k, tkc::CutlassGemmConfig gemmConfig, char* workspacePtr,
        const size_t workspaceBytes, cudaStream_t stream) override;

    // Returns desired workspace size in bytes.
    size_t getWorkspaceSize(const int m, const int n, const int k) override;

    std::vector<tkc::CutlassGemmConfig> getConfigs() const override;

    tkc::CutlassGemmConfig getChosenConfig(const void* A, const void* B, tk::QuantMode quantOption,
        const float* alphaCol, const float* alphaRow, void* C, int m, int n, int k, char* workspacePtr,
        const size_t workspaceBytes, cudaStream_t stream);
};

1.1.3 MoE GEMM

The moe_kernel.h header file defines the Application Programming Interface (API) required to call MoE GEMM.

template<typename T,          /*The type used for activations/scales/compute*/
         typename WeightType, /* The type for the MoE weights */
         cutlass::WeightOnlyQuantOp QuantOp,
         typename Enable = void>
class CutlassMoeFCRunner: public CutlassMoeFCRunnerInterface {
public:
    CutlassMoeFCRunner()           = default;
    ~CutlassMoeFCRunner() override = default;
    //for rtp-llm
    void runMoe(const void* input_activations_void, const float* gating_output, const float *gating_output_with_bias,
        const void* fc1_expert_weights_void, const void* fc1_scales_void, const void* fc1_zeros_void, const void* fc1_expert_biases_void,
        ActivationType fc1_activation_type, const void* fc2_expert_weights_void, const void* fc2_scales_void, const void* fc2_zeros_void,
        const void* fc2_expert_biases_void, const int num_rows, const int hidden_size, const int inter_size,
        const int num_experts, const int k, const int group_size, char* workspace_ptr, void* final_output_void, void* fc2_result_void,
        const bool* finished, const int active_rows, void* expert_scales_void,
        int* expanded_source_row_to_expanded_dest_row, int* expert_for_source_row, MOEParallelismConfig parallelism_config,
        MOEExpertScaleNormalizationMode normalization_mode, cudaStream_t stream) override;
    // for vllm
    void fused_moe(const void* input_activations_void, const float* gating_output, const float* gating_output_with_bias,
                   const void* fc1_expert_weights_void, const void* fc1_scales_void, const void* fc1_zeros_void, const void* fc1_expert_biases_void,
                   void* a1_scale_void, ActivationType fc1_activation_type, const void* fc2_expert_weights_void, const void* fc2_scales_void,
                   const void* fc2_zeros_void, const void* fc2_expert_biases_void, void* a2_scale_void, const int num_tokens, const int hidden_size,
                   const int inter_size, const int num_experts, const int k, const int group_size, char* workspace_ptr, void* final_output_void,
                   void* fc2_result_void, const bool* finished, const int active_rows, const void* expert_scales_void,
                   int* expanded_source_row_to_expanded_dest_row, const int* expert_for_source_row, MOEParallelismConfig parallelism_config,
                   MOEExpertScaleNormalizationMode normalization_mode, cudaStream_t stream) override;

1.2.1 Int8 Weight-only (W8A16)

Int8 weight-only is the most basic quantization algorithm. It quantizes the weight matrix column-wise, which converts weights from a floating-point format to int8. Each column is assigned a scaling factor. If the original weights are in fp16 format, memory usage is reduced by half. Performance improves in memory-bound scenarios compared to native fp16 computation because fewer weight values need to be loaded.

Example code:

// Prepare data: act, weight, scales, and output
......

// Create FpAIntBGemmRunner with A=half, B=Int8, and PER_COLUMN_SCALE_ONLY quantization
auto runner = createFpAIntBGemmRunner<half, WeightOnlyQuantType::Int8b, cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>();

// Preprocess weights offline for mixed-precision GEMM
preprocess_weights_for_mixed_gemm(preprocessed_quantized_weight, row_major_quantized_weight, {k, n}, QuantType::INT8_WEIGHT_ONLY);

// Allocate workspace
char* ws_ptr = nullptr;
const int ws_bytes = runner->getWorkspaceSize(m, n, k);
cudaMalloc(&ws_ptr, ws_bytes);

// Select kernel configuration
gemmConfig = runner->getChosenConfig(p_act, p_weight, p_scales, nullptr, nullptr,
                                    p_out, m, n, k, k, reinterpret_cast<char*>(ws_ptr),
                                    ws_bytes, stream);

// Run quantization algorithm; store result in out
runner->gemm(p_act, p_weight, p_scales, p_out, m, n, k, gemmConfig, ws_ptr, ws_bytes, stream);

1.2.2 Int4 Weight-only (W4A16)

Int4 weight-only uses the same algorithm as Int8 weight-only but stores weights using only 4 bits. This method significantly reduces precision but also reduces memory usage to one-quarter of the original and further boosts performance.

Example code:

// Prepare data: act, weight, scales, and output
......

// Create FpAIntBGemmRunner with A=half, B=Int4, and PER_COLUMN_SCALE_ONLY quantization
auto runner = createFpAIntBGemmRunner<half, WeightOnlyQuantType::Int4b, cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>();

// Preprocess weights offline for mixed-precision GEMM
preprocess_weights_for_mixed_gemm(preprocessed_quantized_weight, row_major_quantized_weight, {k, n}, QuantType::PACKED_INT4_WEIGHT_ONLY);

// Allocate workspace
char* ws_ptr = nullptr;
const int ws_bytes = runner->getWorkspaceSize(m, n, k);
cudaMalloc(&ws_ptr, ws_bytes);

// Select kernel configuration
gemmConfig = runner->getChosenConfig(p_act, p_weight, p_scales, nullptr, nullptr,
                                    p_out, m, n, k, k, reinterpret_cast<char*>(ws_ptr),
                                    ws_bytes, stream);

// Run quantization algorithm; store result in out
runner->gemm(p_act, p_weight, p_scales, p_out, m, n, k, gemmConfig, ws_ptr, ws_bytes, stream);

1.2.3 Int4 Groupwise (W4A16)

Although Int4 weight-only significantly reduces memory and improves performance, its low precision of only 4 bits leads to large quantization errors compared to fp16. Groupwise quantization addresses this issue by dividing each column into groups of a fixed size, such as 64 or 128, and assigning one scale per group instead of one per column.

In practice, for algorithms such as AutoQuant (AWQ) or Generative Pre-trained Transformer Quantization (GPTQ), models require fine-tuning on a dataset to produce quantized int4 weights. Asymmetric quantization may also introduce zero-point offsets, also known as zeros and bias.

Example code:

// Prepare data: act, weight, scales, and output
......

// Create FpAIntBGemmRunner with A=half, B=Int8, and FINEGRAINED_SCALE_ONLY quantization
auto runner = createFpAIntBGemmRunner<half, WeightOnlyQuantType::Int8b, cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY>();

// Preprocess weights offline for mixed-precision GEMM
preprocess_weights_for_mixed_gemm(preprocessed_quantized_weight, row_major_quantized_weight, {k, n}, QuantType::PACKED_INT4_WEIGHT_ONLY);

// Allocate workspace
char* ws_ptr = nullptr;
const int ws_bytes = runner->getWorkspaceSize(m, n, k);
cudaMalloc(&ws_ptr, ws_bytes);

// Select kernel configuration
gemmConfig = runner->getChosenConfig(p_act, p_weight, p_scales, p_zeros, p_bias,
                                    p_out, m, n, k, groupsize, reinterpret_cast<char*>(ws_ptr),
                                    ws_bytes, stream);

// Run quantization algorithm; store result in out
runner->gemm(p_act, p_weight, p_scales, p_zeros, p_bias, p_out, m, n, k, groupsize, gemmConfig, ws_ptr, ws_bytes, stream);

1.2.4 SmoothQuant (W8A8)

fpAintB is a mixed-precision algorithm. Although the weights are integers, they must be dequantized to a floating-point format before matrix operations, which still rely on floating-point Tensor Cores. W8A8 uses integer weights and integer inputs, which enables integer Tensor Core operations for better compute throughput.

SmoothQuant is a quantization algorithm that uses the W8A8 kernel. https://github.com/mit-han-lab/smoothquant

W8A8 variants depend on whether per-token quantization or per-channel quantization is applied to A and B. Per-token quantization uses one set of quantization parameters per row and defaults to one set for the entire matrix. Per-channel quantization uses one set per column and defaults to one set for the entire matrix. This results in four variants: PerTensorPerTensor, PerTensorPerChannel, PerTokenPerTensor, and PerTokenPerChannel.

acext implements all four W8A8 quantization kernels. You can use them as follows:

// Choose one of the following W8A8 variants:
QuantMode quant_mode = QuantMode::fromDescription(false, false, false, false);  // A_PerTensor_W_PerTensor
QuantMode quant_mode = QuantMode::fromDescription(false, false, false, true); // A_PerTensor_W_PerChannel
QuantMode quant_mode = QuantMode::fromDescription(false, false, true, false); // A_PerToken_W_PerTensor
QuantMode quant_mode = QuantMode::fromDescription(false, false, true, true); // A_PerToken_W_PerChannel

// Prepare data: act, weight, scales, out
......

// Create W8A8 runner
auto runner = std::make_shared<acext::kernels::cutlass_kernels::CutlassInt8GemmRunner<half>>();

// Allocate workspace
int ws_bytes = runner->getWorkspaceSize(m, n, k);
char* ws_ptr = nullptr;
if (ws_bytes)
    cudaMalloc(&ws_ptr, ws_bytes);

// Select kernel configuration
const auto gemmConfig = runner->getChosenConfig(p_act, p_weight, quant_mode,
                    reinterpret_cast<const float*>(p_scales_channels),
                    reinterpret_cast<const float*>(p_scales_tokens),
                    p_out, m, n, k, ws_ptr, ws_bytes, stream);

// Execute kernel; store result in out
runner->gemm(p_act, p_weight, quant_mode,
                    reinterpret_cast<const float*>(p_scales_channels),
                    reinterpret_cast<const float*>(p_scales_tokens),
                    p_out, m, n, k, gemmConfig, ws_ptr, ws_bytes, stream);

2.1.1 AutoWOQ

To use quantization, you must convert the large model weights from fp16 or bf16 format to int8 format.

We use our AutoWOQ tool for this conversion. The components of the tool are:

AutoWOQForCausalLM serves as a unified entry point for all large models, similar to the Auto classes in Hugging Face. It accepts the model weight path and automatically detects the model type to initialize the appropriate WOQForCausalLM instance. All instances inherit from BaseWOQForCausalLM. The quantization logic resides in BaseWOQForCausalLM and primarily calls symmetric quantization algorithms that are implemented in acext.

You can use AutoWOQ to quantize a model as follows:

# Set quantization parameters: 8-bit or 4-bit
quantize_config = BaseQuantizeConfig(
    bits=8,  # Quantize model to 8-bit
)

# Load unquantized weights (Hugging Face format). By default, loads onto CPU.
model = AutoWOQForCausalLM.from_pretrained(model_path, quantize_config, trust_remote_code=args.trust_remote_code)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Quantize the model. Applies symmetric quantization to linear layers based on model type.
model.quantize()

# Save quantized model and tokenizer to specified location.
quant_path = args.output
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

2.1.2 Supporting WOQ Quantization

After you generate WOQ-quantized weights, you can add weight-only quantization to vLLM. When you use WOQ weights, you must explicitly specify `quantization="weightonly"` when you create the LLM object.

# Create an LLM.
llm = LLM(model=model_path, quantization="weightonly")

vLLM automatically creates quantized tensor parameters for linear layers and calls mixed-precision quantization at runtime.

As shown in the following figure, a Transformer block in vLLM contains green linear layers that use the WeightOnlyQuant algorithm.

These linear layers use WeightOnlyLinearMethod to initialize weights, load quantized weights, and invoke acext quantization algorithms at runtime. Because acext uses CUTLASS for kernel acceleration, the weights must be pre-reordered. Currently, this reordering occurs once at system startup and accounts for the number of GPUs.

class WeightOnlyLinearMethod(LinearMethodBase):
    # Set quantization configuration
    def __init__(self, quant_config: WeightOnlyConfig):
        self.quant_config = quant_config

    # Initialize weights: create int-type weight parameters and load WOQ-processed weights
    def create_weights(self, input_size_per_partition: int,
                       output_size_per_partition: int, input_size: int,
                       output_size: int,
                       params_dtype: torch.dtype) -> Dict[str, Any]:
        qweight = Parameter(
            torch.empty(
                input_size_per_partition,
                output_size_per_partition,
                device="cuda",
                dtype=torch.int8,
            ),
            requires_grad=False,
        )
        scales = Parameter(
            torch.empty(
                output_size_per_partition,
                device="cuda",
                dtype=params_dtype,
            ),
            requires_grad=False,
        )

        # Instantiate acext's WeightOnlyQuantMatmul class to call mixed-precision kernels later
        try:
            self.weightOnlyMatmul = ops.WeightOnlyQuantMatmul()
        except AttributeError as e:
            raise AttributeError(
                "WeightOnlyQuantMatmul kernels could not be imported. If you built vLLM "
                "from source, make sure BUILD_WEIGHTONLY_KERNELS=1 env var "
                "was set.") from e
        woq_state = WOQState.UNINITIALIZED
        ......

    # Apply weights at runtime using acext quantization algorithm
    def apply_weights(self,
                      weights: Dict[str, torch.Tensor],
                      x: torch.Tensor,
                      bias: Optional[torch.Tensor] = None) -> torch.Tensor:
        # Online weight preprocessing; done once at system startup
        if weights["woq_state"] == WOQState.UNINITIALIZED:
            qweight_cpu = weights["qweight"].cpu()
            qweight_cpu = preprocess_weights_for_mixed_gemm(qweight_cpu, torch.int8, False, False)
            weights["qweight"] = qweight_cpu.cuda()
            weights["woq_state"] = WOQState.READY
        qweight = weights["qweight"]
        scales = weights["scales"]
        return self.weightOnlyMatmul.weightonly_gemm(x, qweight, scales, bias)

2.1.3 Running WOQ Model Inference

To enable weight-only quantization, compile vLLM with BUILD_WEIGHTONLY_KERNELS=1.

# Compile vLLM
cd vllm
BUILD_WEIGHTONLY_KERNELS=1 python setup.py develop

# Run vLLM inference. Assume model is already quantized.
python examples/offline_inference.py --model /path/to/Llama-2-7b-hf-woq/ -q weightonly --trust-remote-code

Results:

2.2.1 Optimizing AWQ/GPTQ Inference

AWQ and GPTQ are two int4 quantization algorithms that are natively supported by vLLM.

https://docs.vllm.ai/en/latest/quantization/auto_awq.html

However, the current AWQ/GPTQ implementation in vLLM is inefficient. Although it reduces the memory footprint, its performance may be worse than fp16.

Using the mixed-precision quantization of acext, we can achieve more efficient kernels and higher performance than the native AWQ/GPTQ. Because acext has strict weight layout requirements, the original AWQ weights must be preprocessed before they can be used.

By adding the `weightonly_groupwise_gemm` backend of acext to vLLM, we can accelerate AWQ inference while retaining the original AWQ weights. Preprocessing occurs just before the first use. The workflow is as follows:

GPTQ follows the same workflow.

Currently, preprocessing runs online when vLLM starts, which causes latency. Support for offline preprocessing will be added in a future release.

2.2.2 Fast AWQ/GPTQ Inference

This optimization is under active development. By default, vLLM does not use acext-accelerated AWQ/GPTQ. To enable this feature, you can set the `USE_FAST_AWQ=1` or `USE_FAST_GPTQ=1` environment variable. If you do not set these variables, vLLM uses its native quantization.

To invoke the FastAWQ flow and accelerate AWQ quantization with acext, you can use the following command line.

USE_FAST_AWQ=1 python offline_inference.py --model /ppusw/share/LLM_Weights/Llama-2-7b-hf-awq -q awq --trust-remote-code

FastAWQ produces the same inference results as AWQ but runs faster.

2.3.1 Construct acext Backend

In acext_moe_func.cu, we added an acext backend that uses the AcextFusedMoe class to call the APIs of the acext library. We then registered it as a Python API through csrc/moe/moe_ops.h and csrc/torch_bindings.cpp.

2.3.2 Modify fused MoE Layer

In the fused_moe.py file, add a call to `acext` within the `fused_experts_impl` implementation. First, use ops.get_acext_fusedmoe_status() to check whether the current model configuration is suitable for the `acext` backend. If it is not suitable, the `triton` backend is called. If the return value is acext_fusedmoe_status = 0, use ops.acext_fusedmoe_warpper() to call the `acext` implementation of `fused_moe`.

2.3.3 Add W8A8Int8MoEMethod (for w8a8-int8-channel)

Add the CompressedTensorsW8A8Int8MoEMethod method to the `compressed_tensors_moe.py` file to load the weights of the w8a8-int8-channel model and execute the corresponding MoE layer.

2.4.1 Known Issues

2.5.1 Detailed Performance Data

2.6.1 Steps

# If your setuptools version is too low, upgrade it.
#pip3 install --upgrade pip setuptools
# Due to the following issue, you must use setuptools version 69.5.1.
#https://github.com/vllm-project/vllm/issues/4961

git clone https://github.com/EleutherAI/lm-evaluation-harness
# If you have issues running the MMLU test, it may be version-related. Try switching to this tested commit:
# git reset --hard e5e5ee0cb629c9c88165292d1b4bf34623392d33
cd lm-evaluation-harness
pip install -e .

# Use a Hugging Face proxy.
export HF_ENDPOINT=https://hf-mirror.com

# Use lm_eval for MMLU inference.
lm_eval --model vllm --model_args pretrained=/ppusw/datasets/huggingface/Llama-2-7b-hf --tasks mmlu --batch_size auto --num_fewshot 5

# If you still have issues downloading the dataset after using the proxy, it is because the following file hardcodes the Hugging Face download URL.
#/root/.cache/huggingface/modules/datasets_modules/datasets/hails--mmlu_no_train/b7d5f7f21003c21be079f11495ee011332b980bd1cd7e70cc740e8c079e5bda2/mmlu_no_train.py
# Modify the URL in this file to point to hf-mirror.

2.6.2 Inference results

You can set `num-fewshot` to control the few-shot parameter. A higher value results in a longer runtime but may also yield a higher score.

If you do not set `num_fewshot`, the run is faster (minutes), and the result is 41.3.

If you set `num_fewshot=5`, the run is slower (hours), but the score is higher, at 45.9.

This result is consistent with the results on the website.

2.6.3 MMLU accuracy for different quantization methods

The MMLU score from the fp16 inference on the unquantized model serves as the baseline. You can set different quantization methods during inference to obtain the MMLU scores for the quantized models. You can then compare these scores with the baseline to determine the MMLU quantization loss.

The following table shows the MMLU scores (few-shot 0) and quantization loss relative to fp16 for different models and quantization methods.

Gray indicates quantization methods that are not supported in the current version (PPU Release 1.3). Red indicates methods that have buggy inference results.

3.2.1 Run Unit Tests

cd build

# Unit tests for a16w8/a16w4 quantized GEMM
python ../tests/test_woq_gemm.py

# Unit tests for a8w8 quantized GEMM
python ../tests/th_int8_gemm.py

# Unit tests for MoE
python ../tests/th_moe.py

3.2.2 Run Performance Tests

cd build
# Run fpAintB benchmark separately. Default MNK is (1, 5120, 5120). Specify M N K in command-line arguments.
./benchmark/benchmark_fpAintB

# Run w8a8 benchmark separately. Default MNK is (1, 5120, 5120). Specify M N K in command-line arguments.
./benchmark/benchmark_w8a8

# Run benchmarks for all M N K combinations and generate a unified CSV report.
cd benchmark
python run_benchmark.py