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.

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