Implementation:Deepspeedai DeepSpeed Quantization Utils

Knowledge Sources	DeepSpeed
Domains	Quantization, Model_Compression, CUDA_Kernels
Last Updated	2026-02-09 00:00 GMT

Overview

Device-side quantization framework providing templated 4-bit and 8-bit symmetric/asymmetric quantization with group-wise parameter computation.

Description

The quantization_utils.h header implements a comprehensive quantization system for reducing FP16 activations to INT4/INT8 representations. It provides a template-based design with Params<qType, numBits> class handling quantization parameter calculation (scale and optional offset) and value transformation, and GroupStats<qType> class tracking running statistics (min/max or abs-max) across quantization groups. The implementation supports both symmetric quantization (scale only, centered at zero) and asymmetric quantization (scale + offset for arbitrary ranges). Quantization occurs in groups rather than per-tensor to preserve granular accuracy, with device functions performing serial reductions to compute group statistics before applying quantization. The framework efficiently packs 4-bit values two per byte using PackedInt4 structure and provides local_array functions that orchestrate the complete quantization pipeline from statistics gathering to parameter computation to value quantization.

Usage

Use these utilities when implementing activation quantization for inference or training, particularly for communication compression in distributed training or model serving. The group-wise quantization provides better accuracy than per-tensor while maintaining computational efficiency through warp-level reductions.

Code Reference

Source Location

• Repository: DeepSpeed
• File: csrc/includes/quantization_utils.h

Signature

namespace quantize {
    enum class Type { Symmetric, Asymmetric };

    template <Type qType, int numBits>
    class Params {
    public:
        DS_D_INLINE int8_t quantize(__half val);

        template <typename T>
        DS_D_INLINE T dequantize(int8_t val);

        DS_D_INLINE void store(float* params, int group_index);
        DS_D_INLINE Params(const float* params, int group_index);
    };

    template <Type qType>
    class GroupStats {
    public:
        DS_D_INLINE void update(__half2 val);

        template <int numBits, int threads_per_group>
        DS_D_INLINE Params<qType, numBits> get_params(
            cg::thread_block& tb, cg::thread_block_tile<hw_warp_size>& warp);
    };

    // Full quantization pipeline
    template <Type qType, int numBits, int numChunks,
              int threads_per_group = 256, int max_threads = 256>
    __device__ void local_array(__half2* local_buffer,
                                float* global_params,
                                int8_t* output_data,
                                const int& elems_per_group,
                                const int& groups);
}

Import

#include "csrc/includes/quantization_utils.h"

I/O Contract

Input	Type	Description
local_buffer	__half2*	FP16 data to quantize in registers/shared mem
elems_per_group	int	Number of elements per quantization group
groups	int	Total number of groups
numBits	int (template)	4 or 8 bits per quantized value
qType	Type (template)	Symmetric or Asymmetric quantization

Output	Type	Description
output_data	int8_t*	Packed quantized values in global memory
global_params	float*	Quantization parameters (scale, offset)

Usage Examples

8-bit Symmetric Quantization:

__global__ void quantize_activations_int8(const __half* input,
                                          int8_t* output,
                                          float* params,
                                          int total_elems,
                                          int group_size) {
    constexpr int numChunks = 8;  // Process 8×16 bytes per thread
    __half2 local_buffer[numChunks * 8];  // 16 bytes → 8 half2

    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int group_id = blockIdx.y;
    int offset = group_id * group_size + tid * numChunks * 16;

    // Load data into local buffer
    for (int i = 0; i < numChunks; i++) {
        mem_access::load_global<16>(&local_buffer[i * 8],
                                    input + offset + i * blockDim.x * 16);
    }

    int groups = total_elems / group_size;
    quantize::local_array<quantize::Type::Symmetric, 8, numChunks>(
        local_buffer, params, output, group_size, groups);
}

4-bit Asymmetric Quantization:

__global__ void quantize_weights_int4(const __half* weights,
                                      int8_t* quantized,
                                      float* scales_offsets,
                                      int rows, int cols) {
    constexpr int group_size = 128;
    constexpr int numChunks = 4;

    __half2 buffer[numChunks * 8];
    int row = blockIdx.x;
    int col_base = threadIdx.x * numChunks * 16;

    // Load weight group
    for (int i = 0; i < numChunks; i++) {
        mem_access::load_global<16>(&buffer[i * 8],
                                    weights + row * cols + col_base + i * 128);
    }

    // Quantize with asymmetric mode (preserves full range)
    quantize::local_array<quantize::Type::Asymmetric, 4, numChunks>(
        buffer, scales_offsets, quantized, group_size, rows);
}

Manual Parameter Computation:

__device__ void manual_quantize_example(__half* data, int8_t* output, int n) {
    cg::thread_block tb = cg::this_thread_block();
    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(tb);

    // Compute statistics
    quantize::GroupStats<quantize::Type::Symmetric> stats;
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        __half2 val = __halves2half2(data[i], data[i+1]);
        stats.update(val);
    }

    // Get quantization parameters via warp reduction
    auto params = stats.get_params<8, 256>(tb, warp);

    // Quantize values
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        output[i] = params.quantize(data[i]);
    }
}

Gradient Communication Compression:

// Compress gradients before all-reduce
void compress_gradients(torch::Tensor& grads,
                       torch::Tensor& quantized,
                       torch::Tensor& params) {
    int total_elems = grads.numel();
    int group_size = 512;  // 512 elements per group
    int groups = (total_elems + group_size - 1) / group_size;

    dim3 block(256);
    dim3 grid((groups + 255) / 256);

    quantize_kernel<quantize::Type::Symmetric, 8><<<grid, block>>>(
        (__half*)grads.data_ptr(),
        (int8_t*)quantized.data_ptr(),
        (float*)params.data_ptr(),
        total_elems, group_size);

    // All-reduce quantized data (8× smaller)
    // Dequantize on receive
}

Related Pages

• Dequantization Utils - Inverse operation for decompression
• Quantization Binding - PyTorch interface for quantization
• Conversion Utils - Type conversions used internally
• Reduction Utils - Reduction primitives for statistics