🛠️ XProf CLI

CLI-first TPU/GPU profile analysis for AI agents and humans — analyze JAX / PyTorch-XLA / TensorFlow profiles via the open-source xprof profiler. Also ships an MCP server for assistants that prefer structured tools (Claude Code, Gemini, Cursor, etc.).

Renamed from xprof-mcp (2026-07). Same repo, full history preserved; GitHub redirects all old URLs, so existing clones, submodules, and MCP client configs keep working unchanged. The MCP server is retained as a fully supported frontend — but the CLI is the recommended interface: it needs no running server, no MCP wiring, picks up new profiles on every invocation, and works identically from any agent framework that can run a shell command.

Tools read .xplane.pb traces and XLA/LLO dump directories directly from disk (in-process conversion); a locally running xprof HTTP server is only needed for the interactive trace-viewer UI, not for analysis.

See also: TPU Performance Optimization Guide — practical guide covering the roofline model, common gotchas (dimension alignment, dtype, fusion failures, KV cache, rematerialization), training and inference optimization strategies, and decision trees for diagnosing bottlenecks.

---

What's here

One shared tool core (tool_registry.py, 38 tools) with two frontends:

Frontend	Entry point	Use it when
CLI (recommended)	xprof-cli <tool> --logdir=... [--run=...]	Any agent or human with a shell. No server, no wiring; every invocation sees the latest captures; per-tool --help from the same docstrings that drive the MCP schemas.
MCP server (preserved)	xprof-mcp / python -m xprof_mcp.server.xprof_mcp_server	Assistants that prefer structured tools (Claude Code agents, Gemini). Identical tool surface.

Functionality highlights (see the full tool table below):

