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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

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

mod paired;

pub use paired::{
    PairedVariableGroupedFp4Launch, PairedVariableGroupedFp4Plan, VariableGroupedFp4Metadata,
    VariableGroupedFp4Operands,
};

/// Capacity and maximum geometry for variable-row grouped FP4 multiplication.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct VariableGroupedFp4Spec {
    groups: usize,
    matrices: usize,
    max_rows: usize,
    n: usize,
    k: usize,
    capacity_rows: usize,
}

impl VariableGroupedFp4Spec {
    /// Validates a compact grouped product with device-resident row counts and offsets.
    pub fn new(
        groups: usize,
        matrices: usize,
        max_rows: usize,
        n: usize,
        k: usize,
        capacity_rows: usize,
    ) -> Result<Self> {
        if groups == 0
            || matrices == 0
            || max_rows == 0
            || n == 0
            || k == 0
            || capacity_rows == 0
            || !k.is_multiple_of(64)
        {
            return Err(Error::InvalidMatmulShape);
        }
        for value in [groups, matrices, max_rows, n, k, capacity_rows] {
            let _ = u32::try_from(value)?;
        }
        capacity_rows.checked_mul(k / 2).ok_or(Error::InvalidMatmulShape)?;
        capacity_rows.checked_mul(n).ok_or(Error::InvalidMatmulShape)?;
        Ok(Self {
            groups,
            matrices,
            max_rows,
            n,
            k,
            capacity_rows,
        })
    }

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

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

    /// Maximum rows accepted for one product.
    #[must_use]
    pub const fn max_rows(self) -> usize {
        self.max_rows
    }

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

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

    /// Total compact A and output row capacity.
    #[must_use]
    pub const fn capacity_rows(self) -> usize {
        self.capacity_rows
    }

    const fn native(self) -> mircuda_sys::VariableGroupedFp4Spec {
        mircuda_sys::VariableGroupedFp4Spec {
            groups: self.groups,
            matrices: self.matrices,
            max_m: self.max_rows,
            n: self.n,
            k: self.k,
            capacity_rows: self.capacity_rows,
        }
    }
}

/// Persistent CUTLASS plan whose row counts and compact offsets stay on CUDA.
#[derive(Debug)]
pub struct VariableGroupedFp4Plan {
    native: mircuda_sys::VariableGroupedFp4Plan,
    spec: VariableGroupedFp4Spec,
}

impl VariableGroupedFp4Plan {
    /// Creates a capacity-bounded plan without synchronizing the stream.
    pub fn new(context: &Context, stream: &Stream, spec: VariableGroupedFp4Spec) -> Result<Self> {
        Ok(Self {
            native: context
                .native
                .create_variable_grouped_fp4_plan(&stream.native, spec.native())?,
            spec,
        })
    }

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

    /// Enqueues all products without reading row counts or offsets on the host.
    ///
    /// `rows[group]` must not exceed `max_rows`. `offsets[group]` must identify
    /// a non-overlapping region of `capacity_rows`. Zero-row groups occupy no
    /// compact rows. A scales use one fixed `max_rows`-sized region per group.
    #[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>,
        rows: &DeviceBuffer<u32>,
        offsets: &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, &rows.native, &offsets.native, &output.native,
        )?)
    }
}