diff --git a/docs/runtime.md b/docs/runtime.md index 460c5e9c..b23451ce 100644 --- a/docs/runtime.md +++ b/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. 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 new file mode 100644 index 00000000..1d104ce3 --- /dev/null +++ b/tests/test_fp8_kv_cache.py @@ -0,0 +1,373 @@ +"""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 ( + FP8_E4M3_MAX, + 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 + + +@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()) + + +@pytest.mark.gpu +def test_triton_fp8_store_skips_masked_and_inactive_rows(): + # The device-masked writes used by speculative verification and bucketed CUDA graphs. + from vllm_rlt.kernels.triton_attention import store_fp8 + + torch.manual_seed(4) + pages, block, kv_heads, head_dim, rows, active = 6, 8, 2, 64, 10, 7 + shape = (pages, 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 + k = torch.randn(rows, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16) + v = torch.randn(rows, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16) + mask = torch.tensor([1, 0, 1, 1, 0, 1, 1, 1, 1, 1], device="cuda", dtype=torch.bool) + count = torch.tensor(active, device="cuda", dtype=torch.int32) + store_fp8(k, v, *codes, *scales, blocks, offsets, FP8_E4M3_MAX, mask=mask, active=count) + written = mask.clone() + written[active:] = False + for x, code, scale in zip((k, v), codes, scales): + want_codes, want_scale = quantize_fp8(x) + assert torch.equal(code[blocks, offsets][written].float(), want_codes[written].float()) + assert torch.equal(scale[blocks, offsets][written], want_scale[written]) + assert not code[blocks, offsets][~written].float().any() + assert bool((scale[blocks, offsets][~written] == 1).all()) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_triton_fp8_prefill_matches_dequantized_prefill(dtype): + from vllm_rlt.kernels.triton_attention import paged_prefill_attention + + torch.manual_seed(5) + # 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) + dense = dequantize_fp8(codes, scale, dtype) + tables = torch.randint(1, pages, (3, 12), device="cuda", dtype=torch.int32) + lengths = torch.tensor([70, 5, 150], device="cuda", dtype=torch.int32) + query_lengths = torch.tensor([70, 3, 64], device="cuda", dtype=torch.int32) + cu = torch.zeros(4, device="cuda", dtype=torch.int32) + cu[1:] = query_lengths.cumsum(0) + q = torch.randn(int(cu[-1]), heads, head_dim, device="cuda", dtype=dtype) + args = (tables, lengths, cu, int(query_lengths.max())) + got = paged_prefill_attention(q, codes, codes, *args, scale, scale) + want = paged_prefill_attention(q, dense, dense, *args) + assert torch.isfinite(got.float()).all() + torch.testing.assert_close(got.float(), want.float(), rtol=1e-2, atol=1e-2) + + +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/attention/backend.py b/vllm_rlt/attention/backend.py index 210e766e..20cb1454 100644 --- a/vllm_rlt/attention/backend.py +++ b/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) diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 648781d7..efb0e87b 100644 --- a/vllm_rlt/config.py +++ b/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 diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index dfcdf0df..dcaf539b 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 @@ -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,6 +846,7 @@ def _attend_prepared( batch.context_lengths, batch.cu_seqlens_q, batch.max_seqlen_q, + **self._layer_scales(layer), ) return self.attention( q, @@ -746,6 +854,7 @@ def _attend_prepared( 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), + ) diff --git a/vllm_rlt/core/memory.py b/vllm_rlt/core/memory.py index e9d2d4cf..b71c0098 100644 --- a/vllm_rlt/core/memory.py +++ b/vllm_rlt/core/memory.py @@ -5,7 +5,7 @@ import torch from vllm_rlt.config import cuda_graph_rows -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 from vllm_rlt.worker.cuda_graph import speculative_capture_sizes @@ -56,13 +56,13 @@ def plan_cache( ): 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} @@ -102,6 +102,7 @@ def plan_cache( 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/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..c354a47a 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -167,6 +167,7 @@ def __init__( parameter.device.type == "cuda" and self.execution_config.async_scheduling ), device_context_lengths=self.async_speculative, + fp8_kv=cache_config.kv_cache_dtype == "fp8_e4m3", ), ) num_blocks, self.memory_plan = plan_cache( @@ -194,6 +195,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, diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py index 883d884b..b16a37a1 100644 --- a/vllm_rlt/entrypoints/runtime_args.py +++ b/vllm_rlt/entrypoints/runtime_args.py @@ -36,6 +36,12 @@ def add_runtime_args(parser): help="Copy metadata to the GPU instead of reading mapped host memory", ) 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) @@ -106,6 +112,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 8dc005d2..239e923f 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,25 +63,54 @@ 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 + ) def triton_paged_prefill( - q, key_cache, value_cache, block_tables, context_lengths, cu_seqlens_q, max_seqlen_q + q, + key_cache, + value_cache, + block_tables, + context_lengths, + cu_seqlens_q, + max_seqlen_q, + key_scale=None, + value_scale=None, ): """Packed causal prefill with one block table per sequence (see triton_attention).""" from vllm_rlt.kernels.triton_attention import paged_prefill_attention return paged_prefill_attention( - q, key_cache, value_cache, block_tables, context_lengths, cu_seqlens_q, max_seqlen_q + q, + key_cache, + value_cache, + block_tables, + context_lengths, + cu_seqlens_q, + max_seqlen_q, + key_scale, + value_scale, ) class TritonPagedAttention: """Triton single-query decode attention plus packed causal prefill.""" - def __call__(self, q, key_cache, value_cache, block_tables, context_lengths): - return triton_paged_attention(q, key_cache, value_cache, block_tables, context_lengths) + def __call__( + self, + q, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale=None, + value_scale=None, + ): + return triton_paged_attention( + q, key_cache, value_cache, block_tables, context_lengths, key_scale, value_scale + ) def prefill(self, *args, **kwargs): return triton_paged_prefill(*args, **kwargs) diff --git a/vllm_rlt/kernels/routing.py b/vllm_rlt/kernels/routing.py index 71a77a1b..d80ee7e2 100644 --- a/vllm_rlt/kernels/routing.py +++ b/vllm_rlt/kernels/routing.py @@ -171,8 +171,10 @@ 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 f1953a92..e5c92996 100644 --- a/vllm_rlt/kernels/triton_attention.py +++ b/vllm_rlt/kernels/triton_attention.py @@ -6,6 +6,9 @@ merges the parts. A row's partition depends only on its own length, so its result does not depend on the batch it runs in. Without splitting, one long row serialized its whole context in one program and set the decode step time. + +FP8 e4m3 pools are read with one float32 scale per (block, token, KV head) +vector; both kernels dequantize on load. """ import torch @@ -27,6 +30,8 @@ def _paged_attention_kernel( PART_M, PART_L, PART_O, + KS, + VS, q_batch_stride: tl.constexpr, q_head_stride: tl.constexpr, q_dim_stride: tl.constexpr, @@ -39,6 +44,10 @@ def _paged_attention_kernel( v_head_stride: tl.constexpr, v_dim_stride: tl.constexpr, table_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, NUM_HEADS: tl.constexpr, @@ -77,29 +86,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 @@ -142,11 +162,18 @@ def _paged_attention_merge_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 split-KV parts over [row, head, split], then merge per [row, head].""" 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) rows, heads = q.shape[0], q.shape[1] block_d = triton.next_power_of_2(q.shape[-1]) part_m = torch.empty((rows, heads, NUM_SPLITS), device=q.device, dtype=torch.float32) @@ -161,10 +188,14 @@ def paged_attention(q, key_cache, value_cache, block_tables, context_lengths): part_m, part_l, part_o, + key_scale, + value_scale, *q.stride(), *key_cache.stride(), *value_cache.stride(), block_tables.stride(0), + *scale_strides, + HAS_SCALE=has_scale, HEAD_DIM=q.shape[-1], HEAD_GROUPS=heads // key_cache.shape[2], NUM_HEADS=heads, @@ -199,6 +230,8 @@ def _paged_prefill_kernel( LENGTHS, CU_Q, OUT, + KS, + VS, q_token_stride, q_head_stride, k_block_stride, @@ -210,6 +243,10 @@ def _paged_prefill_kernel( table_stride, out_token_stride, out_head_stride, + s_block_stride, + s_token_stride, + s_head_stride, + HAS_SCALE: tl.constexpr, HEAD_DIM: tl.constexpr, HEAD_GROUPS: tl.constexpr, PAGE_SIZE: tl.constexpr, @@ -256,15 +293,25 @@ def _paged_prefill_kernel( TABLES + seq * table_stride + keys_pos // PAGE_SIZE, mask=key_valid, other=0 ).to(tl.int64) offsets = keys_pos % PAGE_SIZE - keys = tl.load( + kv_mask = key_valid[:, None] & dim_valid[None, :] + key_pointers = ( K + blocks[:, None] * k_block_stride + offsets[:, None] * k_token_stride + kv_head * k_head_stride - + dims[None, :], - mask=key_valid[:, None] & dim_valid[None, :], - other=0.0, + + dims[None, :] ) + if HAS_SCALE: + # Dequantize to the query dtype for tl.dot. Masked lanes are zeroed + # after widening: an unfilled fp8 lane may hold a NaN code. + 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=key_valid, other=0.0)[:, None] + keys = keys.to(Q.dtype.element_ty) + else: + keys = tl.load(key_pointers, mask=kv_mask, other=0.0) scores = tl.dot(query, tl.trans(keys), input_precision=PRECISION) * SCALE visible = key_valid[None, :] & (keys_pos[None, :] <= positions[:, None]) scores = tl.where(visible, scores, -float("inf")) @@ -273,15 +320,19 @@ def _paged_prefill_kernel( safe = tl.where(next_maximum == -float("inf"), 0.0, next_maximum) correction = tl.exp(maximum - safe) weights = tl.exp(scores - safe[:, None]) - values = tl.load( + value_pointers = ( V + blocks[:, None] * v_block_stride + offsets[:, None] * v_token_stride + kv_head * v_head_stride - + dims[None, :], - mask=key_valid[:, None] & dim_valid[None, :], - other=0.0, + + dims[None, :] ) + 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=key_valid, other=0.0)[:, None] + values = values.to(Q.dtype.element_ty) + else: + values = tl.load(value_pointers, mask=kv_mask, other=0.0) accumulator = accumulator * correction[:, None] + tl.dot( weights.to(values.dtype), values, input_precision=PRECISION ) @@ -299,7 +350,15 @@ def _paged_prefill_kernel( def paged_prefill_attention( - q, key_cache, value_cache, block_tables, context_lengths, cu_seqlens_q, max_seqlen_q + q, + key_cache, + value_cache, + block_tables, + context_lengths, + cu_seqlens_q, + max_seqlen_q, + key_scale=None, + value_scale=None, ): """Packed causal prefill over paged K/V, bottom-right aligned per sequence. @@ -307,6 +366,7 @@ def paged_prefill_attention( cu_seqlens_q[s]:cu_seqlens_q[s + 1] and keys [0, context_lengths[s]) through block_tables[s]. Query i of a sequence sits at position context_length - query_length + i and attends keys up to that position. + FP8 pools pass per-vector scales [physical block, token, KV head]. """ output = torch.empty_like(q) sequences = block_tables.shape[0] @@ -319,6 +379,11 @@ def paged_prefill_attention( raise ValueError("paged prefill requires a contiguous head dimension") if q.stride(2) != 1: raise ValueError("paged prefill requires contiguous query rows") + 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_prefill_kernel[(sequences * q_blocks, q.shape[1])]( q, key_cache, @@ -327,6 +392,8 @@ def paged_prefill_attention( context_lengths, cu_seqlens_q, output, + key_scale, + value_scale, q.stride(0), q.stride(1), key_cache.stride(0), @@ -338,6 +405,8 @@ def paged_prefill_attention( block_tables.stride(0), output.stride(0), output.stride(1), + *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], @@ -352,3 +421,115 @@ def paged_prefill_attention( num_stages=2, ) return output + + +@triton.jit +def _fp8_store_kernel( + K, + V, + K_CODES, + V_CODES, + K_SCALES, + V_SCALES, + BLOCKS, + OFFSETS, + ACTIVE, + MASK, + 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) + # Optional device-side row skipping: rows at or past *ACTIVE (bucketed graph + # padding) and rows whose MASK entry is false write neither codes nor scales. + enabled = row >= 0 + if ACTIVE is not None: + enabled = enabled & (row < tl.load(ACTIVE)) + if MASK is not None: + enabled = enabled & tl.load(MASK + row) + if enabled: + 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, + *, + mask=None, + active=None, +): + """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]. + Optional ``mask`` (bool [row]) and ``active`` (int32 scalar row count) skip rows + on the device, as the unquantized masked and graph KV writers do. + """ + 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, + active, + mask, + 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, + ) diff --git a/vllm_rlt/pd/engine.py b/vllm_rlt/pd/engine.py index 47f58ad1..0a60fb08 100644 --- a/vllm_rlt/pd/engine.py +++ b/vllm_rlt/pd/engine.py @@ -92,6 +92,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 be0d423a..b69fd822 100644 --- a/vllm_rlt/worker/async_state.py +++ b/vllm_rlt/worker/async_state.py @@ -132,7 +132,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 10149516..ba03d99c 100644 --- a/vllm_rlt/worker/cuda_graph.py +++ b/vllm_rlt/worker/cuda_graph.py @@ -107,7 +107,18 @@ def __init__(self, cache): def _write_prepared(self, layer, batch, k, v): cache = self.cache - if batch.active_rows is not None: + if cache.key_scale is not None: + # FP8 pools: one fused quantizing writer that honors the same row masks. + cache._store_layer( + layer, + batch.write_blocks, + batch.write_offsets, + k, + v, + mask=batch.write_mask, + active=batch.active_rows, + ) + elif batch.active_rows is not None: from vllm_rlt.kernels.graph_kv import write_graph_kv write_graph_kv(cache.key_cache[:, layer], cache.value_cache[:, layer], batch, k, v) @@ -130,6 +141,7 @@ def _attend_prepared(self, layer, batch, q): batch.context_lengths, batch.cu_seqlens_q, batch.max_seqlen_q, + **cache._layer_scales(layer), ) return cache.attention( q, @@ -137,6 +149,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 ada23d20..c508581b 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -1338,6 +1338,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), + ) + ]( + routing_buffers.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, + ) routing_buffers.record_done() for request in copies: self.events[request.request_id] = routing_buffers.done