🚨Add an FP8 backend to `QuantizedCache` by jiqing-feng · Pull Request #48677 · huggingface/transformers

jiqing-feng · 2026-09-10T08:23:14Z

What does this PR do?

Adds a third backend to QuantizedCache, next to quanto and hqq, that stores the KV states in torch.float8_e4m3fn:

model.generate(**inputs, cache_implementation="quantized", cache_config={"backend": "fp8"})

Unlike the integer backends, FP8 uses a single scale per tensor. New tokens can therefore be appended to the quantized states directly, so the layer needs neither a full-precision residual cache nor a periodic re-quantization of the whole cache. The scale is calibrated on the first update and then kept fixed.

The states are dequantized in update, i.e. before attention is computed, so every attention implementation keeps working unchanged. They are nonetheless kept as a single contiguous FP8 tensor, so that a follow-up PR can let FP8-capable attention backends consume them directly.

It also requires no extra dependency, contrary to quanto and hqq.

Results

Qwen3-4B-Instruct-2507, 1024-token prompt + 256 generated tokens, A100 80GB, measured with the benchmark script added in this PR (benchmark_v2/benchmark_scripts/quantized_cache.py):

cache	KV cache	peak memory	latency	matching tokens
bf16	175.2 MiB	8098.9 MiB	8396.0 ms	256 / 256
fp8	87.6 MiB	7917.4 MiB	9723.7 ms	256 / 256

The KV cache is exactly 2x smaller, and greedy generation is bit-for-bit identical to the bf16 baseline on this prompt.

Latency is 16% higher and peak memory barely moves, because the whole cache is dequantized to bf16 at every decoding step. Both are addressed by the follow-up PR consuming the FP8 states directly in the attention kernels; this PR is deliberately kept to the cache backend alone to stay easy to review.

Changes

Fp8QuantizedLayer in cache_utils.py, selected by cache_config={"backend": "fp8"}
Exported from the top-level package and documented in kv_cache.md / internal/generation_utils.md
test_fp8_quantized_layer unit test, and fp8 added to the test_quantized_cache_generation parameterization
benchmark_v2/benchmark_scripts/quantized_cache.py, comparing the backends against a bf16 DynamicCache

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.

vasqu · 2026-09-17T12:55:46Z

I have this still on my mind, just going through the fixed quant cache and then revisiting this

…r tests

jiqing-feng · 2026-10-06T09:10:55Z

Hi @vasqu . This PR is ready to be reviewed.

vasqu · 2026-10-06T09:27:08Z

Also cc @SunMarc should be interesting for you as well

IlyasMoutawwakil · 2026-10-15T08:05:55Z

+        # their default values and unused (`nbits` is only informative here).
+        super().__init__(nbits=8, residual_length=0)
+
+        self._fp8_max = torch.finfo(self.fp8_dtype).max


do we need to save it here ?

No, it is a constant derived from fp8_dtype. Made it a class attribute right next to it, same as
integrations/finegrained_fp8.py and integrations/compressed_tensors.py already do.

IlyasMoutawwakil · 2026-10-15T08:06:21Z

+    def _quantize(self, tensor, axis=None, scale=None):
+        """Quantize `tensor` to FP8, calibrating the scale on its amplitude if it is not provided. The `axis`
+        argument is unused, as the scale is per-tensor."""
+        if scale is None:
+            # `clamp` guards against an all-zeros tensor, which would give a null scale
+            scale = (tensor.abs().amax().float() / self._fp8_max).clamp(min=torch.finfo(torch.float32).tiny)
+        return (tensor.float() / scale).to(self.fp8_dtype), scale


would it make sense to @torch.compile ?

Measured it, it is a regression, so I left them as-is.

End-to-end generate, SmolLM2-135M, 128 tokens, best of 5:

bare 2918.4 ms compiled 4510.9 ms (+55%) compiled(dynamic=True) 4797.7 ms (+64%) bare (re-check) 2964.2 ms

Not a recompilation issue (dynamo reported 0 recompiles) — both functions are a few memory-bound elementwise ops, so there is nothing to fuse and only guard overhead is left. They also already get compiled as part of the larger graph when the user compiles the model.

both functions are a few memory-bound elementwise ops, so there is nothing to fuse and only guard overhead is left

