Skip to content

Prepare optional clustered PCIe collectives - #174

Merged
lukealonso merged 2 commits into
mainfrom
astra/clustered-collective-preparation-main
Nov 3, 2026
Merged

lukealonso merged 2 commits into
mainfrom
astra/clustered-collective-preparation-main

Conversation

@voipmonitor

@voipmonitor voipmonitor commented Nov 2, 2026 •

Copy link
Copy Markdown
Collaborator

Phased PCIe writes send to one peer destination at a time. This lets Ignition prepare those reductions using CUDA clusters of eight thread blocks, alongside the existing protocols, so startup can choose a profile on the serving machine.

The launch geometry is fixed before capture. Ordinary profiles in a clustered bucket therefore also run with cluster scheduling; keeping them in the menu does not guarantee unchanged performance. The ordinary catalog remains available and is not replaced.

Prepare the clustered catalog

Requires FlashInfer #70. The preparation dependency and uv.lock pin its official commit 560fef16648e4953bf39b701a226c322a0ab6c7c, so this branch can be prepared before that PR is merged.

After the repository's normal environment setup:

.venv/bin/python -m ignition export deepseek-v4.1-flash \
  --gpus 0-3 --max-concurrency 32 --image-concurrency 1 \
  --mtp-draft-tokens 7 --speculative-token-capacity 128 \
  --speculative-graph-budget 64 --prefill-chunk-size 4096 \
  --mla-compute-mode bf16 --expert-scale-compression native \
  --draft-head-quantization nvfp4 --markov-head-quantization fp8 \
  --vocabulary-parallel-sampling --autotune-effort default \
  --compile-workers 64 --collective-tuning runtime \
  --collective-ladder ignition/preparation/ladders/gb202_pcie_tp4_clustered_full.json

./flare/target/release/flare run deepseek-ai/DeepSeek-V4.1-Flash-1017 \
  --gpus 0-3 --max-concurrency 32 --mtp-draft-tokens 7 \
  --mtp-max-verification-token-budget 128 \
  --mtp-adaptive-verification-alpha 0.6 --scheduler-policy balanced \
  --embedding-table-storage ram --prefix-cache-memory 64G --port 8000

When a compatible weight cache already exists, add --reuse-weights to export.

Catalog Available choices
gb202_pcie_tp4.json Existing ordinary profiles
gb202_pcie_tp4_clustered.json Adds phased RSAG in 36 compatible buckets
gb202_pcie_tp4_clustered_full.json Also adds phased full all-reduce in the bucket containing 80 KiB; this is the measured catalog

The optional cluster_size: 8 field requires 512 threads, 64 physical blocks, SM120/SM121 and at least 64 SMs. Phased profiles use exactly one active cluster. b12x checks compiled residency. All original ordinary profiles and gather recommendations remain available; 256-thread buckets stay ordinary.

Compatibility: no Helios, scheduler, execution ABI 5 or prepared contract 3 changes. Existing packages retain their captured kernels. Re-export is needed to add the clustered candidates. Static collective preparation is unchanged.

Validation

Implemented: shared-component planning, ladder validation, both catalogs, and the reproducible FI dependency pin.

Qualified component checks: 118 CPU tests passed against the FI source fetched from the pinned official Git commit; uv lock --check --offline passes. One TP4 GPU test compiles and captures reductions followed by selected gathers, then alternates 8/16/32-row graphs for 18 replays. It checks exact outputs, preserved inputs, poisoned reusable outputs, consistent startup choices across ranks, and unchanged allocation counters. Hardware: four RTX PRO 6000 Blackwell GPUs. The compatible Heliopause extension used for these Python/component checks was built from Flare 6fa01a4; this PR changes no native code.

Focused CPU tests: tests/preparation/test_collective_ladder.py, tests/components/test_collective_plan_choice.py, tests/preparation/test_collectives.py, and tests/preparation/test_startup_tuning.py. GPU test: tests/components/test_clustered_gather_interleaving_cuda.py.

Serving measurements

