Respect CPU quotas when sizing quantized GEMM compiler workers

.ci/scripts/export_model_artifact.sh

@@ -716,7 +716,11 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then
       # executor drives, plus the vision encoder and forward_from_embeddings.
       OFFGRAPH_FLAG=""
       if [ "$MODE" = "dflash-image-batching" ]; then
-        OFFGRAPH_FLAG="--use-offgraph-kv-cache"
+        # Every static forward_{N} and draft_forward_{N} is its own AOTI
+        # compile and autotune, which the 11-vCPU runner cannot fit in the
+        # job timeout at the default sizes; these still cover static and
+        # dynamic target steps and DFlash drafting.
+        OFFGRAPH_FLAG="--use-offgraph-kv-cache --forward-sizes 1,4 --max-draft-owners 1"
       fi
       EXPORT_START_SECONDS=$SECONDS
       TMPDIR="$INDUCTOR_TMPDIR" \

.ci/scripts/tests/test_muse_glimmer_batching_ci.py

@@ -129,6 +129,26 @@ def test_smoke_cmake_is_opt_in_and_reuses_runtime(self):
         self.assertNotIn("BENCHMARK", smoke)
         self.assertNotIn("batching_bench", smoke)
 
+    def test_full_e2e_runs_only_on_main_pushes(self):
+        workflow = (ROOT / ".github/workflows/cuda.yml").read_text()
+        e2e_if = (
+            workflow.split("  test-muse-glimmer-cuda-e2e:", 1)[1]
+            .split("    if: |\n", 1)[1]
+            .split("    uses:", 1)[0]
+            .strip()
+        )
+        self.assertTrue(
+            e2e_if.startswith(
+                "github.event_name == 'push' && github.ref == 'refs/heads/main' && ("
+            )
+        )
+        self.assertTrue(e2e_if.endswith(")"))
+        self.assertEqual(e2e_if.count("("), e2e_if.count(")"))
+        unit_if = workflow.split("  unittest-cuda-muse-glimmer-batching:", 1)[1].split(
+            "    uses:", 1
+        )[0]
+        self.assertNotIn("github.event_name", unit_if)
+
     def test_owned_shell_and_unit_workflow_syntax(self):
         for script in ("test_model_e2e.sh", "build_muse_glimmer_batching_smoke.sh"):
             result = subprocess.run(["bash", "-n", str(SCRIPTS / script)])

.github/workflows/cuda.yml

@@ -1197,7 +1197,11 @@ jobs:
   test-muse-glimmer-cuda-e2e:
     name: test-muse-glimmer-cuda-e2e-${{ matrix.variant }}-${{ matrix.mode }}
     needs: [changed-files, run-decision]
+    # Hours of A100 time per run, so only after merge to main (not on pull
+    # requests, ciflow tags or manual dispatch); pull requests get the
+    # unittest-cuda-* jobs above.
     if: |
+      (
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/build_muse_glimmer_batching_smoke.sh') ||
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/muse_glimmer_batching_smoke.cpp') ||
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/muse_glimmer_batching_workload.h') ||
@@ -1212,7 +1216,7 @@ jobs:
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/export_model_artifact.sh') ||
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') ||
       contains(needs.changed-files.outputs.changed-files, 'examples/models/muse-glimmer') ||
-      needs.run-decision.outputs.is-full-run == 'true'
+      needs.run-decision.outputs.is-full-run == 'true')
     uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main
     permissions:
       id-token: write

backends/cuda/BUCK

@@ -3,6 +3,18 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 oncall("executorch")
 
+fbcode_target(
+    _kind = runtime.python_library,
+    name = "cpu_count",
+    srcs = [
+        "utils/__init__.py",
+        "utils/cpu_count.py",
+    ],
+    visibility = [
+        "//executorch/backends/cuda/...",
+    ],
+)
+
 fbcode_target(
     _kind = runtime.python_library,
     name = "coalesced_int4_tensor",
@@ -188,6 +200,7 @@ fbcode_target(
     ],
     deps = [
         ":autotune",
+        ":cpu_count",
         "//caffe2:torch",
     ],
 )

backends/cuda/autotune/cuda_graph_timing.py

@@ -31,8 +31,11 @@
 captured).
 
 To keep captures succeeding:
-- each capture's graph and its memory pool are released right after timing,
-  and the allocator's cache is emptied once it holds a lot of unused memory;
+- each capture's graph is released and the allocator's cache emptied right
+  after timing. The cache is not checked first: the allocator stats keep an
+  entry for every private pool ever captured, so reading them costs more with
+  each capture, while emptying a cache that holds little more than the
+  released pool is cheap;
 - on CUDA OOM the cache is emptied and the capture retried; if that fails too,
   the tensors given as ``offload`` (the model's weights moved to the GPU for
   the compile) are parked on the CPU while the candidate is timed and restored
@@ -59,7 +62,6 @@
 _TARGET_GRAPH_US = 2000.0
 _MIN_CALLS = 4
 _MAX_CALLS = 32
-_EMPTY_CACHE_THRESHOLD_BYTES = 1 << 30
 
 
 @dataclass
@@ -116,14 +118,6 @@ def _default_cuda_bench(
     return bench.benchmark_gpu(fn, warmup=warmup, rep=rep, **kwargs)
 
 
-def _release_cached_memory_if_large() -> None:
-    if (
-        torch.cuda.memory_reserved() - torch.cuda.memory_allocated()
-        > _EMPTY_CACHE_THRESHOLD_BYTES
-    ):
-        torch.cuda.empty_cache()
-
-
 _l2_flush_buffers: dict[int, torch.Tensor] = {}
 
 
@@ -180,7 +174,7 @@ def time_with_cuda_graph_us(fn: Callable[[], Any]) -> float:
         with contextlib.suppress(RuntimeError):
             graph.reset()
         del graph
-        _release_cached_memory_if_large()
+        torch.cuda.empty_cache()
 
 
 class ParkedTensorsNotRestored(RuntimeError):

backends/cuda/cuda_backend.py

@@ -1166,7 +1166,7 @@ def materialize_weights_blob(
         blob_path = os.path.splitext(so_path)[0] + "_weights.blob"
         capture = CudaWeightCollector.current_capture()
         capture.artifact = capture.collector.materialize(
-            weights[0], os.path.dirname(blob_path), capture.device_type_for_weight
+            weights[0], os.path.dirname(blob_path), capture
         )
 
         # Keep AotiBackend's existing path contract intact. The compatibility
@@ -1545,14 +1545,30 @@ def move_program_to_device(
         check on the cache update. They are re-synthesized as zeros for the lifted
         graph and the serialized blob. This activates automatically with low-memory
         mode. Other (non-low-memory) exports use the stock pass.
+
+        Low-memory methods share their CPU weights, so the CPU source of each
+        moved tensor is recorded on the active weight capture. ``materialize``
+        then reuses a storage an earlier method already wrote from the same
+        source view. The moved storages must keep their source's bytes until
+        ``materialize`` runs.
         """
         from torch.export.passes import move_to_device_pass
 
         if not cls._is_low_memory_mode(compile_specs):
             moved = move_to_device_pass(edge_program, device)
-        else:
-            moved = _move_to_device_resize_kv(edge_program, device)
+            _track_moved_tensors(moved)
+            return moved
+        # Only parameters are recorded: low-memory compilation can hand AOTI CPU
+        # clones of buffers (_compile_time_cpu_clones).
+        sources = {
+            name: edge_program.state_dict[name]
+            for name in edge_program.graph_signature.parameters
+        }
+        moved = _move_to_device_resize_kv(edge_program, device)
         _track_moved_tensors(moved)
+        CudaWeightCollector.current_capture().record_moved_sources(
+            sources, dict(moved.state_dict)
+        )
         return moved
 
     @classmethod

backends/cuda/cuda_weight_collector.py

@@ -21,6 +21,7 @@
 from executorch.exir._serialize._named_data_store import NamedDataStore
 from executorch.exir.backend.backend_details import PreprocessResult
 from executorch.exir.tensor import scalar_type_enum
+from torch.multiprocessing.reductions import StorageWeakRef
 
 
 CUDA_AOTI_METADATA_MAGIC = b"ETCUDAFQN0"
@@ -46,7 +47,9 @@
     entries: List[CudaWeightEntry]
     storages: Dict[str, FileBackedData]
     # Storage keys an earlier method of this collector already stored with the
-    # same bytes; they are not written again.
+    # same bytes; they are not written again. A key reused because it was moved
+    # from the same CPU source view is not re-read: that relies on the source
+    # not changing between methods in ways its version counter misses.
     stored: Set[str] = field(default_factory=set)
 
 
@@ -68,11 +71,71 @@
     entries: List[CudaWeightEntry]
 
 
+def _source_identity(source: torch.Tensor) -> Tuple[Any, ...]:
+    # Bumped by in-place ops autograd tracks; writes through .data, .numpy() or
+    # the storage do not bump it.
+    return (
+        source.untyped_storage()._cdata,
+        source._version,
+        source.dtype,
+        source.storage_offset(),
+        tuple(source.shape),
+        tuple(source.stride()),
+    )
+
+
+@dataclass(frozen=True)
+class _MovedSource:
+    source: StorageWeakRef
+    identity: Tuple[Any, ...]
+    # The moved tensor's version when it was recorded; a write to it before
+    # materialize makes it be read and checked again.
+    moved_version: int
+
+
 @dataclass
 class _CudaWeightCapture:
     collector: "CudaWeightCollector"
     device_type_for_weight: Callable[[torch.Tensor], int]
     artifact: Optional[CudaWeightArtifact] = None
+    # This method's moved storage -> the CPU source it was copied from. The
+    # StorageWeakRefs (which hash and compare by address) keep both storage
+    # objects, not their data, allocated, so no other storage can take either
+    # address while this capture lives.
+    moved_sources: Dict[StorageWeakRef, "_MovedSource"] = field(default_factory=dict)
+
+    def record_moved_sources(
+        self, sources: Dict[str, torch.Tensor], moved: Dict[str, torch.Tensor]
+    ) -> None:
+        """Record the CPU tensor each tensor of this method was moved from.
+
+        Must run inside ``capture()`` before ``materialize``. A matched storage
+        is reused without being read, so nothing may write the moved tensors
+        between this call and ``materialize`` except through in-place ops
+        their version counter records.
+        """
+        self.moved_sources = {}
+        for name, tensor in moved.items():
+            source = sources.get(name)
+            if (
+                type(source) in (torch.Tensor, torch.nn.Parameter)
+                and type(tensor) in (torch.Tensor, torch.nn.Parameter)
+                and tensor.is_cuda
+                and source.device.type == "cpu"
+                # Weights loaded under inference_mode can change without a
+                # version bump.
+                and not source.is_inference()
+                and source.untyped_storage().nbytes() > 0
+                # Low-memory KV placeholders are emptied on the device.
+                and tensor.untyped_storage().nbytes() > 0
+            ):
+                self.moved_sources[StorageWeakRef(tensor.untyped_storage())] = (
+                    _MovedSource(
+                        source=StorageWeakRef(source.untyped_storage()),
+                        identity=_source_identity(source),
+                        moved_version=tensor._version,
+                    )
+                )
 
 
 def trim_host_memory() -> None:
@@ -384,13 +447,20 @@
     def __init__(self) -> None:
         self._store = NamedDataStore()
         self._results: List[PreprocessResult] = []
-        # The CUDA storages already written to the store, by storage key. Methods
-        # share their weights, so a later method's identical storage is compared
-        # with the first one on the device and not written again. The weights
-        # must not change until finish() releases them.
+        # The CUDA storages already written to the store, by storage key. A later
+        # method's identical storage is compared with the first one on the
+        # device and not written again while the first one is still allocated.
         self._stored_storages: Dict[str, torch.Tensor] = {}
         # The key each storage was written under, which ends in its SHA-256.
         self._content_keys: Dict[str, str] = {}
+        # The CPU source view each storage key was written from. When methods
+        # share their CPU weights (low_memory_mode), each moves the same source
+        # to the GPU and frees the copy afterwards, so a later method's copy of
+        # the same source view is recognized by its source and not written
+        # again. The StorageWeakRef keeps the source's storage object allocated,
+        # so its address, part of the identity, is not reused; once its data is
+        # freed, a later live source takes the key over.
+        self._written_sources: Dict[str, _MovedSource] = {}
 
     @contextlib.contextmanager
     def capture(
@@ -413,13 +483,18 @@
             )
         return capture
 
     def materialize(
         self,
         weights: Any,
         directory: str,
-        device_type_for_weight: Callable[[torch.Tensor], int],
+        capture: _CudaWeightCapture,
     ) -> CudaWeightArtifact:
-        """Turn AOTI ``Weights`` into one independently named blob per FQN."""
+        """Turn AOTI ``Weights`` into one independently named blob per FQN.
+
+        ``capture`` holds the sources ``move_program_to_device`` recorded for
+        this method; storages copied from an already written source view are
+        reused without being read.
+        """
         trim_host_memory()
         entries: List[CudaWeightEntry] = []
         storages: Dict[str, FileBackedData] = {}
@@ -440,7 +515,7 @@
                 and tuple(tensor.stride()) == tuple(getattr(properties, "stride", ()))
             )
             del storage
-            device_type = device_type_for_weight(tensor)
+            device_type = capture.device_type_for_weight(tensor)
             expected_storage_nbytes = int(
                 getattr(properties, "storage_size", None) or 0
             )
@@ -460,9 +535,24 @@
                     raise RuntimeError(f"Duplicate CUDA FQN weight key for {fqn!r}")
                 # Library-local constants are keyed per shared library, so
                 # only FQN weights can match an earlier method's storage.
-                shareable = tensor.is_cuda and not _is_aoti_library_local_fqn(fqn)
+                library_local = _is_aoti_library_local_fqn(fqn)
+                shareable = tensor.is_cuda and not library_local
                 first = self._stored_storages.get(storage_key)
+                moved = (
+                    capture.moved_sources.get(StorageWeakRef(tensor.untyped_storage()))
+                    if shareable
+                    else None
+                )
+                if moved is not None and tensor._version != moved.moved_version:
+                    moved = None
+                written = self._written_sources.get(storage_key)
                 if (
+                    moved is not None
+                    and written is not None
+                    and written.identity == moved.identity
+                ):
+                    stored.add(storage_key)
+                elif (
                     shareable
                     and first is not None
                     and _same_cuda_storage(first, tensor)
@@ -472,6 +562,10 @@
                     storages[storage_key] = _materialize_storage(tensor, directory)
                     if shareable:
                         self._stored_storages.setdefault(storage_key, tensor)
+                    if moved is not None and (
+                        written is None or written.source.expired()
+                    ):
+                        self._written_sources[storage_key] = moved
 
             sizes = tuple(
                 int(size) for size in getattr(properties, "shape", tensor.shape)
@@ -639,3 +733,4 @@
             result.data_store_output = shared_output
         self._stored_storages.clear()
         self._content_keys.clear()
+        self._written_sources.clear()

backends/cuda/tests/autotune_test_utils.py

@@ -17,6 +17,7 @@
 import torch
 import triton
 from executorch.backends.cuda.autotune import cuda_graph_timing as cgt
+from executorch.backends.cuda.utils.cpu_count import available_cpu_count
 from torch._inductor.runtime.triton_heuristics import CachingAutotuner
 from torch._inductor.select_algorithm import AlgorithmSelectorCache
 
@@ -58,13 +59,13 @@ def _spin(stop_at: float) -> None:
 
 
 class SaturatedCpu:
-    """One busy-loop process per core; can be paused to measure a reference."""
+    """One busy-loop process per usable CPU; can be paused to measure a reference."""
 
     def __init__(self) -> None:
         stop_at = time.time() + 3600
         self.workers = [
             multiprocessing.Process(target=_spin, args=(stop_at,), daemon=True)
-            for _ in range(os.cpu_count() or 1)
+            for _ in range(available_cpu_count())
         ]
 
     def __enter__(self) -> "SaturatedCpu":