Metadata-Version: 2.4
Name: llm-flops
Version: 0.0.1
Summary: Accurate FLOPs, memory, latency, and energy estimation for Large Language Models.
Author-email: Soumyajit Basak <soumyajitbasak96@gmail.com>
License: MIT
License-File: LICENSE
Keywords: energy,flops,inference,kv-cache,latency,llm,llm-flops,memory,profiling,roofline,transformer
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.9
Provides-Extra: all
Requires-Dist: huggingface-hub>=0.20; extra == 'all'
Requires-Dist: matplotlib>=3.5; extra == 'all'
Requires-Dist: pytest>=7.0; extra == 'all'
Requires-Dist: torch>=2.1; extra == 'all'
Requires-Dist: transformers>=4.40; extra == 'all'
Provides-Extra: dev
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Provides-Extra: hf
Requires-Dist: huggingface-hub>=0.20; extra == 'hf'
Provides-Extra: keras
Requires-Dist: tensorflow>=2.10; extra == 'keras'
Provides-Extra: torch
Requires-Dist: torch>=2.1; extra == 'torch'
Requires-Dist: transformers>=4.40; extra == 'torch'
Provides-Extra: viz
Requires-Dist: matplotlib>=3.5; extra == 'viz'
Description-Content-Type: text/markdown

# LLM-FLOPS

**Accurate FLOPs, memory, latency, and energy estimation for Large Language Models.**

---

## What is it?

LLM-FLOPS is an open-source Python library that tells you exactly what one pass of
a transformer LLM costs — **compute (FLOPs), memory, latency, and energy** — under
any inference setting, *without downloading a single model weight*.

It is **analytical**: everything is derived from a model's architecture
hyperparameters, so you can size a 70B model on a laptop, offline, in
milliseconds. It is **transformer-aware**: it understands self-attention, KV
caching, grouped-query attention, and gated MLPs — not just generic matmuls. And
it is **precision-aware**: fp32, tf32, fp16, bf16, fp8, int8, int4.

It exists because existing FLOP counters tend to support only CNNs, report a
single opaque total with no per-layer breakdown, ignore transformer-specific
structure (attention cost, KV cache), and say nothing about precision, memory,
latency, or energy. LLM-FLOPS targets exactly those gaps with one unified API for
comparing models.

```python
from llmflops import estimate

result = estimate(
    model="meta-llama/Llama-3-8B",
    batch_size=4,
    sequence_length=2048,
    precision="bf16",
    gpu="A100",
)
print(result)
```

```
==============================================================
 LLM-FLOPS estimate: meta-llama/Meta-Llama-3-8B
  batch=4  seq_len=2048  precision=bf16  gpu=A100-80GB-SXM
==============================================================
 Parameters        : 8.03B  (8,030,261,248)
--------------------------------------------------------------
 FLOPs (single forward pass over batch x seq)
   Forward         : 131.885 TFLOP
   Backward        : 263.770 TFLOP
     - attention   : 30.880 TFLOP
     - mlp         : 92.389 TFLOP
     - embedding   : 0 FLOP
     - layernorm   : 8.724 GFLOP
     - output head : 8.607 TFLOP
--------------------------------------------------------------
 Memory
   Total           : 17.268 GB
     - weights     : 16.061 GB
     - activations : 134.218 MB
     - kv cache    : 1.074 GB
   Fits on GPU     : yes
--------------------------------------------------------------
 Performance (roofline)
   Prefill latency : 845.42 ms  [compute-bound]
   Decode/token    : 10.50 ms  (381 tok/s)
   Forward energy  : 287.44 J (0.0798 Wh)
==============================================================
```

---

## Main Features

- **Transformer-aware FLOP accounting** — forward *and* backward FLOPs, broken
  down by component: embeddings, multi-head attention (Q/K/V/output projections,
  attention scores, softmax·V), feed-forward MLP (standard **and** gated
  SwiGLU/GeGLU), normalization, and the output head. GQA/MQA aware.
- **Full memory model** — exact parameter count, weight memory by precision,
  activation memory, and KV-cache size (which shrinks automatically under
  grouped-query attention).
- **Roofline latency & energy** — separate, physically-motivated models for the
  compute-bound **prefill** phase and the memory-bound **decode** phase, backed
  by a database of real GPUs (peak throughput per precision, bandwidth, TDP).
- **Precision support** — fp32, tf32, fp16, bf16, fp8, int8, int4 (with aliases).
- **30+ bundled models, exact** — decoder, encoder, and encoder-decoder
  architectures; parameter counts validated against public model cards; anything
  else resolves from the Hugging Face Hub or a custom config.
