[WIP] simple MoE inference API

[vkuzo]

Summary:

Not ready for review yet

Test Plan:

pytest test/prototype/moe_inference/test_base.py -s -x

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• -> [WIP] simple MoE inference API #2744

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2744

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 8 New Failures

As of commit af8831b with merge base 317179e ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CPU 2.6, linux.4xlarge, torch==2.6.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t e1f25bc0b4d967ad95eeb6c5121a104c31b1e64eddda3fe99c56e8aaa9f1ceaf /exec failed with exit code 2
• Run Regression Tests / test (CPU 2.7, linux.4xlarge, torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 0b55f5481d5f9794418cd0f5cea60aba6b999bf5945d8fc67a029ae28e5f27d1 /exec failed with exit code 2
• Run Regression Tests / test (CPU 2.8, linux.4xlarge, torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 1826439b644b99dff60a8e24e61468dcd8fea8245dd61afa9937dec77e819dbf /exec failed with exit code 2
• Run Regression Tests / test (CUDA 2.6, linux.g5.12xlarge.nvidia.gpu, torch==2.6.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 6c6b3be510858759e5e361a1b2c8b7ce8f5643084c70e15b0cce3812524e1933 /exec failed with exit code 2
• Run Regression Tests / test (CUDA 2.7, linux.g5.12xlarge.nvidia.gpu, torch==2.7.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 6d14def84dfb365fc12f7142e3713c9a206002bf3b4b06661cb2b7b088f72eb6 /exec failed with exit code 2
• Run Regression Tests / test (CUDA 2.8, linux.g5.12xlarge.nvidia.gpu, torch==2.8.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t adcc675dad19b22d0a1aac1f74e80d5f35c57610566052014c7a4598f78a13a7 /exec failed with exit code 2
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  RuntimeError: Command docker exec -t 2a400dcabb5e68178534d41a6e057c61fa85bc629f97a79f53151014ec454682 /exec failed with exit code 2
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t 712857951b115349b4e6de417f2c98f07e783f3429d7ce0da4207a7afff8b295 /exec failed with exit code 2

This comment was automatically generated by Dr. CI and updates every 15 minutes.