[Feature]: Extend the FP8 KV/Attn Calibration to Per-head

[yiliu30]

Feature Description

Extend FP8 KV/Attention calibration from per-tensor scaling to per-head scaling.

Context

We previously added support for collecting per-tensor scales for FP8 KV/Attention calibration, including usage examples and unit tests:

auto-round/test/test_cpu/export/test_llmc_format.py

Lines 131 to 152 in 5e586bc

	def test_mxfp8_llmcompressor_kv_config(self, tiny_opt_model_path, tmp_path):
	ar = AutoRound(
	model=tiny_opt_model_path,
	iters=0,
	disable_opt_rtn=True,
	scheme="mxfp8",
	static_kv_dtype="fp8",
	)
	_, quantized_model_path = ar.quantize_and_save(output_dir=tmp_path, format="llm_compressor")
	with open(os.path.join(quantized_model_path, "config.json")) as f:
	config = json.load(f)
	kv_cache_scheme = config["quantization_config"]["kv_cache_scheme"]
	assert kv_cache_scheme is not None
	assert kv_cache_scheme["num_bits"] == 8
	assert kv_cache_scheme["type"] == "float"
	assert kv_cache_scheme["strategy"] == "tensor"
	assert kv_cache_scheme["dynamic"] is False
	assert kv_cache_scheme["symmetric"] is True
	def test_mxfp8_llmcompressor_attention_config(self, tiny_opt_model_path, tmp_path):

This feature extends the calibration granularity to per-head scaling, meaning each attention head has its own scale. Compared with per-tensor scaling, per-head scaling can better capture head-wise activation ranges and may improve accuracy while keeping inference overhead minimal, assuming the target backend supports per-head scales efficiently.

Definition of Done

• Support per-head FP8 KV/Attention calibration for:

  • MHA
  • GQA
  • MLA

• Support export formats:

  • LLMC
  • AutoRound

• For the LLMC format, validate accuracy on the vLLM side using the Triton backend.

Additional Context

N/A

[zeel2104]

I would like to work on it. Do I need to get assigned for it?

[yiliu30]

Hi @zeel2104, thanks for showing interest in this task. I’ve updated the description with more details.

We can break this into several PRs. The priority should be MHA/GQA support and the LLMC export format first.
If that works for you, I can assign it to you.

[zeel2104]

"Hi @yiliu30, yes, that sounds like a great plan! I will focus on MHA/GQA support and the LLMC export format first. I'll dive into the codebase and set up the calibration environment once assigned"

[yiliu30]

@zeel2104 great! Here are some PRs adding the per-tensor support:

• Add Static FP8 KV Support #737
• Add static FP8 attention support #1045
• Add FP8 KV Support for DS #1180