libmir-cuda 0.3.0

use mircuda::{Compiler, DeviceBuffer, LaunchConfig, Stream, TypedKernel, bf16, cuda_kernel_files};

use super::{
    super::{
        affine::{AffineGemvSpec, compile_options},
        geometry::{narrow, product, require},
    },
    kernel::{SelectedPairFallbackKernel, SelectedPairInt4Kernel, SelectedPairInt8Kernel},
};
use crate::{Error, Result};

/// Geometry for paired projections over a device-selected expert set.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct SelectedAffinePairSpec {
    pub matrix: AffineGemvSpec,
    pub expert_count: usize,
    pub selected_count: usize,
    pub tokens: usize,
}

impl SelectedAffinePairSpec {
    pub const fn new(
        matrix: AffineGemvSpec,
        expert_count: usize,
        selected_count: usize,
    ) -> Result<Self> {
        Self::new_batch(matrix, expert_count, selected_count, 1)
    }

    pub const fn new_batch(
        matrix: AffineGemvSpec,
        expert_count: usize,
        selected_count: usize,
        tokens: usize,
    ) -> Result<Self> {
        if expert_count == 0 || selected_count == 0 || selected_count > expert_count || tokens == 0
        {
            return Err(Error::InvalidQuantizedGemv("invalid selected expert count"));
        }
        Ok(Self {
            matrix,
            expert_count,
            selected_count,
            tokens,
        })
    }
}

/// Device buffers used by one selected-expert paired projection.
pub struct SelectedAffinePairLaunch<'a> {
    pub input: &'a DeviceBuffer<bf16>,
    pub selected: &'a DeviceBuffer<u32>,
    pub gate_weight: &'a DeviceBuffer<u32>,
    pub gate_scales: &'a DeviceBuffer<bf16>,
    pub gate_biases: &'a DeviceBuffer<bf16>,
    pub up_weight: &'a DeviceBuffer<u32>,
    pub up_scales: &'a DeviceBuffer<bf16>,
    pub up_biases: &'a DeviceBuffer<bf16>,
    pub gate_output: &'a mut DeviceBuffer<bf16>,
    pub up_output: &'a mut DeviceBuffer<bf16>,
}

#[derive(Clone, Debug)]
pub struct SelectedAffinePair {
    kernel: PairKernel,
    spec: SelectedAffinePairSpec,
}

#[derive(Clone, Debug)]
enum PairKernel {
    Int4(TypedKernel<SelectedPairInt4Kernel>),
    Int8(TypedKernel<SelectedPairInt8Kernel>),
    Fallback(TypedKernel<SelectedPairFallbackKernel>),
}

impl SelectedAffinePair {
    pub fn compile(compiler: &Compiler, spec: SelectedAffinePairSpec) -> Result<Self> {
        let source = cuda_kernel_files!(
            "selected_affine_pair_bf16.cu";
            "../../../kernels/affine_packed.cuh",
            "../../../kernels/selected_affine_pair_bf16.cu",
        );
        let module = compiler.compile(source, &compile_options(spec.matrix.bits, true))?;
        let kernel = match spec.matrix.bits {
            4 => PairKernel::Int4(module.kernel()?),
            8 => PairKernel::Int8(module.kernel()?),
            2 | 3 | 5 | 6 => PairKernel::Fallback(module.kernel()?),
            _ => return Err(Error::InvalidQuantizedGemv("unsupported weight precision")),
        };
        Ok(Self { kernel, spec })
    }

    pub fn execute(
        &self,
        stream: &Stream,
        launch: &mut SelectedAffinePairLaunch<'_>,
    ) -> Result<()> {
        self.validate(launch)?;
        let matrix = self.spec.matrix;
        let config = LaunchConfig {
            grid: (
                narrow(matrix.output_features.div_ceil(8))?,
                narrow(self.spec.selected_count)?,
                narrow(self.spec.tokens)?,
            ),
            block: (32, 8, 1),
            shared_memory_bytes: 0,
        };
        let dimensions = (
            narrow(matrix.input_features)?,
            narrow(matrix.output_features)?,
            narrow(matrix.group_size)?,
            narrow(self.spec.expert_count)?,
        );
        Ok(match &self.kernel {
            PairKernel::Int4(kernel) => kernel.launch(
                stream,
                config,
                (
                    launch.input,
                    launch.selected,
                    launch.gate_weight,
                    launch.gate_scales,
                    launch.gate_biases,
                    launch.up_weight,
                    launch.up_scales,
                    launch.up_biases,
                    &mut *launch.gate_output,
                    &mut *launch.up_output,
                    dimensions.0,
                    dimensions.1,
                    dimensions.2,
                    dimensions.3,
                ),
            ),
            PairKernel::Int8(kernel) => kernel.launch(
                stream,
                config,
                (
                    launch.input,
                    launch.selected,
                    launch.gate_weight,
                    launch.gate_scales,
                    launch.gate_biases,
                    launch.up_weight,
                    launch.up_scales,
                    launch.up_biases,
                    &mut *launch.gate_output,
                    &mut *launch.up_output,
                    dimensions.0,
                    dimensions.1,
                    dimensions.2,
                    dimensions.3,
                ),
            ),
            PairKernel::Fallback(kernel) => kernel.launch(
                stream,
                config,
                (
                    launch.input,
                    launch.selected,
                    launch.gate_weight,
                    launch.gate_scales,
                    launch.gate_biases,
                    launch.up_weight,
                    launch.up_scales,
                    launch.up_biases,
                    &mut *launch.gate_output,
                    &mut *launch.up_output,
                    dimensions.0,
                    dimensions.1,
                    dimensions.2,
                    dimensions.3,
                ),
            ),
        }?)
    }

    #[must_use]
    pub const fn spec(&self) -> SelectedAffinePairSpec {
        self.spec
    }

    fn validate(&self, launch: &SelectedAffinePairLaunch<'_>) -> Result<()> {
        let matrix = self.spec.matrix;
        let layout = matrix.layout()?;
        let packed = product(layout.packed_per_matrix, self.spec.expert_count)?;
        let grouped = product(layout.groups_per_matrix, self.spec.expert_count)?;
        let routes = product(self.spec.selected_count, self.spec.tokens)?;
        let output = product(matrix.output_features, routes)?;
        require("input", product(matrix.input_features, self.spec.tokens)?, launch.input.len())?;
        require("selected experts", routes, launch.selected.len())?;
        require("gate weight", packed, launch.gate_weight.len())?;
        require("up weight", packed, launch.up_weight.len())?;
        for (name, actual) in [
            ("gate scales", launch.gate_scales.len()),
            ("gate biases", launch.gate_biases.len()),
            ("up scales", launch.up_scales.len()),
            ("up biases", launch.up_biases.len()),
        ] {
            require(name, grouped, actual)?;
        }
        require("gate output", output, launch.gate_output.len())?;
        require("up output", output, launch.up_output.len())
    }
}