diff --git a/runs/frontier-qwen30-vllm020-profile-v1/jobs_flashattn_full.toml b/runs/frontier-qwen30-vllm020-profile-v1/jobs_flashattn_full.toml index c4bc708..5af0e6c 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/jobs_flashattn_full.toml +++ b/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" diff --git a/runs/frontier-qwen30-vllm020-profile-v1/jobs_smoke.toml b/runs/frontier-qwen30-vllm020-profile-v1/jobs_smoke.toml index 16be440..c811b28 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/jobs_smoke.toml +++ b/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"] diff --git a/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_flashattn.py b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_flashattn.py index 9a37c12..0e25d35 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_flashattn.py +++ b/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, } diff --git a/runs/frontier-qwen30-vllm020-profile-v1/run_flashattn_full.sh b/runs/frontier-qwen30-vllm020-profile-v1/run_flashattn_full.sh index 7188969..5d84abc 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/run_flashattn_full.sh +++ b/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" \ diff --git a/runs/frontier-qwen30-vllm020-profile-v1/run_flashattn_smoke.sh b/runs/frontier-qwen30-vllm020-profile-v1/run_flashattn_smoke.sh index 04318bb..327a261 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/run_flashattn_smoke.sh +++ b/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" \