diff --git a/frontier/config/quantization_manager.py b/frontier/config/quantization_manager.py --- a/frontier/config/quantization_manager.py +++ b/frontier/config/quantization_manager.py @@ -356,9 +356,17 @@ class QuantizationManager: target_precision = self.get_precision(op_name) precision_match = target_precision == profiling_precision quant_match = profiling_quant_signature == expected_quant_signature + # A mixed-precision profiling CSV records the model's output dtype at + # file level. Quantized operators in that CSV were nevertheless + # measured with the exact quantization scheme identified by the + # quant signature. Treating those samples as BF16 and scaling them + # again would double-apply the FP8 speedup. + exact_fp8_profile = quant_match and target_precision == PrecisionType.FP8 with self._lock: - self._operation_profiling_precision[op_name] = profiling_precision - if precision_match: + self._operation_profiling_precision[op_name] = ( + target_precision if exact_fp8_profile else profiling_precision + ) + if precision_match or exact_fp8_profile: self._operation_data_sources[op_name] = "profiling" self._operation_approximation_factors.pop(op_name, None) self._operation_speedup_factors.pop(op_name, None) diff --git a/tests/unit/test_quantization_profile_contract.py b/tests/unit/test_quantization_profile_contract.py new file mode 100644 --- /dev/null +++ b/tests/unit/test_quantization_profile_contract.py @@ -0,0 +1,61 @@ +from unittest.mock import MagicMock + +import pytest + +from frontier.config.model_config import QuantizationConfig +from frontier.config.precision_type import PrecisionType +from frontier.config.quantization_manager import QuantizationManager + + +def test_exact_fp8_quant_signature_does_not_rescale_mixed_profile() -> None: + manager = QuantizationManager() + manager.load_config() + + quant_config = QuantizationConfig( + quant_method="fp8", + activation_scheme="dynamic", + is_checkpoint_fp8_serialized=True, + weight_block_size=(128, 128), + ) + quant_signature = quant_config.get_quant_signature() + model_config = MagicMock() + model_config.get_default_precision.return_value = PrecisionType.BF16 + model_config.get_name.return_value = "Qwen3-235B-A22B-FP8" + model_config.torch_dtype = "bfloat16" + model_config.quantization_config = quant_config + model_config.get_quant_signature.return_value = quant_signature + + manager.configure_from_model_config(model_config) + manager.register_profiling_metadata( + operation_names=["attn_pre_proj"], + profiling_precision=PrecisionType.BF16, + profiling_quant_signature=quant_signature, + expected_quant_signature=quant_signature, + file_path="linear_op.csv", + ) + + metadata = { + item["operation"]: item + for item in manager.get_operation_precision_metadata() + } + assert metadata["attn_pre_proj"]["data_source"] == "profiling" + assert metadata["attn_pre_proj"]["approximation_factor"] is None + assert manager.has_precision_mismatch("attn_pre_proj") is False + assert manager.adjust_compute_time("attn_pre_proj", 1.25) == pytest.approx(1.25) + + +def test_mismatched_fp8_quant_signature_still_uses_approximation() -> None: + manager = QuantizationManager() + manager.load_config() + manager._operation_precisions["attn_pre_proj"] = PrecisionType.FP8 + + manager.register_profiling_metadata( + operation_names=["attn_pre_proj"], + profiling_precision=PrecisionType.BF16, + profiling_quant_signature="none", + expected_quant_signature="method=fp8", + file_path="linear_op.csv", + ) + + assert manager.has_precision_mismatch("attn_pre_proj") is True + assert manager.adjust_compute_time("attn_pre_proj", 1.0) == pytest.approx(0.5)