Skip to content

llmcompressor.utils.triton_utils

JIT-inline primitives shared by observer Triton kernels.

Functions:

calculate_candidate_scale 

calculate_candidate_scale(
    min_value,
    max_value,
    p,
    q_min,
    q_max,
    SCALE_ROUND_TYPE: constexpr,
    QUANT_TYPE: constexpr,
    SYMMETRIC: constexpr,
)

Derive and round a candidate scale from its original min/max range.

Source code in src/llmcompressor/utils/triton_utils.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
@triton.jit
def calculate_candidate_scale(
    min_value,
    max_value,
    p,
    q_min,
    q_max,
    SCALE_ROUND_TYPE: tl.constexpr,
    QUANT_TYPE: tl.constexpr,
    SYMMETRIC: tl.constexpr,
):
    """Derive and round a candidate scale from its original min/max range."""
    scaled_min = min_value * p
    scaled_max = max_value * p
    if SCALE_ROUND_TYPE == 4:
        # MX/E8M0 operates on the group absmax before element-format scaling.
        return round_candidate_scale(
            tl.maximum(tl.abs(scaled_min), tl.abs(scaled_max)),
            SCALE_ROUND_TYPE=SCALE_ROUND_TYPE,
            QUANT_TYPE=QUANT_TYPE,
        )
    if SYMMETRIC:
        scale = tl.maximum(tl.abs(scaled_min), tl.abs(scaled_max)) / (
            (q_max - q_min) * 0.5
        )
    else:
        scale = (scaled_max - scaled_min) / (q_max - q_min)
    return round_candidate_scale(
        scale, SCALE_ROUND_TYPE=SCALE_ROUND_TYPE, QUANT_TYPE=QUANT_TYPE
    )

quantize_dequantize 

quantize_dequantize(
    values,
    scale,
    zero_point,
    q_min,
    q_max,
    QUANT_TYPE: constexpr,
    NUM_BITS: constexpr,
    HAS_ZP: constexpr,
    COMPUTE_DTYPE: constexpr,
)

GPTQ-style QDQ shared by MSE and future observer kernels.

Source code in src/llmcompressor/utils/triton_utils.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
@triton.jit
def quantize_dequantize(
    values,
    scale,
    zero_point,
    q_min,
    q_max,
    QUANT_TYPE: tl.constexpr,
    NUM_BITS: tl.constexpr,
    HAS_ZP: tl.constexpr,
    COMPUTE_DTYPE: tl.constexpr,
):
    """GPTQ-style QDQ shared by MSE and future observer kernels."""
    normalized = tldevice.div_rn(values, scale)
    if COMPUTE_DTYPE == 1:
        normalized = normalized.to(tl.float16).to(tl.float32)
    elif COMPUTE_DTYPE == 2:
        normalized = normalized.to(tl.bfloat16).to(tl.float32)
    if HAS_ZP:
        # fake_quantize adds the zero point in the input dtype
        normalized += zero_point
        if COMPUTE_DTYPE == 1:
            normalized = normalized.to(tl.float16).to(tl.float32)
        elif COMPUTE_DTYPE == 2:
            normalized = normalized.to(tl.bfloat16).to(tl.float32)
    normalized = tl.clamp(normalized, q_min, q_max)
    if QUANT_TYPE == 0:
        rounded = tldevice.rint(normalized)
    elif NUM_BITS == 4:
        rounded = _round_to_fp4(normalized)
    else:
        rounded = normalized.to(tl.float8e4nv).to(tl.float32)
    if HAS_ZP:
        rounded -= zero_point
    return tldevice.mul_rn(rounded, scale)

round_candidate_scale 

round_candidate_scale(
    scale,
    SCALE_ROUND_TYPE: constexpr,
    QUANT_TYPE: constexpr,
)

Round a floating candidate scale as compressed-tensors qparams do.

Source code in src/llmcompressor/utils/triton_utils.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
@triton.jit
def round_candidate_scale(
    scale, SCALE_ROUND_TYPE: tl.constexpr, QUANT_TYPE: tl.constexpr
):
    """Round a floating candidate scale as compressed-tensors qparams do."""
    if SCALE_ROUND_TYPE == 1:
        return tl.minimum(scale, 65504.0).to(tl.float16).to(tl.float32)
    if SCALE_ROUND_TYPE == 2:
        return scale.to(tl.bfloat16).to(tl.float32)
    if SCALE_ROUND_TYPE == 3:
        return tl.minimum(scale, 448.0).to(tl.float8e4nv).to(tl.float32)
    if SCALE_ROUND_TYPE == 4:
        # E8M0 stores the power-of-two scale after accounting for the element
        # format's largest exponent (FP4: 2, FP8: 8).
        is_zero = scale == 0.0
        safe_scale = tl.where(is_zero, 1.0, scale)
        exponent = tl.floor(tldevice.log2(safe_scale))
        power_2 = tldevice.exp2(exponent)
        exponent += tl.where(safe_scale >= 1.75 * power_2, 1.0, 0.0)
        offset = 2.0 if QUANT_TYPE == 1 else 8.0
        encoded = tl.clamp(exponent - offset + 127.0, 0.0, 255.0)
        encoded = tl.where(is_zero, 0.0, encoded)
        return tldevice.exp2(encoded - 127.0)
    return scale