- **GroqCloud provider** — estimate GroqCloud model IDs fully offline, no HF fetch.
- **One API to compare** models across batch size, sequence length, precision,
  and GPU.
- **Encoder & encoder-decoder too** — not just decoders: BERT/RoBERTa
  (bidirectional, no KV cache) and T5 (with cross-attention), alongside the
  decoder LLMs. Exact parameter counts (BERT-base 109.5M, T5-base 222.9M).
- **Mixture-of-Experts** — models total vs. active parameters and routed
  (top-k) FLOPs, so MoE models are estimated correctly rather than overcounted.
  Validated against Mixtral-8x7B (46.7B total / 12.9B active).
- **Training *and* inference memory** — weights, gradients, optimizer state
  (Adam/SGD), activations, and KV cache (the classic ~16 bytes/param for Adam).
- **Detailed backward & training‑step analysis** — decomposes the backward pass
  (weight / input / bias gradients; attention dQ·dK·dV·softmax·RoPE; MLP; norm;
  embedding), adds optimizer / gradient‑clip / mixed‑precision / checkpoint‑recompute
  FLOPs, backward memory traffic, and a per‑phase **bottleneck analysis with
  recommendations** (`training_analysis()`).
- **Energy & cost** — Joules, Wh, kWh, CO₂ emissions (configurable grid
  intensity), and cloud $/run.
- **Reports & exports** — a standalone **HTML or PDF report** with all charts
  (`viz.report(...)`, rendered offline via matplotlib) plus JSON / CSV / HTML
  data export. No server, no notebook — just a file.
- **Multi-GPU parallelism** — tensor / pipeline / data / expert layouts with
  per-GPU memory sharding, all-reduce / pipeline communication cost, and
  total-GPU + fits-per-GPU reporting (`ParallelConfig`).
- **Serving throughput** — continuous-batching capacity, aggregate tokens/s,
  requests/s, and **cost per 1M tokens** (`throughput()`).
- **Calibration** — pin the roofline efficiency to *your* measured throughput so
  latency/energy match your hardware (`calibration.calibrate(...)`).
- **Modern attention** — sliding-window attention (Mistral, Gemma-2) caps the KV
  cache and bands attention FLOPs; DeepSeek **MLA** models the compressed KV cache.
- **Universal profiler** — optional `profile()` measures the real FLOPs / MACs /
  params of *any* model in **PyTorch or TensorFlow/Keras** (CNNs, ViTs, RNNs,
  custom), and can record real latency + peak GPU memory — matching every
  profiler-based tool, across both frameworks, so LLM-FLOPS isn't limited to LLMs.
- **Hybrid validation** — optional empirical check of the analytical FLOPs
  against a real PyTorch forward pass.
- **Zero required dependencies** — the analytical core is pure Python. CLI included.

---

## Installation

```bash
pip install llm-flops                 # analytical core, zero dependencies
pip install "llm-flops[hf]"           # + fetch model configs from the HF Hub
pip install "llm-flops[torch]"        # + empirical validation (torch, transformers)
pip install "llm-flops[all]"          # everything
```

From a local copy of the project folder (editable install, for development):

```bash
cd LLM-FLOPS                 # the folder containing pyproject.toml
pip install -e ".[dev]"      # installs the package + the `llmflops` CLI
```

Optional extras can be combined, e.g. `pip install "llm-flops[hf,viz]"`.

| Extra | Adds | Enables |
|---|---|---|
| _(none)_ | — | analytical core + CLI (pure Python) |
| `hf` | huggingface_hub | fetch any model's `config.json` from the HF Hub |
| `torch` | torch, transformers | `validate()` and `profile()` for PyTorch models |
| `keras` | tensorflow | `profile()` for TensorFlow / Keras models |
| `viz` | matplotlib | HTML / PDF reports and charts |
| `all` | all of the above | everything |

---

## Functionalities (with Python code)

### 1. Estimate a single model

```python
import llmflops as fc

result = fc.estimate(
    model="meta-llama/Llama-3-8B",
    batch_size=4,
    sequence_length=2048,
    precision="bf16",
    gpu="A100",
)

print(result)                 # pretty, human-readable report
print(result["parameters"])   # 8030261248
print(result["forward_flops"])
```

### 2. Work with the result

`estimate()` returns an `EstimationResult` that behaves like a normal `dict`
(so indexing and `json.dumps` both work) but prints as the report above. Values
use SI base units: **FLOPs are raw counts, memory is bytes, latency is seconds,
energy is joules.**

```python
r = fc.estimate("meta-llama/Llama-3-8B", gpu="A100")

r["memory"]        # total bytes
r["kv_cache"]      # bytes
r["latency"]       # seconds (prefill / time-to-first-token)
r["energy"]        # joules

r.core()           # just the 10 headline metrics as a dict
r.to_dict()        # full JSON-serializable dict (incl. nested "details")

import json
json.dumps(r.to_dict(), default=str)
```

