[Feature][QDP] Add MPI-ready distributed amplitude execution scaffolding

[viiccwen]

Related Issues

Closes #1295
Part of #1297

Summary

This PR adds the first QDP-native distributed GPU state execution foundation, with amplitude state construction as the first workload.

The implementation is intentionally MPI-ready but does not require MPI or NCCL as mandatory dependencies. The executable path in this PR is a real single-process fallback that runs through the same rank-local planning, runtime, transport, and state boundaries that a future MPI-backed execution path should use.

In other words, this PR is not claiming full mpirun -n N support yet. It establishes the architecture and a testable fallback boundary that future MPI, CUDA-aware MPI, and NCCL backends can replace or extend.

The main architecture boundary is intentionally split into four layers:

• Planning decides where shards should live.
• Runtime decides what work the current rank performs.
• Transport owns scalar and future GPU-resident collectives.
• State owns local GPU shard buffers and exposes local zero-copy views.

Current Capability

This PR currently supports:

• rank/world execution metadata through DistributedExecutionContext
• rank-local device meshes and placement planning
• shared distributed state planning through DistributedStatePlan
• amplitude-specific execution through DistributedAmplitudePlan
• local GPU shard materialization for shards owned by the current rank
• scalar rank-level reduction through CollectiveCommunicator
• single-process fallback through LocalCollectiveCommunicator
• zero-copy local shard metadata through DistributedStateVector::local_shard_views()
• planned device_id to CUDA handle resolution, so shard metadata, active CUDA device context, and allocated buffers stay aligned
• placeholder CUDA-aware MPI / NCCL device collective types that identify their backend but return NotImplemented for GPU-resident collectives

Architecture Shape

sequenceDiagram
    participant Caller as QdpEngine caller
    participant Engine as QdpEngine
    participant Ctx as DistributedExecutionContext
    participant Planner as PlacementPlanner
    participant StatePlan as DistributedStatePlan
    participant AmpPlan as DistributedAmplitudePlan
    participant Runtime as distributed runtime
    participant Comm as CollectiveCommunicator
    participant State as DistributedStateVector

    Caller->>Engine: submit distributed amplitude request
    Engine->>Engine: validate input and resolve request
    Engine->>Ctx: construct rank-local execution context
    Engine->>Planner: build rank-local placement plan
    Planner-->>Engine: placement + shard ranges
    Engine->>StatePlan: build shared shard/rank contract
    StatePlan-->>AmpPlan: wrap for amplitude workload
    AmpPlan->>Runtime: execute local-rank shard work
    Runtime->>Runtime: compute local norm contribution
    Runtime->>Comm: all_reduce_sum_f64(local_norm_sq)
    Comm-->>Runtime: global norm
    Runtime->>Runtime: allocate and encode local GPU shards
    Runtime-->>State: local distributed state
    State-->>Caller: zero-copy local shard views available

Loading

This PR does not yet support:

• true MPI multi-rank launch
• CUDA-aware MPI collective execution
• NCCL collective execution
• multi-node execution
• cross-rank zero-copy transport

Follow-up Work

• Implement a true MPI scalar communicator for CollectiveCommunicator.
• Add a documented mpirun -n N launch path.
• Add multi-rank smoke tests gated behind an MPI feature or CI environment.
• Implement CUDA-aware MPI and/or NCCL DeviceCollectiveCommunicator backends.
• Define gather/export and downstream consumer contracts for distributed state.

[400Ping]

Nice one, will probably take a look on Thursday.

[ryankert01]

Hi, nice initiative. Curious of can it be plugging in our current lighting.gpu (penny lane) workflow?

mahout.qdp -> lighting.gpu (by zero copy)

[viiccwen]

@ryankert01, after researching, I think the main gap is on the lightning.gpu integration boundary.

From what I can tell, the missing piece is not “can Mahout produce a GPU-resident state?” but “can lightning.gpu ingest an external GPU state buffer through a stable public interface without copying?”.

So my current estimate would be:

1. expose or formalize an external-state ingest path on the PennyLane / lightning.gpu side
2. add a Mahout-side bridge/adapter for mahout.qdp -> lightning.gpu
3. add end-to-end tests and examples for the workflow

So yes: for a proper zero-copy integration, I would expect roughly 3 PRs for a narrow MVP.

[ryankert01]

@viiccwen I think since it's 3 PRs away and so big. Can we postpone it to the next release? It will be more mature at next release and we can think about its detail.

[viiccwen]

@viiccwen I think since it's 3 PRs away and so big. Can we postpone it to the next release? It will be more mature at next release and we can think about its detail.

Sure, it'll be fine.

[400Ping]

@viiccwen do you have a design doc for this?

[400Ping]

Also please solve conflicts

[rich7420]

hey @viiccwen — i think we'll overlap here. i'm planning a follow-up pr for ddp rank-aware sharding in QuantumDataLoader (feistel-based O(1) partitioning).

two quick questions:

1. rough merge eta for [Feature][QDP] Add MPI-ready distributed amplitude execution scaffolding #1296?
2. what does the MPIContext / topology api surface look like? would prefer to consume yours than invent a parallel one.

happy to wait if you're close, otherwise i'll build single-node-first and layer mpi on top later.

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity for 30 days.

It will be closed in another 7 days if no further activity occurs. Thank you for your contribution. If you'd like to keep this open, leave any comment and the stale label will be removed.

You can always ask for help on the Mahout dev mailing list or in GitHub Discussions.

[viiccwen]

This branch supports single-process, single-node multi-GPU shard materialization with zero-copy local shard views. It does not yet implement true MPI multi-rank execution or cross-rank GPU-resident zero-copy transport.