Profile FA3 KV-cache updates separately

runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_flashattn.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 import argparse
 import json
+import statistics
 import subprocess
 import sys
 import types
@@ -37,6 +38,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--warmup-iters", type=int, default=3)
     parser.add_argument("--repeats", type=int, default=5)
     parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--profile-kv-update", action="store_true")
     return parser.parse_args()
 
 
@@ -68,6 +70,7 @@ def main() -> None:
     bench_dir = args.vllm_source / "benchmarks" / "attention_benchmarks"
     sys.path.insert(0, str(bench_dir))
     import runner  # type: ignore[import-not-found]  # noqa: PLC0415
+    from batch_spec import parse_batch_spec  # type: ignore[import-not-found]  # noqa: PLC0415
     from common import BenchmarkConfig  # type: ignore[import-not-found]  # noqa: PLC0415
     from vllm.config import (  # noqa: PLC0415
         CacheConfig,
@@ -78,7 +81,13 @@ def main() -> None:
         ParallelConfig,
         SchedulerConfig,
         VllmConfig,
+        set_current_vllm_config,
     )
+    from vllm.v1.attention.backends.utils import (  # noqa: PLC0415
+        get_kv_cache_layout,
+        set_kv_cache_layout,
+    )
+    from vllm.v1.kv_cache_interface import FullAttentionSpec  # noqa: PLC0415
     from vllm.v1.worker.workspace import init_workspace_manager  # noqa: PLC0415
 
     def create_vllm_config(config: BenchmarkConfig, max_num_blocks: int) -> VllmConfig:
@@ -138,6 +147,77 @@ def main() -> None:
     runner._create_vllm_config = create_vllm_config
     init_workspace_manager(args.device)
 
+    def profile_kv_cache_update(config: BenchmarkConfig) -> dict[str, float]:
+        device = torch.device(config.device)
+        requests = parse_batch_spec(config.batch_spec)
+        total_q = sum(request.q_len for request in requests)
+        max_kv = max(request.kv_len for request in requests)
+        max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size
+        max_num_blocks = len(requests) * max_blocks_per_request
+        vllm_config = create_vllm_config(config, max_num_blocks)
+        dtype = vllm_config.model_config.dtype
+        with set_current_vllm_config(vllm_config):
+            backend_config = runner._get_backend_config(config.backend)
+            backend_class, impl, layer = runner._create_backend_impl(
+                backend_config, config, device, dtype
+            )
+            required_layout = backend_class.get_required_kv_cache_layout()
+            if required_layout is not None:
+                set_kv_cache_layout(required_layout)
+                get_kv_cache_layout.cache_clear()
+            common_metadata = runner._build_common_attn_metadata(
+                [request.q_len for request in requests],
+                [request.kv_len for request in requests],
+                config.block_size,
+                device,
+            )
+            kv_cache_spec = FullAttentionSpec(
+                block_size=config.block_size,
+                num_kv_heads=config.num_kv_heads,
+                head_size=config.head_dim,
+                dtype=dtype,
+            )
+            layer._kv_cache_spec = kv_cache_spec
+            _, key_list, value_list = runner._create_input_tensors(
+                config, total_q, device, dtype
+            )
+            cache_list = runner._create_kv_cache(
+                config, max_num_blocks, backend_class, device, dtype
+            )
+            for _ in range(config.warmup_iters):
+                for layer_index in range(config.num_layers):
+                    impl.do_kv_cache_update(
+                        layer,
+                        key_list[layer_index],
+                        value_list[layer_index],
+                        cache_list[layer_index],
+                        common_metadata.slot_mapping,
+                    )
+            torch.accelerator.synchronize()
+            samples: list[float] = []
+            for _ in range(config.repeats):
+                start = torch.cuda.Event(enable_timing=True)
+                end = torch.cuda.Event(enable_timing=True)
+                start.record()
+                for layer_index in range(config.num_layers):
+                    impl.do_kv_cache_update(
+                        layer,
+                        key_list[layer_index],
+                        value_list[layer_index],
+                        cache_list[layer_index],
+                        common_metadata.slot_mapping,
+                    )
+                end.record()
+                torch.accelerator.synchronize()
+                samples.append(float(start.elapsed_time(end)) / config.num_layers)
+        return {
+            "min_ms": min(samples),
+            "max_ms": max(samples),
+            "mean_ms": statistics.fmean(samples),
+            "median_ms": statistics.median(samples),
+            "std_ms": statistics.pstdev(samples),
+        }
+
     rows: list[dict[str, object]] = []
     for tp in args.tp:
         for batch_spec in args.batch_specs:
@@ -160,6 +240,9 @@ def main() -> None:
             result = runner.run_attention_benchmark(config)
             row = result.to_dict()
             row["tensor_parallel_size"] = tp
+            row["attention_core_excludes_kv_cache_update"] = True
+            if args.profile_kv_update:
+                row["kv_cache_update_time"] = profile_kv_cache_update(config)
             rows.append(row)
             print(
                 json.dumps(
@@ -189,6 +272,7 @@ def main() -> None:
             "dtype": "bfloat16",
             "attention_backend": "FLASH_ATTN",
             "block_size": 16,
+            "profile_kv_update": args.profile_kv_update,
         },
         "rows": rows,
     }