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

vyre_driver_wgpu/engine/
multi_gpu.rs
