From f6286f59672de608f5a75bed62dcb6f49de7b71b Mon Sep 17 00:00:00 2001 From: 0z5a Date: Thu, 15 Oct 2026 20:10:02 +0200 Subject: [PATCH 1/5] [KV] FP8 e4m3 KV cache with per-vector scales Opt-in CacheConfig(kv_cache_dtype="fp8_e4m3") / --kv-cache-dtype stores K/V as float8 e4m3 codes with one float32 scale per (block, layer, token, KV head) vector, quantized on write and dequantized by the Torch and Triton attention backends on read (about 51% of BF16 KV bytes). LAST_EXITED finalize copies (sync and async kernel), CUDA-graph writes, preemption snapshots and memory planning carry the scales. FlashAttention backends and PD transfer reject FP8. Co-Authored-By: Claude Opus 5.5 --- tests/test_fp8_kv_cache.py | 244 +++++++++++++++++++++++++++ vllm_rlt/attention/backend.py | 8 +- vllm_rlt/config.py | 5 + vllm_rlt/core/kv_cache_manager.py | 130 +++++++++++--- vllm_rlt/core/memory.py | 17 +- vllm_rlt/engine/llm_engine.py | 2 + vllm_rlt/entrypoints/runtime_args.py | 7 + vllm_rlt/kernels/paged_attention.py | 22 ++- vllm_rlt/kernels/routing.py | 6 +- vllm_rlt/kernels/triton_attention.py | 53 ++++-- vllm_rlt/pd/engine.py | 4 + vllm_rlt/pd/transport.py | 4 + vllm_rlt/worker/async_state.py | 2 +- vllm_rlt/worker/cuda_graph.py | 5 +- vllm_rlt/worker/model_runner.py | 21 +++ 15 files changed, 472 insertions(+), 58 deletions(-) create mode 100644 tests/test_fp8_kv_cache.py diff --git a/tests/test_fp8_kv_cache.py b/tests/test_fp8_kv_cache.py new file mode 100644 index 00000000..ae0a4c70 --- /dev/null +++ b/tests/test_fp8_kv_cache.py @@ -0,0 +1,244 @@ +"""FP8 e4m3 KV pool: per-vector quantization, attention, LAST_EXITED copies and preemption.""" + +import pytest +import torch + +from tests.helpers import tiny_ouro_config +from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig, SamplingParams, SchedulerConfig +from vllm_rlt.attention import AttentionBackend, BackendCapabilities +from vllm_rlt.core.kv_cache_manager import ( + KVCacheManager, + dequantize_fp8, + kv_bytes_per_block, + quantize_fp8, +) +from vllm_rlt.engine.llm_engine import LLMEngine +from vllm_rlt.kernels.paged_attention import torch_paged_attention +from vllm_rlt.models import OuroForCausalLM + + +def make_cache(kv_cache_dtype="fp8_e4m3", **overrides): + values = dict( + num_layers=2, + num_kv_heads=2, + head_dim=8, + num_blocks=64, + block_size=4, + max_loops=4, + kv_cache_dtype=kv_cache_dtype, + ) + values.update(overrides) + return KVCacheManager(**values) + + +def fill(cache, rid, length, depths=range(4), seed=0): + """Write random K/V for positions [0, length) at every depth and layer.""" + generator = torch.Generator().manual_seed(seed) + for depth in depths: + for layer in range(cache.num_layers): + shape = (length, cache.num_kv_heads, cache.head_dim) + k = torch.randn(shape, generator=generator) * (1 + layer) + v = torch.randn(shape, generator=generator) + cache.write(layer, [rid] * length, [depth] * length, list(range(length)), k, v) + + +def test_quantize_round_trip_is_bounded_per_vector(): + x = torch.randn(64, 4, 32) * torch.logspace(-3, 3, 64).view(64, 1, 1) + codes, scale = quantize_fp8(x) + assert codes.dtype == torch.float8_e4m3fn and scale.shape == x.shape[:-1] + error = (dequantize_fp8(codes, scale, torch.float32) - x).abs() + # e4m3 keeps 3 mantissa bits: rounding error is at most 1/16 of the vector's amax. + assert (error <= x.abs().amax(-1, keepdim=True) / 16 + 1e-30).all() + zeros = quantize_fp8(torch.zeros(2, 8)) + assert torch.isfinite(zeros[1]).all() and not dequantize_fp8(*zeros, torch.float32).any() + + +def test_storage_layout_and_byte_accounting(): + cache = make_cache() + assert cache.key_cache.dtype == torch.float8_e4m3fn + assert cache.key_scale.shape == cache.key_cache.shape[:-1] + assert cache.bytes_per_block == kv_bytes_per_block(2, 4, 2, 8, 1, "fp8_e4m3") + assert cache.bytes_per_block == 2 * 2 * 4 * 2 * (8 + 4) + plain = make_cache("auto", dtype=torch.bfloat16) + assert plain.key_scale is None and plain.bytes_per_block == 2 * 2 * 4 * 2 * 8 * 2 + with pytest.raises(ValueError, match="kv_cache_dtype"): + make_cache("int4") + + +def test_attention_equals_attention_over_dequantized_values(): + cache = make_cache() + reference = make_cache("auto") + assert cache.allocate("r", 11) and reference.allocate("r", 11) + fill(cache, "r", 11) + for layer in range(cache.num_layers): + for depth in range(4): + k, v = cache.read(layer, "r", depth) + reference.write(layer, ["r"] * 11, [depth] * 11, list(range(11)), k, v) + q = torch.randn(3, 2, 8) + rows = (["r"] * 3, [0, 2, 3], [10, 4, 7]) + for layer in range(cache.num_layers): + got = cache.attend(layer, *rows, q) + want = reference.attend(layer, *rows, q) + torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-6) + + +def test_torch_kernel_scales_multiply_each_token_vector(): + codes, scale = quantize_fp8(torch.randn(5, 4, 2, 8)) # [block, token, head, dim] + tables = torch.tensor([[3, 1, 4]], dtype=torch.int32) + lengths = torch.tensor([10], dtype=torch.int32) + q = torch.randn(1, 2, 8) + got = torch_paged_attention(q, codes, codes, tables, lengths, scale, scale) + dense = dequantize_fp8(codes, scale, torch.float32) + want = torch_paged_attention(q, dense, dense, tables, lengths) + torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-6) + + +def test_finalize_copies_codes_and_scales_to_deeper_planes(): + cache = make_cache() + assert cache.allocate("r", 6) + fill(cache, "r", 5) + fill(cache, "r", 6, depths=[0], seed=1) # position 5 exists only at depth 0 + cache.finalize_token("r", 5, 0) + for layer in range(cache.num_layers): + exited = [t[5] for t in cache.read(layer, "r", 0)] + for depth in (1, 2, 3): + deeper = [t[5] for t in cache.read(layer, "r", depth)] + assert all(torch.equal(a, b) for a, b in zip(exited, deeper)) + + +def test_snapshot_restore_is_bitwise_and_checks_dtype(): + cache = make_cache() + assert cache.allocate("r", 7) + fill(cache, "r", 7) + before = [cache.read(layer, "r", depth) for layer in range(2) for depth in range(4)] + snapshot = cache.snapshot("r") + assert snapshot.keys.dtype == torch.float8_e4m3fn and snapshot.key_scales is not None + cache.free("r") + assert cache.allocate("r", 7) + cache.restore("r", snapshot) + after = [cache.read(layer, "r", depth) for layer in range(2) for depth in range(4)] + assert all(torch.equal(a, b) for x, y in zip(before, after) for a, b in zip(x, y)) + plain = make_cache("auto") + assert plain.allocate("r", 7) + with pytest.raises(ValueError, match="dtype"): + plain.restore("r", snapshot) + + +def test_backends_without_fp8_support_are_rejected(): + info = {"backend": "flash", "selected_backend": "flash"} + backend = AttentionBackend(torch_paged_attention, BackendCapabilities(), info) + with pytest.raises(ValueError, match="FP8"): + make_cache(attention=backend) + from vllm_rlt.pd.engine import PDEngine + + torch.manual_seed(0) + with pytest.raises(ValueError, match="FP8"): + PDEngine( + OuroForCausalLM(tiny_ouro_config()), + prefill_cache_config=CacheConfig(kv_cache_dtype="fp8_e4m3"), + attention_backend="triton", + ) + + +def run_engine(model, kv_cache_dtype, prompts, params, exit_config): + engine = LLMEngine( + model, + cache_config=CacheConfig(128, 4, kv_cache_dtype=kv_cache_dtype), + exit_config=exit_config, + ) + assert (engine.cache_manager.key_scale is not None) == (kv_cache_dtype == "fp8_e4m3") + for i, prompt in enumerate(prompts): + engine.add_request(str(i), prompt, params) + done = {} + while engine.has_unfinished_requests(): + for output in engine.step(): + if output.finished: + done[output.request_id] = output + assert engine.cache_manager.num_used_blocks == 0 + return done + + +def tiny_model(): + torch.manual_seed(5) + model = OuroForCausalLM(tiny_ouro_config()) + with torch.no_grad(): + model.lm_head.weight.mul_(8) # decisive logits: FP8 KV error stays below the margins + return model + + +def test_fixed_depth_engine_tokens_match_unquantized_kv(): + model = tiny_model() + prompts = [[2, 3, 4, 5], [7, 8], [9, 10, 11]] + params = SamplingParams(max_tokens=12, ignore_eos=True) + plain = run_engine(model, "auto", prompts, params, ExitConfig("ouro")) + fp8 = run_engine(model, "fp8_e4m3", prompts, params, ExitConfig("ouro")) + assert {k: o.token_ids for k, o in fp8.items()} == {k: o.token_ids for k, o in plain.items()} + + +def test_early_exit_trace_replays_through_fp8_finalize_copies(): + # A fixed exit trace isolates the LAST_EXITED copy path from gate sensitivity. + traces = {"0": [4, 1, 3, 2, 4, 1, 2, 3], "1": [4, 2, 2, 1, 1, 4, 3, 2]} + params = SamplingParams(max_tokens=8, ignore_eos=True, min_loops=1) + done = run_engine( + tiny_model(), + "fp8_e4m3", + [[2, 3, 4, 5], [7, 8]], + params, + ExitConfig("trace", depths_by_request=traces), + ) + assert {k: list(o.exit_depths) for k, o in done.items()} == traces + + +@pytest.mark.gpu +def test_triton_fp8_attention_matches_torch_reference(): + from vllm_rlt.kernels.triton_attention import paged_attention + + torch.manual_seed(3) + codes, scale = quantize_fp8(torch.randn(40, 16, 4, 128, device="cuda")) + tables = torch.randint(0, 40, (6, 9), device="cuda", dtype=torch.int32) + lengths = torch.tensor([1, 17, 50, 100, 144, 3], device="cuda", dtype=torch.int32) + q = torch.randn(6, 8, 128, device="cuda", dtype=torch.bfloat16) + got = paged_attention(q, codes, codes, tables, lengths, scale, scale) + want = torch_paged_attention(q, codes, codes, tables, lengths, scale, scale) + torch.testing.assert_close(got.float(), want.float(), rtol=2e-2, atol=2e-2) + + +@pytest.mark.gpu +@pytest.mark.parametrize("mode", ["async", "multi"]) +@pytest.mark.parametrize("graphs", [False, True]) +def test_fp8_async_and_graph_paths_match_sync_eager(mode, graphs): + # Covers the async finalize kernel (scale copies) and quantized writes inside graphs. + torch.manual_seed(321) + model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda", torch.bfloat16) + traces = {str(i): [4, 1, 3, 2, 4, 2, 1, 3, 2][: 5 + i] for i in range(3)} + + def run(execution): + engine = LLMEngine( + model, + cache_config=CacheConfig(num_blocks=64, block_size=16, kv_cache_dtype="fp8_e4m3"), + scheduler_config=SchedulerConfig( + max_num_seqs=3, max_num_batched_tokens=8, prefill_chunk_size=4 + ), + execution_config=execution, + exit_config=ExitConfig("trace", depths_by_request=traces), + attention_backend="triton", + ) + for i in range(3): + engine.add_request( + str(i), + [2 + i, 3, 4] * (i + 1), + SamplingParams(max_tokens=5 + i, min_loops=1, ignore_eos=True), + ) + done = {} + while engine.has_unfinished_requests(): + for output in engine.step(): + if output.finished: + done[output.request_id] = (output.token_ids, list(output.exit_depths)) + assert engine.cache_manager.num_used_blocks == 0 + return done + + reference = run(ExecutionConfig()) + candidate = run( + ExecutionConfig(async_scheduling=True, multi_stream=mode == "multi", cuda_graphs=graphs) + ) + assert candidate == reference diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 15bde763..75bd8099 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -18,6 +18,8 @@ class BackendCapabilities: packed_prefill: bool = False cuda_graphs: bool = False async_scheduling: bool = False + # Reads an FP8 e4m3 KV pool with per-vector float32 scales. + fp8_kv: bool = False @dataclass @@ -40,8 +42,8 @@ def generation(self): return getattr(self.kernel, "generation", None) -_TORCH_CAPABILITIES = BackendCapabilities() -_TRITON_CAPABILITIES = BackendCapabilities(cuda_graphs=True, async_scheduling=True) +_TORCH_CAPABILITIES = BackendCapabilities(fp8_kv=True) +_TRITON_CAPABILITIES = BackendCapabilities(cuda_graphs=True, async_scheduling=True, fp8_kv=True) _FLASH_CAPABILITIES = { 2: BackendCapabilities(cuda_graphs=True, async_scheduling=True), 3: BackendCapabilities(cuda_graphs=True, async_scheduling=True), @@ -103,7 +105,7 @@ def create_backend( implementation = _create_explicit(candidate, device, dtype, head_dim, block_size) missing = [ name - for name in ("packed_prefill", "cuda_graphs", "async_scheduling") + for name in ("packed_prefill", "cuda_graphs", "async_scheduling", "fp8_kv") if getattr(required_capabilities, name) and not getattr(implementation.capabilities, name) ] diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 2491dcf0..f212eb81 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -19,6 +19,9 @@ class CacheConfig: enable_prefix_caching: bool = False incremental_allocation: bool = False 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"): @@ -33,6 +36,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 diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index 3c4a99c5..611669c1 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/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 @@ -26,6 +30,31 @@ validate_backend_name, ) +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() + scale = (x.abs().amax(-1) / 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: @@ -74,6 +103,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 @dataclass(frozen=True, eq=False) @@ -144,6 +176,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), @@ -161,6 +194,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) @@ -181,9 +217,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))) @@ -214,13 +258,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: @@ -517,17 +561,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: @@ -544,10 +594,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: @@ -830,8 +885,7 @@ def _write_prepared( if isinstance(batch, _DeviceKVBatch): self._validate_tensor(k, len(batch.position_ids), "k") self._validate_tensor(v, len(batch.position_ids), "v") - self.key_cache[batch.write_blocks, layer, batch.write_offsets] = k - self.value_cache[batch.write_blocks, layer, batch.write_offsets] = v + self._store_layer(layer, batch.write_blocks, batch.write_offsets, k, v) return if not batch.writable: raise ValueError("a read-only prepared KV batch cannot be written") @@ -839,11 +893,30 @@ 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) -> None: + """Write one layer's K/V rows, quantizing them for an FP8 pool.""" + if self.key_scale is None: + self.key_cache[blocks, layer, offsets] = k + self.value_cache[blocks, layer, offsets] = v + 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: @@ -881,6 +954,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) @@ -902,6 +976,7 @@ def _attend_prepared( self.value_cache[:, layer], batch.block_tables, batch.context_lengths, + **self._layer_scales(layer), ) @torch.no_grad() @@ -918,18 +993,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() @@ -953,4 +1026,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), + ) diff --git a/vllm_rlt/core/memory.py b/vllm_rlt/core/memory.py index 52ba8a69..81e64064 100644 --- a/vllm_rlt/core/memory.py +++ b/vllm_rlt/core/memory.py @@ -4,7 +4,7 @@ import torch -from vllm_rlt.core.kv_cache_manager import KVCacheManager +from vllm_rlt.core.kv_cache_manager import KVCacheManager, kv_bytes_per_block from vllm_rlt.worker.buffers import execution_buffer_bytes @@ -21,13 +21,13 @@ def budget_blocks(available_bytes, bytes_per_block): def plan_cache(model, cache, scheduler, execution, backend): parameter = next(model.parameters()) config, device = model.config, parameter.device - per_block = ( - 2 - * config.num_hidden_layers - * cache.block_size - * config.num_key_value_heads - * config.head_dim - * parameter.element_size() + per_block = kv_bytes_per_block( + config.num_hidden_layers, + cache.block_size, + config.num_key_value_heads, + config.head_dim, + parameter.element_size(), + cache.kv_cache_dtype, ) if cache.num_blocks is not None: return cache.num_blocks, {"source": "blocks", "bytes_per_block": per_block} @@ -67,6 +67,7 @@ def plan_cache(model, cache, scheduler, execution, backend): dtype=parameter.dtype, backend=backend, recurrent_layers=getattr(model, "recurrent_kv_layers", None), + kv_cache_dtype=cache.kv_cache_dtype, ) ids, positions = [], [] for index, length in enumerate(lengths): diff --git a/vllm_rlt/engine/llm_engine.py b/vllm_rlt/engine/llm_engine.py index 17a22cad..4e63c647 100644 --- a/vllm_rlt/engine/llm_engine.py +++ b/vllm_rlt/engine/llm_engine.py @@ -117,6 +117,7 @@ def __init__( async_scheduling=( parameter.device.type == "cuda" and self.execution_config.async_scheduling ), + fp8_kv=cache_config.kv_cache_dtype == "fp8_e4m3", ), ) num_blocks, self.memory_plan = plan_cache( @@ -142,6 +143,7 @@ def __init__( incremental_allocation=cache_config.incremental_allocation, watermark_ratio=cache_config.watermark_ratio, recurrent_layers=getattr(model, "recurrent_kv_layers", None), + kv_cache_dtype=cache_config.kv_cache_dtype, ) self.scheduler = Scheduler(scheduler_config, self.cache_manager, speculative_config) # One RNG registry per engine, shared by both runners: the CODA and diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py index 9a7df07e..306438f2 100644 --- a/vllm_rlt/entrypoints/runtime_args.py +++ b/vllm_rlt/entrypoints/runtime_args.py @@ -25,6 +25,12 @@ def add_runtime_args(parser): parser.add_argument("--enable-preemption", action="store_true") parser.add_argument("--prefill-uva", action="store_true") parser.add_argument("--kv-layout", choices=["last_exited", "shared"], default="last_exited") + parser.add_argument( + "--kv-cache-dtype", + choices=["auto", "fp8_e4m3"], + default="auto", + help="fp8_e4m3 stores KV as FP8 with per-vector scales (Torch/Triton attention)", + ) parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) parser.add_argument("--kv-cache-memory-bytes", type=int) parser.add_argument("--memory-reserve-bytes", type=int, default=256 * 1024 * 1024) @@ -75,6 +81,7 @@ def runtime_configs(args): num_blocks=args.num_blocks, block_size=args.block_size, layout=args.kv_layout, + kv_cache_dtype=args.kv_cache_dtype, gpu_memory_utilization=args.gpu_memory_utilization, kv_cache_memory_bytes=args.kv_cache_memory_bytes, memory_reserve_bytes=args.memory_reserve_bytes, diff --git a/vllm_rlt/kernels/paged_attention.py b/vllm_rlt/kernels/paged_attention.py index 2d1f2428..d2119e61 100644 --- a/vllm_rlt/kernels/paged_attention.py +++ b/vllm_rlt/kernels/paged_attention.py @@ -15,8 +15,13 @@ def torch_paged_attention( value_cache: torch.Tensor, block_tables: torch.Tensor, context_lengths: torch.Tensor, + key_scale: torch.Tensor | None = None, + value_scale: torch.Tensor | None = None, ) -> torch.Tensor: - """Reference attention over [physical block, token, KV head, dimension].""" + """Reference attention over [physical block, token, KV head, dimension]. + + FP8 pools pass per-vector scales shaped [physical block, token, KV head]. + """ output = torch.empty_like(q) block_size = key_cache.shape[1] head_groups = q.shape[1] // key_cache.shape[2] @@ -28,8 +33,13 @@ def torch_paged_attention( token_positions = torch.arange(length, device=q.device) blocks = block_tables[row, token_positions // block_size] offsets = token_positions % block_size - keys = key_cache[blocks, offsets].repeat_interleave(head_groups, dim=1) - values = value_cache[blocks, offsets].repeat_interleave(head_groups, dim=1) + keys = key_cache[blocks, offsets].float() + values = value_cache[blocks, offsets].float() + if key_scale is not None: + keys = keys * key_scale[blocks, offsets].unsqueeze(-1) + values = values * value_scale[blocks, offsets].unsqueeze(-1) + keys = keys.repeat_interleave(head_groups, dim=1) + values = values.repeat_interleave(head_groups, dim=1) scores = torch.einsum("hd,thd->ht", q[row].float(), keys.float()) * scale probabilities = torch.softmax(scores, dim=-1) output[row] = torch.einsum("ht,thd->hd", probabilities, values.float()).to(q.dtype) @@ -42,6 +52,8 @@ def triton_paged_attention( value_cache: torch.Tensor, block_tables: torch.Tensor, context_lengths: torch.Tensor, + key_scale: torch.Tensor | None = None, + value_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Dispatch explicitly to Triton, with no silent reference fallback.""" if q.device.type != "cuda": @@ -51,4 +63,6 @@ def triton_paged_attention( # Keep Triton optional and avoid loading its runtime for CPU-only callers. from vllm_rlt.kernels.triton_attention import paged_attention - return paged_attention(q, key_cache, value_cache, block_tables, context_lengths) + return paged_attention( + q, key_cache, value_cache, block_tables, context_lengths, key_scale, value_scale + ) diff --git a/vllm_rlt/kernels/routing.py b/vllm_rlt/kernels/routing.py index 570382e5..5547bb6a 100644 --- a/vllm_rlt/kernels/routing.py +++ b/vllm_rlt/kernels/routing.py @@ -103,7 +103,9 @@ def finalize_kernel( offset = (x // CHANNELS) * LAYER_STRIDE + (pos % PAGE) * TOKEN_STRIDE + x % CHANNELS src = source.to(tl.int64) * BLOCK_STRIDE + offset dst = dest.to(tl.int64) * BLOCK_STRIDE + offset - k = tl.load(K + src, x < LAYERS * CHANNELS, 0) - v = tl.load(V + src, x < LAYERS * CHANNELS, 0) + # No fill value: masked lanes are never stored, and an integer fill cannot be + # cast to an FP8 pool's element type. + k = tl.load(K + src, x < LAYERS * CHANNELS) + v = tl.load(V + src, x < LAYERS * CHANNELS) tl.store(K + dst, k, x < LAYERS * CHANNELS) tl.store(V + dst, v, x < LAYERS * CHANNELS) diff --git a/vllm_rlt/kernels/triton_attention.py b/vllm_rlt/kernels/triton_attention.py index cf53e818..bd3c7e2d 100644 --- a/vllm_rlt/kernels/triton_attention.py +++ b/vllm_rlt/kernels/triton_attention.py @@ -2,7 +2,8 @@ One Triton program computes one query head. Unlike a fixed-depth batch, every query row receives its own physical block table. This is a correctness-first -kernel without a performance claim or autotuning. +kernel without a performance claim or autotuning. FP8 e4m3 pools are read with +one float32 scale per (block, token, KV head) vector. """ import torch @@ -18,6 +19,8 @@ def _paged_attention_kernel( TABLES, LENGTHS, OUT, + KS, + VS, q_batch_stride: tl.constexpr, q_head_stride: tl.constexpr, q_dim_stride: tl.constexpr, @@ -33,6 +36,10 @@ def _paged_attention_kernel( out_batch_stride: tl.constexpr, out_head_stride: tl.constexpr, out_dim_stride: tl.constexpr, + s_block_stride: tl.constexpr, + s_token_stride: tl.constexpr, + s_head_stride: tl.constexpr, + HAS_SCALE: tl.constexpr, HEAD_DIM: tl.constexpr, HEAD_GROUPS: tl.constexpr, PAGE_SIZE: tl.constexpr, @@ -64,29 +71,40 @@ def _paged_attention_kernel( # Physical IDs fit in int32, but ID * block_stride can exceed 2**31 # elements in a large KV pool. Promote BEFORE multiplying, for K and V. offsets = positions % PAGE_SIZE - keys = tl.load( + kv_mask = valid_tokens[:, None] & (dims[None, :] < HEAD_DIM) + key_pointers = ( K + blocks[:, None] * k_block_stride + offsets[:, None] * k_token_stride + kv_head * k_head_stride - + dims[None, :] * k_dim_stride, - mask=valid_tokens[:, None] & (dims[None, :] < HEAD_DIM), - other=0, - ).to(tl.float32) + + dims[None, :] * k_dim_stride + ) + if HAS_SCALE: + # An integer fill value cannot be cast to fp8: zero masked lanes after widening. + scale_offsets = ( + blocks * s_block_stride + offsets * s_token_stride + kv_head * s_head_stride + ) + keys = tl.where(kv_mask, tl.load(key_pointers, mask=kv_mask).to(tl.float32), 0.0) + keys = keys * tl.load(KS + scale_offsets, mask=valid_tokens, other=0.0)[:, None] + else: + keys = tl.load(key_pointers, mask=kv_mask, other=0).to(tl.float32) scores = tl.sum(keys * query[None, :], axis=1) * SCALE scores = tl.where(valid_tokens, scores, -float("inf")) next_maximum = tl.maximum(maximum, tl.max(scores, axis=0)) correction = tl.exp(maximum - next_maximum) weights = tl.exp(scores - next_maximum) - values = tl.load( + value_pointers = ( V + blocks[:, None] * v_block_stride + offsets[:, None] * v_token_stride + kv_head * v_head_stride - + dims[None, :] * v_dim_stride, - mask=valid_tokens[:, None] & (dims[None, :] < HEAD_DIM), - other=0, - ).to(tl.float32) + + dims[None, :] * v_dim_stride + ) + if HAS_SCALE: + values = tl.where(kv_mask, tl.load(value_pointers, mask=kv_mask).to(tl.float32), 0.0) + values = values * tl.load(VS + scale_offsets, mask=valid_tokens, other=0.0)[:, None] + else: + values = tl.load(value_pointers, mask=kv_mask, other=0).to(tl.float32) accumulator = accumulator * correction + tl.sum(weights[:, None] * values, axis=0) normalizer = normalizer * correction + tl.sum(weights, axis=0) maximum = next_maximum @@ -97,11 +115,18 @@ def _paged_attention_kernel( ) -def paged_attention(q, key_cache, value_cache, block_tables, context_lengths): +def paged_attention( + q, key_cache, value_cache, block_tables, context_lengths, key_scale=None, value_scale=None +): """Launch over [batch row, query head]; inputs are validated by the manager.""" output = torch.empty_like(q) if q.shape[0] == 0: return output + has_scale = key_scale is not None + # Unused pointer arguments when the pool is not quantized. + key_scale = key_scale if has_scale else context_lengths + value_scale = value_scale if has_scale else context_lengths + scale_strides = key_scale.stride() if has_scale else (0, 0, 0) _paged_attention_kernel[(q.shape[0], q.shape[1])]( q, key_cache, @@ -109,11 +134,15 @@ def paged_attention(q, key_cache, value_cache, block_tables, context_lengths): block_tables, context_lengths, output, + key_scale, + value_scale, *q.stride(), *key_cache.stride(), *value_cache.stride(), block_tables.stride(0), *output.stride(), + *scale_strides, + HAS_SCALE=has_scale, HEAD_DIM=q.shape[-1], HEAD_GROUPS=q.shape[1] // key_cache.shape[2], PAGE_SIZE=key_cache.shape[1], diff --git a/vllm_rlt/pd/engine.py b/vllm_rlt/pd/engine.py index 2951d589..f213cb71 100644 --- a/vllm_rlt/pd/engine.py +++ b/vllm_rlt/pd/engine.py @@ -101,6 +101,10 @@ def __init__( d_cache = decode_cache_config or CacheConfig() if (p_cache.layout, p_cache.block_size) != (d_cache.layout, d_cache.block_size): raise ValueError("P and D must use the same KV layout and block size") + if "fp8_e4m3" in (p_cache.kv_cache_dtype, d_cache.kv_cache_dtype): + raise ValueError( + "PD transfer does not yet carry FP8 KV scales; use kv_cache_dtype='auto'" + ) if attention_backend not in ("triton", "flash_attn", "flash_attn_4"): raise ValueError("PD supports CUDA Triton or paged FA4 attention") self.scheduling_policy = p_scheduler.policy diff --git a/vllm_rlt/pd/transport.py b/vllm_rlt/pd/transport.py index fca044a1..ba61b19f 100644 --- a/vllm_rlt/pd/transport.py +++ b/vllm_rlt/pd/transport.py @@ -72,6 +72,10 @@ class NixlConnector: def __init__(self, name, cache, hidden, config): from nixl._api import nixl_agent, nixl_agent_config + if getattr(cache, "key_scale", None) is not None: + raise ValueError( + "PD transfer does not yet carry FP8 KV scales; use kv_cache_dtype='auto'" + ) self.config = config self.agent = nixl_agent(name, nixl_agent_config(backends=[config.backend])) self.storage = (cache.key_cache, cache.value_cache, hidden) diff --git a/vllm_rlt/worker/async_state.py b/vllm_rlt/worker/async_state.py index 2b289d90..f3c3773c 100644 --- a/vllm_rlt/worker/async_state.py +++ b/vllm_rlt/worker/async_state.py @@ -120,7 +120,7 @@ def __init__(self, cache, config, scheduler, rows, *, use_uva=True): self.planes = config.total_ut_steps if cache.layout == "last_exited" else 1 self.hidden = torch.empty( (scheduler.max_num_seqs, config.hidden_size), - dtype=cache.key_cache.dtype, + dtype=cache.dtype, device=cache.device, ) self.tokens = torch.empty( diff --git a/vllm_rlt/worker/cuda_graph.py b/vllm_rlt/worker/cuda_graph.py index 12959cea..f0953907 100644 --- a/vllm_rlt/worker/cuda_graph.py +++ b/vllm_rlt/worker/cuda_graph.py @@ -27,9 +27,7 @@ def __init__(self, cache): self.cache = cache def _write_prepared(self, layer, batch, k, v): - cache = self.cache - cache.key_cache[batch.write_blocks, layer, batch.write_offsets] = k - cache.value_cache[batch.write_blocks, layer, batch.write_offsets] = v + self.cache._store_layer(layer, batch.write_blocks, batch.write_offsets, k, v) def _attend_prepared(self, layer, batch, q): cache = self.cache @@ -49,6 +47,7 @@ def _attend_prepared(self, layer, batch, q): cache.value_cache[:, layer], batch.block_tables, batch.context_lengths, + **cache._layer_scales(layer), ) diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index b6510533..77e79d06 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -641,6 +641,27 @@ def finalize_many(self, requests): *cache.key_cache.stride()[:3], 256, ) + if cache.key_scale is not None: + # FP8 pools: the same copy for the per-vector scales. + finalize_kernel[ + ( + len(copies), + cache.max_loops, + math.ceil(cache.num_layers * cache.num_kv_heads / 256), + ) + ]( + bank.descriptor, + self.async_state.tables, + cache.key_scale, + cache.value_scale, + self.async_state.width, + cache.max_loops, + cache.block_size, + cache.num_layers, + cache.num_kv_heads, + *cache.key_scale.stride()[:3], + 256, + ) bank.record_done() for request in copies: self.events[request.request_id] = bank.done From ee375cba87126c8088f2d25101f6414d806adb4e Mon Sep 17 00:00:00 2001 From: 0z5a Date: Thu, 15 Oct 2026 20:40:39 +0200 Subject: [PATCH 2/5] [Docs] FP8 KV cache usage, semantics and limits Co-Authored-By: Claude Opus 5.5 --- docs/cdb_runtime.md | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/docs/cdb_runtime.md b/docs/cdb_runtime.md index d421585b..d54533a6 100644 --- a/docs/cdb_runtime.md +++ b/docs/cdb_runtime.md @@ -17,6 +17,7 @@ Implemented 2026-09-15 against base commit `6e8abbc405e1`. CUDA graphs and Hugin | Admission fairness | `Scheduler._admit`: bounded scan and bounded bypass count | Short request admitted behind blocked long request; protection prevents further bypass | | Chunked prefill policy | Per-request chunk ceiling, queue rotation and bounded prefill batches before decode | Short prompt completes before long prefill; existing decode starvation and chunk consistency tests | | Exit-trace replay | `ExitConfig(mode="trace")` bypasses gate compute/readback | Variable-depth replay, gate hook rejects accidental execution, all replay variants match | +| FP8 KV cache | `CacheConfig(kv_cache_dtype="fp8_e4m3")`: e4m3 codes plus per-vector scales, Torch/Triton dequantize on read | `tests/test_fp8_kv_cache.py`: round-trip bound, attention vs dequantized reference, finalize/snapshot, async/graph equality on GPU | This implements runtime mechanisms. It does not reproduce the paper's complete accuracy ablations, workload distributions, roofline fit or published speedups. @@ -129,6 +130,26 @@ CUDA automatic sizing profiles a synthetic maximum-row core/coda pass. It reserv Admission scans up to `admission_scan_limit` waiting items and preserves the order of deferred requests. A later admission increments bypass counts of earlier blocked requests. Once a request reaches `max_admission_bypasses`, later admissions stop bypassing it until enough active work drains. Existing active requests can finish because lifetime KV is already reserved. Chunk sizes and prefill/decode interleaving are independent from admission and apply in both layouts. +## 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. +- Rejected: FlashAttention backends (the `fp8_kv` backend capability; `auto` selects Triton) + and PD transfer, which does not yet move the scales. +- The current token's K/V enter attention unquantized in its own loop, as before; only stored + history is FP8. 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. From 0e530b2f902af1a85cec78701e11e9f5047214f7 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Fri, 16 Oct 2026 10:56:20 +0200 Subject: [PATCH 3/5] [KV] Fused FP8 quantize-and-store kernel for Triton An FP8 layer write on Triton is one launch: per (row, KV head) it computes the vector's scale, quantizes, and scatters codes and scale into the pool, instead of about 20 small Torch kernels (quantize ops plus four index writes). Under CUDA graphs every layer of every loop pays these per decode step. quantize_fp8 now multiplies by the reciprocal of 448, as CUDA division by a scalar already does, so the Torch path on CPU and CUDA and the Triton kernel write the same bits; a GPU test checks codes and scales bitwise. The Torch backend keeps the Torch writer as the reference, and a GPU engine test checks Triton FP8 (fused writes, chunked prefill and decode, eager and CUDA graphs) token for token against the Torch FP8 engine. Co-Authored-By: Claude Opus 5.5 --- tests/test_fp8_kv_cache.py | 75 ++++++++++++++++++++++++ vllm_rlt/core/kv_cache_manager.py | 20 ++++++- vllm_rlt/kernels/triton_attention.py | 85 ++++++++++++++++++++++++++++ 3 files changed, 179 insertions(+), 1 deletion(-) diff --git a/tests/test_fp8_kv_cache.py b/tests/test_fp8_kv_cache.py index ae0a4c70..d3535569 100644 --- a/tests/test_fp8_kv_cache.py +++ b/tests/test_fp8_kv_cache.py @@ -7,6 +7,7 @@ from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig, SamplingParams, SchedulerConfig from vllm_rlt.attention import AttentionBackend, BackendCapabilities from vllm_rlt.core.kv_cache_manager import ( + FP8_E4M3_MAX, KVCacheManager, dequantize_fp8, kv_bytes_per_block, @@ -242,3 +243,77 @@ def run(execution): ExecutionConfig(async_scheduling=True, multi_stream=mode == "multi", cuda_graphs=graphs) ) assert candidate == reference + + +@pytest.mark.gpu +def test_triton_fp8_store_matches_quantize_fp8_bitwise(): + from vllm_rlt.kernels.triton_attention import store_fp8 + + torch.manual_seed(2) + pages, layers, block, kv_heads, head_dim, rows, layer = 12, 3, 8, 4, 128, 40, 1 + shape = (pages, layers, block, kv_heads, head_dim) + codes = [torch.zeros(shape, device="cuda", dtype=torch.float8_e4m3fn) for _ in range(2)] + scales = [torch.ones(shape[:-1], device="cuda") for _ in range(2)] + slots = torch.randperm(pages * block, device="cuda")[:rows] + blocks, offsets = slots // block, slots % block + magnitudes = torch.logspace(-30, 4, rows, device="cuda").view(rows, 1, 1) + k = (torch.randn(rows, kv_heads, head_dim, device="cuda") * magnitudes).to(torch.bfloat16) + k[3] = 0 # A zero vector keeps a finite scale and zero codes. + # V rows strided across heads, as a slice of a wider projection would be. + v = torch.randn(rows, kv_heads, 2 * head_dim, device="cuda", dtype=torch.bfloat16) + v = v[..., :head_dim] + store_fp8( + k, + v, + *(t[:, layer] for t in codes), + *(t[:, layer] for t in scales), + blocks, + offsets, + FP8_E4M3_MAX, + ) + untouched = torch.ones(shape[:-1], dtype=torch.bool, device="cuda") + untouched[blocks, layer, offsets] = False + for x, code, scale in zip((k, v), codes, scales): + want_codes, want_scale = quantize_fp8(x) + assert torch.equal(code[blocks, layer, offsets].float(), want_codes.float()) + assert torch.equal(scale[blocks, layer, offsets], want_scale) + assert not code.float()[untouched].any() and bool((scale[untouched] == 1).all()) + + +def run_cuda_engine(model, backend, execution, traces, prompts, kv_cache_dtype="fp8_e4m3"): + engine = LLMEngine( + model, + cache_config=CacheConfig(num_blocks=64, block_size=16, kv_cache_dtype=kv_cache_dtype), + scheduler_config=SchedulerConfig( + max_num_seqs=3, max_num_batched_tokens=64, prefill_chunk_size=32 + ), + execution_config=execution, + exit_config=ExitConfig("trace", depths_by_request=traces), + attention_backend=backend, + ) + for i, prompt in enumerate(prompts): + params = SamplingParams(max_tokens=len(traces[str(i)]), min_loops=1, ignore_eos=True) + engine.add_request(str(i), prompt, params) + done = {} + while engine.has_unfinished_requests(): + for output in engine.step(): + if output.finished: + done[output.request_id] = (output.token_ids, list(output.exit_depths)) + assert engine.cache_manager.num_used_blocks == 0 + return done + + +@pytest.mark.gpu +@pytest.mark.parametrize("graphs", [False, True]) +def test_triton_fp8_engine_matches_torch_fp8_engine(graphs): + # Fused quantized writes and Triton reads, through chunked prefill and decode, against + # the Torch backend's reference writes and reads; FP32 isolates the KV path. + torch.manual_seed(5) + model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda") + with torch.no_grad(): + model.lm_head.weight.mul_(8) + prompts = [[2 + (7 * i + n) % 60 for i in range(n)] for n in (5, 41, 100)] + traces = {str(i): [4, 1, 3, 2, 4, 2, 1, 3, 2, 4, 3, 1][: 8 + 2 * i] for i in range(3)} + reference = run_cuda_engine(model, "torch", ExecutionConfig(), traces, prompts) + execution = ExecutionConfig(async_scheduling=graphs, cuda_graphs=graphs) + assert run_cuda_engine(model, "triton", execution, traces, prompts) == reference diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index 611669c1..bfbe65c5 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -47,7 +47,9 @@ def kv_bytes_per_block( def quantize_fp8(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Per-vector symmetric e4m3: x [..., head_dim] -> (codes, float32 scale [...]).""" x = x.float() - scale = (x.abs().amax(-1) / FP8_E4M3_MAX).clamp_min(torch.finfo(torch.float32).tiny) + # 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 @@ -904,6 +906,22 @@ def _store_layer(self, layer, blocks, offsets, k, v) -> None: self.key_cache[blocks, layer, offsets] = k self.value_cache[blocks, layer, offsets] = v return + if self.backend == "triton": + # One fused launch instead of about 20 small kernels per layer write. + 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, + ) + return k_codes, k_scale = quantize_fp8(k) v_codes, v_scale = quantize_fp8(v) self.key_cache[blocks, layer, offsets] = k_codes diff --git a/vllm_rlt/kernels/triton_attention.py b/vllm_rlt/kernels/triton_attention.py index bd3c7e2d..e8dc86f2 100644 --- a/vllm_rlt/kernels/triton_attention.py +++ b/vllm_rlt/kernels/triton_attention.py @@ -152,3 +152,88 @@ def paged_attention( num_warps=4, ) return output + + +@triton.jit +def _fp8_store_kernel( + K, + V, + K_CODES, + V_CODES, + K_SCALES, + V_SCALES, + BLOCKS, + OFFSETS, + k_row_stride, + k_head_stride, + v_row_stride, + v_head_stride, + c_block_stride, + c_token_stride, + c_head_stride, + s_block_stride, + s_token_stride, + s_head_stride, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + FP8_MAX: tl.constexpr, + INV_FP8_MAX: tl.constexpr, + TINY: tl.constexpr, +): + row = tl.program_id(0) + head = tl.program_id(1) + dims = tl.arange(0, BLOCK_D) + valid = dims < HEAD_DIM + block = tl.load(BLOCKS + row).to(tl.int64) + offset = tl.load(OFFSETS + row).to(tl.int64) + codes_at = block * c_block_stride + offset * c_token_stride + head * c_head_stride + dims + scale_at = block * s_block_stride + offset * s_token_stride + head * s_head_stride + # The same float32 operations as kv_cache_manager.quantize_fp8, in the same order. + k = tl.load(K + row * k_row_stride + head * k_head_stride + dims, mask=valid, other=0.0) + k = k.to(tl.float32) + k_scale = tl.maximum(tl.max(tl.abs(k), axis=0) * INV_FP8_MAX, TINY) + k_codes = tl.minimum(tl.maximum(tl.div_rn(k, k_scale), -FP8_MAX), FP8_MAX) + tl.store(K_CODES + codes_at, k_codes.to(K_CODES.dtype.element_ty), mask=valid) + tl.store(K_SCALES + scale_at, k_scale) + v = tl.load(V + row * v_row_stride + head * v_head_stride + dims, mask=valid, other=0.0) + v = v.to(tl.float32) + v_scale = tl.maximum(tl.max(tl.abs(v), axis=0) * INV_FP8_MAX, TINY) + v_codes = tl.minimum(tl.maximum(tl.div_rn(v, v_scale), -FP8_MAX), FP8_MAX) + tl.store(V_CODES + codes_at, v_codes.to(V_CODES.dtype.element_ty), mask=valid) + tl.store(V_SCALES + scale_at, v_scale) + + +def store_fp8(k, v, key_codes, value_codes, key_scale, value_scale, blocks, offsets, fp8_max): + """Quantize rows k, v [row, KV head, dim] and scatter codes and scales in one launch. + + Row i goes to (blocks[i], offsets[i]) of one layer's code pools + [physical block, token, KV head, dim] and scale pools [physical block, token, KV head]. + """ + if k.shape[0] == 0: + return + k = k if k.stride(-1) == 1 else k.contiguous() + v = v if v.stride(-1) == 1 else v.contiguous() + if key_codes.stride() != value_codes.stride() or key_scale.stride() != value_scale.stride(): + raise ValueError("FP8 key and value pools must share a layout") + _fp8_store_kernel[(k.shape[0], k.shape[1])]( + k, + v, + key_codes, + value_codes, + key_scale, + value_scale, + blocks, + offsets, + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + *key_codes.stride()[:3], + *key_scale.stride(), + HEAD_DIM=k.shape[-1], + BLOCK_D=triton.next_power_of_2(k.shape[-1]), + FP8_MAX=fp8_max, + INV_FP8_MAX=1.0 / fp8_max, + TINY=torch.finfo(torch.float32).tiny, + num_warps=1, + ) From f36f2f95e6d63c4fd4fd1333c374c0f13675fbb1 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Fri, 16 Oct 2026 12:47:37 +0200 Subject: [PATCH 4/5] [Docs] FP8 KV: fused Triton writer, current token is read in FP8 On every backend each layer writes the current token's K/V before attending, so attention reads the current token in FP8 as well; the previous wording said it stayed unquantized. Also list the fused writer and the new GPU tests. Co-Authored-By: Claude Opus 5.5 --- docs/cdb_runtime.md | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/cdb_runtime.md b/docs/cdb_runtime.md index d54533a6..a2db5730 100644 --- a/docs/cdb_runtime.md +++ b/docs/cdb_runtime.md @@ -17,7 +17,7 @@ Implemented 2026-09-15 against base commit `6e8abbc405e1`. CUDA graphs and Hugin | Admission fairness | `Scheduler._admit`: bounded scan and bounded bypass count | Short request admitted behind blocked long request; protection prevents further bypass | | Chunked prefill policy | Per-request chunk ceiling, queue rotation and bounded prefill batches before decode | Short prompt completes before long prefill; existing decode starvation and chunk consistency tests | | Exit-trace replay | `ExitConfig(mode="trace")` bypasses gate compute/readback | Variable-depth replay, gate hook rejects accidental execution, all replay variants match | -| FP8 KV cache | `CacheConfig(kv_cache_dtype="fp8_e4m3")`: e4m3 codes plus per-vector scales, Torch/Triton dequantize on read | `tests/test_fp8_kv_cache.py`: round-trip bound, attention vs dequantized reference, finalize/snapshot, async/graph equality on GPU | +| FP8 KV cache | `CacheConfig(kv_cache_dtype="fp8_e4m3")`: e4m3 codes plus per-vector scales, Torch/Triton dequantize on read | `tests/test_fp8_kv_cache.py`: round-trip bound, attention vs dequantized reference, finalize/snapshot; on GPU fused writer bitwise, Triton engine vs Torch engine, async/graph equality | This implements runtime mechanisms. It does not reproduce the paper's complete accuracy ablations, workload distributions, roofline fit or published speedups. @@ -144,11 +144,13 @@ model dtype and leaves every existing path unchanged. 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 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. -- The current token's K/V enter attention unquantized in its own loop, as before; only stored - history is FP8. Accuracy should be checked per model and workload; adaptive exits can change - where a token exits because the gate sees slightly different attention. +- 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 From fbfeaad59e4be4edc9db865176ffce765eb656fe Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sun, 1 Nov 2026 23:32:13 +0100 Subject: [PATCH 5/5] Run the merged FP8 tests on GPUs with less shared memory The bucketed graph KV test from main fakes the cache without FP8 scales; give it the None scales KVCacheManager reports for a BF16 pool. The FP32 dense reference prefill needs about 112 KB of shared memory at head dim 128, above the limit of sm_86, sm_89 and sm_120; the FP32 case of the FP8 prefill comparison uses head dim 64. Signed-off-by: 0z5a --- tests/test_cuda_graph.py | 3 ++- tests/test_fp8_kv_cache.py | 5 ++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_cuda_graph.py b/tests/test_cuda_graph.py index 6066c0d4..22eaab7c 100644 --- a/tests/test_cuda_graph.py +++ b/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"), diff --git a/tests/test_fp8_kv_cache.py b/tests/test_fp8_kv_cache.py index a7e17467..1d104ce3 100644 --- a/tests/test_fp8_kv_cache.py +++ b/tests/test_fp8_kv_cache.py @@ -313,7 +313,10 @@ def test_triton_fp8_prefill_matches_dequantized_prefill(dtype): from vllm_rlt.kernels.triton_attention import paged_prefill_attention torch.manual_seed(5) - pages, block, kv_heads, heads, head_dim = 40, 16, 2, 8, 128 + # The FP32 dense reference kernel needs ~112 KB of shared memory at head + # dim 128, above the ~99 KB of sm_86/sm_89/sm_120; 64 fits everywhere. + head_dim = 128 if dtype is torch.bfloat16 else 64 + pages, block, kv_heads, heads = 40, 16, 2, 8 codes, scale = quantize_fp8(torch.randn(pages, block, kv_heads, head_dim, device="cuda")) # Block 0 is never mapped: its NaN codes (e4m3fn 0x7F) must not reach any output. codes[0].view(torch.uint8).fill_(0x7F)