Skip to content

feat(quantization): add activation export handler hook - #82

Open
dengqiaoyu wants to merge 1 commit into
apple:mainfrom
dengqiaoyu:u/qiaoyu_deng/activation-handler-hook
Open

feat(quantization): add activation export handler hook#82
dengqiaoyu wants to merge 1 commit into
apple:mainfrom
dengqiaoyu:u/qiaoyu_deng/activation-handler-hook

Conversation

@dengqiaoyu

Copy link
Copy Markdown
Contributor

feat(quantization): add activation export handler hook

In order to support the export handler, I need to make the following refactor so that handler insertion and testing shares more code and is more clean.

  • Stop finalize from rejecting a dynamic qparams calculator when a handler covers it. _validate_no_persistent_observer_calculators now asks can_export_stateless_fake_quant first, so the restriction applies only where no handler is registered.
  • Share the qparam extraction and the activation validation between the eager and graph export paths, through extract_export_qparams, dequant_output_dtype and validate_activation_export_supported in _export_utils. Both paths previously carried their own copy.
  • Rename FP_DTYPE_TO_MAX_POW2 to E8M0_TARGET_MAX_POW2 and make it the source of truth in QuantizationSpec.validate_scale_dtype. The validator now rejects any dtype with no e8m0 target exponent, so fp16, fp32 and bf16 with an e8m0 scale fail at spec construction rather than later inside _compute_e8m0_scale.
  • Extract the export test workflow into export_utils.run_export_test, a common function now shared by the eager-MIL, eager-MLIR and graph-mode-MLIR test files; it reads the execution mode from the config, so one helper covers both modes.
- Let an out-of-tree package register a per-granularity activation export handler for graph or eager mode, so a granularity the built-in Core AI path cannot express (per-block, dynamic qparams) can be exported without patching this repo.
- Allow an e8m0 `scale_dtype` on integer dtypes, gated in `QuantizationSpec.validate_scale_dtype` to `PerBlockGranularity` + `SYMMETRIC` + `ZP`, with `E8M0_TARGET_MAX_POW2` as the single source of truth for which dtypes have an e8m0 target exponent.
- Extract the export test workflow into `export_utils.run_export_test`, a common function now shared by the eager-MIL, eager-MLIR and graph-mode-MLIR test files; it reads the execution mode from the config, so one helper covers both modes.
- Also reject an e8m0 scale on integer weights in `finalize`, and share the qparam extraction and activation validation between the eager and graph export paths.
@dengqiaoyu dengqiaoyu added the enhancement New feature or request label Sep 1, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

2 participants