Align Frontier piecewise graph profiles

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

@@ -39,6 +39,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--repeats", type=int, default=5)
     parser.add_argument("--device", default="cuda:0")
     parser.add_argument("--profile-kv-update", action="store_true")
+    parser.add_argument(
+        "--profile-method",
+        choices=("cuda_event", "record_function"),
+        default="cuda_event",
+    )
+    parser.add_argument("--frontier-source", type=Path)
     return parser.parse_args()
 
 
@@ -66,12 +72,17 @@ def main() -> None:
         raise SystemExit(f"expected vLLM source {VLLM_COMMIT}, got {source_head}")
     if not args.model.joinpath("config.json").is_file():
         raise SystemExit(f"missing model config: {args.model / 'config.json'}")
+    if args.profile_method == "record_function" and args.frontier_source is None:
+        raise SystemExit("--frontier-source is required for --profile-method record_function")
 
     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 common import (  # type: ignore[import-not-found]  # noqa: PLC0415
+        BenchmarkConfig,
+        BenchmarkResult,
+    )
     from vllm.config import (  # noqa: PLC0415
         CacheConfig,
         CompilationConfig,
@@ -90,6 +101,13 @@ def main() -> None:
     from vllm.v1.kv_cache_interface import FullAttentionSpec  # noqa: PLC0415
     from vllm.v1.worker.workspace import init_workspace_manager  # noqa: PLC0415
 
+    record_function_tracer = None
+    if args.profile_method == "record_function":
+        sys.path.insert(0, str(args.frontier_source.resolve()))
+        from frontier.profiling.utils.record_function_tracer import RecordFunctionTracer
+
+        record_function_tracer = RecordFunctionTracer
+
     def create_vllm_config(config: BenchmarkConfig, max_num_blocks: int) -> VllmConfig:
         model_config = ModelConfig(
             model=str(args.model),
@@ -146,6 +164,9 @@ def main() -> None:
 
     runner._create_vllm_config = create_vllm_config
     init_workspace_manager(args.device)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    if args.profile_method == "record_function":
+        (args.output.parent / "profiler_traces").mkdir(exist_ok=True)
 
     def profile_kv_cache_update(config: BenchmarkConfig) -> dict[str, float]:
         device = torch.device(config.device)
@@ -218,6 +239,137 @@ def main() -> None:
             "std_ms": statistics.pstdev(samples),
         }
 
+    def profile_kernel_only(
+        config: BenchmarkConfig,
+    ) -> tuple[BenchmarkResult, dict[str, float] | None]:
+        """Trace exactly one vLLM FA3 forward/KV-update per annotation.
+
+        `RecordFunctionTracer` is Frontier's actual KERNEL_ONLY collector: it
+        follows CUDA launch correlations and sums kernels under the annotation.
+        The profiling loop therefore contains no CUDA-event value relabeling.
+        """
+        device = torch.device(config.device)
+        torch.accelerator.set_device_index(device)
+        backend_config = runner._get_backend_config(config.backend)
+        requests = parse_batch_spec(config.batch_spec)
+        q_lens = [request.q_len for request in requests]
+        kv_lens = [request.kv_len for request in requests]
+        total_q = sum(q_lens)
+        max_kv = max(kv_lens)
+        max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size
+        max_num_blocks = len(requests) * max_blocks_per_request
+
+        with runner.log_warnings_and_errors_only():
+            vllm_config = create_vllm_config(config, max_num_blocks)
+            dtype = vllm_config.model_config.dtype
+            with set_current_vllm_config(vllm_config):
+                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(
+                    q_lens, kv_lens, 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,
+                )
+                builder = runner._create_metadata_builder(
+                    backend_class, kv_cache_spec, vllm_config, device, config.backend
+                )
+                attn_metadata = builder.build(
+                    common_prefix_len=0, common_attn_metadata=common_metadata
+                )
+                quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
+                    impl, "supports_quant_query_input", False
+                )
+                q_list, k_list, v_list = runner._create_input_tensors(
+                    config, total_q, device, dtype, quantize_query=quantize_query
+                )
+                cache_list = runner._create_kv_cache(
+                    config, max_num_blocks, backend_class, device, dtype
+                )
+                output = torch.empty(
+                    total_q,
+                    config.num_q_heads,
+                    config.head_dim,
+                    device=device,
+                    dtype=dtype,
+                )
+
+                def run_core() -> None:
+                    for layer_index in range(config.num_layers):
+                        impl.forward(
+                            layer,
+                            q_list[layer_index],
+                            k_list[layer_index],
+                            v_list[layer_index],
+                            cache_list[layer_index],
+                            attn_metadata,
+                            output=output,
+                        )
+
+                for _ in range(config.warmup_iters):
+                    run_core()
+                torch.accelerator.synchronize()
+                core_tracer = record_function_tracer(str(args.output.parent))
+                with core_tracer:
+                    for _ in range(config.repeats):
+                        with torch.profiler.record_function("vidur_attention_core"):
+                            run_core()
+                core_stats = core_tracer.get_operation_time_stats()
+                if "attention_core" not in core_stats:
+                    raise RuntimeError("missing KERNEL_ONLY FlashAttention core stats")
+                core = {
+                    name: float(value) / config.num_layers
+                    for name, value in core_stats["attention_core"].items()
+                }
+
+                kv_stats = None
+                if args.profile_kv_update:
+                    def run_kv_cache_update() -> None:
+                        for layer_index in range(config.num_layers):
+                            impl.do_kv_cache_update(
+                                layer,
+                                k_list[layer_index],
+                                v_list[layer_index],
+                                cache_list[layer_index],
+                                common_metadata.slot_mapping,
+                            )
+
+                    for _ in range(config.warmup_iters):
+                        run_kv_cache_update()
+                    torch.accelerator.synchronize()
+                    kv_tracer = record_function_tracer(str(args.output.parent))
+                    with kv_tracer:
+                        for _ in range(config.repeats):
+                            with torch.profiler.record_function("vidur_attn_kv_cache_save"):
+                                run_kv_cache_update()
+                    kv_time_stats = kv_tracer.get_operation_time_stats()
+                    if "attn_kv_cache_save" not in kv_time_stats:
+                        raise RuntimeError("missing KERNEL_ONLY KV-update stats")
+                    kv_stats = {
+                        f"{name}_ms": float(value) / config.num_layers
+                        for name, value in kv_time_stats["attn_kv_cache_save"].items()
+                    }
+
+        result = BenchmarkResult(
+            config=config,
+            mean_time=core["mean"] / 1000.0,
+            std_time=core["std"] / 1000.0,
+            min_time=core["min"] / 1000.0,
+            max_time=core["max"] / 1000.0,
+            throughput_tokens_per_sec=(
+                total_q / (core["mean"] / 1000.0) if core["mean"] > 0 else 0.0
+            ),
+        )
+        return result, kv_stats
+
     rows: list[dict[str, object]] = []
     for tp in args.tp:
         for batch_spec in args.batch_specs:
@@ -237,12 +389,18 @@ def main() -> None:
                 kv_cache_dtype="auto",
                 use_cuda_graphs=False,
             )
-            result = runner.run_attention_benchmark(config)
+            if args.profile_method == "record_function":
+                result, kv_stats = profile_kernel_only(config)
+            else:
+                result = runner.run_attention_benchmark(config)
+                kv_stats = (
+                    profile_kv_cache_update(config) if args.profile_kv_update else None
+                )
             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)
+            if kv_stats is not None:
+                row["kv_cache_update_time"] = kv_stats
             rows.append(row)
             print(
                 json.dumps(
@@ -273,10 +431,10 @@ def main() -> None:
             "attention_backend": "FLASH_ATTN",
             "block_size": 16,
             "profile_kv_update": args.profile_kv_update,
+            "profile_method": args.profile_method,
         },
         "rows": rows,
     }
-    args.output.parent.mkdir(parents=True, exist_ok=True)
     args.output.write_text(
         json.dumps(payload, indent=2, sort_keys=True, default=json_default) + "\n"
     )