TorchTalk

TorchTalk

An MCP server that gives Claude Code deep understanding of PyTorch's cross-language architecture (Python → C++ → CUDA).

What It Does

TorchTalk provides structural knowledge that Claude can't get from just reading code:

• Binding chains: Trace torch.matmul → at::native::matmul → LinearAlgebra.cpp:1996
• Impact analysis: "If I modify GEMM, what breaks?" → Shows all 15 callers with file:line
• Dispatch mapping: Which backend (CPU/CUDA/MPS) handles each operation.
• Call graphs: C++ functions, call edges
• Test discovery: Find existing tests for any operator, browse test utilities

Quick Start

# Install
pip install -e .

# Add to Claude Code (one command)
claude mcp add torchtalk -s user -- torchtalk mcp-serve --pytorch-source /path/to/pytorch

# Add to Cursor: copies .claude/ into the project's .cursor/ and adds the torchtalk MCP
torchtalk cursor-add -C /path/to/your/project -p /path/to/pytorch

Requirements

• PyTorch source code: git clone https://github.com/pytorch/pytorch
• compile_commands.json (optional): For full C++ call graph, build PyTorch once:

  cd /path/to/pytorch && python setup.py develop
  

Available Tools

Tool	Description
get_status()	TorchTalk readiness summary across bindings, call graph, modules, tests
trace(func, focus?)	Trace any PyTorch op: Python → YAML → C++ → file:line
search(query, mode?, backend?)	mode="bindings": dispatch registrations. mode="kernels": CUDA kernel launches
graph(func, mode?, depth?, walk_python?)	mode="callers": inbound. mode="calls": outbound. mode="impact": transitive callers (depth/walk_python apply to impact only)
modules(name, mode?, focus?)	mode="trace": class details (focus="full" adds bases/docstring). mode="list": browse by category ("nn", "optim", "all")
tests(query?, mode?)	mode="find": search tests. mode="utils": list utilities (query ignored). mode="file_info": test file details
affected(funcs, depth?)	Map changed C++ functions (comma-separated) to impacted Python test files

CLI Commands

Command	Description
init --pytorch-source <path>	Save PyTorch source path to config
status	Show config and cache status
mcp-serve	Start the MCP server
index build [--no-wait]	Build or refresh the index and exit (headless)
index update --since <snapshot>	Incrementally refresh bindings for files changed since <snapshot>'s commit
snapshot save <name>	Capture current cache as a named snapshot
snapshot load <name|--nearest> [--force]	Restore a snapshot into the cache
snapshot list	List saved snapshots
snapshot delete <name>	Delete a snapshot
snapshot diff <a> <b> [--json]	Structural diff between two snapshots
snapshot export <name> [-o file]	Package a snapshot into a .tar.gz
snapshot import <archive> [--name new]	Extract a snapshot tarball

Snapshot names may use up to three /-separated components (e.g. main/abc1234/v1), so you can namespace snapshots by branch, commit, or release.

Snapshot Matching

Each snapshot records:

• source_fingerprint — hash of the indexed PyTorch source path (per-checkout).
• git_commit — short HEAD at save time.
• content_fingerprint — BLAKE2b over HEAD^{tree} + uncommitted diff; a Merkle-style content hash that's identical across checkouts of the same code.

snapshot load accepts a snapshot whose content or path fingerprint matches the current source. snapshot load --nearest resolves in tiered order: exact content match → exact commit match → most recent ancestor commit (via git merge-base --is-ancestor).

CI Integration

Snapshots make TorchTalk usable in CI without rebuilding the index per job. Build the index once on a nightly runner, ship the .tar.gz as a build artifact, and pull it into PR jobs.

Nightly job: build and publish

- run: torchtalk init --pytorch-source $GITHUB_WORKSPACE/pytorch
- run: torchtalk index build
- run: torchtalk snapshot save nightly/${{ github.sha }}
- run: torchtalk snapshot export nightly/${{ github.sha }} -o torchtalk-index.tar.gz
- uses: actions/upload-artifact@v4
  with: { name: torchtalk-index, path: torchtalk-index.tar.gz }

PR job: load and use

- uses: actions/download-artifact@v4
  with: { name: torchtalk-index }
- run: torchtalk snapshot import torchtalk-index.tar.gz
- run: torchtalk snapshot load --nearest
- run: torchtalk mcp-serve &

Fast PR refresh with index update

When only a few files changed vs. the baseline, skip the full rebuild:

- run: torchtalk snapshot load baseline --force
- run: torchtalk index update --since baseline