[cuda backend] support multi-SM AOTI PTEs (pytorch#22198)

### Summary

Add one current `ETCUDAFQN0` metadata format that stores target-specific
CUDA AOTI shared libraries with one shared FQN weight manifest. At
runtime, CUDA first selects an exact-SM regular variant. If no exact
regular variant exists, it may use the PTX from one explicitly
designated fallback PTE when that PTX is compatible with the current
GPU. The fallback is no longer inferred from the lowest regular target
SM.

Add a PTE/PTD merge tool with two input classes:

- One or more regular PTEs provide exact-SM native variants. Their
merged metadata sets `ptx_compute=0`, so they are never considered as
PTX fallbacks. Regular exports should use
`CompileSpec("cuda_include_ptx", b"OFF")` to avoid embedding redundant
PTX.
- At most one fallback PTE provides the PTX-capable variant. It is
marked `fallback_only`, is never selected as an exact native variant,
and is considered only when the merged PTE has no regular variant for
the current SM.

The merge validates the ExecuTorch program, CUDA delegate layout,
non-CUDA named data, weight metadata, and streamed PTD weight hashes.
CUDA PTE merging is NVIDIA-only and explicitly rejects ROCm.

### Merged artifact layout

For an A100 (`sm80`) regular export, an RTX 5090 (`sm120`) regular
export, and one PTX fallback export, the merged artifact is structured
as follows:

```text
merged.pte
├── program and CUDA delegates
│   └── `ETCUDAFQN0` CUDA variant metadata (per delegate)
│       ├── regular: target_sm=80,  ptx_compute=0 -> sm80 source SO key
│       ├── regular: target_sm=120, ptx_compute=0 -> sm120 source SO key
│       └── fallback_only: target_sm=<source SM>,
│           ptx_compute=<fallback compute capability> -> fallback source SO key
├── named data
│   ├── sm80 source SO key -> actual AOTI SO containing the prebuilt sm80 cubin
│   ├── sm120 source SO key -> actual AOTI SO containing the prebuilt sm120 cubin
│   └── fallback source SO key -> actual PTX-bearing fallback AOTI SO
└── one shared FQN weight manifest

merged.ptd
└── shared external tensor data reused from the regular base export
```

The tool copies each AOTI SO atomically rather than rewriting its
ELF/fatbin. The regular/fallback distinction is encoded in the merged
metadata: only regular variants participate in exact-SM native
selection, while the fallback SO contributes PTX only to fallback
selection. After merging, the CLI prints a provenance table mapping
every stored SM cubin and the fallback PTX entry to its source PTE.

This PR was authored with OpenAI Codex.

### Test plan

The primary validation is a real cross-hardware CUDA CI pipeline
spanning A100 (`sm80`) and A10G (`sm86`):

1. Export an A100 native PTE and an A10G native PTE independently on
their respective GPUs, both with PTX disabled and isolated Inductor
caches.
2. Export a separate PTX fallback PTE on A100 with the A10G
shared-memory constraint.
3. Verify byte-for-byte that the A100 native PTD, A10G native PTD,
fallback PTD, and final merged PTD are identical.
4. On A10G, verify the A10G native PTE succeeds, the A100 native-only
PTE is rejected with no compatible native/PTX variant, and the merged
PTE succeeds.
5. On A100, verify the A100 native PTE succeeds, the A10G native-only
PTE is rejected with no compatible native/PTX variant, and the same
merged PTE succeeds.

Additional Python coverage exercises the single `ETCUDAFQN0` metadata
format, untargeted ROCm metadata, explicit fallback-only encoding,
PTE/PTD merging, shared-weight validation, duplicate regular SM
rejection, provenance reporting, and ROCm merge rejection. C++ coverage
exercises metadata parsing, exact regular selection, PTX fallback
selection, incompatibility, and duplicate targets.

Locally ran the 17 targeted CUDA metadata/merge tests, Python
lint/format checks, workflow YAML parsing, and `git diff --check`.

Author: Gasoonjia

.github/workflows/cuda.yml

@@ -160,7 +160,7 @@ jobs:
         done
 
   export-cuda-target-smem-cross-arch:
-    name: export-cuda-target-smem-cross-arch-a100
+    name: export-cuda-multi-arch-a100
     needs: [changed-files, run-decision]
     if: |
       contains(needs.changed-files.outputs.changed-files, 'backends/cuda') ||
@@ -178,15 +178,13 @@ jobs:
       gpu-arch-version: "13.0"
       use-custom-docker-registry: false
       submodules: recursive
-      upload-artifact: cuda-target-smem-cross-arch
+      upload-artifact: cuda-multi-arch-a100
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       script: |
         set -eux
 
         PYTHON_EXECUTABLE=python ./install_executorch.sh
         export LD_LIBRARY_PATH="/opt/conda/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
-        # A10G (sm_86) supports 99 KiB of opt-in shared memory per block.
-        export ET_CUDA_TARGET_SMEM_BYTES=101376
         # The A100 exporter has an AVX-512 host CPU while the A10G runner has
         # AVX2. Keep the generated AOTI host wrapper compatible with both.
         export ATEN_CPU_CAPABILITY=avx2
@@ -199,16 +197,33 @@ jobs:
         assert capability == (8, 0), f"Expected A100 sm_80 exporter, got {capability}"
         PY
 
+        # Export an sm80 native PTE with no PTX. This artifact must fail on the
+        # A10G because its metadata has neither an sm86 native variant nor PTX.
+        model_dir="${RUNNER_ARTIFACT_DIR}/native-sm80/linear"
+        mkdir -p "${model_dir}"
+        TORCHINDUCTOR_CACHE_DIR="${RUNNER_TEMP}/inductor-native-sm80" \
+          python -m examples.cuda.scripts.export \
+            --model_name=linear \
+            --output_dir="${model_dir}" \
+            --cuda_include_ptx=OFF \
+            --seed=0
+
+        # Export the explicit PTX fallback separately. A10G supports 99 KiB of
+        # opt-in shared memory per block, so constrain the A100 export to it.
         for model in linear sdpa; do
-          model_dir="${RUNNER_ARTIFACT_DIR}/${model}"
+          model_dir="${RUNNER_ARTIFACT_DIR}/fallback-sm80/${model}"
           mkdir -p "${model_dir}"
-          python -m examples.cuda.scripts.export \
+          ET_CUDA_TARGET_SMEM_BYTES=101376 \
+          TORCHINDUCTOR_CACHE_DIR="${RUNNER_TEMP}/inductor-fallback-sm80-${model}" \
+            python -m examples.cuda.scripts.export \
             --model_name="${model}" \
-            --output_dir "${model_dir}"
+            --output_dir="${model_dir}" \
+            --cuda_include_ptx=ON \
+            --seed=0
         done
 
   test-cuda-target-smem-cross-arch:
-    name: test-cuda-target-smem-cross-arch-a10g
+    name: merge-and-test-cuda-multi-arch-a10g
     needs: [export-cuda-target-smem-cross-arch]
     if: needs.export-cuda-target-smem-cross-arch.result == 'success'
     uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main
@@ -222,7 +237,8 @@ jobs:
       gpu-arch-version: "13.0"
       use-custom-docker-registry: false
       submodules: recursive
-      download-artifact: cuda-target-smem-cross-arch
+      download-artifact: cuda-multi-arch-a100
+      upload-artifact: cuda-multi-arch-merged
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       script: |
         set -eux
@@ -243,20 +259,141 @@ jobs:
         )
         PY
 
+        # Independently export an sm86 native PTE with no PTX. A fixed seed
+        # makes its external weights byte-identical to both A100 exports.
+        model_dir="${RUNNER_ARTIFACT_DIR}/native-sm86/linear"
+        mkdir -p "${model_dir}"
+        TORCHINDUCTOR_CACHE_DIR="${RUNNER_TEMP}/inductor-native-sm86" \
+          python -m examples.cuda.scripts.export \
+            --model_name=linear \
+            --output_dir="${model_dir}" \
+            --cuda_include_ptx=OFF \
+            --seed=0
+
+        sm80_ptd="${RUNNER_ARTIFACT_DIR}/native-sm80/linear/aoti_cuda_blob.ptd"
+        sm86_ptd="${RUNNER_ARTIFACT_DIR}/native-sm86/linear/aoti_cuda_blob.ptd"
+        fallback_ptd="${RUNNER_ARTIFACT_DIR}/fallback-sm80/linear/aoti_cuda_blob.ptd"
+        cmp "${sm80_ptd}" "${sm86_ptd}"
+        cmp "${sm80_ptd}" "${fallback_ptd}"
+        sha256sum "${sm80_ptd}" "${sm86_ptd}" "${fallback_ptd}"
+
+        merged_dir="${RUNNER_ARTIFACT_DIR}/merged"
+        mkdir -p "${merged_dir}"
+        python -m executorch.backends.cuda.merge_ptes \
+          --input-pte "${RUNNER_ARTIFACT_DIR}/native-sm80/linear/linear.pte" \
+          --input-pte "${RUNNER_ARTIFACT_DIR}/native-sm86/linear/linear.pte" \
+          --input-ptd "${sm80_ptd}" \
+          --input-ptd "${sm86_ptd}" \
+          --fallback-pte "${RUNNER_ARTIFACT_DIR}/fallback-sm80/linear/linear.pte" \
+          --fallback-ptd "${fallback_ptd}" \
+          --output-pte "${merged_dir}/linear.pte" \
+          --output-ptd "${merged_dir}/aoti_cuda_blob.ptd" \
+          | tee "${merged_dir}/provenance.txt"
+        cmp "${sm80_ptd}" "${merged_dir}/aoti_cuda_blob.ptd"
+
         cmake -DCMAKE_BUILD_TYPE=Release \
             -DEXECUTORCH_BUILD_CUDA=ON \
             -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
+            -DEXECUTORCH_ENABLE_LOGGING=ON \
             -DPYTHON_EXECUTABLE=python \
             -Bcmake-out .
         cmake --build cmake-out --target executor_runner -j4
 
         for model in linear sdpa; do
-          model_dir="${RUNNER_ARTIFACT_DIR}/${model}"
+          model_dir="${RUNNER_ARTIFACT_DIR}/fallback-sm80/${model}"
           ./cmake-out/executor_runner \
             --model_path "${model_dir}/${model}.pte" \
             --data_path "${model_dir}/aoti_cuda_blob.ptd"
         done
 
+        # The local native PTE and merged PTE must run on sm86.
+        ./cmake-out/executor_runner \
+          --model_path "${RUNNER_ARTIFACT_DIR}/native-sm86/linear/linear.pte" \
+          --data_path "${sm86_ptd}"
+        ./cmake-out/executor_runner \
+          --model_path "${merged_dir}/linear.pte" \
+          --data_path "${merged_dir}/aoti_cuda_blob.ptd"
+
+        # The sm80 native-only PTE must be rejected on sm86.
+        set +e
+        ./cmake-out/executor_runner \
+          --model_path "${RUNNER_ARTIFACT_DIR}/native-sm80/linear/linear.pte" \
+          --data_path "${sm80_ptd}" \
+          >"${RUNNER_ARTIFACT_DIR}/native-sm80-on-sm86.log" 2>&1
+        native_status=$?
+        set -e
+        cat "${RUNNER_ARTIFACT_DIR}/native-sm80-on-sm86.log"
+        test "${native_status}" -ne 0
+        grep -q "no native or PTX variant compatible with sm86" \
+          "${RUNNER_ARTIFACT_DIR}/native-sm80-on-sm86.log"
+
+  test-cuda-merged-pte-a100:
+    name: test-cuda-merged-pte-a100
+    needs: [test-cuda-target-smem-cross-arch]
+    if: needs.test-cuda-target-smem-cross-arch.result == 'success'
+    uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      timeout: 90
+      runner: mt-l-x86iavx512-11-125-a100
+      gpu-arch-type: cuda
+      gpu-arch-version: "13.0"
+      use-custom-docker-registry: false
+      submodules: recursive
+      download-artifact: cuda-multi-arch-merged
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      script: |
+        set -eux
+
+        PYTHON_EXECUTABLE=python ./install_executorch.sh
+        export LD_LIBRARY_PATH="/opt/conda/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
+
+        python - <<'PY'
+        import torch
+
+        capability = torch.cuda.get_device_capability()
+        assert capability == (8, 0), f"Expected A100 sm_80 runner, got {capability}"
+        PY
+
+        cmake -DCMAKE_BUILD_TYPE=Release \
+            -DEXECUTORCH_BUILD_CUDA=ON \
+            -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
+            -DEXECUTORCH_ENABLE_LOGGING=ON \
+            -DPYTHON_EXECUTABLE=python \
+            -Bcmake-out .
+        cmake --build cmake-out --target executor_runner -j4
+
+        sm80_ptd="${RUNNER_ARTIFACT_DIR}/native-sm80/linear/aoti_cuda_blob.ptd"
+        sm86_ptd="${RUNNER_ARTIFACT_DIR}/native-sm86/linear/aoti_cuda_blob.ptd"
+        fallback_ptd="${RUNNER_ARTIFACT_DIR}/fallback-sm80/linear/aoti_cuda_blob.ptd"
+        merged_ptd="${RUNNER_ARTIFACT_DIR}/merged/aoti_cuda_blob.ptd"
+        cmp "${sm80_ptd}" "${sm86_ptd}"
+        cmp "${sm80_ptd}" "${fallback_ptd}"
+        cmp "${sm80_ptd}" "${merged_ptd}"
+
+        # The local native PTE and merged PTE must run on sm80.
+        ./cmake-out/executor_runner \
+          --model_path "${RUNNER_ARTIFACT_DIR}/native-sm80/linear/linear.pte" \
+          --data_path "${sm80_ptd}"
+        ./cmake-out/executor_runner \
+          --model_path "${RUNNER_ARTIFACT_DIR}/merged/linear.pte" \
+          --data_path "${merged_ptd}"
+
+        # The sm86 native-only PTE must be rejected on sm80.
+        set +e
+        ./cmake-out/executor_runner \
+          --model_path "${RUNNER_ARTIFACT_DIR}/native-sm86/linear/linear.pte" \
+          --data_path "${sm86_ptd}" \
+          >"${RUNNER_ARTIFACT_DIR}/native-sm86-on-sm80.log" 2>&1
+        native_status=$?
+        set -e
+        cat "${RUNNER_ARTIFACT_DIR}/native-sm86-on-sm80.log"
+        test "${native_status}" -ne 0
+        grep -q "no native or PTX variant compatible with sm80" \
+          "${RUNNER_ARTIFACT_DIR}/native-sm86-on-sm80.log"
+
   unittest-cuda:
     name: unittest-cuda
     needs: [changed-files, run-decision]
@@ -370,6 +507,11 @@ jobs:
         conda install -y -c conda-forge 'libstdcxx-ng>=12'
         export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH
 
+        python -m pytest \
+          backends/cuda/tests/test_cuda_weight_metadata.py \
+          backends/cuda/tests/test_merge_ptes.py \
+          -v -o "addopts="
+
         cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON
         cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state test_cuda_weight_cache -j$(nproc)
         ctest --test-dir cmake-out -R test_cuda_allocator --output-on-failure -V

backends/cuda/BUCK

@@ -109,6 +109,22 @@ fbcode_target(
     ],
 )
 
+fbcode_target(
+    _kind = runtime.python_library,
+    name = "merge_ptes",
+    srcs = [
+        "merge_ptes.py",
+    ],
+    visibility = ["PUBLIC"],
+    deps = [
+        ":cuda_backend",
+        "//caffe2:torch",
+        "//executorch/exir/_serialize:lib",
+        "//executorch/exir:schema",
+        "//executorch/extension/flat_tensor/serialize:serialize",
+    ],
+)
+
 fbcode_target(
     _kind = runtime.python_library,
     name = "cuda_partitioner",

backends/cuda/cuda_backend.py

@@ -519,6 +519,14 @@ def _setup_cuda_environment_for_fatbin() -> bool:
         except Exception:
             return False
 
+    @classmethod
+    def _should_include_ptx(cls, compile_specs: List[CompileSpec]) -> bool:
+        include_ptx = True
+        for spec in compile_specs:
+            if spec.key == "cuda_include_ptx":
+                include_ptx = _on_off_compile_spec_value(spec)
+        return include_ptx and cls._setup_cuda_environment_for_fatbin()
+
     @classmethod
     def save_data_externally(cls) -> bool:
         """
@@ -539,7 +547,29 @@ def _preprocess_with_weight_collector(
             result = super().preprocess(edge_program, compile_specs)
         if capture.artifact is None:
             raise RuntimeError("CUDA AOTI did not return a structured Weights output")
-        collector.add_preprocess_result(result, capture.artifact, cls.get_device_name())
+        target_sm = None
+        ptx_compute = 0
+        if torch.version.hip is None:
+            from torch._inductor.codegen.cuda.compile_utils import (
+                _nvcc_arch_as_compile_option,
+            )
+
+            compiled_arch = _nvcc_arch_as_compile_option()
+            target_arch = compiled_arch.removesuffix("a")
+            if not target_arch.isdigit():
+                raise RuntimeError(f"Unsupported CUDA architecture: {compiled_arch}")
+            target_sm = int(target_arch)
+            # Architecture-accelerated PTX targets such as compute_90a and
+            # compute_120a are not forward-compatible fallback images.
+            if cls._should_include_ptx(compile_specs) and compiled_arch.isdigit():
+                ptx_compute = int(compiled_arch)
+        collector.add_preprocess_result(
+            result,
+            capture.artifact,
+            cls.get_device_name(),
+            target_sm=target_sm,
+            ptx_compute=ptx_compute,
+        )
         return result
 
     @classmethod
@@ -762,7 +792,7 @@ def get_aoti_compile_options(
 
         # Configure CUDA environment variables based on detected version
 
-        emit_multi_arch_kernel = CudaBackend._setup_cuda_environment_for_fatbin()
+        emit_multi_arch_kernel = cls._should_include_ptx(compile_specs)
         # Base options for all platforms
 
         options: Dict[str, typing.Any] = {