484 lines
17 KiB
Rust
484 lines
17 KiB
Rust
use crate::{CompiledModel, Error, ModelInput, ModelPhase, ModelStatus, Result};
|
|
use calculet_calrt::{
|
|
Calbin, ConfiguredRuntime, DeviceIo, Model, PrimitiveType, TaskSlot, Tensor, TensorBuffer,
|
|
};
|
|
use candle_core::{Device, Tensor as CandleTensor};
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::json;
|
|
use std::collections::BTreeMap;
|
|
|
|
const TOKEN_INPUT: &str = "inputs[0]";
|
|
const POSITION_INPUT: &str = "inputs[1]";
|
|
const LOGITS_OUTPUT: &str = "outputs[0]";
|
|
const CUR_SEQUENCE_CSR: &str = "cur_seq_len[0]";
|
|
const PAST_AND_CURRENT_CSR: &str = "past_kv_cur_seq_len[0]";
|
|
|
|
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
|
|
pub struct CalrtModelPlan {
|
|
pub prefill_model: String,
|
|
pub decode_model: String,
|
|
pub vocabulary_size: usize,
|
|
pub maximum_sequence_length: usize,
|
|
pub maximum_batch_size: usize,
|
|
}
|
|
|
|
impl CalrtModelPlan {
|
|
pub fn from_calbin(calbin: &Calbin) -> Result<Self> {
|
|
let prefill = select_prefill(&calbin.models)
|
|
.ok_or_else(|| Error::InvalidConfiguration("Calbin has no prefill submodel".into()))?;
|
|
let decode = calbin
|
|
.models
|
|
.iter()
|
|
.find(|model| {
|
|
let name = model.name.to_ascii_lowercase();
|
|
name.contains("decode") && !name.contains("prefill")
|
|
})
|
|
.ok_or_else(|| Error::InvalidConfiguration("Calbin has no decode submodel".into()))?;
|
|
validate_model_contract(prefill)?;
|
|
validate_model_contract(decode)?;
|
|
let vocabulary_size = output_vocabulary_size(decode)?;
|
|
let prefill_vocabulary = output_vocabulary_size(prefill)?;
|
|
if prefill_vocabulary != vocabulary_size {
|
|
return Err(Error::VocabularyMismatch {
|
|
expected: vocabulary_size,
|
|
actual: prefill_vocabulary,
|
|
});
|
|
}
|
|
let maximum_sequence_length = usize::try_from(calbin.llm.max_sequence_length)
|
|
.map_err(|_| Error::InvalidConfiguration("max_seq_len exceeds usize".into()))?;
|
|
if maximum_sequence_length == 0 {
|
|
return Err(Error::InvalidConfiguration(
|
|
"Calbin does not declare max_seq_len".into(),
|
|
));
|
|
}
|
|
let maximum_batch_size = usize::try_from(calbin.llm.max_batch_size)
|
|
.map_err(|_| Error::InvalidConfiguration("n_batch exceeds usize".into()))?;
|
|
if maximum_batch_size == 0 {
|
|
return Err(Error::InvalidConfiguration(
|
|
"Calbin does not declare n_batch".into(),
|
|
));
|
|
}
|
|
Ok(Self {
|
|
prefill_model: prefill.name.clone(),
|
|
decode_model: decode.name.clone(),
|
|
vocabulary_size,
|
|
maximum_sequence_length,
|
|
maximum_batch_size,
|
|
})
|
|
}
|
|
}
|
|
|
|
pub struct CalrtCompiledModel<D> {
|
|
runtime: ConfiguredRuntime<D>,
|
|
plan: CalrtModelPlan,
|
|
next_slot: TaskSlot,
|
|
sequence_lengths: BTreeMap<u32, usize>,
|
|
}
|
|
|
|
impl<D: std::fmt::Debug> std::fmt::Debug for CalrtCompiledModel<D> {
|
|
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
formatter
|
|
.debug_struct("CalrtCompiledModel")
|
|
.field("runtime", &self.runtime)
|
|
.field("plan", &self.plan)
|
|
.field("next_slot", &self.next_slot)
|
|
.field("sequence_lengths", &self.sequence_lengths)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl<D: DeviceIo> CalrtCompiledModel<D> {
|
|
pub fn new(runtime: ConfiguredRuntime<D>) -> Result<Self> {
|
|
let plan = CalrtModelPlan::from_calbin(runtime.calbin())?;
|
|
Ok(Self {
|
|
runtime,
|
|
plan,
|
|
next_slot: TaskSlot::Ping,
|
|
sequence_lengths: BTreeMap::new(),
|
|
})
|
|
}
|
|
|
|
pub fn deploy_parameters(&mut self) -> Result<()> {
|
|
self.runtime.deploy_parameters()?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn plan(&self) -> &CalrtModelPlan {
|
|
&self.plan
|
|
}
|
|
|
|
pub fn runtime(&self) -> &ConfiguredRuntime<D> {
|
|
&self.runtime
|
|
}
|
|
|
|
pub fn runtime_mut(&mut self) -> &mut ConfiguredRuntime<D> {
|
|
&mut self.runtime
|
|
}
|
|
|
|
fn run(&mut self, input: &ModelInput) -> Result<CandleTensor> {
|
|
validate_sequence_input(input, self.plan.maximum_sequence_length)?;
|
|
let known_length = self
|
|
.sequence_lengths
|
|
.get(&input.sequence_id)
|
|
.copied()
|
|
.unwrap_or_default();
|
|
if input.past_tokens != known_length {
|
|
return Err(Error::TensorContract(format!(
|
|
"sequence {} expected past length {known_length}, received {}",
|
|
input.sequence_id, input.past_tokens
|
|
)));
|
|
}
|
|
|
|
let model_name = match input.phase {
|
|
ModelPhase::Prefill => &self.plan.prefill_model,
|
|
ModelPhase::Decode => &self.plan.decode_model,
|
|
};
|
|
let mut inputs = self.runtime.input_buffer(model_name)?;
|
|
validate_csr_contract(&inputs)?;
|
|
fill_i32(inputs.tensor_mut(TOKEN_INPUT)?, &input.token_ids)?;
|
|
fill_i32(inputs.tensor_mut(POSITION_INPUT)?, &input.positions)?;
|
|
|
|
let mut outputs = self.runtime.output_buffer(model_name)?;
|
|
self.runtime.write_inputs(&inputs, self.next_slot)?;
|
|
|
|
// The pure-Rust CALRT intentionally returns HardwareExecutionUnavailable here until
|
|
// CCU relocation, launch, job queue completion, and device KV reset are implemented.
|
|
// Keeping this call in the real path makes the remaining boundary explicit.
|
|
self.runtime.submit(model_name, self.next_slot)?;
|
|
self.runtime.read_outputs(&mut outputs, self.next_slot)?;
|
|
let logits = decode_logits(
|
|
outputs.tensor(LOGITS_OUTPUT)?,
|
|
input,
|
|
self.plan.vocabulary_size,
|
|
)?;
|
|
|
|
self.sequence_lengths.insert(
|
|
input.sequence_id,
|
|
input
|
|
.past_tokens
|
|
.checked_add(input.token_ids.len())
|
|
.ok_or_else(|| Error::TensorContract("sequence length overflow".into()))?,
|
|
);
|
|
self.next_slot = match self.next_slot {
|
|
TaskSlot::Ping => TaskSlot::Pong,
|
|
TaskSlot::Pong => TaskSlot::Ping,
|
|
};
|
|
Ok(logits)
|
|
}
|
|
}
|
|
|
|
impl<D: DeviceIo> CompiledModel for CalrtCompiledModel<D> {
|
|
fn status(&self) -> ModelStatus {
|
|
ModelStatus {
|
|
ready: false,
|
|
reason: "Rust CALRT still lacks CCU launch, job completion, and device KV reset".into(),
|
|
details: json!({
|
|
"prefill_model": self.plan.prefill_model,
|
|
"decode_model": self.plan.decode_model,
|
|
"maximum_sequence_length": self.plan.maximum_sequence_length,
|
|
"maximum_batch_size": self.plan.maximum_batch_size,
|
|
"vocabulary_size": self.plan.vocabulary_size,
|
|
"host_pipeline_ready": true,
|
|
"hardware_submission_ready": false,
|
|
}),
|
|
}
|
|
}
|
|
|
|
fn vocabulary_size(&self) -> usize {
|
|
self.plan.vocabulary_size
|
|
}
|
|
|
|
fn maximum_sequence_length(&self) -> usize {
|
|
self.plan.maximum_sequence_length
|
|
}
|
|
|
|
fn reset_sequence(&mut self, sequence_id: u32) -> Result<()> {
|
|
self.sequence_lengths.remove(&sequence_id);
|
|
Ok(())
|
|
}
|
|
|
|
fn forward(&mut self, input: &ModelInput) -> Result<CandleTensor> {
|
|
self.run(input)
|
|
}
|
|
}
|
|
|
|
fn select_prefill(models: &[Model]) -> Option<&Model> {
|
|
models
|
|
.iter()
|
|
.filter(|model| model.name.to_ascii_lowercase().contains("prefill"))
|
|
.min_by_key(|model| {
|
|
let name = model.name.to_ascii_lowercase();
|
|
!name.contains("by_ids")
|
|
})
|
|
}
|
|
|
|
fn validate_model_contract(model: &Model) -> Result<()> {
|
|
let input_names = model
|
|
.input_tensors()
|
|
.map(|tensor| tensor.name.as_str())
|
|
.collect::<Vec<_>>();
|
|
if !input_names.contains(&TOKEN_INPUT) || !input_names.contains(&POSITION_INPUT) {
|
|
return Err(Error::TensorContract(format!(
|
|
"model {:?} must expose {TOKEN_INPUT} and {POSITION_INPUT}",
|
|
model.name
|
|
)));
|
|
}
|
|
if !model
|
|
.output_tensors()
|
|
.any(|tensor| tensor.name == LOGITS_OUTPUT)
|
|
{
|
|
return Err(Error::TensorContract(format!(
|
|
"model {:?} must expose {LOGITS_OUTPUT}",
|
|
model.name
|
|
)));
|
|
}
|
|
let csr_names = model
|
|
.csr_offsets()
|
|
.map(|(name, _)| name)
|
|
.collect::<Vec<_>>();
|
|
if !csr_names.contains(&CUR_SEQUENCE_CSR) || !csr_names.contains(&PAST_AND_CURRENT_CSR) {
|
|
return Err(Error::TensorContract(format!(
|
|
"model {:?} lacks required LLM CSR fields",
|
|
model.name
|
|
)));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn validate_csr_contract(inputs: &TensorBuffer) -> Result<()> {
|
|
if inputs.csr(CUR_SEQUENCE_CSR).is_none() || inputs.csr(PAST_AND_CURRENT_CSR).is_none() {
|
|
return Err(Error::TensorContract(
|
|
"CALRT input buffer lacks required LLM CSR offsets".into(),
|
|
));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn output_vocabulary_size(model: &Model) -> Result<usize> {
|
|
let output = model
|
|
.output_tensors()
|
|
.find(|tensor| tensor.name == LOGITS_OUTPUT)
|
|
.ok_or_else(|| Error::TensorContract(format!("model {:?} has no logits", model.name)))?;
|
|
match output.shape.as_slice() {
|
|
[_, _, vocabulary] => usize::try_from(*vocabulary)
|
|
.map_err(|_| Error::TensorContract("vocabulary exceeds usize".into())),
|
|
[_, vocabulary_blocks, _, _] => {
|
|
let bytes = usize::from(output.data_type.bit_size()).div_ceil(8);
|
|
if bytes == 0 || 512 % (16 * bytes) != 0 {
|
|
return Err(Error::TensorContract(
|
|
"unsupported tiled logits primitive type".into(),
|
|
));
|
|
}
|
|
usize::try_from(*vocabulary_blocks)
|
|
.ok()
|
|
.and_then(|blocks| blocks.checked_mul(512 / 16 / bytes))
|
|
.ok_or_else(|| Error::TensorContract("vocabulary size overflow".into()))
|
|
}
|
|
shape => Err(Error::TensorContract(format!(
|
|
"unsupported logits shape {shape:?}"
|
|
))),
|
|
}
|
|
}
|
|
|
|
fn validate_sequence_input(input: &ModelInput, maximum: usize) -> Result<()> {
|
|
if input.token_ids.is_empty() {
|
|
return Err(Error::EmptyPrompt);
|
|
}
|
|
if input.token_ids.len() != input.positions.len() {
|
|
return Err(Error::TensorContract(
|
|
"token and position counts differ".into(),
|
|
));
|
|
}
|
|
let requested =
|
|
input
|
|
.past_tokens
|
|
.checked_add(input.token_ids.len())
|
|
.ok_or(Error::ContextOverflow {
|
|
requested: usize::MAX,
|
|
maximum,
|
|
})?;
|
|
if requested > maximum {
|
|
return Err(Error::ContextOverflow { requested, maximum });
|
|
}
|
|
if input.phase == ModelPhase::Decode && input.token_ids.len() != 1 {
|
|
return Err(Error::TensorContract(
|
|
"decode accepts exactly one token".into(),
|
|
));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn fill_i32(tensor: &mut Tensor, values: &[u32]) -> Result<()> {
|
|
if !matches!(
|
|
tensor.info().data_type,
|
|
PrimitiveType::S32 | PrimitiveType::U32 | PrimitiveType::Token
|
|
) {
|
|
return Err(Error::TensorContract(format!(
|
|
"tensor {:?} must use a 32-bit token type, found {}",
|
|
tensor.info().name,
|
|
tensor.info().data_type
|
|
)));
|
|
}
|
|
let byte_size = values
|
|
.len()
|
|
.checked_mul(size_of::<u32>())
|
|
.ok_or_else(|| Error::TensorContract("input size overflow".into()))?;
|
|
if byte_size > tensor.data().len() {
|
|
return Err(Error::TensorContract(format!(
|
|
"tensor {:?} has {} bytes but needs {byte_size}",
|
|
tensor.info().name,
|
|
tensor.data().len()
|
|
)));
|
|
}
|
|
tensor.data_mut().fill(0);
|
|
for (target, value) in tensor.data_mut()[..byte_size]
|
|
.chunks_exact_mut(size_of::<u32>())
|
|
.zip(values)
|
|
{
|
|
target.copy_from_slice(&value.to_le_bytes());
|
|
}
|
|
tensor.slice(0, byte_size)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn decode_logits(tensor: &Tensor, input: &ModelInput, vocabulary: usize) -> Result<CandleTensor> {
|
|
let info = tensor.info();
|
|
let element_bytes = usize::from(info.data_type.bit_size()).div_ceil(8);
|
|
if !matches!(info.data_type, PrimitiveType::Bf16 | PrimitiveType::F32) {
|
|
return Err(Error::TensorContract(format!(
|
|
"logits must be bf16 or f32, found {}",
|
|
info.data_type
|
|
)));
|
|
}
|
|
let indices = match info.shape.as_slice() {
|
|
[_, _, declared_vocabulary] => {
|
|
let declared = usize::try_from(*declared_vocabulary)
|
|
.map_err(|_| Error::TensorContract("vocabulary exceeds usize".into()))?;
|
|
if declared < vocabulary {
|
|
return Err(Error::VocabularyMismatch {
|
|
expected: vocabulary,
|
|
actual: declared,
|
|
});
|
|
}
|
|
(0..vocabulary).collect::<Vec<_>>()
|
|
}
|
|
[sequence_blocks, vocabulary_blocks, _, _] => {
|
|
let d0 = 16_usize;
|
|
let d2 = 512_usize
|
|
.checked_div(d0 * element_bytes)
|
|
.ok_or_else(|| Error::TensorContract("invalid tiled logits type".into()))?;
|
|
let sequence_length = usize::try_from(*sequence_blocks)
|
|
.ok()
|
|
.and_then(|blocks| blocks.checked_mul(d0))
|
|
.ok_or_else(|| Error::TensorContract("sequence shape overflow".into()))?;
|
|
let row = input.token_ids.len() - 1;
|
|
if row >= sequence_length {
|
|
return Err(Error::TensorContract(
|
|
"prefill row exceeds tiled logits shape".into(),
|
|
));
|
|
}
|
|
let declared = usize::try_from(*vocabulary_blocks)
|
|
.ok()
|
|
.and_then(|blocks| blocks.checked_mul(d2))
|
|
.ok_or_else(|| Error::TensorContract("vocabulary shape overflow".into()))?;
|
|
if declared < vocabulary {
|
|
return Err(Error::VocabularyMismatch {
|
|
expected: vocabulary,
|
|
actual: declared,
|
|
});
|
|
}
|
|
let row_block = row / d0;
|
|
let row_inside = row % d0;
|
|
let row_stride = vocabulary
|
|
.checked_mul(d0)
|
|
.ok_or_else(|| Error::TensorContract("tiled row stride overflow".into()))?;
|
|
let block_stride = d0 * d2;
|
|
let base = row_block
|
|
.checked_mul(row_stride)
|
|
.and_then(|offset| offset.checked_add(row_inside * d2))
|
|
.ok_or_else(|| Error::TensorContract("tiled logits offset overflow".into()))?;
|
|
(0..vocabulary)
|
|
.map(|column| base + (column / d2) * block_stride + column % d2)
|
|
.collect()
|
|
}
|
|
shape => {
|
|
return Err(Error::TensorContract(format!(
|
|
"unsupported logits shape {shape:?}"
|
|
)));
|
|
}
|
|
};
|
|
let mut logits = Vec::with_capacity(vocabulary);
|
|
for index in indices {
|
|
let offset = index
|
|
.checked_mul(element_bytes)
|
|
.ok_or_else(|| Error::TensorContract("logits byte offset overflow".into()))?;
|
|
let bytes = tensor
|
|
.data()
|
|
.get(offset..offset + element_bytes)
|
|
.ok_or_else(|| Error::TensorContract("logits buffer is truncated".into()))?;
|
|
logits.push(match info.data_type {
|
|
PrimitiveType::Bf16 => {
|
|
let bits = u16::from_le_bytes([bytes[0], bytes[1]]);
|
|
f32::from_bits(u32::from(bits) << 16)
|
|
}
|
|
PrimitiveType::F32 => f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]),
|
|
_ => unreachable!(),
|
|
});
|
|
}
|
|
Ok(CandleTensor::from_vec(logits, vocabulary, &Device::Cpu)?)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use calculet_calrt::{BufferDirection, TensorInfo};
|
|
use std::path::Path;
|
|
|
|
#[test]
|
|
fn finds_the_captured_qwen_prefill_and_decode_contracts() {
|
|
let root = Path::new(env!("CARGO_MANIFEST_DIR")).join(
|
|
"../../npu_features/snapshot_20260801/remote/data/models/\
|
|
Qwen3-30B-A3B-dynamic-W8A8-W4AF16-full_layers_merged_2_chips_40960_fa_2026-05-22",
|
|
);
|
|
let calbin = Calbin::load(root).unwrap();
|
|
let plan = CalrtModelPlan::from_calbin(&calbin).unwrap();
|
|
assert!(plan.prefill_model.contains("prefill"));
|
|
assert!(plan.decode_model.contains("decode"));
|
|
assert_eq!(plan.vocabulary_size, 151_936);
|
|
assert_eq!(plan.maximum_sequence_length, 40_960);
|
|
}
|
|
|
|
#[test]
|
|
fn converts_calrt_bf16_logits_to_a_candle_tensor() {
|
|
let mut tensor = Tensor::new(
|
|
TensorInfo {
|
|
name: LOGITS_OUTPUT.into(),
|
|
shape: vec![1, 1, 3],
|
|
data_type: PrimitiveType::Bf16,
|
|
ping_address: 0x1000,
|
|
pong_address: 0x2000,
|
|
byte_size: 6,
|
|
},
|
|
BufferDirection::DeviceToHost,
|
|
)
|
|
.unwrap();
|
|
let values = [1.0_f32, -2.0, 3.5];
|
|
for (target, value) in tensor.data_mut().chunks_exact_mut(2).zip(values) {
|
|
target.copy_from_slice(&((value.to_bits() >> 16) as u16).to_le_bytes());
|
|
}
|
|
let logits = decode_logits(
|
|
&tensor,
|
|
&ModelInput {
|
|
sequence_id: 0,
|
|
phase: ModelPhase::Decode,
|
|
past_tokens: 0,
|
|
token_ids: vec![1],
|
|
positions: vec![0],
|
|
},
|
|
3,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(logits.to_vec1::<f32>().unwrap(), values);
|
|
}
|
|
}
|