ferrum-kernels 0.8.1

Unified compute kernels (CUDA/Metal/CPU) and model runner for Ferrum inference
Documentation 
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

//! Ferrum unified compute kernels for high-performance inference.
//!
//! Provides the `Backend` trait and implementations for CUDA, Metal, and CPU.
//! On CUDA builds, kernels are compiled to PTX during `cargo build` and loaded
//! on demand at runtime.

pub fn configure_native_profile_sink(
    config: &ferrum_bench_core::ProfileSinkConfig,
) -> std::io::Result<()> {
    #[cfg(all(feature = "cuda", feature = "vllm-moe-marlin"))]
    backend::cuda::marlin::configure_vllm_moe_profile_sink(config)?;
    #[cfg(not(all(feature = "cuda", feature = "vllm-moe-marlin")))]
    let _ = config;
    Ok(())
}

#[cfg(feature = "cuda")]
pub fn cuda_device_count() -> Result<usize, String> {
    cudarc::driver::CudaContext::device_count()
        .map_err(|error| format!("failed to query CUDA device count: {error}"))
        .and_then(|count| {
            usize::try_from(count)
                .map_err(|_| format!("CUDA driver returned a negative device count: {count}"))
        })
}

#[cfg(feature = "cuda")]
pub fn cuda_device_name(ordinal: usize) -> Result<String, String> {
    cudarc::driver::CudaContext::new(ordinal)
        .map_err(|error| format!("failed to open CUDA device {ordinal}: {error}"))?
        .name()
        .map_err(|error| format!("failed to query CUDA device {ordinal} name: {error}"))
}

pub mod backend;
pub mod native_ops;

pub mod linear;
pub use linear::{Linear, LinearMetadata, LinearProjectionRole};

pub mod stacked_expert;
pub use stacked_expert::StackedExpertGgufLinear;

pub mod marlin_expert_stack;
pub mod marlin_fp8_materializer;
pub mod marlin_repack;
pub use marlin_expert_stack::MarlinExpertStack;

pub mod quant_linear;

pub mod attention;

pub mod moe_host;

// Audit #9: Metal GGUF k-quant kernels (q4_k_*, q6_k_*, moe_*) physically
// live in `backend/metal/` now. Re-exported here so external callers'
// `ferrum_kernels::q4_k_gemm::*` paths + internal `crate::q4_k_*::*` paths
// keep working unchanged. (`moe_host` stays top-level — it's the CPU
// reference impl used from `ferrum-models`, not a Metal kernel.)
#[cfg(all(target_os = "macos", feature = "metal"))]
pub use backend::metal::{
    moe_post_ops, moe_post_ops_batched, moe_router, q4_k, q4_k_gemm, q4_k_gemv, q4_k_gemv_v2,
    q4_k_moe_id_gate_up_silu, q4_k_moe_id_gate_up_silu_batched, q4_k_moe_id_gemm, q4_k_moe_id_gemv,
    q4_k_moe_id_gemv_batched, q6_k_gemm, q6_k_gemv, q6_k_moe_id_gemm, q6_k_moe_id_gemv,
    q6_k_moe_id_gemv_batched,
};

#[cfg(feature = "cuda")]
pub(crate) mod ptx {
    // Generated by build.rs from all .cu sources. Some kernels (e.g.
    // SOFTMAX, BATCHED_FLASH_DECODE_ATTENTION) are emitted unconditionally
    // but only loaded behind specific code paths, so dead_code fires in
    // configs that don't hit them.
    #![allow(dead_code)]
    include!(concat!(env!("OUT_DIR"), "/ptx.rs"));
}

// Audit #9: CUDA kernels physically live under `backend/cuda/` now.
// Re-exports preserve the historical `ferrum_kernels::foo::*` public
// surface + internal `crate::foo::*` paths.
//
// Two files stay at the crate root for now because they would otherwise
// collide with same-named files already under `backend/cuda/` (which
// host the Backend-trait impls, not the kernel launchers):
//   - `int8_kv.rs` (top-level launchers `launch_int8_paged_decode_*`)
//   - `quant.rs`   (top-level `dequant_int4` legacy path)

#[cfg(feature = "cuda")]
pub mod int8_kv;
#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
pub mod quant;

#[cfg(feature = "cuda")]
pub use backend::cuda::{cublas, decode_buffers, gpu_paged_kv, marlin, nccl_comm};

#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
pub use backend::cuda::{cuda_decode, cuda_graph, tp_decode, weight_store};

#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
pub use backend::cuda::decode_attention::decode_attention;
#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
pub use backend::cuda::fused_add_rms_norm::fused_add_rms_norm;
#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
pub use backend::cuda::fused_silu_mul::fused_silu_mul;
#[cfg(feature = "cuda")]
pub use backend::cuda::gated_delta_rule::recurrent_gated_delta_rule_f32;
#[cfg(feature = "cuda")]
pub use backend::cuda::linear_attention::{gated_rms_norm_f32, linear_attention_prepare_f32};
#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
pub use backend::cuda::residual_add::residual_add;
#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
pub use backend::cuda::rms_norm::rms_norm;
#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
pub use backend::cuda::rope::rope;

// Preserve `crate::triton_ptx` / `crate::triton_meta` paths for in-crate
// callers (e.g. `quant_linear::cuda_marlin::CudaMarlinLinear::forward`).
// These are NOT part of the kernels-crate public API.
#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub(crate) use backend::cuda::{triton_meta, triton_ptx};

#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub use backend::cuda::triton_add_bias::add_bias_triton;
#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub use backend::cuda::triton_fused_add_rms_norm::fused_add_rms_norm_triton;
#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub use backend::cuda::triton_fused_silu_mul::fused_silu_mul_triton;
#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub use backend::cuda::triton_gelu::gelu_triton;
#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub use backend::cuda::triton_layer_norm::layer_norm_triton;
#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub use backend::cuda::triton_residual_add::residual_add_triton;
#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub use backend::cuda::triton_residual_add_inplace::residual_add_inplace_triton;
#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub use backend::cuda::triton_rms_norm::rms_norm_triton;
#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub use backend::cuda::triton_softmax::softmax_triton;
#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
pub use backend::cuda::{triton_fused_moe, triton_w4a16};

// vLLM gptq_marlin port (Phase 12). Behind its own feature for opt-in
// while we validate correctness + perf vs ferrum's existing IST-DASLab Marlin.
#[cfg(feature = "vllm-marlin")]
pub use backend::cuda::vllm_marlin;