Support fp8 quantization for KV cache in Inference perf models

[devalshahamd]

Summary

Adds first-class handling of KV-cache dtype to InferenceAttention so perf predictions are correct for the common BF16-Q / FP8-KV deployment, and adds a new perf model for aiter::paged_attention_v1.

Motivation

Until now InferenceAttention.bytes_func charged every read/write at a single bytes_per_element. So the bytes-moved estimate was ~2x too high for quantized KV cache. We also had no perf model at all for the leaf op aiter::paged_attention_v1.

Changes

1. InferenceAttention base — track dtype_KV separately from dtype_Q (attention_perf_model_extensions.py)

• Annotation parsing factored into _parse_chunk_stats(event) so subclasses can reuse it without paying for the base's input-dims layout.
• bytes_func(...) now takes an optional bytes_per_element_KV (defaults to bytes_per_element). The c_sq/g_sq tokens (current chunk just produced this iteration in Q-dtype) are charged at bpe_Q; the cached remainder c_sk - c_sq / g_sk - g_sq is charged at bpe_KV. The four read/write terms are kept as separate ctx_elems_q / ctx_elems_kv / gen_elems_q / gen_elems_kv intermediates. When dtype_KV == dtype_Q the formula reduces exactly to the previous one (verified numerically), so every existing subclass keeps its current predictions.
• dtype_KV is resolved with explicit priority: perf_meta.KCache_dtype → Input type[1] → dtype_Q (see "tree_perf change" below for why perf_meta is needed).

2. New subclass aiter_paged_attention_v1

• Reads Q/K/V from Input Dims[2..4] (positional signature is (out, workspace, query, key_cache, value_cache, …)).
• Pulls dtype_Q / dtype_KV directly from Input type[2..3] — no perf_meta lookup needed since the leaf op's K-cache tensor already carries the true FP8 dtype.
• Overrides get_compute_precision() to return dtype_KV. The kernel's MFMAs run on the KV-cache dtype, so the FP8 throughput peak — not the BF16 one — is the correct denominator for Pct Roofline.

3. Pseudo-op wiring

• pseudo_ops_perf_utils.py: "aiter::paged_attention_v1" → new class.
• generate_perf_report_pytorch_inference.py: added to apply_annotation name_filters so the leaf inherits the user-annotation chunk stats.
• extensions/__init__.py: re-exported.

4. New Trace2Tree extension paged_attn_perf_meta.py (and auto-registration in pseudo_ops_utils.py)

This is the part that makes the FP8 numbers actually flow through. Two parent-traversal helpers are added; both write to a brand-new top-level event["perf_meta"] dict so existing event["args"] consumers are untouched.

• mark_paged_attn_v1_parents: when the trace contains aiter::paged_attention_v1, walk up its ancestor chain and set perf_meta.exclude_perf_model = True on every cpu_op. Why: vllm::unified_attention_with_output is normally the parent and also has a perf model. Without this flag, collect_unified_perf_events stops at the parent and never reaches the more granular leaf, so we'd lose the per-kernel breakdown for paged decode. Marking the parent excluded lets the leaf win.
• mark_rocm_paged_attn_kvcache_dtype: when the trace contains _rocm_C::paged_attention, copy Input type[5..6] (K-cache, V-cache dtypes) onto every cpu_op ancestor as perf_meta.KCache_dtype / perf_meta.VCache_dtype. Why: _rocm_C::paged_attention itself has no perf model — its parent vllm::unified_attention_with_output does. But the parent op only sees a BF16 staging tensor at its Input type[1]; the true FP8 cache dtype is buried in the ROCm leaf. Propagating it up via perf_meta lets InferenceAttention.get_param_details pick it up and pass FP8 bytes/sec into the parent's perf model.

Both extensions are early-return no-ops on traces that don't contain the respective leaf ops, so existing inference traces are unaffected.

5. One-line guard in tree_perf.py::_has_perf_model

if event.get("perf_meta", {}).get("exclude_perf_model"):
    return False