| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Benchmark script comparing torch.cat vs direct copy for k_nope/k_pe concatenation |
| 4 | in MLA (Multi-head Latent Attention) prefill. |
| 5 | |
| 6 | This validates that the optimization from commit 8d4142bd is beneficial across |
| 7 | various batch sizes, not just the originally tested batch size of 32768. |
| 8 | """ |
| 9 | |
| 10 | import time |
| 11 | from collections.abc import Callable |
| 12 | |
| 13 | import torch |
| 14 | |
| 15 | # DeepSeek-V3 MLA dimensions |
| 16 | NUM_HEADS = 128 |
| 17 | QK_NOPE_HEAD_DIM = 128 |
| 18 | PE_DIM = 64 |
| 19 | |
| 20 | |
| 21 | def cat_method(k_nope: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor: |
| 22 | """Original torch.cat approach with expand.""" |
| 23 | return torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1) |
| 24 | |
| 25 | |
| 26 | def direct_copy_method(k_nope: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor: |
| 27 | """Optimized direct copy approach (avoids expand + cat overhead).""" |
| 28 | k = torch.empty( |
| 29 | (*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]), |
| 30 | dtype=k_nope.dtype, |
| 31 | device=k_nope.device, |
| 32 | ) |
| 33 | k[..., : k_nope.shape[-1]] = k_nope |
| 34 | k[..., k_nope.shape[-1] :] = k_pe |
| 35 | return k |
| 36 | |
| 37 | |
| 38 | def benchmark_method( |
| 39 | method: Callable, |
| 40 | k_nope: torch.Tensor, |
| 41 | k_pe: torch.Tensor, |
| 42 | num_warmup: int = 10, |
| 43 | num_iters: int = 100, |
| 44 | ) -> float: |
| 45 | """Benchmark a concatenation method and return mean latency in ms.""" |
| 46 | # Warmup |
| 47 | for _ in range(num_warmup): |
| 48 | _ = method(k_nope, k_pe) |
| 49 | torch.accelerator.synchronize() |
| 50 | |
| 51 | # Benchmark |
| 52 | start = time.perf_counter() |
| 53 | for _ in range(num_iters): |
| 54 | _ = method(k_nope, k_pe) |
| 55 | torch.accelerator.synchronize() |
| 56 | end = time.perf_counter() |
| 57 | |
| 58 | return (end - start) / num_iters * 1000 # Convert to ms |
| 59 | |
| 60 | |
| 61 | @torch.inference_mode() |
| 62 | def run_benchmark(dtype: torch.dtype, dtype_name: str): |
| 63 | """Run benchmark for a specific dtype.""" |
| 64 | torch.set_default_device("cuda") |
| 65 | |
| 66 | # Batch sizes to test (powers of 2 from 32 to 65536) |
| 67 | batch_sizes = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536] |
| 68 | |
| 69 | print("=" * 80) |
| 70 | print("Benchmark: torch.cat vs direct copy for MLA k_nope/k_pe concatenation") |
| 71 | print("=" * 80) |
| 72 | print( |
| 73 | f"Tensor shapes: k_nope=[B, {NUM_HEADS}, {QK_NOPE_HEAD_DIM}], " |
| 74 | f"k_pe=[B, 1, {PE_DIM}]" |
| 75 | ) |
| 76 | print(f"dtype: {dtype_name}") |
| 77 | print() |
| 78 | print( |
| 79 | f"{'Batch Size':>12} | {'cat (ms)':>10} | {'direct (ms)':>12} | " |
| 80 | f"{'Speedup':>8} | {'Reduction':>10}" |
| 81 | ) |
| 82 | print("-" * 70) |
| 83 | |
| 84 | results = [] |
| 85 | for batch_size in batch_sizes: |
| 86 | # Create input tensors (generate in float32 then convert for FP8 compatibility) |
| 87 | k_nope = torch.randn( |
| 88 | batch_size, NUM_HEADS, QK_NOPE_HEAD_DIM, dtype=torch.float32, device="cuda" |
| 89 | ).to(dtype) |
| 90 | k_pe = torch.randn( |
| 91 | batch_size, 1, PE_DIM, dtype=torch.float32, device="cuda" |
| 92 | ).to(dtype) |
| 93 | |
| 94 | # Benchmark both methods |
| 95 | cat_time = benchmark_method(cat_method, k_nope, k_pe) |
| 96 | direct_time = benchmark_method(direct_copy_method, k_nope, k_pe) |
| 97 | |
| 98 | speedup = cat_time / direct_time |
| 99 | reduction = (1 - direct_time / cat_time) * 100 |
| 100 | |
| 101 | results.append((batch_size, cat_time, direct_time, speedup, reduction)) |
| 102 | |
| 103 | print( |
| 104 | f"{batch_size:>12} | {cat_time:>10.3f} | {direct_time:>12.3f} | " |
| 105 | f"{speedup:>7.2f}x | {reduction:>9.1f}%" |
| 106 | ) |
| 107 | |
| 108 | print("=" * 80) |
| 109 | |
| 110 | # Summary statistics |
| 111 | speedups = [r[3] for r in results] |
| 112 | print("\nSpeedup summary:") |
| 113 | print(f" Min: {min(speedups):.2f}x") |
| 114 | print(f" Max: {max(speedups):.2f}x") |
| 115 | print(f" Mean: {sum(speedups) / len(speedups):.2f}x") |
| 116 | |
| 117 | # Find crossover point |
| 118 | crossover_batch = None |
| 119 | for batch_size, _, _, speedup, _ in results: |
| 120 | if speedup >= 1.0: |
| 121 | crossover_batch = batch_size |
| 122 | break |
| 123 | |
| 124 | print("\nConclusion:") |
| 125 | if crossover_batch: |
| 126 | print(f" - Direct copy becomes beneficial at batch size >= {crossover_batch}") |
| 127 | # Filter for large batches (>= 512 which is typical for prefill) |
| 128 | large_batch_speedups = [r[3] for r in results if r[0] >= 512] |
| 129 | if large_batch_speedups: |
| 130 | avg_large = sum(large_batch_speedups) / len(large_batch_speedups) |
| 131 | print(f" - For batch sizes >= 512: avg speedup = {avg_large:.2f}x") |
| 132 | print(" - MLA prefill typically uses large batches, so optimization is effective") |
| 133 | |
| 134 | return results |
| 135 | |
| 136 | |
| 137 | @torch.inference_mode() |
| 138 | def main(): |
| 139 | # Test bfloat16 |
| 140 | print("\n") |
| 141 | run_benchmark(torch.bfloat16, "bfloat16") |
| 142 | |
| 143 | # Test float8_e4m3fn |
| 144 | print("\n") |
| 145 | run_benchmark(torch.float8_e4m3fn, "float8_e4m3fn") |
| 146 | |
| 147 | |
| 148 | if __name__ == "__main__": |
| 149 | main() |