Initial commit

2026-08-02 15:26:10 +08:00
parent fcc5d31137
commit eee7fed161
76 changed files with 10686 additions and 3564 deletions
@@ -0,0 +1,22 @@
+[package]
+name = "agentos-npu"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+agentos-candle.workspace = true
+agentos-inference.workspace = true
+agentos-protocol.workspace = true
+async-trait.workspace = true
+calculet-calrt.workspace = true
+calculet-pcie-abi.workspace = true
+rustix.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+tokio.workspace = true
+url.workspace = true
+
+[lints]
+workspace = true
@@ -0,0 +1,427 @@
+//! CALCULET NPU discovery, the Rust Candle runner, and the legacy llama-server bridge.
+
+use agentos_candle::{
+    CalrtCompiledModel, CalrtModelPlan, CandleChatBackend, CandleEngine, Error as CandleError,
+    HuggingFaceTokenizer, Result as CandleResult, TokenCodec,
+};
+use agentos_inference::{
+    NtexJsonClient, SubprocessBackend, chat_completion_messages, parse_chat_completions,
+};
+use agentos_protocol::{BackendError, BackendStatus, ChatBackend, ChatRequest, ChatResponse};
+use async_trait::async_trait;
+use rustix::fs::Access;
+use serde::{Deserialize, Serialize};
+use serde_json::{Value, json};
+use std::path::{Path, PathBuf};
+use std::time::Duration;
+use url::Url;
+
+use calculet_calrt::{Calbin, ConfiguredRuntime, PcieRuntimeDevice};
+
+pub use agentos_candle::GenerationConfig;
+
+pub const DRIVER_PACKAGE_VERSION: &str = calculet_pcie_abi::DRIVER_PACKAGE_VERSION;
+pub const DRIVER_ABI_VERSION: &str = calculet_pcie_abi::DRIVER_ABI_VERSION;
+pub const CALRT_VERSION: &str = calculet_calrt::CALRT_COMPATIBLE_VERSION;
+pub const H2C_DMA_CHANNELS: u8 = calculet_pcie_abi::WRITE_CHANNELS;
+pub const C2H_DMA_CHANNELS: u8 = calculet_pcie_abi::READ_CHANNELS;
+pub const MAX_DMA_TRANSFER_BYTES: usize = calculet_pcie_abi::MAX_DMA_TRANSFER_BYTES;
+
+pub type RustCandleNpuBackend = CandleChatBackend<CalrtCompiledModel<PcieRuntimeDevice>>;
+
+#[derive(Clone, Debug)]
+pub struct CandleNpuConfig {
+    pub calbin_path: PathBuf,
+    pub tokenizer_path: PathBuf,
+    pub device_index: u8,
+    pub read_channel: u8,
+    pub write_channel: u8,
+    pub deploy_parameters: bool,
+    pub generation: GenerationConfig,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+pub struct CandleNpuManifest {
+    pub plan: CalrtModelPlan,
+    pub tokenizer_vocabulary_size: usize,
+    pub padded_logits: usize,
+    pub compatible: bool,
+}
+
+pub fn inspect_candle_npu(config: &CandleNpuConfig) -> CandleResult<CandleNpuManifest> {
+    let calbin = Calbin::load(&config.calbin_path)?;
+    let plan = CalrtModelPlan::from_calbin(&calbin)?;
+    let tokenizer = HuggingFaceTokenizer::load(&config.tokenizer_path)?;
+    let tokenizer_vocabulary_size = tokenizer.vocabulary_size();
+    let compatible = tokenizer_vocabulary_size <= plan.vocabulary_size;
+    Ok(CandleNpuManifest {
+        padded_logits: plan
+            .vocabulary_size
+            .saturating_sub(tokenizer_vocabulary_size),
+        compatible,
+        plan,
+        tokenizer_vocabulary_size,
+    })
+}
+
+pub fn open_candle_npu(config: &CandleNpuConfig) -> CandleResult<RustCandleNpuBackend> {
+    let calbin = Calbin::load(&config.calbin_path)?;
+    let tokenizer = std::sync::Arc::new(HuggingFaceTokenizer::load(&config.tokenizer_path)?);
+    let device = PcieRuntimeDevice::open_index(config.device_index)?
+        .with_channels(config.read_channel, config.write_channel)?;
+    let runtime = ConfiguredRuntime::new(device, calbin)?;
+    let mut model = CalrtCompiledModel::new(runtime)?;
+    if config.deploy_parameters {
+        model.deploy_parameters()?;
+    }
+    if tokenizer.vocabulary_size() > model.plan().vocabulary_size {
+        return Err(CandleError::VocabularyMismatch {
+            expected: model.plan().vocabulary_size,
+            actual: tokenizer.vocabulary_size(),
+        });
+    }
+    let engine = CandleEngine::new(tokenizer, model, config.generation.clone())?;
+    Ok(CandleChatBackend::new("npu_candle", engine))
+}
+
+#[derive(Clone, Debug)]
+pub struct NpuProbe {
+    pub device_root: PathBuf,
+    pub sysfs_root: PathBuf,
+    pub module_path: PathBuf,
+    pub proc_root: PathBuf,
+    pub calrt_library: Option<PathBuf>,
+}
+
+impl Default for NpuProbe {
+    fn default() -> Self {
+        Self {
+            device_root: PathBuf::from("/dev"),
+            sysfs_root: PathBuf::from("/sys/class/calculet_chardev"),
+            module_path: PathBuf::from("/sys/module/calculet_pci"),
+            proc_root: PathBuf::from("/proc/calculet_monitor"),
+            calrt_library: None,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+pub struct NpuDeviceStatus {
+    pub path: String,
+    pub readable: bool,
+    pub writable: bool,
+    pub driver_version: Option<String>,
+    pub poll_mode: Option<String>,
+    pub board_state: Option<String>,
+    pub health_status: Option<String>,
+    pub ready: bool,
+    pub warnings: Vec<String>,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+pub struct NpuStatus {
+    pub module_loaded: bool,
+    pub proc_monitor_available: bool,
+    pub devices: Vec<NpuDeviceStatus>,
+    pub calrt_library: Option<String>,
+    pub hardware_ready: bool,
+    pub known_abi: KnownAbi,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+pub struct KnownAbi {
+    pub driver_package_version: String,
+    pub driver_abi_version: String,
+    pub calrt_version: String,
+    pub h2c_dma_channels: u8,
+    pub c2h_dma_channels: u8,
+    pub max_transfer_bytes: usize,
+}
+
+impl NpuProbe {
+    pub fn probe(&self) -> NpuStatus {
+        let mut devices = Vec::new();
+        if let Ok(entries) = std::fs::read_dir(&self.device_root) {
+            for entry in entries.flatten() {
+                let name = entry.file_name();
+                let name = name.to_string_lossy();
+                if let Some(suffix) = name.strip_prefix("calculet")
+                    && !suffix.is_empty()
+                    && suffix.bytes().all(|byte| byte.is_ascii_digit())
+                {
+                    devices.push(self.probe_device(&entry.path()));
+                }
+            }
+        }
+        devices.sort_by(|left, right| left.path.cmp(&right.path));
+        let calrt_library = self
+            .calrt_library
+            .as_ref()
+            .filter(|path| path.is_file())
+            .cloned()
+            .or_else(find_calrt_library)
+            .map(|path| path.to_string_lossy().into_owned());
+        let module_loaded = self.module_path.is_dir();
+        let hardware_ready =
+            module_loaded && calrt_library.is_some() && devices.iter().any(|device| device.ready);
+        NpuStatus {
+            module_loaded,
+            proc_monitor_available: self.proc_root.join("summary").is_file(),
+            devices,
+            calrt_library,
+            hardware_ready,
+            known_abi: KnownAbi {
+                driver_package_version: DRIVER_PACKAGE_VERSION.into(),
+                driver_abi_version: DRIVER_ABI_VERSION.into(),
+                calrt_version: CALRT_VERSION.into(),
+                h2c_dma_channels: H2C_DMA_CHANNELS,
+                c2h_dma_channels: C2H_DMA_CHANNELS,
+                max_transfer_bytes: MAX_DMA_TRANSFER_BYTES,
+            },
+        }
+    }
+
+    fn probe_device(&self, path: &Path) -> NpuDeviceStatus {
+        let readable = rustix::fs::access(path, Access::READ_OK).is_ok();
+        let writable = rustix::fs::access(path, Access::WRITE_OK).is_ok();
+        let sysfs = self.sysfs_root.join(path.file_name().unwrap_or_default());
+        let driver_version = read_small(sysfs.join("driver_version"));
+        let poll_mode = read_small(sysfs.join("poll_mode"));
+        let board_state = read_small(sysfs.join("board_state"));
+        let health_status = read_small(sysfs.join("health_status"));
+        let mut warnings = Vec::new();
+        if !readable || !writable {
+            warnings.push("current process lacks read/write access".into());
+        }
+        if board_state
+            .as_deref()
+            .is_some_and(|state| state != "active")
+        {
+            warnings.push("board state is not active".into());
+        }
+        if health_status
+            .as_deref()
+            .is_some_and(|health| !health.to_ascii_lowercase().starts_with("good"))
+        {
+            warnings.push("driver health status is not good".into());
+        }
+        let ready = readable
+            && writable
+            && board_state.as_deref() == Some("active")
+            && health_status
+                .as_deref()
+                .is_some_and(|health| health.to_ascii_lowercase().starts_with("good"));
+        NpuDeviceStatus {
+            path: path.to_string_lossy().into_owned(),
+            readable,
+            writable,
+            driver_version,
+            poll_mode,
+            board_state,
+            health_status,
+            ready,
+            warnings,
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct NpuBackend {
+    bridge: SubprocessBackend,
+    probe: NpuProbe,
+}
+
+impl NpuBackend {
+    pub fn new(command: Vec<String>, timeout: Duration, probe: NpuProbe) -> Self {
+        Self {
+            bridge: SubprocessBackend {
+                name: "npu".into(),
+                command,
+                timeout,
+                serialized: true,
+            },
+            probe,
+        }
+    }
+}
+
+#[async_trait]
+impl ChatBackend for NpuBackend {
+    fn name(&self) -> &'static str {
+        "npu"
+    }
+
+    fn serialized(&self) -> bool {
+        true
+    }
+
+    fn status(&self) -> BackendStatus {
+        let bridge = self.bridge.status();
+        let npu = self.probe.probe();
+        BackendStatus {
+            name: self.name().into(),
+            ready: bridge.ready && npu.hardware_ready,
+            reason: if bridge.ready && npu.hardware_ready {
+                "NPU inference chain is ready".into()
+            } else {
+                "NPU bridge, driver, device, or calrt is unavailable".into()
+            },
+            details: json!({"bridge": bridge, "npu": npu}),
+        }
+    }
+
+    async fn chat(&self, request: &ChatRequest) -> Result<ChatResponse, BackendError> {
+        let status = self.status();
+        if !status.ready {
+            return Err(BackendError::NotReady(status.reason));
+        }
+        let npu = self.probe.probe();
+        let device = npu
+            .devices
+            .iter()
+            .find(|device| device.ready)
+            .ok_or_else(|| BackendError::NotReady("NPU device became unavailable".into()))?;
+        let mut request = request.clone();
+        request.parallel_tool_calls = false;
+        request
+            .metadata
+            .insert("backend".into(), Value::String("npu".into()));
+        request
+            .metadata
+            .insert("device_path".into(), Value::String(device.path.clone()));
+        if let Some(version) = &device.driver_version {
+            request
+                .metadata
+                .insert("driver_version".into(), Value::String(version.clone()));
+        }
+        if let Some(library) = npu.calrt_library {
+            request
+                .metadata
+                .insert("calrt_library".into(), Value::String(library));
+        }
+        self.bridge.chat(&request).await
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct LoopbackServerConfig {
+    pub endpoint: String,
+    pub model: String,
+    pub timeout: Duration,
+}
+
+pub async fn bridge_chat(
+    config: &LoopbackServerConfig,
+    request: &ChatRequest,
+) -> Result<ChatResponse, BackendError> {
+    let endpoint = validate_loopback_endpoint(&config.endpoint)?;
+    let client = NtexJsonClient::new(
+        config.timeout,
+        0,
+        concat!("agentos-npu-bridge/", env!("CARGO_PKG_VERSION")),
+        None,
+    )?;
+    let payload = json!({
+        "model": config.model,
+        "messages": chat_completion_messages(&request.messages)?,
+        "tools": request.tools.iter().map(|tool| json!({
+            "type": "function",
+            "function": {
+                "name": tool.name,
+                "description": tool.description,
+                "parameters": tool.input_schema,
+                "strict": tool.strict,
+            }
+        })).collect::<Vec<_>>(),
+        "tool_choice": request.tool_choice,
+        "parallel_tool_calls": false,
+        "max_tokens": request.max_output_tokens,
+        "temperature": 0,
+        "stream": false,
+    });
+    let response = client.post(endpoint.as_str(), &payload).await?;
+    if !response.is_success() {
+        return Err(BackendError::Transport(format!(
+            "llama-server HTTP {}: {}",
+            response.status,
+            response
+                .body
+                .pointer("/error/message")
+                .and_then(Value::as_str)
+                .unwrap_or("request rejected")
+        )));
+    }
+    parse_chat_completions(&response.body)
+}
+
+fn validate_loopback_endpoint(value: &str) -> Result<Url, BackendError> {
+    let url = Url::parse(value).map_err(|error| BackendError::NotReady(error.to_string()))?;
+    let loopback = url
+        .host_str()
+        .is_some_and(|host| matches!(host, "localhost" | "127.0.0.1" | "::1"));
+    if url.scheme() != "http"
+        || !loopback
+        || url.username() != ""
+        || url.password().is_some()
+        || url.query().is_some()
+        || url.fragment().is_some()
+        || !url.path().ends_with("/chat/completions")
+    {
+        return Err(BackendError::NotReady(
+            "NPU server URL must be a loopback HTTP chat endpoint".into(),
+        ));
+    }
+    Ok(url)
+}
+
+fn read_small(path: PathBuf) -> Option<String> {
+    let bytes = std::fs::read(path).ok()?;
+    let value = String::from_utf8_lossy(&bytes[..bytes.len().min(4096)]);
+    let value = value.trim();
+    (!value.is_empty()).then(|| value.to_owned())
+}
+
+fn find_calrt_library() -> Option<PathBuf> {
+    [
+        "/usr/lib/libcalrt-linux-x86_64.so",
+        "/usr/lib64/libcalrt-linux-x86_64.so",
+        "/usr/local/lib/libcalrt-linux-x86_64.so",
+        "/opt/calculet/lib/libcalrt-linux-x86_64.so",
+    ]
+    .into_iter()
+    .map(PathBuf::from)
+    .find(|path| path.is_file())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn bridge_never_accepts_a_remote_plaintext_endpoint() {
+        assert!(validate_loopback_endpoint("http://127.0.0.1:8031/v1/chat/completions").is_ok());
+        assert!(validate_loopback_endpoint("http://npu.example/v1/chat/completions").is_err());
+    }
+
+    #[test]
+    fn inspects_the_captured_candle_deployment_without_opening_the_device() {
+        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 manifest = inspect_candle_npu(&CandleNpuConfig {
+            tokenizer_path: root.join("__Tokenizer/tokenizer.json"),
+            calbin_path: root,
+            device_index: 0,
+            read_channel: 0,
+            write_channel: 0,
+            deploy_parameters: false,
+            generation: GenerationConfig::default(),
+        })
+        .unwrap();
+        assert!(manifest.compatible);
+        assert_eq!(manifest.plan.vocabulary_size, 151_936);
+        assert_eq!(manifest.tokenizer_vocabulary_size, 151_669);
+        assert_eq!(manifest.padded_logits, 267);
+    }
+}