acext User Guide (v1.6.1)

1.2.1 fpAintB GEMM

The fpA_intB_gemm.h header file implements mixed-precision quantization where matrix A is floating-point and matrix B is an integer. Three template parameters allow you to specify the exact types for A and B, and the quantization algorithm to use.

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.2.2 w8a8 GEMM

The int8_gemm.h header file implements quantization where both matrices 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.2.3 MoE GEMM

The moe_kernel.h header file provides the APIs required to call MoE GEMM kernels.

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.3.1 Int8 Weight-only (W8A16)

Int8 weight-only is the simplest quantization method. It quantizes weights column-wise into int8 values, with one scale factor per column. If the original weights are fp16, this method reduces memory usage by half. Because less data is loaded from memory, performance improves in memory-bound scenarios compared to native fp16 computation.

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 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);

// Choose optimal kernel config
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 quantized kernel; result stored in out
runner->gemm(p_act, p_weight, p_scales, p_out, m, n, k, gemmConfig, ws_ptr, ws_bytes, stream);

1.3.2 Int4 Weight-only (W4A16)

Int4 weight-only uses the same principle as Int8 weight-only but stores weights in a 4-bit format. This further reduces memory usage to one-quarter of the fp16 memory usage and boosts performance, but at the cost of lower accuracy.

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 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);

// Choose optimal kernel config
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 quantized kernel; result stored in out
runner->gemm(p_act, p_weight, p_scales, p_out, m, n, k, gemmConfig, ws_ptr, ws_bytes, stream);

1.3.3 Int4 Groupwise (W4A16)

While Int4 weight-only saves memory and speeds up computation, its low precision causes large quantization errors. Groupwise quantization solves this issue by dividing each column into groups (for example, a group size of 64 or 128) and assigning one scale factor per group instead of per column.

In practice (for example, with AWQ or GPTQ), models are fine-tuned on datasets to produce quantized int4 weights. Asymmetric quantization may also introduce zero-point offsets (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 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);

// Choose optimal kernel config
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 quantized kernel; result stored 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.3.4 SmoothQuant (W8A8)

fpAintB uses mixed-precision quantization where weights are integers, but they are dequantized to float before computation and then passed to the TensorCore. In contrast, W8A8 keeps both weights and activations as integers, which enables full integer TensorCore execution and delivers better compute throughput.

SmoothQuant is the quantization algorithm that uses W8A8 kernels. For more information, see https://github.com/mit-han-lab/smoothquant.

W8A8 can be divided into four variants based on whether matrix A uses PerToken quantization (each row shares a set of quantization parameters) and matrix B uses PerChannel quantization (each column shares a set), as opposed to the default PerTensor method where the entire matrix shares one set of parameters. The four variants are PerTensorPerTensor, PerTensorPerChannel, PerTokenPerTensor, and PerTokenPerChannel.

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

// Choose one of the four 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, and output
......

// 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);

// Choose kernel config
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; result stored 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 apply quantization, you must convert the original fp16 or bf16 model weights to the int8 format.

We provide AutoWOQ for this conversion. Its components are shown below:

AutoWOQForCausalLM serves as a unified entry point for all large language models, similar to the Auto classes in Hugging Face. It accepts the model weight path and automatically detects the model type to instantiate the correct WOQForCausalLM subclass. All subclasses inherit from BaseWOQForCausalLM, which implements the core quantization logic, including symmetric quantization using acext.

You can use AutoWOQ as follows:

# Configure quantization: 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 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 Enable WOQ Quantization

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

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

vLLM will then automatically create quantized tensor parameters for linear layers and invoke mixed-precision kernels at runtime.

As shown in the following figure, a Transformer block in vLLM contains green linear layers that use weight-only quantization.

These linear layers are initialized using WeightOnlyLinearMethod. This method loads quantized weights and invokes the acext quantization kernels at runtime. Because acext relies on CUTLASS for kernel acceleration, weights must be pre-reordered. Currently, this reordering happens 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 int8 weight tensors and load quantized weights from WOQ
    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 was set.") from e
        woq_state = WOQState.UNINITIALIZED
        ......

    # Apply weights at runtime using acext's quantization kernel
    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 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 Run WOQ Models

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

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

