ferrum_kernels/
lib.rs

1//! Ferrum unified compute kernels for high-performance inference.
2//!
3//! Provides the `Backend` trait and implementations for CUDA, Metal, and CPU.
4//! On CUDA builds, kernels are compiled to PTX during `cargo build` and loaded
5//! on demand at runtime.
6
7pub fn configure_native_profile_sink(
8    config: &ferrum_bench_core::ProfileSinkConfig,
9) -> std::io::Result<()> {
10    #[cfg(all(feature = "cuda", feature = "vllm-moe-marlin"))]
11    backend::cuda::marlin::configure_vllm_moe_profile_sink(config)?;
12    #[cfg(not(all(feature = "cuda", feature = "vllm-moe-marlin")))]
13    let _ = config;
14    Ok(())
15}
16
17#[cfg(feature = "cuda")]
18pub fn cuda_device_count() -> Result<usize, String> {
19    cudarc::driver::CudaContext::device_count()
20        .map_err(|error| format!("failed to query CUDA device count: {error}"))
21        .and_then(|count| {
22            usize::try_from(count)
23                .map_err(|_| format!("CUDA driver returned a negative device count: {count}"))
24        })
25}
26
27#[cfg(feature = "cuda")]
28pub fn cuda_device_name(ordinal: usize) -> Result<String, String> {
29    cudarc::driver::CudaContext::new(ordinal)
30        .map_err(|error| format!("failed to open CUDA device {ordinal}: {error}"))?
31        .name()
32        .map_err(|error| format!("failed to query CUDA device {ordinal} name: {error}"))
33}
34
35pub mod backend;
36pub mod native_ops;
37
38pub mod linear;
39pub use linear::{Linear, LinearMetadata, LinearProjectionRole};
40
41pub mod stacked_expert;
42pub use stacked_expert::StackedExpertGgufLinear;
43
44pub mod marlin_expert_stack;
45pub mod marlin_fp8_materializer;
46pub mod marlin_repack;
47pub use marlin_expert_stack::MarlinExpertStack;
48
49pub mod quant_linear;
50
51pub mod attention;
52
53pub mod moe_host;
54
55// Audit #9: Metal GGUF k-quant kernels (q4_k_*, q6_k_*, moe_*) physically
56// live in `backend/metal/` now. Re-exported here so external callers'
57// `ferrum_kernels::q4_k_gemm::*` paths + internal `crate::q4_k_*::*` paths
58// keep working unchanged. (`moe_host` stays top-level — it's the CPU
59// reference impl used from `ferrum-models`, not a Metal kernel.)
60#[cfg(all(target_os = "macos", feature = "metal"))]
61pub use backend::metal::{
62    moe_post_ops, moe_post_ops_batched, moe_router, q4_k, q4_k_gemm, q4_k_gemv, q4_k_gemv_v2,
63    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,
64    q4_k_moe_id_gemv_batched, q6_k_gemm, q6_k_gemv, q6_k_moe_id_gemm, q6_k_moe_id_gemv,
65    q6_k_moe_id_gemv_batched,
66};
67
68#[cfg(feature = "cuda")]
69pub(crate) mod ptx {
70    // Generated by build.rs from all .cu sources. Some kernels (e.g.
71    // SOFTMAX, BATCHED_FLASH_DECODE_ATTENTION) are emitted unconditionally
72    // but only loaded behind specific code paths, so dead_code fires in
73    // configs that don't hit them.
74    #![allow(dead_code)]
75    include!(concat!(env!("OUT_DIR"), "/ptx.rs"));
76}
77
78// Audit #9: CUDA kernels physically live under `backend/cuda/` now.
79// Re-exports preserve the historical `ferrum_kernels::foo::*` public
80// surface + internal `crate::foo::*` paths.
81//
82// Two files stay at the crate root for now because they would otherwise
83// collide with same-named files already under `backend/cuda/` (which
84// host the Backend-trait impls, not the kernel launchers):
85//   - `int8_kv.rs` (top-level launchers `launch_int8_paged_decode_*`)
86//   - `quant.rs`   (top-level `dequant_int4` legacy path)
87
88#[cfg(feature = "cuda")]
89pub mod int8_kv;
90#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
91pub mod quant;
92
93#[cfg(feature = "cuda")]
94pub use backend::cuda::{cublas, decode_buffers, gpu_paged_kv, marlin, nccl_comm};
95
96#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
97pub use backend::cuda::{cuda_decode, cuda_graph, tp_decode, weight_store};
98
99#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
100pub use backend::cuda::decode_attention::decode_attention;
101#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
102pub use backend::cuda::fused_add_rms_norm::fused_add_rms_norm;
103#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
104pub use backend::cuda::fused_silu_mul::fused_silu_mul;
105#[cfg(feature = "cuda")]
106pub use backend::cuda::gated_delta_rule::recurrent_gated_delta_rule_f32;
107#[cfg(feature = "cuda")]
108pub use backend::cuda::linear_attention::{gated_rms_norm_f32, linear_attention_prepare_f32};
109#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
110pub use backend::cuda::residual_add::residual_add;
111#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
112pub use backend::cuda::rms_norm::rms_norm;
113#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
114pub use backend::cuda::rope::rope;
115
116// Preserve `crate::triton_ptx` / `crate::triton_meta` paths for in-crate
117// callers (e.g. `quant_linear::cuda_marlin::CudaMarlinLinear::forward`).
118// These are NOT part of the kernels-crate public API.
119#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
120pub(crate) use backend::cuda::{triton_meta, triton_ptx};
121
122#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
123pub use backend::cuda::triton_add_bias::add_bias_triton;
124#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
125pub use backend::cuda::triton_fused_add_rms_norm::fused_add_rms_norm_triton;
126#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
127pub use backend::cuda::triton_fused_silu_mul::fused_silu_mul_triton;
128#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
129pub use backend::cuda::triton_gelu::gelu_triton;
130#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
131pub use backend::cuda::triton_layer_norm::layer_norm_triton;
132#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
133pub use backend::cuda::triton_residual_add::residual_add_triton;
134#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
135pub use backend::cuda::triton_residual_add_inplace::residual_add_inplace_triton;
136#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
137pub use backend::cuda::triton_rms_norm::rms_norm_triton;
138#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
139pub use backend::cuda::triton_softmax::softmax_triton;
140#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
141pub use backend::cuda::{triton_fused_moe, triton_w4a16};
142
143// vLLM gptq_marlin port (Phase 12). Behind its own feature for opt-in
144// while we validate correctness + perf vs ferrum's existing IST-DASLab Marlin.
145#[cfg(feature = "vllm-marlin")]
146pub use backend::cuda::vllm_marlin;
ferrum_kernels/lib.rs

ferrum_kernels/
lib.rs
