From a0ac07d73baf2bf5e97e8d03466ef75875b4ebda Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:58:30 +0000 Subject: [PATCH 1/5] Add opt-in FP8 attention and context caching to TabICLv2 --- sdm/_kernels/fp8_attention.py | 111 ++++++++++++ sdm/_kernels/triton/fp8_attention.py | 200 +++++++++++++++++++++ sdm/cache.py | 36 ++++ sdm/models/tabiclv2/model.py | 23 ++- sdm/nn/attention.py | 163 +++++++++++++---- test/models/tabiclv2/test_model.py | 36 ++++ test/models/test_attention_quantization.py | 33 ++++ test/nn/test_attention_quantization.py | 89 +++++++++ test/nn/test_fp8_attention.py | 151 ++++++++++++++++ 9 files changed, 802 insertions(+), 40 deletions(-) create mode 100644 sdm/_kernels/fp8_attention.py create mode 100644 sdm/_kernels/triton/fp8_attention.py create mode 100644 test/models/test_attention_quantization.py create mode 100644 test/nn/test_attention_quantization.py create mode 100644 test/nn/test_fp8_attention.py diff --git a/sdm/_kernels/fp8_attention.py b/sdm/_kernels/fp8_attention.py new file mode 100644 index 0000000000..4b0b95cb52 --- /dev/null +++ b/sdm/_kernels/fp8_attention.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import ModuleType + +import torch +from torch import Tensor + +from sdm.cache import QuantizedKVCacheEntry + +_triton: ModuleType | None = None +try: + from sdm._kernels.triton import fp8_attention as _triton_impl +except ImportError: + pass +else: + _triton = _triton_impl + + +def supports_fp8(query: Tensor) -> bool: + return ( + _triton is not None + and not torch.is_grad_enabled() + and not torch.compiler.is_compiling() + and query.is_cuda + and ( + query.dtype in {torch.float16, torch.bfloat16} + or ( + query.dtype == torch.float32 + and torch.is_autocast_enabled("cuda") + ) + ) + and query.size(-1) in {32, 64, 128, 256} + and torch.cuda.get_device_capability(query.device) in {(8, 9), (12, 0)} + ) + + +def fp8_attention( + query: Tensor, + key: Tensor, + value: Tensor, + cache: QuantizedKVCacheEntry | None = None, + scale: float | None = None, +) -> tuple[Tensor, QuantizedKVCacheEntry]: + assert _triton is not None + # RMSNorm keeps Q/K in FP32. Match SDPA's autocast input conversion before + # quantizing, without ever dequantizing cached FP8 keys and values. + if torch.is_autocast_enabled("cuda"): + dtype = torch.get_autocast_dtype("cuda") + query = query.to(dtype) + if cache is None: + key, value = key.to(dtype), value.to(dtype) + batch_shape = torch.broadcast_shapes( + query.shape[:-3], key.shape[:-3], value.shape[:-3] + ) + + def heads(x: Tensor) -> Tensor: + return ( + x.expand(batch_shape + x.shape[-3:]) + .reshape(-1, *x.shape[-3:]) + .transpose(1, 2) + ) + + def rows(x: Tensor) -> Tensor: + x = x.transpose(1, 2) + return x.reshape(batch_shape + x.shape[-3:]) + + q = heads(query) + if cache is None: + context = q[..., : key.size(-3), :] + qs = context.abs().amax((-2, -1), keepdim=True).float() + qs = qs.clamp_min(1e-12) / 448.0 + k, ks = _triton.quantize(heads(key)) + v, vs = _triton.quantize(heads(value)) + else: + k = heads(cache.key).contiguous() + v = heads(cache.value).contiguous() + ks, vs, qs = ( + heads(t).contiguous() + for t in (cache.key_scale, cache.value_scale, cache.query_scale) + ) + q, qs = _triton.quantize(q, qs) + blackwell = torch.cuda.get_device_capability(query.device) == (12, 0) + wide = query.size(-1) == 256 + result = _triton.quantized_attention( + q, + k, + v.transpose(-1, -2).contiguous(), + qs, + ks, + vs, + query.dtype, + 64 if wide else 128, + scale, + tile=64 if wide or blackwell or k.size(-2) < 32768 else 128, + warps=8 if wide else 4, + stages=4 if blackwell and not wide else 2, + accumulation_chunk=0 if blackwell else 16, + lift_exp=blackwell, + fused_softmax=blackwell, + ) + if cache is None: + cache = QuantizedKVCacheEntry( + key=rows(k), + value=rows(v), + key_scale=rows(ks), + value_scale=rows(vs), + query_scale=rows(qs), + dtype=query.dtype, + ) + return rows(result), cache diff --git a/sdm/_kernels/triton/fp8_attention.py b/sdm/_kernels/triton/fp8_attention.py new file mode 100644 index 0000000000..30fea4e747 --- /dev/null +++ b/sdm/_kernels/triton/fp8_attention.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, cast + +import torch +import triton +import triton.language as _tl + +tl: Any = _tl + + +@triton.jit +def attention_kernel( + Q, + K, + V, + QS, + KS, + VS, + Out, + HQ: tl.constexpr, + HK: tl.constexpr, + M: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + SCALE: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + ACC_CHUNK: tl.constexpr = 0, + LIFT_EXP: tl.constexpr = False, + FUSED_SOFTMAX: tl.constexpr = False, +): + block, head = tl.program_id(0), tl.program_id(1) + kv_head = (head // HQ) * HK + (head % HQ) // (HQ // HK) + rows = block * BM + tl.arange(0, BM) + cols = tl.arange(0, BN) + dims = tl.arange(0, D) + q = tl.load( + Q + head * M * D + rows[:, None] * D + dims[None, :], + rows[:, None] < M, + 0.0, + ) + scale = ( + tl.load(QS + head) * tl.load(KS + kv_head) * SCALE * 1.4426950408889634 + ) + vmax = tl.load(VS + kv_head) + maximum = tl.full((BM,), -float("inf"), tl.float32) + denominator = tl.full((BM,), 0, tl.float32) + acc = tl.full((BM, D), 0, tl.float32) + if ACC_CHUNK > 0: + outer_acc = tl.full((BM, D), 0, tl.float32) + outer_max = tl.full((BM,), -float("inf"), tl.float32) + for start in range(tl.cdiv(N, BN)): + indices = start * BN + cols + k = tl.load( + K + kv_head * N * D + indices[None, :] * D + dims[:, None], + indices[None, :] < N, + 0.0, + ) + score = tl.dot(q, k, max_num_imprecise_acc=32) + if FUSED_SOFTMAX and SCALE > 0: + score = tl.where(indices[None, :] < N, score, -float("inf")) + new_max = tl.maximum(maximum, tl.max(score, 1) * scale) + shift = new_max - 8.0 if LIFT_EXP else new_max + p = tl.exp2(score * scale - shift[:, None]) + else: + score = tl.where( + indices[None, :] < N, score * scale, -float("inf") + ) + new_max = tl.maximum(maximum, tl.max(score, 1)) + if LIFT_EXP: + p = tl.exp2(score - (new_max[:, None] - 8.0)) + else: + p = tl.exp2(score - new_max[:, None]) + correction = tl.exp2(maximum - new_max) + denominator = denominator * correction + tl.sum(p, 1) + acc *= correction[:, None] + v_offsets = indices[:, None] + dims[None, :] * N + v = tl.load( + V + kv_head * N * D + v_offsets, + indices[:, None] < N, + 0.0, + ) + # Lift softmax probabilities by 256 to retain precision in FP8. + acc = tl.dot( + (p if LIFT_EXP else p * 256.0).to(tl.float8e4nv), + v, + acc, + max_num_imprecise_acc=32, + ) + # Periodically flush the dot accumulator to FP32 for long L4 contexts. + if ACC_CHUNK > 0: # noqa: SIM102 - constexpr guard avoids modulo zero. + if ((start + 1) % ACC_CHUNK == 0) | (start + 1 == tl.cdiv(N, BN)): + outer_acc = ( + outer_acc * tl.exp2(outer_max - new_max)[:, None] + acc + ) + outer_max = new_max + acc = tl.full((BM, D), 0, tl.float32) + maximum = new_max + if ACC_CHUNK > 0: + acc = outer_acc + out = acc * vmax / denominator[:, None] + if not LIFT_EXP: + out *= 1.0 / 256.0 + tl.store( + Out + head * M * D + rows[:, None] * D + dims[None, :], + out, + rows[:, None] < M, + ) + + +@triton.jit +def quantize_kernel( + X, + S, + Y, + LENGTH: tl.constexpr, + BLOCK: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + S0: tl.constexpr, + S1: tl.constexpr, + S2: tl.constexpr, + S3: tl.constexpr, +): + offset = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + head = tl.program_id(1) + source = ( + (head // H) * S0 + + (head % H) * S1 + + (offset // D) * S2 + + (offset % D) * S3 + ) + x = tl.load(X + source, offset < LENGTH, 0.0).to(tl.float32) + scale = tl.load(S + head) + y = tl.minimum(tl.maximum(x / scale, -448.0), 448.0) + tl.store(Y + head * LENGTH + offset, y, offset < LENGTH) + + +def quantize( + x: torch.Tensor, scale: torch.Tensor | None = None +) -> tuple[torch.Tensor, torch.Tensor]: + if scale is None: + scale = ( + x.abs().amax(dim=(-2, -1), keepdim=True).float().clamp_min(1e-12) + / 448.0 + ) + output = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn) + length = x.size(-2) * x.size(-1) + cast(Any, quantize_kernel)[ + (((length + 1023) // 1024), x.size(0) * x.size(1)) + ](x, scale, output, length, 1024, x.size(1), x.size(-1), *x.stride()) + return output, scale + + +def quantized_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + qs: torch.Tensor, + ks: torch.Tensor, + vs: torch.Tensor, + dtype: torch.dtype, + block: int, + scale: float | None = None, + tile: int = 64, + warps: int = 4, + stages: int = 2, + accumulation_chunk: int = 0, + lift_exp: bool = False, + fused_softmax: bool = False, +) -> torch.Tensor: + b, h, m, d = q.shape + n = k.size(-2) + # Q/K are [batch, heads, sequence, channels]; V is transposed. + output = torch.empty((b, h, m, d), device=q.device, dtype=dtype) + cast(Any, attention_kernel)[(((m + block - 1) // block), b * h)]( + q, + k, + v, + qs, + ks, + vs, + output, + h, + k.size(1), + m, + n, + d, + d**-0.5 if scale is None else scale, + block, + tile, + accumulation_chunk, + lift_exp, + fused_softmax, + num_warps=warps, + num_stages=stages, + ) + return output diff --git a/sdm/cache.py b/sdm/cache.py index e8ee2bb631..1cf6b365c4 100644 --- a/sdm/cache.py +++ b/sdm/cache.py @@ -9,9 +9,11 @@ Mapping, MutableMapping, ) +from dataclasses import dataclass from enum import StrEnum from typing import NamedTuple, Self +import torch from torch import Tensor from sdm.tensor.mixin import DeviceMixin @@ -38,6 +40,40 @@ def _apply_tensor(self, fn: Callable[[Tensor], Tensor]) -> Self: return self.__class__(key=fn(self.key), value=fn(self.value)) +@dataclass +class QuantizedKVCacheEntry(DeviceMixin): + """FP8 attention projections and their per-head dequantization scales. + + Keys and values have shape ``[..., rows, heads, channels]``. Scales have + shape ``[..., 1, heads, 1]``. The query scale is learned from context + queries and reused for prediction. ``dtype`` is the attention output dtype. + """ + + key: Tensor + value: Tensor + key_scale: Tensor + value_scale: Tensor + query_scale: Tensor + dtype: torch.dtype + + def _tensors(self) -> Iterator[Tensor]: + yield self.key + yield self.value + yield self.key_scale + yield self.value_scale + yield self.query_scale + + def _apply_tensor(self, fn: Callable[[Tensor], Tensor]) -> Self: + return self.__class__( + key=fn(self.key), + value=fn(self.value), + key_scale=fn(self.key_scale), + value_scale=fn(self.value_scale), + query_scale=fn(self.query_scale), + dtype=self.dtype, + ) + + class Cache(MutableMapping[Hashable, object], DeviceMixin): r"""A mutable mapping of model cache values.""" diff --git a/sdm/models/tabiclv2/model.py b/sdm/models/tabiclv2/model.py index 495c17b27f..811b470c80 100644 --- a/sdm/models/tabiclv2/model.py +++ b/sdm/models/tabiclv2/model.py @@ -9,7 +9,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any, ClassVar, cast +from typing import Any, ClassVar, Literal, cast import torch from torch import Tensor @@ -19,6 +19,7 @@ from sdm.cache import Cache from sdm.models import ICLModel from sdm.models._huggingface import download_checkpoint +from sdm.models.tabiclv2.block import TabICLv2TransformerBlock from sdm.models.tabiclv2.ckpt import remap_ckpt from sdm.models.tabiclv2.icl import ICLBlock from sdm.models.tabiclv2.recipe import default_recipe @@ -116,6 +117,14 @@ class TabICLv2(ICLModel): model are initialized. pretrained: Whether to load the pretrained checkpoint. device: The device. + attention_quantization: Set to ``"fp8"`` to quantize eligible ICL + attention and its cache during FP16/BF16 CUDA inference on Ada + and RTX Blackwell GPUs. Requires Triton and more than 8192 + context rows. + Weights and the final ICL layer remain unquantized. Other inputs + use standard attention. Fitted FP8 caches require the same + inference dtype and do not support ``torch.compile``. + Defaults to ``None``. Cached prediction is not always faster. """ supported_feature_stypes: ClassVar[frozenset[Stype]] = frozenset( @@ -132,10 +141,12 @@ def __init__( task: TaskLike | Iterable[TaskLike] | None = None, pretrained: bool = True, device: torch.device | str | None = None, + *, + attention_quantization: Literal["fp8"] | None = None, ) -> None: super().__init__(task=task) - self.models: ModuleDict[TaskLike, torch.nn.Module] = ModuleDict() + self.models: ModuleDict = ModuleDict() for task in self.tasks: self.models[task] = _TabICLv2( num_classes=10 if task == Task.classification else 0, @@ -147,6 +158,14 @@ def __init__( if pretrained: self._load_from_pretrained(device=device) + if attention_quantization == "fp8": + for network in self.models.values(): + for layer in cast(_TabICLv2, network).icl_block.layers[:-1]: + cast( + TabICLv2TransformerBlock, layer + ).attn.attention_quantization = "fp8" + else: + assert attention_quantization is None self.eval() @classmethod diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py index b2020d9ff3..4390df3b3e 100644 --- a/sdm/nn/attention.py +++ b/sdm/nn/attention.py @@ -12,9 +12,12 @@ from torch import Tensor from torch.nn import Linear -from sdm.cache import KVCacheEntry +from sdm._kernels.fp8_attention import fp8_attention, supports_fp8 +from sdm.cache import KVCacheEntry, QuantizedKVCacheEntry from sdm.nn import QueryScaling +_CacheEntry = KVCacheEntry | QuantizedKVCacheEntry + class SDPA(torch.nn.Module): r"""Scaled Dot-Product Attention (SDPA). @@ -219,6 +222,8 @@ def __init__( self.q_dim = num_query_heads * self.head_dim # == channels self.kv_dim = num_key_value_heads * self.head_dim + self.attention_quantization: Literal["fp8"] | None = None + self.qkv_lin = Linear( channels, self.q_dim + 2 * self.kv_dim, bias=bias, **factory_kwargs ) @@ -240,7 +245,7 @@ def __init__( def forward( self, query: Tensor, - key_value: Tensor | KVCacheEntry | None = None, + key_value: Tensor | _CacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, *, @@ -251,33 +256,33 @@ def forward( def forward( self, query: Tensor, - key_value: Tensor | KVCacheEntry | None = None, + key_value: Tensor | _CacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, *, return_key_value: Literal[True], - ) -> tuple[Tensor, KVCacheEntry]: ... + ) -> tuple[Tensor, _CacheEntry]: ... @overload def forward( self, query: Tensor, - key_value: Tensor | KVCacheEntry | None = None, + key_value: Tensor | _CacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, *, return_key_value: bool, - ) -> Tensor | tuple[Tensor, KVCacheEntry]: ... + ) -> Tensor | tuple[Tensor, _CacheEntry]: ... def forward( self, query: Tensor, # [..., Q, C] - key_value: Tensor | KVCacheEntry | None = None, # [..., KV, C] + key_value: Tensor | _CacheEntry | None = None, # [..., KV, C] seqused_key_value: Tensor | None = None, # [...] attn_mask: Tensor | None = None, # [..., Q, KV] *, return_key_value: bool = False, - ) -> Tensor | tuple[Tensor, KVCacheEntry]: # [..., Q, C] + ) -> Tensor | tuple[Tensor, _CacheEntry]: # [..., Q, C] r"""The forward pass. Args: @@ -302,7 +307,7 @@ def forward( Otherwise, a tuple of the output tensor and a :class:`~sdm.cache.KVCacheEntry`. """ - if isinstance(key_value, KVCacheEntry): + if isinstance(key_value, _CacheEntry): query = F.linear( query, weight=self.qkv_lin.weight[: self.q_dim], @@ -310,7 +315,12 @@ def forward( if self.qkv_lin.bias is not None else None, ) - if ( + if isinstance(key_value, QuantizedKVCacheEntry): + if query.dtype != key_value.dtype: + raise ValueError( + "Query dtype differs from fitted FP8 cache" + ) + elif ( key_value.key.dtype != query.dtype or key_value.value.dtype != query.dtype ): @@ -336,7 +346,7 @@ def forward( # [..., S, C] -> [..., S, H, C // H], with separate query/kv heads. query = query.unflatten(-1, [self.num_query_heads, self.head_dim]) - if not isinstance(key_value, KVCacheEntry): + if not isinstance(key_value, _CacheEntry): key = key.unflatten(-1, [self.num_key_value_heads, self.head_dim]) value = value.unflatten( -1, [self.num_key_value_heads, self.head_dim] @@ -345,22 +355,63 @@ def forward( if self.query_transform is not None: query = self.query_transform(query) if ( - not isinstance(key_value, KVCacheEntry) + not isinstance(key_value, _CacheEntry) and self.key_transform is not None ): key = self.key_transform(key) - out = self.sdpa( - query=query, # [..., Q, Hq, C // Hq] - key=key, # [..., KV, Hkv, C // Hq] - value=value, # [..., KV, Hkv, C // Hq] - seqused_key_value=seqused_key_value, # [...] - attn_mask=attn_mask, # [..., Q, KV] - ) # [..., Q, Hq, C // Hq] - + quantized_cache = ( + key_value if isinstance(key_value, QuantizedKVCacheEntry) else None + ) + use_fp8 = ( + self.attention_quantization == "fp8" + and not isinstance(key_value, _CacheEntry) + and query.numel() > 0 + and key.size(-3) > 8192 + and query.size(-3) >= key.size(-3) + and seqused_key_value is None + and attn_mask is None + and supports_fp8(query) + ) + if quantized_cache is not None: + if ( + not supports_fp8(query) + or attn_mask is not None + or seqused_key_value is not None + ): + raise ValueError( + "Fitted FP8 attention requires supported CUDA inference " + "without attention masks or torch.compile" + ) + use_fp8 = True + if quantized_cache is not None and query.numel() == 0: + out = query + elif use_fp8: + scaled_query = query + if self.sdpa.query_scaling is not None: + scaled_query = self.sdpa.query_scaling( + query, key_len=key.size(-3) + ) + out, quantized_cache = fp8_attention( + query=scaled_query, + key=key, + value=value, + cache=quantized_cache, + scale=self.sdpa.scale, + ) + else: + out = self.sdpa( + query=query, # [..., Q, Hq, C // Hq] + key=key, # [..., KV, Hkv, C // Hq] + value=value, # [..., KV, Hkv, C // Hq] + seqused_key_value=seqused_key_value, # [...] + attn_mask=attn_mask, # [..., Q, KV] + ) # [..., Q, Hq, C // Hq] out = out.flatten(-2, -1) # [..., Q, C] out = self.out_lin(out) # [..., Q, C] if return_key_value: + if quantized_cache is not None: + return out, quantized_cache # CUDA autocast runs RMSNorm in fp32, so cast explicitly before # caching: https://github.com/pytorch/pytorch/blob/v2.13.0/aten/src/ATen/autocast_mode.h#L875 return out, KVCacheEntry( @@ -439,7 +490,7 @@ def __init__( def forward( self, query: Tensor, - key_value: Tensor | KVCacheEntry | None = None, + key_value: Tensor | _CacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, *, @@ -452,39 +503,39 @@ def forward( def forward( self, query: Tensor, - key_value: Tensor | KVCacheEntry | None = None, + key_value: Tensor | _CacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, *, return_key_value: Literal[True], batch_size_limit: int | Literal["auto"] | None = None, out: Tensor | None = None, - ) -> tuple[Tensor, KVCacheEntry]: ... + ) -> tuple[Tensor, _CacheEntry]: ... @overload def forward( self, query: Tensor, - key_value: Tensor | KVCacheEntry | None = None, + key_value: Tensor | _CacheEntry | None = None, seqused_key_value: Tensor | None = None, attn_mask: Tensor | None = None, *, return_key_value: bool, batch_size_limit: int | Literal["auto"] | None = None, out: Tensor | None = None, - ) -> Tensor | tuple[Tensor, KVCacheEntry]: ... + ) -> Tensor | tuple[Tensor, _CacheEntry]: ... def forward( self, query: Tensor, # [..., Q, C] - key_value: Tensor | KVCacheEntry | None = None, # [..., KV, C] + key_value: Tensor | _CacheEntry | None = None, # [..., KV, C] seqused_key_value: Tensor | None = None, # [...] attn_mask: Tensor | None = None, # [..., Q, KV] *, return_key_value: bool = False, batch_size_limit: int | Literal["auto"] | None = None, out: Tensor | None = None, - ) -> Tensor | tuple[Tensor, KVCacheEntry]: # [..., Q, C] + ) -> Tensor | tuple[Tensor, _CacheEntry]: # [..., Q, C] r"""The forward pass. Args: @@ -516,7 +567,7 @@ def forward( raise RuntimeError( "'out' is only supported when gradients are disabled" ) - if return_key_value and isinstance(key_value, KVCacheEntry): + if return_key_value and isinstance(key_value, _CacheEntry): raise ValueError( "'return_key_value=True' is not supported when 'key_value' is " "already cached" @@ -538,7 +589,7 @@ def forward( key_value_length: int | None = None if isinstance(key_value, Tensor): key_value_length = key_value.size(-2) - elif isinstance(key_value, KVCacheEntry): + elif isinstance(key_value, _CacheEntry): key_value_length = key_value.key.size(-3) bytes_per_example = self.peak_bytes_per_example( @@ -574,7 +625,7 @@ def forward( if batch_size <= batch_size_limit or query.size(-2) == 0: disable_chunking = True if return_key_value: - assert not isinstance(key_value, KVCacheEntry) + assert not isinstance(key_value, _CacheEntry) if batch_shape != ( query.size()[:-2] if key_value is None @@ -595,6 +646,8 @@ def forward( flat_out: Tensor | None = None flat_key: Tensor | None = None flat_value: Tensor | None = None + flat_scales: list[Tensor] | None = None + cache_dtype: torch.dtype | None = None if out is not None and (out.dim() <= 3 or out.is_contiguous()): flat_out = out.view(batch_size, *query.size()[-2:]) @@ -638,6 +691,20 @@ def forward( flat_key[start:end] = chunk_kv.key flat_value[start:end] = chunk_kv.value + if isinstance(chunk_kv, QuantizedKVCacheEntry): + scales = ( + chunk_kv.key_scale, + chunk_kv.value_scale, + chunk_kv.query_scale, + ) + if flat_scales is None: + flat_scales = [ + t.new_empty(batch_size, *t.shape[-3:]) + for t in scales + ] + cache_dtype = chunk_kv.dtype + for dest, source in zip(flat_scales, scales): + dest[start:end] = source del chunk_kv else: @@ -666,6 +733,22 @@ def forward( assert flat_key is not None assert flat_value is not None + if flat_scales is not None: + assert cache_dtype is not None + return out, QuantizedKVCacheEntry( + key=flat_key.view(*batch_shape, *flat_key.shape[-3:]), + value=flat_value.view(*batch_shape, *flat_value.shape[-3:]), + key_scale=flat_scales[0].view( + *batch_shape, *flat_scales[0].shape[-3:] + ), + value_scale=flat_scales[1].view( + *batch_shape, *flat_scales[1].shape[-3:] + ), + query_scale=flat_scales[2].view( + *batch_shape, *flat_scales[2].shape[-3:] + ), + dtype=cache_dtype, + ) return out, KVCacheEntry( key=flat_key.view(*batch_shape, *flat_key.size()[-3:]), value=flat_value.view(*batch_shape, *flat_value.size()[-3:]), @@ -674,13 +757,13 @@ def forward( def _forward( self, query: Tensor, # [..., Q, C] - key_value: Tensor | KVCacheEntry | None = None, # [..., KV, C] + key_value: Tensor | _CacheEntry | None = None, # [..., KV, C] seqused_key_value: Tensor | None = None, # [...] attn_mask: Tensor | None = None, # [..., Q, KV] *, return_key_value: bool = False, out: Tensor | None = None, - ) -> Tensor | tuple[Tensor, KVCacheEntry]: # [..., Q, C] + ) -> Tensor | tuple[Tensor, _CacheEntry]: # [..., Q, C] if self.key_value_norm is None: pass @@ -731,14 +814,14 @@ def peak_bytes_per_example( def _batch_shape( query: Tensor, - key_value: Tensor | KVCacheEntry | None, + key_value: Tensor | _CacheEntry | None, seqused_key_value: Tensor | None, attn_mask: Tensor | None, ) -> torch.Size: shapes = [query.size()[:-2]] if isinstance(key_value, Tensor): shapes.append(key_value.size()[:-2]) - elif isinstance(key_value, KVCacheEntry): + elif isinstance(key_value, _CacheEntry): shapes += [key_value.key.size()[:-3], key_value.value.size()[:-3]] if seqused_key_value is not None: shapes.append(seqused_key_value.size()) @@ -759,12 +842,12 @@ def _chunk( @overload def _chunk( - tensor: KVCacheEntry, + tensor: _CacheEntry, batch_shape: torch.Size, trailing_dims: int, start: int, end: int, -) -> KVCacheEntry: ... +) -> _CacheEntry: ... @overload @@ -778,16 +861,20 @@ def _chunk( def _chunk( - tensor: Tensor | KVCacheEntry | None, + tensor: Tensor | _CacheEntry | None, batch_shape: torch.Size, trailing_dims: int, start: int, end: int, -) -> Tensor | KVCacheEntry | None: +) -> Tensor | _CacheEntry | None: if tensor is None: return None + if isinstance(tensor, QuantizedKVCacheEntry): + return tensor._apply_tensor( + lambda value: _chunk(value, batch_shape, 3, start, end) + ) if isinstance(tensor, KVCacheEntry): return KVCacheEntry( key=_chunk(tensor.key, batch_shape, 3, start, end), diff --git a/test/models/tabiclv2/test_model.py b/test/models/tabiclv2/test_model.py index ead2881c70..66707c911b 100644 --- a/test/models/tabiclv2/test_model.py +++ b/test/models/tabiclv2/test_model.py @@ -1,10 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from typing import cast + import pytest import torch from sdm import Recipe +from sdm.cache import Cache, QuantizedKVCacheEntry from sdm.models import TabICLv2 from sdm.models.tabiclv2.model import _TabICLv2 from sdm.models.tabiclv2.row_embedding import RowEmbedding @@ -258,3 +261,36 @@ def test_compile(dtype: torch.dtype) -> None: predicted = model.predict(x_query) assert predicted.allclose(expected, atol=5e-4, rtol=5e-3) assert torch.is_inference(predicted) + + +@onlyCUDA +def test_fp8_fit_predict() -> None: + if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}: + pytest.skip("FP8 integration supports Ada and RTX Blackwell") + model = TabICLv2( + task="regression", + pretrained=False, + device="cuda", + attention_quantization="fp8", + ) + x = torch.randn(8193, 3, device="cuda") + y = torch.randn(8193, 1, device="cuda") + query = torch.randn(17, 3, device="cuda") + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16): + expected = model(x, y, query, recipe=Recipe(), num_estimators=1) + model.fit(x, y, recipe=Recipe(), num_estimators=1) + actual = model.predict(query) + assert actual.numerical.isfinite().all() + torch.testing.assert_close( + actual.numerical, expected.numerical, atol=0.01, rtol=0.01 + ) + assert model._cache is not None + # The public lifecycle must retain real quantized tensors, not merely + # accept an option while silently executing ordinary attention. + entries = cast(Cache, model._cache[0]) + assert any( + isinstance(entry, QuantizedKVCacheEntry) + for entry in entries.values() + ) + model.clear() + assert model._cache is None diff --git a/test/models/test_attention_quantization.py b/test/models/test_attention_quantization.py new file mode 100644 index 0000000000..a664caad03 --- /dev/null +++ b/test/models/test_attention_quantization.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from sdm.models import TabICLv2 +from sdm.nn import Attention + + +@pytest.mark.parametrize("model_cls", [TabICLv2]) +def test_model_fp8_opt_in(model_cls: type[torch.nn.Module]) -> None: + model = model_cls( + task="regression", + pretrained=False, + device="meta", + attention_quantization="fp8", + ) + ordinary = model_cls(task="regression", pretrained=False, device="meta") + assert ordinary.state_dict().keys() == model.state_dict().keys() + quantized = [ + name + for name, module in model.named_modules() + if isinstance(module, Attention) + and module.attention_quantization == "fp8" + ] + assert quantized + assert all("icl_block.layers" in name for name in quantized) + assert all( + module.attention_quantization is None + for module in ordinary.modules() + if isinstance(module, Attention) + ) diff --git a/test/nn/test_attention_quantization.py b/test/nn/test_attention_quantization.py new file mode 100644 index 0000000000..9b317e6540 --- /dev/null +++ b/test/nn/test_attention_quantization.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import copy + +import pytest +import torch + +from sdm.cache import Cache, KVCacheEntry, QuantizedKVCacheEntry +from sdm.models.tabiclv2.block import TabICLv2TransformerBlock +from sdm.nn import Attention +from sdm.testing import onlyCUDA, withCUDA + + +@withCUDA +def test_fp8_small_context_fallback(device: torch.device) -> None: + module = Attention(128, 2, device=device) + torch.nn.init.normal_(module.out_lin.weight, std=0.05) + reference = copy.deepcopy(module) + module.attention_quantization = "fp8" + query = torch.randn(2, 19, 128, device=device) + with torch.inference_mode(): + actual, cache = module(query, return_key_value=True) + expected = reference(query) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert isinstance(cache, KVCacheEntry) + assert module.state_dict().keys() == reference.state_dict().keys() + + +@onlyCUDA +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_fp8_context_cache_and_chunking(dtype: torch.dtype) -> None: + if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}: + pytest.skip("FP8 integration supports Ada and RTX Blackwell") + module = TabICLv2TransformerBlock( + channels=128, + num_heads=2, + norm_bias=True, + qassmax=True, + device="cuda", + ) + torch.nn.init.normal_(module.attn.out_lin.weight, std=0.05) + reference = copy.deepcopy(module) + module.attn.attention_quantization = "fp8" + context = torch.randn(2, 1, 8193, 128, device="cuda") + query = torch.randn(2, 1, 129, 128, device="cuda") + joined = torch.cat([context, query], dim=-2) + with torch.inference_mode(), torch.autocast("cuda", dtype=dtype): + expected = reference(joined, context) + full = module(joined, context) + _, cache = module( + context, context, return_key_value=True, batch_size_limit=1 + ) + assert isinstance(cache, QuantizedKVCacheEntry) + actual = module(query, cache, batch_size_limit=1) + unchunked = module(query, cache) + split = torch.cat( + [ + module(query[..., :64, :], cache), + module(query[..., 64:, :], cache), + ], + dim=-2, + ) + empty = module(query[..., :0, :], cache) + migrated = Cache(kv=cache).cpu().cuda()["kv"] + replay = module(query, migrated) + assert empty.size(-2) == 0 + assert actual.isfinite().all() + torch.testing.assert_close(actual, unchunked, atol=2e-3, rtol=2e-3) + torch.testing.assert_close(actual, split, atol=2e-3, rtol=2e-3) + torch.testing.assert_close(actual, replay, atol=2e-3, rtol=2e-3) + torch.testing.assert_close( + actual, full[..., -129:, :], atol=0.03, rtol=0.03 + ) + error = (full - expected).square().mean() / expected.square().mean() + assert error.sqrt() < 0.03 + assert Cache(kv=cache).size() == sum( + t.numel() * t.element_size() for t in cache._tensors() + ) + assert cache.key.dtype == torch.float8_e4m3fn + assert cache.value.dtype == torch.float8_e4m3fn + assert cache.query_scale.dtype == torch.float32 + with torch.inference_mode(), pytest.raises(ValueError, match="dtype"): + module(query, cache) + with ( + torch.autocast("cuda", dtype=dtype), + pytest.raises(ValueError, match="inference"), + ): + module(query, cache) diff --git a/test/nn/test_fp8_attention.py b/test/nn/test_fp8_attention.py new file mode 100644 index 0000000000..5f396ea11a --- /dev/null +++ b/test/nn/test_fp8_attention.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +import pytest +import torch + +pytest.importorskip("triton") +from sdm._kernels.triton.fp8_attention import ( + quantize, + quantized_attention, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.get_device_capability() < (8, 9), + reason="FP8 attention requires an Ada or newer CUDA GPU", +) + + +@pytest.mark.parametrize( + ("batch", "queries", "context"), [(2, 129, 513), (1, 129, 8192)] +) +def test_fp8_partial_tiles(batch: int, queries: int, context: int) -> None: + with torch.inference_mode(): + q = torch.randn( + batch, 8, queries, 64, device="cuda", dtype=torch.float16 + ) + k = torch.randn( + batch, 8, context, 64, device="cuda", dtype=torch.float16 + ) + v = torch.randn_like(k) + q8, qs = quantize(q) + k8, ks = quantize(k) + v8, vs = quantize(v) + actual = quantized_attention( + q8, + k8, + v8.transpose(-1, -2).contiguous(), + qs, + ks, + vs, + q.dtype, + 128, + ) + expected = torch.nn.functional.scaled_dot_product_attention(q, k, v) + assert actual.isfinite().all() + relative_rms = ( + actual.float() - expected.float() + ).square().mean().sqrt() / expected.float().square().mean().sqrt() + assert relative_rms < 0.08 + + +@pytest.mark.parametrize("scale", [None, 1.0]) +@pytest.mark.parametrize("tile", [64, 128]) +def test_grouped_heads_and_model_scale(scale: float | None, tile: int) -> None: + with torch.inference_mode(): + q = ( + torch.randn(2, 8, 129, 64, device="cuda", dtype=torch.float16) + * 0.3 + ) + k = ( + torch.randn(2, 2, 8192, 64, device="cuda", dtype=torch.float16) + * 0.3 + ) + v = torch.randn_like(k) + q8, qs = quantize(q) + k8, ks = quantize(k) + v8, vs = quantize(v) + actual = quantized_attention( + q8, + k8, + v8.transpose(-1, -2).contiguous(), + qs, + ks, + vs, + q.dtype, + 128, + scale, + tile=tile, + ) + expected = torch.nn.functional.scaled_dot_product_attention( + q, k, v, enable_gqa=True, scale=scale + ) + assert actual.isfinite().all() + relative_rms = ( + actual.float() - expected.float() + ).square().mean().sqrt() / expected.float().square().mean().sqrt() + assert relative_rms < 0.08 + + +@pytest.mark.parametrize("accumulation_chunk", [0, 16]) +@pytest.mark.parametrize("optimized", [False, True]) +@pytest.mark.parametrize("tile", [64, 128]) +@pytest.mark.parametrize("context", [8193, 65536]) +def test_periodic_attention_matches_cached_query_slice( + context: int, tile: int, optimized: bool, accumulation_chunk: int +) -> None: + with torch.inference_mode(): + q = torch.randn(1, 1, context, 64, device="cuda", dtype=torch.float16) + k, v = torch.randn_like(q), torch.randn_like(q) + q8, qs = quantize(q) + k8, ks = quantize(k) + v8, vs = quantize(v) + v8 = v8.transpose(-1, -2).contiguous() + opts: dict[str, Any] = { + "accumulation_chunk": accumulation_chunk, + "tile": tile, + "lift_exp": optimized, + "fused_softmax": optimized, + } + full = quantized_attention( + q8, k8, v8, qs, ks, vs, q.dtype, 128, **opts + ) + cached = quantized_attention( + q8[:, :, -129:].contiguous(), + k8, + v8, + qs, + ks, + vs, + q.dtype, + 128, + **opts, + ) + torch.testing.assert_close( + cached, full[:, :, -129:], rtol=1e-3, atol=2e-5 + ) + reference = torch.nn.functional.scaled_dot_product_attention( + q[:, :, -129:], k, v + ) + relative_rms = ( + (cached.float() - reference.float()).square().mean() + / reference.float().square().mean() + ).sqrt() + assert relative_rms < 0.08 + assert cached.isfinite().all() + + +def test_quantization_reads_strided_heads_without_fp16_copy() -> None: + x = torch.randn( + 2, 513, 8, 64, device="cuda", dtype=torch.float16 + ).transpose(1, 2) + actual, actual_scale = quantize(x) + expected, expected_scale = quantize(x.contiguous()) + assert actual.is_contiguous() + torch.testing.assert_close( + actual.float(), expected.float(), rtol=0, atol=0 + ) + torch.testing.assert_close(actual_scale, expected_scale, rtol=0, atol=0) From 896e376453d924ce635813c6c739793deda35c5f Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:43:08 +0000 Subject: [PATCH 2/5] Enable Hopper and FP32 inference for FP8 attention --- sdm/_kernels/fp8_attention.py | 12 ++---------- sdm/models/tabiclv2/model.py | 2 +- test/models/tabiclv2/test_model.py | 12 ++++++++---- test/nn/test_attention_quantization.py | 23 +++++++++++++++++------ 4 files changed, 28 insertions(+), 21 deletions(-) diff --git a/sdm/_kernels/fp8_attention.py b/sdm/_kernels/fp8_attention.py index 4b0b95cb52..0b698bfb83 100644 --- a/sdm/_kernels/fp8_attention.py +++ b/sdm/_kernels/fp8_attention.py @@ -23,15 +23,9 @@ def supports_fp8(query: Tensor) -> bool: and not torch.is_grad_enabled() and not torch.compiler.is_compiling() and query.is_cuda - and ( - query.dtype in {torch.float16, torch.bfloat16} - or ( - query.dtype == torch.float32 - and torch.is_autocast_enabled("cuda") - ) - ) and query.size(-1) in {32, 64, 128, 256} - and torch.cuda.get_device_capability(query.device) in {(8, 9), (12, 0)} + and torch.cuda.get_device_capability(query.device) + in {(8, 9), (9, 0), (12, 0)} ) @@ -43,8 +37,6 @@ def fp8_attention( scale: float | None = None, ) -> tuple[Tensor, QuantizedKVCacheEntry]: assert _triton is not None - # RMSNorm keeps Q/K in FP32. Match SDPA's autocast input conversion before - # quantizing, without ever dequantizing cached FP8 keys and values. if torch.is_autocast_enabled("cuda"): dtype = torch.get_autocast_dtype("cuda") query = query.to(dtype) diff --git a/sdm/models/tabiclv2/model.py b/sdm/models/tabiclv2/model.py index 811b470c80..2ec4652ef6 100644 --- a/sdm/models/tabiclv2/model.py +++ b/sdm/models/tabiclv2/model.py @@ -118,7 +118,7 @@ class TabICLv2(ICLModel): pretrained: Whether to load the pretrained checkpoint. device: The device. attention_quantization: Set to ``"fp8"`` to quantize eligible ICL - attention and its cache during FP16/BF16 CUDA inference on Ada + attention and its cache during CUDA inference on Ada, Hopper, and RTX Blackwell GPUs. Requires Triton and more than 8192 context rows. Weights and the final ICL layer remain unquantized. Other inputs diff --git a/test/models/tabiclv2/test_model.py b/test/models/tabiclv2/test_model.py index 66707c911b..4d3f32aafd 100644 --- a/test/models/tabiclv2/test_model.py +++ b/test/models/tabiclv2/test_model.py @@ -264,9 +264,10 @@ def test_compile(dtype: torch.dtype) -> None: @onlyCUDA -def test_fp8_fit_predict() -> None: - if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}: - pytest.skip("FP8 integration supports Ada and RTX Blackwell") +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +def test_fp8_fit_predict(dtype: torch.dtype) -> None: + if torch.cuda.get_device_capability() not in {(8, 9), (9, 0), (12, 0)}: + pytest.skip("FP8 integration supports Ada, Hopper, and RTX Blackwell") model = TabICLv2( task="regression", pretrained=False, @@ -276,7 +277,10 @@ def test_fp8_fit_predict() -> None: x = torch.randn(8193, 3, device="cuda") y = torch.randn(8193, 1, device="cuda") query = torch.randn(17, 3, device="cuda") - with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16): + with ( + torch.inference_mode(), + torch.autocast("cuda", dtype=dtype, enabled=dtype != torch.float32), + ): expected = model(x, y, query, recipe=Recipe(), num_estimators=1) model.fit(x, y, recipe=Recipe(), num_estimators=1) actual = model.predict(query) diff --git a/test/nn/test_attention_quantization.py b/test/nn/test_attention_quantization.py index 9b317e6540..b7831cbf06 100644 --- a/test/nn/test_attention_quantization.py +++ b/test/nn/test_attention_quantization.py @@ -28,10 +28,12 @@ def test_fp8_small_context_fallback(device: torch.device) -> None: @onlyCUDA -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "dtype", [torch.float16, torch.bfloat16, torch.float32] +) def test_fp8_context_cache_and_chunking(dtype: torch.dtype) -> None: - if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}: - pytest.skip("FP8 integration supports Ada and RTX Blackwell") + if torch.cuda.get_device_capability() not in {(8, 9), (9, 0), (12, 0)}: + pytest.skip("FP8 integration supports Ada, Hopper, and RTX Blackwell") module = TabICLv2TransformerBlock( channels=128, num_heads=2, @@ -45,7 +47,10 @@ def test_fp8_context_cache_and_chunking(dtype: torch.dtype) -> None: context = torch.randn(2, 1, 8193, 128, device="cuda") query = torch.randn(2, 1, 129, 128, device="cuda") joined = torch.cat([context, query], dim=-2) - with torch.inference_mode(), torch.autocast("cuda", dtype=dtype): + with ( + torch.inference_mode(), + torch.autocast("cuda", dtype=dtype, enabled=dtype != torch.float32), + ): expected = reference(joined, context) full = module(joined, context) _, cache = module( @@ -80,10 +85,16 @@ def test_fp8_context_cache_and_chunking(dtype: torch.dtype) -> None: assert cache.key.dtype == torch.float8_e4m3fn assert cache.value.dtype == torch.float8_e4m3fn assert cache.query_scale.dtype == torch.float32 - with torch.inference_mode(), pytest.raises(ValueError, match="dtype"): + with ( + torch.inference_mode(), + torch.autocast( + "cuda", dtype=torch.float16, enabled=dtype == torch.float32 + ), + pytest.raises(ValueError, match="dtype"), + ): module(query, cache) with ( - torch.autocast("cuda", dtype=dtype), + torch.autocast("cuda", dtype=dtype, enabled=dtype != torch.float32), pytest.raises(ValueError, match="inference"), ): module(query, cache) From d0d599896c09ee12f082640167cdffcbaa0ef769 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:58:30 +0000 Subject: [PATCH 3/5] Enable FP8 attention and grouped-query caching in KumoTabular --- sdm/cache.py | 18 ++++ sdm/models/kumo/tabular/icl.py | 27 +++--- sdm/models/kumo/tabular/model.py | 26 +++++- test/models/test_attention_quantization.py | 104 ++++++++++++++++++++- 4 files changed, 155 insertions(+), 20 deletions(-) diff --git a/sdm/cache.py b/sdm/cache.py index 1cf6b365c4..62f85689ed 100644 --- a/sdm/cache.py +++ b/sdm/cache.py @@ -32,6 +32,13 @@ class KVCacheEntry(_KVCacheEntry, DeviceMixin): value: Cached value projection tensor. """ + def select_heads(self, num_heads: int) -> Self: + """Retain the leading key/value heads for grouped-query attention.""" + return self.__class__( + key=self.key[..., :num_heads, :].clone(), + value=self.value[..., :num_heads, :].clone(), + ) + def _tensors(self) -> Iterator[Tensor]: yield self.key yield self.value @@ -56,6 +63,17 @@ class QuantizedKVCacheEntry(DeviceMixin): query_scale: Tensor dtype: torch.dtype + def select_heads(self, num_heads: int) -> Self: + """Retain leading K/V heads and scales, keeping all query scales.""" + return self.__class__( + key=self.key[..., :num_heads, :].clone(), + value=self.value[..., :num_heads, :].clone(), + key_scale=self.key_scale[..., :num_heads, :].clone(), + value_scale=self.value_scale[..., :num_heads, :].clone(), + query_scale=self.query_scale, + dtype=self.dtype, + ) + def _tensors(self) -> Iterator[Tensor]: yield self.key yield self.value diff --git a/sdm/models/kumo/tabular/icl.py b/sdm/models/kumo/tabular/icl.py index b7c3a385c5..68a097e8f8 100644 --- a/sdm/models/kumo/tabular/icl.py +++ b/sdm/models/kumo/tabular/icl.py @@ -9,7 +9,7 @@ from torch import Tensor from torch.nn import GELU, Embedding, Linear, ModuleList, RMSNorm, Sequential -from sdm.cache import Cache, KVCacheEntry +from sdm.cache import Cache, KVCacheEntry, QuantizedKVCacheEntry from sdm.models.kumo.tabular.block import KumoTabularTransformerBlock from sdm.nn import LogScale @@ -83,7 +83,10 @@ def forward( result = layer( query=x[..., R_train:, :] if last_layer else x, key_value=( - cast(KVCacheEntry, cache[cache_key]) + cast( + KVCacheEntry | QuantizedKVCacheEntry, + cache[cache_key], + ) if cache is not None and cache.is_replaying else x[..., :R_train, :] ), @@ -96,18 +99,19 @@ def forward( ) if cache is not None and cache.is_recording: - x, (key, value) = result - if self.kv_heads is not None: - key = key[..., : self.kv_heads, :].contiguous() - value = value[..., : self.kv_heads, :].contiguous() - cache[cache_key] = KVCacheEntry(key, value) - del key, value + x, entry = result + cache[cache_key] = ( + entry.select_heads(self.kv_heads) + if self.kv_heads is not None + else entry + ) + del entry else: x = result del result continue - x_context, (key, value) = layer( + x_context, entry = layer( query=x[..., :0, :] if last_layer else x[..., :R_train, :], key_value=x[..., :R_train, :], return_key_value=True, @@ -119,10 +123,7 @@ def forward( ) x_query = layer( query=x[..., R_train:, :], - key_value=KVCacheEntry( - key=key[..., : self.kv_heads, :].contiguous(), - value=value[..., : self.kv_heads, :].contiguous(), - ), + key_value=entry.select_heads(self.kv_heads), out=None if torch.is_grad_enabled() else x[..., R_train:, :], ) if last_layer: diff --git a/sdm/models/kumo/tabular/model.py b/sdm/models/kumo/tabular/model.py index 6fde07b205..5b692451d3 100644 --- a/sdm/models/kumo/tabular/model.py +++ b/sdm/models/kumo/tabular/model.py @@ -18,6 +18,7 @@ from sdm.models.kumo.tabular.icl import ICLBlock from sdm.models.kumo.tabular.recipe import default_recipe from sdm.models.kumo.tabular.row_embedding import RowEmbedding +from sdm.nn import TransformerBlock from sdm.tensor.table import TableSchema MODEL_KWARGS: dict[str, dict[str, Any]] = { @@ -60,6 +61,13 @@ class KumoTabular(ICLModel): pretrained: Whether to load pretrained checkpoints. device: The device for model parameters. If ``None``, uses PyTorch's default device. + attention_quantization: Set to ``"fp8"`` for eligible ICL attention + and its cache during FP16/BF16 CUDA inference on Ada or RTX + Blackwell. Requires Triton and more than 8192 context rows. + Weights and the final ICL layer remain unquantized. Defaults to + ``None``. Other inputs use ordinary attention; fitted FP8 caches + require supported inference with the same dtype and do not + support ``torch.compile``. Cached prediction is not always faster. """ supported_feature_stypes: ClassVar[frozenset[Stype]] = frozenset( @@ -77,10 +85,12 @@ def __init__( size: Literal["small", "large"] = "large", pretrained: bool = True, device: torch.device | str | None = None, + *, + attention_quantization: Literal["fp8"] | None = None, ) -> None: super().__init__(task=task) - self.models: ModuleDict[TaskLike, torch.nn.Module] = ModuleDict() + self.models: ModuleDict = ModuleDict() for task in self.tasks: self.models[task] = _KumoTabular( num_classes=10 if task == Task.classification else 0, @@ -90,8 +100,16 @@ def __init__( ) if pretrained: - self.models[task] = self._load_from_pretrained(size, device=device) + self._load_from_pretrained(size, device=device) + if attention_quantization == "fp8": + for network in self.models.values(): + for layer in cast(_KumoTabular, network).icl_block.layers[:-1]: + cast( + TransformerBlock, layer + ).attn.attention_quantization = "fp8" + else: + assert attention_quantization is None self.eval() @classmethod @@ -103,7 +121,7 @@ def _load_from_pretrained( self, size: Literal["small", "large"], device: torch.device | str | None, - ) -> _KumoTabular: + ) -> None: device = torch.get_default_device() if device is None else device for task, model in self.models.items(): @@ -126,8 +144,6 @@ def _load_from_pretrained( ) model.load_state_dict(ckpt, assign=True) - return model - def forward(self, *args: Any, **kwargs: Any) -> TableTensor: r""":meta private:""" # noqa: D415 x_context = kwargs["x_context"] if "x_context" in kwargs else args[0] diff --git a/test/models/test_attention_quantization.py b/test/models/test_attention_quantization.py index a664caad03..4bb61d11d2 100644 --- a/test/models/test_attention_quantization.py +++ b/test/models/test_attention_quantization.py @@ -1,14 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import copy +from typing import Any, cast + import pytest import torch -from sdm.models import TabICLv2 +from sdm import Recipe +from sdm.cache import Cache, QuantizedKVCacheEntry +from sdm.models import KumoTabular, TabICLv2 +from sdm.models.kumo.tabular.icl import ICLBlock as KumoICLBlock from sdm.nn import Attention +from sdm.testing import onlyCUDA -@pytest.mark.parametrize("model_cls", [TabICLv2]) +@pytest.mark.parametrize("model_cls", [TabICLv2, KumoTabular]) def test_model_fp8_opt_in(model_cls: type[torch.nn.Module]) -> None: model = model_cls( task="regression", @@ -31,3 +38,96 @@ def test_model_fp8_opt_in(model_cls: type[torch.nn.Module]) -> None: for module in ordinary.modules() if isinstance(module, Attention) ) + + +@onlyCUDA +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("heads", [None, 1, 2]) +def test_fp8_icl_cache(heads: int | None, dtype: torch.dtype) -> None: + if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}: + pytest.skip("FP8 integration supports Ada and RTX Blackwell") + block = KumoICLBlock( + num_classes=0, + out_channels=5, + channels=256, + num_layers=2, + num_heads=4, + num_key_value_heads_for_query=heads, + device="cuda", + ) + for parameter in block.parameters(): + if not parameter.any(): + torch.nn.init.normal_(parameter, std=0.02) + reference = copy.deepcopy(block) + cast(Attention, block.layers[0].attn).attention_quantization = "fp8" + x = torch.randn(2, 8226, 256, device="cuda") + y = torch.randn(2, 8193, device="cuda") + with torch.inference_mode(), torch.autocast("cuda", dtype=dtype): + expected = reference(x.clone(), y) + actual = block(x.clone(), y) + cache = Cache() + block(x[:, :8193].clone(), y, cache=cache) + replayed = block(x[:, 8193:].clone(), y[:, :0], cache=cache.freeze()) + split = torch.cat( + [ + block(x[:, 8193:8209].clone(), y[:, :0], cache=cache), + block(x[:, 8209:].clone(), y[:, :0], cache=cache), + ], + dim=-2, + ) + entry = cache["icl_block.layer0"] + assert isinstance(entry, QuantizedKVCacheEntry) + assert entry.key.dtype == torch.float8_e4m3fn + assert entry.value.dtype == torch.float8_e4m3fn + assert entry.query_scale.size(-2) == 4 + if heads is not None: + assert entry.key.size(-2) == heads + assert entry.key.untyped_storage().nbytes() == entry.key.numel() + assert entry.value.untyped_storage().nbytes() == entry.value.numel() + assert actual.isfinite().all() + torch.testing.assert_close(replayed, actual, atol=0.01, rtol=0.03) + torch.testing.assert_close(split, replayed, atol=0.01, rtol=0.03) + relative_rms = ( + (actual.float() - expected.float()).square().mean() + / expected.float().square().mean() + ).sqrt() + assert relative_rms < 0.08 + + +@onlyCUDA +@pytest.mark.parametrize("kind", ["kumo-small", "kumo-large"]) +def test_fp8_model_fit_predict(kind: str) -> None: + if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}: + pytest.skip("FP8 integration supports Ada and RTX Blackwell") + kwargs: dict[str, Any] = { + "task": "regression", + "pretrained": False, + "device": "cuda", + "attention_quantization": "fp8", + } + model = KumoTabular( + size="small" if kind == "kumo-small" else "large", **kwargs + ) + x = torch.randn(8193, 3, device="cuda") + y = torch.randn(8193, 1, device="cuda") + query = torch.randn(17, 3, device="cuda") + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16): + expected = model(x, y, query, recipe=Recipe(), num_estimators=1) + model.fit(x, y, recipe=Recipe(), num_estimators=1) + actual = model.predict(query) + assert actual.numerical.isfinite().all() + torch.testing.assert_close( + actual.numerical, expected.numerical, atol=0.01, rtol=0.03 + ) + assert model._cache is not None + entries = cast(Cache, model._cache[0]) + quantized = [ + entry + for entry in entries.values() + if isinstance(entry, QuantizedKVCacheEntry) + ] + assert quantized + if kind == "kumo-large": + assert all(entry.key.size(-2) == 2 for entry in quantized) + model.clear() + assert model._cache is None From 9f445c22440835645fda74eab6b84de10f6f56d8 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:43:08 +0000 Subject: [PATCH 4/5] Cover FP32 KumoTabular attention and document Hopper support --- sdm/models/kumo/tabular/model.py | 2 +- test/models/test_attention_quantization.py | 17 +++++++++++------ 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/sdm/models/kumo/tabular/model.py b/sdm/models/kumo/tabular/model.py index 5b692451d3..080bd5c7e1 100644 --- a/sdm/models/kumo/tabular/model.py +++ b/sdm/models/kumo/tabular/model.py @@ -62,7 +62,7 @@ class KumoTabular(ICLModel): device: The device for model parameters. If ``None``, uses PyTorch's default device. attention_quantization: Set to ``"fp8"`` for eligible ICL attention - and its cache during FP16/BF16 CUDA inference on Ada or RTX + and its cache during CUDA inference on Ada, Hopper, or RTX Blackwell. Requires Triton and more than 8192 context rows. Weights and the final ICL layer remain unquantized. Defaults to ``None``. Other inputs use ordinary attention; fitted FP8 caches diff --git a/test/models/test_attention_quantization.py b/test/models/test_attention_quantization.py index 4bb61d11d2..d1232728ab 100644 --- a/test/models/test_attention_quantization.py +++ b/test/models/test_attention_quantization.py @@ -41,11 +41,13 @@ def test_model_fp8_opt_in(model_cls: type[torch.nn.Module]) -> None: @onlyCUDA -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "dtype", [torch.float16, torch.bfloat16, torch.float32] +) @pytest.mark.parametrize("heads", [None, 1, 2]) def test_fp8_icl_cache(heads: int | None, dtype: torch.dtype) -> None: - if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}: - pytest.skip("FP8 integration supports Ada and RTX Blackwell") + if torch.cuda.get_device_capability() not in {(8, 9), (9, 0), (12, 0)}: + pytest.skip("FP8 integration supports Ada, Hopper, and RTX Blackwell") block = KumoICLBlock( num_classes=0, out_channels=5, @@ -62,7 +64,10 @@ def test_fp8_icl_cache(heads: int | None, dtype: torch.dtype) -> None: cast(Attention, block.layers[0].attn).attention_quantization = "fp8" x = torch.randn(2, 8226, 256, device="cuda") y = torch.randn(2, 8193, device="cuda") - with torch.inference_mode(), torch.autocast("cuda", dtype=dtype): + with ( + torch.inference_mode(), + torch.autocast("cuda", dtype=dtype, enabled=dtype != torch.float32), + ): expected = reference(x.clone(), y) actual = block(x.clone(), y) cache = Cache() @@ -97,8 +102,8 @@ def test_fp8_icl_cache(heads: int | None, dtype: torch.dtype) -> None: @onlyCUDA @pytest.mark.parametrize("kind", ["kumo-small", "kumo-large"]) def test_fp8_model_fit_predict(kind: str) -> None: - if torch.cuda.get_device_capability() not in {(8, 9), (12, 0)}: - pytest.skip("FP8 integration supports Ada and RTX Blackwell") + if torch.cuda.get_device_capability() not in {(8, 9), (9, 0), (12, 0)}: + pytest.skip("FP8 integration supports Ada, Hopper, and RTX Blackwell") kwargs: dict[str, Any] = { "task": "regression", "pretrained": False, From 30ed9374eec46af367ad0085e13454f2afeb3c0b Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:58:30 +0000 Subject: [PATCH 5/5] Enable FP8 attention in KumoRelational and TabFM --- sdm/models/kumo/relational/model.py | 24 +++- sdm/models/tabfm/model.py | 22 +++- test/models/test_attention_quantization.py | 128 +++++++++++++++++---- 3 files changed, 149 insertions(+), 25 deletions(-) diff --git a/sdm/models/kumo/relational/model.py b/sdm/models/kumo/relational/model.py index 16026f3387..974ef17daf 100644 --- a/sdm/models/kumo/relational/model.py +++ b/sdm/models/kumo/relational/model.py @@ -5,7 +5,7 @@ from __future__ import annotations from collections.abc import Iterable, Sequence -from typing import Any, ClassVar, cast +from typing import Any, ClassVar, Literal, cast import torch from torch import Tensor @@ -29,6 +29,7 @@ from sdm.models.kumo.relational.task import TaskGraph from sdm.models.tabiclv2.icl import ICLBlock from sdm.models.tabiclv2.row_embedding import RowEmbedding +from sdm.nn import TransformerBlock from sdm.processing import Recipe, Standardize @@ -141,6 +142,13 @@ class KumoRelational(ICLModel): model are initialized. pretrained: Whether to load the pretrained checkpoint. device: The device. + attention_quantization: Set to ``"fp8"`` for eligible ICL attention + and its cache during CUDA inference on Ada, Hopper, or RTX + Blackwell. Requires Triton and more than 8192 context rows. + Weights and the final ICL layer remain unquantized. Defaults to + ``None``. Other inputs use ordinary attention; fitted FP8 caches + require supported inference with the same dtype and do not + support ``torch.compile``. Cached prediction is not always faster. """ supported_feature_stypes: ClassVar[frozenset[Stype]] = frozenset( @@ -157,10 +165,12 @@ def __init__( task: TaskLike | Iterable[TaskLike] | None = None, pretrained: bool = True, device: torch.device | str | None = None, + *, + attention_quantization: Literal["fp8"] | None = None, ) -> None: super().__init__(task=task) - self.models: ModuleDict[TaskLike, torch.nn.Module] = ModuleDict() + self.models: ModuleDict = ModuleDict() for task in self.tasks: self.models[task] = _KumoRelational( num_classes=10 if task == Task.classification else 0, @@ -172,6 +182,16 @@ def __init__( if pretrained: self._load_from_pretrained(device=device) + if attention_quantization == "fp8": + for network in self.models.values(): + for layer in cast(_KumoRelational, network).icl_block.layers[ + :-1 + ]: + cast( + TransformerBlock, layer + ).attn.attention_quantization = "fp8" + else: + assert attention_quantization is None self.eval() def _load_from_pretrained( diff --git a/sdm/models/tabfm/model.py b/sdm/models/tabfm/model.py index 8b53228bf5..f3dcc3c70f 100644 --- a/sdm/models/tabfm/model.py +++ b/sdm/models/tabfm/model.py @@ -20,7 +20,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any, ClassVar, cast +from typing import Any, ClassVar, Literal, cast import torch from torch import Tensor @@ -34,6 +34,7 @@ from sdm.models.tabfm.icl import ICLBlock from sdm.models.tabfm.recipe import default_recipe from sdm.models.tabfm.row_embedding import RowEmbedding +from sdm.nn import TransformerBlock from sdm.tensor.table import TableSchema @@ -73,6 +74,13 @@ class TabFM(ICLModel): v1.0 `__ without showing the interactive license prompt. device: The device. + attention_quantization: Set to ``"fp8"`` for eligible ICL attention + and its cache during CUDA inference on Ada, Hopper, or RTX + Blackwell. Requires Triton and more than 8192 context rows. + Weights and the final ICL layer remain unquantized. Defaults to + ``None``. Other inputs use ordinary attention; fitted FP8 caches + require supported inference with the same dtype and do not + support ``torch.compile``. Cached prediction is not always faster. """ supported_feature_stypes: ClassVar[frozenset[Stype]] = frozenset( @@ -90,10 +98,12 @@ def __init__( pretrained: bool = True, accept_license: bool = False, device: torch.device | str | None = None, + *, + attention_quantization: Literal["fp8"] | None = None, ) -> None: super().__init__(task=task) - self.models: ModuleDict[TaskLike, torch.nn.Module] = ModuleDict() + self.models: ModuleDict = ModuleDict() for task in self.tasks: self.models[task] = _TabFM( num_classes=10 if task == Task.classification else 0, @@ -103,6 +113,14 @@ def __init__( if pretrained: self._load_from_pretrained(accept_license, device=device) + if attention_quantization == "fp8": + for network in self.models.values(): + for layer in cast(_TabFM, network).icl_block.layers[:-1]: + cast( + TransformerBlock, layer + ).attn.attention_quantization = "fp8" + else: + assert attention_quantization is None self.eval() @classmethod diff --git a/test/models/test_attention_quantization.py b/test/models/test_attention_quantization.py index d1232728ab..4f020d7b0a 100644 --- a/test/models/test_attention_quantization.py +++ b/test/models/test_attention_quantization.py @@ -2,20 +2,25 @@ # SPDX-License-Identifier: Apache-2.0 import copy +import functools from typing import Any, cast import pytest import torch -from sdm import Recipe +from sdm import ColumnarTensor, Recipe, RelatedTables, TableTensor from sdm.cache import Cache, QuantizedKVCacheEntry -from sdm.models import KumoTabular, TabICLv2 +from sdm.models import KumoRelational, KumoTabular, TabFM, TabICLv2 from sdm.models.kumo.tabular.icl import ICLBlock as KumoICLBlock +from sdm.models.tabfm import model as tabfm_module +from sdm.models.tabfm.icl import ICLBlock as TabFMICLBlock from sdm.nn import Attention from sdm.testing import onlyCUDA -@pytest.mark.parametrize("model_cls", [TabICLv2, KumoTabular]) +@pytest.mark.parametrize( + "model_cls", [TabICLv2, KumoTabular, KumoRelational, TabFM] +) def test_model_fp8_opt_in(model_cls: type[torch.nn.Module]) -> None: model = model_cls( task="regression", @@ -44,25 +49,42 @@ def test_model_fp8_opt_in(model_cls: type[torch.nn.Module]) -> None: @pytest.mark.parametrize( "dtype", [torch.float16, torch.bfloat16, torch.float32] ) -@pytest.mark.parametrize("heads", [None, 1, 2]) -def test_fp8_icl_cache(heads: int | None, dtype: torch.dtype) -> None: +@pytest.mark.parametrize( + ("kind", "heads"), + [("kumo", None), ("kumo", 1), ("kumo", 2), ("tabfm", None)], +) +def test_fp8_icl_cache( + kind: str, heads: int | None, dtype: torch.dtype +) -> None: if torch.cuda.get_device_capability() not in {(8, 9), (9, 0), (12, 0)}: pytest.skip("FP8 integration supports Ada, Hopper, and RTX Blackwell") - block = KumoICLBlock( - num_classes=0, - out_channels=5, - channels=256, - num_layers=2, - num_heads=4, - num_key_value_heads_for_query=heads, - device="cuda", - ) + if kind == "kumo": + block = KumoICLBlock( + num_classes=0, + out_channels=5, + channels=256, + num_layers=2, + num_heads=4, + num_key_value_heads_for_query=heads, + device="cuda", + ) + else: + # Stock TabFM has 256 channels per ICL attention head. + block = TabFMICLBlock( + num_classes=0, + out_channels=5, + channels=512, + num_layers=2, + num_heads=2, + device="cuda", + ) for parameter in block.parameters(): if not parameter.any(): torch.nn.init.normal_(parameter, std=0.02) reference = copy.deepcopy(block) cast(Attention, block.layers[0].attn).attention_quantization = "fp8" - x = torch.randn(2, 8226, 256, device="cuda") + channels = 256 if kind == "kumo" else 512 + x = torch.randn(2, 8226, channels, device="cuda") y = torch.randn(2, 8193, device="cuda") with ( torch.inference_mode(), @@ -84,7 +106,7 @@ def test_fp8_icl_cache(heads: int | None, dtype: torch.dtype) -> None: assert isinstance(entry, QuantizedKVCacheEntry) assert entry.key.dtype == torch.float8_e4m3fn assert entry.value.dtype == torch.float8_e4m3fn - assert entry.query_scale.size(-2) == 4 + assert entry.query_scale.size(-2) == (4 if kind == "kumo" else 2) if heads is not None: assert entry.key.size(-2) == heads assert entry.key.untyped_storage().nbytes() == entry.key.numel() @@ -100,8 +122,10 @@ def test_fp8_icl_cache(heads: int | None, dtype: torch.dtype) -> None: @onlyCUDA -@pytest.mark.parametrize("kind", ["kumo-small", "kumo-large"]) -def test_fp8_model_fit_predict(kind: str) -> None: +@pytest.mark.parametrize("kind", ["kumo-small", "kumo-large", "tabfm"]) +def test_fp8_model_fit_predict( + kind: str, monkeypatch: pytest.MonkeyPatch +) -> None: if torch.cuda.get_device_capability() not in {(8, 9), (9, 0), (12, 0)}: pytest.skip("FP8 integration supports Ada, Hopper, and RTX Blackwell") kwargs: dict[str, Any] = { @@ -110,9 +134,23 @@ def test_fp8_model_fit_predict(kind: str) -> None: "device": "cuda", "attention_quantization": "fp8", } - model = KumoTabular( - size="small" if kind == "kumo-small" else "large", **kwargs - ) + if kind.startswith("kumo-"): + model = KumoTabular( + size="small" if kind == "kumo-small" else "large", **kwargs + ) + else: + # Keep the stock 256-channel ICL heads while reducing model depth. + monkeypatch.setattr( + tabfm_module, + "_TabFM", + functools.partial( + tabfm_module._TabFM, + num_icl_layers=2, + num_embedding_layers=1, + num_embedding_repeats=1, + ), + ) + model = TabFM(**kwargs) x = torch.randn(8193, 3, device="cuda") y = torch.randn(8193, 1, device="cuda") query = torch.randn(17, 3, device="cuda") @@ -136,3 +174,51 @@ def test_fp8_model_fit_predict(kind: str) -> None: assert all(entry.key.size(-2) == 2 for entry in quantized) model.clear() assert model._cache is None + + +@onlyCUDA +def test_fp8_relational_fit_predict() -> None: + if torch.cuda.get_device_capability() not in {(8, 9), (9, 0), (12, 0)}: + pytest.skip("FP8 integration supports Ada, Hopper, and RTX Blackwell") + model = KumoRelational( + task="regression", + pretrained=False, + device="cuda", + attention_quantization="fp8", + ) + table = TableTensor( + id=ColumnarTensor((torch.arange(8210, device="cuda"),)), + numerical=torch.randn(8210, 1, device="cuda"), + columns={"id": ["id"], "numerical": ["value"]}, + ) + related = RelatedTables( + tables={"entities": table}, + relationships=[], + task_links=[ + {"task_column": "id", "table": "entities", "table_column": "id"} + ], + ) + x, query = table[:8193], table[8193:] + y = torch.randn(8193, 1, device="cuda") + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16): + expected = model( + x, + y, + query, + related_context_tables=related, + related_query_tables=related, + num_hops=0, + recipe=Recipe(), + num_estimators=1, + ) + model.fit(x, y, related, num_hops=0, recipe=Recipe(), num_estimators=1) + actual = model.predict(query, related_tables=related) + assert actual.numerical.isfinite().all() + torch.testing.assert_close( + actual.numerical, expected.numerical, atol=0.01, rtol=0.03 + ) + assert model._cache is not None + entries = cast(Cache, model._cache[0]) + assert any( + isinstance(entry, QuantizedKVCacheEntry) for entry in entries.values() + )