### 3. Layer-wise FLOPs breakdown

```python
r = fc.estimate("meta-llama/Llama-3-8B", batch_size=1, sequence_length=4096, gpu=None)

fwd = r["details"]["flops"]["forward"]
fwd["attention"], fwd["mlp"], fwd["output_head"], fwd["total"]

# even finer: inside attention
detail = r["details"]["flops"]["forward_attention_detail"]
detail["qkv_projection"], detail["attention_scores"], detail["output_projection"]
```

### 4. Memory and KV cache

```python
r = fc.estimate("meta-llama/Llama-3-8B", batch_size=8, sequence_length=8192,
                precision="bf16", gpu="A100")

mem = r["details"]["memory"]
mem["weights"], mem["activations"], mem["kv_cache"], mem["total"]
mem["fits_on_gpu"]        # True / False
```

### 5. Compare precisions

```python
for prec in ("fp32", "bf16", "int8", "int4"):
    r = fc.estimate("meta-llama/Llama-3-8B", precision=prec, gpu="A100")
    print(prec, round(r["memory"] / 1e9, 2), "GB",
          round(r["latency"] * 1e3, 1), "ms")
```

### 6. Latency & energy (roofline, prefill vs decode)

```python
r = fc.estimate("meta-llama/Llama-3-8B", batch_size=1, sequence_length=2048,
                precision="bf16", gpu="H100", num_generated_tokens=256)

perf = r["details"]["performance"]
perf["prefill_latency_s"], perf["prefill_bound"]          # e.g. "compute"
perf["decode_latency_per_token_s"], perf["decode_tokens_per_second"]
perf["total_request_latency_s"]                            # prefill + generation
perf["energy_total_joules"], perf["energy_total_wh"]
```

### 7. Compare models — and sweep settings

`models` and any of `batch_size` / `sequence_length` / `precision` / `gpu` can be
a single value **or a list**. `compare()` estimates the full Cartesian product
(model-major) and prints one row per combination, adding a column for each
dimension that varies.

```python
from llmflops import compare

# sweep precision across three models
print(compare(
    ["gpt2", "meta-llama/Llama-3-8B", "mistralai/Mistral-7B-v0.1"],
    batch_size=1, sequence_length=4096, precision=["bf16", "tf32"], gpu="A100",
))

# sweep several dimensions at once
print(compare("meta-llama/Llama-3-8B",
              precision=["fp16", "int8", "int4"],
              gpu=["A100", "H100"], sequence_length=[2048, 8192]))
```

Example output (precision sweep):

```
model            precision  params  fwd FLOPs     memory      kv cache    latency    fits
---------------  ---------  ------  ------------  ----------  ----------  ---------  ----
gpt2             bf16       124M    1.643 TFLOP   412.457 MB  150.995 MB  10.53 ms   yes
gpt2             tf32       124M    1.643 TFLOP   824.915 MB  301.990 MB  21.07 ms   yes
Meta-Llama-3-8B  bf16       8.03B   70.384 TFLOP  16.665 GB   536.871 MB  451.18 ms  yes
Meta-Llama-3-8B  tf32       8.03B   70.384 TFLOP  33.329 GB   1.074 GB    902.35 ms  yes
Mistral-7B-v0.1  bf16       7.24B   67.154 TFLOP  15.087 GB   536.871 MB  430.47 ms  yes
Mistral-7B-v0.1  tf32       7.24B   67.154 TFLOP  30.175 GB   1.074 GB    860.95 ms  yes
```

(tf32 is stored in 4 bytes vs bf16's 2, so memory and KV cache double — the
sweep makes that trade-off visible at a glance.)

### 8. Which GPUs can serve a model?

```python
for gpu in ("A100-80GB-SXM", "H100-SXM", "H200-SXM", "RTX4090"):
    r = fc.estimate("meta-llama/Llama-2-70b-hf", precision="fp16", gpu=gpu)
    fits = r["details"]["memory"]["fits_on_gpu"]
    print(gpu, round(r["memory"] / 1e9, 1), "GB ->", "fits" if fits else "does NOT fit")
```

### 9. Custom models (no download)

```python
cfg = fc.ModelConfig(
    name="my-model", hidden_size=4096, num_layers=32, num_attention_heads=32,
    num_key_value_heads=8, intermediate_size=14336, vocab_size=128256,
    mlp_type="gated", norm_type="rmsnorm", positional="rope",
)
print(fc.estimate(cfg, gpu="H100", precision="fp8"))

# ...or register it so it resolves by name later
fc.register_model("my-org/my-model", cfg, alias="my-model")
fc.estimate("my-model", gpu="A100")
```

