Prepare optional clustered PCIe collectives

docs/model-preparation.md

@@ -104,10 +104,10 @@ patches before starting their process:
 uv run --locked --extra prepare python -m ignition.patches
 ```
 
-`python scripts/setup_preparation.py` updates FlashInfer to the tip of its `main`
-branch, syncs preparation/test dependencies, rebuilds FlashInfer and Heliopause,
-and applies the bundled patches before exiting. It updates `uv.lock` when the
-branch advances.
+`python scripts/setup_preparation.py` resolves the FlashInfer revision declared
+in `pyproject.toml`, syncs preparation/test dependencies, rebuilds FlashInfer and
+Heliopause, and applies the bundled patches before exiting. It updates `uv.lock`
+when the declared dependency resolves differently.
 
 Build dependencies include a C/C++ compiler, CMake, Ninja, pkg-config, Protobuf's
 `protoc`, OpenSSL headers, Python headers, and liburing headers. On Ubuntu:
@@ -140,10 +140,10 @@ together. Refresh and commit `helios/Cargo.lock`, `flare/Cargo.lock`, and
 and Heliopause against the pinned revision. Compare rendered prompts and token
 IDs with upstream when changing encoder behavior.
 
-The preparation dependency tracks the `main` branch of
-`local-inference-lab/flashinfer`; `uv.lock` pins the resolved commit. The setup
-script updates it automatically. To update the lock alone, run
-`uv lock --upgrade-package flashinfer-python`, then sync the environment.
+The preparation dependency pins a revision of `local-inference-lab/flashinfer`
+in `pyproject.toml`; `uv.lock` records the same commit. To update FlashInfer,
+change that revision and run `uv lock --upgrade-package flashinfer-python`,
+then sync the environment.
 For local FlashInfer development, install the sibling
 checkout explicitly with `uv pip install --python .venv/bin/python --no-deps -e ../flashinfer`;
 `uv sync --locked --extra prepare --extra test` restores the Git dependency.

ignition/components/collectives.py

@@ -143,10 +143,15 @@ def prepare(state):
                              inplace=source.data_ptr() == output.data_ptr())
                 self.runtime_profiles = tuple((mode, active) for mode, counts in row["profiles"].items() for active in counts)
                 self.runtime_cache_policy = ladder["cache_policy"]
-                # One active CTA is legal at every payload; runtime selects from the supplied profiles.
+                cluster_size = row.get("cluster_size", 1)
+                mode, active = self.runtime_profiles[0]
+                # Phased modes need a complete cluster, so use a declared clustered profile.
+                initial_active = active if cluster_size > 1 else 1
                 plan_options = dict(joint_only=True, override=CollectiveConfig(
                     implementation=row["implementation"], threads=row["threads"], blocks=row["blocks"],
-                    submode=self.runtime_profiles[0][0], active_blocks=1))
+                    submode=mode, active_blocks=initial_active))
+                if cluster_size > 1:
+                    plan_options["joint_cluster_size"] = cluster_size
             self.plan = runtime.plan(source, out=output, **plan_options)
             group = runtime.exchange_group
 

ignition/preparation/collective_ladder.py

@@ -12,6 +12,10 @@
     "joint_latency": ("full", "payload_rsag", "pull"),
     "joint_bandwidth": ("payload_rsag", "bulk_push", "bulk_pull", "pull"),
 }
+CLUSTER_MODES = {
+    "joint_latency": ("full_phased", "payload_rsag_phased"),
+    "joint_bandwidth": ("payload_rsag_phased",),
+}
 GATHER_MODES = {
     "gather_joint": ("barrier", "word_exchange"),
     "gather_selected_barrier": ("barrier",),
@@ -89,7 +93,8 @@ def require(condition, message):
     require(isinstance(data["buckets"], list) and bool(data["buckets"]), "buckets must be nonempty")
     previous = data["min_bytes"] - 16
     for row in data["buckets"]:
-        require(isinstance(row, dict) and set(row) == {"max_bytes", "implementation", "threads", "blocks", "profiles"},
+        required = {"max_bytes", "implementation", "threads", "blocks", "profiles"}
+        require(isinstance(row, dict) and required <= set(row) <= required | {"cluster_size"},
                 "invalid bucket fields")
         upper = row["max_bytes"]
         require(type(upper) is int and upper > previous and upper % 16 == 0,
@@ -98,12 +103,23 @@ def require(condition, message):
         require(isinstance(row["implementation"], str) and row["implementation"] in MODES, "unknown joint kernel")
         require(type(row["threads"]) is int and row["threads"] in (32, 64, 128, 256, 512), "invalid threads")
         require(type(row["blocks"]) is int and 1 <= row["blocks"] <= 64, "invalid physical grid")
+        cluster_size = row.get("cluster_size", 1)
+        require(type(cluster_size) is int and cluster_size in (1, 8), "cluster_size must be 1 or 8")
+        if cluster_size == 8:
+            require(tuple(data["compute_capability"]) in ((12, 0), (12, 1)),
+                    "clustered joint reductions require SM120 or SM121")
+            require(row["threads"] == 512 and row["blocks"] == 64,
+                    "clustered joint reductions require 512 threads and 64 physical CTAs")
         require(isinstance(row["profiles"], dict) and bool(row["profiles"]), "profiles must be nonempty")
         for mode, counts in row["profiles"].items():
-            require(mode in MODES[row["implementation"]], "mode does not belong to joint kernel")
+            phased = mode in CLUSTER_MODES[row["implementation"]]
+            require(mode in MODES[row["implementation"]] or (cluster_size == 8 and phased),
+                    "mode does not belong to the prepared joint kernel")
             require(isinstance(counts, list) and bool(counts)
                     and all(type(v) is int and 1 <= v <= row["blocks"] for v in counts), "invalid active CTA counts")
             require(len(set(counts)) == len(counts), "duplicate active CTA counts")
+            if phased:
+                require(counts == [cluster_size], "phased profiles require exactly one complete active cluster")
     _validate_gathers(data["all_gather"], data["world_size"])
     return data
 

ignition/preparation/ladders/README.md

@@ -77,6 +77,55 @@ validate exact payload bytes and select among the supplied profiles. Gathers
 with equal byte sizes use the same recommendations regardless of row width.
 RoCE gathers and reduction-only ladders retain their preparation policy.
 
+### Thread-block clusters
+
+A bucket may set `"cluster_size": 8` to launch eight-CTA CUDA thread-block
+clusters. Omitting the field, or setting it to `1`, retains an ordinary launch.
+Clustered buckets require SM120 or SM121 with at least 64 SMs, 512 threads per
+CTA, and 64 physical CTAs. Preparation checks that the compiled function can
+keep eight clusters resident. Ordinary profiles retain their active-CTA counts.
+Phased full all-reduce is named `full_phased`; phased reduce-scatter followed
+by all-gather is named `payload_rsag_phased`. The latency family admits both;
+the bandwidth family admits `payload_rsag_phased`. Each phased profile must select
+exactly eight active CTAs.
+
+Phased execution uses the first complete cluster to write one peer destination
+at a time. All its CTAs participate in each phase barrier; other clusters return
+without changing collective state. System-scope publication fences remain
+required for peer-visible writes. This uses existing kernel launch attributes
+in execution ABI 5 and does not require cooperative launch.
+
+`gb202_pcie_tp4_clustered.json` adds phased RSAG to 36 compatible buckets used
+by packed DeepSeek V4.1 TP4 verification graphs with eight, sixteen, or thirty-two
+requests, seven draft positions, and up to 128 target rows. It retains every
+ordinary candidate, kernel family, thread count, and physical grid from
+`gb202_pcie_tp4.json`. Its 256-thread buckets use ordinary launches.
+`gb202_pcie_tp4_clustered_full.json` adds one further choice: phased Full in the
+bucket containing 80 KiB. This also covers the target reductions for one request
+verified at eight rows.
+
+The compiled cluster geometry also applies when startup selects an ordinary
+profile; startup selection cannot disable clustering. Cluster scheduling can
+affect those profiles, so retaining every ordinary candidate does not guarantee
+unchanged latency. Qualify these catalogs on the intended GPU topology and
+complete serving workload; component improvements do not establish a throughput gain.
+`gb202_pcie_tp4.json` remains the ordinary catalog.
+
+Select the catalog with both phased modes during export:
+
+```sh
+.venv/bin/ignition export deepseek-v4.1-flash --gpus 0,1,2,3 \
+  --collective-tuning runtime \
+  --collective-ladder ignition/preparation/ladders/gb202_pcie_tp4_clustered_full.json
+```
+
+Use `gb202_pcie_tp4_clustered.json` in the same flag to exclude the phased Full
+choice at 80 KiB. `--collective-tuning static` uses preparation-time autotuning
+and accepts no ladder; these catalogs require the explicit `runtime` policy.
+Changing the catalog requires package export because cluster geometry is compiled
+into its kernels and captured launch attributes. Existing packages retain their
+prepared geometry and profiles.
+
 ## GB202 PCIe TP4 recommendations
 
 `gb202_pcie_tp4.json` contains 210 buckets from 16 bytes through 8 MiB.