diff --git a/sdm/_kernels/fp8_attention.py b/sdm/_kernels/fp8_attention.py new file mode 100644 index 0000000000..0b698bfb83 --- /dev/null +++ b/sdm/_kernels/fp8_attention.py @@ -0,0 +1,103 @@ +# 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.size(-1) in {32, 64, 128, 256} + and torch.cuda.get_device_capability(query.device) + in {(8, 9), (9, 0), (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 + 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..62f85689ed 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 @@ -30,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 @@ -38,6 +47,51 @@ 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 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 + 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/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/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..080bd5c7e1 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 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( @@ -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/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/sdm/models/tabiclv2/model.py b/sdm/models/tabiclv2/model.py index 495c17b27f..2ec4652ef6 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 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 + 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..4d3f32aafd 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,40 @@ 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 +@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, + 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=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) + 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..4f020d7b0a --- /dev/null +++ b/test/models/test_attention_quantization.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import copy +import functools +from typing import Any, cast + +import pytest +import torch + +from sdm import ColumnarTensor, Recipe, RelatedTables, TableTensor +from sdm.cache import Cache, QuantizedKVCacheEntry +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, KumoRelational, TabFM] +) +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) + ) + + +@onlyCUDA +@pytest.mark.parametrize( + "dtype", [torch.float16, torch.bfloat16, torch.float32] +) +@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") + 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" + 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(), + torch.autocast("cuda", dtype=dtype, enabled=dtype != torch.float32), + ): + 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 kind == "kumo" else 2) + 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", "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] = { + "task": "regression", + "pretrained": False, + "device": "cuda", + "attention_quantization": "fp8", + } + 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") + 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 + + +@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() + ) diff --git a/test/nn/test_attention_quantization.py b/test/nn/test_attention_quantization.py new file mode 100644 index 0000000000..b7831cbf06 --- /dev/null +++ b/test/nn/test_attention_quantization.py @@ -0,0 +1,100 @@ +# 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, torch.float32] +) +def test_fp8_context_cache_and_chunking(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") + 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, enabled=dtype != torch.float32), + ): + 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(), + torch.autocast( + "cuda", dtype=torch.float16, enabled=dtype == torch.float32 + ), + pytest.raises(ValueError, match="dtype"), + ): + module(query, cache) + with ( + torch.autocast("cuda", dtype=dtype, enabled=dtype != torch.float32), + 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)