### 10. Custom GPU

```python
gpu = fc.GPUSpec(
    name="MI300X", memory_gb=192, memory_bandwidth_gbps=5300, tdp_watts=750,
    peak_tflops={"bf16": 1307, "fp16": 1307, "fp8": 2615},
)
print(fc.estimate("meta-llama/Llama-3-8B", gpu=gpu, precision="bf16"))
```

### 11. Hugging Face Hub models

Any model not bundled is fetched from the Hub automatically (its `config.json`):

```python
fc.estimate("tiiuae/falcon-7b", gpu="H100")
fc.estimate("deepseek-ai/deepseek-llm-7b-base", gpu="A100")

fc.list_models()      # bundled model names
fc.get_config("meta-llama/Llama-3-8B")   # inspect the resolved ModelConfig
```

### 12. GroqCloud models (offline, no HF)

```python
from llmflops import groq

groq.register()                          # register Groq IDs into the registry
fc.estimate("llama-3.3-70b-versatile", gpu=None, allow_hf_fetch=False)

# convenience wrapper: registers on demand and attaches Groq's published throughput
groq.estimate("llama-3.1-8b-instant", num_generated_tokens=512)
groq.list_models()
```

### 13. Empirical validation (optional, needs torch)

```python
cfg = fc.ModelConfig(name="tiny", hidden_size=256, num_layers=4,
                     num_attention_heads=8, num_key_value_heads=4,
                     vocab_size=1024, intermediate_size=688, model_type="llama")
print(fc.validate(cfg, batch_size=1, sequence_length=64))
# {'measured_forward_flops': ..., 'analytical_forward_flops': ..., 'ratio': ~1.0, ...}
```

### 14. Command line

Run these in a **terminal / shell** (macOS/Linux terminal, or Anaconda Prompt /
Command Prompt on Windows) — *not* inside a Python, Jupyter, or Spyder console:

```bash
llmflops estimate meta-llama/Llama-3-8B -b 4 -s 2048 -p bf16 -g A100
llmflops compare gpt2 meta-llama/Llama-3-8B -p bf16 -g A100
llmflops estimate meta-llama/Llama-3-8B --json      # machine-readable
llmflops list-models
llmflops list-gpus
```

If you are inside a Python/IPython/Jupyter/Spyder console, either prefix the
command with `!` (e.g. `!llmflops list-models`) or — better — use the Python
API shown above (`fc.estimate(...)`, `fc.compare(...)`, `fc.list_models()`).
Pasting a bare `llmflops ...` line into a Python console raises
`SyntaxError: invalid decimal literal`, because it is a shell command, not Python.

### 15. Mixture-of-Experts models

MoE models report both the **total** parameter count (what you must store) and
the **active** count (what runs per token); FLOPs and decode latency use the
active path, memory uses the total.

```python
r = fc.estimate("mixtral-8x7b", batch_size=1, sequence_length=4096,
                precision="bf16", gpu="H100")
r["parameters"]         # 46_702_792_704  (total, drives memory)
r["active_parameters"]  # 12_879_925_248  (top-2 of 8 experts, drives FLOPs/latency)

# custom MoE (e.g. shared experts + some leading dense layers, DeepSeek-style)
cfg = fc.ModelConfig(name="my-moe", hidden_size=2048, num_layers=24,
                     num_attention_heads=16, vocab_size=32000, intermediate_size=8192,
                     num_experts=64, num_experts_per_tok=6, moe_intermediate_size=1408,
                     num_shared_experts=2, dense_layers=1)
fc.estimate(cfg, gpu="H100")
```

### 16. Profile any model — PyTorch or Keras (universal)

Measure the real graph of an arbitrary model. PyTorch (`pip install
"llm-flops[torch]"`) and TensorFlow/Keras (`pip install "llm-flops[keras]"`) are
auto-detected:

```python
import torchvision.models as models
import llmflops as fc

out = fc.profile(models.resnet50(), input_shape=(1, 3, 224, 224))
out["flops"], out["macs"], out["parameters"], out["by_module"], out["framework"]

# also measure real latency + peak GPU memory (like torch-flops)
fc.profile(models.resnet50(), input_shape=(1, 3, 224, 224), device="cuda", measure=True)
#   -> adds "latency_ms" and "peak_memory_bytes"

# TensorFlow / Keras models are detected automatically (like keras-flops)
import tensorflow as tf
fc.profile(tf.keras.applications.ResNet50())
```

### 17. Training memory (weights + gradients + optimizer state)

