quantized_weight_export

Functional export of quantized weight tensors.

Functions

build_hf_quantization_config	Build canonical HF configuration, rejecting unrepresentable activation settings.
capture_kv_cache_export_state	Capture the k_scale / v_scale of one attention module, or None when not quantized.
capture_quantized_weight_export_state	Capture detached state for one weight, or None when it is not quantized.
export_quantized_weight_tensors	Pack a logical weight into canonical ModelOpt checkpoint tensors.
get_kv_cache_export_spec	Return the KV-cache format of one attention module, or None when it is not quantized.
get_quantized_weight_export_spec	Return stable deployment metadata without retaining quantizer tensors.
merge_quantized_weight_export_states	Merge state for logical weight shards concatenated along weight_dim.
restore_quantized_weight_export_state	Restore opaque export state after its tensor values have been transported.
select_quantized_weight_export_state	Select logical indices that preserve complete quantization blocks.
split_quantized_weight_export_state	Separate opaque state metadata from tensor values for typed transport.

build_hf_quantization_config(named_states, kv_cache_formats=())

Build canonical HF configuration, rejecting unrepresentable activation settings.

kv_cache_formats holds get_kv_cache_export_spec() of the attention modules; they must agree, and None (an unquantized cache) is allowed alongside one format.

Parameters:

• named_states (Mapping[str, _QuantizedWeightExportState | _QuantizedWeightExportSpec | None] | Iterable[tuple[str, _QuantizedWeightExportState | _QuantizedWeightExportSpec | None]])

• kv_cache_formats (Iterable[str | None])

Return type:

dict[str, Any]

capture_kv_cache_export_state(attention, *, clamp_fp8_scales=True, cpu=True)

Capture the k_scale / v_scale of one attention module, or None when not quantized.

clamp_fp8_scales raises FP8 cache scales to at least 1.0, as the other exporters do.

Parameters:

• attention (Module)

• clamp_fp8_scales (bool)

• cpu (bool)

Return type:

_KVCacheExportState | None

capture_quantized_weight_export_state(module, weight_name='weight', *, cpu=True)

Capture detached state for one weight, or None when it is not quantized.

Parameters:

• module (Module)

• weight_name (str)

• cpu (bool)

Return type:

_QuantizedWeightExportState | None

export_quantized_weight_tensors(weight, state, dtype, weight_name='weight')

Pack a logical weight into canonical ModelOpt checkpoint tensors.

Parameters:

• weight (Tensor)

• state (_QuantizedWeightExportState)

• dtype (dtype)

• weight_name (str)

Return type:

OrderedDict[str, Tensor]

get_kv_cache_export_spec(attention)

Return the KV-cache format of one attention module, or None when it is not quantized.

Parameters:

attention (Module)

Return type:

str | None

get_quantized_weight_export_spec(module, weight_name='weight')

Return stable deployment metadata without retaining quantizer tensors.

Parameters:

• module (Module)

• weight_name (str)

Return type:

_QuantizedWeightExportSpec | None

merge_quantized_weight_export_states(states, weight_dim)

Merge state for logical weight shards concatenated along weight_dim.

Parameters:

• states (Sequence[_QuantizedWeightExportState])

• weight_dim (int)

Return type:

_QuantizedWeightExportState

restore_quantized_weight_export_state(metadata, tensors)

Restore opaque export state after its tensor values have been transported.

Parameters:

• metadata (object)

• tensors (Sequence[Tensor])

Return type:

_QuantizedWeightExportState

select_quantized_weight_export_state(state, weight_dim, indices)

Select logical indices that preserve complete quantization blocks.

Parameters:

• state (_QuantizedWeightExportState)

• weight_dim (int)

• indices (Iterable[int] | Tensor)

Return type:

_QuantizedWeightExportState

split_quantized_weight_export_state(state)

Separate opaque state metadata from tensor values for typed transport.

Parameters:

state (_QuantizedWeightExportState)

Return type:

tuple[object, tuple[Tensor, …]]