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
108
109
110
111
112
113

use super::{
    DenseGateUpLayout, SelectedDenseDispatch, SelectedDenseGateLaunch, SelectedDenseMoe,
    SelectedDenseReduceLaunch,
};
use crate::{
    Result,
    kernels::geometry::{product, require},
};

impl SelectedDenseMoe {
    pub(super) fn validate_dispatch(
        &self,
        selected: &mircuda::DeviceBuffer<u32>,
        dispatch: &SelectedDenseDispatch<'_>,
    ) -> Result<()> {
        let spec = self.spec();
        let assignments = product(spec.tokens, spec.selected_count)?;
        require("dense dispatch selected", assignments, selected.len())?;
        require("dense dispatch counts", spec.expert_count, dispatch.counts.len())?;
        require("dense dispatch offsets", spec.expert_count, dispatch.offsets.len())?;
        require("dense dispatch cursors", spec.expert_count, dispatch.cursors.len())?;
        require("dense dispatch assignments", assignments, dispatch.assignments.len())?;
        require("dense dispatch experts", assignments, dispatch.experts.len())
    }

    pub(super) fn validate_gated(&self, launch: &SelectedDenseGateLaunch<'_>) -> Result<()> {
        let spec = self.spec();
        require(
            "dense expert input",
            product(spec.tokens, spec.input_features)?,
            launch.input.len(),
        )?;
        require(
            "dense selected experts",
            product(spec.tokens, spec.selected_count)?,
            launch.selected.len(),
        )?;
        let rows = if spec.gate_up_layout == DenseGateUpLayout::Separate {
            spec.output_features
        } else {
            product(spec.output_features, 2)?
        };
        let fused = product(product(spec.expert_count, rows)?, spec.input_features)?;
        let separate =
            product(product(spec.expert_count, spec.output_features)?, spec.input_features)?;
        require("dense gate weight", fused, launch.gate_weight.len())?;
        require(
            "dense up weight",
            if spec.gate_up_layout == DenseGateUpLayout::Separate {
                separate
            } else {
                fused
            },
            launch.up_weight.len(),
        )?;
        self.validate_bias("dense gate bias", launch.gate_bias.len(), rows, spec.gate_bias)?;
        self.validate_bias("dense up bias", launch.up_bias.len(), rows, spec.up_bias)?;
        require(
            "dense gated output",
            product(product(spec.tokens, spec.selected_count)?, spec.output_features)?,
            launch.output.len(),
        )
    }

    pub(super) fn validate_reduce(&self, launch: &SelectedDenseReduceLaunch<'_>) -> Result<()> {
        let spec = self.spec();
        require(
            "dense activated experts",
            product(product(spec.tokens, spec.selected_count)?, spec.output_features)?,
            launch.input.len(),
        )?;
        let selections = product(spec.tokens, spec.selected_count)?;
        require("dense selected experts", selections, launch.selected.len())?;
        require("dense routing weights", selections, launch.routing.len())?;
        if (spec.down_transposed && spec.tokens == 1) || self.prefers_expert_major() {
            require(
                "dense down partials",
                product(selections, spec.input_features)?,
                launch.partial.len(),
            )?;
        }
        require(
            "dense down weight",
            product(product(spec.expert_count, spec.input_features)?, spec.output_features)?,
            launch.weight.len(),
        )?;
        self.validate_bias(
            "dense down bias",
            launch.bias.len(),
            spec.input_features,
            spec.down_bias,
        )?;
        require(
            "dense expert output",
            product(spec.tokens, spec.input_features)?,
            launch.output.len(),
        )
    }

    fn validate_bias(
        &self,
        name: &'static str,
        actual: usize,
        rows: usize,
        present: bool,
    ) -> Result<()> {
        if present {
            require(name, product(self.spec().expert_count, rows)?, actual)
        } else {
            Ok(())
        }
    }
}