```python
r = fc.estimate("meta-llama/Llama-3-8B", precision="bf16", gpu="H100",
                training=True, optimizer="adam")   # optimizer: adam|adamw|sgd|none
t = r["details"]["memory"]["training"]
t["weights"], t["gradients"], t["optimizer_state"], t["activations"]
t["bytes_per_parameter"]     # ~16 for mixed-precision Adam
```

### 18. Energy, CO₂, and cloud cost

```python
r = fc.estimate("meta-llama/Llama-3-8B", batch_size=8, sequence_length=2048,
                precision="bf16", gpu="A100", num_generated_tokens=256,
                carbon_intensity=380,     # gCO2/kWh for your region's grid
                cost_per_hour=1.80)       # $/GPU-hour (defaults to the GPU's)
p = r["details"]["performance"]
p["energy_total_kwh"], p["co2_grams"], p["cost_usd"]
```

### 19. Export reports (JSON / CSV / HTML)

```python
r = fc.estimate("meta-llama/Llama-3-8B", gpu="A100")
r.to_json("report.json")
r.to_csv("report.csv")
r.to_html("report.html")     # standalone, self-contained report
```

### 20. Visualization report (HTML or PDF)

Visualization is delivered as a **standalone file** (LLM-FLOPS is a library, not
an app). Charts are rendered offline with matplotlib — the HTML embeds them as
images (no CDN), and the PDF is multi-page.

The report reflects **the mode you selected** — inference vs. training,
precision, GPU, batch, sequence length. The cleanest way is to make your
selections in `estimate()` and then extract the report from the result:

```python
import llmflops as fc

# choose your mode / settings, then extract a matching report
r = fc.estimate("meta-llama/Llama-3-8B", gpu="H100", precision="bf16", training=True)
r.report("training_report.pdf")     # PDF (mode=training reflected throughout)
r.report("report.html")             # HTML page

# one-shot equivalents (format inferred from the extension)
fc.report("meta-llama/Llama-3-8B", "report.pdf",  gpu="A100", precision="int4")
fc.report("meta-llama/Llama-3-8B", "report.html", gpu="A100", training=False)

# or grab individual matplotlib figures to embed yourself
fc.viz.fig_layerwise_flops(r)
fc.viz.fig_precision_comparison("meta-llama/Llama-3-8B", gpu="A100")
```

The report includes the metrics table plus six charts: layer-wise FLOPs,
attention-vs-MLP, memory breakdown (weights/gradients/optimizer in training
mode; weights/activations/KV in inference mode), precision comparison,
latency-vs-batch, and sequence-length scaling. Install with
`pip install "llm-flops[viz]"`.

### 21. Multi-GPU parallelism

```python
from llmflops import ParallelConfig

# Llama-2-70B in fp16 (~138 GB) doesn't fit one A100 — shard it across 8
r = fc.estimate("meta-llama/Llama-2-70b-hf", precision="fp16", gpu="A100-80GB-SXM",
                parallelism=ParallelConfig(tensor_parallel=8, interconnect_gbps=600))
p = r["details"]["parallelism"]
p["total_gpus"], p["per_gpu_memory_bytes"], p["fits_per_gpu"]
p["prefill_latency_s"], p["tp_comm_s"]              # communication folded into latency

# combine TP + PP + DP (dict form works too)
fc.estimate("meta-llama/Llama-3-8B", gpu="H100",
            parallelism={"tensor_parallel": 2, "pipeline_parallel": 2, "data_parallel": 4})
```

### 22. Serving throughput (continuous batching, $/1M tokens)

```python
t = fc.throughput("meta-llama/Llama-3-8B", gpu="A100", precision="bf16",
                  sequence_length=2048, output_tokens=256)
t["max_concurrent_batch"]      # sequences that fit the KV budget
t["tokens_per_second"], t["requests_per_second"]
t["cost_per_1m_tokens_usd"]

# large models need parallelism to fit before they can serve
fc.throughput("meta-llama/Llama-2-70b-hf", gpu="A100", precision="fp16",
              parallelism=ParallelConfig(tensor_parallel=4))
```

### 23. Calibrate to your measured hardware

```python
# you measured, say, 115 tok/s single-stream decoding Llama-3-8B on your A100:
fc.calibration.calibrate("meta-llama/Llama-3-8B", "A100", "bf16",
                         decode_tokens_per_second=115)

# now every estimate for that (model, gpu, precision) matches your hardware
fc.estimate("meta-llama/Llama-3-8B", gpu="A100")   # decode latency/energy calibrated

# or set the efficiency factors directly; explicit estimate() args always win
fc.calibration.set_efficiency("meta-llama/Llama-3-8B", "A100", "bf16", mem_efficiency=0.92)
```

### 24. Detailed backward / training-step analysis

