🚨Add an FP8 backend to `QuantizedCache` (#48677)
* Add an FP8 backend to `QuantizedCache`
* Reorder the FP8 states for beam search
* Use an early return in `Fp8QuantizedLayer.reorder_cache`
* Measure the quantized cache backends accurately
The memory of the quantized states was under-reported: only `fp8` stores
them as a plain tensor, `hqq` uses a `(tensor, metadata)` tuple and quanto
a tensor subclass whose payload sits in its attributes, neither of which
`numel() * element_size()` accounts for.
Also report the KL divergence of the next-token distribution against the
bf16 one, as the matching token count saturates on short generations, and
copy the `cache_config` per call, which `generate` consumes.
* Reset the FP8 scales, and cover the FP8 backend in the quantized layer tests
* Address review on the FP8 quantized cache layer
- `fp8_max` becomes a class attribute next to `fp8_dtype`, instead of being
recomputed and stored on every instance. This matches how the FP8 constants
are already defined in `integrations/finegrained_fp8.py` and
`integrations/compressed_tensors.py`.
- `Fp8QuantizedLayer` now overrides `lazy_initialization` to calibrate the
scales, rather than doing it inline in the first `update` call.
- `_quantized_keys` and `_quantized_values` are initialized to `None` in
`QuantizedLayer.__init__`, so that they are always defined before the first
`update`. This was missing for the quanto and hqq layers as well.
- `.contiguous()` is now only called in `lazy_initialization`. The states
passed to `update` come straight from a transpose in the attention layer, so
they are genuinely non-contiguous during prefill, and `_quantize` propagates
that layout. Every subsequent append goes through a `cat`, whose output is
contiguous by construction, so the call is redundant there.
- `_cat_fp8` is removed in favour of a plain `torch.cat`, which handles
`float8_e4m3fn` natively on CPU and CUDA, for contiguous and non-contiguous
inputs alike (checked on torch 2.8 through 2.14).
* Trim the comments down to one line
* Store the FP8 scale alongside the states it quantizes
`_quantize` already returned a `(states, scale)` pair, but the two halves were
then split across `_quantized_keys` and a separate `_key_scale` attribute, and
zipped back together at every `_dequantize` call site.
Keeping the pair as the layer's single quantized payload matches what
`HQQQuantizedLayer` does with `(qtensor, meta)`, lets `_dequantize` be called on
`_quantized_keys` directly like the base class does, and removes the `reset`
override, since `QuantizedLayer.reset` already drops `_quantized_keys`.
It also keeps the scale's rank internal to `_quantize` / `_dequantize`, so a
future per-block backend can make it a tensor without changing the layer's
attributes.
* Say what the `.contiguous()` in `lazy_initialization` is actually for
Nothing in `_quantize` needs a contiguous input: on strided states it computes
the same scale and gives bit-identical dequantized values.
What the call is for is the layout of the stored buffer. The states handed to
`lazy_initialization` come from a transpose in the attention layer, while the
ones `update` appends come out of `cat`, and the two disagree: without the call
the buffer is laid out as `[B, S, H, D]` after the prefill and as `[B, H, S, D]`
from the first decode step on. Normalizing it once keeps the strides the same
for the buffer's whole life, which matters for the follow-up feeding these
states to FP8 attention kernels.
The test was building its states with `torch.randn(1, 2, 5, 8)`, so they were
already contiguous and the assertion held whether or not the call was there. It
now builds them the way the attention layer does.
* Address @vasqu's review
Drop the benchmark script: it belongs in the PR description rather than in the
code base.
Drop the `.contiguous()` in `lazy_initialization`. Nothing in `_quantize` needs
it, and nothing reads the quantized states while they are still strided, so the
layout only matters once something actually consumes the FP8 buffer. Better
added then, with a consumer to justify it.
Trim the class docstring down to what is specific to this backend, and the
`_quantize` and `reorder_cache` ones down to one point each.
The test now builds its states the way the attention layer does, so the strided
path is covered.
* One-line the docstrings and comments that fit on one line
* Keep the FP8 state initialization out of the shared quantized layer
* Split the quantized layer into a base and a group-wise class
The group-wise options only belong to the integer backends, so they move to
`GroupWiseQuantizedLayer` and `Fp8QuantizedLayer` no longer inherits them.
`QuantizedCache` now rejects them for the `fp8` backend instead of silently
replacing them.
* Note which subtree a future fp8 variation belongs to
Grouping along the head dimension keeps the states appendable, so it has no
use for the residual window `GroupWiseQuantizedLayer` exists for.
* Keep the group-wise signature on QuantizedLayer and name the new base
The split moved the group-wise options out of QuantizedLayer, which breaks
subclasses that pass them to super().__init__. Keep that signature where it
was and introduce BaseQuantizedLayer for what the backends share instead.
The vibevoice guard now skips on the new base, so FP8 layers keep being
skipped there.
* Declare the quantized states in the base layer, and group the quantized tests
The states were only assigned in `reset` and `lazy_initialization`, so a freshly
built layer had no such attribute at all.
The FP8 tests move next to the other quantized ones, and the beam reordering and
reset tests are parameterized over the backends instead of duplicated.
