Skip to content

Instantly share code, notes, and snippets.

@HDCharles

HDCharles/gist:51c392795d109078cb89b2f45ccfa11c

Last active September 14, 2026 14:21
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save HDCharles/51c392795d109078cb89b2f45ccfa11c to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save HDCharles/51c392795d109078cb89b2f45ccfa11c to your computer and use it in GitHub Desktop.
Download ZIP
mse_bench.py
Raw
gistfile1.txt
"""
Benchmark for MSE observer grid search. (results at bottom)
Profiles the grid search that finds optimal min/max ranges for quantization.
The hot path is `_calculate_error` (called once per shrink step: default 20
steps), which runs `calculate_qparams` + `fake_quantize` on the full weight
tensor each time.
Usage:
python benchmarks/bench_mse_observer.py [--device cuda] [--rows 4096] [--cols 4096]
"""
import argparse
import time
import torch
import triton
import triton.language as tl
from compressed_tensors.quantization import (
QuantizationArgs,
QuantizationScheme,
QuantizationStrategy,
QuantizationType,
)
from compressed_tensors.quantization.lifecycle import fake_quantize
from compressed_tensors.quantization.utils import calculate_qparams, calculate_range
from llmcompressor.modifiers.quantization.calibration import (
initialize_observer,
observe,
)
from llmcompressor.observers.helpers import flatten_for_calibration
WARMUP = 2
ITERS = 5
# Configured by CLI in main(). Kept separate from the common benchmark function
# signature so the existing variants remain directly comparable.
HIER_CONFIG = {
"n": 16,
"keep_num": 1,
"keep_den": 4,
"m": 2,
"k": 3,
"p_min": 0.8,
"p_max": 1.8,
"fast_pow": False,
}
# ── Inlined grid search variants ──────────────────────────────────────────────
def grid_search_eager(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Eager grid search — direct port from mse_quant.py."""
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
best_error = torch.full_like(min_val, torch.finfo(min_val.dtype).max)
best_min_val = min_val.clone()
best_max_val = max_val.clone()
total_steps = int(maxshrink * grid)
no_improve_count = 0
for i in range(total_steps):
p = 1 - i / grid
shrinked_min = min_val * p
shrinked_max = max_val * p
candidate_scales, candidate_zero_points = calculate_qparams(
min_vals=shrinked_min,
max_vals=shrinked_max,
quantization_args=args,
global_scale=None,
)
q = fake_quantize(
observed,
candidate_scales.unsqueeze(-1),
candidate_zero_points.unsqueeze(-1),
token_args,
).to(observed.dtype)
err = torch.sum((q - observed).abs().pow(norm), dim=(0, -1))
del q
improved = err < best_error
if torch.any(improved):
best_error[improved] = err[improved]
best_min_val[improved] = shrinked_min[improved]
best_max_val[improved] = shrinked_max[improved]
no_improve_count = 0
else:
no_improve_count += 1
if no_improve_count >= patience:
break
return best_min_val, best_max_val
def grid_search_eager_no_patience(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Eager grid search that always evaluates the complete candidate grid."""
del patience
total_steps = int(maxshrink * grid)
return grid_search_eager(
observed,
args,
token_args,
maxshrink,
total_steps + 1,
grid,
norm,
chunk_size,
)
@torch.compile(dynamic=True)
def _compute_chunk_compiled(
observed,
args,
token_args,
min_val,
max_val,
ps,
chunk_size,
norm,
best_error,
best_min_val,
best_max_val,
):
for j in range(chunk_size):
shrinked_min = min_val * ps[j]
shrinked_max = max_val * ps[j]
candidate_scales, candidate_zero_points = calculate_qparams(
min_vals=shrinked_min,
max_vals=shrinked_max,
quantization_args=args,
global_scale=None,
)
q = fake_quantize(
observed,
candidate_scales.unsqueeze(-1),
candidate_zero_points.unsqueeze(-1),
token_args,
).to(observed.dtype)
err = torch.sum((q - observed).abs().pow(norm), dim=(0, -1))
del q
improved = err < best_error
best_error = torch.where(improved, err, best_error)
best_min_val = torch.where(improved, shrinked_min, best_min_val)
best_max_val = torch.where(improved, shrinked_max, best_max_val)
return best_error, best_min_val, best_max_val
def grid_search_compiled(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Compiled grid search — chunked torch.compile inner loop."""
import torch._dynamo.config
import torch._dynamo.decorators
torch._dynamo.config.capture_scalar_outputs = True
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
best_error = torch.full_like(min_val, torch.finfo(min_val.dtype).max)
best_min_val = min_val.clone()
best_max_val = max_val.clone()
total_steps = int(maxshrink * grid)
no_improve_count = 0
observed = observed.clone()
torch._dynamo.decorators.mark_unbacked(observed, observed.ndim - 1)
idx = 0
while idx < total_steps:
chunk_end = min(idx + chunk_size, total_steps)
current_chunk = chunk_end - idx
ps = torch.tensor(
[1.0 - (idx + j) / grid for j in range(current_chunk)],
dtype=observed.dtype,
device=observed.device,
)
prev_best = best_error.clone()
best_error, best_min_val, best_max_val = _compute_chunk_compiled(
observed,
args,
token_args,
min_val,
max_val,
ps,
current_chunk,
norm,
best_error,
best_min_val,
best_max_val,
)
if torch.equal(prev_best, best_error):
no_improve_count += current_chunk
if no_improve_count >= patience:
break
else:
no_improve_count = 0
idx = chunk_end
return best_min_val, best_max_val
# ── Triton fused implementation ───────────────────────────────────────────────
@triton.jit
def _fused_grid_search_kernel(
observed_ptr,
all_scales_ptr,
all_zps_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
stride_obs_row,
stride_obs_group,
stride_scale_step,
stride_scale_row,
q_min: tl.constexpr,
q_max: tl.constexpr,
norm,
BLOCK_G: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
):
"""Precomputed-qparam family: serial exhaustive search, one program/group.
Unlike the GPTQ-arithmetic variants, this uses the original simple
clamp/round/dequant path. Every candidate scale and zero point is computed
by PyTorch before launch, then grid steps are scanned serially in-kernel.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
)
best_err = float("inf")
best_s = 0
for step in range(TOTAL_STEPS):
if step < total_steps:
scale = tl.load(
all_scales_ptr
+ step * stride_scale_step
+ row * stride_scale_row
+ group
)
zp = tl.load(
all_zps_ptr + step * stride_scale_step + row * stride_scale_row + group
)
# fused fake_quantize: round(x/scale + zp), clamp, dequant
q_float = obs / scale + zp
q_int = tl.extra.cuda.libdevice.nearbyint(q_float)
q_int = tl.minimum(tl.maximum(q_int, q_min), q_max)
q = (q_int - zp) * scale
diff = tl.abs(q - obs)
norm_f = norm.to(tl.float32)
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm_f)
err = tl.sum(tl.where(g_mask, diff_pow, 0.0), axis=0)
is_better = err < best_err
best_err = tl.where(is_better, err, best_err)
best_s = tl.where(is_better, step, best_s)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
def grid_search_triton(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Fused Triton grid search — one kernel per (row, group), all steps in registers."""
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
# Precompute all candidate scales and zero_points
all_scales = []
all_zps = []
for i in range(total_steps):
p = 1 - i / grid
s, zp = calculate_qparams(
min_vals=min_val * p,
max_vals=max_val * p,
quantization_args=args,
global_scale=None,
)
all_scales.append(s)
all_zps.append(zp)
all_scales = torch.stack(all_scales).contiguous() # (total_steps, *qparam_shape)
all_zps = torch.stack(all_zps).to(dtype=torch.float32).contiguous()
_, num_rows, num_groups, group_size = observed.shape
best_step = torch.zeros(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.full(
(num_rows, num_groups),
float("inf"),
device=observed.device,
dtype=torch.float32,
)
observed_contig = observed.contiguous()
BLOCK_G = triton.next_power_of_2(group_size)
TOTAL_STEPS = triton.next_power_of_2(total_steps)
num_programs = num_rows * num_groups
grid_launch = (num_programs,)
bit_range = 2**args.num_bits - 1
if args.symmetric:
q_min_val = -(2 ** (args.num_bits - 1))
q_max_val = 2 ** (args.num_bits - 1) - 1
else:
q_min_val = 0
q_max_val = 2**args.num_bits - 1
_fused_grid_search_kernel[grid_launch](
observed_contig,
all_scales,
all_zps,
best_step,
best_error,
num_rows,
num_groups,
group_size,
total_steps,
observed_contig.stride(1),
observed_contig.stride(2),
all_scales.stride(0),
all_scales.stride(1),
float(q_min_val),
float(q_max_val),
norm,
BLOCK_G=BLOCK_G,
TOTAL_STEPS=TOTAL_STEPS,
)
# Reconstruct best_min/max from best step indices
ps = torch.tensor(
[1.0 - i / grid for i in range(total_steps)],
dtype=min_val.dtype,
device=min_val.device,
)
best_p = ps[best_step.long()]
best_min_val = min_val * best_p
best_max_val = max_val * best_p
return best_min_val, best_max_val
@triton.jit
def _gptq_quantize_dequantize(
values,
scale,
zp,
q_min,
q_max,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
HAS_ZP: tl.constexpr,
):
"""GPTQ-family inline utility implementing its exact QDQ arithmetic.
Shared by exhaustive, patience, vectorized, and hierarchical GPTQ kernels;
this is an inlined helper rather than an independently launched kernel.
"""
normalized = tl.extra.cuda.libdevice.div_rn(values, scale)
if HAS_ZP:
normalized = normalized + zp
clamped = tl.clamp(normalized, q_min, q_max)
if QUANT_TYPE == 0:
rounded = tl.extra.cuda.libdevice.rint(clamped)
elif QUANT_TYPE == 1:
absolute = tl.abs(clamped)
magnitude = tl.where(
absolute <= 0.25,
0.0,
tl.where(
absolute < 0.75,
0.5,
tl.where(
absolute <= 1.25,
1.0,
tl.where(
absolute < 1.75,
1.5,
tl.where(
absolute <= 2.5,
2.0,
tl.where(
absolute < 3.5,
3.0,
tl.where(absolute <= 5.0, 4.0, 6.0),
),
),
),
),
),
)
rounded = tl.where(clamped < 0.0, -magnitude, magnitude)
else:
rounded = clamped.to(tl.float8e4nv).to(tl.float32)
if DEQUANT_DTYPE == 1:
rounded = rounded.to(tl.bfloat16)
scale_value = scale.to(tl.bfloat16)
if HAS_ZP:
rounded = (rounded - zp.to(tl.bfloat16)).to(tl.bfloat16)
quantized = (rounded * scale_value).to(tl.bfloat16).to(tl.float32)
elif DEQUANT_DTYPE == 2:
rounded = rounded.to(tl.float16)
scale_value = scale.to(tl.float16)
if HAS_ZP:
rounded = (rounded - zp.to(tl.float16)).to(tl.float16)
quantized = (rounded * scale_value).to(tl.float16).to(tl.float32)
else:
if HAS_ZP:
quantized = tl.extra.cuda.libdevice.mul_rn(
tl.extra.cuda.libdevice.sub_rn(rounded, zp), scale
)
else:
quantized = tl.extra.cuda.libdevice.mul_rn(rounded, scale)
return quantized
@triton.jit
def _fused_grid_search_gptq_quant_kernel(
observed_ptr,
all_scales_ptr,
all_zps_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
stride_obs_row,
stride_obs_group,
stride_scale_step,
stride_scale_row,
q_min,
q_max,
norm,
patience,
BLOCK_G: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
OBSERVED_DTYPE: tl.constexpr,
HAS_ZP: tl.constexpr,
USE_PATIENCE: tl.constexpr,
):
"""Precomputed-qparam GPTQ family: serial search, optionally with patience.
The quantize/dequantize sequence intentionally mirrors the fused GPTQ
kernel rather than relying on Triton's default arithmetic:
normalized = div_rn(x, scale) + zero_point
rounded = rint(clamp(normalized))
dequantized = mul_rn(rounded - zero_point, scale)
``QUANT_TYPE`` additionally supports the FP4 E2M1 and FP8 E4M3 paths
used by the GPTQ kernel. The grid-search reduction is otherwise the
same as ``_fused_grid_search_kernel``. Unlike the base-scale family, all
candidate qparams are materialized before launch; unlike the vectorized
sibling, steps are serial so per-group patience can skip trailing work.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
).to(tl.float32)
if OBSERVED_DTYPE == 1:
best_err = tl.full([], float("inf"), dtype=tl.bfloat16)
elif OBSERVED_DTYPE == 2:
best_err = tl.full([], float("inf"), dtype=tl.float16)
else:
best_err = tl.full([], float("inf"), dtype=tl.float32)
best_s = 0
patience_ctr = 0
for step in range(TOTAL_STEPS):
if step < total_steps:
active = 1
if USE_PATIENCE:
active = patience_ctr < patience
if active:
scale = tl.load(
all_scales_ptr
+ step * stride_scale_step
+ row * stride_scale_row
+ group
).to(tl.float32)
scale = tl.maximum(scale, 1.1754943508222875e-38)
zp = scale * 0.0
if HAS_ZP:
zp = tl.load(
all_zps_ptr
+ step * stride_scale_step
+ row * stride_scale_row
+ group
).to(tl.float32)
quantized = _gptq_quantize_dequantize(
obs,
scale,
zp,
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=HAS_ZP,
)
if OBSERVED_DTYPE == 1:
obs_for_error = obs.to(tl.bfloat16)
quantized = quantized.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
obs_for_error = obs.to(tl.float16)
quantized = quantized.to(tl.float16)
else:
obs_for_error = obs
diff = tl.abs(quantized - obs_for_error).to(tl.float32)
norm_f = norm.to(tl.float32)
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm_f)
if OBSERVED_DTYPE == 1:
diff_pow = diff_pow.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
diff_pow = diff_pow.to(tl.float16)
err = tl.sum(tl.where(g_mask, diff_pow, 0.0), axis=0).to(best_err.dtype)
is_better = err < best_err
best_err = tl.where(is_better, err, best_err)
best_s = tl.where(is_better, step, best_s)
if USE_PATIENCE:
patience_ctr = tl.where(is_better, 0, patience_ctr + 1)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
@triton.jit
def _round_candidate_scale(
scale,
SCALE_ROUND_TYPE: tl.constexpr,
QUANT_TYPE: tl.constexpr,
):
"""Base-scale-family inline utility matching qparam scale rounding.
It lets the patience kernel derive qparams without precomputing every grid
point, including float16, bfloat16, FP8, and MX E8M0 scale encodings.
"""
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:
# MX E8M0: round the group maximum to a power of two, shift it by
# the element format's exponent offset, encode, then decode to the
# effective floating-point scale consumed by QDQ.
is_zero = scale == 0.0
safe_scale = tl.where(is_zero, 1.0, scale)
exponent = tl.floor(tl.extra.cuda.libdevice.log2(safe_scale))
power_2 = tl.extra.cuda.libdevice.exp2(exponent)
exponent += tl.where(safe_scale >= 1.75 * power_2, 1.0, 0.0)
element_offset = 2.0 if QUANT_TYPE == 1 else 8.0
encoded = tl.clamp(exponent - element_offset + 127.0, 0.0, 255.0)
encoded = tl.where(is_zero, 0.0, encoded)
return tl.extra.cuda.libdevice.exp2(encoded - 127.0)
return scale
@triton.jit
def _base_scale_grid_search_gptq_quant_exhaustive_kernel(
observed_ptr,
scale_base_ptr,
zp_base_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
stride_obs_row,
stride_obs_group,
stride_scale_row,
inv_grid,
q_min,
q_max,
norm,
scale_eps,
BLOCK_G: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
OBSERVED_DTYPE: tl.constexpr,
HAS_ZP: tl.constexpr,
):
"""Base-scale GPTQ family: fully unrolled exhaustive search.
Candidate scales are derived in-kernel from one qparam per group. This is
the lean exhaustive sibling: exact compile-time step count, no padded loop,
activity mask, patience counter, scale rounding, or stopping bookkeeping.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
).to(tl.float32)
qparam_offset = row * stride_scale_row + group
scale_base = tl.load(scale_base_ptr + qparam_offset).to(tl.float32)
zp = scale_base * 0.0
if HAS_ZP:
zp = tl.load(zp_base_ptr + qparam_offset).to(tl.float32)
if OBSERVED_DTYPE == 1:
best_err = tl.full([], float("inf"), dtype=tl.bfloat16)
elif OBSERVED_DTYPE == 2:
best_err = tl.full([], float("inf"), dtype=tl.float16)
else:
best_err = tl.full([], float("inf"), dtype=tl.float32)
best_s = 0
inv_grid_f = inv_grid.to(tl.float32)
norm_f = norm.to(tl.float32)
for step in tl.static_range(0, TOTAL_STEPS):
p = (1.0 - step * inv_grid_f).to(tl.float32)
scale = tl.maximum(scale_base * p, scale_eps)
quantized = _gptq_quantize_dequantize(
obs,
scale,
zp,
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=HAS_ZP,
)
if OBSERVED_DTYPE == 1:
obs_for_error = obs.to(tl.bfloat16)
quantized = quantized.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
obs_for_error = obs.to(tl.float16)
quantized = quantized.to(tl.float16)
else:
obs_for_error = obs
diff = tl.abs(quantized - obs_for_error).to(tl.float32)
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm_f)
if OBSERVED_DTYPE == 1:
diff_pow = diff_pow.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
diff_pow = diff_pow.to(tl.float16)
err = tl.sum(tl.where(g_mask, diff_pow, 0.0), axis=0).to(best_err.dtype)
is_better = err < best_err
best_err = tl.where(is_better, err, best_err)
best_s = tl.where(is_better, step, best_s)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
@triton.jit
def _grid_major_base_scale_gptq_quant_kernel(
observed_ptr,
scale_base_ptr,
zp_base_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
stride_obs_row,
stride_obs_group,
stride_scale_row,
step,
inv_grid,
q_min,
q_max,
norm,
scale_eps,
BLOCK_G: tl.constexpr,
TILE_GROUPS: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
OBSERVED_DTYPE: tl.constexpr,
HAS_ZP: tl.constexpr,
):
"""One candidate per launch, with one or more qparam groups per program.
This is the grid-major counterpart to the serial base-scale kernel: all
groups evaluate candidate ``step`` before the host launches the next grid
point. ``TILE_GROUPS`` controls the qparam-group tile. In particular,
using ``128 // BLOCK_G`` gives eight 16-value NVFP4 groups per program and
one 128-value group per program.
"""
pid = tl.program_id(0)
linear_groups = pid * TILE_GROUPS + tl.arange(0, TILE_GROUPS)
row = linear_groups // num_groups
group = linear_groups % num_groups
group_mask = linear_groups < num_rows * num_groups
group_idx = tl.arange(0, BLOCK_G)
value_mask = group_mask[:, None] & (group_idx[None, :] < group_size)
obs = tl.load(
observed_ptr
+ row[:, None] * stride_obs_row
+ group[:, None] * stride_obs_group
+ group_idx[None, :],
mask=value_mask,
other=0.0,
).to(tl.float32)
qparam_offsets = row * stride_scale_row + group
scale_base = tl.load(scale_base_ptr + qparam_offsets, mask=group_mask, other=1.0)
scale = tl.maximum(scale_base.to(tl.float32) * (1.0 - step * inv_grid), scale_eps)[
:, None
]
zp = scale * 0.0
if HAS_ZP:
zp = tl.load(zp_base_ptr + qparam_offsets, mask=group_mask, other=0.0).to(
tl.float32
)[:, None]
quantized = _gptq_quantize_dequantize(
obs,
scale,
zp,
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=HAS_ZP,
)
if OBSERVED_DTYPE == 1:
obs_for_error = obs.to(tl.bfloat16)
quantized = quantized.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
obs_for_error = obs.to(tl.float16)
quantized = quantized.to(tl.float16)
else:
obs_for_error = obs
diff = tl.abs(quantized - obs_for_error).to(tl.float32)
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm.to(tl.float32))
if OBSERVED_DTYPE == 1:
diff_pow = diff_pow.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
diff_pow = diff_pow.to(tl.float16)
err = tl.sum(tl.where(value_mask, diff_pow, 0.0), axis=1)
old_error = tl.load(
best_error_ptr + linear_groups, mask=group_mask, other=float("inf")
)
is_better = err < old_error
tl.store(
best_error_ptr + linear_groups,
tl.where(is_better, err, old_error),
mask=group_mask,
)
tl.store(
best_step_ptr + linear_groups,
tl.where(is_better, step, 0),
mask=group_mask & is_better,
)
@triton.jit
def _base_scale_grid_search_gptq_quant_exhaustive_tile128_kernel(
observed_ptr,
scale_base_ptr,
zp_base_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
stride_obs_row,
stride_obs_group,
stride_scale_row,
inv_grid,
q_min,
q_max,
norm,
scale_eps,
BLOCK_G: tl.constexpr,
TILE_GROUPS: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
OBSERVED_DTYPE: tl.constexpr,
HAS_ZP: tl.constexpr,
):
"""Serial exhaustive base-scale search, tiled to a fixed value budget.
This retains the existing base kernel's group-major algorithm (every
program evaluates the full grid), but gives each program approximately 128
value lanes. With group size 16, a program owns eight groups; with group
size 128, it owns one group.
"""
pid = tl.program_id(0)
linear_groups = pid * TILE_GROUPS + tl.arange(0, TILE_GROUPS)
row = linear_groups // num_groups
group = linear_groups % num_groups
group_mask = linear_groups < num_rows * num_groups
group_idx = tl.arange(0, BLOCK_G)
value_mask = group_mask[:, None] & (group_idx[None, :] < group_size)
obs = tl.load(
observed_ptr
+ row[:, None] * stride_obs_row
+ group[:, None] * stride_obs_group
+ group_idx[None, :],
mask=value_mask,
other=0.0,
).to(tl.float32)
qparam_offsets = row * stride_scale_row + group
scale_base = tl.load(scale_base_ptr + qparam_offsets, mask=group_mask, other=1.0)
zp = scale_base * 0.0
if HAS_ZP:
zp = tl.load(zp_base_ptr + qparam_offsets, mask=group_mask, other=0.0).to(
tl.float32
)
best_err = tl.full((TILE_GROUPS,), float("inf"), tl.float32)
best_step = tl.zeros((TILE_GROUPS,), tl.int32)
for step in tl.static_range(0, TOTAL_STEPS):
scale = tl.maximum(
scale_base.to(tl.float32)[:, None] * (1.0 - step * inv_grid),
scale_eps,
)
quantized = _gptq_quantize_dequantize(
obs,
scale,
zp[:, None],
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=HAS_ZP,
)
if OBSERVED_DTYPE == 1:
obs_for_error = obs.to(tl.bfloat16)
quantized = quantized.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
obs_for_error = obs.to(tl.float16)
quantized = quantized.to(tl.float16)
else:
obs_for_error = obs
diff = tl.abs(quantized - obs_for_error).to(tl.float32)
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm.to(tl.float32))
if OBSERVED_DTYPE == 1:
diff_pow = diff_pow.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
diff_pow = diff_pow.to(tl.float16)
err = tl.sum(tl.where(value_mask, diff_pow, 0.0), axis=1)
is_better = err < best_err
best_err = tl.where(is_better, err, best_err)
best_step = tl.where(is_better, step, best_step)
tl.store(best_step_ptr + linear_groups, best_step, mask=group_mask)
tl.store(best_error_ptr + linear_groups, best_err, mask=group_mask)
@triton.jit
def _base_scale_grid_search_gptq_quant_patience_kernel(
observed_ptr,
scale_base_ptr,
zp_base_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
stride_obs_row,
stride_obs_group,
stride_scale_row,
inv_grid,
q_min,
q_max,
norm,
patience,
error_buffer,
scale_eps,
BLOCK_G: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
OBSERVED_DTYPE: tl.constexpr,
HAS_ZP: tl.constexpr,
SCALE_ROUND_TYPE: tl.constexpr,
USE_PATIENCE: tl.constexpr,
):
"""Base-scale GPTQ family: serial search with per-group early stopping.
Compared with the exhaustive sibling, this keeps a patience counter and can
reset it for errors within ``error_buffer`` of the best. It also supports
per-candidate scale-dtype rounding (including FP8 and E8M0 scale formats).
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
).to(tl.float32)
qparam_offset = row * stride_scale_row + group
scale_base = tl.load(scale_base_ptr + qparam_offset).to(tl.float32)
zp = scale_base * 0.0
if HAS_ZP:
zp = tl.load(zp_base_ptr + qparam_offset).to(tl.float32)
if OBSERVED_DTYPE == 1:
best_err = tl.full([], float("inf"), dtype=tl.bfloat16)
elif OBSERVED_DTYPE == 2:
best_err = tl.full([], float("inf"), dtype=tl.float16)
else:
best_err = tl.full([], float("inf"), dtype=tl.float32)
best_s = 0
patience_ctr = 0
inv_grid_f = inv_grid.to(tl.float32)
norm_f = norm.to(tl.float32)
error_buffer_f = error_buffer.to(tl.float32)
for step in range(TOTAL_STEPS):
if step < total_steps:
active = True
if USE_PATIENCE:
active = patience_ctr < patience
if active:
p = (1.0 - step * inv_grid_f).to(tl.float32)
scale = _round_candidate_scale(
scale_base * p,
SCALE_ROUND_TYPE=SCALE_ROUND_TYPE,
QUANT_TYPE=QUANT_TYPE,
)
scale = tl.where(scale == 0.0, scale_eps, scale)
scale = tl.maximum(scale, 1.1754943508222875e-38)
quantized = _gptq_quantize_dequantize(
obs,
scale,
zp,
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=HAS_ZP,
)
if OBSERVED_DTYPE == 1:
obs_for_error = obs.to(tl.bfloat16)
quantized = quantized.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
obs_for_error = obs.to(tl.float16)
quantized = quantized.to(tl.float16)
else:
obs_for_error = obs
diff = tl.abs(quantized - obs_for_error).to(tl.float32)
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm_f)
if OBSERVED_DTYPE == 1:
diff_pow = diff_pow.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
diff_pow = diff_pow.to(tl.float16)
err = tl.sum(tl.where(g_mask, diff_pow, 0.0), axis=0).to(best_err.dtype)
previous_best_err = best_err
is_better = err < previous_best_err
best_err = tl.where(is_better, err, best_err)
best_s = tl.where(is_better, step, best_s)
if USE_PATIENCE:
within_buffer = err <= previous_best_err * (1.0 + error_buffer_f)
patience_ctr = tl.where(within_buffer, 0, patience_ctr + 1)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
@triton.jit
def _base_scale_grid_search_gptq_quant_patience_tiled_kernel(
observed_ptr,
scale_base_ptr,
zp_base_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
stride_obs_row,
stride_obs_group,
stride_scale_row,
inv_grid,
q_min,
q_max,
norm,
patience,
error_buffer,
scale_eps,
BLOCK_G: tl.constexpr,
TILE_GROUPS: tl.constexpr,
BLOCK_TILE_GROUPS: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
OBSERVED_DTYPE: tl.constexpr,
HAS_ZP: tl.constexpr,
SCALE_ROUND_TYPE: tl.constexpr,
):
"""Buffered group-major MSE search tiled across adjacent qparam groups.
Each lane has independent best-error and patience state. Unlike the
scalar-group implementation, inactive lanes still participate in vector
QDQ; Triton cannot predicate away just those lanes' arithmetic. This
variant therefore isolates the throughput benefit of a larger tile from
the work-elision benefit of per-group patience.
"""
pid = tl.program_id(0)
tile_offsets = tl.arange(0, BLOCK_TILE_GROUPS)
linear_groups = pid * TILE_GROUPS + tile_offsets
row = linear_groups // num_groups
group = linear_groups % num_groups
group_mask = (tile_offsets < TILE_GROUPS) & (linear_groups < num_rows * num_groups)
group_idx = tl.arange(0, BLOCK_G)
value_mask = group_mask[:, None] & (group_idx[None, :] < group_size)
obs = tl.load(
observed_ptr
+ row[:, None] * stride_obs_row
+ group[:, None] * stride_obs_group
+ group_idx[None, :],
mask=value_mask,
other=0.0,
).to(tl.float32)
qparam_offsets = row * stride_scale_row + group
scale_base = tl.load(scale_base_ptr + qparam_offsets, mask=group_mask, other=1.0)
zp = scale_base * 0.0
if HAS_ZP:
zp = tl.load(zp_base_ptr + qparam_offsets, mask=group_mask, other=0.0).to(
tl.float32
)
if OBSERVED_DTYPE == 1:
best_err = tl.full((BLOCK_TILE_GROUPS,), float("inf"), tl.bfloat16)
elif OBSERVED_DTYPE == 2:
best_err = tl.full((BLOCK_TILE_GROUPS,), float("inf"), tl.float16)
else:
best_err = tl.full((BLOCK_TILE_GROUPS,), float("inf"), tl.float32)
best_step = tl.zeros((BLOCK_TILE_GROUPS,), tl.int32)
stale = tl.zeros((BLOCK_TILE_GROUPS,), tl.int32)
for step in range(TOTAL_STEPS):
if step < total_steps:
active = stale < patience
scale = _round_candidate_scale(
scale_base.to(tl.float32) * (1.0 - step * inv_grid),
SCALE_ROUND_TYPE=SCALE_ROUND_TYPE,
QUANT_TYPE=QUANT_TYPE,
)
scale = tl.maximum(
tl.where(scale == 0.0, scale_eps, scale), 1.1754943508222875e-38
)[:, None]
quantized = _gptq_quantize_dequantize(
obs,
scale,
zp[:, None],
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=HAS_ZP,
)
if OBSERVED_DTYPE == 1:
obs_for_error = obs.to(tl.bfloat16)
quantized = quantized.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
obs_for_error = obs.to(tl.float16)
quantized = quantized.to(tl.float16)
else:
obs_for_error = obs
diff = tl.abs(quantized - obs_for_error).to(tl.float32)
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm.to(tl.float32))
if OBSERVED_DTYPE == 1:
diff_pow = diff_pow.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
diff_pow = diff_pow.to(tl.float16)
err = tl.sum(tl.where(value_mask, diff_pow, 0.0), axis=1).to(best_err.dtype)
previous_best = best_err
is_better = active & (err < previous_best)
best_err = tl.where(is_better, err, best_err)
best_step = tl.where(is_better, step, best_step)
within_buffer = err <= previous_best * (1.0 + error_buffer)
stale = tl.where(
active & within_buffer,
0,
tl.where(active, stale + 1, stale),
)
tl.store(best_step_ptr + linear_groups, best_step, mask=group_mask)
tl.store(best_error_ptr + linear_groups, best_err, mask=group_mask)
@triton.jit
def _base_scale_grid_search_gptq_quant_split_error_kernel(
observed_ptr,
scale_base_ptr,
zp_base_ptr,
partial_errors_ptr,
num_rows,
num_groups,
group_size,
stride_obs_row,
stride_obs_group,
stride_scale_row,
inv_grid,
q_min,
q_max,
norm,
scale_eps,
BLOCK_VALUES: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
NUM_CHUNKS: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
OBSERVED_DTYPE: tl.constexpr,
HAS_ZP: tl.constexpr,
SCALE_ROUND_TYPE: tl.constexpr,
):
"""Evaluate all candidate errors for one tile inside a large qparam group."""
pid = tl.program_id(0)
qparam = pid // NUM_CHUNKS
chunk = pid % NUM_CHUNKS
total_qparams = num_rows * num_groups
qparam_mask = qparam < total_qparams
row = qparam // num_groups
group = qparam % num_groups
offsets = chunk * BLOCK_VALUES + tl.arange(0, BLOCK_VALUES)
value_mask = qparam_mask & (offsets < group_size)
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + offsets,
mask=value_mask,
other=0.0,
).to(tl.float32)
qparam_offset = row * stride_scale_row + group
scale_base = tl.load(scale_base_ptr + qparam_offset, mask=qparam_mask, other=1.0)
zp = scale_base * 0.0
if HAS_ZP:
zp = tl.load(zp_base_ptr + qparam_offset, mask=qparam_mask, other=0.0).to(
tl.float32
)
for step in tl.static_range(0, TOTAL_STEPS):
scale = _round_candidate_scale(
scale_base.to(tl.float32) * (1.0 - step * inv_grid),
SCALE_ROUND_TYPE=SCALE_ROUND_TYPE,
QUANT_TYPE=QUANT_TYPE,
)
scale = tl.maximum(
tl.where(scale == 0.0, scale_eps, scale), 1.1754943508222875e-38
)
quantized = _gptq_quantize_dequantize(
obs,
scale,
zp,
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=HAS_ZP,
)
if OBSERVED_DTYPE == 1:
obs_for_error = obs.to(tl.bfloat16)
quantized = quantized.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
obs_for_error = obs.to(tl.float16)
quantized = quantized.to(tl.float16)
else:
obs_for_error = obs
diff = tl.abs(quantized - obs_for_error).to(tl.float32)
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm.to(tl.float32))
if OBSERVED_DTYPE == 1:
diff_pow = diff_pow.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
diff_pow = diff_pow.to(tl.float16)
error = tl.sum(tl.where(value_mask, diff_pow, 0.0))
tl.store(
partial_errors_ptr + (step * total_qparams + qparam) * NUM_CHUNKS + chunk,
error,
mask=qparam_mask,
)
@triton.jit
def _base_scale_grid_search_gptq_quant_reduce_split_errors_kernel(
partial_errors_ptr,
best_step_ptr,
best_error_ptr,
total_qparams,
patience,
error_buffer,
BLOCK_CHUNKS: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
NUM_CHUNKS: tl.constexpr,
):
"""Reduce value-tile errors, then apply buffered per-qparam patience."""
qparam = tl.program_id(0)
qparam_mask = qparam < total_qparams
chunk_offsets = tl.arange(0, BLOCK_CHUNKS)
chunk_mask = chunk_offsets < NUM_CHUNKS
best_error = tl.full([], float("inf"), tl.float32)
best_step = 0
stale = 0
for step in tl.static_range(0, TOTAL_STEPS):
partials = tl.load(
partial_errors_ptr
+ (step * total_qparams + qparam) * NUM_CHUNKS
+ chunk_offsets,
mask=qparam_mask & chunk_mask,
other=0.0,
)
error = tl.sum(partials)
active = stale < patience
previous_best = best_error
is_better = active & (error < previous_best)
best_error = tl.where(is_better, error, best_error)
best_step = tl.where(is_better, step, best_step)
within_buffer = error <= previous_best * (1.0 + error_buffer)
stale = tl.where(active & within_buffer, 0, tl.where(active, stale + 1, stale))
tl.store(best_step_ptr + qparam, best_step, mask=qparam_mask)
tl.store(best_error_ptr + qparam, best_error, mask=qparam_mask)
@triton.jit
def _vectorized_grid_search_gptq_quant_kernel(
observed_ptr,
all_scales_ptr,
all_zps_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
stride_obs_row,
stride_obs_group,
stride_scale_step,
stride_scale_row,
q_min,
q_max,
norm,
patience,
BLOCK_G: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
OBSERVED_DTYPE: tl.constexpr,
HAS_ZP: tl.constexpr,
USE_PATIENCE: tl.constexpr,
):
"""Precomputed-qparam GPTQ family: vectorize grid steps with group values.
Unlike the serial GPTQ sibling, this evaluates one
``[TOTAL_STEPS, BLOCK_G]`` tile and reduces over all steps afterward. It
exposes more parallelism but cannot save quantization work via patience.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
step_idx = tl.arange(0, TOTAL_STEPS)
group_idx = tl.arange(0, BLOCK_G)
step_offsets = step_idx[:, None]
group_offsets = group_idx[None, :]
step_mask = step_offsets < total_steps
group_mask = group_offsets < group_size
value_mask = step_mask & group_mask
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + group_offsets,
mask=group_mask,
other=0.0,
).to(tl.float32)
scale = tl.load(
all_scales_ptr
+ step_offsets * stride_scale_step
+ row * stride_scale_row
+ group,
mask=step_mask,
other=1.0,
).to(tl.float32)
scale = tl.maximum(scale, 1.1754943508222875e-38)
zp = scale * 0.0
if HAS_ZP:
zp = tl.load(
all_zps_ptr
+ step_offsets * stride_scale_step
+ row * stride_scale_row
+ group,
mask=step_mask,
other=0.0,
).to(tl.float32)
quantized = _gptq_quantize_dequantize(
obs,
scale,
zp,
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=HAS_ZP,
)
if OBSERVED_DTYPE == 1:
obs_for_error = obs.to(tl.bfloat16)
quantized = quantized.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
obs_for_error = obs.to(tl.float16)
quantized = quantized.to(tl.float16)
else:
obs_for_error = obs
diff = tl.abs(quantized - obs_for_error).to(tl.float32)
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm.to(tl.float32))
if OBSERVED_DTYPE == 1:
diff_pow = diff_pow.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
diff_pow = diff_pow.to(tl.float16)
errors = tl.sum(tl.where(value_mask, diff_pow, 0.0), axis=1)
if OBSERVED_DTYPE == 1:
errors = errors.to(tl.bfloat16)
elif OBSERVED_DTYPE == 2:
errors = errors.to(tl.float16)
errors = tl.where(step_idx < total_steps, errors, float("inf"))
best_err, best_s = tl.min(
errors,
axis=0,
return_indices=True,
return_indices_tie_break_left=True,
)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
def _gptq_quant_type(args):
"""Return the quantization mode constants used by the GPTQ Triton kernel."""
if args.type == QuantizationType.INT:
return 0
if args.type == QuantizationType.FLOAT and args.num_bits == 4:
return 1
if args.type == QuantizationType.FLOAT and args.num_bits == 8:
return 2
raise ValueError(
"GPTQ-style MSE quantization supports INT, FP4 E2M1, and FP8 E4M3; "
f"got type={args.type}, num_bits={args.num_bits}"
)
def _dequant_dtype_code(dtype):
"""Return the GPTQ kernel's dequantization dtype code."""
if dtype == torch.float32:
return 0
if dtype == torch.bfloat16:
return 1
if dtype == torch.float16:
return 2
raise ValueError(
"GPTQ-style MSE quantization supports float32, bfloat16, or float16 "
f"scale dtypes, got {dtype}"
)
def _observed_dtype_code(dtype):
"""Return the observed dtype code used for MSE error evaluation."""
if dtype == torch.float32:
return 0
if dtype == torch.bfloat16:
return 1
if dtype == torch.float16:
return 2
raise ValueError(
"GPTQ-style MSE quantization supports float32, bfloat16, or float16 "
f"observed dtypes, got {dtype}"
)
def _launch_gptq_quant_grid_search(
kernel,
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
use_patience=False,
):
"""Shared launcher for GPTQ-style grid-search kernels."""
del token_args, chunk_size
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
all_scales = []
all_zps = []
for i in range(total_steps):
p = 1 - i / grid
s, zp = calculate_qparams(
min_vals=min_val * p,
max_vals=max_val * p,
quantization_args=args,
global_scale=None,
)
all_scales.append(s)
all_zps.append(zp)
all_scales = torch.stack(all_scales).contiguous()
# The PR's GPTQ kernel widens zero points before entering Triton. This
# also handles FP8 zero-point tensors used by floating-point formats.
all_zps = torch.stack(all_zps).to(dtype=torch.float32).contiguous()
_, num_rows, num_groups, group_size = observed.shape
best_step = torch.zeros(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.full(
(num_rows, num_groups),
float("inf"),
device=observed.device,
dtype=torch.float32,
)
observed_contig = observed.contiguous()
BLOCK_G = triton.next_power_of_2(group_size)
TOTAL_STEPS = triton.next_power_of_2(total_steps)
grid_launch = (num_rows * num_groups,)
q_min, q_max = calculate_range(args, observed.device)
quant_type = _gptq_quant_type(args)
dequant_dtype = _dequant_dtype_code(all_scales.dtype)
observed_dtype = _observed_dtype_code(observed.dtype)
has_zp = not args.symmetric
# A dummy zero-point allocation keeps the pointer valid when symmetric
# quantization compiles out the zero-point path.
zp_ptr = all_zps if has_zp else all_scales
kernel[grid_launch](
observed_contig,
all_scales,
zp_ptr,
best_step,
best_error,
num_rows,
num_groups,
group_size,
total_steps,
observed_contig.stride(1),
observed_contig.stride(2),
all_scales.stride(0),
all_scales.stride(1),
float(q_min),
float(q_max),
norm,
patience,
BLOCK_G=BLOCK_G,
TOTAL_STEPS=TOTAL_STEPS,
QUANT_TYPE=quant_type,
DEQUANT_DTYPE=dequant_dtype,
OBSERVED_DTYPE=observed_dtype,
HAS_ZP=has_zp,
USE_PATIENCE=use_patience,
)
ps = torch.tensor(
[1.0 - i / grid for i in range(total_steps)],
dtype=min_val.dtype,
device=min_val.device,
)
best_p = ps[best_step.long()]
return min_val * best_p, max_val * best_p
def grid_search_triton_gptq_quant(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Serial-step grid search with GPTQ quantize/dequantize arithmetic."""
return _launch_gptq_quant_grid_search(
_fused_grid_search_gptq_quant_kernel,
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
)
def grid_search_triton_gptq_quant_vec(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Fully vectorized-step grid search with GPTQ quantization arithmetic."""
return _launch_gptq_quant_grid_search(
_vectorized_grid_search_gptq_quant_kernel,
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
)
def grid_search_triton_gptq_quant_p(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Exact GPTQ-style grid search with per-group patience stopping."""
return _launch_gptq_quant_grid_search(
_fused_grid_search_gptq_quant_kernel,
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
use_patience=True,
)
@triton.jit
def _hierarchical_gptq_quant_kernel(
observed_ptr,
scale_base_ptr,
best_p_ptr,
num_rows,
num_groups,
group_size,
stride_obs_row,
stride_obs_group,
stride_scale_row,
q_min,
q_max,
norm,
p_min: tl.constexpr,
p_max: tl.constexpr,
N: tl.constexpr,
KEEP_NUM: tl.constexpr,
KEEP_DEN: tl.constexpr,
M: tl.constexpr,
K: tl.constexpr,
MAX_KEEP: tl.constexpr,
BLOCK_G: tl.constexpr,
BLOCK_C: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
USE_L2: tl.constexpr,
USE_FAST_POW: tl.constexpr,
):
"""Hierarchical GPTQ family: configurable coarse-to-fine search per group.
It evaluates N coarse candidates, retains a fraction, and repeatedly
refines survivors inside one program. Unlike exhaustive/patience families,
it reduces candidate count geometrically but introduces sequential
generations and survivor-selection bookkeeping.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
).to(tl.float32)
scale_base = tl.load(scale_base_ptr + row * stride_scale_row + group).to(tl.float32)
lanes = tl.arange(0, BLOCK_C)
candidates = p_min + lanes * ((p_max - p_min) / (N - 1))
active_count = N
spacing = (p_max - p_min) / (N - 1)
best_error = float("inf")
best_p = p_max
norm_f = norm.to(tl.float32)
# K is the number of refinement generations after the initial N points.
for level in range(K + 1):
candidate_mask = lanes < active_count
scales = tl.maximum(scale_base * candidates[:, None], 1.1754943508222875e-38)
quantized = _gptq_quantize_dequantize(
obs[None, :],
scales,
scales * 0.0,
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=False,
)
diff = tl.abs(quantized - obs[None, :]).to(tl.float32)
if USE_L2:
diff_pow = diff * diff
elif USE_FAST_POW:
diff_pow = diff * diff * tl.exp2((norm_f - 2.0) * tl.log2(diff))
else:
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm_f)
errors = tl.sum(tl.where(g_mask[None, :], diff_pow, 0.0), axis=1)
errors = tl.where(candidate_mask, errors, float("inf"))
stage_error = tl.min(errors, axis=0)
stage_index = tl.argmin(errors, axis=0)
stage_p = tl.sum(tl.where(lanes == stage_index, candidates, 0.0), axis=0)
improved = stage_error < best_error
best_error = tl.where(improved, stage_error, best_error)
best_p = tl.where(improved, stage_p, best_p)
if level < K:
keep_count = active_count * KEEP_NUM // KEEP_DEN
retained = tl.zeros((BLOCK_C,), tl.float32)
remaining_errors = errors
for rank in range(MAX_KEEP):
if rank < keep_count:
selected_index = tl.argmin(remaining_errors, axis=0)
selected_p = tl.sum(
tl.where(lanes == selected_index, candidates, 0.0), axis=0
)
retained = tl.where(lanes == rank, selected_p, retained)
remaining_errors = tl.where(
lanes == selected_index, float("inf"), remaining_errors
)
parent_index = lanes // M
child_index = lanes % M
parent = tl.sum(
tl.where(
lanes[None, :] == parent_index[:, None],
retained[None, :],
0.0,
),
axis=1,
)
offset = ((2 * child_index + 1 - M) * spacing) / (2 * M)
candidates = tl.maximum(p_min, tl.minimum(p_max, parent + offset))
active_count = keep_count * M
spacing = spacing / (2 * M)
tl.store(best_p_ptr + row * num_groups + group, best_p)
def grid_search_triton_hierarchical(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Configurable fused hierarchical search using GPTQ QDQ arithmetic."""
del token_args, maxshrink, patience, grid, chunk_size
config = HIER_CONFIG
n = config["n"]
m = config["m"]
keep_num = config["keep_num"]
keep_den = config["keep_den"]
k = config["k"]
active_count = n
max_candidates = n
max_keep = 0
for _ in range(k):
keep_count = active_count * keep_num // keep_den
if keep_count < 1 or active_count * keep_num % keep_den:
raise ValueError(
"Each hierarchical generation must retain a positive integer "
"number of candidates"
)
max_keep = max(max_keep, keep_count)
active_count = keep_count * m
max_candidates = max(max_candidates, active_count)
if m <= 0 or m % 2:
raise ValueError("hier-m must be a positive even integer")
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
scale_base, _ = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=args,
global_scale=None,
)
_, num_rows, num_groups, group_size = observed.shape
best_p = torch.empty(
(num_rows, num_groups), dtype=torch.float32, device=observed.device
)
q_min, q_max = calculate_range(args, observed.device)
block_g = triton.next_power_of_2(group_size)
block_c = triton.next_power_of_2(max_candidates)
observed_contig = observed.contiguous()
_hierarchical_gptq_quant_kernel[(num_rows * num_groups,)](
observed_contig,
scale_base.contiguous(),
best_p,
num_rows,
num_groups,
group_size,
observed_contig.stride(1),
observed_contig.stride(2),
scale_base.stride(0),
float(q_min),
float(q_max),
norm,
p_min=config["p_min"],
p_max=config["p_max"],
N=n,
KEEP_NUM=keep_num,
KEEP_DEN=keep_den,
M=m,
K=k,
MAX_KEEP=max_keep,
BLOCK_G=block_g,
BLOCK_C=block_c,
QUANT_TYPE=_gptq_quant_type(args),
DEQUANT_DTYPE=_dequant_dtype_code(scale_base.dtype),
USE_L2=norm == 2.0,
USE_FAST_POW=config["fast_pow"],
)
return min_val * best_p, max_val * best_p
def grid_search_triton_hierarchical_fast(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Hierarchical search that always uses the fast approximate power."""
previous = HIER_CONFIG["fast_pow"]
HIER_CONFIG["fast_pow"] = True
try:
return grid_search_triton_hierarchical(
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
)
finally:
HIER_CONFIG["fast_pow"] = previous
@triton.jit
def _insert_top4(error, candidate, e0, p0, e1, p1, e2, p2, e3, p3):
"""Hierarchical-family inline utility inserting one scalar into a top-4."""
swap = error < e0
error, e0 = tl.where(swap, e0, error), tl.where(swap, error, e0)
candidate, p0 = (
tl.where(swap, p0, candidate),
tl.where(swap, candidate, p0),
)
swap = error < e1
error, e1 = tl.where(swap, e1, error), tl.where(swap, error, e1)
candidate, p1 = (
tl.where(swap, p1, candidate),
tl.where(swap, candidate, p1),
)
swap = error < e2
error, e2 = tl.where(swap, e2, error), tl.where(swap, error, e2)
candidate, p2 = (
tl.where(swap, p2, candidate),
tl.where(swap, candidate, p2),
)
swap = error < e3
e3 = tl.where(swap, error, e3)
p3 = tl.where(swap, candidate, p3)
return e0, p0, e1, p1, e2, p2, e3, p3
@triton.jit
def _hierarchical_candidate_error(
obs,
g_mask,
scale,
q_min,
q_max,
norm_f,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
USE_L2: tl.constexpr,
):
"""Serial-hierarchical inline utility evaluating one candidate's error."""
quantized = _gptq_quantize_dequantize(
obs,
scale,
scale * 0.0,
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=False,
)
diff = tl.abs(quantized - obs).to(tl.float32)
if USE_L2:
diff_pow = diff * diff
else:
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm_f)
return tl.sum(tl.where(g_mask, diff_pow, 0.0), axis=0)
@triton.jit
def _hierarchical_serial_32_4_8_kernel(
observed_ptr,
scale_base_ptr,
best_p_ptr,
num_rows,
num_groups,
group_size,
stride_obs_row,
stride_obs_group,
stride_scale_row,
q_min,
q_max,
norm,
p_min: tl.constexpr,
p_max: tl.constexpr,
BLOCK_G: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
USE_L2: tl.constexpr,
):
"""Hierarchical GPTQ family: scalar specialized 32 -> top-4 -> 32 search.
This hard-codes the favored configuration and evaluates candidates serially
to minimize tile pressure. It differs from the generic kernel by removing
configurability and from the tiled sibling by avoiding candidate vectors.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
).to(tl.float32)
scale_base = tl.load(scale_base_ptr + row * stride_scale_row + group).to(tl.float32)
norm_f = norm.to(tl.float32)
e0 = float("inf")
e1 = float("inf")
e2 = float("inf")
e3 = float("inf")
p0 = p_min
p1 = p_min
p2 = p_min
p3 = p_min
spacing: tl.constexpr = (p_max - p_min) / 31.0
for coarse in range(32):
candidate = p_min + coarse * spacing
scale = tl.maximum(scale_base * candidate, 1.1754943508222875e-38)
error = _hierarchical_candidate_error(
obs,
g_mask,
scale,
q_min,
q_max,
norm_f,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
USE_L2=USE_L2,
)
e0, p0, e1, p1, e2, p2, e3, p3 = _insert_top4(
error, candidate, e0, p0, e1, p1, e2, p2, e3, p3
)
best_error = e0
best_p = p0
for parent_index in range(4):
if parent_index == 0:
parent = p0
elif parent_index == 1:
parent = p1
elif parent_index == 2:
parent = p2
else:
parent = p3
for child in range(8):
offset = (2 * child + 1 - 8) * spacing / 16.0
candidate = tl.maximum(p_min, tl.minimum(p_max, parent + offset))
scale = tl.maximum(scale_base * candidate, 1.1754943508222875e-38)
error = _hierarchical_candidate_error(
obs,
g_mask,
scale,
q_min,
q_max,
norm_f,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
USE_L2=USE_L2,
)
improved = error < best_error
best_error = tl.where(improved, error, best_error)
best_p = tl.where(improved, candidate, best_p)
tl.store(best_p_ptr + row * num_groups + group, best_p)
def grid_search_triton_hierarchical_serial(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Register-efficient specialized hierarchical NVFP4 search."""
del token_args, maxshrink, patience, grid, chunk_size
if not args.symmetric:
raise ValueError("The specialized hierarchical kernel requires symmetry")
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
scale_base, _ = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=args,
global_scale=None,
)
_, num_rows, num_groups, group_size = observed.shape
best_p = torch.empty(
(num_rows, num_groups), dtype=torch.float32, device=observed.device
)
observed_contig = observed.contiguous()
scale_base = scale_base.contiguous()
q_min, q_max = calculate_range(args, observed.device)
_hierarchical_serial_32_4_8_kernel[(num_rows * num_groups,)](
observed_contig,
scale_base,
best_p,
num_rows,
num_groups,
group_size,
observed_contig.stride(1),
observed_contig.stride(2),
scale_base.stride(0),
float(q_min),
float(q_max),
norm,
p_min=HIER_CONFIG["p_min"],
p_max=HIER_CONFIG["p_max"],
BLOCK_G=triton.next_power_of_2(group_size),
QUANT_TYPE=_gptq_quant_type(args),
DEQUANT_DTYPE=_dequant_dtype_code(scale_base.dtype),
USE_L2=norm == 2.0,
)
return min_val * best_p, max_val * best_p
@triton.jit
def _hierarchical_tile_errors(
obs,
g_mask,
scales,
q_min,
q_max,
norm_f,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
USE_L2: tl.constexpr,
):
"""Tiled-hierarchical inline utility evaluating a candidate vector."""
quantized = _gptq_quantize_dequantize(
obs[None, :],
scales[:, None],
scales[:, None] * 0.0,
q_min,
q_max,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
HAS_ZP=False,
)
diff = tl.abs(quantized - obs[None, :]).to(tl.float32)
if USE_L2:
diff_pow = diff * diff
else:
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm_f)
return tl.sum(tl.where(g_mask[None, :], diff_pow, 0.0), axis=1)
@triton.jit
def _hierarchical_tiled_32_4_8_kernel(
observed_ptr,
scale_base_ptr,
best_p_ptr,
num_rows,
num_groups,
group_size,
stride_obs_row,
stride_obs_group,
stride_scale_row,
q_min,
q_max,
norm,
p_min: tl.constexpr,
p_max: tl.constexpr,
BLOCK_G: tl.constexpr,
QUANT_TYPE: tl.constexpr,
DEQUANT_DTYPE: tl.constexpr,
USE_L2: tl.constexpr,
):
"""Hierarchical GPTQ family: tiled specialized 32 -> top-4 -> 32 search.
This evaluates candidates in parallel tiles of eight, then performs scalar
top-4 selection. It trades higher register pressure for more candidate
parallelism compared with the serial specialized sibling.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
).to(tl.float32)
scale_base = tl.load(scale_base_ptr + row * stride_scale_row + group).to(tl.float32)
norm_f = norm.to(tl.float32)
lanes = tl.arange(0, 8)
spacing: tl.constexpr = (p_max - p_min) / 31.0
e0 = float("inf")
e1 = float("inf")
e2 = float("inf")
e3 = float("inf")
p0 = p_min
p1 = p_min
p2 = p_min
p3 = p_min
for tile in range(4):
candidates = p_min + (tile * 8 + lanes) * spacing
scales = tl.maximum(scale_base * candidates, 1.1754943508222875e-38)
errors = _hierarchical_tile_errors(
obs,
g_mask,
scales,
q_min,
q_max,
norm_f,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
USE_L2=USE_L2,
)
for lane in range(8):
error = tl.sum(tl.where(lanes == lane, errors, 0.0), axis=0)
candidate = p_min + (tile * 8 + lane) * spacing
e0, p0, e1, p1, e2, p2, e3, p3 = _insert_top4(
error, candidate, e0, p0, e1, p1, e2, p2, e3, p3
)
best_error = e0
best_p = p0
child_offsets = (2 * lanes + 1 - 8) * spacing / 16.0
for parent_index in range(4):
if parent_index == 0:
parent = p0
elif parent_index == 1:
parent = p1
elif parent_index == 2:
parent = p2
else:
parent = p3
candidates = tl.maximum(p_min, tl.minimum(p_max, parent + child_offsets))
scales = tl.maximum(scale_base * candidates, 1.1754943508222875e-38)
errors = _hierarchical_tile_errors(
obs,
g_mask,
scales,
q_min,
q_max,
norm_f,
QUANT_TYPE=QUANT_TYPE,
DEQUANT_DTYPE=DEQUANT_DTYPE,
USE_L2=USE_L2,
)
stage_error = tl.min(errors, axis=0)
stage_index = tl.argmin(errors, axis=0)
stage_p = tl.sum(tl.where(lanes == stage_index, candidates, 0.0), axis=0)
improved = stage_error < best_error
best_error = tl.where(improved, stage_error, best_error)
best_p = tl.where(improved, stage_p, best_p)
tl.store(best_p_ptr + row * num_groups + group, best_p)
def grid_search_triton_hierarchical_tiled(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Specialized 32 -> 32 hierarchical search with 8-candidate tiles."""
del token_args, maxshrink, patience, grid, chunk_size
if not args.symmetric:
raise ValueError("The tiled hierarchical kernel requires symmetry")
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
scale_base, _ = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=args,
global_scale=None,
)
_, num_rows, num_groups, group_size = observed.shape
best_p = torch.empty(
(num_rows, num_groups), dtype=torch.float32, device=observed.device
)
observed_contig = observed.contiguous()
scale_base = scale_base.contiguous()
q_min, q_max = calculate_range(args, observed.device)
_hierarchical_tiled_32_4_8_kernel[(num_rows * num_groups,)](
observed_contig,
scale_base,
best_p,
num_rows,
num_groups,
group_size,
observed_contig.stride(1),
observed_contig.stride(2),
scale_base.stride(0),
float(q_min),
float(q_max),
norm,
p_min=HIER_CONFIG["p_min"],
p_max=HIER_CONFIG["p_max"],
BLOCK_G=triton.next_power_of_2(group_size),
QUANT_TYPE=_gptq_quant_type(args),
DEQUANT_DTYPE=_dequant_dtype_code(scale_base.dtype),
USE_L2=norm == 2.0,
)
return min_val * best_p, max_val * best_p
def _scale_round_config(dtype):
"""Return kernel rounding mode and zero-scale replacement for scale dtype."""
if dtype == torch.float32:
return 0, torch.finfo(dtype).eps
if dtype == torch.float16:
return 1, torch.finfo(dtype).eps
if dtype == torch.bfloat16:
return 2, torch.finfo(dtype).eps
if dtype == torch.float8_e4m3fn:
return 3, 0.125
if dtype == torch.uint8:
return 4, 2.0**-127
raise ValueError(f"unsupported scale dtype for in-kernel rounding: {dtype}")
def _launch_gptq_base_scale_patience(
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
round_scale,
use_patience=True,
error_buffer=0.0,
):
"""Launch a search deriving candidate scales from one base scale."""
del token_args, chunk_size
if round_scale:
if args.scale_dtype is None:
raise ValueError("scale-rounding variant requires args.scale_dtype")
scale_round_type, scale_eps = _scale_round_config(args.scale_dtype)
if scale_round_type == 4 and not (
args.symmetric
and args.type == QuantizationType.FLOAT
and args.num_bits in (4, 8)
and args.group_size == 32
):
raise ValueError(
"E8M0 scale generation requires symmetric FLOAT4/FLOAT8 "
"quantization with group_size=32"
)
qparam_args = args.model_copy(update={"scale_dtype": None})
else:
if args.scale_dtype is not None:
raise ValueError(
"base-scale variant requires scale_dtype=None; use the "
"scale-rounding variant instead"
)
qparam_args = args
scale_round_type = 0
scale_eps = None
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
scale_base, zp_base = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=qparam_args,
global_scale=None,
)
if scale_round_type == 4:
# MX scale generation starts from the candidate group maximum, not
# from the ordinary max/qmax scale used by other formats.
scale_base = torch.maximum(torch.abs(min_val), torch.abs(max_val))
scale_base = scale_base.contiguous()
zp_base = zp_base.to(dtype=torch.float32).contiguous()
if scale_eps is None:
_, scale_eps = _scale_round_config(scale_base.dtype)
_, num_rows, num_groups, group_size = observed.shape
best_step = torch.zeros(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.full(
(num_rows, num_groups),
float("inf"),
device=observed.device,
dtype=torch.float32,
)
observed_contig = observed.contiguous()
q_min, q_max = calculate_range(args, observed.device)
quant_type = _gptq_quant_type(args)
dequant_dtype = _dequant_dtype_code(scale_base.dtype)
observed_dtype = _observed_dtype_code(observed.dtype)
has_zp = not args.symmetric
zp_ptr = zp_base if has_zp else scale_base
BLOCK_G = triton.next_power_of_2(group_size)
grid_launch = (num_rows * num_groups,)
if not use_patience:
_base_scale_grid_search_gptq_quant_exhaustive_kernel[grid_launch](
observed_contig,
scale_base,
zp_ptr,
best_step,
best_error,
num_rows,
num_groups,
group_size,
observed_contig.stride(1),
observed_contig.stride(2),
scale_base.stride(0),
1.0 / grid,
float(q_min),
float(q_max),
norm,
scale_eps,
BLOCK_G=BLOCK_G,
TOTAL_STEPS=total_steps,
QUANT_TYPE=quant_type,
DEQUANT_DTYPE=dequant_dtype,
OBSERVED_DTYPE=observed_dtype,
HAS_ZP=has_zp,
)
else:
TOTAL_STEPS = triton.next_power_of_2(total_steps)
_base_scale_grid_search_gptq_quant_patience_kernel[grid_launch](
observed_contig,
scale_base,
zp_ptr,
best_step,
best_error,
num_rows,
num_groups,
group_size,
total_steps,
observed_contig.stride(1),
observed_contig.stride(2),
scale_base.stride(0),
1.0 / grid,
float(q_min),
float(q_max),
norm,
patience,
error_buffer,
scale_eps,
BLOCK_G=BLOCK_G,
TOTAL_STEPS=TOTAL_STEPS,
QUANT_TYPE=quant_type,
DEQUANT_DTYPE=dequant_dtype,
OBSERVED_DTYPE=observed_dtype,
HAS_ZP=has_zp,
SCALE_ROUND_TYPE=scale_round_type,
USE_PATIENCE=True,
)
ps = torch.tensor(
[1.0 - i / grid for i in range(total_steps)],
dtype=min_val.dtype,
device=min_val.device,
)
best_p = ps[best_step.long()]
return min_val * best_p, max_val * best_p
def grid_search_triton_gptq_base_p(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""GPTQ quantization with base-scale derivation and per-group patience."""
return _launch_gptq_base_scale_patience(
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
round_scale=False,
)
def grid_search_triton_gptq_base(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Exhaustive GPTQ quantization deriving candidates from one base scale."""
return _launch_gptq_base_scale_patience(
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
round_scale=False,
use_patience=False,
)
def _launch_gptq_base_grid_major(
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
tile_values,
):
"""Run an exhaustive base-scale search one grid point at a time in Triton.
Unlike ``triton_gptq_base``, which gives each group one long program that
scans every candidate, this launches the whole tensor for each candidate.
It deliberately mirrors eager's grid-major scheduling while retaining
fused QDQ/error/update work. The benchmark versions are exhaustive: the
global eager-patience check would require a device-to-host synchronization
after every candidate and obscure the scheduling comparison.
"""
del token_args, patience, chunk_size
if args.scale_dtype is not None:
raise ValueError("grid-major base-scale variants require scale_dtype=None")
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
scale_base, zp_base = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=args,
global_scale=None,
)
scale_base = scale_base.contiguous()
zp_base = zp_base.to(dtype=torch.float32).contiguous()
_, scale_eps = _scale_round_config(scale_base.dtype)
_, num_rows, num_groups, group_size = observed.shape
best_step = torch.zeros(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.full(
(num_rows, num_groups),
float("inf"),
dtype=torch.float32,
device=observed.device,
)
observed = observed.contiguous()
block_g = triton.next_power_of_2(group_size)
tile_groups = max(1, tile_values // block_g)
grid_launch = (triton.cdiv(num_rows * num_groups, tile_groups),)
q_min, q_max = calculate_range(args, observed.device)
has_zp = not args.symmetric
for step in range(total_steps):
_grid_major_base_scale_gptq_quant_kernel[grid_launch](
observed,
scale_base,
zp_base if has_zp else scale_base,
best_step,
best_error,
num_rows,
num_groups,
group_size,
observed.stride(1),
observed.stride(2),
scale_base.stride(0),
step,
1.0 / grid,
float(q_min),
float(q_max),
norm,
scale_eps,
BLOCK_G=block_g,
TILE_GROUPS=tile_groups,
QUANT_TYPE=_gptq_quant_type(args),
DEQUANT_DTYPE=_dequant_dtype_code(scale_base.dtype),
OBSERVED_DTYPE=_observed_dtype_code(observed.dtype),
HAS_ZP=has_zp,
)
ps = torch.arange(total_steps, device=min_val.device, dtype=min_val.dtype)
best_p = (1.0 - ps / grid)[best_step.long()]
return min_val * best_p, max_val * best_p
def grid_search_triton_gptq_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Grid-major exhaustive search, one qparam group per Triton program."""
return _launch_gptq_base_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 16
)
def grid_search_triton_gptq_grid_major_tile64(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Grid-major exhaustive search tiled to 64 values per program."""
return _launch_gptq_base_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 64
)
def grid_search_triton_gptq_grid_major_tile128(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Grid-major exhaustive search tiled to 128 values per program."""
return _launch_gptq_base_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 128
)
def grid_search_triton_gptq_grid_major_tile256(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Grid-major exhaustive search tiled to 256 values per program."""
return _launch_gptq_base_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 256
)
def grid_search_triton_gptq_grid_major_tile512(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Grid-major exhaustive search tiled to 512 values per program."""
return _launch_gptq_base_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 512
)
def grid_search_triton_gptq_grid_major_tile1024(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Grid-major exhaustive search tiled to 1024 values per program."""
return _launch_gptq_base_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 1024
)
def grid_search_triton_gptq_grid_major_tile2048(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Grid-major exhaustive search tiled to 2048 values per program."""
return _launch_gptq_base_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 2048
)
def grid_search_triton_gptq_grid_major_tile4096(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Grid-major exhaustive search tiled to 4096 values per program."""
return _launch_gptq_base_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 4096
)
def grid_search_triton_gptq_grid_major_tile8192(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Grid-major exhaustive search tiled to 8192 values per program."""
return _launch_gptq_base_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 8192
)
def grid_search_triton_gptq_grid_major_tile16384(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Grid-major exhaustive search tiled to 16384 values per program."""
return _launch_gptq_base_grid_major(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 16384
)
def grid_search_triton_gptq_base_tile128(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Group-major exhaustive search with 128 value lanes per Triton program.
For group size 16 this processes eight contiguous groups per program; for
group size 128 it processes one group per program. Unlike
``triton_gptq_grid_major``, each program still scans the complete grid.
"""
del token_args, patience, chunk_size
if args.scale_dtype is not None:
raise ValueError("base tile128 variant requires scale_dtype=None")
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
scale_base, zp_base = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=args,
global_scale=None,
)
scale_base = scale_base.contiguous()
zp_base = zp_base.to(dtype=torch.float32).contiguous()
_, scale_eps = _scale_round_config(scale_base.dtype)
_, num_rows, num_groups, group_size = observed.shape
observed = observed.contiguous()
block_g = triton.next_power_of_2(group_size)
tile_groups = max(1, 128 // block_g)
best_step = torch.empty(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.empty(
num_rows, num_groups, dtype=torch.float32, device=observed.device
)
has_zp = not args.symmetric
q_min, q_max = calculate_range(args, observed.device)
_base_scale_grid_search_gptq_quant_exhaustive_tile128_kernel[
(triton.cdiv(num_rows * num_groups, tile_groups),)
](
observed,
scale_base,
zp_base if has_zp else scale_base,
best_step,
best_error,
num_rows,
num_groups,
group_size,
observed.stride(1),
observed.stride(2),
scale_base.stride(0),
1.0 / grid,
float(q_min),
float(q_max),
norm,
scale_eps,
BLOCK_G=block_g,
TILE_GROUPS=tile_groups,
TOTAL_STEPS=total_steps,
QUANT_TYPE=_gptq_quant_type(args),
DEQUANT_DTYPE=_dequant_dtype_code(scale_base.dtype),
OBSERVED_DTYPE=_observed_dtype_code(observed.dtype),
HAS_ZP=has_zp,
)
ps = torch.arange(total_steps, device=min_val.device, dtype=min_val.dtype)
best_p = (1.0 - ps / grid)[best_step.long()]
return min_val * best_p, max_val * best_p
def grid_search_triton_gptq_base_p_b(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Buffered base-scale patience, optionally rounding candidate scales.
With ``scale_dtype=None`` this uses ordinary floating-point scales. When a
scale dtype is configured, each candidate is rounded in-kernel through the
same FP16/BF16/FP8/E8M0 path as the scale-rounding patience variants.
"""
return _launch_gptq_base_scale_patience(
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
round_scale=args.scale_dtype is not None,
use_patience=True,
error_buffer=0.30,
)
def _launch_gptq_base_p_b_tiled(
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
tile_values,
):
"""Launch the buffered base-scale search over a fixed-size value tile."""
del token_args, chunk_size
if args.scale_dtype is None:
scale_round_type = 0
qparam_args = args
else:
scale_round_type, _ = _scale_round_config(args.scale_dtype)
qparam_args = args.model_copy(update={"scale_dtype": None})
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
scale_base, zp_base = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=qparam_args,
global_scale=None,
)
if scale_round_type == 4:
scale_base = torch.maximum(torch.abs(min_val), torch.abs(max_val))
scale_base = scale_base.contiguous()
zp_base = zp_base.to(dtype=torch.float32).contiguous()
_, scale_eps = _scale_round_config(
args.scale_dtype if args.scale_dtype is not None else scale_base.dtype
)
_, num_rows, num_groups, group_size = observed.shape
observed = observed.contiguous()
block_g = triton.next_power_of_2(group_size)
tile_groups = max(1, tile_values // block_g)
block_tile_groups = triton.next_power_of_2(tile_groups)
best_step = torch.empty(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.empty(
num_rows, num_groups, dtype=torch.float32, device=observed.device
)
has_zp = not args.symmetric
q_min, q_max = calculate_range(args, observed.device)
_base_scale_grid_search_gptq_quant_patience_tiled_kernel[
(triton.cdiv(num_rows * num_groups, tile_groups),)
](
observed,
scale_base,
zp_base if has_zp else scale_base,
best_step,
best_error,
num_rows,
num_groups,
group_size,
total_steps,
observed.stride(1),
observed.stride(2),
scale_base.stride(0),
1.0 / grid,
float(q_min),
float(q_max),
norm,
patience,
0.30,
scale_eps,
BLOCK_G=block_g,
TILE_GROUPS=tile_groups,
BLOCK_TILE_GROUPS=block_tile_groups,
TOTAL_STEPS=triton.next_power_of_2(total_steps),
QUANT_TYPE=_gptq_quant_type(args),
DEQUANT_DTYPE=_dequant_dtype_code(scale_base.dtype),
OBSERVED_DTYPE=_observed_dtype_code(observed.dtype),
HAS_ZP=has_zp,
SCALE_ROUND_TYPE=scale_round_type,
)
ps = torch.arange(total_steps, device=min_val.device, dtype=min_val.dtype)
best_p = (1.0 - ps / grid)[best_step.long()]
return min_val * best_p, max_val * best_p
def grid_search_triton_gptq_base_p_b_tile512_adaptive(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Keep every program at roughly 512 values, packing or splitting groups.
Groups up to 512 values use the ordinary tiled kernel, which packs multiple
complete groups into a program. Larger groups are split into independent
512-value programs; a second kernel reduces their candidate errors and
applies the same per-group patience and 30% buffer policy.
"""
tile_values = 512
group_size = observed.shape[-1]
if group_size <= tile_values:
return _launch_gptq_base_p_b_tiled(
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
tile_values,
)
del token_args, chunk_size
if args.scale_dtype is None:
scale_round_type = 0
qparam_args = args
else:
scale_round_type, _ = _scale_round_config(args.scale_dtype)
qparam_args = args.model_copy(update={"scale_dtype": None})
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
scale_base, zp_base = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=qparam_args,
global_scale=None,
)
if scale_round_type == 4:
scale_base = torch.maximum(torch.abs(min_val), torch.abs(max_val))
scale_base = scale_base.contiguous()
zp_base = zp_base.to(dtype=torch.float32).contiguous()
_, scale_eps = _scale_round_config(
args.scale_dtype if args.scale_dtype is not None else scale_base.dtype
)
observed = observed.contiguous()
_, num_rows, num_groups, group_size = observed.shape
total_qparams = num_rows * num_groups
num_chunks = triton.cdiv(group_size, tile_values)
partial_errors = torch.empty(
total_steps,
total_qparams,
num_chunks,
dtype=torch.float32,
device=observed.device,
)
best_step = torch.empty(total_qparams, dtype=torch.int32, device=observed.device)
best_error = torch.empty(total_qparams, dtype=torch.float32, device=observed.device)
q_min, q_max = calculate_range(args, observed.device)
has_zp = not args.symmetric
_base_scale_grid_search_gptq_quant_split_error_kernel[
(total_qparams * num_chunks,)
](
observed,
scale_base,
zp_base if has_zp else scale_base,
partial_errors,
num_rows,
num_groups,
group_size,
observed.stride(1),
observed.stride(2),
scale_base.stride(0),
1.0 / grid,
float(q_min),
float(q_max),
norm,
scale_eps,
BLOCK_VALUES=tile_values,
TOTAL_STEPS=total_steps,
NUM_CHUNKS=num_chunks,
QUANT_TYPE=_gptq_quant_type(args),
DEQUANT_DTYPE=_dequant_dtype_code(scale_base.dtype),
OBSERVED_DTYPE=_observed_dtype_code(observed.dtype),
HAS_ZP=has_zp,
SCALE_ROUND_TYPE=scale_round_type,
)
_base_scale_grid_search_gptq_quant_reduce_split_errors_kernel[(total_qparams,)](
partial_errors,
best_step,
best_error,
total_qparams,
patience,
0.30,
BLOCK_CHUNKS=triton.next_power_of_2(num_chunks),
TOTAL_STEPS=total_steps,
NUM_CHUNKS=num_chunks,
)
ps = torch.arange(total_steps, device=min_val.device, dtype=min_val.dtype)
best_p = (1.0 - ps / grid)[best_step.long()].reshape(min_val.shape)
return min_val * best_p, max_val * best_p
def grid_search_triton_gptq_base_p_b_tile64(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Buffered base-scale search tiled to 64 values per program."""
return _launch_gptq_base_p_b_tiled(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 64
)
def grid_search_triton_gptq_base_p_b_tile128(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Buffered base-scale search tiled to 128 values per program."""
return _launch_gptq_base_p_b_tiled(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 128
)
def grid_search_triton_gptq_base_p_b_tile256(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Buffered base-scale search tiled to 256 values per program."""
return _launch_gptq_base_p_b_tiled(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 256
)
def grid_search_triton_gptq_base_p_b_tile512(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Buffered base-scale search tiled to 512 values per program."""
return _launch_gptq_base_p_b_tiled(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 512
)
def grid_search_triton_gptq_base_p_b_tile640(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Buffered base-scale search tiled to 640 values per program."""
return _launch_gptq_base_p_b_tiled(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 640
)
def grid_search_triton_gptq_base_p_b_tile768(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Buffered base-scale search tiled to 768 values per program."""
return _launch_gptq_base_p_b_tiled(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 768
)
def grid_search_triton_gptq_base_p_b_tile896(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Buffered base-scale search tiled to 896 values per program."""
return _launch_gptq_base_p_b_tiled(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 896
)
def grid_search_triton_gptq_base_p_b_tile1024(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Buffered base-scale search tiled to 1024 values per program."""
return _launch_gptq_base_p_b_tiled(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size, 1024
)
def grid_search_triton_gptq_scale_round_p(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Base-scale patience search with per-candidate scale-dtype rounding."""
return _launch_gptq_base_scale_patience(
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
round_scale=True,
)
def grid_search_triton_gptq_e8_p(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""MX E8M0 scale generation with GPTQ QDQ and per-group patience."""
if args.scale_dtype != torch.uint8:
raise ValueError("E8M0 variant requires scale_dtype=torch.uint8")
return grid_search_triton_gptq_scale_round_p(
observed,
args,
token_args,
maxshrink,
patience,
grid,
norm,
chunk_size,
)
# ── Triton codebook/cutoff implementation (format-agnostic) ──────────────────
@triton.jit
def _fused_grid_search_cutoff_kernel(
observed_ptr,
scale_base_ptr,
codes_ptr,
cutoffs_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
stride_obs_row,
stride_obs_group,
inv_grid,
norm,
BLOCK_G: tl.constexpr,
LOG_C: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
):
"""Codebook family: exhaustive serial search using decision cutoffs.
The codebook contains the normalized representable values for the
quantization format (e.g. integers for INT, irregular values for FP4).
Cutoffs are midpoints between adjacent codebook entries.
Shrinking by p just scales both codebook and cutoffs — equivalently,
we divide observed by (scale_base * p) and bin against the fixed cutoffs.
Binning uses binary search over sorted cutoffs: O(log n) comparisons
per element instead of O(n). Unlike ``mindist``, quantization selects bins
by precomputed midpoint cutoffs rather than comparing neighboring codes.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
num_cutoffs = num_codes - 1
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
)
scale_base = tl.load(scale_base_ptr + row * num_groups + group)
ig = inv_grid.to(tl.float32)
norm_f = norm.to(tl.float32)
best_err = tl.full([], float("inf"), dtype=tl.float32)
best_s = 0
for step in range(TOTAL_STEPS):
if step < total_steps:
p = (1.0 - step * ig).to(tl.float32)
eff_scale = tl.maximum(scale_base * p, 1e-38).to(tl.float32)
obs_norm = obs / eff_scale # [BLOCK_G]
# Binary search over sorted cutoffs to find bin index
lo = tl.zeros([BLOCK_G], dtype=tl.int32)
hi = tl.full([BLOCK_G], num_cutoffs, dtype=tl.int32)
for _ in range(LOG_C):
mid = (lo + hi) >> 1
cut_mid = tl.load(cutoffs_ptr + mid) # gather [BLOCK_G]
ge = obs_norm >= cut_mid
lo = tl.where(ge, mid + 1, lo)
hi = tl.where(ge, hi, mid)
bin_idx = lo
q_norm = tl.load(codes_ptr + bin_idx) # gather [BLOCK_G]
q = q_norm * eff_scale
diff = tl.abs(q - obs).to(tl.float32)
diff_pow = tl.extra.cuda.libdevice.pow(diff, norm_f)
err = tl.sum(tl.where(g_mask, diff_pow, 0.0), axis=0)
is_better = err < best_err
best_err = tl.where(is_better, err, best_err)
best_s = tl.where(is_better, step, best_s)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
def _build_normalized_codebook(args):
"""Build normalized codebook and cutoffs for any quantization format.
Returns (codes, cutoffs) where codes are the normalized representable
values and cutoffs are midpoints between adjacent codes. For INT these
are just integers; for FP4/FP8 they would be the format's representable
values.
"""
if args.symmetric:
q_min = -(2 ** (args.num_bits - 1))
q_max = 2 ** (args.num_bits - 1) - 1
else:
q_min = 0
q_max = 2**args.num_bits - 1
codes = torch.arange(q_min, q_max + 1, dtype=torch.float32)
cutoffs = (codes[:-1] + codes[1:]) * 0.5
return codes, cutoffs
def grid_search_triton_cutoff(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Format-agnostic Triton grid search using codebook + cutoffs."""
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
# Compute per-group scale at p=1.0
scale_base, _ = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=args,
global_scale=None,
)
# Build format-agnostic codebook
codes, cutoffs = _build_normalized_codebook(args)
codes = codes.to(device=observed.device)
cutoffs = cutoffs.to(device=observed.device)
_, num_rows, num_groups, group_size = observed.shape
observed_contig = observed.contiguous()
scale_base = scale_base.contiguous()
best_step = torch.zeros(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.full(
(num_rows, num_groups),
float("inf"),
device=observed.device,
dtype=torch.float32,
)
import math
num_codes = codes.shape[0]
BLOCK_G = triton.next_power_of_2(group_size)
num_cutoffs = num_codes - 1
LOG_C = max(1, math.ceil(math.log2(num_cutoffs))) if num_cutoffs > 0 else 1
TOTAL_STEPS = triton.next_power_of_2(total_steps)
grid_launch = (num_rows * num_groups,)
_fused_grid_search_cutoff_kernel[grid_launch](
observed_contig,
scale_base,
codes,
cutoffs,
best_step,
best_error,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
observed_contig.stride(1),
observed_contig.stride(2),
1.0 / grid,
norm,
BLOCK_G=BLOCK_G,
LOG_C=LOG_C,
TOTAL_STEPS=TOTAL_STEPS,
)
# Reconstruct best_min/max from best step indices
ps = torch.tensor(
[1.0 - i / grid for i in range(total_steps)],
dtype=min_val.dtype,
device=min_val.device,
)
best_p = ps[best_step.long()]
best_min_val = min_val * best_p
best_max_val = max_val * best_p
return best_min_val, best_max_val
# ── Triton codebook min-distance (format-agnostic, no cutoffs) ───────────────
@triton.jit
def _fused_grid_search_mindist_kernel(
observed_ptr,
scale_base_ptr,
codes_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
stride_obs_row,
stride_obs_group,
inv_grid,
norm,
BLOCK_G: tl.constexpr,
LOG_C: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
):
"""Codebook family: exhaustive serial search by nearest-code distance.
For each element, binary searches the sorted codebook to find the
insertion point, then compares distances to the two adjacent codes.
No cutoffs needed — just the codebook itself. Compared with ``cutoff``, it
avoids a cutoff table but performs the final adjacent-distance comparison.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
)
scale_base = tl.load(scale_base_ptr + row * num_groups + group)
ig = inv_grid.to(tl.float32)
norm_f = norm.to(tl.float32)
best_err = tl.full([], float("inf"), dtype=tl.float32)
best_s = 0
for step in range(TOTAL_STEPS):
if step < total_steps:
p = (1.0 - step * ig).to(tl.float32)
eff_scale = tl.maximum(scale_base * p, 1e-38).to(tl.float32)
obs_norm = obs / eff_scale # [BLOCK_G]
# Binary search for insertion point in sorted codes
lo = tl.zeros([BLOCK_G], dtype=tl.int32)
hi = tl.full([BLOCK_G], num_codes, dtype=tl.int32)
for _ in range(LOG_C):
mid = (lo + hi) >> 1
code_mid = tl.load(codes_ptr + mid) # gather [BLOCK_G]
ge = obs_norm >= code_mid
lo = tl.where(ge, mid + 1, lo)
hi = tl.where(ge, hi, mid)
# Nearest is codes[lo-1] or codes[lo], check both
idx_left = tl.maximum(lo - 1, 0)
idx_right = tl.minimum(lo, num_codes - 1)
code_left = tl.load(codes_ptr + idx_left)
code_right = tl.load(codes_ptr + idx_right)
dist_left = tl.abs(obs_norm - code_left)
dist_right = tl.abs(obs_norm - code_right)
min_dist = tl.minimum(dist_left, dist_right) * eff_scale
min_dist_pow = tl.extra.cuda.libdevice.pow(min_dist.to(tl.float32), norm_f)
err = tl.sum(tl.where(g_mask, min_dist_pow, 0.0), axis=0)
is_better = err < best_err
best_err = tl.where(is_better, err, best_err)
best_s = tl.where(is_better, step, best_s)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
def grid_search_triton_mindist(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Format-agnostic Triton grid search using codebook min-distance."""
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
# Compute per-group scale at p=1.0
scale_base, _ = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=args,
global_scale=None,
)
# Build format-agnostic codebook (only need codes, not cutoffs)
codes, _ = _build_normalized_codebook(args)
codes = codes.to(device=observed.device)
_, num_rows, num_groups, group_size = observed.shape
observed_contig = observed.contiguous()
scale_base = scale_base.contiguous()
best_step = torch.zeros(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.full(
(num_rows, num_groups),
float("inf"),
device=observed.device,
dtype=torch.float32,
)
import math
num_codes = codes.shape[0]
BLOCK_G = triton.next_power_of_2(group_size)
LOG_C = max(1, math.ceil(math.log2(num_codes)))
TOTAL_STEPS = triton.next_power_of_2(total_steps)
grid_launch = (num_rows * num_groups,)
_fused_grid_search_mindist_kernel[grid_launch](
observed_contig,
scale_base,
codes,
best_step,
best_error,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
observed_contig.stride(1),
observed_contig.stride(2),
1.0 / grid,
norm,
BLOCK_G=BLOCK_G,
LOG_C=LOG_C,
TOTAL_STEPS=TOTAL_STEPS,
)
# Reconstruct best_min/max from best step indices
ps = torch.tensor(
[1.0 - i / grid for i in range(total_steps)],
dtype=min_val.dtype,
device=min_val.device,
)
best_p = ps[best_step.long()]
best_min_val = min_val * best_p
best_max_val = max_val * best_p
return best_min_val, best_max_val
# ── Triton incremental codebook search (2-check and 3-check variants) ────────
@triton.jit
def _fused_grid_search_incr2_kernel(
observed_ptr,
scale_base_ptr,
codes_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
stride_obs_row,
stride_obs_group,
inv_grid,
norm,
BLOCK_G: tl.constexpr,
LOG_C: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
):
"""Incremental-codebook family: track bins with one directional neighbor.
Performs a full binary search at step 0, then checks
2 codes (current + 1 directional neighbor) per subsequent step.
Since p decreases monotonically, |obs_norm| increases — positive values
shift right in the codebook, negative values shift left.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
)
scale_base = tl.load(scale_base_ptr + row * num_groups + group)
ig = inv_grid.to(tl.float32)
norm_f = norm.to(tl.float32)
# Precompute shift direction: +1 for positive obs, -1 for negative
shift_dir = tl.where(obs >= 0, 1, -1)
best_err = tl.full([], float("inf"), dtype=tl.float32)
best_s = 0
bin_idx = tl.zeros([BLOCK_G], dtype=tl.int32)
for step in range(TOTAL_STEPS):
if step < total_steps:
p = (1.0 - step * ig).to(tl.float32)
eff_scale = tl.maximum(scale_base * p, 1e-38).to(tl.float32)
obs_norm = obs / eff_scale
if step == 0:
# Full binary search for insertion point
lo = tl.zeros([BLOCK_G], dtype=tl.int32)
hi = tl.full([BLOCK_G], num_codes, dtype=tl.int32)
for _ in range(LOG_C):
mid = (lo + hi) >> 1
code_mid = tl.load(codes_ptr + mid)
ge = obs_norm >= code_mid
lo = tl.where(ge, mid + 1, lo)
hi = tl.where(ge, hi, mid)
idx_left = tl.maximum(lo - 1, 0)
idx_right = tl.minimum(lo, num_codes - 1)
d_left = tl.abs(obs_norm - tl.load(codes_ptr + idx_left))
d_right = tl.abs(obs_norm - tl.load(codes_ptr + idx_right))
left_wins = d_left <= d_right
bin_idx = tl.where(left_wins, idx_left, idx_right)
min_dist = tl.minimum(d_left, d_right) * eff_scale
else:
# Incremental: check current and 1 neighbor in shift direction
d_cur = tl.abs(obs_norm - tl.load(codes_ptr + bin_idx))
shift_idx = tl.maximum(
tl.minimum(bin_idx + shift_dir, num_codes - 1), 0
)
d_shift = tl.abs(obs_norm - tl.load(codes_ptr + shift_idx))
better = d_shift < d_cur
bin_idx = tl.where(better, shift_idx, bin_idx)
min_dist = tl.minimum(d_cur, d_shift) * eff_scale
min_dist_pow = tl.extra.cuda.libdevice.pow(min_dist.to(tl.float32), norm_f)
err = tl.sum(tl.where(g_mask, min_dist_pow, 0.0), axis=0)
is_better = err < best_err
best_err = tl.where(is_better, err, best_err)
best_s = tl.where(is_better, step, best_s)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
@triton.jit
def _fused_grid_search_incr3_kernel(
observed_ptr,
scale_base_ptr,
codes_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
stride_obs_row,
stride_obs_group,
inv_grid,
norm,
BLOCK_G: tl.constexpr,
LOG_C: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
):
"""Incremental-codebook family: track bins with two directional neighbors.
Performs a full binary search at step 0, then checks
3 codes (current + 2 in shift direction) per subsequent step.
Since |obs_norm| only increases, the shift is always away from zero:
positive values shift right, negative shift left.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
)
scale_base = tl.load(scale_base_ptr + row * num_groups + group)
ig = inv_grid.to(tl.float32)
norm_f = norm.to(tl.float32)
shift_dir = tl.where(obs >= 0, 1, -1)
best_err = tl.full([], float("inf"), dtype=tl.float32)
best_s = 0
bin_idx = tl.zeros([BLOCK_G], dtype=tl.int32)
for step in range(TOTAL_STEPS):
if step < total_steps:
p = (1.0 - step * ig).to(tl.float32)
eff_scale = tl.maximum(scale_base * p, 1e-38).to(tl.float32)
obs_norm = obs / eff_scale
if step == 0:
lo = tl.zeros([BLOCK_G], dtype=tl.int32)
hi = tl.full([BLOCK_G], num_codes, dtype=tl.int32)
for _ in range(LOG_C):
mid = (lo + hi) >> 1
code_mid = tl.load(codes_ptr + mid)
ge = obs_norm >= code_mid
lo = tl.where(ge, mid + 1, lo)
hi = tl.where(ge, hi, mid)
idx_left = tl.maximum(lo - 1, 0)
idx_right = tl.minimum(lo, num_codes - 1)
d_left = tl.abs(obs_norm - tl.load(codes_ptr + idx_left))
d_right = tl.abs(obs_norm - tl.load(codes_ptr + idx_right))
left_wins = d_left <= d_right
bin_idx = tl.where(left_wins, idx_left, idx_right)
min_dist = tl.minimum(d_left, d_right) * eff_scale
else:
# Check current + 2 neighbors in shift direction
idx_s1 = tl.maximum(tl.minimum(bin_idx + shift_dir, num_codes - 1), 0)
idx_s2 = tl.maximum(
tl.minimum(bin_idx + shift_dir * 2, num_codes - 1), 0
)
d_cur = tl.abs(obs_norm - tl.load(codes_ptr + bin_idx))
d_s1 = tl.abs(obs_norm - tl.load(codes_ptr + idx_s1))
d_s2 = tl.abs(obs_norm - tl.load(codes_ptr + idx_s2))
new_bin = bin_idx
new_d = d_cur
better_s1 = d_s1 < new_d
new_bin = tl.where(better_s1, idx_s1, new_bin)
new_d = tl.where(better_s1, d_s1, new_d)
better_s2 = d_s2 < new_d
new_bin = tl.where(better_s2, idx_s2, new_bin)
new_d = tl.where(better_s2, d_s2, new_d)
bin_idx = new_bin
min_dist = new_d * eff_scale
min_dist_pow = tl.extra.cuda.libdevice.pow(min_dist.to(tl.float32), norm_f)
err = tl.sum(tl.where(g_mask, min_dist_pow, 0.0), axis=0)
is_better = err < best_err
best_err = tl.where(is_better, err, best_err)
best_s = tl.where(is_better, step, best_s)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
@triton.jit
def _fused_grid_search_incrN_patience_kernel(
observed_ptr,
scale_base_ptr,
codes_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
stride_obs_row,
stride_obs_group,
inv_grid,
norm,
patience,
BLOCK_G: tl.constexpr,
LOG_C: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
NUM_NEIGHBORS: tl.constexpr,
):
"""Incremental-codebook family: configurable neighbors plus patience.
NUM_NEIGHBORS controls how many codes to check in the shift direction
after the initial binary search at step 0. Unlike the exhaustive incr2/3
siblings, it stops each group after its error fails to improve for the
requested patience.
"""
pid = tl.program_id(0)
row = pid // num_groups
group = pid % num_groups
if row >= num_rows:
return
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
)
scale_base = tl.load(scale_base_ptr + row * num_groups + group)
ig = inv_grid.to(tl.float32)
norm_f = norm.to(tl.float32)
shift_dir = tl.where(obs >= 0, 1, -1)
best_err = tl.full([], float("inf"), dtype=tl.float32)
best_s = 0
bin_idx = tl.zeros([BLOCK_G], dtype=tl.int32)
patience_ctr = 0
for step in range(TOTAL_STEPS):
if step < total_steps:
if patience_ctr < patience:
p = (1.0 - step * ig).to(tl.float32)
eff_scale = tl.maximum(scale_base * p, 1e-38).to(tl.float32)
obs_norm = obs / eff_scale
if step == 0:
lo = tl.zeros([BLOCK_G], dtype=tl.int32)
hi = tl.full([BLOCK_G], num_codes, dtype=tl.int32)
for _ in range(LOG_C):
mid = (lo + hi) >> 1
code_mid = tl.load(codes_ptr + mid)
ge = obs_norm >= code_mid
lo = tl.where(ge, mid + 1, lo)
hi = tl.where(ge, hi, mid)
idx_left = tl.maximum(lo - 1, 0)
idx_right = tl.minimum(lo, num_codes - 1)
d_left = tl.abs(obs_norm - tl.load(codes_ptr + idx_left))
d_right = tl.abs(obs_norm - tl.load(codes_ptr + idx_right))
left_wins = d_left <= d_right
bin_idx = tl.where(left_wins, idx_left, idx_right)
min_dist = tl.minimum(d_left, d_right) * eff_scale
else:
# Check current + NUM_NEIGHBORS in shift direction
new_bin = bin_idx
new_d = tl.abs(obs_norm - tl.load(codes_ptr + bin_idx))
for k in tl.static_range(1, NUM_NEIGHBORS + 1):
idx_k = tl.maximum(
tl.minimum(bin_idx + shift_dir * k, num_codes - 1), 0
)
d_k = tl.abs(obs_norm - tl.load(codes_ptr + idx_k))
better_k = d_k < new_d
new_bin = tl.where(better_k, idx_k, new_bin)
new_d = tl.where(better_k, d_k, new_d)
bin_idx = new_bin
min_dist = new_d * eff_scale
min_dist_pow = tl.extra.cuda.libdevice.pow(
min_dist.to(tl.float32), norm_f
)
err = tl.sum(tl.where(g_mask, min_dist_pow, 0.0), axis=0)
is_better = err < best_err
best_err = tl.where(is_better, err, best_err)
best_s = tl.where(is_better, step, best_s)
patience_ctr = tl.where(is_better, 0, patience_ctr + 1)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
def _launch_incrN_patience(
observed, args, maxshrink, patience, grid, norm, num_neighbors
):
"""Shared launcher for incrN patience kernel with variable neighbor count."""
import math
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
scale_base, _ = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=args,
global_scale=None,
)
codes, _ = _build_normalized_codebook(args)
codes = codes.to(device=observed.device)
_, num_rows, num_groups, group_size = observed.shape
observed_contig = observed.contiguous()
scale_base = scale_base.contiguous()
best_step = torch.zeros(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.full(
(num_rows, num_groups),
float("inf"),
device=observed.device,
dtype=torch.float32,
)
num_codes = codes.shape[0]
BLOCK_G = triton.next_power_of_2(group_size)
LOG_C = max(1, math.ceil(math.log2(num_codes)))
TOTAL_STEPS = triton.next_power_of_2(total_steps)
grid_launch = (num_rows * num_groups,)
_fused_grid_search_incrN_patience_kernel[grid_launch](
observed_contig,
scale_base,
codes,
best_step,
best_error,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
observed_contig.stride(1),
observed_contig.stride(2),
1.0 / grid,
norm,
patience,
BLOCK_G=BLOCK_G,
LOG_C=LOG_C,
TOTAL_STEPS=TOTAL_STEPS,
NUM_NEIGHBORS=num_neighbors,
)
ps = torch.tensor(
[1.0 - i / grid for i in range(total_steps)],
dtype=min_val.dtype,
device=min_val.device,
)
best_p = ps[best_step.long()]
return min_val * best_p, max_val * best_p
def grid_search_triton_incrNp1(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
return _launch_incrN_patience(observed, args, maxshrink, patience, grid, norm, 1)
def grid_search_triton_incrNp2(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
return _launch_incrN_patience(observed, args, maxshrink, patience, grid, norm, 2)
def grid_search_triton_incrNp3(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
return _launch_incrN_patience(observed, args, maxshrink, patience, grid, norm, 3)
def grid_search_triton_incrNp4(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
return _launch_incrN_patience(observed, args, maxshrink, patience, grid, norm, 4)
def _launch_incremental_kernel(kernel, observed, args, maxshrink, grid, norm):
"""Shared launcher for incremental kernels."""
import math
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
scale_base, _ = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=args,
global_scale=None,
)
codes, _ = _build_normalized_codebook(args)
codes = codes.to(device=observed.device)
_, num_rows, num_groups, group_size = observed.shape
observed_contig = observed.contiguous()
scale_base = scale_base.contiguous()
best_step = torch.zeros(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.full(
(num_rows, num_groups),
float("inf"),
device=observed.device,
dtype=torch.float32,
)
num_codes = codes.shape[0]
BLOCK_G = triton.next_power_of_2(group_size)
LOG_C = max(1, math.ceil(math.log2(num_codes)))
TOTAL_STEPS = triton.next_power_of_2(total_steps)
grid_launch = (num_rows * num_groups,)
kernel[grid_launch](
observed_contig,
scale_base,
codes,
best_step,
best_error,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
observed_contig.stride(1),
observed_contig.stride(2),
1.0 / grid,
norm,
BLOCK_G=BLOCK_G,
LOG_C=LOG_C,
TOTAL_STEPS=TOTAL_STEPS,
)
ps = torch.tensor(
[1.0 - i / grid for i in range(total_steps)],
dtype=min_val.dtype,
device=min_val.device,
)
best_p = ps[best_step.long()]
return min_val * best_p, max_val * best_p
def grid_search_triton_incr2(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Incremental search: binary search at step 0, 2 checks after."""
return _launch_incremental_kernel(
_fused_grid_search_incr2_kernel, observed, args, maxshrink, grid, norm
)
def grid_search_triton_incr3(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Incremental search: binary search at step 0, 3 checks after."""
return _launch_incremental_kernel(
_fused_grid_search_incr3_kernel, observed, args, maxshrink, grid, norm
)
# ── Triton multi-group codebook min-distance ─────────────────────────────────
@triton.jit
def _fused_grid_search_multigroup_kernel(
observed_ptr,
scale_base_ptr,
codes_ptr,
best_step_ptr,
best_error_ptr,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
stride_obs_row,
stride_obs_group,
inv_grid,
norm,
total_work,
BLOCK_G: tl.constexpr,
LOG_C: tl.constexpr,
TOTAL_STEPS: tl.constexpr,
GROUPS_PER_PID: tl.constexpr,
):
"""Codebook family: nearest-distance search batching groups per program.
Quantization matches ``mindist``; the difference is launch geometry. Each
program processes ``GROUPS_PER_PID`` groups serially to amortize program
overhead, rather than assigning exactly one program to each group.
"""
pid = tl.program_id(0)
base_work = pid * GROUPS_PER_PID
g_idx = tl.arange(0, BLOCK_G)
g_mask = g_idx < group_size
ig = inv_grid.to(tl.float32)
norm_f = norm.to(tl.float32)
for g_off in range(GROUPS_PER_PID):
work_idx = base_work + g_off
valid = work_idx < total_work
if valid:
row = work_idx // num_groups
group = work_idx % num_groups
obs = tl.load(
observed_ptr + row * stride_obs_row + group * stride_obs_group + g_idx,
mask=g_mask,
other=0.0,
)
scale_base = tl.load(scale_base_ptr + row * num_groups + group)
best_err = tl.full([], float("inf"), dtype=tl.float32)
best_s = 0
for step in range(TOTAL_STEPS):
if step < total_steps:
p = (1.0 - step * ig).to(tl.float32)
eff_scale = tl.maximum(scale_base * p, 1e-38).to(tl.float32)
obs_norm = obs / eff_scale
lo = tl.zeros([BLOCK_G], dtype=tl.int32)
hi = tl.full([BLOCK_G], num_codes, dtype=tl.int32)
for _ in range(LOG_C):
mid = (lo + hi) >> 1
code_mid = tl.load(codes_ptr + mid)
ge = obs_norm >= code_mid
lo = tl.where(ge, mid + 1, lo)
hi = tl.where(ge, hi, mid)
idx_left = tl.maximum(lo - 1, 0)
idx_right = tl.minimum(lo, num_codes - 1)
code_left = tl.load(codes_ptr + idx_left)
code_right = tl.load(codes_ptr + idx_right)
dist_left = tl.abs(obs_norm - code_left)
dist_right = tl.abs(obs_norm - code_right)
min_dist = tl.minimum(dist_left, dist_right) * eff_scale
min_dist_pow = tl.extra.cuda.libdevice.pow(
min_dist.to(tl.float32), norm_f
)
err = tl.sum(tl.where(g_mask, min_dist_pow, 0.0), axis=0)
is_better = err < best_err
best_err = tl.where(is_better, err, best_err)
best_s = tl.where(is_better, step, best_s)
tl.store(best_step_ptr + row * num_groups + group, best_s)
tl.store(best_error_ptr + row * num_groups + group, best_err)
def grid_search_triton_multigroup(
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size
):
"""Format-agnostic Triton grid search — multiple groups per program."""
min_val = torch.amin(observed, dim=(0, -1))
max_val = torch.amax(observed, dim=(0, -1))
total_steps = int(maxshrink * grid)
scale_base, _ = calculate_qparams(
min_vals=min_val,
max_vals=max_val,
quantization_args=args,
global_scale=None,
)
codes, _ = _build_normalized_codebook(args)
codes = codes.to(device=observed.device)
_, num_rows, num_groups, group_size = observed.shape
observed_contig = observed.contiguous()
scale_base = scale_base.contiguous()
best_step = torch.zeros(
num_rows, num_groups, dtype=torch.int32, device=observed.device
)
best_error = torch.full(
(num_rows, num_groups),
float("inf"),
device=observed.device,
dtype=torch.float32,
)
import math
num_codes = codes.shape[0]
BLOCK_G = triton.next_power_of_2(group_size)
LOG_C = max(1, math.ceil(math.log2(num_codes)))
TOTAL_STEPS = triton.next_power_of_2(total_steps)
total_work = num_rows * num_groups
GROUPS_PER_PID = 4
num_pids = math.ceil(total_work / GROUPS_PER_PID)
grid_launch = (num_pids,)
_fused_grid_search_multigroup_kernel[grid_launch](
observed_contig,
scale_base,
codes,
best_step,
best_error,
num_rows,
num_groups,
group_size,
total_steps,
num_codes,
observed_contig.stride(1),
observed_contig.stride(2),
1.0 / grid,
norm,
total_work,
BLOCK_G=BLOCK_G,
LOG_C=LOG_C,
TOTAL_STEPS=TOTAL_STEPS,
GROUPS_PER_PID=GROUPS_PER_PID,
)
ps = torch.tensor(
[1.0 - i / grid for i in range(total_steps)],
dtype=min_val.dtype,
device=min_val.device,
)
best_p = ps[best_step.long()]
best_min_val = min_val * best_p
best_max_val = max_val * best_p
return best_min_val, best_max_val
# ── Benchmark infrastructure ─────────────────────────────────────────────────
def make_observer_inputs(
rows,
cols,
strategy,
group_size,
num_bits,
device,
scale_dtype=None,
quant_type=QuantizationType.INT,
):
"""Create a weight tensor and flatten it for calibration."""
torch.manual_seed(42)
module = torch.nn.Linear(cols, rows, bias=False, device=device)
if strategy == "group":
quant_args = QuantizationArgs(
num_bits=num_bits,
type=quant_type,
symmetric=True,
strategy=QuantizationStrategy.GROUP,
group_size=group_size,
scale_dtype=scale_dtype,
)
elif strategy == "channel":
quant_args = QuantizationArgs(
num_bits=num_bits,
type=quant_type,
symmetric=True,
strategy=QuantizationStrategy.CHANNEL,
scale_dtype=scale_dtype,
)
elif strategy == "tensor":
quant_args = QuantizationArgs(
num_bits=num_bits,
type=quant_type,
symmetric=True,
strategy=QuantizationStrategy.TENSOR,
scale_dtype=scale_dtype,
)
else:
raise ValueError(f"Unknown strategy: {strategy}")
token_args = quant_args.model_copy(update={"strategy": QuantizationStrategy.TOKEN})
observed = flatten_for_calibration(module.weight, "weight", quant_args)
return observed, quant_args, token_args
def time_fn(fn, inputs, warmup, iters):
"""Time a grid search function."""
observed, args, token_args, maxshrink, patience, grid, norm, chunk_size = inputs
is_cuda = observed.is_cuda
for _ in range(warmup):
fn(observed, args, token_args, maxshrink, patience, grid, norm, chunk_size)
if is_cuda:
torch.cuda.synchronize()
if is_cuda:
torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
mem_before = torch.cuda.max_memory_allocated()
times = []
for _ in range(iters):
if is_cuda:
torch.cuda.synchronize()
t0 = time.perf_counter()
fn(observed, args, token_args, maxshrink, patience, grid, norm, chunk_size)
if is_cuda:
torch.cuda.synchronize()
times.append(time.perf_counter() - t0)
peak_delta_mb = 0.0
if is_cuda:
peak_delta_mb = (torch.cuda.max_memory_allocated() - mem_before) / (1024 * 1024)
return times, peak_delta_mb
def main():
parser = argparse.ArgumentParser(description="Benchmark MSE observer grid search")
parser.add_argument(
"--device", default="cuda" if torch.cuda.is_available() else "cpu"
)
parser.add_argument("--rows", type=int, default=14336)
parser.add_argument("--cols", type=int, default=256 * 16)
parser.add_argument("--group-size", type=int, default=16)
parser.add_argument("--num-bits", type=int, default=4)
parser.add_argument("--quant-type", default="float", choices=["int", "float"])
parser.add_argument(
"--scale-dtype",
default="none",
choices=["none", "float16", "bfloat16", "float8_e4m3fn", "uint8"],
)
parser.add_argument(
"--strategy", default="group", choices=["group", "channel", "tensor"]
)
parser.add_argument("--maxshrink", type=float, default=0.20)
parser.add_argument("--patience", type=int, default=5)
parser.add_argument("--grid", type=float, default=100.0)
parser.add_argument("--norm", type=float, default=2.4)
parser.add_argument("--chunk-size", type=int, default=5)
parser.add_argument("--hier-n", type=int, default=16)
parser.add_argument("--hier-keep-num", type=int, default=1)
parser.add_argument("--hier-keep-den", type=int, default=4)
parser.add_argument("--hier-m", type=int, default=2)
parser.add_argument(
"--hier-k",
type=int,
default=3,
help="Number of refinement generations after the initial candidates",
)
parser.add_argument("--hier-p-min", type=float, default=0.8)
parser.add_argument("--hier-p-max", type=float, default=1.8)
parser.add_argument("--hier-fast-pow", action="store_true")
parser.add_argument("--warmup", type=int, default=WARMUP)
parser.add_argument("--iters", type=int, default=ITERS)
parser.add_argument(
"--variants",
nargs="+",
default=None,
help="Optional benchmark variant names to run",
)
args = parser.parse_args()
HIER_CONFIG.update(
n=args.hier_n,
keep_num=args.hier_keep_num,
keep_den=args.hier_keep_den,
m=args.hier_m,
k=args.hier_k,
p_min=args.hier_p_min,
p_max=args.hier_p_max,
fast_pow=args.hier_fast_pow,
)
scale_dtype = {
"none": None,
"float16": torch.float16,
"bfloat16": torch.bfloat16,
"float8_e4m3fn": torch.float8_e4m3fn,
"uint8": torch.uint8,
}[args.scale_dtype]
quant_type = {
"int": QuantizationType.INT,
"float": QuantizationType.FLOAT,
}[args.quant_type]
total_steps = int(args.maxshrink * args.grid)
print(f"Device: {args.device}")
print(f"Weight: ({args.rows}, {args.cols})")
print(f"Strategy: {args.strategy}")
print(f"Group size: {args.group_size}")
print(f"Bits: {args.num_bits}")
print(f"Quant type: {args.quant_type}")
print(f"Scale dtype: {args.scale_dtype}")
print(f"Maxshrink: {args.maxshrink}")
print(f"Grid: {args.grid}")
print(f"Total steps: {total_steps}")
print(f"Patience: {args.patience}")
print(f"Norm: {args.norm}")
print(f"Chunk size: {args.chunk_size}")
hierarchical_counts = [args.hier_n]
for _ in range(args.hier_k):
hierarchical_counts.append(
hierarchical_counts[-1]
* args.hier_keep_num
// args.hier_keep_den
* args.hier_m
)
print(
"Hierarchical: "
f"counts={hierarchical_counts}, total={sum(hierarchical_counts)}, "
f"range=[{args.hier_p_min}, {args.hier_p_max}]"
)
print(f"Hier fast pow: {args.hier_fast_pow}")
print(f"Warmup: {args.warmup} Iters: {args.iters}")
print()
observed, quant_args, token_args = make_observer_inputs(
args.rows,
args.cols,
args.strategy,
args.group_size,
args.num_bits,
args.device,
scale_dtype,
quant_type,
)
print(f"Observed shape: {observed.shape}")
print()
inputs = (
observed,
quant_args,
token_args,
args.maxshrink,
args.patience,
args.grid,
args.norm,
args.chunk_size,
)
if scale_dtype is None:
scale_variant = ("triton_gptq_base_p", grid_search_triton_gptq_base_p)
elif scale_dtype == torch.uint8:
scale_variant = ("triton_gptq_e8_p", grid_search_triton_gptq_e8_p)
else:
scale_variant = (
"triton_gptq_scale_round_p",
grid_search_triton_gptq_scale_round_p,
)
variants = [
("eager", grid_search_eager),
("eager_no_patience", grid_search_eager_no_patience),
# ("compiled", grid_search_compiled),
("triton", grid_search_triton),
("triton_gptq_quant", grid_search_triton_gptq_quant),
("triton_gptq_quant_p", grid_search_triton_gptq_quant_p),
("triton_gptq_vec", grid_search_triton_gptq_quant_vec),
("triton_gptq_base", grid_search_triton_gptq_base),
("triton_gptq_grid_major", grid_search_triton_gptq_grid_major),
("triton_gptq_grid_major_tile64", grid_search_triton_gptq_grid_major_tile64),
("triton_gptq_grid_major_tile128", grid_search_triton_gptq_grid_major_tile128),
("triton_gptq_grid_major_tile256", grid_search_triton_gptq_grid_major_tile256),
("triton_gptq_grid_major_tile512", grid_search_triton_gptq_grid_major_tile512),
(
"triton_gptq_grid_major_tile1024",
grid_search_triton_gptq_grid_major_tile1024,
),
(
"triton_gptq_grid_major_tile2048",
grid_search_triton_gptq_grid_major_tile2048,
),
(
"triton_gptq_grid_major_tile4096",
grid_search_triton_gptq_grid_major_tile4096,
),
(
"triton_gptq_grid_major_tile8192",
grid_search_triton_gptq_grid_major_tile8192,
),
(
"triton_gptq_grid_major_tile16384",
grid_search_triton_gptq_grid_major_tile16384,
),
("triton_gptq_base_tile128", grid_search_triton_gptq_base_tile128),
scale_variant,
("triton_gptq_base_p_b", grid_search_triton_gptq_base_p_b),
("triton_gptq_base_p_b_tile64", grid_search_triton_gptq_base_p_b_tile64),
("triton_gptq_base_p_b_tile128", grid_search_triton_gptq_base_p_b_tile128),
("triton_gptq_base_p_b_tile256", grid_search_triton_gptq_base_p_b_tile256),
("triton_gptq_base_p_b_tile512", grid_search_triton_gptq_base_p_b_tile512),
(
"triton_gptq_base_p_b_tile512_adaptive",
grid_search_triton_gptq_base_p_b_tile512_adaptive,
),
# ("triton_gptq_base_p_b_tile640", grid_search_triton_gptq_base_p_b_tile640),
# ("triton_gptq_base_p_b_tile768", grid_search_triton_gptq_base_p_b_tile768),
# ("triton_gptq_base_p_b_tile896", grid_search_triton_gptq_base_p_b_tile896),
("triton_gptq_base_p_b_tile1024", grid_search_triton_gptq_base_p_b_tile1024),
# ("triton_hierarchical", grid_search_triton_hierarchical),
# ("triton_hierarchical_fast", grid_search_triton_hierarchical_fast),
# ("triton_hierarchical_serial", grid_search_triton_hierarchical_serial),
# ("triton_hierarchical_tiled", grid_search_triton_hierarchical_tiled),
# ("triton_cutoff", grid_search_triton_cutoff),
# ("triton_mindist", grid_search_triton_mindist),
# ("triton_multigrp", grid_search_triton_multigroup),
# ("triton_incr2", grid_search_triton_incr2),
# ("triton_incr3", grid_search_triton_incr3),
# ("triton_incrNp1", grid_search_triton_incrNp1),
# ("triton_incrNp2", grid_search_triton_incrNp2),
# ("triton_incrNp3", grid_search_triton_incrNp3),
# ("triton_incrNp4", grid_search_triton_incrNp4),
]
if args.variants is not None:
requested = set(args.variants)
known = {name for name, _ in variants}
if unknown := requested - known:
raise ValueError(f"Unknown variants: {sorted(unknown)}")
variants = [(name, fn) for name, fn in variants if name in requested]
results = {}
for name, fn in variants:
print(f"Running {name} ...")
times, peak_mb = time_fn(fn, inputs, args.warmup, args.iters)
results[name] = {"times": times, "peak_mb": peak_mb}
medians = {n: sorted(r["times"])[len(r["times"]) // 2] for n, r in results.items()}
name_width = max(16, *(len(name) for name in results))
print()
print(
f"{'':>{name_width}} {'median':>10} {'min':>10} "
f"{'max':>10} {'peak_mem':>10}"
)
for name in results:
t = results[name]["times"]
med = medians[name]
peak = results[name]["peak_mb"]
mem_str = f"{peak:.1f}MB" if peak > 0 else "n/a"
print(
f"{name:>{name_width}} {med:>10.4f}s "
f"{min(t):>10.4f}s {max(t):>10.4f}s "
f"{mem_str:>10}"
)
baseline = "eager_no_patience"
if len(results) > 1 and baseline in medians:
print()
base = medians[baseline]
for name in results:
if name == baseline:
continue
if medians[name] > 0:
print(f"Speedup {name} vs {baseline}: " f"{base / medians[name]:.2f}x")
if __name__ == "__main__":
with torch.no_grad():
main()
@HDCharles

HDCharles commented Sep 14, 2026

Copy link
Copy Markdown
Author 
                                          median        min        max   peak_mem
                                eager     0.0578s     0.0578s     0.0581s   1029.0MB
                    eager_no_patience     0.0588s     0.0585s     0.0591s   1029.0MB
                               triton     0.0485s     0.0485s     0.0486s    742.0MB
                    triton_gptq_quant     0.0547s     0.0545s     0.0551s    742.0MB
                  triton_gptq_quant_p     0.0370s     0.0363s     0.0373s    742.0MB
                      triton_gptq_vec     0.0161s     0.0161s     0.0162s    742.0MB
                     triton_gptq_base     0.0474s     0.0474s     0.0476s    126.0MB
               triton_gptq_grid_major     0.0729s     0.0729s     0.0730s    126.0MB
        triton_gptq_grid_major_tile64     0.0208s     0.0208s     0.0237s    126.0MB
       triton_gptq_grid_major_tile128     0.0124s     0.0123s     0.0124s    126.0MB
       triton_gptq_grid_major_tile256     0.0092s     0.0092s     0.0092s    126.0MB
       triton_gptq_grid_major_tile512     0.0077s     0.0077s     0.0078s    126.0MB
      triton_gptq_grid_major_tile1024     0.0076s     0.0076s     0.0076s    126.0MB
      triton_gptq_grid_major_tile2048     0.0076s     0.0075s     0.0076s    126.0MB
      triton_gptq_grid_major_tile4096     0.0094s     0.0094s     0.0095s    126.0MB
      triton_gptq_grid_major_tile8192     0.0208s     0.0207s     0.0208s    126.0MB
     triton_gptq_grid_major_tile16384     0.0224s     0.0223s     0.0225s    126.0MB
             triton_gptq_base_tile128     0.0071s     0.0071s     0.0072s    126.0MB
                   triton_gptq_base_p     0.0304s     0.0304s     0.0305s    126.0MB
                 triton_gptq_base_p_b     0.0408s     0.0408s     0.0408s    126.0MB
          triton_gptq_base_p_b_tile64     0.0136s     0.0136s     0.0137s    126.0MB
         triton_gptq_base_p_b_tile128     0.0073s     0.0073s     0.0073s    126.0MB
         triton_gptq_base_p_b_tile256     0.0063s     0.0063s     0.0064s    126.0MB
         triton_gptq_base_p_b_tile512     0.0059s     0.0059s     0.0060s    126.0MB
triton_gptq_base_p_b_tile512_adaptive     0.0059s     0.0059s     0.0059s    126.0MB
        triton_gptq_base_p_b_tile1024     0.0060s     0.0059s     0.0060s    126.0MB

Speedup eager vs eager_no_patience: 1.02x
Speedup triton vs eager_no_patience: 1.21x
Speedup triton_gptq_quant vs eager_no_patience: 1.07x
Speedup triton_gptq_quant_p vs eager_no_patience: 1.59x
Speedup triton_gptq_vec vs eager_no_patience: 3.64x
Speedup triton_gptq_base vs eager_no_patience: 1.24x
Speedup triton_gptq_grid_major vs eager_no_patience: 0.81x
Speedup triton_gptq_grid_major_tile64 vs eager_no_patience: 2.83x
Speedup triton_gptq_grid_major_tile128 vs eager_no_patience: 4.75x
Speedup triton_gptq_grid_major_tile256 vs eager_no_patience: 6.39x
Speedup triton_gptq_grid_major_tile512 vs eager_no_patience: 7.61x
Speedup triton_gptq_grid_major_tile1024 vs eager_no_patience: 7.74x
Speedup triton_gptq_grid_major_tile2048 vs eager_no_patience: 7.79x
Speedup triton_gptq_grid_major_tile4096 vs eager_no_patience: 6.23x
Speedup triton_gptq_grid_major_tile8192 vs eager_no_patience: 2.83x
Speedup triton_gptq_grid_major_tile16384 vs eager_no_patience: 2.63x
Speedup triton_gptq_base_tile128 vs eager_no_patience: 8.25x
Speedup triton_gptq_base_p vs eager_no_patience: 1.93x
Speedup triton_gptq_base_p_b vs eager_no_patience: 1.44x
Speedup triton_gptq_base_p_b_tile64 vs eager_no_patience: 4.32x
Speedup triton_gptq_base_p_b_tile128 vs eager_no_patience: 8.05x
Speedup triton_gptq_base_p_b_tile256 vs eager_no_patience: 9.30x
Speedup triton_gptq_base_p_b_tile512 vs eager_no_patience: 9.92x
Speedup triton_gptq_base_p_b_tile512_adaptive vs eager_no_patience: 9.97x <-
Speedup triton_gptq_base_p_b_tile1024 vs eager_no_patience: 9.86x

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment