pytorch-mcp

pytorch-mcp is a full-featured MCP server for PyTorch documentation workflows. It indexes the repository's local docs/ tree and exposes search, page retrieval, symbol lookup, code-example extraction, troubleshooting, and question-answering tools that an LLM can use to help developers work with PyTorch.

Features

• Uses docs/ as the source of truth for documentation-aware tools.

• Indexes local Markdown and reStructuredText PyTorch docs.

• Searches by workflow, concept, topic, or exact symbol such as torch.compile.

• Returns grounded snippets, headings, and page metadata for follow-up exploration.

• Extracts code examples from relevant docs pages.

• Understands declared Torch ecosystem libraries from pyproject.toml and inspects runtime availability.

• Recommends reading paths for tasks like training, compilation, data loading, profiling, and troubleshooting.

• Exposes MCP tools, resources, prompts, plus HTTP health/readiness routes.

Related MCP server: Documentation Retrieval MCP Server (DOCRET)

Server Instructions

The MCP server is intended to behave like a PyTorch development copilot:

• Use local docs/ content as the authoritative source for documentation-aware answers.

• Inspect declared and installed Torch libraries before making environment-specific recommendations.

• Use planning, template-generation, code-inspection, and runtime-validation tools to help developers build models faster.

• Keep debugging and optimization advice grounded in retrieved docs, parsed traces, profiler data, and runtime checks when available.

Tools

• list_doc_topics

• search_docs

• get_doc_page

• get_symbol_reference

• extract_code_examples

• answer_pytorch_question

• recommend_docs

• troubleshoot_pytorch

• plan_model_build

• assemble_training_stack

• generate_training_loop_template

• generate_task_specific_template

• generate_training_project_template

• review_training_code

• suggest_model_architecture

• choose_loss_and_optimizer

• optimize_data_pipeline

• diagnose_training_issue

• inspect_pytorch_code

• inspect_runtime_environment

• execute_pytorch_snippet

• run_forward_pass_check

• benchmark_compile_candidate

• validate_training_setup

• list_torch_libraries

• inspect_torch_library

• recommend_torch_libraries

• audit_torch_stack

• parse_stack_trace

• analyze_stack_trace

• analyze_shape_mismatch

• parse_profiler_export

• analyze_profiler_summary

Resources

• pytorch://server/capabilities

• pytorch://project/settings

• pytorch://docs/index

• pytorch://docs/categories

• pytorch://docs/page/{doc_path}

• pytorch://docs/category/{category}

• pytorch://docs/search/{query}?limit=5

• pytorch://reference/overview

Prompts

• explain pytorch topic

• plan pytorch implementation

• debug pytorch issue

• compare pytorch approaches

• build pytorch model

• review pytorch training code

• choose pytorch training objective

• diagnose pytorch training issue

• inspect pytorch code

• analyze pytorch stack trace

• recommend torch libraries

Run

Install dependencies:

uv sync

Run over stdio:

uv run python mcp_server.py --transport stdio

Run over HTTP:

uv run python mcp_server.py --transport http --host 127.0.0.1 --port 8000

Health endpoints:

• GET /healthz

• GET /readyz

Configuration

Important environment variables:

• PYTORCH_MCP_DOCS_ROOT

• PYTORCH_MCP_MAX_SEARCH_RESULTS

• PYTORCH_MCP_MAX_PAGE_CHARACTERS

• PYTORCH_MCP_MAX_CODE_EXAMPLES

• PYTORCH_MCP_TRANSPORT

• PYTORCH_MCP_HOST

• PYTORCH_MCP_PORT

Example:

PYTORCH_MCP_DOCS_ROOT=/path/to/pytorch/docs \
uv run python mcp_server.py --transport stdio

By default the server reads from this repository's docs/ directory. If you package or deploy the server elsewhere, point PYTORCH_MCP_DOCS_ROOT at a local PyTorch documentation checkout.

Testing

just test