ferrum-kernels 0.7.7

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

//! Source-built FlashAttention-2 bridge.
//!
//! This module is compiled only with the `fa2-source` cargo feature. The C ABI
//! symbol is linked from a static library built by `build.rs` from the Ferrum
//! FA2 shim plus vendored FlashAttention source templates, so runtime does not
//! need vLLM, Torch, Python, or a user-provided shim `.so`.

use std::ffi::{c_char, c_int, c_void};
use std::sync::Arc;

use cudarc::driver::{CudaSlice, CudaStream};
use ferrum_types::Result;
use half::f16;

use super::fa2_ffi::call_paged_varlen_fn;

extern "C" {
    fn ferrum_fa2_paged_varlen_fwd(
        q: *const c_void,
        k: *const c_void,
        v: *const c_void,
        out: *mut c_void,
        lse: *mut c_void,
        cu_seqlens_q: *const c_void,
        seq_lens: *const c_void,
        block_tables: *const c_void,
        num_seqs: c_int,
        total_q_tokens: c_int,
        max_q_len: c_int,
        max_kv_len: c_int,
        num_heads: c_int,
        num_kv_heads: c_int,
        head_dim: c_int,
        block_size: c_int,
        max_blocks_per_seq: c_int,
        stream: *mut c_void,
        err_buf: *mut c_char,
        err_buf_len: usize,
    ) -> c_int;
}

pub(crate) struct Fa2SourcePagedVarlenArgs<'a> {
    pub(crate) stream: &'a Arc<CudaStream>,
    pub(crate) q: &'a CudaSlice<f16>,
    pub(crate) k_pool: &'a CudaSlice<f16>,
    pub(crate) v_pool: &'a CudaSlice<f16>,
    pub(crate) out: &'a mut CudaSlice<f16>,
    pub(crate) lse: &'a mut CudaSlice<f32>,
    pub(crate) cu_seqlens_q: &'a CudaSlice<u32>,
    pub(crate) seq_lens: &'a CudaSlice<u32>,
    pub(crate) block_tables: &'a CudaSlice<u32>,
    pub(crate) num_seqs: usize,
    pub(crate) total_q_tokens: usize,
    pub(crate) max_q_len: usize,
    pub(crate) max_kv_len: usize,
    pub(crate) num_heads: usize,
    pub(crate) num_kv_heads: usize,
    pub(crate) head_dim: usize,
    pub(crate) block_size: usize,
    pub(crate) max_blocks_per_seq: usize,
}

pub(crate) fn paged_varlen_attention_fa2_source(args: Fa2SourcePagedVarlenArgs<'_>) -> Result<()> {
    let Fa2SourcePagedVarlenArgs {
        stream,
        q,
        k_pool,
        v_pool,
        out,
        lse,
        cu_seqlens_q,
        seq_lens,
        block_tables,
        num_seqs,
        total_q_tokens,
        max_q_len,
        max_kv_len,
        num_heads,
        num_kv_heads,
        head_dim,
        block_size,
        max_blocks_per_seq,
    } = args;

    unsafe {
        call_paged_varlen_fn(
            ferrum_fa2_paged_varlen_fwd,
            stream,
            q,
            k_pool,
            v_pool,
            out,
            lse,
            cu_seqlens_q,
            seq_lens,
            block_tables,
            num_seqs,
            total_q_tokens,
            max_q_len,
            max_kv_len,
            num_heads,
            num_kv_heads,
            head_dim,
            block_size,
            max_blocks_per_seq,
        )
    }
}