mircuda 0.1.0

Native, explicit-stream Rust gateway to NVIDIA CUDA
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
114
115
116
117
118
119
120
121
122
123
124

use crate::{Context, DeviceBuffer, Error, Result, Stream, bf16};

/// Fixed geometry for device-indexed grouped block-scaled FP4 multiplication.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct IndexedGroupedFp4Spec {
    groups: usize,
    matrices: usize,
    m: usize,
    n: usize,
    k: usize,
    broadcast_input: bool,
}

impl IndexedGroupedFp4Spec {
    /// Validates fixed-row `M x K` by `N x K` products selected from a matrix bank.
    pub fn new(groups: usize, matrices: usize, m: usize, n: usize, k: usize) -> Result<Self> {
        if groups == 0 || matrices == 0 || m == 0 || n == 0 || k == 0 || !k.is_multiple_of(64) {
            return Err(Error::InvalidMatmulShape);
        }
        for value in [groups, matrices, m, n, k] {
            let _ = i32::try_from(value)?;
        }
        Ok(Self {
            groups,
            matrices,
            m,
            n,
            k,
            broadcast_input: false,
        })
    }

    /// Reuses one A matrix and its block scales for every grouped product.
    #[must_use]
    pub const fn with_broadcast_input(mut self) -> Self {
        self.broadcast_input = true;
        self
    }

    /// Number of independent products.
    #[must_use]
    pub const fn groups(self) -> usize {
        self.groups
    }

    /// Number of matrices in the device-resident B bank.
    #[must_use]
    pub const fn matrices(self) -> usize {
        self.matrices
    }

    /// Rows in each grouped product.
    #[must_use]
    pub const fn m(self) -> usize {
        self.m
    }

    /// Output columns in each product.
    #[must_use]
    pub const fn n(self) -> usize {
        self.n
    }

    /// Reduction dimension in each product.
    #[must_use]
    pub const fn k(self) -> usize {
        self.k
    }

    const fn native(self) -> mircuda_sys::IndexedGroupedFp4Spec {
        mircuda_sys::IndexedGroupedFp4Spec {
            groups: self.groups,
            matrices: self.matrices,
            m: self.m,
            n: self.n,
            k: self.k,
            broadcast_input: self.broadcast_input,
        }
    }
}

/// Persistent SM120 CUTLASS plan with device-resident matrix selection.
#[derive(Debug)]
pub struct IndexedGroupedFp4Plan {
    native: mircuda_sys::IndexedGroupedFp4Plan,
    spec: IndexedGroupedFp4Spec,
}

impl IndexedGroupedFp4Plan {
    /// Enqueues persistent CUTLASS metadata and workspace creation on `stream`.
    pub fn new(context: &Context, stream: &Stream, spec: IndexedGroupedFp4Spec) -> Result<Self> {
        Ok(Self {
            native: context
                .native
                .create_indexed_grouped_fp4_plan(&stream.native, spec.native())?,
            spec,
        })
    }

    /// Returns the immutable geometry associated with this plan.
    #[must_use]
    pub const fn spec(&self) -> IndexedGroupedFp4Spec {
        self.spec
    }

    /// Enqueues all indexed products without host selection, allocation, or synchronization.
    #[allow(clippy::too_many_arguments)]
    pub fn execute(
        &mut self,
        stream: &Stream,
        a: &DeviceBuffer<u8>,
        a_scales: &DeviceBuffer<u8>,
        b: &DeviceBuffer<u8>,
        b_scales: &DeviceBuffer<u8>,
        alphas: &DeviceBuffer<f32>,
        indices: &DeviceBuffer<u32>,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        Ok(self.native.execute(
            &stream.native, &a.native, &a_scales.native, &b.native, &b_scales.native,
            &alphas.native, &indices.native, &output.native,
        )?)
    }
}