gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Profile FA3 KV-cache updates separately

Browse Source
This commit is contained in:
Gahow Wang
2026-07-16 21:37:01 +08:00
parent a9aed87518
commit 9bc38d8851
5 changed files with 100 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
runs/frontier-qwen30-vllm020-profile-v1/jobs_flashattn_full.toml
View File

@@ -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"

2
runs/frontier-qwen30-vllm020-profile-v1/jobs_smoke.toml
View File

@@ -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"]

84
runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_flashattn.py
View File

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

5
runs/frontier-qwen30-vllm020-profile-v1/run_flashattn_full.sh
View File

@@ -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" \

5
runs/frontier-qwen30-vllm020-profile-v1/run_flashattn_smoke.sh
View File

@@ -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" \