vyre_driver_wgpu/engine/
multi_gpu.rs

1//! Multi-GPU adapter probing, backend acquisition, and work partitioning.
2//!
3//! The pure partitioner stays separately testable, but production callers can
4//! now derive device loads from live wgpu adapters and acquire one backend per
5//! selected adapter instead of stopping at scheduling math.
6//!
7//! ## Two allocation modes
8//!
9//! 1. **Batch + cost-aware** (`partition_work_stealing`): caller
10//!    knows every work item + its cost up front; LPT greedy assigns
11//!    the heaviest item to the least-loaded device.
12//! 2. **Stream + content-addressed**
13//!    (`shard_by_blake3` + `StreamShardAllocator`): caller yields
14//!    `(key, cost)` pairs one at a time from a walker. The initial
15//!    device is `blake3(key)[0] % n_gpus` for deterministic
16//!    affinity  -  files with the same path always land on the same
17//!    GPU across runs, which enables cache-warm re-scans. Overflow
18//!    (queue on the target GPU is already loaded above threshold)
19//!    spills to the least-loaded neighbor to keep tail latency
20//!    bounded.
21
22mod partition;
23mod stream_shard;
24
25use crate::staging_reserve::{reserve_multi_gpu_vec, reserve_smallvec, reserve_vec};
26
27pub use partition::{partition_work_stealing, DeviceLoad, Partition, WeightedWorkItem};
28pub use stream_shard::{shard_by_blake3, StreamShardAllocator};
29/// Error returned by public stream-sharding operations.
30pub type StreamShardError = stream_shard::StreamShardError;
31
32fn empty_gpu_work_result_slots(
33    len: usize,
34) -> Result<Vec<Option<Result<GpuWorkOutput, vyre_driver::BackendError>>>, vyre_driver::BackendError>
35{
36    let mut slots = Vec::new();
37    reserve_vec(
38        &mut slots,
39        len,
40        "multi-GPU executor",
41        "borrowed result slot",
42        "split the multi-GPU batch before dispatch",
43    )?;
44    slots.resize_with(len, || None);
45    Ok(slots)
46}
47
48fn finalize_gpu_work_results(
49    slots: Vec<Option<Result<GpuWorkOutput, vyre_driver::BackendError>>>,
50) -> Result<Vec<Result<GpuWorkOutput, vyre_driver::BackendError>>, vyre_driver::BackendError> {
51    let mut results = Vec::new();
52    reserve_vec(
53        &mut results,
54        slots.len(),
55        "multi-GPU executor",
56        "final borrowed result",
57        "split the multi-GPU batch before dispatch",
58    )?;
59    for slot in slots {
60        results.push(slot.unwrap_or_else(|| {
61            Err(vyre_driver::BackendError::new(
62                "multi-GPU borrowed dispatch result slot was not filled. Fix: ensure partitioning assigns every job exactly once.",
63            ))
64        }));
65    }
66    Ok(results)
67}
68
69/// Enumerate live wgpu GPU adapters as zero-load scheduling targets.
70///
71/// # Errors
72///
73/// Returns an error when wgpu exposes no real GPU adapters. On supported
74/// production hosts that means adapter probing or driver setup is broken.
75pub fn live_gpu_loads() -> Result<Vec<DeviceLoad>, String> {
76    let adapters = crate::runtime::device::enumerate_adapters();
77    let mut loads = Vec::new();
78    vyre_driver::allocation::try_reserve_vec_to_capacity(&mut loads, adapters.len()).map_err(
79        |source| {
80            format!(
81                "live GPU load enumeration could not reserve {} adapter slot(s): {source}. Fix: reduce adapter fanout or repair driver memory pressure before scheduling.",
82                adapters.len()
83            )
84        },
85    )?;
86    loads.extend(
87        adapters
88            .iter()
89            .enumerate()
90            .filter_map(|(device_index, info)| {
91                crate::capabilities::is_real_gpu(info).then_some(DeviceLoad {
92                    device_index,
93                    queued_cost: 0,
94                })
95            }),
96    );
97    if loads.is_empty() {
98        return Err(format!(
99            "wgpu enumerated {} adapters but none were real GPU execution targets. Fix: inspect driver setup and adapter filtering.",
100            adapters.len()
101        ));
102    }
103    Ok(loads)
104}
105
106/// A live GPU selected for multi-device scheduling.
107#[derive(Clone, Debug, Eq, PartialEq)]
108pub struct LiveGpu {
109    /// Adapter index accepted by `runtime::device::init_device_for_adapter`.
110    pub adapter_index: usize,
111    /// Stable adapter information captured during enumeration.
112    pub info: wgpu::AdapterInfo,
113}
114
115/// One executable multi-GPU work packet.
116pub struct GpuWorkItem {
117    /// Stable work identifier returned with the output.
118    pub id: usize,
119    /// Relative scheduling cost.
120    pub cost: u64,
121    /// Program to compile and dispatch on the selected adapter.
122    pub program: vyre_foundation::ir::Program,
123    /// Input buffers for the dispatch.
124    pub inputs: Vec<Vec<u8>>,
125    /// Dispatch policy for this item.
126    pub config: vyre_driver::DispatchConfig,
127}
128
129/// Output from one dispatched [`GpuWorkItem`].
130#[derive(Debug)]
131pub struct GpuWorkOutput {
132    /// Work identifier.
133    pub id: usize,
134    /// Adapter index that executed the work.
135    pub adapter_index: usize,
136    /// Dispatch outputs.
137    pub outputs: Vec<Vec<u8>>,
138}
139
140/// Borrowed multi-GPU work packet for hot scan paths.
141pub struct BorrowedGpuWorkItem<'a> {
142    /// Stable work identifier returned with the output.
143    pub id: usize,
144    /// Relative scheduling cost.
145    pub cost: u64,
146    /// Program to compile and dispatch on the selected adapter.
147    pub program: &'a vyre_foundation::ir::Program,
148    /// Input buffers for the dispatch.
149    pub inputs: &'a [&'a [u8]],
150    /// Dispatch policy for this item.
151    pub config: &'a vyre_driver::DispatchConfig,
152}
153
154/// Real multi-GPU executor backed by one [`crate::WgpuBackend`] per adapter.
155pub struct MultiGpuExecutor {
156    devices: Vec<ExecutorDevice>,
157}
158
159struct ExecutorDevice {
160    adapter_index: usize,
161    backend: crate::WgpuBackend,
162    queued_cost: u64,
163}
164
165impl MultiGpuExecutor {
166    /// Enumerate real GPU adapters visible to wgpu.
167    #[must_use]
168    pub fn enumerate_live_gpus() -> Vec<LiveGpu> {
169        let adapters = crate::runtime::device::enumerate_adapters();
170        let mut live = Vec::new();
171        let _ = vyre_driver::allocation::try_reserve_vec_to_capacity(&mut live, adapters.len());
172        for (adapter_index, info) in adapters.into_iter().enumerate() {
173            if crate::capabilities::is_real_gpu(&info) {
174                live.push(LiveGpu {
175                    adapter_index,
176                    info,
177                });
178            }
179        }
180        live
181    }
182
183    /// Build an executor over every real adapter that can create a device.
184    ///
185    /// # Errors
186    ///
187    /// Returns a backend error when no real adapters are visible or when a
188    /// visible adapter fails device creation.
189    pub fn acquire_all() -> Result<Self, vyre_driver::BackendError> {
190        let live = Self::enumerate_live_gpus();
191        if live.is_empty() {
192            return Err(vyre_driver::BackendError::new(
193                "no real GPU adapters found for multi-GPU execution. Fix: expose at least one discrete, integrated, or virtual GPU through wgpu before calling MultiGpuExecutor::acquire_all.",
194            ));
195        }
196        let mut devices = Vec::new();
197        reserve_multi_gpu_vec(&mut devices, live.len(), "executor device")?;
198        for gpu in live {
199            let backend = crate::WgpuBackend::acquire_adapter(gpu.adapter_index)?;
200            devices.push(ExecutorDevice {
201                adapter_index: gpu.adapter_index,
202                backend,
203                queued_cost: 0,
204            });
205        }
206        Ok(Self { devices })
207    }
208
209    /// Build an executor over explicit adapter indices.
210    ///
211    /// # Errors
212    ///
213    /// Returns a backend error if the list is empty, duplicated, or contains an
214    /// adapter that cannot create a real GPU device.
215    pub fn acquire_indices(indices: &[usize]) -> Result<Self, vyre_driver::BackendError> {
216        if indices.is_empty() {
217            return Err(vyre_driver::BackendError::new(
218                "no adapter indices supplied for multi-GPU execution. Fix: pass indices returned by runtime::device::enumerate_adapters().",
219            ));
220        }
221        let mut seen = rustc_hash::FxHashSet::default();
222        vyre_foundation::allocation::try_reserve_hash_set_to_capacity(&mut seen, indices.len())
223            .map_err(|source| {
224                vyre_driver::BackendError::new(format!(
225                    "multi-GPU adapter-index validation could not reserve {} seen slot(s): {source}. Fix: reduce adapter fanout before acquisition.",
226                    indices.len()
227                ))
228            })?;
229        let mut devices = Vec::new();
230        reserve_multi_gpu_vec(&mut devices, indices.len(), "executor device")?;
231        for &index in indices {
232            if !seen.insert(index) {
233                return Err(vyre_driver::BackendError::new(format!(
234                    "duplicate adapter index {index} supplied for multi-GPU execution. Fix: pass each adapter once."
235                )));
236            }
237            let backend = crate::WgpuBackend::acquire_adapter(index)?;
238            devices.push(ExecutorDevice {
239                adapter_index: index,
240                backend,
241                queued_cost: 0,
242            });
243        }
244        Ok(Self { devices })
245    }
246
247    /// Number of live device backends owned by this executor.
248    #[must_use]
249    pub fn len(&self) -> usize {
250        self.devices.len()
251    }
252
253    /// Whether the executor owns no devices.
254    #[must_use]
255    pub fn is_empty(&self) -> bool {
256        self.devices.is_empty()
257    }
258
259    /// Adapter indices owned by this executor.
260    #[must_use]
261    pub fn adapter_indices(&self) -> Vec<usize> {
262        let mut indices = Vec::new();
263        let _ =
264            vyre_driver::allocation::try_reserve_vec_to_capacity(&mut indices, self.devices.len());
265        indices.extend(self.devices.iter().map(|device| device.adapter_index));
266        indices
267    }
268
269    /// Dispatch a batch across the selected adapters using the LPT partitioner.
270    ///
271    /// Dispatches are submitted from one host thread per selected adapter. wgpu
272    /// devices are independent, so separate physical adapters can compile,
273    /// submit, and read back concurrently.
274    pub fn dispatch_batch(
275        &mut self,
276        items: Vec<GpuWorkItem>,
277    ) -> Result<Vec<GpuWorkOutput>, vyre_driver::BackendError> {
278        let mut devices = smallvec::SmallVec::<[DeviceLoad; 8]>::new();
279        reserve_smallvec(
280            &mut devices,
281            self.devices.len(),
282            "multi-GPU executor",
283            "device-load descriptor",
284            "split the multi-GPU batch before dispatch",
285        )?;
286        devices.extend(self.devices.iter().map(|device| DeviceLoad {
287            device_index: device.adapter_index,
288            queued_cost: device.queued_cost,
289        }));
290        let mut work = smallvec::SmallVec::<[WeightedWorkItem; 32]>::new();
291        reserve_smallvec(
292            &mut work,
293            items.len(),
294            "multi-GPU executor",
295            "weighted work descriptor",
296            "split the multi-GPU batch before partitioning",
297        )?;
298        work.extend(items.iter().map(|item| WeightedWorkItem {
299            id: item.id,
300            cost: item.cost,
301        }));
302        let partitions =
303            partition_work_stealing(&devices, &work).map_err(vyre_driver::BackendError::new)?;
304        let mut by_id = rustc_hash::FxHashMap::default();
305        vyre_foundation::allocation::try_reserve_hash_map_to_capacity(&mut by_id, items.len())
306            .map_err(|source| {
307                vyre_driver::BackendError::new(format!(
308                    "multi-GPU work-item lookup could not reserve {} owned item slot(s): {source}. Fix: split the multi-GPU batch.",
309                    items.len()
310                ))
311            })?;
312        by_id.extend(items.into_iter().map(|item| (item.id, item)));
313        let mut outputs = Vec::new();
314        reserve_multi_gpu_vec(&mut outputs, by_id.len(), "owned output")?;
315        std::thread::scope(|scope| {
316            let mut handles = smallvec::SmallVec::<[_; 8]>::new();
317            reserve_smallvec(
318                &mut handles,
319                partitions.len(),
320                "multi-GPU executor",
321                "worker thread handle",
322                "split the multi-GPU batch before dispatch",
323            )?;
324            for (partition, device) in partitions.into_iter().zip(self.devices.iter_mut()) {
325                if device.adapter_index != partition.device_index {
326                    return Err(vyre_driver::BackendError::new(format!(
327                        "partition targeted missing adapter {}. Fix: keep partition device indices synchronized with executor devices.",
328                        partition.device_index
329                    )));
330                }
331                device.queued_cost = partition.total_cost;
332                let backend = device.backend.clone();
333                let adapter_index = device.adapter_index;
334                let mut assigned = smallvec::SmallVec::<[_; 8]>::new();
335                reserve_smallvec(
336                    &mut assigned,
337                    partition.item_ids.len(),
338                    "multi-GPU executor",
339                    "assigned owned work item",
340                    "split the multi-GPU batch before dispatch",
341                )?;
342                for id in partition.item_ids {
343                    let item = by_id.remove(&id).ok_or_else(|| {
344                        vyre_driver::BackendError::new(format!(
345                            "partition referenced unknown work item {id}. Fix: partition only ids from the submitted batch."
346                        ))
347                    })?;
348                    assigned.push(item);
349                }
350                handles.push(scope.spawn(move || {
351                    let mut local = Vec::new();
352                    reserve_multi_gpu_vec(&mut local, assigned.len(), "worker-local output")?;
353                    for item in assigned {
354                        let outputs = vyre_driver::VyreBackend::dispatch(
355                            &backend,
356                            &item.program,
357                            &item.inputs,
358                            &item.config,
359                        )?;
360                        local.push(GpuWorkOutput {
361                            id: item.id,
362                            adapter_index,
363                            outputs,
364                        });
365                    }
366                    Ok::<_, vyre_driver::BackendError>(local)
367                }));
368            }
369            for handle in handles {
370                let mut local = handle.join().map_err(|_| {
371                    vyre_driver::BackendError::new(
372                        "multi-GPU worker thread panicked. Fix: inspect adapter-specific dispatch failure handling.",
373                    )
374                })??;
375                outputs.append(&mut local);
376            }
377            Ok::<_, vyre_driver::BackendError>(())
378        })?;
379        if !by_id.is_empty() {
380            return Err(vyre_driver::BackendError::new(
381                "multi-GPU partition left unassigned work items. Fix: partition every submitted item exactly once.",
382            ));
383        }
384        outputs.sort_by_key(|output| output.id);
385        Ok(outputs)
386    }
387
388    /// Dispatch a borrowed batch across selected adapters without cloning
389    /// input buffers into owned work packets.
390    ///
391    /// Each adapter receives one backend-local batched submission, so wgpu's
392    /// per-device command buffers stay valid while independent adapters run on
393    /// separate host threads.
394    pub fn dispatch_borrowed_batch(
395        &mut self,
396        items: &[BorrowedGpuWorkItem<'_>],
397    ) -> Result<Vec<Result<GpuWorkOutput, vyre_driver::BackendError>>, vyre_driver::BackendError>
398    {
399        if items.is_empty() {
400            return Ok(Vec::new());
401        }
402        let mut devices = smallvec::SmallVec::<[DeviceLoad; 8]>::new();
403        reserve_smallvec(
404            &mut devices,
405            self.devices.len(),
406            "multi-GPU executor",
407            "borrowed device-load descriptor",
408            "split the multi-GPU batch before dispatch",
409        )?;
410        devices.extend(self.devices.iter().map(|device| DeviceLoad {
411            device_index: device.adapter_index,
412            queued_cost: device.queued_cost,
413        }));
414        let mut work = smallvec::SmallVec::<[WeightedWorkItem; 32]>::new();
415        reserve_smallvec(
416            &mut work,
417            items.len(),
418            "multi-GPU executor",
419            "borrowed weighted work descriptor",
420            "split the multi-GPU batch before partitioning",
421        )?;
422        work.extend(items.iter().map(|item| WeightedWorkItem {
423            id: item.id,
424            cost: item.cost,
425        }));
426        let partitions =
427            partition_work_stealing(&devices, &work).map_err(vyre_driver::BackendError::new)?;
428        let mut by_id = rustc_hash::FxHashMap::default();
429        vyre_foundation::allocation::try_reserve_hash_map_to_capacity(&mut by_id, items.len())
430            .map_err(|source| {
431                vyre_driver::BackendError::new(format!(
432                    "multi-GPU work-item lookup could not reserve {} borrowed item slot(s): {source}. Fix: split the multi-GPU batch.",
433                    items.len()
434                ))
435            })?;
436        by_id.extend(items.iter().enumerate().map(|(slot, item)| (item.id, slot)));
437        let mut results = empty_gpu_work_result_slots(items.len())?;
438
439        std::thread::scope(|scope| {
440            let mut handles = smallvec::SmallVec::<[_; 8]>::new();
441            reserve_smallvec(
442                &mut handles,
443                partitions.len(),
444                "multi-GPU executor",
445                "borrowed worker thread handle",
446                "split the multi-GPU batch before dispatch",
447            )?;
448            for (partition, device) in partitions.into_iter().zip(self.devices.iter_mut()) {
449                if device.adapter_index != partition.device_index {
450                    return Err(vyre_driver::BackendError::new(format!(
451                        "partition targeted missing adapter {}. Fix: keep partition device indices synchronized with executor devices.",
452                        partition.device_index
453                    )));
454                }
455                device.queued_cost = partition.total_cost;
456                let backend = device.backend.clone();
457                let adapter_index = device.adapter_index;
458                let mut assigned_slots = smallvec::SmallVec::<[_; 8]>::new();
459                reserve_smallvec(
460                    &mut assigned_slots,
461                    partition.item_ids.len(),
462                    "multi-GPU executor",
463                    "assigned borrowed work slot",
464                    "split the multi-GPU batch before dispatch",
465                )?;
466                for id in partition.item_ids {
467                    let slot = by_id.remove(&id).ok_or_else(|| {
468                        vyre_driver::BackendError::new(format!(
469                            "partition referenced unknown work item {id}. Fix: partition only ids from the submitted batch."
470                        ))
471                    })?;
472                    assigned_slots.push(slot);
473                }
474                handles.push(scope.spawn(move || {
475                    let mut backend_jobs = smallvec::SmallVec::<
476                        [(
477                            &vyre_foundation::ir::Program,
478                            &[&[u8]],
479                            &vyre_driver::DispatchConfig,
480                        ); 8],
481                    >::new();
482                    reserve_smallvec(
483                        &mut backend_jobs,
484                        assigned_slots.len(),
485                        "multi-GPU executor",
486                        "backend-local borrowed job descriptor",
487                        "split the multi-GPU batch before dispatch",
488                    )?;
489                    backend_jobs.extend(assigned_slots.iter().map(|&slot| {
490                        let item = &items[slot];
491                        (item.program, item.inputs, item.config)
492                    }));
493                    let local = backend.dispatch_borrowed_batch(&backend_jobs)?;
494                    Ok::<_, vyre_driver::BackendError>((adapter_index, assigned_slots, local))
495                }));
496            }
497            for handle in handles {
498                let (adapter_index, assigned_slots, local) = handle.join().map_err(|_| {
499                    vyre_driver::BackendError::new(
500                        "multi-GPU borrowed worker thread panicked. Fix: inspect adapter-specific dispatch failure handling.",
501                    )
502                })??;
503                if assigned_slots.len() != local.len() {
504                    return Err(vyre_driver::BackendError::new(format!(
505                        "adapter {adapter_index} returned {} results for {} assigned jobs. Fix: keep backend batch metadata synchronized.",
506                        local.len(),
507                        assigned_slots.len()
508                    )));
509                }
510                for (slot, output_result) in assigned_slots.into_iter().zip(local) {
511                    let id = items[slot].id;
512                    results[slot] = Some(output_result.map(|outputs| GpuWorkOutput {
513                        id,
514                        adapter_index,
515                        outputs,
516                    }));
517                }
518            }
519            Ok::<_, vyre_driver::BackendError>(())
520        })?;
521        if !by_id.is_empty() {
522            return Err(vyre_driver::BackendError::new(
523                "multi-GPU borrowed partition left unassigned work items. Fix: partition every submitted item exactly once.",
524            ));
525        }
526
527        finalize_gpu_work_results(results)
528    }
529}
530
531#[cfg(test)]
532mod tests {
533    use super::*;
534
535    #[test]
536    fn live_gpu_enumeration_uses_wgpu_adapters() {
537        let live = MultiGpuExecutor::enumerate_live_gpus();
538        assert!(
539            !live.is_empty(),
540            "Fix: multi-GPU runtime must enumerate at least one real GPU adapter on this fleet host."
541        );
542        for gpu in live {
543            assert!(
544                crate::capabilities::is_real_gpu(&gpu.info),
545                "Fix: multi-GPU executor must filter CPU/Other adapters before scheduling: {:?}",
546                gpu.info
547            );
548        }
549    }
550
551    #[test]
552    fn acquire_indices_rejects_duplicate_live_ordinals_before_dispatch() {
553        let live = MultiGpuExecutor::enumerate_live_gpus();
554        let first = live
555            .first()
556            .expect("Fix: duplicate-index test requires the live GPU adapter promised by the fleet")
557            .adapter_index;
558        let error = match MultiGpuExecutor::acquire_indices(&[first, first]) {
559            Ok(_) => panic!("Fix: duplicate adapter indices must be rejected"),
560            Err(error) => error,
561        };
562        assert!(
563            error.to_string().contains("duplicate adapter index"),
564            "Fix: duplicate-index diagnostic must be actionable, got: {error}"
565        );
566    }
567
568    #[test]
569    fn generated_borrowed_result_finalization_preserves_slots_and_reports_missing_work() {
570        for case in 0..4096usize {
571            let len = (case % 17) + 1;
572            let mut slots = empty_gpu_work_result_slots(len)
573                .expect("Fix: generated multi-GPU slot test must reserve result slots");
574            for slot in 0..len {
575                if (slot + case) % 5 == 0 {
576                    continue;
577                }
578                slots[slot] = Some(Ok(GpuWorkOutput {
579                    id: slot,
580                    adapter_index: case % 3,
581                    outputs: vec![vec![slot as u8, case as u8]],
582                }));
583            }
584
585            let finalized = finalize_gpu_work_results(slots)
586                .expect("Fix: generated multi-GPU finalization must reserve final results");
587            assert_eq!(
588                finalized.len(),
589                len,
590                "generated multi-GPU case {case} must preserve slot count"
591            );
592            for (slot, result) in finalized.into_iter().enumerate() {
593                if (slot + case) % 5 == 0 {
594                    let error = result
595                        .expect_err("Fix: unfilled generated multi-GPU slot must be an error");
596                    assert!(
597                        error.to_string().contains("result slot was not filled"),
598                        "Fix: missing generated multi-GPU slot must explain partition coverage, got {error}"
599                    );
600                } else {
601                    let output =
602                        result.expect("Fix: filled generated multi-GPU slot must stay successful");
603                    assert_eq!(output.id, slot);
604                    assert_eq!(output.adapter_index, case % 3);
605                    assert_eq!(output.outputs, vec![vec![slot as u8, case as u8]]);
606                }
607            }
608        }
609    }
610}
vyre_driver_wgpu/engine/multi_gpu.rs

vyre_driver_wgpu/engine/
multi_gpu.rs
