libmir-cuda 0.3.0

CUDA inference backend for libmir
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use std::time::Duration;

use mircuda::{DeviceBuffer, bf16};
use runtime::tuning::{select_fastest_candidate, select_robust_candidate};

use super::AutoNvFp4Experts;
use crate::{
    Error, ExecutionPhase, Result,
    kernels::{RoutePattern, RoutePatternGenerator, RoutePatternSpec},
};

struct RoutePatterns {
    balanced: DeviceBuffer<u32>,
    hot_set: DeviceBuffer<u32>,
}

impl AutoNvFp4Experts {
    #[allow(clippy::cast_precision_loss)]
    pub(super) fn measure(
        &mut self,
        input: &DeviceBuffer<bf16>,
        selected: &DeviceBuffer<u32>,
        routing: &DeviceBuffer<bf16>,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<(usize, Duration, Duration)> {
        let patterns = (self.request.phase == ExecutionPhase::Prefill)
            .then(|| self.route_patterns())
            .transpose()?;
        let mut route_sets = vec![("resident", selected)];
        if let Some(patterns) = &patterns {
            route_sets.extend([("balanced", &patterns.balanced), ("hot_set", &patterns.hot_set)]);
        }
        let (warmup, iterations) = self.backend.auto_tuner().iterations(self.request.tokens);
        let mut timings = vec![Vec::with_capacity(route_sets.len()); self.candidates.len()];
        let mut elapsed = Duration::ZERO;
        for (distribution, routes) in route_sets {
            for (index, candidate) in self.candidates.iter_mut().enumerate() {
                for _ in 0..warmup {
                    candidate.plan.execute(input, routes, routing, output)?;
                }
                let started = self.backend.context().create_event(true)?;
                let completed = self.backend.context().create_event(true)?;
                started.record(self.backend.stream())?;
                for _ in 0..iterations {
                    candidate.plan.execute(input, routes, routing, output)?;
                }
                completed.record(self.backend.stream())?;
                completed.synchronize()?;
                let average = Duration::from_secs_f32(
                    started.elapsed_ms(&completed)? / (iterations as f32 * 1_000.0),
                );
                tracing::debug!(
                    target: "libmir::cuda::tuning",
                    distribution,
                    execution = ?candidate.execution,
                    average_us = average.as_secs_f64() * 1_000_000.0,
                    "measured CUDA MoE route distribution"
                );
                elapsed = elapsed
                    .saturating_add(average.saturating_mul(iterations.saturating_add(warmup)));
                timings[index].push(average);
            }
        }
        let winner = if patterns.is_some() {
            select_robust_candidate(
                self.fallback,
                &timings,
                self.backend.auto_tuner().minimum_improvement_bps(),
            )
        } else {
            let fastest = timings
                .iter()
                .enumerate()
                .min_by_key(|(_, values)| values[0])
                .map(|value| value.0)
                .ok_or(Error::InvalidExecutionPlan("MoE tuner has no candidates"))?;
            let scalar = timings.iter().map(|values| values[0]).collect::<Vec<_>>();
            select_fastest_candidate(
                fastest,
                self.fallback,
                &scalar,
                self.backend.auto_tuner().minimum_improvement_bps(),
            )
        };
        let total = timings[winner].iter().copied().fold(Duration::ZERO, Duration::saturating_add);
        let scenarios = u32::try_from(timings[winner].len())?;
        Ok((winner, total / scenarios, elapsed))
    }

    fn route_patterns(&self) -> Result<RoutePatterns> {
        let spec = RoutePatternSpec {
            tokens: self.request.tokens,
            experts: self.request.experts,
            top_k: self.request.top_k,
        };
        let generator = RoutePatternGenerator::compile(self.backend.compiler(), spec)?;
        let routes = spec
            .tokens
            .checked_mul(spec.top_k)
            .ok_or(Error::InvalidRouter("route pattern size overflow"))?;
        let mut balanced = self.backend.pool().allocate(self.backend.stream(), routes)?;
        let mut hot_set = self.backend.pool().allocate(self.backend.stream(), routes)?;
        generator.execute(self.backend.stream(), RoutePattern::Balanced, &mut balanced)?;
        generator.execute(self.backend.stream(), RoutePattern::HotSet, &mut hot_set)?;
        Ok(RoutePatterns { balanced, hot_set })
    }
}