that's the opposite argument for kernel fusion 😅 fusion is most suited for "a few mem bound elementwise ops", because it avoids the hbm roundtrip.

They also already get compiled as part of the larger graph when the user compiles the model.

okaaay this part makes sense

You're right about fusion and I had the argument backwards — sorry. I went and looked at what inductor actually emits, and the two functions behave differently.

_dequantize fuses into a single kernel, exactly as you'd expect, and it really is faster per call. Measured inside a real generate (Qwen3-4B, 8192-token prompt, 4536 calls): 0.143 ms -> 0.095 ms.

_quantize does not fuse, because inductor does not lower fp8 stores — a bare x.to(torch.float8_e4m3fn) compiles to no kernel at all, just a fallback:

buf0 = torch.ops.prims.convert_element_type.default(arg0_1, torch.float8_e4m3fn)

So it keeps an fp32 intermediate and does the cast eagerly, which costs an extra roundtrip the eager path never pays.

What stopped me from adding the decorator is that the per-call win does not survive end to end. Qwen3-4B, 8192-token prompt + 64 tokens, 5 interleaved reps, compiling _dequantize only:

eager min 2786 median 2837 max 2973 compiled min 3056 median 3122 max 3162 +10%

Same direction on SmolLM2 with a 4096-token prompt: +14.4% for _dequantize, +42.2% for _quantize. No recompiles in any of these runs.

I don't have a satisfying explanation for why ~220 ms saved inside the function turns into ~280 ms lost overall, so I'd rather not add the decorators on a result I can't account for. For _quantize the fp8-store fallback is a concrete reason to leave it alone regardless. Happy to revisit if you think the per-call win is worth chasing.

yeh pretty sure it's torch dynamo or inductor still having problems with fp8 (pytorch/pytorch#189522 for example), thanks for the investigation

kind of unexpected lol but it is what it is

IlyasMoutawwakil · 2026-10-15T08:06:34Z

+    def _dequantize(self, qtensor):
+        quant_tensor, scale = qtensor
+        return (quant_tensor.float() * scale).to(self.dtype)


same question

IlyasMoutawwakil · 2026-10-15T08:07:56Z

+            self.lazy_initialization(key_states, value_states)
+            self._quantized_keys, self._key_scale = self._quantize(key_states.contiguous())
+            self._quantized_values, self._value_scale = self._quantize(value_states.contiguous())
+            return key_states, value_states


i think it makes sense to override lazy_initialization to init scales as well

Done. update is now only the steady-state path.

IlyasMoutawwakil · 2026-10-15T08:08:21Z

+        new_keys, _ = self._quantize(key_states.contiguous(), scale=self._key_scale)
+        new_values, _ = self._quantize(value_states.contiguous(), scale=self._value_scale)


does it have to be contiguous ? and why wouldn't it be ?

Kept it in lazy_initialization, dropped it from update.

It is not contiguous during prefill: the states come straight out of .view(...).transpose(1, 2) in the attention layer. Instrumenting a real generate gives stride=(1152, 64, 192, 1) for every prefill call, under both eager and sdpa. And _quantize propagates that layout, so the call is not a no-op — without it the stored tensor keeps the transposed strides, which defeats the point of keeping the quantized states as one contiguous FP8 buffer for FP8 attention kernels to consume later.

In update it was redundant, since cat output is contiguous regardless of its inputs.

IlyasMoutawwakil · 2026-10-15T08:09:06Z

+        self._quantized_keys = _cat_fp8(self._quantized_keys, new_keys)
+        self._quantized_values = _cat_fp8(self._quantized_values, new_values)


these were not None initialized

Fixed in QuantizedLayer.__init__ rather than here, since quanto and hqq had the same gap.

IlyasMoutawwakil · 2026-10-15T08:09:48Z

+        self._quantized_keys = self._quantized_keys.index_select(0, beam_idx)
+        self._quantized_values = self._quantized_values.index_select(0, beam_idx)


why do we name them _quantized_keys / _quantized_values and not just keys and values like other cache layers ?

