[FP8 3/3] Enable FP8 attention in KumoRelational and TabFM

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

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

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."""