pytorch/docs/source
anwang cd68559d04 [Inductor] Support native Inductor as backend for MTIA (#158526)
This diff/PR includes the changes to support native Inductor integration for MTIA. The goal is to support `torch.compile(backend="inductor")` for MTIA. Inductor should generate code(triton kernel + python wrapper code) similar to CUDA. And the triton kernels can be launched eagerly.

The changes include:
- Add MTIA device interfaces used by Dynamo and Inductor, including APIs on device, stream, event, etc.
- Add required torch.mtia APIs, like is_bf16_supported, memory_allocated, set_stream_by_id, etc.
- MTIA specific codegen logic, for example, loading MTIA dynamic_library.
- Other necessary changes to integrate with Inductor codegn, following other devices like CUDA, XPU.
- Integrate with the [empty_strided_mtia](https://www.internalfb.com/code/fbsource/[0d017d3a4a1bdff7253f9c66a9f38e77bd62166b]/fbcode/caffe2/aten/src/ATen/native/mtia/EmptyTensor.cpp?lines=49%2C63%2C71%2C74%2C78) API that we’ve added for the new MTIA ATen backend.
- A change in Inductor runtime to avoid re-initialize MTIADriver.
- BUCK changes to include ATen-mtia in Inductor, and to use -USE_MTIA preprocessor flag.
- Update `test_mnist_e2e.py` to cover native Inductor as backend, using the `--use_native_inductor` flag.
- Add a personal script(`scripts/anwang/run_native_inductor_script.py`) for testing purpose.

Note:
- This approach(option 3) aims to provide a pytorch native approach of Inductor integration for MTIA, minimizing the onboarding overhead. The downside of this approach is that it doesn't leverage MTIA specific graph optimization, and is limited to eagerly launch overhead.
- MTIA will support another approach(option 2) to provide best performance, based on WrapperFxCodegen. We should be able to reuse the fundamental changes of this diff for option 2, like the device interfaces, steam/event APIs, etc, especially as WrapperFxCodegen inherits PythonWrapperCodegen.

Internal:
References:
- [post for context](https://fb.workplace.com/groups/mtiasw/permalink/1718377262384606/)
- [Inductor integration discussion(option 1/2/3)](https://docs.google.com/document/d/1p6363OXtVIRv1hPoaKlRSK3j-iir3QIbDd5bjyqCNig/edit?tab=t.0#heading=h.7s4ns6wcnhmb)
- [Project design doc(option 3)](https://docs.google.com/document/d/1jXUmhgoV9WvkMf-bcY3Od_kK9K_RDOdgHdt1LoQ5Tc4/edit?tab=t.0#heading=h.y43gwdqlv46w)
- [early prototying diff](https://www.internalfb.com/diff/D75110196)
- [MPS integration PR](https://github.com/pytorch/pytorch/pull/153959)
- [empty_strided_xpu PR](https://github.com/pytorch/pytorch/pull/126678)

Differential Revision: [D78458745](https://our.internmc.facebook.com/intern/diff/D78458745/)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/158526
Approved by: https://github.com/blaine-rister, https://github.com/jansel, https://github.com/eellison
2025-07-26 08:16:34 +00:00
..
_static [ONNX] Filter out torchscript sentences (#158850) 2025-07-24 20:59:06 +00:00
_templates Migrate to new theme (#149331) 2025-04-16 21:35:19 +00:00
community [doc] add weifengpy to torch distributed pocs (#158989) 2025-07-24 04:42:33 +00:00
elastic NUMA binding integration with elastic agent and torchrun (#149334) 2025-07-25 21:19:49 +00:00
export [export] Add runnable code to export docs (#158506) 2025-07-17 20:15:22 +00:00
notes [doc] remove FSDP1 developer note (#158991) 2025-07-24 08:21:54 +00:00
rpc Fix broken URLs (#152237) 2025-04-27 09:56:42 +00:00
scripts [ONNX] Filter out torchscript sentences (#158850) 2025-07-24 20:59:06 +00:00
accelerator.md Revert "Add unified memory APIs for torch.accelerator (#152932)" 2025-07-22 01:01:41 +00:00
amp.md [Docs] Convert to markdown: accelerator.rst, amp.rst, autograd.rst, backends.rst, benchmark_utils.rst (#155762) 2025-06-12 02:55:06 +00:00
autograd.md [Docs] Convert to markdown: accelerator.rst, amp.rst, autograd.rst, backends.rst, benchmark_utils.rst (#155762) 2025-06-12 02:55:06 +00:00
backends.md [ROCm] add flag torch.backends.miopen.immediate (#158951) 2025-07-25 04:01:51 +00:00
benchmark_utils.md [Docs] Convert to markdown: accelerator.rst, amp.rst, autograd.rst, backends.rst, benchmark_utils.rst (#155762) 2025-06-12 02:55:06 +00:00
bottleneck.rst
checkpoint.md Convert to markdown: checkpoint.rst (#156009) 2025-06-16 17:48:23 +00:00
complex_numbers.md Convert complex_numbers.rst to markdown (#156039) 2025-06-16 17:24:37 +00:00
cond.md [Docs] Convert to markdown cond.rst, config_mod.rst (#155653) 2025-06-13 20:58:57 +00:00
conf.py [BE] More torch.nn docs coverage test (except for torch.nn.parallel) (#158654) 2025-07-25 22:03:55 +00:00
config_mod.md [Docs] Convert to markdown cond.rst, config_mod.rst (#155653) 2025-06-13 20:58:57 +00:00
cpp_extension.rst xpu: support sycl with torch.utils.cpp_extension APIs (#132945) 2025-02-16 16:50:59 +00:00
cpp_index.rst [3/n] Remove references to TorchScript in PyTorch docs (#158315) 2025-07-15 21:14:18 +00:00
cpu.rst
cuda_environment_variables.rst
cuda._sanitizer.rst
cuda.md Fix #155016 for Docathon - convert rst to markdown (#155198) 2025-06-13 20:24:34 +00:00
cuda.tunable.md Fix #155016 for Docathon - convert rst to markdown (#155198) 2025-06-13 20:24:34 +00:00
cudnn_persistent_rnn.rst
cudnn_rnn_determinism.rst Fix broken URLs (#152237) 2025-04-27 09:56:42 +00:00
data.md Fix #155016 for Docathon - convert rst to markdown (#155198) 2025-06-13 20:24:34 +00:00
ddp_comm_hooks.md DOC: Convert to markdown: ddp_comm_hooks.rst, debugging_environment_variables.rst, deploy.rst, deterministic.rst, distributed.algorithms.join.rst (#155298) 2025-06-06 22:44:50 +00:00
debugging_environment_variables.md DOC: Convert to markdown: ddp_comm_hooks.rst, debugging_environment_variables.rst, deploy.rst, deterministic.rst, distributed.algorithms.join.rst (#155298) 2025-06-06 22:44:50 +00:00
deploy.md Revert "[BE] Remove torch deploy | remove torch deploy specific files (#158290)" 2025-07-25 16:09:39 +00:00
deterministic.md DOC: Convert to markdown: ddp_comm_hooks.rst, debugging_environment_variables.rst, deploy.rst, deterministic.rst, distributed.algorithms.join.rst (#155298) 2025-06-06 22:44:50 +00:00
distributed._dist2.md dist2: cleanup non-option methods on PG (missing, timeouts) (#158123) 2025-07-12 00:06:37 +00:00
distributed.algorithms.join.md DOC: Convert to markdown: ddp_comm_hooks.rst, debugging_environment_variables.rst, deploy.rst, deterministic.rst, distributed.algorithms.join.rst (#155298) 2025-06-06 22:44:50 +00:00
distributed.checkpoint.md [DCP] OSS Zero Overhead Checkpointing Implementation (#156207) 2025-06-29 03:19:48 +00:00
distributed.elastic.md NUMA binding integration with elastic agent and torchrun (#149334) 2025-07-25 21:19:49 +00:00
distributed.fsdp.fully_shard.md [FSDP2] explain user contract for fully_shard (#156070) 2025-06-17 10:03:19 +00:00
distributed.md [doc] Updates to distributed.md for XCCL backend (#155834) 2025-07-22 21:01:43 +00:00
distributed.optim.md Fix #155018 (convert distributed rst to markdown) (#155528) 2025-06-16 20:46:09 +00:00
distributed.pipelining.md Fix #155018 (convert distributed rst to markdown) (#155528) 2025-06-16 20:46:09 +00:00
distributed.tensor.md [BE] fix typos in docs/ (#156080) 2025-06-21 02:47:32 +00:00
distributed.tensor.parallel.md Convert to markdown: distributed.tensor.parallel.rst, distributed.tensor.rst, distributions.rst, dlpack.rst (#155297) 2025-06-13 22:08:37 +00:00
distributions.md Convert to markdown: distributed.tensor.parallel.rst, distributed.tensor.rst, distributions.rst, dlpack.rst (#155297) 2025-06-13 22:08:37 +00:00
dlpack.md Convert to markdown: distributed.tensor.parallel.rst, distributed.tensor.rst, distributions.rst, dlpack.rst (#155297) 2025-06-13 22:08:37 +00:00
docutils.conf
export.md [export] Add runnable code to export docs (#158506) 2025-07-17 20:15:22 +00:00
fft.md Convert to .md: draft_export.rst, export.ir_spec.rst, fft.rst (#155567) 2025-06-13 05:19:43 +00:00
fsdp.md Convert rst files to md (#155369) 2025-06-11 23:00:52 +00:00
func.api.md Convert rst files to md (#155369) 2025-06-11 23:00:52 +00:00
func.batch_norm.md Convert rst files to md (#155369) 2025-06-11 23:00:52 +00:00
func.md Convert rst files to md (#155369) 2025-06-11 23:00:52 +00:00
func.migrating.md Convert rst files to md (#155369) 2025-06-11 23:00:52 +00:00
func.ux_limitations.md Fix #155022 rst to markdown conversion (#155540) 2025-06-12 00:21:22 +00:00
func.whirlwind_tour.md Fix #155022 rst to markdown conversion (#155540) 2025-06-12 00:21:22 +00:00
future_mod.md Fix #155022 rst to markdown conversion (#155540) 2025-06-12 00:21:22 +00:00
futures.md Fix #155022 rst to markdown conversion (#155540) 2025-06-12 00:21:22 +00:00
fx.experimental.md Fix #155022 rst to markdown conversion (#155540) 2025-06-12 00:21:22 +00:00
fx.md [3/n] Remove references to TorchScript in PyTorch docs (#158315) 2025-07-15 21:14:18 +00:00
hub.md Convert hub.rst to hub.md (#155483) 2025-06-13 04:39:55 +00:00
index.md Update index.md (#155143) 2025-06-20 18:53:32 +00:00
jit_builtin_functions.rst [4/n] Remove references to TorchScript in PyTorch docs (#158317) 2025-07-16 20:01:34 +00:00
jit_language_reference_v2.md [1/n] Remove references to TorchScript in PyTorch docs (#158305) 2025-07-15 20:16:53 +00:00
jit_language_reference.md [2/n] Remove references to TorchScript in PyTorch docs (#158306) 2025-07-15 20:57:23 +00:00
jit_python_reference.md [3/n] Remove references to TorchScript in PyTorch docs (#158315) 2025-07-15 21:14:18 +00:00
jit_unsupported.md [4/n] Remove references to TorchScript in PyTorch docs (#158317) 2025-07-16 20:01:34 +00:00
jit_utils.md Convert to markdown: jit_python_reference.rst, jit_unsupported.rst, jit_utils.rst, library.rst (#155404) 2025-06-26 21:09:46 +00:00
jit.rst [4/n] Remove references to TorchScript in PyTorch docs (#158317) 2025-07-16 20:01:34 +00:00
library.md Convert to markdown: jit_python_reference.rst, jit_unsupported.rst, jit_utils.rst, library.rst (#155404) 2025-06-26 21:09:46 +00:00
linalg.md [Docs] Convert to markdown to fix 155025 (#155789) 2025-06-17 15:08:14 +00:00
logging.md [Docs] Convert to markdown to fix 155025 (#155789) 2025-06-17 15:08:14 +00:00
masked.md [Docs] Convert to markdown to fix 155025 (#155789) 2025-06-17 15:08:14 +00:00
math-quantizer-equation.png
meta.md [Docs] Convert to markdown to fix 155025 (#155789) 2025-06-17 15:08:14 +00:00
miscellaneous_environment_variables.md [Docs] Convert to markdown to fix 155025 (#155789) 2025-06-17 15:08:14 +00:00
mobile_optimizer.md DOC: Convert to markdown: mobile_optimizer.rst, model_zoo.rst, module_tracker.rst, monitor.rst, mps_environment_variables.rst (#155702) 2025-06-11 22:16:04 +00:00
model_zoo.md DOC: Convert to markdown: mobile_optimizer.rst, model_zoo.rst, module_tracker.rst, monitor.rst, mps_environment_variables.rst (#155702) 2025-06-11 22:16:04 +00:00
module_tracker.md DOC: Convert to markdown: mobile_optimizer.rst, model_zoo.rst, module_tracker.rst, monitor.rst, mps_environment_variables.rst (#155702) 2025-06-11 22:16:04 +00:00
monitor.md DOC: Convert to markdown: mobile_optimizer.rst, model_zoo.rst, module_tracker.rst, monitor.rst, mps_environment_variables.rst (#155702) 2025-06-11 22:16:04 +00:00
mps_environment_variables.md DOC: Convert to markdown: mobile_optimizer.rst, model_zoo.rst, module_tracker.rst, monitor.rst, mps_environment_variables.rst (#155702) 2025-06-11 22:16:04 +00:00
mps.md Fix/issue #155027 (#155252) 2025-06-08 21:17:31 +00:00
mtia.md Fix/issue #155027 (#155252) 2025-06-08 21:17:31 +00:00
mtia.memory.md [Inductor] Support native Inductor as backend for MTIA (#158526) 2025-07-26 08:16:34 +00:00
multiprocessing.md Fix/issue #155027 (#155252) 2025-06-08 21:17:31 +00:00
name_inference.md Fix/issue #155027 (#155252) 2025-06-08 21:17:31 +00:00
named_tensor.md Convert to markdown: named_tensor.rst, nested.rst, nn.attention.bias.rst, nn.attention.experimental.rst, nn.attention.flex_attention.rst #155028 (#155696) 2025-06-14 03:32:00 +00:00
nested.md Convert to markdown: named_tensor.rst, nested.rst, nn.attention.bias.rst, nn.attention.experimental.rst, nn.attention.flex_attention.rst #155028 (#155696) 2025-06-14 03:32:00 +00:00
nn.aliases.md [BE] Use .md instead of .rst for nn.aliases doc (#158666) 2025-07-25 22:03:55 +00:00
nn.attention.bias.md Convert to markdown: named_tensor.rst, nested.rst, nn.attention.bias.rst, nn.attention.experimental.rst, nn.attention.flex_attention.rst #155028 (#155696) 2025-06-14 03:32:00 +00:00
nn.attention.experimental.md Convert to markdown: named_tensor.rst, nested.rst, nn.attention.bias.rst, nn.attention.experimental.rst, nn.attention.flex_attention.rst #155028 (#155696) 2025-06-14 03:32:00 +00:00
nn.attention.flex_attention.md Add kernel options to flex docs (#158875) 2025-07-23 19:05:19 +00:00
nn.attention.rst
nn.functional.rst
nn.init.rst
nn.rst [BE] More torch.nn docs coverage test (except for torch.nn.parallel) (#158654) 2025-07-25 22:03:55 +00:00
notes.md Migrate to new theme (#149331) 2025-04-16 21:35:19 +00:00
onnx_export.md [ONNX] Filter out torchscript sentences (#158850) 2025-07-24 20:59:06 +00:00
onnx_ops.md [ONNX] Implement Attention-23 (#156431) 2025-06-20 23:54:57 +00:00
onnx_verification.md Convert to .md: onnx_verification.rst, onnx.rst, package.rst, (#155556) 2025-06-10 21:40:40 +00:00
onnx.md [ONNX] Filter out torchscript sentences (#158850) 2025-07-24 20:59:06 +00:00
optim.aliases.md Document the rest of the specific optimizer module APIs (#158669) 2025-07-19 07:27:15 +00:00
optim.md Document the rest of the specific optimizer module APIs (#158669) 2025-07-19 07:27:15 +00:00
package.md [3/n] Remove references to TorchScript in PyTorch docs (#158315) 2025-07-15 21:14:18 +00:00
profiler.md Convert rst to markdown - profiler.rst #155031 (#155559) 2025-06-13 05:02:54 +00:00
pytorch-api.md Migrate to new theme (#149331) 2025-04-16 21:35:19 +00:00
quantization-support.md Convert to markdown: quantization-accuracy-debugging.rst, quantization-backend-configuration.rst, quantization-support.rst, random.rst (#155520) 2025-06-18 18:46:04 +00:00
quantization.rst Remove pytorch quant docs since we are moving to torchao (#157766) 2025-07-11 03:21:47 +00:00
random.md Convert to markdown: quantization-accuracy-debugging.rst, quantization-backend-configuration.rst, quantization-support.rst, random.rst (#155520) 2025-06-18 18:46:04 +00:00
rpc.md RPC tutorial audit (#157938) 2025-07-10 14:15:37 +00:00
signal.md Convert rst to md: rpc.rst, signal.rst, size.rst, special.rst (#155430) 2025-06-18 01:27:04 +00:00
size.md Convert rst to md: rpc.rst, signal.rst, size.rst, special.rst (#155430) 2025-06-18 01:27:04 +00:00
sparse.rst [BE] fix typos in docs/ (#156080) 2025-06-21 02:47:32 +00:00
special.md Convert rst to md: rpc.rst, signal.rst, size.rst, special.rst (#155430) 2025-06-18 01:27:04 +00:00
storage.rst Super tiny fix typo (#151212) 2025-04-14 16:47:40 +00:00
tensor_attributes.rst revamp dtype documentation for 2025 (#156087) 2025-06-27 13:10:23 +00:00
tensor_view.rst [docs] fix numpy docs reference (#147697) 2025-02-26 01:30:03 +00:00
tensorboard.rst
tensors.rst revamp dtype documentation for 2025 (#156087) 2025-06-27 13:10:23 +00:00
testing.md Convert to markdown: testing.rst, threading_environment_variables.rst, torch_cuda_memory.rst, torch_environment_variables.rst, torch_nccl_environment_variables.rst (#155523) 2025-06-10 20:38:36 +00:00
threading_environment_variables.md Convert to markdown: testing.rst, threading_environment_variables.rst, torch_cuda_memory.rst, torch_environment_variables.rst, torch_nccl_environment_variables.rst (#155523) 2025-06-10 20:38:36 +00:00
torch_cuda_memory.md Convert to markdown: testing.rst, threading_environment_variables.rst, torch_cuda_memory.rst, torch_environment_variables.rst, torch_nccl_environment_variables.rst (#155523) 2025-06-10 20:38:36 +00:00
torch_environment_variables.md Convert to markdown: testing.rst, threading_environment_variables.rst, torch_cuda_memory.rst, torch_environment_variables.rst, torch_nccl_environment_variables.rst (#155523) 2025-06-10 20:38:36 +00:00
torch_nccl_environment_variables.md Convert to markdown: testing.rst, threading_environment_variables.rst, torch_cuda_memory.rst, torch_environment_variables.rst, torch_nccl_environment_variables.rst (#155523) 2025-06-10 20:38:36 +00:00
torch.aliases.md Remove torch.functional entries from the doc ignore list (#158581) 2025-07-25 17:19:01 +00:00
torch.compiler_aot_inductor_minifier.md Converting .rst files to .md files (#155377) 2025-06-13 22:54:27 +00:00
torch.compiler_aot_inductor.md [export] Update docs (#157750) 2025-07-16 19:53:12 +00:00
torch.compiler_api.md Implement guard collectives (optimized version) (#156562) 2025-06-24 04:59:49 +00:00
torch.compiler_backward.md Add AOTDispatcher config to set backward autocast behavior (#156356) 2025-06-27 14:58:58 +00:00
torch.compiler_cudagraph_trees.md DOC: Convert to markdown: torch.compiler_best_practices_for_backends.rst, torch.compiler_cudagraph_trees.rst, torch.compiler_custom_backends.rst, torch.compiler_dynamic_shapes.rst, torch.compiler_dynamo_deepdive.rst (#155137) 2025-06-10 20:51:05 +00:00
torch.compiler_custom_backends.md DOC: Convert to markdown: torch.compiler_best_practices_for_backends.rst, torch.compiler_cudagraph_trees.rst, torch.compiler_custom_backends.rst, torch.compiler_dynamic_shapes.rst, torch.compiler_dynamo_deepdive.rst (#155137) 2025-06-10 20:51:05 +00:00
torch.compiler_dynamic_shapes.md DOC: Convert to markdown: torch.compiler_best_practices_for_backends.rst, torch.compiler_cudagraph_trees.rst, torch.compiler_custom_backends.rst, torch.compiler_dynamic_shapes.rst, torch.compiler_dynamo_deepdive.rst (#155137) 2025-06-10 20:51:05 +00:00
torch.compiler_dynamo_deepdive.md DOC: Convert to markdown: torch.compiler_best_practices_for_backends.rst, torch.compiler_cudagraph_trees.rst, torch.compiler_custom_backends.rst, torch.compiler_dynamic_shapes.rst, torch.compiler_dynamo_deepdive.rst (#155137) 2025-06-10 20:51:05 +00:00
torch.compiler_dynamo_overview.md convert: rst to myst pr 1/2 (#155840) 2025-06-13 18:02:28 +00:00
torch.compiler_fake_tensor.md convert: rst to myst pr 1/2 (#155840) 2025-06-13 18:02:28 +00:00
torch.compiler_faq.md convert: rst to myst pr2/2 (#155911) 2025-06-16 00:44:44 +00:00
torch.compiler_fine_grain_apis.md convert: rst to myst pr2/2 (#155911) 2025-06-16 00:44:44 +00:00
torch.compiler_get_started.md convert: rst to myst pr2/2 (#155911) 2025-06-16 00:44:44 +00:00
torch.compiler_inductor_profiling.md Convert compiler rst files to markdown (#155335) 2025-06-10 01:12:11 +00:00
torch.compiler_inductor_provenance.rst Update provenance tracking doc (#154062) 2025-05-23 17:09:52 +00:00
torch.compiler_ir.md [export] Update docs (#157750) 2025-07-16 19:53:12 +00:00
torch.compiler_nn_module.md Convert compiler rst files to markdown (#155335) 2025-06-10 01:12:11 +00:00
torch.compiler_performance_dashboard.md Convert compiler rst files to markdown (#155335) 2025-06-10 01:12:11 +00:00
torch.compiler_profiling_torch_compile.md [Docs] Update PT2 Profiler Torch-Compiled Region Image (#158066) 2025-07-11 07:56:45 +00:00
torch.compiler_transformations.md [Docs] Convert to markdown: torch.compiler_transformations.rst, torch.compiler.config.rst (#155347) 2025-06-11 18:55:30 +00:00
torch.compiler_troubleshooting_old.md Revert "Add torch compile force disable caches alias (#158072)" 2025-07-18 04:58:24 +00:00
torch.compiler_troubleshooting.md [Docs] Convert to markdown: torch.compiler_troubleshooting.rst (#155514) 2025-06-10 15:41:31 +00:00
torch.compiler.config.md [Docs] Convert to markdown: torch.compiler_transformations.rst, torch.compiler.config.rst (#155347) 2025-06-11 18:55:30 +00:00
torch.compiler.md [ONNX] Remove legacy Dort (#158258) 2025-07-15 19:14:06 +00:00
torch.overrides.md DOC: Convert to markdown: torch.overrides.rst, type_info.rst, utils.rst, xpu.rst (#155088) 2025-06-06 20:16:13 +00:00
torch.rst Remove torch.functional entries from the doc ignore list (#158581) 2025-07-25 17:19:01 +00:00
type_info.md DOC: Convert to markdown: torch.overrides.rst, type_info.rst, utils.rst, xpu.rst (#155088) 2025-06-06 20:16:13 +00:00
utils.md DOC: Convert to markdown: torch.overrides.rst, type_info.rst, utils.rst, xpu.rst (#155088) 2025-06-06 20:16:13 +00:00
xpu.md DOC: Convert to markdown: torch.overrides.rst, type_info.rst, utils.rst, xpu.rst (#155088) 2025-06-06 20:16:13 +00:00