That name is not mine — it comes from QuantizedLayer, which has used it since the per-layer cache refactor (#39106). I am only following the existing convention of the class I inherit from.

The two names mean different things there: self.keys is the full-precision residual window of the last residual_length tokens, while _quantized_keys holds everything older, quantized. KIVI needs both, since the residual is what gets flushed into the quantized states in one go.

More importantly, _quantized_keys is not a [B, H, S, D] tensor at all — it is a quanto QTensor for the quanto backend, a (tensor, meta) tuple for hqq, and a (tensor, scale) tuple here. DynamicLayer documents keys as a plain tensor and several base methods rely on that, which is also why QuantizedLayer overrides get_seq_length to use cumulative_length instead of self.keys.shape[-2]. Hqq even reads self.keys.device / self.keys.dtype inside _quantize.

So renaming would mean finding another name for the residual buffer, touching all three backends, and revisiting the base-class assumptions. Happy to do it if you want, but it feels like its own PR rather than part of the FP8 backend — what do you think?

thanks for the explanation, yes i wasnt asking for a change of names, just wondering what the other attrs are used for

IlyasMoutawwakil · 2026-10-15T08:10:23Z

+    def reset(self) -> None:
+        """Resets the cache values while preserving the objects."""
+        super().reset()
+        # The scales are dropped along with the states they quantize, so that the next `update` recalibrates


some comments are self evident so no need (i understand the struggle haha)

Dropped the one in update, and trimmed the rest down to one line each.

IlyasMoutawwakil · 2026-10-15T08:11:15Z

+        # The `fp8` backend quantizes per-tensor, so it does not use any of the group-wise options
+        layer_kwargs = {}
+        if backend != "fp8":
+            layer_kwargs = {
+                "nbits": nbits,
+                "axis_key": axis_key,
+                "axis_value": axis_value,
+                "q_group_size": q_group_size,
+                "residual_length": residual_length,
+            }


why not support block dynamic too ? for accuracy sensitive workloads

Agreed it is worth having — and I found a concrete argument for it. With the scale frozen on the prefill, 15% of the decode steps in a short-prompt run exceeded the calibrated amax (worst case 1.56x), silently saturating 0.039% of the elements. That is also why the long benchmark prompt stays bit-for-bit identical to bf16 while short prompts diverge after 30-90 tokens: a 991-token prefill calibrates the scale well enough that nothing overflows.

Still, I would rather do it as a follow-up, for two reasons.

It changes the storage contract rather than adding a flag: the scale becomes a tensor that grows with the sequence, appending has to handle block boundaries, and reorder_cache has to reorder the scales too instead of just index_select-ing the quantized states.

And the follow-up PR that feeds these states to FP8 attention kernels will constrain how the blocks should be laid out, so picking a granularity now risks picking the wrong one.

agree with the follow up pr, but the api should be built taking into consideration the requirements of the future work, to not have to break too many things between iterations 🤗

Agreed. QuantizedCache already exposes axis_key / axis_value / q_group_size / residual_length, and the fp8 path strips them today, so per-block would mostly stop ignoring existing parameters rather than add new ones. q_group_size maps directly onto the block size; axis_* I'm less sure about, since quanto constrains them to 0/-1 and the blocks would run along the sequence dim. Per-tensor stays the default either way.

Worth flagging now: blocking along the sequence dim re-introduces a residual window, since a block's scale can't be finalised until the block is full. That is the residual_length machinery QuantizedLayer already has and that this layer bypasses with residual_length=0 — so the follow-up re-enables the base class rather than working around it.

I did make one change now (05ef372): the scale is stored inside _quantized_keys as a (states, scale) pair instead of in its own attribute, like HQQQuantizedLayer does with (qtensor, meta). That keeps the scale's rank internal to _quantize / _dequantize, so turning it into a per-block tensor won't touch the layer's attributes. It also drops the reset override, since the base class already clears _quantized_keys.

hmm ok so we keep the shared layout for block dynamic in the future?

That was the idea, yes — and it's a separate question from the class split above, since QuantizedCache keeps this signature either way.

One wart worth fixing: the dropped arguments don't get ignored, they get replaced. QuantizedCache(backend="fp8", q_group_size=128, nbits=2, residual_length=512) leaves the layer holding 64 / 8 / 0. Nothing reads them on the fp8 path, so it's harmless today, but it's confusing if you go looking. Happy to raise instead.

IlyasMoutawwakil · 2026-10-15T08:15:23Z

+def _cat_fp8(tensor: torch.Tensor, other: torch.Tensor, dim: int = -2) -> torch.Tensor:
+    """Concatenate two FP8 tensors, as `torch.cat` is not implemented for FP8 dtypes on all backends."""
+    return torch.cat([tensor.view(torch.uint8), other.view(torch.uint8)], dim=dim).view(tensor.dtype)


could be a @torch.compiled helper in the class like _quantize and _dequantize

Turns out it should not exist at all — torch.cat handles FP8 natively, so I deleted it. I had added the uint8-view workaround defensively without checking.

oh really, did you check it's not a version dependent behaviour ?

Verified on 2.8, 2.10, 2.11, 2.12, 2.13 and 2.14, CPU and CUDA, contiguous and non-contiguous inputs. setup.py says torch>=2.5 and I had no container for 2.5–2.7, so that range is untested — tell me if you would rather keep the helper for safety.

i tested locally, it works on 2.6 but not 2.5 but we already moved to 2.6 on main

- `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).

IlyasMoutawwakil · 2026-10-16T08:20:50Z

not sure what's the policy for adding bench scripts @vasqu

Honestly I'd remove it from this PR and have a gist or details section in the PR description with the script. Shouldnt live in the code base tho

@jiqing-feng

Done, removed.

IlyasMoutawwakil · 2026-10-16T08:22:51Z

+        self._quantized_keys, self._key_scale = self._quantize(key_states.contiguous())
+        self._quantized_values, self._value_scale = self._quantize(value_states.contiguous())


but why do we need to contiguate them for quantize ? does it have ops that don't support strided tensors ?

Not for _quantize — you're right, nothing in it needs contiguity, and on strided states it returns the same scale and bit-identical values. Bad wording on my side.

It's there so the buffer keeps one layout for its whole life: the prefill states come from a transpose, the appended ones from cat, and the two disagree.

without .contiguous() with .contiguous() prefill (524288, 128, 1024, 1) False (524288, 65536, 128, 1) True decode 1 (525312, 65664, 128, 1) True (525312, 65664, 128, 1) True

Nothing notices today since update hands back dequantized states, but the follow-up feeding this buffer to an FP8 kernel would, and a prefill-only stride discrepancy is a nasty thing to debug. It costs one copy per prefill per layer (0.009 ms at 512 tokens, 0.077 ms at 16k), so I kept it and fixed the comment to say this instead.

The test was using freshly allocated contiguous states, so it never covered any of this — it now builds them the way the attention layer does, torch.randn(1, 5, 2, 8).transpose(1, 2).

Nothing notices today since update hands back dequantized states

so it's unnecessary for now, we can add the contiguous when strided tensors do break smth.

Fair, dropped it. The test keeps the transposed states, so the strided path stays covered.

IlyasMoutawwakil

LGTM ! one last nit from my side

`_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.

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.

github-actions · 2026-10-16T09:20:41Z

Thank you for your contribution 🤗!

CI Security Gate — automatic approval blocked

This PR was not automatically approved for CI because the security gate failed.

Possible reasons:

The PR touches 50 or more files — only PRs with fewer than 50 changed files are automatically approved
A changed file is outside the allowed directories (src/, tests/, docs/, utils/), has a disallowed extension (only .py, .txt, .md permitted outside tests/ and docs/), or is not .md/.yml inside docs/ — this covers files the PR deletes or renames, not only the ones it edits
A new security issue was detected in the changed Python files (Bandit check) — a high-severity one, shell execution, exec/eval, a URL fetch, or disabled TLS verification
The PR touches a path this repository protects from untrusted PRs, such as the file that decides who reviews it — a maintainer must make that change in a separate PR

See the workflow run for the exact violations.

A maintainer can review and manually approve CI if a finding is a false positive.

vasqu

Just got some questions on the shared keys that are currently not used, I only saw at the end that they could still be useful so maybe some comments might not make sense

vasqu · 2026-10-16T11:40:10Z

Honestly I'd remove it from this PR and have a gist or details section in the PR description with the script. Shouldnt live in the code base tho

vasqu · 2026-10-16T11:46:52Z

        self.q_group_size = q_group_size
        self.residual_length = residual_length
        self.cumulative_length = 0
+        self._quantized_keys = self._quantized_values = None


hmm maybe we should rather repeat things a bit where actually needed instead of storing things alongside, e.g. fp8 does not need q_group_size - quantized_key not needed in hqq

Done, moved it into Fp8QuantizedLayer.__init__. It was never in QuantizedLayer.__init__ on main, so quanto and hqq are untouched again.

vasqu · 2026-10-16T11:48:01Z

+    Contrary to the integer backends, FP8 uses a single scale per tensor, which allows new tokens to be
+    appended to the quantized states directly. This layer therefore needs neither a full-precision residual
+    cache, nor a re-quantization of the whole cache every `residual_length` tokens. The scales are calibrated
+    on the states the layer is initialized with, and then kept fixed.
+
+    As the other backends do, the quantized states are stored as a single payload, here a `(states, scale)`
+    pair, so that the scale stays internal to `_quantize` and `_dequantize`.
+
+    The states are dequantized in `update`, i.e. before attention is computed, so that every attention
+    implementation keeps working unchanged. They are nonetheless stored as a single contiguous FP8 tensor, so
+    that attention backends supporting FP8 inputs can later consume them directly.


Too much yapping to my liking, we can shorten this significantly as the interestin part is the non quant/dequant need imo

vasqu · 2026-10-16T11:48:21Z

+    fp8_dtype = torch.float8_e4m3fn
+    fp8_max = torch.finfo(fp8_dtype).max


do we need to guard for older torch versions e.g. 2.6?

No. setup.py on main requires torch>=2.6, and Ilyas tested the fp8 path on 2.6 earlier in this review — it only breaks on 2.5. I also ran the torch.cat check on 2.8 / 2.10 / 2.11 / 2.12 / 2.13 / 2.14, CPU and CUDA.

vasqu · 2026-10-16T11:49:08Z

+        # The base class options are all specific to group-wise integer quantization, so they are left to
+        # their default values and unused (`nbits` is only informative here).
+        super().__init__(nbits=8, residual_length=0)


as mentioned before id rather decouple a bit

Part guesswork on my side — I kept them expecting a per-block follow-up to reuse q_group_size and residual_length, since a sequence-dim block can't finalise its scale until it's full. The axis_* ones may well not map at all.

Happy to split instead, and it's cheaper than I assumed: a minimal base above QuantizedLayer, which keeps the group-wise machinery and its current signature, so only Fp8QuantizedLayer re-points and quanto and hqq don't change. The one thing that has to follow is the isinstance(layer, QuantizedLayer) check in generation_vibevoice.py, which would otherwise stop skipping fp8 layers and index into their empty keys. Your call — I can do it here rather than as a separate PR.

vasqu · 2026-10-16T11:51:52Z

+    def _quantize(self, tensor, axis=None, scale=None):
+        """Quantize `tensor` to FP8, calibrating the scale on its amplitude if it is not provided. The `axis`
+        argument is unused, as the scale is per-tensor."""
+        if scale is None:
+            # `clamp` guards against an all-zeros tensor, which would give a null scale
+            scale = (tensor.abs().amax().float() / self._fp8_max).clamp(min=torch.finfo(torch.float32).tiny)
+        return (tensor.float() / scale).to(self.fp8_dtype), scale


kind of unexpected lol but it is what it is

vasqu · 2026-10-16T11:53:54Z

+        keys_to_return = self._dequantize(self._quantized_keys)
+        values_to_return = self._dequantize(self._quantized_values)


ig at some point we could involve FA with fp8 support right? so no real dequant etc

Yes, and the kernel is already there — kernels-community/flash-attn3 builds the *_e4m3_sm90 instantiations and takes q_descale/k_descale/v_descale. The missing piece is ours: modeling_flash_attention_utils.py has no way to pass a descale, so there's nothing to hand fp8 states to yet. SDPA can't take them either ("mul_cuda" not implemented for 'Float8_e4m3fn'), hence the dequantization.

Two notes for whoever wires it up. FA3 wants the descale as a (batch_size, num_heads_k) tensor, so a per-tensor scale goes in expanded — the kernel is already finer-grained than this layer. And it's Hopper-only; Intel XPU falls back to flash-attn2, which has no fp8 kernels. I'm looking at the XPU side next, so I'd rather the storage layout here didn't end up CUDA-shaped.

vasqu · 2026-10-16T11:54:40Z

+        """Reorders this layer's cache for beam search. The scale being per-tensor, the FP8 states can be
+        reordered as-is, without any dequantization. The reordering is therefore applied right away, contrary
+        to the base class which defers it to the next `update`."""


can also be shortened; or rather if we split the quant cache a bit more it's more obv

vasqu · 2026-10-16T11:55:24Z

+        # The `fp8` backend quantizes per-tensor, so it does not use any of the group-wise options
+        layer_kwargs = {}
+        if backend != "fp8":
+            layer_kwargs = {
+                "nbits": nbits,
+                "axis_key": axis_key,
+                "axis_value": axis_value,
+                "q_group_size": q_group_size,
+                "residual_length": residual_length,
+            }


hmm ok so we keep the shared layout for block dynamic in the future?

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.

vasqu

Ok imo we should do the proper split then to avoid having unnecessary kwargs for fp8. It's harmless but confusing and we usually want to be explicit in transformers

It will also help to flag for the future fp8 variations then 🤔 wdyt @jiqing-feng

vasqu · 2026-10-19T15:12:33Z

I dont mind a separate PR or here tbh

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.

Grouping along the head dimension keeps the states appendable, so it has no use for the residual window `GroupWiseQuantizedLayer` exists for.

jiqing-feng · 2026-10-20T05:56:15Z

Hi @vasqu

The split

QuantizedLayer is now the base with only what all three backends share. The group-wise options move into a new GroupWiseQuantizedLayer that quanto and hqq inherit, so Fp8QuantizedLayer sits directly on the base with no __init__ at all, and axis is gone from its _quantize.

I kept the name QuantizedLayer for the base rather than adding a new one, which is the opposite of what I suggested above. It means the isinstance check in generation_vibevoice.py doesn't need the fix I mentioned, and it reads better — it's meant to skip quantized layers, fp8 included.

The kwargs

QuantizedCache raises for fp8 now instead of replacing them. The five options default to None and are only forwarded when set, so quanto and hqq keep their defaults. I left the check explicit rather than deriving it from the layer signature; a group-wise fp8 backend is what would make it per-backend, which is the next section.

One thing that came out of it

reset now drops the quantized states for the integer backends too. They never did on main, so the old payload stayed alive until the next update overwrote it — 2 MiB per layer on a 4x8x512x64 cache.

Where a future fp8 variation goes

Not implemented here, but I prototyped it to check: a group-wise fp8 layer would subclass Fp8QuantizedLayer, not GroupWiseQuantizedLayer.

What GroupWiseQuantizedLayer really holds is the residual window, and that exists for quantized states you can't append to. Grouping along the head dimension can be appended to — each token calibrates its own scale and cats straight in, and beam reordering stays a lossless index_select. Blocking along the sequence dimension can't, since a block's scale isn't final until it's full. So the grouping axis decides, not fp8 versus int. b138c83 notes this on both classes.

vasqu

1 last point we need to check because we might be breaking

vasqu · 2026-10-20T10:18:33Z

@@ -690,6 +690,34 @@ def reorder_cache(self, beam_idx: torch.LongTensor) -> None:


 class QuantizedLayer(DynamicLayer):


This is slightly breaking thos if users relied on QuantizedLayer so let's add 🚨 to the title and describe this

Or do you think we can make a proper BC version without making it too ugly? 😅

It's a trade-off between BC and naming, pick whichever you prefer. 780cc9d has the BC one.

QuantizedLayer keeps the exact signature it has on main, down to _quantize(self, tensor, axis), and the new BaseQuantizedLayer holds what all three backends share. A subclass written against main needs no change.

What it costs is that QuantizedLayer now means "group-wise quantized layer" rather than "any quantized layer", so Fp8QuantizedLayer is not an instance of it. The isinstance check in generation_vibevoice.py has to move to the new base, otherwise FP8 layers stop being skipped and the code indexes into their empty keys. That's the only call site in the repo, and every new backend now has to pick the right one of the two to inherit from.

Outside the repo nothing breaks either way: isinstance(layer, QuantizedLayer) keeps matching quanto and hqq, and couldn't have been handed an FP8 layer before this PR.

If you'd rather keep the clearer naming, I'll revert to b138c83 and flag the break instead.

Let me know your choice (these 2 above or others) so I can rewrite the title and PR statement.

I think the current option is quite nice, maybe slightly confusing but not by much imo

We can remove the 🚨 then as well no?

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.

vasqu

Last nits then we can merge

vasqu · 2026-10-21T14:00:08Z

+        """Resets the cache values while preserving the objects."""
+        super().reset()
+        # The quantized states are dropped instead of zeroed, so that the next `update` quantizes from scratch
+        self._quantized_keys = self._quantized_values = None


shouldnt these also not be declared in init?

Right, they were only ever assigned in reset and lazy_initialization, so a freshly built layer had no such attribute at all. Declared both in BaseQuantizedLayer.__init__ next to cumulative_length, and kept the reset assignment, which is what actually drops the states.

vasqu · 2026-10-21T14:03:55Z

+    def test_fp8_quantized_layer(self):
+        """Tests that `Fp8QuantizedLayer` appends new tokens using the scale calibrated on the first `update`."""
+        layer = Fp8QuantizedLayer()
+
+        # As in the attention layer, the states come out of a transpose, so they are not contiguous
+        keys, values = torch.randn(1, 5, 2, 8).transpose(1, 2), torch.randn(1, 5, 2, 8).transpose(1, 2)
+        layer.update(keys, values)
+        quantized_keys, calibrated_scale = layer._quantized_keys
+        self.assertEqual(quantized_keys.dtype, torch.float8_e4m3fn)
+        self.assertEqual(quantized_keys.shape, keys.shape)
+        calibrated_scale = calibrated_scale.clone()
+
+        # Even states with a much larger amplitude do not change the scale, which is frozen after calibration
+        cached_keys, _ = layer.update(torch.randn(1, 2, 1, 8) * 100, torch.randn(1, 2, 1, 8))
+        self.assertEqual(cached_keys.shape, (1, 2, 6, 8))
+        self.assertEqual(cached_keys.dtype, keys.dtype)
+        quantized_keys, scale = layer._quantized_keys
+        self.assertTrue(torch.equal(scale, calibrated_scale))
+        self.assertEqual(quantized_keys.shape, (1, 2, 6, 8))
+
+    def test_fp8_quantized_cache_rejects_group_wise_options(self):
+        """Tests that the `fp8` backend, which quantizes per-tensor, does not silently drop the group-wise options."""
+        config = LlamaConfig(num_hidden_layers=2)
+        with self.assertRaisesRegex(ValueError, "does not support"):
+            QuantizedCache(backend="fp8", config=config, q_group_size=128, nbits=2)
+
+        layer = QuantizedCache(backend="fp8", config=config).layers[0]
+        self.assertIsInstance(layer, Fp8QuantizedLayer)
+        self.assertFalse(hasattr(layer, "q_group_size"))


lets move these closer to the other quant tests?

Moved, they now sit right after test_quantized_layer_reset.

vasqu · 2026-10-21T14:04:17Z

+        keys, _ = layer.update(torch.rand(4, 2, 3, 8), torch.rand(4, 2, 3, 8))
+        self.assertEqual(keys.shape[-2], 3)
+        self.assertEqual(layer.get_seq_length(), 3)
+        self.assertLess(layer._quantized_keys[1].item(), calibrated_scale.item())


in general i feel like we might be able to parametrize more?

Done for the two that were straight duplicates: test_quantized_layer_beam_reorder and test_quantized_layer_reset are now parameterized over quanto and fp8, and the FP8-only copies are gone. The quanto case keeps its @require_optimum_quanto as a skipTest inside, since the decorator would skip the FP8 case too.

I left test_fp8_quantized_layer and test_fp8_quantized_cache_rejects_group_wise_options standalone, as they assert the frozen per-tensor scale and the rejected group-wise kwargs, which have no counterpart on the other backends.

…ed 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.

github-actions · 2026-10-22T01:13:23Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: vibevoice

vasqu

Thank you, very nice direction! Merging

HuggingFaceDocBuilderDev · 2026-10-22T11:12:03Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

jiqing-feng added 3 commits September 10, 2026 16:12

Add an FP8 backend to QuantizedCache

9789b12

Reorder the FP8 states for beam search

64cfc5d

Use an early return in Fp8QuantizedLayer.reorder_cache

87ab14a

jiqing-feng mentioned this pull request Sep 11, 2026

[Feature request] FP8 KV cache + flash-attention: is this in scope? #48608

Open

vasqu requested review from vasqu and zucchini-nlp September 11, 2026 10:01

jiqing-feng added 2 commits September 14, 2026 06:49

Merge branch 'main' into fp8-kv-cache

8882b42

jiqing-feng added 2 commits October 6, 2026 16:39

Merge branch 'main' into fp8-kv-cache

2b7a9eb

Reset the FP8 scales, and cover the FP8 backend in the quantized laye…

bc126c1

…r tests

jiqing-feng marked this pull request as ready for review October 6, 2026 09:09

Merge branch 'main' into fp8-kv-cache

5eaeb32

Merge branch 'main' into fp8-kv-cache

4fc4f00

IlyasMoutawwakil reviewed Oct 15, 2026

View reviewed changes

jiqing-feng added 2 commits October 16, 2026 14:55

Trim the comments down to one line

e9e889d

IlyasMoutawwakil reviewed Oct 16, 2026

View reviewed changes

IlyasMoutawwakil approved these changes Oct 16, 2026

View reviewed changes

jiqing-feng added 2 commits October 16, 2026 16:51

vasqu reviewed Oct 16, 2026

View reviewed changes

jiqing-feng added 4 commits October 19, 2026 09:05

One-line the docstrings and comments that fit on one line

7018cc4

Merge branch 'main' into fp8-kv-cache

862a643

Keep the FP8 state initialization out of the shared quantized layer

b4e172b

vasqu reviewed Oct 19, 2026

View reviewed changes

jiqing-feng added 2 commits October 20, 2026 12:53

Note which subtree a future fp8 variation belongs to

b138c83

Grouping along the head dimension keeps the states appendable, so it has no use for the residual window `GroupWiseQuantizedLayer` exists for.

vasqu reviewed Oct 20, 2026

View reviewed changes

jiqing-feng changed the title ~~Add an FP8 backend to QuantizedCache~~ 🚨Add an FP8 backend to QuantizedCache Oct 21, 2026

vasqu approved these changes Oct 21, 2026

View reviewed changes

vasqu approved these changes Oct 22, 2026

View reviewed changes

vasqu enabled auto-merge October 22, 2026 11:05

vasqu added this pull request to the merge queue Oct 22, 2026

Merged via the queue into huggingface:main with commit 43cc0ab Oct 22, 2026
116 checks passed

stevhliu mentioned this pull request Oct 23, 2026

[docs] Translation workflow #49889

Closed

qgallouedec mentioned this pull request Nov 9, 2026

Allocate the StaticCache values with the value states' dtype #49554

Open

		new_keys, _ = self._quantize(key_states.contiguous(), scale=self._key_scale)
		new_values, _ = self._quantize(value_states.contiguous(), scale=self._value_scale)

		self._quantized_keys = _cat_fp8(self._quantized_keys, new_keys)
		self._quantized_values = _cat_fp8(self._quantized_values, new_values)

		self._quantized_keys = self._quantized_keys.index_select(0, beam_idx)
		self._quantized_values = self._quantized_values.index_select(0, beam_idx)

		self._quantized_keys, self._key_scale = self._quantize(key_states.contiguous())
		self._quantized_values, self._value_scale = self._quantize(value_states.contiguous())

		fp8_dtype = torch.float8_e4m3fn
		fp8_max = torch.finfo(fp8_dtype).max

		keys_to_return = self._dequantize(self._quantized_keys)
		values_to_return = self._dequantize(self._quantized_values)

		@@ -690,6 +690,34 @@ def reorder_cache(self, beam_idx: torch.LongTensor) -> None:


		class QuantizedLayer(DynamicLayer):

Conversation

jiqing-feng commented Sep 10, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Results

Changes

Uh oh!

vasqu commented Sep 17, 2026

Uh oh!

jiqing-feng commented Oct 6, 2026

Uh oh!

vasqu commented Oct 6, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jiqing-feng Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

jiqing-feng commented Sep 10, 2026 •

edited by github-actions Bot

Loading

jiqing-feng Oct 16, 2026 •

edited

Loading

IlyasMoutawwakil Oct 16, 2026 •

edited

Loading

IlyasMoutawwakil left a comment •

edited

Loading