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

mod attention;
mod block;
mod clamped_routed;
mod decoder;
mod dense;
mod embedding;
mod gated_delta;
mod gated_full_attention;
mod hybrid_layer;
mod init;
mod kv;
mod linear;
mod model;
mod output;
mod planning;
mod prepare;
mod router;
mod runtime;
mod sampling;
mod shared_moe;
mod shared_routed;
mod task;
mod vision;

use std::sync::Arc;

pub use attention::{
    BatchedDecodeAttentionBf16, CapturedDecodeAttentionBf16, DecodeAttentionBf16,
    DecodeAttentionConfig, DecodeAttentionOutputWeight, DecodeAttentionWeights, DecodeQkvWeights,
    PrefillAttentionBf16,
};
pub use block::{
    BatchedDecodeMoeBlockBf16, BatchedDecodeMoeLayer, CapturedDecodeMoeBlockBf16,
    DecodeGraphAction, DecodeMoeBlockBf16, DecodeMoeBlockConfig, DecodeMoeBlockExecutor,
    DecodeMoeBlockWeights, DecodeMoeLayerTemplate, PrefillMoeBlockBf16,
};
pub use clamped_routed::{CudaClampedRoutedModelSession, CudaClampedRoutedModelTemplate};
pub use decoder::{RmsNormBf16, RopeBf16};
pub use dense::{
    DecodeDenseSwiGlu, DenseDownSource, DenseDownWeight, DenseGateUpSource, DenseGateUpWeights,
    DenseOutputSource, DenseQkvSource, DenseSwiGluConfig, DenseSwiGluLayerTemplate,
    DenseSwiGluWeights, DenseWeightSource, PrefillDenseSwiGlu,
};
pub use embedding::{AffineQuantizedEmbedding, Bf16Embedding};
pub use gated_delta::{
    AffineGatedDeltaLayerConfig, AffineGatedDeltaLayerWeights, CudaAffineGatedDeltaExecution,
    CudaAffineGatedDeltaLayer, CudaGatedDeltaState, GatedDeltaInputs, GatedDeltaStateConfig,
};
pub use gated_full_attention::{
    AffineGatedFullAttentionConfig, AffineGatedFullAttentionMoeLayerConfig,
    AffineGatedFullAttentionWeights, CudaAffineGatedFullAttention,
    CudaAffineGatedFullAttentionExecution, CudaAffineGatedFullAttentionMoeExecution,
    CudaAffineGatedFullAttentionMoeLayer, CudaAffineGatedFullAttentionState,
};
pub use hybrid_layer::{
    AffineGatedDeltaMoeLayerConfig, CudaAffineGatedDeltaMoeExecution, CudaAffineGatedDeltaMoeLayer,
};
pub use kv::{BatchedPagedAttentionBf16, PagedAttentionBf16, PagedDecodeBatch, PagedKvCache};
pub use linear::{
    AffineQuantizedBf16Linear, AffineQuantizedBf16Qmm, AffineQuantizedConfig,
    AffineQuantizedPairTensors, AffineQuantizedTensors, AffineQuantizedWeight, Bf16Linear,
    Bf16LinearPack, Bf16LinearPackWeights, Bf16LinearPair, Bf16LinearPairWeights, Bf16Projection,
    Bf16VectorLinear, BlockFp8LinearWeight, BucketedNvFp4MoeBf16, DirectNvFp4MoeBf16,
    Fp8ResidualLinearWeight, GatedActivation, GroupedNvFp4MoeBf16, HybridNvFp4MoeBf16,
    NvFp4Bf16Linear, NvFp4Bf16Pack, NvFp4Config, NvFp4ExpertBank, NvFp4ExpertBankConfig,
    NvFp4ExpertSource, NvFp4LinearWeight, NvFp4Tensors, ProjectionFormat,
    SelectedAffineGatedBf16Linear, SelectedAffinePairBf16Linear, SelectedAffineReduceBf16Linear,
    SelectedNvFp4LinearBf16, SelectedNvFp4MoeBf16, SelectedNvFp4TensorCoreMoeBf16,
};
use mircuda::{Compiler, Context, DeviceInfo, MemoryPool, Stream};
pub use model::{
    CudaDecodeBatch, CudaModelSessionConfig, CudaMoeModelSession, CudaMoeModelTemplate,
};
pub use output::{CudaAffineOutputHead, CudaOutputHead};
pub use planning::{
    AttentionExecution, AttentionPlan, AttentionPlanRequest, CudaAttentionPolicy,
    CudaDenseVectorPolicy, CudaDenseWeightPolicy, CudaExecutionPlanner, CudaHardwareProfile,
    CudaKernelAdmission, CudaMemoryArchitecture, CudaMoeBatchPolicy, CudaMoeFusionPolicy,
    CudaNumericalPolicy, CudaOutputHeadPolicy, CudaPlanningPolicy, DenseExecution, DensePlan,
    DensePlanRequest, DenseRole, ExecutionPhase, MoeExecution, MoePlan, MoePlanRequest,
    MoeQuantization, OutputHeadExecution, OutputHeadPlan, OutputHeadPlanRequest, PlanSource,
};
pub use router::{AffineRouterBf16, RouterBf16, RouterSelection, RouterTensors};
pub use sampling::{DeviceBatchSamplerBf16, DeviceSamplerBf16};
pub use shared_moe::{
    AffineSharedExpertMoeConfig, AffineSharedExpertMoeWeights, CudaAffineSharedExpertMoe,
    CudaAffineSharedExpertMoeExecution,
};
pub use shared_routed::{
    CudaSharedRoutedLayerState, CudaSharedRoutedModelSession, CudaSharedRoutedModelTemplate,
};
pub use task::{CudaSequenceScoringModel, CudaTextEmbeddingModel};
pub use vision::{CudaPooledVisionTower, CudaSpatialMergeVisionTower};

/// Initialized native CUDA backend resources.
#[derive(Clone, Debug)]
pub struct CudaBackend {
    inner: Arc<CudaRuntime>,
}

#[derive(Debug)]
struct CudaRuntime {
    device: DeviceInfo,
    context: Context,
    stream: Stream,
    pool: MemoryPool,
    compiler: Compiler,
    planner: CudaExecutionPlanner,
}