Profile FA3 KV-cache updates separately

runs/frontier-qwen30-vllm020-profile-v1/jobs_flashattn_full.toml

@@ -1,55 +1,55 @@
 version = 1
 
 [[jobs]]
-name = "qwen30-vllm020-flashattn-full-tp1-20260716-v1"
+name = "qwen30-vllm020-flashattn-kv-full-tp1-20260716-v2"
 gpus = 1
 gpu_model = "H20"
 hosts = ["dash0"]
 command = "cd /home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-qwen30-vllm020-profile-v1 && timeout --signal=TERM --kill-after=30s 1320 bash run_flashattn_full.sh"
-artifacts = ["artifacts/flashattn-full-tp1"]
+artifacts = ["artifacts/flashattn-kv-full-v2-tp1"]
 
 [jobs.env]
 HOME = "/tmp/wjh"
 XDG_CACHE_HOME = "/tmp/wjh/.cache"
 VLLM_CACHE_ROOT = "/tmp/wjh/.cache/vllm"
 TP = "1"
-OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/flashattn-full-tp1"
+OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/flashattn-kv-full-v2-tp1"
 VENV_ROOT = "/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1"
 VLLM_SOURCE = "/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build"
 MODEL = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B"
 
 [[jobs]]
-name = "qwen30-vllm020-flashattn-full-tp2-20260716-v1"
+name = "qwen30-vllm020-flashattn-kv-full-tp2-20260716-v2"
 gpus = 1
 gpu_model = "H20"
 hosts = ["dash0"]
 command = "cd /home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-qwen30-vllm020-profile-v1 && timeout --signal=TERM --kill-after=30s 1320 bash run_flashattn_full.sh"
-artifacts = ["artifacts/flashattn-full-tp2"]
+artifacts = ["artifacts/flashattn-kv-full-v2-tp2"]
 
 [jobs.env]
 HOME = "/tmp/wjh"
 XDG_CACHE_HOME = "/tmp/wjh/.cache"
 VLLM_CACHE_ROOT = "/tmp/wjh/.cache/vllm"
 TP = "2"
-OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/flashattn-full-tp2"
+OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/flashattn-kv-full-v2-tp2"
 VENV_ROOT = "/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1"
 VLLM_SOURCE = "/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build"
 MODEL = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B"
 
 [[jobs]]
-name = "qwen30-vllm020-flashattn-full-tp4-20260716-v1"
+name = "qwen30-vllm020-flashattn-kv-full-tp4-20260716-v2"
 gpus = 1
 gpu_model = "H20"
 hosts = ["dash0"]
 command = "cd /home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-qwen30-vllm020-profile-v1 && timeout --signal=TERM --kill-after=30s 1320 bash run_flashattn_full.sh"
-artifacts = ["artifacts/flashattn-full-tp4"]
+artifacts = ["artifacts/flashattn-kv-full-v2-tp4"]
 
 [jobs.env]
 HOME = "/tmp/wjh"
 XDG_CACHE_HOME = "/tmp/wjh/.cache"
 VLLM_CACHE_ROOT = "/tmp/wjh/.cache/vllm"
 TP = "4"
-OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/flashattn-full-tp4"
+OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/flashattn-kv-full-v2-tp4"
 VENV_ROOT = "/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1"
 VLLM_SOURCE = "/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build"
 MODEL = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B"

runs/frontier-qwen30-vllm020-profile-v1/jobs_smoke.toml

@@ -1,7 +1,7 @@
 version = 1
 
 [[jobs]]
-name = "qwen30-vllm020-flashattn-smoke-20260716-v2-cu129"
+name = "qwen30-vllm020-flashattn-kv-smoke-20260716-v3-cu129"
 gpus = 1
 gpu_model = "H20"
 hosts = ["dash0"]

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,
     }

runs/frontier-qwen30-vllm020-profile-v1/run_flashattn_full.sh

@@ -51,7 +51,7 @@ if [[ "${#GPU_IDS[@]}" -ne 1 ]]; then
   exit 1
 fi
 
-echo "PROFILE_LAUNCH_ECHO host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES} model=${MODEL} runtime=vLLM-0.20.0+cu129 operator=FlashAttention3 tp=${TP} cases=${#BATCH_SPECS[@]} regimes=prefill,prefix_extend,decode,true_mixed dtype=BF16 block=16 output=${OUTPUT_ROOT} expected_wall=4-10m hard_wall=1200s hard_gpu_cap=0.33_H20h"
+echo "PROFILE_LAUNCH_ECHO host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES} model=${MODEL} runtime=vLLM-0.20.0+cu129 operators=FlashAttention3,reshape_and_cache_flash tp=${TP} cases=${#BATCH_SPECS[@]} regimes=prefill,prefix_extend,decode,true_mixed dtype=BF16 block=16 output=${OUTPUT_ROOT} expected_wall=4-10m hard_wall=1200s hard_gpu_cap=0.33_H20h"
 date -u +"START_UTC=%Y-%m-%dT%H:%M:%SZ"
 nvidia-smi --query-gpu=index,name,driver_version,memory.used,utilization.gpu --format=csv,noheader
 
@@ -78,7 +78,8 @@ timeout --signal=TERM --kill-after=30s 1080 \
   --tp "${TP}" \
   --batch-specs "${BATCH_SPECS[@]}" \
   --warmup-iters 5 \
-  --repeats 10
+  --repeats 10 \
+  --profile-kv-update
 
 test -s "${OUTPUT_ROOT}/raw/flashattn-tp${TP}.json"
 sha256sum "${OUTPUT_ROOT}/raw/flashattn-tp${TP}.json" \

runs/frontier-qwen30-vllm020-profile-v1/run_flashattn_smoke.sh

@@ -22,7 +22,7 @@ if [[ "${#GPU_IDS[@]}" -ne 1 ]]; then
   exit 1
 fi
 
-echo "PROFILE_LAUNCH_ECHO host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES} model=${MODEL} runtime=vLLM-0.20.0+cu129 operator=FlashAttention3 tp_local_shapes=1,2,4 specs=q128,4q1s128,q128_4q1s128 dtype=BF16 block=16 output=${OUTPUT_ROOT} expected_wall=3-8m hard_wall=900s hard_gpu_cap=0.15_H20h"
+echo "PROFILE_LAUNCH_ECHO host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES} model=${MODEL} runtime=vLLM-0.20.0+cu129 operators=FlashAttention3,reshape_and_cache_flash tp_local_shapes=1,2,4 specs=q128,4q1s128,q128_4q1s128 dtype=BF16 block=16 output=${OUTPUT_ROOT} expected_wall=3-8m hard_wall=900s hard_gpu_cap=0.15_H20h"
 date -u +"START_UTC=%Y-%m-%dT%H:%M:%SZ"
 nvidia-smi --query-gpu=index,name,driver_version,memory.used,utilization.gpu --format=csv,noheader
 
@@ -48,7 +48,8 @@ timeout --signal=TERM --kill-after=30s 780 \
   --tp 1 2 4 \
   --batch-specs q128 4q1s128 q128_4q1s128 \
   --warmup-iters 3 \
-  --repeats 5
+  --repeats 5 \
+  --profile-kv-update
 
 test -s "${OUTPUT_ROOT}/raw/flashattn-smoke.json"
 sha256sum "${OUTPUT_ROOT}/raw/flashattn-smoke.json" \