ferrum_kernels/backend/types.rs
1//! Backend data types shared by the capability traits and model code.
2
3use half::{bf16, f16};
4
5use super::traits::{Backend, BackendKvDtype};
6use ferrum_interfaces::kv_dtype::{KvDtypeKind, KvFp16};
7
8/// Number of score slots to reserve for a paged-attention launch.
9///
10/// CUDA graphs freeze dynamic shared-memory size at capture time. Rounding a
11/// short live KV length up to a power-of-two bucket keeps a captured launch
12/// safe for later lengths in the same bucket without reserving the model's
13/// entire configured context window on every decode step. Above 16K, retain
14/// the exact length: rounding 16,385 scores to 32K would require 128 KiB of
15/// dynamic shared memory and make otherwise valid launches fail on GPUs with
16/// the common 96-100 KiB per-block opt-in limit.
17pub fn attention_score_capacity_bucket(kv_len: usize) -> usize {
18 let live_len = kv_len.max(1);
19 if live_len > 16_384 {
20 live_len
21 } else {
22 live_len.checked_next_power_of_two().unwrap_or(live_len)
23 }
24}
25
26/// Source dtype for a weight tensor read straight from safetensors mmap.
27///
28/// Passed to `Backend::from_weight_bytes` so each backend can choose whether
29/// to upcast to its compute dtype or store as-is.
30#[derive(Clone, Copy, Debug, PartialEq, Eq)]
31pub enum SrcDtype {
32 F32,
33 F16,
34 BF16,
35}
36
37impl SrcDtype {
38 /// Number of bytes per element in the raw on-disk representation.
39 pub const fn bytes_per_elem(self) -> usize {
40 match self {
41 SrcDtype::F32 => 4,
42 SrcDtype::F16 | SrcDtype::BF16 => 2,
43 }
44 }
45
46 /// Materialise the raw byte slice into a `Vec<f32>`. Used by the default
47 /// `Backend::from_weight_bytes` impl; fp16-preferring backends bypass it.
48 pub fn to_f32_vec(self, raw: &[u8]) -> Vec<f32> {
49 match self {
50 SrcDtype::F32 => {
51 debug_assert_eq!(raw.len() % 4, 0);
52 let n = raw.len() / 4;
53 let mut out = vec![0f32; n];
54 for i in 0..n {
55 let b = [raw[i * 4], raw[i * 4 + 1], raw[i * 4 + 2], raw[i * 4 + 3]];
56 out[i] = f32::from_le_bytes(b);
57 }
58 out
59 }
60 SrcDtype::F16 => {
61 debug_assert_eq!(raw.len() % 2, 0);
62 let n = raw.len() / 2;
63 let mut out = vec![0f32; n];
64 for i in 0..n {
65 out[i] = f16::from_le_bytes([raw[i * 2], raw[i * 2 + 1]]).to_f32();
66 }
67 out
68 }
69 SrcDtype::BF16 => {
70 debug_assert_eq!(raw.len() % 2, 0);
71 let n = raw.len() / 2;
72 let mut out = vec![0f32; n];
73 for i in 0..n {
74 out[i] = bf16::from_le_bytes([raw[i * 2], raw[i * 2 + 1]]).to_f32();
75 }
76 out
77 }
78 }
79 }
80}
81
82/// Quantization flavour discriminator for `Backend::gemm_quant`.
83///
84/// Distinct schemes need distinct kernels. Carried as a parameter so the
85/// Backend trait does not explode with one method per quantization type.
86#[derive(Clone, Debug)]
87pub enum QuantKind {
88 /// GPTQ: group-wise int4/int8 with scales + zeros (asymmetric) + optional g_idx.
89 Gptq {
90 bits: u32,
91 group_size: usize,
92 desc_act: bool,
93 },
94 /// AWQ: activation-aware int4 with scales + zeros, different packing from GPTQ.
95 Awq { bits: u32, group_size: usize },
96 /// GGUF: one of k-quants / legacy quants, fully specified by the inner type.
97 Gguf { quant_type: GgufQuantType },
98}
99
100/// GGUF quantization sub-type (expand as kernels are added).
101#[derive(Clone, Copy, Debug)]
102pub enum GgufQuantType {
103 Q4_0,
104 Q4_1,
105 Q4K,
106 Q5K,
107 Q6K,
108 Q8_0,
109}
110
111/// Packed quantized weight buffers passed to `Backend::gemm_quant`.
112///
113/// Not every field is used by every `QuantKind` — e.g. GGUF packs scales
114/// inside `qweight`, so `scales` / `zeros` may be dummies. The Backend
115/// implementation is expected to validate the shape for the kind it handles.
116pub struct QuantWeights<'a, B: Backend> {
117 pub qweight: &'a B::Buffer,
118 pub scales: Option<&'a B::Buffer>,
119 pub zeros: Option<&'a B::Buffer>,
120 pub g_idx: Option<&'a B::Buffer>,
121}
122
123/// Collective-op reduction kind for TP all_reduce.
124#[derive(Clone, Copy, Debug)]
125pub enum ReduceOp {
126 Sum,
127 Max,
128 Min,
129}
130
131/// Configuration for attention dispatch.
132#[derive(Clone, Debug)]
133pub struct AttnConfig {
134 pub num_heads: usize,
135 pub num_kv_heads: usize,
136 pub head_dim: usize,
137 pub causal: bool,
138 pub scale: f32,
139 /// Stride (in rows) between head blocks in the KV buffer.
140 /// `0` means contiguous (use `kv_len`, legacy behaviour).
141 /// Set to `cache_capacity` when flashing against a pre-allocated cache
142 /// that only has `kv_len` valid slots out of `cache_capacity`.
143 pub kv_seq_stride: usize,
144 /// Sliding-window attention size (Mistral v0.1, Gemma).
145 /// `0` = disabled (full causal attention).
146 /// `w > 0` = each query position attends to the previous `w` KV positions
147 /// (still bounded by `causal` + `pos_offset + qi + 1` as the upper end).
148 pub sliding_window: usize,
149}
150
151impl Default for AttnConfig {
152 fn default() -> Self {
153 Self {
154 num_heads: 0,
155 num_kv_heads: 0,
156 head_dim: 0,
157 causal: false,
158 scale: 1.0,
159 kv_seq_stride: 0,
160 sliding_window: 0,
161 }
162 }
163}
164
165/// Per-layer KV cache. Each model owns its own `Vec<KvCache<B, K>>` per
166/// sequence. The `K: KvDtypeKind` parameter selects the cache element
167/// type — defaults to [`KvFp16`] so existing call sites that wrote
168/// `KvCache<B>` keep compiling unchanged.
169///
170/// Two layouts are supported, selected at allocation time:
171/// 1. **Contiguous** (default): `k`/`v` are `[num_kv_heads, capacity, head_dim]`
172/// f32 buffers. `block_size == 0` and `block_table` / `context_lens` are
173/// `None`. Original ferrum layout — used when `FERRUM_METAL_PAGED_KV` is
174/// unset.
175/// 2. **Paged** (vLLM-style): `k`/`v` are `[num_blocks, num_kv_heads,
176/// block_size, head_dim]` block pools. `block_size > 0` and
177/// `block_table` (`u32[max_num_blocks_per_seq]`) + `context_lens`
178/// (`u32[1]` single-seq for now) are populated. Multi-seq sharing
179/// is a Phase 4 concern; today every paged cache_id has its own
180/// pool but the kernel-level indirection works.
181///
182/// The `K` parameter is currently a phantom-type marker — the buffer
183/// fields stay `B::Buffer` regardless. Future PRs will switch backends
184/// to `BackendKvDtype<KvInt8>` etc. and the kernel dispatch will read
185/// `K::NAME` / `K::BYTES_PER_ELEM` to pick the right append / attention
186/// kernel without any `KvCache` struct change.
187pub struct KvCache<B: Backend, K: KvDtypeKind = KvFp16> {
188 pub k: B::Buffer,
189 pub v: B::Buffer,
190 pub len: usize,
191 pub capacity: usize,
192 pub num_kv_heads: usize,
193 pub head_dim: usize,
194 /// Paged: KV positions per physical block. `0` => contiguous layout.
195 pub block_size: usize,
196 /// Paged: `[max_num_blocks_per_seq]` u32 — logical → physical block.
197 pub block_table: Option<B::Buffer>,
198 /// Paged: `[1]` u32 — current context length for the kernel to read.
199 pub context_lens: Option<B::Buffer>,
200 /// Paged: host-side mirror of the physical block indices owned by
201 /// this cache. Lets the model's release path return blocks to the
202 /// shared allocator without reading them back from device.
203 pub paged_block_indices: Vec<u32>,
204 /// Marker — KV cache element type. Zero-sized.
205 pub _kv_dtype: std::marker::PhantomData<K>,
206}
207
208/// Quantized-KV cache (Dim 5 INT8 / future FP8 paths). Sibling of
209/// [`KvCache`] for backends that store K/V in a non-FP16 element type
210/// plus per-token per-kv-head scales.
211///
212/// Why a separate struct: the FP16 `KvCache<B, K>` uses `B::Buffer`
213/// uniformly, which is FP16 on every concrete backend. Stuffing INT8
214/// storage into that buffer would require unsafe transmutes; making
215/// the FP16 struct generic over the storage type would force every
216/// existing call site (4 model files, ~20 functions) to pick up an
217/// equality-bound on the associated type. Keeping a parallel struct
218/// for INT8 is the cheaper trade — the kernel launchers in
219/// [`crate::int8_kv`] take cudarc primitives directly anyway.
220///
221/// `KStorage` and `ScaleStorage` come from `BackendKvDtype<K>::KvBuffer`
222/// and `BackendKvDtype<K>::KvScales`. On CUDA they wrap `CudaSlice<i8>`
223/// and `CudaSlice<f16>`.
224pub struct KvCacheQuant<B: BackendKvDtype<K>, K: KvDtypeKind> {
225 pub k: <B as BackendKvDtype<K>>::KvBuffer,
226 pub v: <B as BackendKvDtype<K>>::KvBuffer,
227 pub k_scales: <B as BackendKvDtype<K>>::KvScales,
228 pub v_scales: <B as BackendKvDtype<K>>::KvScales,
229 pub len: usize,
230 pub capacity: usize,
231 pub num_kv_heads: usize,
232 pub head_dim: usize,
233 pub block_size: usize,
234 pub block_table: Option<B::Buffer>,
235 pub context_lens: Option<B::Buffer>,
236 pub paged_block_indices: Vec<u32>,
237 pub _kv_dtype: std::marker::PhantomData<K>,
238}
239
240/// Routing buffers consumed by `moe_gemm_phase_vllm` — held by the
241/// caller across phase 1 and phase 3 of one MoE forward. All three
242/// fields are i32 device tensors in disguise (`Self::Buffer = fp16` on
243/// CUDA; the backend reinterprets the underlying device pointer).
244pub struct MoeRouting<B: Backend + ?Sized> {
245 pub sorted_token_ids: B::Buffer,
246 pub expert_ids: B::Buffer,
247 pub num_tokens_past_padded: B::Buffer,
248}
249
250#[cfg(test)]
251mod tests {
252 use super::attention_score_capacity_bucket;
253
254 #[test]
255 fn attention_score_capacity_tracks_live_kv_by_power_of_two() {
256 assert_eq!(attention_score_capacity_bucket(0), 1);
257 assert_eq!(attention_score_capacity_bucket(1), 1);
258 assert_eq!(attention_score_capacity_bucket(128), 128);
259 assert_eq!(attention_score_capacity_bucket(129), 256);
260 assert_eq!(attention_score_capacity_bucket(300), 512);
261 assert_eq!(attention_score_capacity_bucket(16_384), 16_384);
262 assert_eq!(attention_score_capacity_bucket(16_385), 16_385);
263 assert_eq!(attention_score_capacity_bucket(25_000), 25_000);
264 }
265}