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

ferrum_kernels/
lib.rs
