[KV] Opt-in FP8 e4m3 KV cache with per-vector scales

docs/runtime.md

@@ -304,6 +304,31 @@ zero watermark in non-incremental mode, see [scheduler admission](scheduler.md#a
 Pool sizing establishes physical capacity; scheduler admission decides how much
 of that capacity can safely be promised to new work.
 
+## FP8 KV cache
+
+`CacheConfig(kv_cache_dtype="fp8_e4m3")` (CLI `--kv-cache-dtype fp8_e4m3`) stores K and V as
+float8 e4m3 codes with one float32 scale per (block, layer, token, KV head) vector
+(`scale = amax / 448`). Values are quantized when written and dequantized by the attention
+backend when read, so a block costs `head_dim + 4` bytes per vector instead of
+`head_dim * element_size`: 51.6% of BF16 at `head_dim=128`, and automatic sizing or
+`--kv-cache-memory-bytes` admits about 1.94x as many tokens. The default `"auto"` stores the
+model dtype and leaves every existing path unchanged.
+
+- Supported: Torch (reference) and Triton attention; synchronous and asynchronous scheduling;
+  CUDA graphs (quantized writes are captured); LAST_EXITED finalize copies, including the
+  asynchronous finalize kernel, copy codes and scales; preemption snapshots carry the scales;
+  `KVCacheManager.read` returns dequantized values.
+- Triton's split-KV decode and packed paged-prefill kernels both apply the scales on load
+  (prefill dequantizes to the query dtype before `tl.dot`). Device-masked writes (speculative
+  verification, bucketed CUDA graphs) use the same fused writer with the row mask applied.
+- Triton writes with one fused quantize-and-store kernel per layer. The Torch backend quantizes
+  with Torch ops and is the reference; both writers produce the same codes and scales.
+- Rejected: FlashAttention backends (the `fp8_kv` backend capability; `auto` selects Triton)
+  and PD transfer, which does not yet move the scales.
+- Each layer writes the current token's K/V before attending, so attention reads the current
+  token in FP8 as well. Accuracy should be checked per model and workload; adaptive exits can
+  change where a token exits because the gate sees slightly different attention.
+
 ## Replay and validation commands
 
 Trace mode expects `depths_by_request`, with one depth per output: full prefill depth first, then each decode input's exit depth. Admission validates coverage and loop bounds. Request IDs must match the trace. The public engine replay fixes depth decisions; token generation still follows the configured sampler.

tests/test_cuda_graph.py

@@ -455,7 +455,8 @@ def test_bucketed_kv_graph_combines_active_rows_and_pipeline_write_mask():
 
     keys = torch.full((4, 1, 2, 1, 4), -1.0, device="cuda")
     values = torch.full_like(keys, -2.0)
-    cache = SimpleNamespace(key_cache=keys, value_cache=values)
+    # A BF16 pool has no FP8 scales, as KVCacheManager reports them.
+    cache = SimpleNamespace(key_cache=keys, value_cache=values, key_scale=None, value_scale=None)
     batch = SimpleNamespace(
         active_rows=torch.zeros((), dtype=torch.int32, device="cuda"),
         write_mask=torch.ones(5, dtype=torch.bool, device="cuda"),

vllm_rlt/attention/backend.py

@@ -21,6 +21,8 @@ class BackendCapabilities:
     # Packed rows reproduce single-query decode exactly; speculative verification
     # relies on this to match serial decoding token for token.
     packed_matches_decode: bool = False
+    # Reads an FP8 e4m3 KV pool with per-vector float32 scales.
+    fp8_kv: bool = False
     # Consume GPU-owned context lengths without a host read between submissions.
     device_context_lengths: bool = False
 
@@ -45,9 +47,9 @@ def generation(self):
         return getattr(self.kernel, "generation", None)
 
 
-_TORCH_CAPABILITIES = BackendCapabilities()
+_TORCH_CAPABILITIES = BackendCapabilities(fp8_kv=True)
 _TRITON_CAPABILITIES = BackendCapabilities(
-    cuda_graphs=True, async_scheduling=True, device_context_lengths=True
+    cuda_graphs=True, async_scheduling=True, device_context_lengths=True, fp8_kv=True
 )
 _FLASH_CAPABILITIES = {
     2: BackendCapabilities(cuda_graphs=True, async_scheduling=True, device_context_lengths=True),
@@ -128,6 +130,7 @@ def create_backend(
                     "cuda_graphs",
                     "async_scheduling",
                     "device_context_lengths",
+                    "fp8_kv",
                 )
                 if getattr(required_capabilities, name)
                 and not getattr(implementation.capabilities, name)

vllm_rlt/config.py

@@ -20,6 +20,9 @@ class CacheConfig:
     incremental_allocation: bool = False
     # Admission headroom for incremental allocation; otherwise effective watermark is zero.
     watermark_ratio: float = 0.0
+    # "auto" stores KV in the model dtype; "fp8_e4m3" stores float8 e4m3 codes with
+    # one float32 scale per (block, layer, token, KV head) vector (Torch/Triton only).
+    kv_cache_dtype: str = "auto"
 
     def __post_init__(self):
         for name in ("enable_prefix_caching", "incremental_allocation"):
@@ -34,6 +37,8 @@ def __post_init__(self):
         _positive("block_size", self.block_size)
         if self.layout not in ("last_exited", "shared"):
             raise ValueError("layout must be 'last_exited' or 'shared'")
+        if self.kv_cache_dtype not in ("auto", "fp8_e4m3"):
+            raise ValueError("kv_cache_dtype must be 'auto' or 'fp8_e4m3'")
         if (
             not math.isfinite(self.gpu_memory_utilization)
             or not 0 < self.gpu_memory_utilization <= 1

vllm_rlt/core/kv_cache_manager.py

@@ -4,6 +4,10 @@
 SHARED overwrites a single plane and needs no exit copies. Optional incremental
 allocation separates logical capacity from physical pages. Immutable full-depth
 prompt blocks can be shared through the prefix cache.
+
+With ``kv_cache_dtype="fp8_e4m3"`` the pool stores K/V as float8 e4m3 with one
+float32 scale per (block, layer, token, KV head) vector: values are quantized on
+write and dequantized by the attention backends on read.
 """
 
 import hashlib
@@ -27,6 +31,33 @@
 from vllm_rlt.worker.buffers import stage_metadata
 from vllm_rlt.worker.metadata import DeviceKVMetadata, KVMetadata, MetadataBuilder
 
+KV_CACHE_DTYPES = ("auto", "fp8_e4m3")
+FP8_E4M3_MAX = 448.0
+
+
+def kv_bytes_per_block(
+    num_layers, block_size, num_kv_heads, head_dim, element_size, kv_cache_dtype
+):
+    """Physical bytes of one block (K and V for all layers), including FP8 scales."""
+    if kv_cache_dtype not in KV_CACHE_DTYPES:
+        raise ValueError(f"kv_cache_dtype must be one of {KV_CACHE_DTYPES}")
+    vector = head_dim * element_size if kv_cache_dtype == "auto" else head_dim + 4
+    return 2 * num_layers * block_size * num_kv_heads * vector
+
+
+def quantize_fp8(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-vector symmetric e4m3: x [..., head_dim] -> (codes, float32 scale [...])."""
+    x = x.float()
+    # A reciprocal multiply, as CUDA division by a scalar does anyway: CPU, CUDA and
+    # the fused Triton writer then produce the same scale bits.
+    scale = (x.abs().amax(-1) * (1.0 / FP8_E4M3_MAX)).clamp_min(torch.finfo(torch.float32).tiny)
+    codes = (x / scale.unsqueeze(-1)).clamp(-FP8_E4M3_MAX, FP8_E4M3_MAX)
+    return codes.to(torch.float8_e4m3fn), scale
+
+
+def dequantize_fp8(codes: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    return (codes.float() * scale.unsqueeze(-1)).to(dtype)
+
 
 @dataclass
 class _WrittenPositions:
@@ -75,6 +106,9 @@ class KVSnapshot:
     written: list[list[_WrittenPositions]]
     keys: torch.Tensor
     values: torch.Tensor
+    # FP8 pools only: per-vector scales, laid out like keys/values minus head_dim.
+    key_scales: torch.Tensor | None = None
+    value_scales: torch.Tensor | None = None
 
 
 class KVCacheManager:
@@ -103,6 +137,7 @@ def __init__(
         recurrent_layers: Sequence[int] | None = None,
         *,
         attention: AttentionBackend | None = None,
+        kv_cache_dtype: str = "auto",
     ):
         for name, value in (
             ("num_layers", num_layers),
@@ -120,6 +155,9 @@ def __init__(
             raise ValueError("KV dtype must be float32, float16, or bfloat16")
         if layout not in {"last_exited", "shared"}:
             raise ValueError("unsupported KV layout")
+        if kv_cache_dtype not in KV_CACHE_DTYPES:
+            raise ValueError(f"kv_cache_dtype must be one of {KV_CACHE_DTYPES}")
+        self.kv_cache_dtype = kv_cache_dtype
         self.layout = layout
         self.recurrent_layers = (
             tuple(range(num_layers)) if recurrent_layers is None else tuple(recurrent_layers)
@@ -140,9 +178,17 @@ def __init__(
         )
         self.attention_info = self.attention.info
         self.backend = self.attention_info["selected_backend"]
+        if kv_cache_dtype != "auto" and not getattr(self.attention.capabilities, "fp8_kv", False):
+            raise ValueError(f"the {self.attention_info['backend']} backend does not read FP8 KV")
         shape = (num_blocks, num_layers, block_size, num_kv_heads, head_dim)
-        self.key_cache = torch.empty(shape, device=self.device, dtype=dtype)
+        storage = torch.float8_e4m3fn if kv_cache_dtype == "fp8_e4m3" else dtype
+        self.key_cache = torch.empty(shape, device=self.device, dtype=storage)
         self.value_cache = torch.empty_like(self.key_cache)
+        # One scale per stored vector; None for an unquantized pool.
+        self.key_scale = self.value_scale = None
+        if kv_cache_dtype == "fp8_e4m3":
+            self.key_scale = torch.ones(shape[:-1], device=self.device, dtype=torch.float32)
+            self.value_scale = torch.ones_like(self.key_scale)
         # Resolve an implicit CUDA index to the actual storage device once.
         self.device = self.key_cache.device
         self._free_blocks = list(reversed(range(num_blocks)))
@@ -174,13 +220,13 @@ def num_used_blocks(self) -> int:
 
     @property
     def bytes_per_block(self) -> int:
-        return (
-            2
-            * self.num_layers
-            * self.block_size
-            * self.num_kv_heads
-            * self.head_dim
-            * self.key_cache.element_size()
+        return kv_bytes_per_block(
+            self.num_layers,
+            self.block_size,
+            self.num_kv_heads,
+            self.head_dim,
+            self.key_cache.element_size(),
+            self.kv_cache_dtype,
         )
 
     def required_blocks(self, max_tokens: int) -> int:
@@ -477,17 +523,23 @@ def snapshot(self, request_id: str) -> KVSnapshot:
         if allocation.transfer_leases:
             raise RuntimeError("cannot snapshot KV during a transfer")
         blocks = [b for table in allocation.block_tables for b in table]
-        keys = torch.empty((len(blocks), *self.key_cache.shape[1:]), dtype=self.dtype)
+        keys = torch.empty((len(blocks), *self.key_cache.shape[1:]), dtype=self.key_cache.dtype)
         values = torch.empty_like(keys)
         for row, block in enumerate(blocks):
             keys[row].copy_(self.key_cache[block])
             values[row].copy_(self.value_cache[block])
+        key_scales = value_scales = None
+        if self.key_scale is not None:
+            key_scales = self.key_scale[blocks].cpu()
+            value_scales = self.value_scale[blocks].cpu()
         return KVSnapshot(
             max_tokens=allocation.max_tokens,
             pages=len(allocation.block_tables[0]),
             written=deepcopy(allocation.written),
             keys=keys,
             values=values,
+            key_scales=key_scales,
+            value_scales=value_scales,
         )
 
     def restore(self, request_id: str, snapshot: KVSnapshot) -> None:
@@ -504,10 +556,15 @@ def restore(self, request_id: str, snapshot: KVSnapshot) -> None:
             raise ValueError("snapshot page count does not match the allocation")
         if any(w.prefix or w.pending for plane in allocation.written for w in plane):
             raise RuntimeError("cannot restore into an allocation that already holds KV")
+        if (self.key_scale is None) != (snapshot.key_scales is None):
+            raise ValueError("snapshot KV dtype does not match the pool")
         blocks = [b for table in allocation.block_tables for b in table]
         for row, block in enumerate(blocks):
             self.key_cache[block].copy_(snapshot.keys[row])
             self.value_cache[block].copy_(snapshot.values[row])
+            if self.key_scale is not None:
+                self.key_scale[block].copy_(snapshot.key_scales[row])
+                self.value_scale[block].copy_(snapshot.value_scales[row])
         allocation.written = deepcopy(snapshot.written)
 
     def _plane(self, depth: int) -> int:
@@ -669,7 +726,11 @@ def _write_prepared(
         if isinstance(batch, DeviceKVMetadata):
             self._validate_tensor(k, len(batch.position_ids), "k")
             self._validate_tensor(v, len(batch.position_ids), "v")
-            if batch.write_mask is not None:
+            if self.key_scale is not None:
+                self._store_layer(
+                    layer, batch.write_blocks, batch.write_offsets, k, v, mask=batch.write_mask
+                )
+            elif batch.write_mask is not None:
                 from vllm_rlt.kernels.routing import masked_kv_store
 
                 masked_kv_store(self.key_cache[:, layer], self.value_cache[:, layer], k, v, batch)
@@ -683,11 +744,56 @@ def _write_prepared(
         self._validate_tensor(v, len(batch.position_ids), "v")
         if not batch.rows:
             return
-        self.key_cache[batch.write_blocks, layer, batch.write_offsets] = k[: len(batch.rows)]
-        self.value_cache[batch.write_blocks, layer, batch.write_offsets] = v[: len(batch.rows)]
+        rows = len(batch.rows)
+        self._store_layer(layer, batch.write_blocks, batch.write_offsets, k[:rows], v[:rows])
         for allocation, depth, position in batch.rows:
             allocation.written[self._plane(depth)][layer].add(position)
 
+    def _store_layer(self, layer, blocks, offsets, k, v, *, mask=None, active=None) -> None:
+        """Write one layer's K/V rows, quantizing them for an FP8 pool.
+
+        FP8 only: ``mask`` (bool [rows]) and ``active`` (int32 row count) skip
+        rows on the device, like the unquantized masked and graph writers.
+        """
+        if self.key_scale is None:
+            if mask is not None or active is not None:
+                raise ValueError("row masks are handled by the unquantized device writers")
+            self.key_cache[blocks, layer, offsets] = k
+            self.value_cache[blocks, layer, offsets] = v
+            return
+        if self.backend == "triton" or mask is not None or active is not None:
+            # One fused launch instead of about 20 small kernels per layer write.
+            # Masked writes always take it, as the unquantized masked writers are
+            # Triton too and boolean indexing would read the mask on the host.
+            from vllm_rlt.kernels.triton_attention import store_fp8
+
+            store_fp8(
+                k,
+                v,
+                self.key_cache[:, layer],
+                self.value_cache[:, layer],
+                self.key_scale[:, layer],
+                self.value_scale[:, layer],
+                blocks,
+                offsets,
+                FP8_E4M3_MAX,
+                mask=mask,
+                active=active,
+            )
+            return
+        k_codes, k_scale = quantize_fp8(k)
+        v_codes, v_scale = quantize_fp8(v)
+        self.key_cache[blocks, layer, offsets] = k_codes
+        self.value_cache[blocks, layer, offsets] = v_codes
+        self.key_scale[blocks, layer, offsets] = k_scale
+        self.value_scale[blocks, layer, offsets] = v_scale
+
+    def _layer_scales(self, layer) -> dict:
+        """Extra attention arguments for one layer: its scales, or nothing."""
+        if self.key_scale is None:
+            return {}
+        return dict(key_scale=self.key_scale[:, layer], value_scale=self.value_scale[:, layer])
+
     def _require_prefix(self, allocation, layer, depth, length):
         written = allocation.written[self._plane(depth)][layer]
         if length > written.prefix:
@@ -725,6 +831,7 @@ def _attend_prepared(
                 self.value_cache[:, layer],
                 batch.block_tables,
                 batch.context_lengths,
+                **self._layer_scales(layer),
             )
         if not batch.rows:
             return torch.empty_like(q)
@@ -739,13 +846,15 @@ def _attend_prepared(
                 batch.context_lengths,
                 batch.cu_seqlens_q,
                 batch.max_seqlen_q,
+                **self._layer_scales(layer),
             )
         return self.attention(
             q,
             self.key_cache[:, layer],
             self.value_cache[:, layer],
             batch.block_tables,
             batch.context_lengths,
+            **self._layer_scales(layer),
         )
 
     @torch.no_grad()
@@ -762,18 +871,16 @@ def finalize_token(self, request_id: str, position: int, exit_depth: int) -> Non
         ]
         # Basic-index copies avoid a blocking host->device index tensor on the
         # boundary stream. That transfer would serialize final core and routing.
+        tensors = [self.key_cache, self.value_cache]
+        if self.key_scale is not None:
+            tensors += [self.key_scale, self.value_scale]
         for destination in destinations:
-            if len(self.recurrent_layers) == self.num_layers:
-                self.key_cache[destination, :, offset].copy_(self.key_cache[source, :, offset])
-                self.value_cache[destination, :, offset].copy_(self.value_cache[source, :, offset])
-            else:
-                for layer in self.recurrent_layers:
-                    self.key_cache[destination, layer, offset].copy_(
-                        self.key_cache[source, layer, offset]
-                    )
-                    self.value_cache[destination, layer, offset].copy_(
-                        self.value_cache[source, layer, offset]
-                    )
+            for tensor in tensors:
+                if len(self.recurrent_layers) == self.num_layers:
+                    tensor[destination, :, offset].copy_(tensor[source, :, offset])
+                else:
+                    for layer in self.recurrent_layers:
+                        tensor[destination, layer, offset].copy_(tensor[source, layer, offset])
         self._record_finalized(allocation, position, exit_depth)
 
     @torch.no_grad()
@@ -797,4 +904,11 @@ def read(self, layer: int, request_id: str, depth: int, length: int | None = Non
         )
         blocks = table[positions // self.block_size]
         offsets = positions % self.block_size
-        return self.key_cache[blocks, layer, offsets], self.value_cache[blocks, layer, offsets]
+        keys = self.key_cache[blocks, layer, offsets]
+        values = self.value_cache[blocks, layer, offsets]
+        if self.key_scale is None:
+            return keys, values
+        return (
+            dequantize_fp8(keys, self.key_scale[blocks, layer, offsets], self.dtype),
+            dequantize_fp8(values, self.value_scale[blocks, layer, offsets], self.dtype),
+        )