# Run offline inference with a quantized model
python examples/offline_inference.py --model /path/to/Llama-2-7b-hf-woq/ -q weightonly --trust-remote-code

Output:

2.2.1 Optimize AWQ/GPTQ Inference

AWQ and GPTQ are natively supported in vLLM and both use int4 quantization.

For more information, see https://docs.vllm.ai/en/latest/quantization/auto_awq.html.

However, the current AWQ/GPTQ implementation in vLLM only reduces the memory footprint and does not improve the inference speed compared to fp16.

We improved AWQ/GPTQ inference using the mixed-precision kernels of acext. These kernels deliver measurable speedups. Because acext requires specific weight layouts, original AWQ weights must be preprocessed before you can call acext.

By adding acext weightonly_groupwise_gemm support to vLLM, you can accelerate AWQ inference while retaining the original AWQ weights. The preprocessing runs once before the weights are first used.

The same flow also applies to GPTQ quantization.

Similar to WOQ, weight processing is currently performed online, which causes a long startup delay for vLLM. An offline processing method will be supported in a future release.

2.2.2 Fast AWQ/GPTQ Inference

This optimization is still under active development. It is disabled by default in vLLM. To enable it, set the `USE_FAST_AWQ=1` or `USE_FAST_GPTQ=1` environment variable.

You can use the following command line to invoke the FastAWQ flow, which accelerates AWQ quantization with acext.

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 outputs that are identical to standard AWQ, but with faster performance.

2.3.1 Build acext Backend

In `acext_moe_func.cu`, we added the acext backend using the AcextFusedMoe class to call acext APIs. We then registered these as Python APIs using csrc/moe/moe_ops.h and csrc/torch_bindings.cpp.

2.3.2 Modify fused MoE Layer

In the fused_moe.py file, you can add a call to `acext` within the implementation of `fused_experts_impl`. First, you can 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, you can use ops.acext_fusedmoe_wrapper() to call the `acext` implementation of `fused_moe`.

2.3.3 Add W8A8Int8MoEMethod (for w8a8-int8-channel)

In `compressed_tensors_moe.py`, we added the CompressedTensorsW8A8Int8MoEMethod class to load w8a8-int8-channel weights and run the corresponding MoE layer.

2.4.1 Known Issues

2.5.1 Detailed Performance Data

2.6.1 Running steps

# If your setuptools version is too low, upgrade it first.
# pip3 install --upgrade pip setuptools
# For the following reason, we need to restrict setuptools to version 69.5.1:
# https://github.com/vllm-project/vllm/issues/4961

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

# Use the 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 increases the runtime but may yield a higher score.

Without `num_fewshot`, the test runs quickly (in minutes) and scores 41.3.

With `num_fewshot=5`, the test runs slowly (in hours) but scores higher at 45.9.

This matches the result shown on the website.

2.6.3 MMLU accuracy for different quantization methods

We use the MMLU score from fp16 inference as the baseline. By applying different quantization methods during inference, we can obtain post-quantization MMLU scores and compare them to the baseline to calculate the MMLU loss.

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

Gray cells indicate methods that are not supported in the current version (PPU Release 1.3). Red cells indicate methods with inference bugs.

3.2.1 Run unit tests

cd build

# Unit test for a16w8/a16w4 quant gemm
python ../tests/test_woq_gemm.py

# Unit test for a8w8 quant gemm
python ../tests/th_int8_gemm.py

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

3.2.2 Run performance tests

cd build
# Run FpAintB benchmark. Default M, N, K is (1, 5120, 5120). You can specify M, N, K as command-line arguments.
./benchmark/benchmark_fpAintB

# Run w8a8 benchmark. Default M, N, K is (1, 5120, 5120). You can specify M, N, K as command-line arguments.
./benchmark/benchmark_w8a8

# Iterate through all M, N, K values to run the benchmark and generate a unified CSV report.
cd benchmark
python run_benchmark.py