The serving comparison uses an ordinary-ladder package and a package prepared with gb202_pcie_tp4_clustered_full.json. Serving measurements used DeepSeek V4.1 Flash 1017, TP4/C32/K7, maximum 128 verification rows, prefill 4096, native expert scales, BF16 MLA, RAM Engram tables and a 64 GiB prefix cache. The two machines, frank1 and frank2, have AMD EPYC 9575F CPUs and RTX PRO 6000 Blackwell 96 GB GPUs. Memory offsets were +6000 MHz; graphics clocks remained automatic. Both packages used the balanced scheduler, temperature/top-p 1 and, in this table, alpha 0.6.

Placement Host Requests Ordinary Clustered Change
Separate host bridges frank1 1 334.8 tok/s 337.3 tok/s +0.7%
Separate host bridges frank1 8 1179.7 tok/s 1289.7 tok/s +9.3%
Separate host bridges frank1 16 1655.6 tok/s 1799.8 tok/s +8.7%
Separate host bridges frank1 32 2313.6 tok/s 2379.5 tok/s +2.8%
Separate host bridges frank2 1 350.5 tok/s 330.6 tok/s -5.7%
Separate host bridges frank2 8 1159.7 tok/s 1252.1 tok/s +8.0%
Separate host bridges frank2 16 1701.3 tok/s 1784.1 tok/s +4.9%
Separate host bridges frank2 32 2306.0 tok/s 2301.1 tok/s -0.2%
Switch group frank1 32 2529.3 tok/s 2495.9 tok/s -1.3%
Switch group frank2 32 2544.8 tok/s 2460.6 tok/s -3.3%

Each host-bridge value is the mean of five 30-second LLM Bench cells on GPUs 0,4,8,12. Each switch value combines three cells on 0–3 and three on 4–7. Switch measurements ran two active services per host; host-bridge measurements ran one. These are within-placement comparisons, not an isolated measurement of topology effects. Cells used context 0/max 8192 output tokens and shared adaptive state within each native process; no IID or fixed-work GPU-latency claim is made.

All 352 accepted samples, including alpha 0.8 retain emitted-token throughput, completed steps/s, maximum-rank GPU time per step, verified draft depth and width-one fraction. Host-bridge C8/C16 improved across both machines; switch/C1 cases include lower means, so the clustered catalog remains explicit opt-in. All ordinary profiles remain available, but their compiled cluster geometry is still active.

Two synchronized groups of four cells were excluded and rerun with unchanged clients/guards: a C1 request-rollover gap violated full concurrency, and one C32 failure counter increased after the timed client finished. The client reported zero errors; an existing Flare unit test confirms a closed output channel can increment that counter, but the incident-specific error was not logged.

The serving packages used Flare 6fa01a4 and FlashInfer 361e45fa plus these clustered changes. Full-model performance on the Flare PR base 881eccd has not been rerun. The signed candidate runtime contains the measured preparation implementation; promotion preserved its execution descriptors and it passed 192 operational requests in the Flare 6fa Docker image.

Research-only / unqualified: complete-model performance on base 881eccd, broader hardware, and complete-model performance of the RSAG-only catalog. The component tests do not establish those results. Other cluster geometries and phased active counts are unsupported.

Tracking and graduation criteria: voipmonitor#7.

Implemented eight-CTA cluster geometry in Ignition's collective ladder
validation and plan declaration. Preserve ordinary profiles and provide
explicit catalogs with and without phased full all-reduce at 80 KiB.
Pin FlashInfer's cluster-capable preparation API at 560fef16648e4953.
Cluster geometry remains active when an ordinary startup profile wins;
Helios scheduling, execution ABI 5, and prepared contract 3 are unchanged.
Other cluster sizes and phased active-CTA counts are unsupported.

Qualified 118 CPU tests against the fetched official dependency, locked
resolution, and a TP4 captured reduction/gather interleaving test on RTX PRO
6000 GPUs. Replay preserves inputs, checks exact output, and allocates no
storage across 18 graph transitions at 8, 16, and 32 rows.

Retain 352 serving measurements from Flare 6fa01a4 and FlashInfer
361e45fad3c9ed0c; host and topology cohorts remain distinct. Full-model
performance on source base 881eccd is unmeasured. The retained
measurements include switched and concurrency-one regressions.

Design-Signed-Off-By: Codex (GPT-6); design-review=1c01
@lukealonso
lukealonso merged commit 18b16e1 into main Nov 3, 2026
1 of 3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants