| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import argparse |
| 4 | import math |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm import _custom_ops as ops |
| 9 | from vllm.triton_utils import triton |
| 10 | |
| 11 | # DeepSeek V3 MLA dimensions |
| 12 | NOPE_DIM = 512 |
| 13 | ROPE_DIM = 64 |
| 14 | HEAD_DIM = NOPE_DIM + ROPE_DIM # 576 BF16 output elements per token |
| 15 | ENTRY_BYTES = 656 # 512 FP8 + 16 scales + 128 BF16 RoPE |
| 16 | BLOCK_SIZE = 64 # tokens per physical cache block - get_supported_kernel_block_sizes |
| 17 | |
| 18 | # Realistic prefill scenarios: |
| 19 | # - 1 long prefill: single request, 16K-96K tokens |
| 20 | # - 4 medium prefills: 4 requests, 4K-24K tokens each |
| 21 | # - 16 shorter prefills: 16 requests, 1K-6K tokens each |
| 22 | SCENARIOS = [ |
| 23 | # (label, num_reqs, total_tokens_list) |
| 24 | ("1-req", 1, [8192, 16384, 32768, 65536, 98304]), |
| 25 | ("4-reqs", 4, [8192, 16384, 32768, 65536, 98304]), |
| 26 | ("16-reqs", 16, [8192, 16384, 32768, 65536, 98304]), |
| 27 | ] |
| 28 | |
| 29 | |
| 30 | def make_inputs(total_tokens, num_reqs, block_size): |
| 31 | """Create synthetic FP8 cache, block table, and output buffer. |
| 32 | |
| 33 | Fills the cache with random bytes (we only measure throughput, |
| 34 | not correctness). Block table maps each request to contiguous |
| 35 | physical blocks. |
| 36 | """ |
| 37 | # Divide tokens evenly across requests |
| 38 | base_len = total_tokens // num_reqs |
| 39 | remainder = total_tokens % num_reqs |
| 40 | seq_lens = [base_len + (1 if r < remainder else 0) for r in range(num_reqs)] |
| 41 | |
| 42 | # workspace_starts: cumulative sum of seq_lens |
| 43 | workspace_starts = [0] * num_reqs |
| 44 | for r in range(1, num_reqs): |
| 45 | workspace_starts[r] = workspace_starts[r - 1] + seq_lens[r - 1] |
| 46 | |
| 47 | # Physical blocks needed per request |
| 48 | blocks_per_req = [math.ceil(s / block_size) for s in seq_lens] |
| 49 | total_blocks = sum(blocks_per_req) |
| 50 | max_blocks = max(blocks_per_req) |
| 51 | |
| 52 | # Allocate cache with random data (content doesn't matter for perf) |
| 53 | cache = torch.randint( |
| 54 | 0, |
| 55 | 256, |
| 56 | (total_blocks, block_size, ENTRY_BYTES), |
| 57 | dtype=torch.uint8, |
| 58 | device="cuda", |
| 59 | ) |
| 60 | |
| 61 | # Block table: contiguous block assignments |
| 62 | block_table = torch.zeros(num_reqs, max_blocks, dtype=torch.int32, device="cuda") |
| 63 | block_idx = 0 |
| 64 | for r in range(num_reqs): |
| 65 | for b in range(blocks_per_req[r]): |
| 66 | block_table[r, b] = block_idx |
| 67 | block_idx += 1 |
| 68 | |
| 69 | # Output workspace |
| 70 | dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda") |
| 71 | |
| 72 | workspace_starts_t = torch.tensor( |
| 73 | workspace_starts, dtype=torch.int32, device="cuda" |
| 74 | ) |
| 75 | |
| 76 | return cache, dst, block_table, workspace_starts_t |
| 77 | |
| 78 | |
| 79 | def bench_scenario(label, num_reqs, total_tokens_list, save_path): |
| 80 | """Run benchmark for a specific (num_reqs, total_tokens) scenario.""" |
| 81 | |
| 82 | @triton.testing.perf_report( |
| 83 | triton.testing.Benchmark( |
| 84 | x_names=["total_tokens"], |
| 85 | x_vals=total_tokens_list, |
| 86 | line_arg="provider", |
| 87 | line_vals=["cuda_kernel"], |
| 88 | line_names=["cp_gather_fp8 (CUDA)"], |
| 89 | styles=[("green", "-")], |
| 90 | ylabel="Latency (us)", |
| 91 | plot_name=f"cp_gather_fp8-{label}-bs{BLOCK_SIZE}", |
| 92 | args={"num_reqs": num_reqs}, |
| 93 | ) |
| 94 | ) |
| 95 | def bench_fn(total_tokens, provider, num_reqs): |
| 96 | cache, dst, block_table, ws_starts = make_inputs( |
| 97 | total_tokens, num_reqs, BLOCK_SIZE |
| 98 | ) |
| 99 | |
| 100 | quantiles = [0.5, 0.2, 0.8] |
| 101 | |
| 102 | ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( |
| 103 | lambda: ops.cp_gather_and_upconvert_fp8_kv_cache( |
| 104 | cache, dst, block_table, ws_starts, num_reqs |
| 105 | ), |
| 106 | quantiles=quantiles, |
| 107 | rep=500, |
| 108 | ) |
| 109 | |
| 110 | return ms * 1000, max_ms * 1000, min_ms * 1000 # us |
| 111 | |
| 112 | seq_len_per_req = total_tokens_list[0] // num_reqs |
| 113 | seq_len_per_req_max = total_tokens_list[-1] // num_reqs |
| 114 | print( |
| 115 | f"\n--- {label}: {num_reqs} request(s), " |
| 116 | f"~{seq_len_per_req}-{seq_len_per_req_max} tokens/req ---" |
| 117 | ) |
| 118 | bench_fn.run(print_data=True, save_path=save_path) |
| 119 | |
| 120 | |
| 121 | if __name__ == "__main__": |
| 122 | parser = argparse.ArgumentParser( |
| 123 | description="Benchmark cp_gather_and_upconvert_fp8_kv_cache" |
| 124 | ) |
| 125 | parser.add_argument( |
| 126 | "--save-path", |
| 127 | type=str, |
| 128 | default=None, |
| 129 | help="Path to save benchmark results as CSV", |
| 130 | ) |
| 131 | args = parser.parse_args() |
| 132 | |
| 133 | # Print data volume info for bandwidth analysis |
| 134 | read_per_token = ENTRY_BYTES # 656 bytes from cache |
| 135 | write_per_token = HEAD_DIM * 2 # 576 * 2 = 1152 bytes to workspace |
| 136 | total_per_token = read_per_token + write_per_token # 1808 bytes |
| 137 | |
| 138 | print("\n" + "=" * 70) |
| 139 | print("CP_GATHER_AND_UPCONVERT_FP8_KV_CACHE BENCHMARKS") |
| 140 | print("=" * 70) |
| 141 | print(f"Cache entry: {ENTRY_BYTES} bytes (512 FP8 + 16 scales + 128 RoPE)") |
| 142 | print(f"Output row: {HEAD_DIM} BF16 = {HEAD_DIM * 2} bytes") |
| 143 | print(f"Per token: {total_per_token} bytes (read + write)") |
| 144 | print(f"Block size: {BLOCK_SIZE} tokens/block") |
| 145 | print("=" * 70) |
| 146 | |
| 147 | for label, num_reqs, total_tokens_list in SCENARIOS: |
| 148 | bench_scenario(label, num_reqs, total_tokens_list, args.save_path) |
| 149 | |
| 150 | print("\n" + "=" * 70) |
| 151 | print("Benchmarking complete!") |
| 152 | print("=" * 70) |