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

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

/// One operand set for a paired variable-row grouped FP4 dispatch.
pub struct VariableGroupedFp4Operands<'a> {
    /// Compact packed FP4 activation rows.
    pub a: &'a DeviceBuffer<u8>,
    /// CUTLASS-formatted activation scales.
    pub a_scales: &'a DeviceBuffer<u8>,
    /// Packed matrix bank.
    pub b: &'a DeviceBuffer<u8>,
    /// CUTLASS-formatted matrix scales.
    pub b_scales: &'a DeviceBuffer<u8>,
    /// Per-matrix FP32 epilogue multipliers.
    pub alphas: &'a DeviceBuffer<f32>,
    /// Compact BF16 result rows.
    pub output: &'a mut DeviceBuffer<bf16>,
}

/// Shared device-resident scheduling metadata for a paired dispatch.
#[derive(Clone, Copy)]
pub struct VariableGroupedFp4Metadata<'a> {
    /// Matrix index selected for every logical group.
    pub indices: &'a DeviceBuffer<u32>,
    /// Active row count for every logical group.
    pub rows: &'a DeviceBuffer<u32>,
    /// Compact row offset for every logical group.
    pub offsets: &'a DeviceBuffer<u32>,
}

/// Two independent grouped products sharing only scheduling metadata.
pub struct PairedVariableGroupedFp4Launch<'a> {
    /// First grouped product.
    pub left: VariableGroupedFp4Operands<'a>,
    /// Second grouped product.
    pub right: VariableGroupedFp4Operands<'a>,
    /// Metadata shared by both products.
    pub metadata: VariableGroupedFp4Metadata<'a>,
}

/// Persistent CUTLASS plan issuing two same-geometry grouped products through
/// one device scheduler.
#[derive(Debug)]
pub struct PairedVariableGroupedFp4Plan {
    native: mircuda_sys::PairedVariableGroupedFp4Plan,
    spec: VariableGroupedFp4Spec,
}

impl PairedVariableGroupedFp4Plan {
    /// Creates a paired plan without allocating or joining operand banks.
    pub fn new(context: &Context, stream: &Stream, spec: VariableGroupedFp4Spec) -> Result<Self> {
        Ok(Self {
            native: context
                .native
                .create_paired_variable_grouped_fp4_plan(&stream.native, spec.native())?,
            spec,
        })
    }

    /// Returns the geometry of each logical product set.
    #[must_use]
    pub const fn spec(&self) -> VariableGroupedFp4Spec {
        self.spec
    }

    /// Enqueues both products in one grouped CUTLASS dispatch.
    pub fn execute(
        &mut self,
        stream: &Stream,
        launch: &mut PairedVariableGroupedFp4Launch<'_>,
    ) -> Result<()> {
        Ok(self.native.execute(
            &stream.native,
            &launch.left.a.native,
            &launch.left.a_scales.native,
            &launch.left.b.native,
            &launch.left.b_scales.native,
            &launch.left.alphas.native,
            &launch.right.a.native,
            &launch.right.a_scales.native,
            &launch.right.b.native,
            &launch.right.b_scales.native,
            &launch.right.alphas.native,
            &launch.metadata.indices.native,
            &launch.metadata.rows.native,
            &launch.metadata.offsets.native,
            &launch.left.output.native,
            &launch.right.output.native,
        )?)
    }
}