xprof-profiling-analysis

XProf Profiling Analysis

TPU/GPU training performance analysis. Primary workflow uses XProf MCP tools (xprof_*) for live queries; appendix sections provide deep domain knowledge for offline analysis and result interpretation.

Step 0: Connection Check (MUST do first)

Before any analysis, verify that XProf MCP tools are available:

1. Try calling xprof_list_runs().

2. If it succeeds — skip to Step 1.

3. If the tool is not found — the plugin is not enabled. Tell the user:

   XProf MCP tools are not available. Enable the plugin and restart Claude Code:

   claude settings set enabledPlugins.xprof-profiling-analysis@primatrix-skills true
   

   Stop here — wait for the user to restart before continuing.

4. If the tool exists but returns an error — ask the user to configure XPROF_URL to point to their XProf instance. They need to either:

   • Run kubectl port-forward svc/xprof-service <local-port>:8080 and set XPROF_URL=http://localhost:<local-port>
   • Or point to an XProf instance directly if running locally

   The URL is configured in the plugin's .mcp.json env or via the XPROF_URL environment variable.

Analysis Workflow

Step 1: Identify Runs

xprof_list_runs()

Pick the runs to analyze. If comparing optimizations, identify:

• Baseline run: the main branch or previous best
• Experiment run: the optimization branch

Step 2: Overview

xprof_overview(run="<run_name>")

Check key metrics:

• MXU utilization — target >50% for compute-heavy workloads
• Device idle time — high idle = pipeline stall or host bottleneck
• Top ops by time — identify the dominant operations

Step 3: Framework Op Breakdown

xprof_framework_ops(run="<run_name>", top=30)

This gives JAX-level operation statistics. Filter by category to focus:

xprof_framework_ops(run="<run_name>", category="gla")       # GLA kernels
xprof_framework_ops(run="<run_name>", category="moe")       # MoE/GMM kernels
xprof_framework_ops(run="<run_name>", category="matmul")    # Dense matmuls
xprof_framework_ops(run="<run_name>", category="collective") # Communication

Categories: gla, moe, attention, norm, embedding, collective, matmul, other

Step 4: A/B Comparison

xprof_compare(run_a="<baseline>", run_b="<experiment>", top=20)

Focus on:

• Ops with largest absolute delta (biggest time savings/regressions)
• Ops where delta_pct > 5% (significant relative change)
• Category-level trends (did GLA get faster? Did communication increase?)

Step 5: Memory Analysis

xprof_memory(run="<run_name>")

Check:

• Peak HBM usage vs capacity — how much headroom?
• Fragmentation — >10% may indicate allocation pattern issues
• Stack reserved vs heap allocated — stack is XLA buffers, heap is dynamic

Step 6: Detailed Inspection

For deeper analysis, get browser URLs for human inspection:

xprof_trace_url(run="<run_name>")

Share these URLs with the user for:

• Trace viewer — timeline of all operations (zoom into specific steps)
• Memory viewer — allocation/deallocation timeline
• Pod viewer — cross-host communication patterns

Step 7: Report

Generate a comparison table:

| Metric | Baseline | Experiment | Delta |
|--------|----------|------------|-------|
| Step time | X.XXs | X.XXs | -X.X% |
| MXU utilization | XX% | XX% | +X.X% |
| Top hotspot | op_name (XX%) | op_name (XX%) | -X.X% |
| Peak HBM | XX.X GiB | XX.X GiB | -X.X GiB |

Include:

1. Summary of what changed and why
2. Per-category breakdown (GLA, MoE, collective, etc.)
3. Top 5 ops by time change
4. XProf UI links for human verification

Quick Recipes

"How fast is the latest profiling run?"

runs = xprof_list_runs()
xprof_overview(run=runs[-1])

"What's the GLA kernel time?"

xprof_framework_ops(run="<run>", category="gla")

"Did my optimization help?"

xprof_compare(run_a="<baseline>", run_b="<experiment>")

"Is there a memory issue?"

xprof_memory(run="<run>")
xprof_trace_url(run="<run>")  # memory viewer URL for human inspection

"Show me the trace timeline"

xprof_trace_url(run="<run>")  # returns browser URLs for all XProf views

Roofline Analysis

Bound Type	Indicator	Optimization Direction
Compute bound	High FLOP rate, arithmetic intensity > 313 FLOPs/byte (TPUv7x)	Reduce FLOPs (algorithmic), increase MXU utilization
Memory bound	High bandwidth utilization, low arithmetic intensity	Reduce data movement, fuse ops, quantize
Latency bound	Low both FLOP rate and bandwidth	Pipeline better, reduce sync barriers, batch ops

Ridge Point = Peak_FLOPS / HBM_BW: TPU v7x = 2307T / 7380 GB/s = 313 FLOPs/byte

Common Bottleneck Patterns

Pattern	Symptom	Action
Stalled communication	High device idle, large collective ops	Overlap compute + comms, check XLA flags (see Appendix D)
Kernel regression	Same op category significantly slower after change	Check pallas-kernel version, chunk_size, block_size
Memory pressure	Peak HBM near capacity, high fragmentation	Enable remat, reduce batch size, chunked CE (see Appendix E)
Load imbalance	Cross-device step time variance >5%	Check MoE expert routing balance, PP stage sizing
Host bottleneck	High host idle time	Check data pipeline, increase grain workers

---

Appendix: Deep Analysis Reference

Domain knowledge for interpreting MCP results and offline analysis when MCP data is insufficient.

Appendix A: HLO Operator Classification

XLA-compiled op names lose JAX semantics. Classification by hlo_category + name patterns:

Category	hlo_category / Name Pattern	Meaning
matmul	convolution fusion	Matrix multiply (XLA uses convolution for matmul)
attention	splash_mha, flash_attention	Attention kernel
communication	all-reduce, all-gather, reduce-scatter, all-to-all (incl. async-start/async-done)	Collective comms
custom_kernel	custom-call (excl. GMM/TGMM/splash_mha)	Pallas kernel, offload
custom_fusion	custom fusion	XLA custom fusions
elementwise_fusion	loop fusion, non-fusion elementwise	Elementwise fusions
data_formatting	data formatting	Tensor layout conversion (FSDP sharding copies)
routing	sort	MoE TopK routing

Key: async-done events for communication (e.g., all-reduce.*.call-done) have duration = stall time (compute engine waiting). This is the core metric for exposed communication cost.

Key Operator Patterns

Pattern	Meaning	Phase
convolution_bitcast_fusion.*	MatMul (fwd or bwd dlhs)	fwd/bwd
tgmm.*	MoE weight gradient (Transposed GMM)	backward only
splash_mha_fwd_*	SplashAttention forward	forward
splash_mha_dkv_* / splash_mha_dq_*	SplashAttention gradients	backward
ragged-all-to-all	MoE expert routing (EP)	fwd/bwd

Forward/Backward Identification

Via tf_op field (most reliable):

• Contains transpose(jvp(...)) → backward
• Contains jvp(...) without transpose → forward
• Also extracts layer number (e.g., moe_layers_0)

Via op name (quick): tgmm.* → always backward; splash_mha_d* → backward; splash_mha_fwd_* → forward.

Appendix B: MFU Calculation

MFU = F_useful / (Peak_cluster x T_step)

F_useful = GBS x seq_len x FLOPs_per_token(fwd+bwd)
Peak_cluster = num_devices x peak_flops_per_device

MoE caveat: Per-token FLOPs much lower than equivalent dense model (only top-k experts activated). MoE MFU is naturally low (~10%) — this is architectural, not poor hardware utilization.

Pallas kernel caveat: SplashAttention model_flops reports 0. FLOPs from trace accumulation underestimates actual compute. GMM/TGMM model_flops correct in newer versions.

Appendix C: Communication Analysis

Primitive-to-Strategy Mapping

Primitive	Parallelism	Purpose
AllGather	FSDP	Collect sharded weights
ReduceScatter	FSDP	Aggregate and distribute gradients
AllReduce	DP (non-FSDP)	Gradient sync
AllToAll	EP	Expert routing data exchange

Overlap Measurement

Single TPU device executes all ops serially. Communication overlap = between communication engine (ICI/DCN) and compute engine (TensorCore):

• async-done with duration = 0: Communication completed, no stall → effectively overlapped
• async-done with duration > 0: Compute engine waiting = exposed communication cost

Wrong: Don't analyze interval overlaps on single-device timeline (events are serial). Measure async-done duration instead.

Data Formatting

data formatting ops = tensor layout conversions. When >5% of step time: check for unnecessary FSDP sharding layout copies.

Appendix D: XLA Flags for Communication Overlap

Check LIBTPU_INIT_ARGS when communication is not overlapped.

Continuation Fusion (CF):

Flag	Purpose
xla_tpu_enable_async_collective_fusion=true	Enable async collective fusion
xla_tpu_enable_async_collective_fusion_fuse_all_gather=true	AllGather CF
xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true	AllReduce CF
xla_tpu_enable_async_collective_fusion_fuse_reduce_scatter=true	ReduceScatter CF
xla_tpu_overlap_compute_collective_tc=true	TensorCore overlap compute/comm

SparseCore Offloading (v7x recommended):

