Skip to content

feat: add per-head fp8 kv attention calibration#2044

Open
zeel2104 wants to merge 3 commits into
intel:mainfrom
zeel2104:feat/fp8-per-head-kv-attn
Open

feat: add per-head fp8 kv attention calibration#2044
zeel2104 wants to merge 3 commits into
intel:mainfrom
zeel2104:feat/fp8-per-head-kv-attn

Conversation

@zeel2104

Copy link
Copy Markdown

Description

Adds per-head FP8 KV/attention calibration support on top of the existing per-tensor static FP8 paths.

Main changes:

  • Adds static_kv_granularity and static_attention_granularity options with "tensor" as the default and "head" as the new per-head mode.
  • Collects per-head k_scale, v_scale, and q_scale for attention tensors shaped as [batch, heads, seq, head_dim].
  • Exports LLMC metadata with attn_head strategy for per-head KV/attention scales.
  • Preserves existing per-tensor behavior by default.
  • Adds tests for helper MHA/GQA scale shapes, LLMC export, and AutoRound export.
  • Fixes a Windows-only test file-lock issue by closing safe_open handles before reusing the output path.

Type of Change

New feature

Related Issues

Relates to #2026

Checklist Before Submitting

  • My code has been tested locally.
  • Documentation has been updated as needed.
  • New or updated tests are included where applicable.
  • The CUDA CI has passed. You can trigger it by commenting /azp run Unit-Test-CUDA-AutoRound.

Tested locally:

  • python -m pytest test/test_cpu/test_experimental_fp8_granularity.py -q
  • python -m pytest test/test_cpu/export/test_llmc_format.py -q
  • python -m pytest test/test_cpu/export/test_export.py -k "per_head or static_afp8_export or static_fp8_attn" -q
  • python -m pytest test/test_cpu/quantization/test_mxfp_nvfp.py -k "fp8_kv_attn" -q

Note:

  • test/test_cpu/integrations/test_inc_integration.py was not run because neural_compressor is not installed locally.

@zeel2104

Copy link
Copy Markdown
Author

/azp run Unit-Test-CUDA-AutoRound

@azure-pipelines

Copy link
Copy Markdown
Commenter does not have sufficient privileges for PR 2044 in repo intel/auto-round

@zeel2104
zeel2104 force-pushed the feat/fp8-per-head-kv-attn branch from 04e19a6 to 43cd2cb Compare July 10, 2026 17:21
Signed-off-by: Zeel <desaizeel2128@gmail.com>
@zeel2104
zeel2104 force-pushed the feat/fp8-per-head-kv-attn branch from 43cd2cb to d2d95fb Compare July 10, 2026 17:21
@yiliu30

yiliu30 commented Jul 12, 2026

Copy link
Copy Markdown
Contributor

/azp run Unit-Test-CUDA-AutoRound

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 1 pipeline(s).

@yiliu30

yiliu30 commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

/azp run Unit-Test-CUDA-AutoRound

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 1 pipeline(s).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants