Skip to content

Add the analyze_async_overlap CLI subcommand to measure how much of each asynchronous XLA operation (collectives, async copies, send/recv) is hidden behind compute on the same device, and how much time multiple engines on the same TPU chip (TensorCore and SparseCore) spend busy concurrently: - #3830

Open
copybara-service[bot] wants to merge 1 commit into
masterfrom
test_997089732

Conversation

@copybara-service

Copy link
Copy Markdown

Add the analyze_async_overlap CLI subcommand to measure how much of each asynchronous XLA operation (collectives, async copies, send/recv) is hidden behind compute on the same device, and how much time multiple engines on the same TPU chip (TensorCore and SparseCore) spend busy concurrently:

  • cli/internal/async_overlap.py: Pure-Python interval analysis over device TimelineEvents (XLA Ops, Async XLA Ops, Sparse Core Ops). Uses Async XLA Ops windows when present and falls back to FIFO <op>-start / <op>-done pairing on the op lines when absent. Classifies each async window as hidden (covered fraction >= hidden_threshold, default 0.9), exposed (<= exposed_threshold, default 0.1), or partial, and computes pairwise concurrent busy time across engines per TPU chip.
  • cli/internal/{google,oss}/xplane_tools.py: Adds iter_timeline_events to extract device timeline events with absolute picosecond start times across both XPlane backends.
  • cli/tools/analyze_async_overlap_tool.py: CLI/MCP tool wrapper registered as analyze_async_overlap in xprof_cli.py and server.py.
  • skills/xprof/: Documents analyze_async_overlap in references/analyze_async_overlap.md and SKILL.md.

… each asynchronous XLA operation (collectives, async copies, `send`/`recv`) is hidden behind compute on the same device, and how much time multiple engines on the same TPU chip (TensorCore and SparseCore) spend busy concurrently:

* `cli/internal/async_overlap.py`: Pure-Python interval analysis over device `TimelineEvent`s (`XLA Ops`, `Async XLA Ops`, `Sparse Core Ops`). Uses `Async XLA Ops` windows when present and falls back to FIFO `<op>-start` / `<op>-done` pairing on the op lines when absent. Classifies each async window as `hidden` (covered fraction >= `hidden_threshold`, default 0.9), `exposed` (<= `exposed_threshold`, default 0.1), or `partial`, and computes pairwise concurrent busy time across engines per TPU chip.
* `cli/internal/{google,oss}/xplane_tools.py`: Adds `iter_timeline_events` to extract device timeline events with absolute picosecond start times across both XPlane backends.
* `cli/tools/analyze_async_overlap_tool.py`: CLI/MCP tool wrapper registered as `analyze_async_overlap` in `xprof_cli.py` and `server.py`.
* `skills/xprof/`: Documents `analyze_async_overlap` in `references/analyze_async_overlap.md` and `SKILL.md`.

PiperOrigin-RevId: 997089732
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant