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 mod mxfp4_marlin_materializer;
48pub use marlin_expert_stack::MarlinExpertStack;
49
50pub mod quant_linear;
51
52pub mod attention;
53
54pub mod moe_host;
55
56// Audit #9: Metal GGUF k-quant kernels (q4_k_*, q6_k_*, moe_*) physically
57// live in `backend/metal/` now. Re-exported here so external callers'
58// `ferrum_kernels::q4_k_gemm::*` paths + internal `crate::q4_k_*::*` paths
59// keep working unchanged. (`moe_host` stays top-level — it's the CPU
60// reference impl used from `ferrum-models`, not a Metal kernel.)
61#[cfg(all(target_os = "macos", feature = "metal"))]
62pub use backend::metal::{
63    moe_post_ops, moe_post_ops_batched, moe_router, q4_k, q4_k_gemm, q4_k_gemv, q4_k_gemv_v2,
64    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,
65    q4_k_moe_id_gemv_batched, q6_k_gemm, q6_k_gemv, q6_k_moe_id_gemm, q6_k_moe_id_gemv,
66    q6_k_moe_id_gemv_batched,
67};
68
69#[cfg(feature = "cuda")]
70pub(crate) mod ptx {
71    // Generated by build.rs from all .cu sources. Some kernels (e.g.
72    // SOFTMAX, BATCHED_FLASH_DECODE_ATTENTION) are emitted unconditionally
73    // but only loaded behind specific code paths, so dead_code fires in
74    // configs that don't hit them.
75    #![allow(dead_code)]
76    include!(concat!(env!("OUT_DIR"), "/ptx.rs"));
77}
78
79// Audit #9: CUDA kernels physically live under `backend/cuda/` now.
80// Re-exports preserve the historical `ferrum_kernels::foo::*` public
81// surface + internal `crate::foo::*` paths.
82//
83// Two files stay at the crate root for now because they would otherwise
84// collide with same-named files already under `backend/cuda/` (which
85// host the Backend-trait impls, not the kernel launchers):
86//   - `int8_kv.rs` (top-level launchers `launch_int8_paged_decode_*`)
87//   - `quant.rs`   (top-level `dequant_int4` legacy path)
88
89#[cfg(feature = "cuda")]
90pub mod int8_kv;
91#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
92pub mod quant;
93
94#[cfg(feature = "cuda")]
95pub use backend::cuda::{cublas, decode_buffers, gpu_paged_kv, marlin, nccl_comm};
96
97#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
98pub use backend::cuda::{cuda_decode, cuda_graph, tp_decode, weight_store};
99
100#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
101pub use backend::cuda::decode_attention::decode_attention;
102#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
103pub use backend::cuda::fused_add_rms_norm::fused_add_rms_norm;
104#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
105pub use backend::cuda::fused_silu_mul::fused_silu_mul;
106#[cfg(feature = "cuda")]
107pub use backend::cuda::gated_delta_rule::recurrent_gated_delta_rule_f32;
108#[cfg(feature = "cuda")]
109pub use backend::cuda::linear_attention::{gated_rms_norm_f32, linear_attention_prepare_f32};
110#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
111pub use backend::cuda::residual_add::residual_add;
112#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
113pub use backend::cuda::rms_norm::rms_norm;
114#[cfg(all(feature = "cuda", feature = "candle-cuda-compat"))]
115pub use backend::cuda::rope::rope;
116
117// Preserve `crate::triton_ptx` / `crate::triton_meta` paths for in-crate
118// callers (e.g. `quant_linear::cuda_marlin::CudaMarlinLinear::forward`).
119// These are NOT part of the kernels-crate public API.
120#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
121pub(crate) use backend::cuda::{triton_meta, triton_ptx};
122
123#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
124pub use backend::cuda::triton_add_bias::add_bias_triton;
125#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
126pub use backend::cuda::triton_fused_add_rms_norm::fused_add_rms_norm_triton;
127#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
128pub use backend::cuda::triton_fused_silu_mul::fused_silu_mul_triton;
129#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
130pub use backend::cuda::triton_gelu::gelu_triton;
131#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
132pub use backend::cuda::triton_layer_norm::layer_norm_triton;
133#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
134pub use backend::cuda::triton_residual_add::residual_add_triton;
135#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
136pub use backend::cuda::triton_residual_add_inplace::residual_add_inplace_triton;
137#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
138pub use backend::cuda::triton_rms_norm::rms_norm_triton;
139#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
140pub use backend::cuda::triton_softmax::softmax_triton;
141#[cfg(all(feature = "cuda", feature = "triton-kernels"))]
142pub use backend::cuda::{triton_fused_moe, triton_w4a16};
143
144// vLLM gptq_marlin port (Phase 12). Behind its own feature for opt-in
145// while we validate correctness + perf vs ferrum's existing IST-DASLab Marlin.
146#[cfg(feature = "vllm-marlin")]
147pub use backend::cuda::vllm_marlin;
ferrum_kernels/lib.rs

ferrum_kernels/
lib.rs