Decompose a full training step — the backward pass by gradient type, plus
optimizer, clipping, mixed-precision, and checkpoint-recompute FLOPs — with a
memory-traffic and bottleneck analysis:

```python
ta = fc.training_analysis("meta-llama/Llama-3-8B", batch_size=1, sequence_length=2048,
                          precision="bf16", gpu="H100", optimizer="adam")
print(ta)                       # readable report (shown below)

f = ta["flops"]
f["forward"], f["backward"]["total"], f["optimizer"], f["total_step"]
f["backward"]["weight_gradient"], f["backward"]["input_gradient"]
f["backward"]["attention"]      # dQ/dK/dV/scores/softmax/RoPE/projection grads
ta["memory_traffic"]            # read/write bytes + arithmetic intensity
ta["bottleneck"]                # per-phase compute/memory-bound + recommendations
```

```
============================================================
 Training-step analysis: meta-llama/Meta-Llama-3-8B
============================================================
 Forward           : 31.861 TFLOP
 Backward          : 63.714 TFLOP
   weight gradient : 30.739 TFLOP
   input gradient  : 30.739 TFLOP
   attention bwd   : 13.218 TFLOP  (dQ/dK/dV/softmax/RoPE)
   mlp backward    : 46.187 TFLOP
   norm backward   : 5.453 GFLOP
 Optimizer (adam)   : 144.545 GFLOP
 Gradient clipping : 32.121 GFLOP
 Mixed-precision   : 16.061 GFLOP
 Recompute (ckpt)  : 31.861 TFLOP
 TOTAL step        : 127.628 TFLOP
------------------------------------------------------------
 Backward memory   : read 17.134 GB, write 17.134 GB
 Arithmetic inten. : 1859.2 FLOPs/byte
------------------------------------------------------------
 Backward bottleneck
   weight gradient       compute-bound    62.16 ms
   input gradient        compute-bound    62.16 ms
   attention backward    compute-bound    26.73 ms
   optimizer              memory-bound    71.91 ms
   activation recompute  compute-bound    64.43 ms
   ✓ Use a fused optimizer (fused AdamW) to cut optimizer memory traffic
   ✓ Shard optimizer state with ZeRO (Stage 1–2)
============================================================
```

The same backward breakdown is attached to any `estimate(..., training=True)`
result under `result["details"]["backward"]`.

---

## API reference

### `estimate(model, ...) -> EstimationResult`

The single entry point. Only `model` is required; everything else has sensible
defaults.

| Parameter | Default | Meaning |
|---|---|---|
| `model` | — | Registry name / HF Hub id / short alias, a `ModelConfig`, or a config `dict`. |
| `batch_size` | `1` | Inference batch size. |
| `sequence_length` | `2048` | Prompt / context length (tokens). |
| `precision` | `"bf16"` | Compute dtype: `fp32, tf32, fp16, bf16, fp8, int8, int4` (aliases accepted). |
| `gpu` | `"A100"` | GPU name / alias / `GPUSpec`, or `None` to skip latency & energy. |
| `training` | `False` | Report training memory (weights + grads + optimizer + activations) instead of inference memory. |
| `optimizer` | `"adam"` | Training optimizer: `adam`, `adamw`, `sgd`, `none`. |
| `gradient_checkpointing` | `True` | Assume activation checkpointing in the training-memory estimate. |
| `num_generated_tokens` | `128` | Decode horizon for the full-request latency/energy. |
| `causal` | `False` | Count only causal (lower-triangular) attention. |
| `kv_cache_precision` | auto | KV-cache dtype (defaults to `precision`; fp16 when weights are int8/int4). |
| `flash_attention` | `True` | Whether the activation estimate assumes FlashAttention. |
| `activation_factor` | `2.0` | Multiplier for the (rough) inference activation estimate. |
| `compute_efficiency` | calib/`0.5` | Prefill MFU (fraction of peak FLOP/s). Falls back to a calibrated value if registered, else the default. |
| `mem_efficiency` | calib/`0.8` | Decode's achieved fraction of peak memory bandwidth (calibratable). |
| `power_utilization` | calib/`0.85` | Mean power draw as a fraction of TDP (for energy; calibratable). |
| `carbon_intensity` | `400.0` | Grid carbon intensity in gCO₂/kWh (for CO₂). |
| `cost_per_hour` | GPU's | Cloud USD/GPU-hour (defaults to the GPU's bundled rate). |
| `parallelism` | `None` | A `ParallelConfig` (or dict) for multi-GPU memory sharding + communication. |
| `allow_hf_fetch` | `True` | Allow fetching unknown models from the HF Hub. |

### Result keys

Top level (SI base units — FLOPs are raw counts, memory in **bytes**, latency in
**seconds**, energy in **joules**):

