kv_cache_auto_quant

Layer-wise KV-cache AutoQuant using isolated forward KL sensitivity.

Functions

auto_quantize_kv_cache

Select one supplied K/V format per attention layer using isolated forward KL.
auto_quantize_kv_cache(model, constraints, quantization_formats, data_loader, forward_step, *, num_calib_steps, num_score_steps, disabled_layers=None, verbose=False, checkpoint=None)

Select one supplied K/V format per attention layer using isolated forward KL.

Candidate formats are format-agnostic QuantizeConfig dictionaries. Each must configure K and V together and declare effective_bits matching its packed storage per K-or-V scalar, including scale overhead. Candidate calibration is scoped to the candidate K/V quantizers, while pre-existing fixed quantizers keep executing with frozen state. Persistent constant_amax formats may skip calibration forwards.

Parameters:

  • model (nn.Module)

  • constraints (dict[str, Any])

  • quantization_formats (list[tuple[dict[str, Any], str]])

  • data_loader (Iterable)

  • forward_step (Callable[[nn.Module, Any], torch.Tensor])

  • num_calib_steps (int)

  • num_score_steps (int)

  • disabled_layers (list[str] | str | None)

  • verbose (bool)

  • checkpoint (str | None)

Return type:

tuple[nn.Module, dict[str, Any]]