feat(quantization): add qparam_qspec for double-scale quantization

changelog.d/161.added

@@ -0,0 +1 @@
+Add `QuantizationSpec.qparam_qspec` to simulate double-scale quantization in PyTorch, such as NVFP4 and GGUF-style double scales, and `scale_dtype=torch.float16` to round scales to FP16.

docs/architecture_notes/graph_annotation.md

@@ -44,12 +44,12 @@ The fields are exactly the settable inputs of `coreai_opt...QuantizationSpec`, p
 
 `_qspec_constraints._FIELD_POLICY` is the authority. The reasoning:
 
-| Field                                                                                                                         | Policy                             | Why                                                                                                                                                                                                                                                                                                                                 |
-| ----------------------------------------------------------------------------------------------------------------------------- | ---------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `DTYPE`                                                                                                                       | Highest-priority proposal wins     | Two tensors sharing an observer can't have different dtypes, and there's no safe join — int4 and int8 aren't compatible in either direction. Config precedence decides.                                                                                                                                                             |
-| `QFORMULATION`, `FAKE_QUANTIZE_CLS`, `QPARAM_CALCULATOR_CLS`, `RANGE_CALCULATOR_CLS`, `SCALE_DTYPE`, `QSCHEME`, `GRANULARITY` | Highest-priority proposal wins     | Ordinary config choices.                                                                                                                                                                                                                                                                                                            |
-| `FLOAT_RANGE`                                                                                                                 | Widest of the proposals, per bound | A shared observer must cover every member's data, so a pin survives only while every member agrees on it. Priority is deliberately not consulted: covering the data is a correctness constraint, not a preference. This is also where "fixed qparams" lives — a fixed observer is one whose `float_range` is pinned on both bounds. |
-| `QUANTIZATION_TARGET`                                                                                                         | Highest-priority proposal wins     | This is not actually set as part of the config, but comes from whether the tensor is a weight or an activation. If we share a qspec between a weight and an activation, using the higher-priority one creates a coherent set of fields that matches at least one target's original provisional qspec.                               |
+| Field                                                                                                                                         | Policy                             | Why                                                                                                                                                                                                                                                                                                                                 |
+| --------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `DTYPE`                                                                                                                                       | Highest-priority proposal wins     | Two tensors sharing an observer can't have different dtypes, and there's no safe join — int4 and int8 aren't compatible in either direction. Config precedence decides.                                                                                                                                                             |
+| `QFORMULATION`, `FAKE_QUANTIZE_CLS`, `QPARAM_CALCULATOR_CLS`, `RANGE_CALCULATOR_CLS`, `SCALE_DTYPE`, `QPARAM_QSPEC`, `QSCHEME`, `GRANULARITY` | Highest-priority proposal wins     | Ordinary config choices.                                                                                                                                                                                                                                                                                                            |
+| `FLOAT_RANGE`                                                                                                                                 | Widest of the proposals, per bound | A shared observer must cover every member's data, so a pin survives only while every member agrees on it. Priority is deliberately not consulted: covering the data is a correctness constraint, not a preference. This is also where "fixed qparams" lives — a fixed observer is one whose `float_range` is pinned on both bounds. |
+| `QUANTIZATION_TARGET`                                                                                                                         | Highest-priority proposal wins     | This is not actually set as part of the config, but comes from whether the tensor is a weight or an activation. If we share a qspec between a weight and an activation, using the higher-priority one creates a coherent set of fields that matches at least one target's original provisional qspec.                               |
 
 Some properties are deliberately *not* fields. `quant_min` / `quant_max` / `n_bits` / `target_dtype` are computed by `QuantizationSpec` from `dtype` and `qscheme`, so they fall out of reassembly; reconciling them alongside their own inputs would let them contradict those inputs. `ch_axis` lives inside `GRANULARITY`.
 

docs/src/quantization/config.md

@@ -27,8 +27,10 @@ We will first take a look at the {class}`~coreai_opt.quantization.spec.Quantizat
 - `qscheme`: symmetric or asymmetric quantization scheme (default: {class}`~coreai_opt.quantization.spec.QuantizationScheme`.SYMMETRIC). See [Symmetric vs asymmetric quantization](advanced.md#symmetric-vs-asymmetric-quantization) for how these schemes differ.
 - `granularity`: granularity for quantization scales/zp, per-block, per-channel, per-tensor (default: {class}`~coreai_opt.quantization.spec.PerTensorGranularity()`)
 - `qparam_calculator_cls`: algorithm used for calculating scales/zp, e.g. static, global min/max, moving average min/max etc. (defaults: {class}`~coreai_opt.quantization.spec.MovingAverageQParamsCalculator` for activations)
-- `scale_dtype`: dtype for the quantization scale to be used. (Default: None, picked automatically based on quantization `dtype`)
-  Example:
+- `scale_dtype`: how the quantization scale is stored. `None` (default) keeps the computed scale, except for FP4, which defaults to `torch.float8_e8m0fnu`. `torch.float8_e8m0fnu` rounds scales to powers of two (OCP MX), for FP4 and FP8. `torch.float16` rounds the scale to FP16 whatever the model dtype, for integer and FP8 dtypes; Core AI and Core ML export of it is not supported yet.
+- `qparam_qspec`: a nested `QuantizationSpec` that quantizes the per-block scales themselves (and the per-block minvals, with `qformulation="minval"`). See [Quantizing the quantization parameters](#example-quantizing-the-quantization-parameters-qparam_qspec). (Default: None)
+
+Example:
 
 ```python
 from coreai_opt.quantization import QuantizationSpec
