ferrum_interfaces/vnext/numerical/
kv_storage.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use super::{invalid, ElementType, KvStorageFormat, ModelFamilyId, StateSpec, VNextError};
4use crate::vnext::{
5    ModelProgram, NumericalExecutionProfile, ResolvedTensorLayout, StateCapacityDemand,
6    StateCheckpointCapability, StateCheckpointContents, StateId, StateLifetime,
7    CAUSAL_PAGED_ATTENTION_F32_MASTER_INT8_KV_OPERATION_ID,
8    CAUSAL_PAGED_ATTENTION_F32_MASTER_OPERATION_ID, CAUSAL_PAGED_ATTENTION_INT8_KV_OPERATION_ID,
9    CAUSAL_PAGED_ATTENTION_OPERATION_ID, GPT_OSS_CAUSAL_PAGED_ATTENTION_OPERATION_ID,
10    HYBRID_VNORM_CAUSAL_PAGED_ATTENTION_OPERATION_ID,
11};
12use serde::{Deserialize, Serialize};
13
14/// Semantic K/V storage declared by the family, never inferred from a name or
15/// from a provider's private buffers. Each referenced state belongs to the same
16/// sequence and participates in the ordinary allocation/completion lifecycle.
17#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case", deny_unknown_fields)]
19pub enum KvStateStorage {
20    F16 {
21        state: StateId,
22    },
23    Int8PerTokenHeadF32ScaleV1 {
24        payload_state: StateId,
25        scale_state: StateId,
26    },
27}
28
29impl KvStateStorage {
30    pub const fn format(&self) -> KvStorageFormat {
31        match self {
32            Self::F16 { .. } => KvStorageFormat::F16,
33            Self::Int8PerTokenHeadF32ScaleV1 { .. } => KvStorageFormat::Int8PerTokenHeadF32ScaleV1,
34        }
35    }
36
37    pub fn payload_state(&self) -> &StateId {
38        match self {
39            Self::F16 { state } => state,
40            Self::Int8PerTokenHeadF32ScaleV1 { payload_state, .. } => payload_state,
41        }
42    }
43
44    pub fn scale_state(&self) -> Option<&StateId> {
45        match self {
46            Self::F16 { .. } => None,
47            Self::Int8PerTokenHeadF32ScaleV1 { scale_state, .. } => Some(scale_state),
48        }
49    }
50}
51
52pub(super) fn validate_kv_storage(
53    family: &ModelFamilyId,
54    declarations: &[KvStateStorage],
55    states: &[StateSpec],
56) -> Result<Option<KvStorageFormat>, VNextError> {
57    let states: BTreeMap<_, _> = states.iter().map(|state| (&state.id, state)).collect();
58    let mut referenced = BTreeSet::new();
59    let mut format = None;
60    for declaration in declarations {
61        if format.is_some_and(|format| format != declaration.format()) {
62            return Err(invalid(family, "a profile cannot mix KV storage formats"));
63        }
64        format = Some(declaration.format());
65        let mut resolve = |id: &StateId| {
66            if !referenced.insert(id.clone()) {
67                return Err(invalid(family, "KV declarations cannot share a state"));
68            }
69            states
70                .get(id)
71                .copied()
72                .ok_or_else(|| invalid(family, "KV storage references an undeclared state"))
73        };
74        let payload = resolve(declaration.payload_state())?;
75        let shape = &payload.tensor.dimensions;
76        let payload_type = match declaration.format() {
77            KvStorageFormat::F16 => ElementType::F16,
78            KvStorageFormat::Int8PerTokenHeadF32ScaleV1 => ElementType::I8,
79        };
80        if shape.len() != 3
81            || shape[0] != 2
82            || shape[1] == 0
83            || shape[2] == 0
84            || payload.tensor.element_type != payload_type
85        {
86            return Err(invalid(
87                family,
88                "KV payload must have its declared dtype and shape [2, heads, head_dim]",
89            ));
90        }
91        let maximum_tokens = validate_token_state(family, payload)?;
92        if let Some(id) = declaration.scale_state() {
93            let scale = resolve(id)?;
94            if scale.tensor.dimensions != [2, shape[1]]
95                || scale.tensor.element_type != ElementType::F32
96                || validate_token_state(family, scale)? != maximum_tokens
97                || scale.initialization != payload.initialization
98                || scale.checkpoint != payload.checkpoint
99            {
100                return Err(invalid(family, "INT8 KV scales must be F32 [2, heads] with the same token capacity, initialization and checkpoint contract as the payload"));
101            }
102        }
103    }
104    Ok(format)
105}
106
107/// Validate coverage using standard operation identities and required ports,
108/// rather than inferring KV from tensor shapes or trusting a partial state list.
109/// These ports are part of the corresponding versioned operation ABIs.
110pub(super) fn validate_program_kv_storage(
111    profile: &NumericalExecutionProfile,
112    program: &ModelProgram,
113) -> Result<(), VNextError> {
114    let by_value: BTreeMap<_, _> = program.states().iter().map(|s| (&s.value_id, s)).collect();
115    let mut consumed = BTreeSet::new();
116    for node in program.blocks().iter().flat_map(|block| &block.nodes) {
117        let (format, payload_port, scale_port) = match node.operation_id.as_str() {
118            CAUSAL_PAGED_ATTENTION_OPERATION_ID
119            | CAUSAL_PAGED_ATTENTION_F32_MASTER_OPERATION_ID
120            | HYBRID_VNORM_CAUSAL_PAGED_ATTENTION_OPERATION_ID => (KvStorageFormat::F16, 8, None),
121            GPT_OSS_CAUSAL_PAGED_ATTENTION_OPERATION_ID => (KvStorageFormat::F16, 11, None),
122            CAUSAL_PAGED_ATTENTION_INT8_KV_OPERATION_ID
123            | CAUSAL_PAGED_ATTENTION_F32_MASTER_INT8_KV_OPERATION_ID => {
124                (KvStorageFormat::Int8PerTokenHeadF32ScaleV1, 8, Some(9))
125            }
126            _ => continue,
127        };
128        let state_at = |port: usize| {
129            node.inputs
130                .get(port)
131                .and_then(|value| by_value.get(value).copied())
132                .ok_or_else(|| {
133                    invalid(
134                        &profile.family_id,
135                        format!(
136                            "operation {} requires a declared KV state at input {port}",
137                            node.operation_id
138                        ),
139                    )
140                })
141        };
142        let payload = state_at(payload_port)?;
143        let declaration = profile
144            .kv_storage
145            .iter()
146            .find(|storage| storage.payload_state() == &payload.id)
147            .ok_or_else(|| {
148                invalid(
149                    &profile.family_id,
150                    format!(
151                        "operation {} KV state {} is missing from the numerical storage contract",
152                        node.operation_id, payload.id
153                    ),
154                )
155            })?;
156        if declaration.format() != format
157            || declaration.scale_state()
158                != scale_port
159                    .map(|port| state_at(port).map(|state| &state.id))
160                    .transpose()?
161        {
162            return Err(invalid(
163                &profile.family_id,
164                format!(
165                    "operation {} state ports conflict with the declared KV format",
166                    node.operation_id
167                ),
168            ));
169        }
170        consumed.insert(declaration.payload_state());
171    }
172    if profile
173        .kv_storage
174        .iter()
175        .any(|storage| !consumed.contains(storage.payload_state()))
176    {
177        return Err(invalid(
178            &profile.family_id,
179            "KV storage is not consumed by a supported attention operation contract",
180        ));
181    }
182    Ok(())
183}
184
185#[cfg(test)]
186mod tests;
187
188fn validate_token_state(family: &ModelFamilyId, state: &StateSpec) -> Result<u64, VNextError> {
189    let tensor_bytes = state.tensor.byte_len()?;
190    state.capacity_demand.validate(tensor_bytes)?;
191    if state.lifetime != StateLifetime::Sequence
192        || state.tensor.layout != ResolvedTensorLayout::Contiguous
193    {
194        return Err(invalid(
195            family,
196            "KV storage must be contiguous Sequence state",
197        ));
198    }
199    if matches!(state.checkpoint, StateCheckpointCapability::CompletedBoundary(contract)
200        if contract.contents() != StateCheckpointContents::PrefixPositions)
201    {
202        return Err(invalid(
203            family,
204            "KV checkpoint state must retain every valid prefix position",
205        ));
206    }
207    match state.capacity_demand {
208        StateCapacityDemand::TokenScaled {
209            bytes_per_token,
210            maximum_tokens,
211        } if bytes_per_token == tensor_bytes && maximum_tokens > 0 => Ok(maximum_tokens),
212        _ => Err(invalid(
213            family,
214            "KV capacity must be exactly one typed tensor per token",
215        )),
216    }
217}
ferrum_interfaces/vnext/numerical/kv_storage.rs

ferrum_interfaces/vnext/numerical/
kv_storage.rs