`parameters`, `active_parameters`, `forward_flops`, `backward_flops`,
`attention_flops`, `mlp_flops`, `embedding_flops`, `layernorm_flops`,
`output_head_flops`, `memory`, `kv_cache`, `latency`, `energy`, plus context
(`model`, `precision`, `gpu`, `batch_size`, `sequence_length`, `note`).

`result["details"]` holds the full breakdown: `config`, `parameters`, `flops`
(with `forward`, `forward_attention_detail`, `backward_total`), `memory`
(`weights`, `activations`, `kv_cache`, `fits_on_gpu`, and `training` when
`training=True`), `performance` (prefill/decode latency, `energy_total_kwh`,
`co2_grams`, `cost_usd`, …), and `settings` (the mode you selected).

### Functions & classes

| Name | Purpose |
|---|---|
| `estimate(...)` | Estimate one model → `EstimationResult`. |
| `compare([...], **kw)` | Estimate several models under identical settings; prints a table. |
| `throughput(model, gpu=..., ...)` | Serving throughput: max batch, tokens/s, requests/s, $/1M tokens. |
| `training_analysis(model, ...)` | Detailed backward + optimizer FLOPs, memory traffic, bottleneck analysis. |
| `report(model_or_result, path, **kw)` | Write an HTML/PDF visual report (also `result.report(path)`). |
| `profile(model, input_shape=...)` | Measure any PyTorch model's real FLOPs/MACs/params (needs torch). |
| `validate(config, ...)` | Check analytical FLOPs against a real forward pass (needs torch). |
| `list_models()`, `list_gpus()` | Names available in the bundled registries. |
| `get_config(name)` | Resolve a name to a `ModelConfig`. |
| `register_model(name, config, alias=...)` | Add a custom model to the registry. |
| `ModelConfig(...)` | Describe a custom architecture (incl. MoE, SWA, MLA). |
| `GPUSpec(...)` | Describe a custom accelerator. |
| `ParallelConfig(...)` | Tensor/pipeline/data/expert parallel layout. |
| `calibration.calibrate(...)` / `.set_efficiency(...)` | Pin roofline efficiency to measured throughput. |
| `groq.register()`, `groq.estimate(id, ...)` | GroqCloud model IDs, fully offline. |
| `EstimationResult` | Dict-like result: `.core()`, `.to_dict()`, `.to_json/csv/html()`, `.report()`. |

Utility exports: `list_gpus()`, `get_gpu`, `register_gpu` (add a custom GPU by
name), `fetch_hf_config`, `SUPPORTED_PRECISIONS`, `normalize_precision`, and the
raw `MODEL_REGISTRY` / `GPU_DATABASE` tables. `compare()` returns a
`ComparisonResult` (a list with a tabular `__str__`).

---

## Supported models, GPUs, and precisions

**Models (bundled, exact).** Decoder LLMs: GPT-2 (base/medium/large/xl), Llama-2
(7B/13B/70B), Llama-3 / 3.1 / 3.2 (1B/3B/8B/70B), Mistral-7B,
Mixtral-8x7B / 8x22B (MoE), Qwen2.5 (0.5B/7B/72B), Phi-2, Gemma (2B/7B),
Falcon-7B, GPT-NeoX-20B. Encoders: BERT (base/large), RoBERTa (base/large).
Encoder-decoders: T5 (small/base/large). Anything else resolves via the Hugging
Face Hub (encoder/enc-dec/MoE fields included), a custom `ModelConfig`, or
`register_model(...)`. For non-transformer models, use `profile()`.

**GPUs (bundled).** A100 (40/80GB, SXM/PCIe), H100 (SXM/PCIe), H200, V100, T4,
L4, L40S, A10G, A6000, RTX 4090, RTX 3090. Define a `GPUSpec` for anything else.

**Precisions.** fp32, tf32, fp16, bf16, fp8, int8, int4.

---

## How it works

**FLOPs.** A matmul `[m,k]·[k,n]` costs `2·m·k·n` FLOPs (the Kaplan/Chinchilla
convention), so results match published figures. Per layer LLM-FLOPS counts the
Q/K/V/output projections, the `O(s²)` attention terms, and the MLP matmuls (2 for
a standard MLP, 3 for a gated SwiGLU). This reproduces `forward ≈ 2·N·D` per
token; the backward pass is `≈ 2×` forward, giving the `6·N·D` training estimate.

**Memory.** The parameter count is exact. Runtime memory = weights
(`params × bytes`) + KV cache
(`2 · batch · seq · layers · kv_heads · head_dim · bytes`) + an activation
estimate. GQA shrinks the KV cache by `num_heads / num_kv_heads`.