Flag	Purpose
xla_tpu_enable_sparse_core_collective_offload_all_gather=true	AllGather via SparseCore
xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true	ReduceScatter via SparseCore
xla_tpu_enable_sparse_core_collective_offload_all_reduce=true	AllReduce via SparseCore

CF and SparseCore are mutually exclusive for the same primitive.

Symptom	Likely Missing Flags
AllReduce stall high	fuse_all_reduce=true or SparseCore offload
AllGather stall high	fuse_all_gather=true or SparseCore offload
All comms not overlapped	async_collective_fusion=true base flag
MoE comms not optimized	DATA_PARALLEL_OVERLAP flags

Appendix E: HBM Memory Deep Analysis

Data Sources (priority order)

Source	Best For	Notes
xprof_memory() MCP	Quick peak/capacity/fragmentation	Use first
memory_viewer.json	Peak composition, optimization targets	Static upper bound
xplane.pb /host:CPU	Runtime actual HBM	actual = reserved - available
hlo_proto.pb	Buffer-level detail	Large files, slow

Critical: bytes_allocated only tracks dynamic I/O buffers, NOT internal workspace (~70 GiB). Actual HBM = bytes_reserved - bytes_available.

Buffer Classification (memory_viewer.json maxHeap)

Category	Rule (tfOpName / groupName)	Typical Shape
Parameters	groupName=="Parameter" + state.params	Weight shapes
Optimizer	groupName=="Parameter" + opt_state	Same as weights, f32
Logits/CE	logits_dense / cross_entropy / softmax	[B,T,vocab]
MoE experts	moe / expert / gmm / router	[num_experts,dim,dim]
Attention/MLA	attention / mla / wq_ / wkv_	[B,T,H,D]

Optimization Targets

Optimization	Expected Savings	Impact
Chunked cross-entropy	~18 GiB (3x vocab logits)	No accuracy impact
save_out_proj remat	Large activation savings	~33% recompute increase
Increase FSDP	Linear per-chip param reduction	More communication
bf16 optimizer (mu_dtype)	~2 GiB (mu f32 -> bf16)	May affect stability

Vocab-sized buffer detection: Search maxHeap for any dimension >= 100K → chunked CE targets.

Memory Anomaly Diagnostics

From memory_viewer.json heapSizes timeline:

1. Spike detection: Find largest delta jumps — usually multiple vocab-sized buffers briefly overlapping
2. Buffer overlap: Group peak-time buffers by shape to quantify per-group contribution
3. Lifetime anomaly: Temporary buffers with >80% lifetime + >100 MiB = critical (but MoE expert AllGather is a known false positive)
4. Phase curve: Sample heapSizes by schedule % to see memory across compute phases

Appendix F: MoE Backward Pass Structure

GMM / TGMM Relationship

Forward:  Y = gmm(X, W)
Backward:
  dlhs = gmm(dY, W^T)   # activation gradient → convolution_bitcast_fusion
  drhs = tgmm(X^T, dY)  # weight gradient → tgmm.*

• dlhs: On critical path, serial across layers
• drhs: Independent per layer

Layer count: MoE_layers = GMM_count / 4 = TGMM_count / 2

XLA schedules backward in two loops: (1) all dlhs (serial), then (2) all tgmm (parallelizable).

Appendix G: Offline CLI Tool

When MCP is insufficient, use tpu-profiling/scripts/xprof.py:

xprof memory peak --run-dir <run>           # Peak HBM composition
xprof memory diagnose --run-dir <run>       # Anomaly detection
xprof compute breakdown --run-dir <run>     # Op category breakdown
xprof compute mfu --run-dir <run> --gbs 5120 --num-chips 128
xprof comm overlap --run-dir <run>          # Async-done stall analysis
xprof compare --runs 115,183               # A/B comparison
xprof audit flags --flags-file launch.sh   # XLA flags check

Data source priority: xplane.pb (complete) > trace.json.gz (1M event hard limit — check Complete (X) events).

Common Pitfalls

Pitfall	Correct
XLA convolution = CNN	XLA uses convolution for all matmul
hlo_category distinguishes fwd/bwd	Cannot — use op name patterns or tf_op
Low MFU = poor utilization	MoE MFU naturally low; understand per-token FLOPs
bytes_allocated = actual HBM	Only dynamic I/O buffers; actual = reserved - available
memory_viewer peak = runtime peak	Static upper bound; runtime 5-15% lower
trace.json.gz is complete	1M event hard limit — use xplane.pb if truncated
Comms not overlapped = need algorithm change	Check XLA flags first (CF/SparseCore)
Single-device interval overlap = comm overlap	Single device is serial; measure async-done duration
Logits size is fixed	Chunked CE: [B,T,vocab] → [B,chunk,vocab], saving ~18 GiB