scripts/roofline_analysis.py

"""
Roofline / arithmetic-intensity analysis for decode-time attention.

Reconstructs (independently of the TileSparse paper's code, which is not public)
the mechanism behind Claim 3: that Multi-head Latent Attention (MLA) KV-cache
compression, combined with Multi-Token Prediction (MTP) batching multiple query
tokens per decode step, pushes attention arithmetic intensity (AI = FLOPs/byte)
across the GPU roofline ridge point, moving decode-time attention from the
memory-bound regime (classic single-token autoregressive decoding) into the
compute-bound regime.

Derivation (standard causal decode attention, KV cache of length L, query block
of D tokens computed per forward step, n_h heads, head dim d_h, element size
`elt_size` bytes):

  FLOPs   = 4 * D * L * n_h * d_h        (QK^T + softmax(QK^T)V, 2 flops/MAC each)
  Bytes   = 2 * L * n_h * d_h * elt_size  (reading K and V from cache once, reused
                                            across the D queries in the block)
  AI = FLOPs / Bytes = 2*D / elt_size

Notably L, n_h, d_h all cancel: decode-attention AI depends only on the query
block size D (the MTP degree) and element size. This is why MTP alone raises AI
linearly in D regardless of context length.

For MLA, the KV cache stores a compressed latent (kv_lora_rank d_c + decoupled
RoPE dim d_r) shared across all heads instead of full per-head K/V, so:

  Bytes_MLA = (d_c + d_r) * elt_size   (no n_h factor -- latent is shared)

We keep the FLOPs estimate for the "absorbed" MLA decode matmuls at the same
order as standard MHA (the absorption trick folds the up-projection into the
per-head query/output projections, so FLOPs are comparable while bytes shrink
drastically) -- this is the documented mechanism in the DeepSeek-V2 report and
is the honest, reproducible part of the analysis; we do not have access to the
TileSparse paper's own FLOP accounting since the OpenReview PDF is behind a
Cloudflare Turnstile challenge we could not solve programmatically, and no
arXiv preprint or GitHub repo is public for this paper (verified via HF papers
search, OpenReview API, and the ICML 2026 virtual site).
"""
import json
import math

ELT_SIZE_BYTES = 2  # bf16

# DeepSeek-V2 published MLA dimensions (public architecture, used as a
# representative real MLA config since TileSparse's own model config is not
# available to us).
N_HEADS = 128
HEAD_DIM = 128
KV_LORA_RANK = 512  # d_c
QK_ROPE_HEAD_DIM = 64  # d_r (shared decoupled RoPE dim)

# GPU roofline specs: (peak bf16 TFLOPS, peak HBM bandwidth GB/s)
GPUS = {
    "A10G (a10g-large)": (125.0, 600.0),
    "L4 (l4x1)": (121.0, 300.0),
    "L40S (l40sx1)": (362.0, 864.0),
    "A100-80GB (a100-large)": (312.0, 2039.0),
    "H100 (h200 class)": (989.0, 3350.0),
}


def ai_mha(d):
    return 2 * d / ELT_SIZE_BYTES


def ai_mla(d):
    mha_bytes_per_token = 2 * N_HEADS * HEAD_DIM * ELT_SIZE_BYTES
    mla_bytes_per_token = (KV_LORA_RANK + QK_ROPE_HEAD_DIM) * ELT_SIZE_BYTES
    compression_ratio = mha_bytes_per_token / mla_bytes_per_token
    # FLOPs held at the same order as MHA (absorption trick); bytes shrink by
    # compression_ratio -> AI scales up by that same ratio.
    return ai_mha(d) * compression_ratio, compression_ratio


def ridge_point(gpu_tflops, gpu_bw_gbs):
    return (gpu_tflops * 1e12) / (gpu_bw_gbs * 1e9)  # FLOPs/byte


def main():
    results = {"mha_bytes_per_token": 2 * N_HEADS * HEAD_DIM * ELT_SIZE_BYTES,
               "mla_bytes_per_token": (KV_LORA_RANK + QK_ROPE_HEAD_DIM) * ELT_SIZE_BYTES}
    _, compression_ratio = ai_mla(1)
    results["mla_kv_compression_ratio_x"] = compression_ratio

    print(f"MLA KV-cache compression vs full MHA: {compression_ratio:.1f}x fewer bytes/token")
    print(f"  MHA:  {results['mha_bytes_per_token']} bytes/token  (n_h={N_HEADS}, d_h={HEAD_DIM})")
    print(f"  MLA:  {results['mla_bytes_per_token']} bytes/token  (d_c={KV_LORA_RANK}, d_r={QK_ROPE_HEAD_DIM})")
    print()

    print(f"{'GPU':<24}{'ridge (FLOPs/B)':<18}{'D* for MHA':<14}{'D* for MLA':<14}")
    gpu_rows = []
    for name, (tflops, bw) in GPUS.items():
        rp = ridge_point(tflops, bw)
        d_star_mha = rp * ELT_SIZE_BYTES / 2  # solve ai_mha(D) = rp
        d_star_mla = d_star_mha / compression_ratio  # solve ai_mla(D) = rp
        print(f"{name:<24}{rp:<18.1f}{d_star_mha:<14.2f}{d_star_mla:<14.2f}")
        gpu_rows.append({
            "gpu": name, "peak_tflops_bf16": tflops, "peak_bw_gbs": bw,
            "ridge_flops_per_byte": rp,
            "D_star_plain_MHA_decode": d_star_mha,
            "D_star_MLA_decode": d_star_mla,
        })
    results["gpus"] = gpu_rows

    print()
    print("Interpretation: D* is the query-block size (MTP degree) at which decode")
    print("attention crosses from memory-bound (AI < ridge) to compute-bound (AI > ridge).")
    print("Plain single-token MHA decoding (D=1) is memory-bound on every GPU listed")
    print(f"(needs D >= {min(r['D_star_plain_MHA_decode'] for r in gpu_rows):.1f} to become compute-bound, i.e. far beyond typical MTP degrees of 1-4).")
    print(f"MLA cuts the required D* by {compression_ratio:.0f}x; on the fastest GPU (H100-class) MLA decode")
    print(f"crosses into compute-bound territory at D* ~= {min(r['D_star_MLA_decode'] for r in gpu_rows):.2f}, i.e. MTP degree as low as 1-2 tokens.")

    import os
    os.makedirs("outputs", exist_ok=True)
    with open("outputs/roofline_results.json", "w") as f:
        json.dump(results, f, indent=2)

    make_plot(gpu_rows, compression_ratio)


def make_plot(gpu_rows, compression_ratio):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    d_vals = [2 ** i for i in range(0, 8)]  # 1..128
    fig, ax = plt.subplots(figsize=(6.5, 5.8))
    ax.plot(d_vals, [ai_mha(d) for d in d_vals], "o-", label="Plain MHA decode AI(D)", linewidth=2, zorder=3)
    ax.plot(d_vals, [ai_mha(d) * compression_ratio for d in d_vals], "s-", label="MLA decode AI(D)", linewidth=2, zorder=3)

    # Only annotate the fastest/slowest GPU ridge (5 lines was illegible at
    # poster scale); the rest are in outputs/roofline_results.json.
    endpoints = {r["gpu"]: r for r in gpu_rows}
    for name in ("A100-80GB (a100-large)", "H100 (h200 class)"):
        row = endpoints[name]
        ax.axhline(row["ridge_flops_per_byte"], linestyle="--", color="gray", alpha=0.6, zorder=1)
        ax.annotate(f"{row['gpu'].split(' (')[0]} ridge", xy=(d_vals[0], row["ridge_flops_per_byte"]),
                    xytext=(2, 4), textcoords="offset points", fontsize=8, color="dimgray")

    ax.set_xscale("log", base=2)
    ax.set_yscale("log")
    ax.set_xlabel("Query block size D (MTP degree)")
    ax.set_ylabel("Arithmetic intensity (FLOPs/byte)")
    ax.set_title("Decode-attention AI vs MTP degree D")
    ax.legend(fontsize=9, loc="lower right", framealpha=0.9)
    fig.tight_layout()
    fig.savefig("outputs/roofline_plot.png", dpi=150)
    print("\nSaved outputs/roofline_plot.png and outputs/roofline_results.json")


if __name__ == "__main__":
    main()