Add tensor and context parallel inference to the core runtime

[jmccaffrey-nv]

Summary

• Add reusable tensor- and context-parallel inference primitives, including autograd-aware sequence collect and parallel linear, attention, and normalization modules.
• Add device-mesh construction and runtime state with independently selectable backends for model-parallel axes.
• Coordinate multi-GPU API v2 sessions at model-step boundaries while keeping presentation and metrics ownership on rank zero.
• Document the runtime architecture and add CPU fault-injection coverage for distributed primitives, device ownership, and session lifecycle behavior.

Review follow-up

• Clean up process groups created by init_parallel(): successful v2 CLI runs use coordinated shutdown, failures avoid peer waits, and the owned-world exit fallback skips potentially blocking destruction after an uncaught rank failure.
• Preserve the caller-selected CUDA device when LOCAL_RANK is unavailable.
• Keep replacement requests alive across session cleanup and application deadlines.
• Keep an unfinished UI open after inference while workers poll rank zero once per UI tick.
• Bypass session-result collectives on failures so the original exception reaches the process supervisor promptly.
• Make rank zero's terminal or replacement decision authoritative across all runtime ranks.

Validation

• GitHub CPU CI (pre-commit/type checks and ci_cpu pytest): passed.
• GitHub GPU CI: passed.
• GitHub MP4 demo and WorldLens canary: passed.
• GitHub docs, REUSE, and OSRB checks: passed.
• Greptile review: 5/5, with no outstanding findings or actionable regressions.
• Local CI-pinned Ruff lint/format, Python bytecode compilation, and git diff --check: passed.

Release review

Reviewed for public release. No secrets, internal infrastructure references, or non-public benchmark results were found in the final GitHub diff.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[jmccaffrey-nv]

/ok to test 128732b

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge, with no outstanding findings or actionable new regressions identified.

Findings

1.  Exit cleanup blocks failed ranks ▶

Summary

• Introduces mesh planning, uneven token gathering, and sharded linear projections.
• Coordinates model-step admission, input replication, and terminal or replacement decisions.
• The latest change skips process-group exit cleanup after an uncaught exception and adds focused regression tests.
• No actionable new findings were identified.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Initialize tensor/context mesh] --> B[Create runtime control group]
    B --> C[Agree on step admission and inputs]
    C --> D[Execute model shards on all ranks]
    D --> E[Rank zero publishes output]
    E --> F{Model execution complete?}
    F -- No --> C
    F -- Yes --> G[Poll rank-zero session decision each UI tick]
    G -- Continue viewing --> G
    G -- Replace --> H[Clean up and enter replacement together]
    H --> B
    G -- Stop --> I[Clean up and exit]
    D -- Failure --> J[Propagate error without result agreement]

Loading

_{Reviews (8) · Last reviewed commit: "Avoid blocking cleanup after rank failur..."}

[jmccaffrey-nv]

/ok to test 3de71c5

[jarcherNV]

Thanks for doing this, the overall approach makes sense to me. Just a few comments (mostly from me asking codex about this PR):

P1 - Could you please make sure the process group created by init_parallel() is cleaned up? It looks like this path bypasses the existing atexit registration, and the v2 CLI never calls the distributed shutdown helper. Perhaps it could follow the shutdown pattern used by flashdreams-run.

P2 - It looks like reusing an existing process group without LOCAL_RANK resets every rank to cuda:0, even if the caller already selected a device. Do we need to preserve the current CUDA device or accept an explicit device?

P2 - It looks like request_new_session() is discarded during distributed execution, causing the application to exit instead of starting the replacement session. Is that intended behavior?

P2 - It looks like distributed sessions skip the normal unfinished-UI check once inference completes. Should rank zero retain the existing post-inference UI lifecycle while the worker ranks wait for a coordinated exit?

Up to you if you think the P2s are worth addressing.

[jmccaffrey-nv]

/ok to test 823746f

[jmccaffrey-nv]

/ok to test 1950038

[jmccaffrey-nv]

/ok to test 2b29590

[jmccaffrey-nv]

/ok to test 4c996ff

[jmccaffrey-nv]

/ok to test 066df61

[jmccaffrey-nv]

/ok to test d4446e4

[jmccaffrey-nv]

Thanks - I agreed these were worth addressing and pushed the fixes through d4446e41.

• Process-group cleanup: init_parallel() registers cleanup only when it creates the default group. The v2 CLI now follows the existing runner's bounded shutdown policy: synchronized cleanup after success and no peer wait/destruction after a rank-local failure. The exit fallback is also failure-aware for programmatic callers, so an uncaught exception cannot block the failed rank in destroy_process_group().
• Existing CUDA worlds: when LOCAL_RANK is absent, init_parallel() preserves the caller's current CUDA device.
• Replacement sessions: rank zero's replacement request and deadline decision are coordinated after cleanup and returned on every rank.
• Post-inference UI: rank zero retains the normal unfinished-UI lifecycle; workers participate in a lightweight coordinated decision once per UI tick.

I also incorporated the follow-on distributed edge cases Greptile found (idle UI timeout, failure-path collectives, and deadline divergence). Final-head CPU, GPU, MP4, WorldLens, docs, REUSE, and OSRB checks all pass. Greptile is now 5/5 with no outstanding/actionable findings.