**Latency & energy (roofline).** Prefill is compute-bound
(`FLOPs / (peak · MFU)`); decode is memory-bound (each step re-reads all weights
+ KV cache from HBM). Each phase takes `max(compute, memory)`; energy =
`time × TDP × utilization`. The efficiency knobs (`compute_efficiency`,
`mem_efficiency`, `power_utilization`) are exposed as parameters.

---

## Validation

Bundled parameter counts match public model cards **exactly**:

| Model | LLM-FLOPS | Published |
|---|---|---|
| GPT-2 | 124,439,808 | 124M |
| Llama-2-7B | 6,738,415,616 | 6.74B |
| Mistral-7B | 7,241,732,096 | 7.24B |
| Llama-3-8B | 8,030,261,248 | 8.03B |
| Qwen2.5-7B | 7,615,716,864 | 7.61B |
| Gemma-7B | 8,537,680,896 | 8.54B |
| Falcon-7B | 6,921,720,704 | 6.92B |
| Llama-2-70B | 68,976,648,192 | ~69B |
| Mixtral-8x7B | 46,702,792,704 (12.9B active) | 46.7B / 12.9B |

**Cross-tool accuracy.** LLM-FLOPS's *analytical* FLOP counts match what
profiler-based tools get by *tracing the real model*. `calflops` reports
Llama-2-7B at **1700.06 GFLOP** (forward, 1×128); LLM-FLOPS computes
**1700.78 GFLOP** — a 0.04% difference, with identical parameter and MAC counts
— but without torch and without downloading any weights.

---

## Comparison with related tools

The established tools are **runtime profilers**: you instantiate a model and
they trace a forward pass to report FLOPs (PyTorch: `calflops`, `ptflops`,
`fvcore`, `flopth`, `torch-flops`; TensorFlow: `keras-flops`). LLM-FLOPS does
that too — its `profile()` covers **both PyTorch and Keras** and can measure
real latency + GPU memory — but it is primarily an **analytical cost model** for
transformer LLMs that answers questions the profilers can't.

| Tool | Framework | FLOPs / params | Per‑layer | Measures time + GPU mem | Precision · KV · memory | Latency · energy · cost | MoE · SWA · MLA | Multi‑GPU · serving | Analytical (no model) |
|---|:--:|:--:|:--:|:--:|:--:|:--:|:--:|:--:|:--:|
| calflops | PyTorch | ✅ | ✅ | — | — | — | — | — | partial¹ |
| ptflops | PyTorch | ✅ | ✅ | — | — | — | — | — | — |
| fvcore | PyTorch | ✅ | ✅ | — | — | — | — | — | — |
| flopth | PyTorch | ✅ | ✅ | — | — | — | — | — | — |
| keras‑flops | TensorFlow | ✅ | — | — | — | — | — | — | — |
| torch‑flops | PyTorch | ✅ | ✅ | ✅ | — | — | — | — | — |
| **LLM-FLOPS** | **PyTorch + TF** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |

¹ `calflops` can compute a model's FLOPs from a Hugging Face name without
downloading weights (meta-device), but still requires torch and reports FLOPs
only. LLM-FLOPS's analytical core needs no framework at all.

Note on MoE: profilers count whichever experts a specific input happens to
route to (input-dependent); LLM-FLOPS reports both **total** and
expected-**active** parameters/FLOPs deterministically.

**When to use which.** Reach for a profiler to measure a specific model you
already have in memory (and LLM-FLOPS's `profile()` now does this for PyTorch
*and* Keras, with optional real latency/memory). Reach for LLM-FLOPS's
analytical API to *plan* — size memory, KV cache, latency, energy, cost,
multi-GPU sharding, and serving throughput across precisions, batch sizes,
sequence lengths, and GPUs, including for models you haven't downloaded.

## Accuracy & limitations

Estimates are analytical, first-order, and meant for planning and comparison, not
cycle-accurate profiling. Real throughput depends on kernel quality and achieved
MFU — tune the efficiency knobs, or use `calibration.calibrate(...)` to pin them
to your measured numbers. Activation memory (inference and training) is a rough
estimate. Mixture-of-Experts is modelled at the routing level (total vs. active
params, top-k FLOPs); sliding-window and MLA attention are modelled for KV-cache
size and attention cost. Multi-GPU parallelism and serving throughput use
first-order Megatron / roofline / continuous-batching models — good for sizing,
not a replacement for a real benchmark. Cloud $/hr and CO₂ intensity are
configurable ballparks. Speculative decoding and cross-node network topology are
not modelled. The universal `profile()` path requires torch. Parameter counts and
matmul FLOP counts are exact.

---

## License

MIT — see [LICENSE](LICENSE).
