Add FP8 KV Support

[yiliu30]

No description provided.

[Copilot]

Pull Request Overview

This PR adds FP8 KV cache quantization support to the AutoRound quantization library. This feature enables quantizing the key-value cache in attention mechanisms to FP8 format for improved memory efficiency.

• Adds a new enable_fp8_kv parameter to control FP8 KV cache quantization
• Implements FP8 KV cache infrastructure with calibration and quantization context
• Updates test coverage to verify FP8 KV cache serialization and functionality

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
auto_round/autoround.py	Adds enable_fp8_kv parameter and context manager integration
auto_round/experimental/fp8_kv_cache.py	Core FP8 KV cache implementation with quantization logic
test/test_cpu/test_export.py	Test updates to verify FP8 KV cache export functionality

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Copilot]

Adding a logger to sys.stderr at TRACE level in a module import can interfere with application-level logging configuration. Consider removing this global logger configuration or making it conditional.

Suggested change

[Copilot]

The FIXME comment indicates uncertainty about the intended behavior. This should be resolved before production use, as it could affect the correctness of the KV cache implementation.

[Copilot]

The attention module detection logic uses a simple string check and attribute existence. This brittle approach should be improved with a more robust detection mechanism.

Suggested change

	)
	"""
	Robustly detect if a module is an attention module.
	Checks for known attention base classes and required attributes.
	"""
	# Check for known attention base classes
	attention_bases = tuple(
	cls for cls in (MultiheadAttention, LlamaAttention) if cls is not None
	)
	if attention_bases and isinstance(module, attention_bases):
	return True
	# Fallback: check for common attention attributes and method signatures
	has_proj = any(
	hasattr(module, attr)
	for attr in ("k_proj", "v_proj", "qkv_proj", "in_proj_weight", "out_proj")
	)
	has_forward = hasattr(module, "forward")
	# Optionally, check for typical attention input signatures
	return has_proj and has_forward

[Copilot]

The TODO comment is incomplete and doesn't provide meaningful information. It should either be completed with specific details or removed.

Suggested change

# TODO: Triggered by adding kv_cache_scheme in ...

[wenhuach21]

Please note that kv cache does not affect tuning as we only use forward in tuning. So the first change is moving the arg from init to save_quantized?