libmir-cuda 0.1.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

use mircuda::{DeviceBuffer, LaunchConfig, Stream, bf16};

use super::{GroupedQuantize, NvFp4GroupedPreparation, narrow, product, scale_elements};
use crate::{Error, GatedActivation, Result};

impl NvFp4GroupedPreparation {
    #[allow(clippy::too_many_arguments)]
    pub fn quantize_gated(
        &self,
        stream: &Stream,
        gate: &DeviceBuffer<bf16>,
        up: &DeviceBuffer<bf16>,
        selected: &DeviceBuffer<u32>,
        global_scales: &DeviceBuffer<f32>,
        packed: &mut DeviceBuffer<u8>,
        scales: &mut DeviceBuffer<u8>,
        geometry: GroupedQuantize,
        activation: GatedActivation,
    ) -> Result<()> {
        if !geometry.ranked {
            return Err(Error::InvalidNvFp4("gated quantization requires ranked inputs"));
        }
        geometry.validate(gate, selected, global_scales, packed, scales)?;
        geometry.validate(up, selected, global_scales, packed, scales)?;
        let blocks = product(geometry.groups, geometry.columns / 16)?;
        let scale_stride = scale_elements(1, geometry.columns)?;
        Ok(self.gated_quantize.launch(
            stream,
            LaunchConfig {
                grid: (narrow(blocks)?, 1, 1),
                block: (32, 1, 1),
                shared_memory_bytes: 0,
            },
            (
                gate,
                up,
                selected,
                global_scales,
                packed,
                scales,
                narrow(geometry.groups)?,
                narrow(geometry.columns)?,
                narrow(scale_stride)?,
                activation.code(),
            ),
        )?)
    }
}