• Whole-model analyses: overview, top HLO ops, op profile, memory profile, device info, consolidated get_kpi_metrics.
• Roofline (get_roofline_model): per-op compute- vs memory-bound classification with ridge points — caveat-annotated in the output (FLOPs/bytes are XLA cost-model estimates; custom calls are invisible to it; there is no communication-bound class; no HW counters on TPU v5p/v6e). Read the caveats field before acting on the numbers.
• Communication: get_pod_viewer (ICI/step breakdown), get_megascale_stats (DCN/multi-slice) — the comm attribution the roofline model structurally lacks.
• Memory attribution: get_memory_viewer — per-buffer HBM map (which tensor holds the peak), beyond get_memory_profile's scalar.
• Host/input pipeline: get_input_pipeline, plus get_framework_op_stats (device time by framework op name) and get_smart_suggestions (xprof's automated triage).
• HLO: module listing/content, instruction-neighborhood BFS, XLA dump reading + stage diffing (--xla_dump_to dirs, no trace needed).
• Pallas / Mosaic kernel internals (unique to this repo): LLO per-functional-unit utilization, Mosaic pipeline stage breakdown with DMA wait_ratio, VLIW bundle inspection, lowered-MLIR audit, and a capture-flag sanity check (check_kernel_profiling).
• In-process by default: tools read .xplane.pb and dump dirs directly and run the OSS xprof converters in this process (XPROF_MODE=local). The old server-backed path is kept as XPROF_MODE=http.
• CLI result cache: per-user SQLite cache (1h TTL) keyed on args and the profile directory's mtime — re-captured runs are never served stale; errors are never cached; --bypass_cache=True forces recompute.

Repo history

This repo was vlasenkoalexey/xprof-mcp until 2026-07: an MCP-only wrapper. It was renamed (not forked — full commit history, issues, and stars carried over) when the CLI became the primary interface; GitHub permanently redirects all old URLs, so existing clones, submodules, and MCP client configs continue to work without changes. The Python package directory is still xprof_mcp/ and the MCP entry point is unchanged — all MCP functionality is preserved; the CLI is simply the recommended way to consume the same tools.

---

Quick Start (CLI — recommended)

1. Generate a profile

With JAX (see JAX profiling guide):

import jax
import jax.numpy as jnp

# Collect a profile into /tmp/profiles/
with jax.profiler.trace("/tmp/profiles/", create_perfetto_link=False):
    y = jnp.dot(jnp.ones((1024, 1024)), jnp.ones((1024, 1024))).block_until_ready()

With PyTorch/XLA (see scaling-book profiling guide):

import torch_xla.debug.profiler as xp
server = xp.start_server(9012)
xp.trace('localhost:9012', '/tmp/profiles/', duration_ms=2000)

2. Install and run

cd /path/to/xprof-cli
pip install -e '.[full]'    # xprof converters + tensorflow-cpu

xprof-cli list_runs --logdir=/tmp/profiles
xprof-cli get_overview --logdir=/tmp/profiles --run=<run>
xprof-cli get_roofline_model --logdir=/tmp/profiles --run=<run>
xprof-cli get_llo_utilization --logdir=/tmp/profiles --run=<run> --kernel=<name>

xprof-cli                       # list all 38 commands
xprof-cli get_overview -- --help   # per-tool help

No server required. Set XPROF_LOGDIR instead of passing --logdir each time. Results print to stdout as JSON; failures exit non-zero with a JSON error on stderr — script/agent friendly.

3. (Optional) the interactive UI

The xprof web UI is the one thing that still wants the server:

pip install xprof
xprof --logdir=/tmp/profiles --port=8791   # browse http://localhost:8791

---

MCP server (preserved frontend)

The same 38 tools over the Model Context Protocol. Two modes; HTTP mode is recommended for active development — you can restart the MCP server without restarting your assistant.

Mode A: HTTP (recommended — restart-friendly)

The MCP server runs as a standalone HTTP process. Your AI assistant connects to it via URL. Restart or edit the server anytime without touching your assistant.

Start the server (run once; re-run after any code change):

PYTHONPATH=/path/to/xprof_mcp/.. \
XPROF_URL=http://localhost:8791 \
XPROF_LOGDIR=/tmp/profiles \
XLA_HLO_DUMP_DIR=/tmp/hlo_dumps \  # optional: enables HLO dump tools
MCP_PORT=8792 \
python -m xprof_mcp.server.xprof_mcp_server --transport http \
  > /tmp/xprof_mcp.log 2>&1 &

To restart after a code change:

kill $(pgrep -f "xprof_mcp_server --transport http") 2>/dev/null
# then re-run the start command above

Claude Code — run once:

claude mcp add --transport http --scope user xprof http://localhost:8792/mcp

No restart needed — Claude Code connects immediately.

Gemini CLI — edit ~/.gemini/settings.json:

{
  "mcpServers": {
    "xprof": {
      "url": "http://localhost:8792/sse"
    }
  }
}

Restart Gemini CLI once to pick up the config change. After that, MCP server restarts are transparent — no assistant restart needed.

---

Mode B: stdio (simpler setup, assistant manages the process)

The assistant spawns the MCP server as a subprocess. Requires an assistant restart whenever you update the MCP server code.

Claude Code — edit ~/.claude/.mcp.json:

{
  "xprof": {
    "command": "python",
    "args": ["-m", "xprof_mcp.server.xprof_mcp_server"],
    "env": {
      "PYTHONPATH": "/path/to/xprof_mcp/..",
      "XPROF_URL": "http://localhost:8791",
      "XPROF_LOGDIR": "/tmp/profiles",
      "XLA_HLO_DUMP_DIR": "/tmp/hlo_dumps"
    }
  }
}

Gemini CLI — edit ~/.gemini/settings.json:

{
  "mcpServers": {
    "xprof": {
      "command": "python",
      "args": ["-m", "xprof_mcp.server.xprof_mcp_server"],
      "env": {
        "PYTHONPATH": "/path/to/xprof_mcp/..",
        "XPROF_URL": "http://localhost:8791",
        "XPROF_LOGDIR": "/tmp/profiles",
        "XLA_HLO_DUMP_DIR": "/tmp/hlo_dumps"
      }
    }
  }
}

Tip: If python resolves to the wrong interpreter, use the full path (e.g. /usr/bin/python3 or the path to your venv's Python).

---

Directory Structure

xprof_mcp/                 # package dir keeps its pre-rename name on purpose
├── tool_registry.py       # SINGLE source of truth: the 38-tool surface
├── cli/
│   ├── main.py            # xprof-cli entry point (fire over the registry)
│   └── cache.py           # per-user SQLite result cache (mtime-salted TTL)
├── internal/
│   ├── xprof_client.py    # HTTP client (XPROF_MODE=http) + disk access + mode switch
│   ├── local_client.py    # in-process converters (XPROF_MODE=local, default)
│   ├── xprof_data.py      # get_profile_summary, get_op_profile, get_hosts, ...
│   ├── hlo_tools.py       # list_hlo_modules, get_hlo_module_content,
│   │                      #   get_hlo_neighborhood
│   ├── xplane_tools.py    # list_xplane_events, aggregate_xplane_events,
│   │                      #   get_xspace_proto  (require tensorflow)
│   ├── kernel_profiling_tools.py  # check_kernel_profiling, list_kernel_invocations,
│   │                      #   get_llo_utilization, get_kernel_stage_breakdown
│   ├── llo_dump_tools.py  # list_llo_programs, get_llo_schedule_analysis,
│   │                      #   get_llo_static_utilization, get_llo_bundles
│   └── mosaic_tools.py    # get_custom_call_mlir
├── tools/
│   ├── analysis_tools.py  # roofline / pod / megascale / memory_viewer /
│   │                      #   input_pipeline / framework_op_stats /
│   │                      #   smart_suggestions / kpi_metrics
│   ├── list_runs_tool.py           # list_runs
│   ├── get_overview_tool.py        # get_overview
│   ├── get_memory_profile_tool.py  # get_memory_profile
│   └── get_top_hlo_ops_tool.py     # get_top_hlo_ops
├── server/
│   └── xprof_mcp_server.py  # FastMCP entry point (iterates the registry)
└── tests/                 # pytest suite + real v6e trace/dump fixtures

---

Available Tools

Every tool is simultaneously a CLI subcommand and an MCP tool — one registry (tool_registry.py), two frontends.

Tool	Description	Needs TF?
list_runs	List profiling sessions in the logdir	No
get_hosts	List hosts in a run	No
get_overview	Step time, device utilization, run environment	No
get_memory_profile	Peak HBM usage, heap/stack breakdown	No
get_top_hlo_ops	Top ops by time, FLOPs, bytes accessed	No
get_profile_summary	Text summary of top ops	No
get_device_information	Accelerator specs from Roofline Model	No
get_kpi_metrics	Consolidated headline KPIs (step time, duty cycle, MXU, peak HBM)	No
get_roofline_model	Per-op compute/memory-bound classification + ridge points — caveat-annotated (cost-model FLOPs; custom-call-blind; no comm class)	No
get_pod_viewer	Pod-level step breakdown + ICI collective stats	No
get_megascale_stats	Multi-slice DCN collective stats (per-rendezvous)	No
get_memory_viewer	Per-buffer HBM attribution for one HLO module (which tensor holds peak)	No
get_input_pipeline	Host-vs-device input-pipeline stall decomposition	No
get_framework_op_stats	Device time by framework-level op name (JAX/PyTorch/TF)	No
get_utilization_viewer	Sampled utilization timeline (achieved vs peak over time)	No
get_perf_counters	Measured HW counters (TPU v7+/Ironwood; empty on v5p/v6e — noted in output)	No
get_smart_suggestions	xprof's automated bottleneck triage	No
detect_unfused_reshapes	Automated audit: standalone reshape/copy/transpose ops forcing HBM intermediates	No
list_hlo_modules	List compiled HLO programs in a run	No
get_hlo_module_content	Full HLO text for a module	No
get_hlo_neighborhood	BFS neighborhood of an HLO instruction	No
list_xplane_events	Filter timeline events by regex	Yes
aggregate_xplane_events	Stats (count/avg/stddev) per event type	Yes
get_xspace_proto	Raw XSpace proto bytes or text	Yes
list_hlo_dump_modules	List modules and stages in an XLA dump dir	No
get_hlo_dump	Read HLO text at a specific compilation stage	No
diff_hlo_stages	Unified diff between two compilation stages	No
get_hlo_dump_neighborhood	BFS neighborhood from a dump file	No
check_kernel_profiling	Were the kernel-profiling (LLO) flags active at capture?	Yes
list_kernel_invocations	Pallas/Mosaic custom-call executions + duration stats	Yes
get_llo_utilization	Per-functional-unit LLO % util (MXU/ALU/loads/...) — kernel bottleneck verdict	Yes
get_kernel_stage_breakdown	Mosaic pipeline stage times (ep_*) + DMA wait_ratio	Yes
list_llo_programs	Programs & pass checkpoints in an --xla_jf_dump_to LLO dump dir	No
get_llo_schedule_analysis	Bundle counts per HLO op / opcode (static attribution)	No
get_llo_static_utilization	Per-bundle slot occupancy vs capacity + hot ranges	No
get_llo_bundles	Windowed VLIW bundle listing (by address range / grep)	No
get_custom_call_mlir	Lowered Mosaic MLIR for a Pallas kernel (noop audit)	No

"Needs TF?" = requires tensorflow-cpu and XPROF_LOGDIR to be set.

Kernel profiling / LLO workflow

Capture with the XProf Kernel Profiling flags to light up the LLO-level tools:

# Trace-side (LLO utilization tracks, bundle markers, Mosaic ep_* stages):
LIBTPU_INIT_ARGS="--xla_enable_custom_call_region_trace=true \
                  --xla_xprof_register_llo_debug_info=true" \
python your_pallas_workload.py   # with jax.profiler.trace(logdir)

# Dump-side (per-pass LLO IR: VLIW bundles, per-bundle slot utilization):
LIBTPU_INIT_ARGS="--xla_jf_dump_to=/tmp/jf_dump --xla_jf_dump_llo_text=true \
                  --xla_mosaic_dump_to=/tmp/mosaic_dump" \
python your_pallas_workload.py
export XLA_JF_DUMP_DIR=/tmp/jf_dump

Note: the LLO dumper uses its own --xla_jf_dump_to flag — XLA's --xla_dump_to does not receive LLO dumps. Analysis flow: check_kernel_profiling → list_kernel_invocations → get_llo_utilization / get_kernel_stage_breakdown → (drilldown) list_llo_programs → get_llo_static_utilization → get_llo_bundles, with get_custom_call_mlir as the structural did-it-lower-as-planned audit. The Tensor Core trace markers carry bundle addresses that plug directly into get_llo_bundles' address_range. Trace % util values are LLO static slot occupancy over measured time windows; raw runtime counter sampling requires TPU v7+. Full guide (capture recipes, failure signatures, flag discovery): docs/KERNEL_PROFILING.md.

---

Configuration

Env Var	Default	Description
XPROF_MODE	(auto)	local = in-process converters, no server (CLI default); http = talk to a running xprof server at XPROF_URL (legacy mode); unset = local when converters + logdir are available, else http.
XPROF_LOGDIR	(auto-detected)	Root of the profile captures (the path you'd pass to xprof --logdir=...). Required for local mode and disk-based tools; the CLI's --logdir flag sets it per-call. Auto-detected from a localhost xprof server process when one is running. GCS paths work with tensorflow installed.
XPROF_URL	http://localhost:8791	URL of the running xprof server — only used in http mode.
XLA_JF_DUMP_DIR	(unset)	--xla_jf_dump_to directory for the LLO dump tools (list_llo_programs, get_llo_bundles, ...).
XLA_HLO_DUMP_DIR	(empty)	Directory the MCP dump tools read from (set to the same path as --xla_dump_to). Optional — can also be passed per-call as dump_dir.

---

Profile File Structure

<logdir>/
└── plugins/
    └── profile/
        └── <run_name>/
            ├── host0.xplane.pb
            ├── host1.xplane.pb
            └── <module_name>.hlo_proto.pb  (if exported)

The run_name is what you pass as run to all tools.

---

XLA HLO Dump Workflow

HLO dumps let you inspect the XLA compiler's work without running the xprof server and at every compilation stage, including per-pass diffs.

Enable dumps

# Step 1 — tell XLA to write dumps when running your program:
export XLA_FLAGS="--xla_dump_to=/tmp/hlo_dumps \
                  --xla_dump_hlo_as_text \
                  --xla_dump_hlo_pass_re=.*"
python your_script.py

# Step 2 — point the tools at those files (CLI):
xprof-cli list_hlo_dump_modules --dump_dir=/tmp/hlo_dumps
# or set XLA_HLO_DUMP_DIR=/tmp/hlo_dumps (used by both CLI and MCP server)

Recommended format: --xla_dump_hlo_as_text (default above). This produces human-readable HLO text files that all tools here support. Binary proto (--xla_dump_hlo_as_proto) requires a protobuf parser and is not needed for text-based analysis.

Two file naming formats are supported automatically:

JAX / TensorFlow (classic .hlo files):

/tmp/hlo_dumps/
├── module_0001.jit_my_fn.before_optimizations.hlo  ← raw JAX/TF output
├── module_0001.jit_my_fn.after_optimizations.hlo   ← final compiled HLO
├── module_0001.jit_my_fn.after_pass_HloCSE.hlo     ← per-pass (with --xla_dump_hlo_pass_re=.*)
└── ...

PyTorch/XLA (torch.compile with openxla backend — .txt files):

/tmp/hlo_dumps/
├── module_0006.jit_my_fn.0000.<pipeline>.after_<X>.before_<Y>.txt  ← first stage
├── module_0006.jit_my_fn.0001.<pipeline>.after_<X>.before_<Y>.txt  ← intermediate
├── module_0006.jit_my_fn.0005.<pipeline>.after_<X>.before_<Y>.txt  ← last stage
└── ...

Stage aliases: "before_optimizations" = first file, "after_optimizations" = last file. Intermediate stages accessible as "pass_NNNN" (e.g. "pass_0003").

Analysis workflow

list_hlo_dump_modules()                        # discover modules + stages (both formats)
get_hlo_dump("my_fn", "before_optimizations")  # first/pre-opt stage
get_hlo_dump("my_fn", "after_optimizations")   # last/final stage
diff_hlo_stages("my_fn",                       # what did the optimizer change?
    "before_optimizations", "after_optimizations")
diff_hlo_stages("my_fn",                       # what did one pass do? (JAX format)
    "after_pass_HloCSE", "after_pass_AlgebraicSimplifier")
diff_hlo_stages("my_fn",                       # two consecutive passes (PyTorch/XLA)
    "pass_0003", "pass_0004")
get_hlo_dump_neighborhood("fusion.3", "my_fn") # root-cause a specific op

For PyTorch/XLA with many instances of the same module name (e.g. multiple jit_splash_fn compilations), include the module number to target a specific one:

get_hlo_dump("module_0006.jit_my_fn", "after_optimizations")

When to use dumps vs. xprof

Situation	Use
No profiling yet, just want to see compiled HLO	HLO dumps
Want to see pre-optimization HLO (what JAX emitted)	HLO dumps
Debugging a compiler regression (which pass changed something)	HLO dumps
Want timing data (which op is slow)	xprof server
Want memory profile, step time breakdown	xprof server
Want timeline events (kernel durations, gaps)	xprof server + tensorflow

---

Recommended Analysis Workflow

Same tool names in both frontends; CLI shown (xprof-cli <tool> --logdir=... --run=...).

1. Orient — list_runs → get_kpi_metrics (headline numbers) → get_overview (bottleneck category, run environment).

2. Attribute — get_top_hlo_ops (time / FLOPs / bytes leaderboards) and get_roofline_model (per-op compute- vs memory-bound + ridge points; read its caveats field — cost-model FLOPs, custom-call blind, no comm class).

3. Branch on the bottleneck:

Signal	Next tools
Device idle / host-bound	get_input_pipeline (host stall decomposition), get_utilization_viewer (utilization over time)
Collective-bound (multi-chip)	get_pod_viewer (ICI/step breakdown); multi-slice: get_megascale_stats (DCN per-rendezvous)
Memory-bound / OOM hunting	get_memory_profile (peak + timeline) → get_memory_viewer (which buffer holds the peak, per module)
Large copy/reshape/transpose in top ops	detect_unfused_reshapes (automated HBM-materialization audit)
Compute-bound but low efficiency	get_framework_op_stats (map to framework op names), then the HLO lane below
A Pallas/Mosaic kernel dominates	kernel lane below — roofline is blind here

4. HLO root-cause — list_hlo_modules → get_hlo_module_content → get_hlo_neighborhood(instruction); compiler-stage questions: list_hlo_dump_modules / get_hlo_dump / diff_hlo_stages (from --xla_dump_to, no trace needed).

5. Kernel lane (Pallas/Mosaic) — check_kernel_profiling (were the capture flags on? ALWAYS first) → list_kernel_invocations → get_llo_utilization (per-unit bottleneck verdict) / get_kernel_stage_breakdown (DMA wait_ratio) → drilldown list_llo_programs → get_llo_static_utilization → get_llo_bundles, with get_custom_call_mlir as the lowering audit. See docs/KERNEL_PROFILING.md.

6. Timeline deep-dive (anything else) — list_xplane_events / aggregate_xplane_events with event regexes.

get_smart_suggestions is a free second opinion at any point.

---