@@ -267,6 +269,133 @@ flowchart LR
     style output fill:#d4edda,stroke:#28a745
 ```
 
+### Example: Quantizing the quantization parameters (`qparam_qspec`)
+
+`qparam_qspec` quantizes the quantization parameters themselves. Quantizing a tensor per block produces a grid of qparams with one scale per block, plus one minval per block with `qformulation="minval"`. The nested `QuantizationSpec` quantizes that grid like any other tensor, and the tensor is then quantized against the decoded qparams, as a format that stores its qparams at low precision reconstructs them.
+
+Every `QuantizationSpec` field works in the nested spec as usual:
+
+- `granularity`: `PerBlockGranularity(block_size=8)` shares one set of nested qparams across 8 consecutive blocks. Left at `axis=None`, it follows the outer spec's block axis.
+- `scale_dtype` is the dtype of the super-scales; for example, `torch.float16` gives FP16 super-scales like GGUF Q4_K's `d` and `dmin`.
+- Calibration, observer toggles and QAT schedules apply to the nested qparams the same way as to any other qparams.
+
+For example, int4 blocks whose scales are int8 codes under one scale per output channel:
+
+```python
+from coreai_opt.quantization import QuantizationSpec
+from coreai_opt.quantization.spec import PerBlockGranularity, PerChannelGranularity
+
+int4_weight = QuantizationSpec(
+    dtype=torch.int4,
+    granularity=PerBlockGranularity(block_size=32),
+    qparam_qspec=QuantizationSpec(
+        dtype=torch.int8,
+        granularity=PerChannelGranularity(axis=0),
+    ),
+)
+```
+
+Layers whose blocks do not divide into the nested groups are left unquantized, with a warning. Core AI and Core ML export of a spec with `qparam_qspec` is not supported yet: `finalize()` raises `NotImplementedError`. Use the prepared model for torch inference.
+
+Double scale can be used to specify several popular quantization formats, such as GGUF's Q4_K and NVFP4.
+
+#### GGUF double scale
+
+GGUF k-quants such as Q4_K store 4-bit weights in blocks of 32, each with a scale and a min (`qformulation="minval"`). The scales and the mins are quantized in turn, under separate super-scales shared by 8 blocks (256 weights), so layers whose input channels are not a multiple of 256 are left unquantized. The super-scales are FP16 (`scale_dtype=torch.float16` in the nested spec), whatever the model dtype. Q4_K itself stores 6-bit codes, which needs `torch.uint6`; it is not supported yet, so this example uses uint8 codes.
+
+```python
+# programmatic
+
+from coreai_opt.quantization.config import QuantizerConfig, ModuleQuantizerConfig
+from coreai_opt.quantization import QuantizationSpec
+from coreai_opt.quantization.spec import PerBlockGranularity
+
+gguf_weight = QuantizationSpec(
+    dtype=torch.uint4,
+    qscheme="asymmetric",
+    qformulation="minval",
+    granularity=PerBlockGranularity(block_size=32),
+    qparam_qspec=QuantizationSpec(
+        dtype=torch.uint8,
+        qscheme="asymmetric",
+        granularity=PerBlockGranularity(block_size=8),
+        scale_dtype=torch.float16,
+    ),
+)
+
+config = QuantizerConfig(
+    global_config=ModuleQuantizerConfig(
+        op_input_spec=None,
+        op_output_spec=None,
+        op_state_spec={"weight": gguf_weight},
+    )
+)
+```
+
+```yaml
+# yaml config
+quantization_spec:
+  spec1: &gguf_weight
+    dtype: uint4
+    qscheme: asymmetric
+    qformulation: minval
+    granularity: { type: per_block, block_size: 32 }
+    qparam_qspec:
+      dtype: uint8
+      qscheme: asymmetric
+      granularity: { type: per_block, block_size: 8 }
+      scale_dtype: float16
+
+quantization_config:
+  global_config:
+    op_input_spec: null
+    op_output_spec: null
+    op_state_spec:
+      weight: *gguf_weight
+```
+
+#### NVFP4
+
+NVFP4 stores FP4 E2M1 values in blocks of 16, each block with an FP8 E4M3 scale, under one FP32 scale per tensor (4.5 bits per weight, plus 32 bits per tensor). For weights and activations:
+
+```python
+# programmatic
+
+from coreai_opt.quantization.config import QuantizerConfig, ModuleQuantizerConfig
+from coreai_opt.quantization import QuantizationSpec
+from coreai_opt.quantization.spec import PerBlockGranularity, PerTensorGranularity
+
+nvfp4_weight = QuantizationSpec(
+    dtype=torch.float4_e2m1fn_x2,
+    granularity=PerBlockGranularity(block_size=16),
+    # FP8 E4M3 block scales, under one FP32 scale per tensor
+    qparam_qspec=QuantizationSpec(
+        dtype=torch.float8_e4m3fn,
+        granularity=PerTensorGranularity(),
+    ),
+)
+
+# NVFP4 activations: block scales recomputed every forward, under a tensor scale
+# that is calibrated and then frozen
+nvfp4_activation = QuantizationSpec(
+    dtype=torch.float4_e2m1fn_x2,
+    granularity=PerBlockGranularity(axis=-1, block_size=16),
+    qparam_calculator_cls="dynamic",
+    qparam_qspec=QuantizationSpec(
+        dtype=torch.float8_e4m3fn,
+        granularity=PerTensorGranularity(),
+    ),
+)
+
+config = QuantizerConfig(
+    global_config=ModuleQuantizerConfig(
+        op_input_spec={"*": nvfp4_activation},
+        op_output_spec=None,
+        op_state_spec={"weight": nvfp4_weight},
+    )
+)
+```
+
 ### Example: Skip quantization for a specific layer type
 
 Let's say we want to skip all `linear` layers in the model from being quantized. Since linear is commonly used as `nn.Module` in model definition, there are two equivalent ways to do this:

docs/src/quantization/overview.md

@@ -164,6 +164,8 @@ Broadly speaking, the following dtypes can be used for quantization:
 
 A few other dtypes (e.g. FP8_E5M2) may also be available to experiment with; however, they may have limited support on the Core AI runtime.
 
+The quantization parameters can themselves be quantized with a nested spec (`qparam_qspec`), for weights and activations. These are available for simulation in PyTorch; Core AI and Core ML export is not supported yet. See [Quantizing the quantization parameters](config.md#example-quantizing-the-quantization-parameters-qparam_qspec).
+
 ## Two Execution Modes: graph and eager
 
 The {class}`~coreai_opt.quantization.Quantizer` has two internal implementations:

docs/src/utils/mixed_precision.md

@@ -103,6 +103,7 @@ The utility supports:
 
 - **Eager-mode integer weight quantization**: any spec-supported integer dtype, any quantization scheme (symmetric or asymmetric), at any granularity.
 - **Eager-mode floating-point weight quantization**: FP8 (`e4m3`, `e5m2`) and FP4 (`e2m1`).
+- **Quantized qparams** (`qparam_qspec`): each block's scale (and minval) costs its code width, plus the nested spec's own qparams, counted the same way. For example, NVFP4 (FP4 elements, FP8 E4M3 scales per block of 16, one FP32 tensor scale) comes to 4.5 BPW, plus 32 bits per tensor.
 - **Palettization** at any spec-supported `n_bits`, including a quantized LUT.
 
 Sub-byte payloads and zero-points are packed at `n_bits` with no padding.

src/coreai_opt/_utils/torch_utils.py

@@ -33,6 +33,9 @@
 E8M0_EXPONENT_BIAS: Final[int] = 127
 F32_MIN_NORMAL: Final[float] = 2**-126  # ~1.175e-38
 
+# Scale dtypes that the scale is rounded to before element codes are chosen.
+ROUNDED_SCALE_DTYPES: Final[frozenset[torch.dtype]] = frozenset({torch.float16})
+
 # Maps a graph-mode aten ``OpOverload`` to the ``nn.Module`` type it corresponds
 # to.
 ATEN_OP_TO_MODULE_TYPE: dict[torch._ops.OpOverload, type[torch.nn.Module]] = {
@@ -305,6 +308,11 @@ def is_tensor_positive(tensor: torch.Tensor) -> bool:
     return bool(tensor.gt(0).all())
 
 
+def nonzero_scale(scale: torch.Tensor) -> torch.Tensor:
+    """Return ``scale`` with zeros replaced by 1, as a safe divisor."""
+    return torch.where(scale == 0, 1.0, scale)
+
+
 def is_float_quant_dtype(dtype: torch.dtype) -> bool:
     """Check if dtype is a floating-point (FP4/FP8) quantization dtype.