diff --git a/.github/workflows/queue-unsloth-layer-packages.yml b/.github/workflows/queue-unsloth-layer-packages.yml index aadfac881..38b6b922a 100644 --- a/.github/workflows/queue-unsloth-layer-packages.yml +++ b/.github/workflows/queue-unsloth-layer-packages.yml @@ -31,7 +31,7 @@ on: required: false default: "main" quant_preference: - description: "Comma-separated 4-bit quant preference" + description: "Comma-separated quant preference" required: false default: "UD-Q4_K_XL,UD-Q4_K_M,Q4_K_XL,Q4_K_M" split_candidate_vram_gib: @@ -64,6 +64,8 @@ jobs: steps: - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0 + with: + persist-credentials: false - uses: dtolnay/rust-toolchain@4cda84d5c5c54efe2404f9d843567869ab1699d4 # stable 2026-07-16 diff --git a/Cargo.lock b/Cargo.lock index c6b6e722b..310c7eb5d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7395,6 +7395,7 @@ dependencies = [ "serde_json", "sha2 0.10.9", "skippy-ffi", + "skippy-protocol", "skippy-runtime", "tokio", ] @@ -7438,6 +7439,7 @@ dependencies = [ "serde", "serde_json", "skippy-ffi", + "skippy-runtime", ] [[package]] diff --git a/Justfile b/Justfile index 48af49ecf..4d3c92c5e 100644 --- a/Justfile +++ b/Justfile @@ -266,11 +266,15 @@ skippy-quantize-release-build: # llama.cpp quantization ABI linked into the executable. [unix] skippy-quantize-standalone-build backend="cpu": - LLAMA_STAGE_BACKEND="{{ backend }}" LLAMA_STAGE_LINK_MODE=static just with-lld cargo build -p skippy-quantize + scripts/prepare-llama.sh pinned + LLAMA_STAGE_BACKEND="{{ backend }}" LLAMA_STAGE_LINK_MODE=static scripts/build-llama.sh + LLAMA_STAGE_BACKEND="{{ backend }}" LLAMA_STAGE_LINK_MODE=static just with-lld cargo build -p skippy-quantize --no-default-features [unix] skippy-quantize-standalone-release-build backend="cpu": - LLAMA_STAGE_BACKEND="{{ backend }}" LLAMA_STAGE_LINK_MODE=static just with-lld cargo build --release --locked -p skippy-quantize + scripts/prepare-llama.sh pinned + LLAMA_STAGE_BACKEND="{{ backend }}" LLAMA_STAGE_LINK_MODE=static scripts/build-llama.sh + LLAMA_STAGE_BACKEND="{{ backend }}" LLAMA_STAGE_LINK_MODE=static just with-lld cargo build --release --locked -p skippy-quantize --no-default-features # Generate a reproducible benchmark corpus for skippy bench tooling. bench-corpus tier="smoke" *ARGS="": diff --git a/crates/llama-quant-ffi/src/lib.rs b/crates/llama-quant-ffi/src/lib.rs index 40666e378..8f29e81d4 100644 --- a/crates/llama-quant-ffi/src/lib.rs +++ b/crates/llama-quant-ffi/src/lib.rs @@ -38,6 +38,7 @@ pub enum LlamaFileType { MostlyMxfp4Moe = 38, MostlyNvfp4 = 39, MostlyQ1_0 = 40, + MostlyQ2_0 = 41, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/crates/mesh-llm-cli/src/models.rs b/crates/mesh-llm-cli/src/models.rs index 135a9a220..b13d432e2 100644 --- a/crates/mesh-llm-cli/src/models.rs +++ b/crates/mesh-llm-cli/src/models.rs @@ -39,6 +39,9 @@ pub enum ModelsCommand { /// Branch or tag of mesh-llm to build in the job. #[arg(long, default_value = "main")] mesh_llm_ref: String, + /// Publish a public package marked experimental and open an unmerged HF catalog PR. + #[arg(long)] + experimental: bool, /// Explicitly keep this as a dry run. This is the default unless --confirm is set. #[arg(long)] dry_run: bool, diff --git a/crates/mesh-llm-cli/src/parser.rs b/crates/mesh-llm-cli/src/parser.rs index 7f058a82c..166ab20e6 100644 --- a/crates/mesh-llm-cli/src/parser.rs +++ b/crates/mesh-llm-cli/src/parser.rs @@ -98,7 +98,8 @@ pub fn assert_mesh_requirements_docs_examples_parse() { #[cfg(test)] mod tests { - use super::Cli; + use super::{Cli, Command}; + use crate::models::ModelsCommand; use clap::Parser; #[test] @@ -138,4 +139,29 @@ mod tests { Some(std::path::PathBuf::from("topology.json")) ); } + + #[test] + fn models_package_parses_experimental_publication() { + let cli = Cli::parse_from([ + "mesh-llm", + "models", + "package", + "unsloth/inkling-GGUF:UD-Q2_K_XL", + "--experimental", + "--dry-run", + ]); + + match cli.command.expect("models command expected") { + Command::Models { + command: + ModelsCommand::Package { + source_repo: Some(source_repo), + experimental: true, + dry_run: true, + .. + }, + } => assert_eq!(source_repo, "unsloth/inkling-GGUF:UD-Q2_K_XL"), + other => panic!("unexpected command: {other:?}"), + } + } } diff --git a/crates/mesh-llm-commands/src/model_package.rs b/crates/mesh-llm-commands/src/model_package.rs index 15c52163c..8414ac911 100644 --- a/crates/mesh-llm-commands/src/model_package.rs +++ b/crates/mesh-llm-commands/src/model_package.rs @@ -3,7 +3,7 @@ use tokio_stream::StreamExt; use ::model_package::jobs::HfJobsClient; use ::model_package::permissions; -use ::model_package::prepare::{self, DiscoveredQuant, PrepareParams}; +use ::model_package::prepare::{self, DiscoveredQuant, PrepareJob, PrepareParams}; use ::model_package::script; use serde_json::json; @@ -16,6 +16,7 @@ pub struct ModelPrepareArgs<'a> { pub flavor: &'a str, pub timeout: &'a str, pub mesh_llm_ref: &'a str, + pub experimental: bool, pub dry_run: bool, pub confirm: bool, pub follow: bool, @@ -37,6 +38,7 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> { flavor, timeout, mesh_llm_ref, + experimental, dry_run, confirm, follow, @@ -95,7 +97,13 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> { // If no quant specified, list available quants and exit. // This path doesn't need HF_TOKEN — works for public repos. if source_quant.is_none() { - return run_list_quants(&hf_client, source_repo, json).await; + return run_list_quants( + &hf_client, + source_repo, + source_model_ref.revision.as_deref(), + json, + ) + .await; } let submitting = confirm && !dry_run; @@ -116,12 +124,14 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> { eprintln!("🔍 Resolving source..."); let params = PrepareParams { source_repo: source_repo.to_string(), + source_revision: source_model_ref.revision.clone(), quant: source_quant.map(|s| s.to_string()), target: target.map(|s| s.to_string()), model_id: model_id.map(|s| s.to_string()), flavor: flavor.to_string(), timeout_seconds, mesh_llm_ref: mesh_llm_ref.to_string(), + experimental, hf_token: jobs_client .as_ref() .map(|client| client.token().to_string()), @@ -129,58 +139,7 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> { let job = prepare::resolve(&hf_client, params, &perms).await?; - // Print resolved info. - let shard_info = model_ref::split_gguf_shard_info(&job.source_file); - let shard_str = if let Some(shard) = shard_info { - format!(" ({} shards)", shard.total) - } else { - String::new() - }; - - eprintln!(" Repo: {}", job.source_repo); - eprintln!(" File: {}{}", job.source_file, shard_str); - eprintln!(); - eprintln!( - "🔑 Permissions: {} ({})", - perms.username, - if perms.is_meshllm_member { - "meshllm org member" - } else { - "not in meshllm org" - } - ); - eprintln!(" Target: {}", job.target_repo); - eprintln!( - " Catalog: meshllm/catalog ({})", - if job.catalog_create_pr { - "will open PR" - } else { - "direct commit" - } - ); - eprintln!(); - eprintln!( - "📋 Job: {}, timeout {}, mesh-llm@{}", - job.spec.flavor, - format_timeout(job.spec.timeout_seconds), - job.spec - .environment - .get("MESH_LLM_REF") - .map(|s| s.as_str()) - .unwrap_or("main") - ); - eprintln!( - " Hardware: {} {} ({})", - job.job_plan.pretty_name, - hardware_label(job.job_plan.cpu.as_deref(), job.job_plan.ram.as_deref()), - job.job_plan.selection_reason - ); - eprintln!( - " Pricing: ${:.6}/{}, max {}", - job.job_plan.unit_cost_usd, - job.job_plan.unit_label, - format_cost(job.job_plan.max_cost_usd) - ); + print_prepare_job(&job, &perms); if !submitting { let redacted = redacted_spec(&job.spec); @@ -191,9 +150,12 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> { "dryRun": true, "confirmRequired": true, "sourceRepo": job.source_repo, + "sourceRevision": job.source_revision, "sourceFile": job.source_file, + "projectors": job.projectors, "targetRepo": job.target_repo, "modelId": job.model_id, + "experimental": job.experimental, "jobPlan": job.job_plan, "spec": redacted, }))? @@ -232,9 +194,12 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> { "jobUrl": job_url, "namespace": job.namespace, "sourceRepo": job.source_repo, + "sourceRevision": job.source_revision, "sourceFile": job.source_file, + "projectors": job.projectors, "targetRepo": job.target_repo, "modelId": job.model_id, + "experimental": job.experimental, "jobPlan": job.job_plan, }))? ); @@ -251,19 +216,89 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> { Ok(()) } +fn print_prepare_job(job: &PrepareJob, perms: &permissions::PermissionCheck) { + let shard_info = model_ref::split_gguf_shard_info(&job.source_file); + let shard_str = if let Some(shard) = shard_info { + format!(" ({} shards)", shard.total) + } else { + String::new() + }; + + eprintln!(" Repo: {}", job.source_repo); + eprintln!(" Commit: {}", job.source_revision); + eprintln!(" File: {}{}", job.source_file, shard_str); + for projector in &job.projectors { + eprintln!(" MMProj: {}", projector.path); + } + eprintln!(); + eprintln!( + "🔑 Permissions: {} ({})", + perms.username, + if perms.is_meshllm_member { + "meshllm org member" + } else { + "not in meshllm org" + } + ); + eprintln!(" Target: {}", job.target_repo); + eprintln!( + " Release: {}", + if job.experimental { + "experimental (public, not cataloged until HF PR merge)" + } else { + "stable" + } + ); + eprintln!( + " Catalog: meshllm/catalog ({})", + if job.catalog_create_pr { + "will open PR" + } else { + "direct commit" + } + ); + eprintln!(); + eprintln!( + "📋 Job: {}, timeout {}, mesh-llm@{}", + job.spec.flavor, + format_timeout(job.spec.timeout_seconds), + job.spec + .environment + .get("MESH_LLM_REF") + .map(|s| s.as_str()) + .unwrap_or("main") + ); + eprintln!( + " Hardware: {} {} ({})", + job.job_plan.pretty_name, + hardware_label(job.job_plan.cpu.as_deref(), job.job_plan.ram.as_deref()), + job.job_plan.selection_reason + ); + eprintln!( + " Pricing: ${:.6}/{}, max {}", + job.job_plan.unit_cost_usd, + job.job_plan.unit_label, + format_cost(job.job_plan.max_cost_usd) + ); +} + async fn run_list_quants( client: &hf_hub::HFClient, source_repo: &str, + source_revision: Option<&str>, json_output: bool, ) -> Result<()> { - let quants = prepare::list_quants(client, source_repo).await?; + let inventory = prepare::list_inventory(client, source_repo, source_revision).await?; + let quants = inventory.quants; if json_output { println!( "{}", serde_json::to_string_pretty(&json!({ "sourceRepo": source_repo, + "sourceRevision": source_revision, "quants": quants, + "projectors": inventory.projectors, }))? ); return Ok(()); @@ -280,13 +315,20 @@ async fn run_list_quants( eprintln!(); eprintln!("Specify one as a model ref, e.g.:"); eprintln!( - " mesh-llm models package {}:{}", - source_repo, quants[0].name + " mesh-llm models package {}", + source_quant_ref(source_repo, source_revision, &quants[0].name) ); Ok(()) } +fn source_quant_ref(source_repo: &str, source_revision: Option<&str>, quant: &str) -> String { + source_revision.map_or_else( + || format!("{source_repo}:{quant}"), + |revision| format!("{source_repo}@{revision}:{quant}"), + ) +} + fn print_quant_table(quants: &[DiscoveredQuant]) { // Find the longest name for alignment. let max_name = quants.iter().map(|q| q.name.len()).max().unwrap_or(0); @@ -639,4 +681,20 @@ mod tests { fn parse_timeout_mixed() { assert_eq!(parse_timeout("1h30m45s").unwrap(), 5445); } + + #[test] + fn source_quant_ref_preserves_revision() { + assert_eq!( + source_quant_ref("poolside/Laguna-S-2.1-GGUF", Some("abc123"), "Q4_K_M"), + "poolside/Laguna-S-2.1-GGUF@abc123:Q4_K_M" + ); + } + + #[test] + fn source_quant_ref_omits_absent_revision() { + assert_eq!( + source_quant_ref("poolside/Laguna-S-2.1-GGUF", None, "Q4_K_M"), + "poolside/Laguna-S-2.1-GGUF:Q4_K_M" + ); + } } diff --git a/crates/mesh-llm-config/src/model_validation.rs b/crates/mesh-llm-config/src/model_validation.rs index a564a2d1b..0505452fb 100644 --- a/crates/mesh-llm-config/src/model_validation.rs +++ b/crates/mesh-llm-config/src/model_validation.rs @@ -9,6 +9,8 @@ use crate::model::{ SkippyConfig, SpeculativeConfig, StringOrStringList, merge_hardware, merge_model_fit, merge_multimodal, merge_throughput, }; +use skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH; + use crate::validation_support::{ looks_like_model_identifier, validate_allowed, validate_bool_or_auto, validate_hf_pair, validate_model_identifier, validate_non_empty, validate_non_negative_f64, @@ -630,7 +632,7 @@ fn validate_verify_window_controls( config.verify_window_pipeline_depth, &format!("{base_path}.verify_window_pipeline_depth"), 1, - 1_024, + u32::try_from(MAX_VERIFY_WINDOW_PIPELINE_DEPTH).expect("verify depth limit fits u32"), ) } @@ -961,6 +963,34 @@ strategy = "mystery-oracle" assert!(validate_config_diagnostics(&config).is_empty()); } + #[test] + fn verify_window_pipeline_depth_matches_native_retention_bound() { + let accepted: MeshConfig = toml::from_str( + r#" +[defaults.speculative] +verify_window_pipeline_depth = 64 +"#, + ) + .expect("bounded depth should parse"); + validate_config(&accepted).expect("native retention boundary should validate"); + + let rejected: MeshConfig = toml::from_str( + r#" +[defaults.speculative] +verify_window_pipeline_depth = 65 +"#, + ) + .expect("out-of-range depth should parse before validation"); + let diagnostics = validate_config_diagnostics(&rejected); + let text = legacy_validation_error_text(&diagnostics); + + assert!(text.contains("verify_window_pipeline_depth")); + assert!( + text.contains("between 1 and 64"), + "unexpected diagnostic: {text}" + ); + } + #[test] fn duplicate_model_with_same_profile_is_rejected() { let config: MeshConfig = toml::from_str( diff --git a/crates/mesh-llm-host-runtime/src/api/split_readiness.rs b/crates/mesh-llm-host-runtime/src/api/split_readiness.rs index 1ec3327e6..0b3382371 100644 --- a/crates/mesh-llm-host-runtime/src/api/split_readiness.rs +++ b/crates/mesh-llm-host-runtime/src/api/split_readiness.rs @@ -537,7 +537,7 @@ fn split_readiness_recommendations( { recommendations.push(format!( "Use lower-latency peers for split serving; direct stage RTT must be at or below {}ms.", - crate::mesh::MAX_SPLIT_RTT_MS + crate::mesh::max_split_rtt_ms() )); } recommendations @@ -1163,7 +1163,7 @@ mod tests { report .recommendations .iter() - .any(|item| item.contains("80ms")) + .any(|item| item.contains(&format!("{}ms", crate::mesh::max_split_rtt_ms()))) ); } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs index 572d36c14..e11dca324 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs @@ -125,7 +125,7 @@ pub(crate) fn stage0_config( }; config.kv_cache = context .family_policy - .stage_kv_cache_config_for_stage(&config); + .stage_kv_cache_config_for_package(&config, &context.package.package_dir); config } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs index 2ea9f02c4..fd822bce8 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs @@ -9,6 +9,7 @@ use crate::models::gguf::{GgufCompactMeta, scan_gguf_compact_meta}; #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) struct FamilyPolicy { pub(crate) activation_wire_dtype: StageWireDType, + pub(crate) default_kv_cache_type: Option<&'static str>, pub(crate) prefix_cache: FamilyPrefixCachePolicy, } @@ -50,6 +51,24 @@ impl FamilyPolicy { pub(crate) fn stage_kv_cache_config_for_stage( &self, config: &StageConfig, + ) -> Option { + self.stage_kv_cache_config_for_stage_with_meta(config, None) + } + + pub(crate) fn stage_kv_cache_config_for_package( + &self, + config: &StageConfig, + package_dir: &Path, + ) -> Option { + let metadata_path = package_dir.join("shared/metadata.gguf"); + let metadata = scan_gguf_compact_meta(&metadata_path); + self.stage_kv_cache_config_for_stage_with_meta(config, metadata.as_ref()) + } + + fn stage_kv_cache_config_for_stage_with_meta( + &self, + config: &StageConfig, + package_meta: Option<&GgufCompactMeta>, ) -> Option { match self.prefix_cache { FamilyPrefixCachePolicy::Disabled { .. } => None, @@ -58,7 +77,14 @@ impl FamilyPolicy { min_tokens, max_entries, } => { - let max_bytes = derive_stage_cache_max_bytes(config)?; + // Layer-package configs can be resolved before their GGUF + // parts are materialized, so there may be no scannable model + // metadata here yet. Keep the certified family cache enabled + // in that case: the resident cache still enforces its + // ctx-derived token budget, while zero means no additional + // byte cap. Disabling the cache entirely made every packaged + // model silently miss the family default. + let max_bytes = derive_stage_cache_max_bytes(config, package_meta).unwrap_or(0); // The family policy's `max_entries` is a generous // upper bound on cache cardinality. The real ceiling // is the unified KV cell pool size: each resident @@ -135,6 +161,18 @@ pub(crate) fn family_policy_for_stage_config(config: &StageConfig) -> FamilyPoli .unwrap_or_else(|| family_policy_for_model_id(&config.model_id)) } +/// Family policy derived from already-scanned GGUF metadata. +/// +/// Split topology planning uses this so it applies the same family K/V +/// defaults that stage loading will apply, instead of re-deriving a +/// size-tiered guess that can badly under-estimate the KV budget. +pub(crate) fn family_policy_for_compact_meta( + meta: &GgufCompactMeta, + model_id: Option<&str>, +) -> FamilyPolicy { + family_policy_for_gguf_meta(meta, model_id) +} + pub(crate) fn family_policy_for_model_path( path: impl AsRef, model_id: Option<&str>, @@ -159,10 +197,14 @@ fn family_policy_for_gguf_meta(meta: &GgufCompactMeta, model_id: Option<&str>) - } fn family_policy_for_capability(capability: &FamilyCapabilityRecord) -> FamilyPolicy { - family_policy_for_normalized_family_id( + let mut policy = family_policy_for_normalized_family_id( capability.family_id.as_str(), wire_dtype_from_capability(capability.default_wire_dtype), - ) + ); + if capability.family_id == "inkling" { + policy.default_kv_cache_type = Some("q4_0"); + } + policy } fn family_policy_for_model_id(model_id: &str) -> FamilyPolicy { @@ -239,6 +281,7 @@ fn family_policy_for_normalized_family_id( fn resident_kv_policy(activation_wire_dtype: StageWireDType) -> FamilyPolicy { FamilyPolicy { activation_wire_dtype, + default_kv_cache_type: None, prefix_cache: FamilyPrefixCachePolicy::Auto { payload: FamilyPrefixCachePayload::ResidentKv, min_tokens: 256, @@ -276,6 +319,7 @@ fn resident_kv_policy(activation_wire_dtype: StageWireDType) -> FamilyPolicy { fn kv_recurrent_policy(activation_wire_dtype: StageWireDType) -> FamilyPolicy { FamilyPolicy { activation_wire_dtype, + default_kv_cache_type: None, prefix_cache: FamilyPrefixCachePolicy::Auto { payload: FamilyPrefixCachePayload::KvRecurrent, min_tokens: 256, @@ -304,6 +348,7 @@ fn disabled_family_policy( ) -> FamilyPolicy { FamilyPolicy { activation_wire_dtype, + default_kv_cache_type: None, prefix_cache: FamilyPrefixCachePolicy::Disabled { reason }, } } @@ -316,7 +361,16 @@ fn wire_dtype_from_capability(dtype: WireDType) -> StageWireDType { } } -fn derive_stage_cache_max_bytes(config: &StageConfig) -> Option { +fn derive_stage_cache_max_bytes( + config: &StageConfig, + package_meta: Option<&GgufCompactMeta>, +) -> Option { + if let Some(max_bytes) = + package_meta.and_then(|meta| estimate_stage_cache_max_bytes(config, meta)) + { + return Some(max_bytes); + } + [ config.materialized_path.as_deref(), config.source_model_path.as_deref(), @@ -324,10 +378,15 @@ fn derive_stage_cache_max_bytes(config: &StageConfig) -> Option { ] .into_iter() .flatten() - .find_map(|path| scan_gguf_compact_meta(Path::new(path))) + .find_map(|path| scan_stage_cache_meta(Path::new(path))) .and_then(|meta| estimate_stage_cache_max_bytes(config, &meta)) } +fn scan_stage_cache_meta(path: &Path) -> Option { + scan_gguf_compact_meta(path) + .or_else(|| scan_gguf_compact_meta(&path.join("shared/metadata.gguf"))) +} + fn estimate_stage_cache_max_bytes(config: &StageConfig, meta: &GgufCompactMeta) -> Option { let stage_layers = config.layer_end.checked_sub(config.layer_start)?; if stage_layers == 0 { @@ -403,6 +462,8 @@ fn ggml_block_bytes(elements: u64, block_size: u64, type_size: u64) -> Option, value: &str) { + bytes.extend_from_slice(&(value.len() as u64).to_le_bytes()); + bytes.extend_from_slice(value.as_bytes()); + } + + fn push_gguf_u32(bytes: &mut Vec, key: &str, value: u32) { + push_gguf_string(bytes, key); + bytes.extend_from_slice(&4u32.to_le_bytes()); + bytes.extend_from_slice(&value.to_le_bytes()); + } + + fn push_gguf_string_kv(bytes: &mut Vec, key: &str, value: &str) { + push_gguf_string(bytes, key); + bytes.extend_from_slice(&8u32.to_le_bytes()); + push_gguf_string(bytes, value); + } + + fn write_package_metadata(package_dir: &Path) { + let mut bytes = Vec::new(); + bytes.extend_from_slice(b"GGUF"); + bytes.extend_from_slice(&2u32.to_le_bytes()); + bytes.extend_from_slice(&0i64.to_le_bytes()); + bytes.extend_from_slice(&8i64.to_le_bytes()); + push_gguf_string_kv(&mut bytes, "general.architecture", "llama"); + push_gguf_u32(&mut bytes, "llama.context_length", 8192); + push_gguf_u32(&mut bytes, "llama.embedding_length", 4096); + push_gguf_u32(&mut bytes, "llama.block_count", 32); + push_gguf_u32(&mut bytes, "llama.attention.head_count", 32); + push_gguf_u32(&mut bytes, "llama.attention.head_count_kv", 8); + push_gguf_u32(&mut bytes, "llama.attention.key_length", 128); + push_gguf_u32(&mut bytes, "llama.attention.value_length", 128); + let shared_dir = package_dir.join("shared"); + fs::create_dir_all(&shared_dir).expect("create package shared directory"); + fs::write(shared_dir.join("metadata.gguf"), bytes).expect("write package metadata"); + } + #[test] fn qwen_policy_comes_from_gguf_architecture() { let policy = family_policy_for_gguf_meta(&meta("qwen3"), None); @@ -608,7 +705,7 @@ mod tests { | "minicpm3" | "plamo" | "plamo3" | "plm" | "refact" | "smallthinker" | "smollm3" | "arcee" | "chatglm" | "codeshell" | "deci" | "xverse" | "apertus" | "bitnet" | "command_r" | "starcoder" | "ernie4_5" | "ernie4_5_moe" | "qwen" - | "jais" | "jais2" | "nemotron" | "llama4" | "mistral4" | "seed_oss" => { + | "jais" | "jais2" | "nemotron" | "llama4" | "mistral4" | "seed_oss" | "laguna" => { assert_eq!( policy.prefix_cache, FamilyPrefixCachePolicy::Auto { @@ -621,7 +718,7 @@ mod tests { } "qwen3next" | "falcon_h1" | "jamba" | "lfm2" | "mamba" | "mamba2" | "rwkv6" | "rwkv7" | "granite_hybrid" | "qwen35" | "qwen35moe" | "plamo2" | "nemotron_h" - | "nemotron_h_moe" | "lfm2moe" | "kimi_linear" => assert_eq!( + | "nemotron_h_moe" | "lfm2moe" | "kimi_linear" | "inkling" => assert_eq!( policy.prefix_cache, FamilyPrefixCachePolicy::Auto { payload: FamilyPrefixCachePayload::KvRecurrent, @@ -768,4 +865,30 @@ mod tests { assert!(estimate_stage_cache_max_bytes(&config, &kv_meta()).is_none()); } + + #[test] + fn package_metadata_enables_cache_for_remote_package_paths() { + let package_dir = tempfile::tempdir().expect("package directory"); + write_package_metadata(package_dir.path()); + let mut config = stage_config(); + config.materialized_path = None; + config.source_model_path = Some("/source/not-downloaded/model.gguf".to_string()); + config.model_path = Some("hf://mesh-llm/laguna-layers".to_string()); + let policy = FamilyPolicy { + activation_wire_dtype: StageWireDType::F16, + default_kv_cache_type: None, + prefix_cache: FamilyPrefixCachePolicy::Auto { + payload: FamilyPrefixCachePayload::ResidentKv, + min_tokens: 256, + max_entries: 16, + }, + }; + + let cache = policy + .stage_kv_cache_config_for_package(&config, package_dir.path()) + .expect("package metadata should provide the cache byte budget"); + + assert_eq!(cache.payload, StageKvCachePayload::ResidentKv); + assert_eq!(cache.max_bytes, 3_211_264); + } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs index 477a31038..00434eeac 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs @@ -43,7 +43,9 @@ use skippy_server::{ pub use certification::{ CertificationGateStatus, SkippyCertificationRequest, certify_layer_package, }; -pub(crate) use family_policy::{family_policy_for_model_path, family_policy_for_stage_config}; +pub(crate) use family_policy::{ + family_policy_for_compact_meta, family_policy_for_model_path, family_policy_for_stage_config, +}; pub(crate) use hooks::MeshAutoHookPolicy; pub(crate) use kv_cache::KvCachePolicy; pub use materialization::{ diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs index df1b58fe6..172acf54d 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs @@ -34,6 +34,7 @@ fn native_mtp_generation() -> PackageGenerationInfo { initial_window: 1, min_window: 1, max_window: 1, + pipeline_depth: None, }), proposer: Some("mtp".to_string()), primary: None, @@ -89,6 +90,7 @@ fn native_mtp_cache_generation() -> PackageGenerationInfo { initial_window: 2, min_window: 1, max_window: 6, + pipeline_depth: None, }), proposer: None, primary: Some("mtp".to_string()), @@ -131,6 +133,7 @@ fn ngram_cache_generation() -> PackageGenerationInfo { initial_window: 6, min_window: 1, max_window: 6, + pipeline_depth: None, }), proposer: Some("cache".to_string()), primary: None, @@ -173,6 +176,7 @@ fn ngram_suffix_generation() -> PackageGenerationInfo { initial_window: 32, min_window: 1, max_window: 32, + pipeline_depth: Some(2), }), proposer: Some("suffix".to_string()), primary: None, @@ -239,6 +243,7 @@ fn speculative_strategy_auto_detects_direct_gguf_native_mtp_tensors() { assert_eq!(resolved.speculative.strategy, "auto"); assert!(resolved.speculative.native_mtp_enabled); + assert_eq!(resolved.speculative.decode.verify_window.pipeline_depth, 1); let load_options = resolved .to_model_load_options(SkippyTelemetryOptions::off()) .expect("model load options should build"); @@ -451,12 +456,7 @@ strategy = "ngram-cache" #[test] fn package_suffix_strategy_resolves_as_a_standalone_proposer() { - let mesh_config = parse_config( - r#" -[defaults.speculative] -strategy = "ngram-suffix" -"#, - ); + let mesh_config = parse_config(""); let model_file = temp_model_file(); let generation = ngram_suffix_generation(); @@ -476,10 +476,12 @@ strategy = "ngram-suffix" resolved.speculative.decode.effective_strategy, "ngram-suffix" ); + assert_eq!(resolved.speculative.decode.verify_window.pipeline_depth, 2); let openai = resolved .to_embedded_openai_args(4096, true) .expect("package suffix strategy should build OpenAI args"); assert_eq!(openai.speculative_window, 48); + assert_eq!(openai.speculative.verify_window.pipeline_depth, 2); assert_eq!( openai.speculative.ngram.as_ref().map(|ngram| ngram.kind), Some(skippy_server::NgramProposerKind::Suffix) diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs index ce071f959..211184079 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs @@ -30,13 +30,12 @@ pub(crate) fn resolve_skippy_config( validate_supported_hardware_controls(&context)?; let kv_policy = KvCachePolicy::for_model_size(context.request.model_bytes); - - let model_fit = resolve_model_fit_config(&context, kv_policy)?; let hardware = resolve_hardware_config(&context)?; let family_policy = family_policy_for_model_path( &hardware.resolved_model_path, Some(context.request.model_id), ); + let model_fit = resolve_model_fit_config(&context, kv_policy, &family_policy)?; let throughput = resolve_throughput_config(&context); let skippy = resolve_execution_config(&context, family_policy.activation_wire_dtype); let speculative = resolve_speculative_config( @@ -153,6 +152,7 @@ fn validate_supported_hardware_controls(context: &ResolverContext<'_>) -> Result fn resolve_model_fit_config( context: &ResolverContext<'_>, kv_policy: KvCachePolicy, + family_policy: &super::super::family_policy::FamilyPolicy, ) -> Result { let kv = resolve_kv_defaults(context, kv_policy); let throughput = resolve_throughput_defaults(context); @@ -188,8 +188,8 @@ fn resolve_model_fit_config( .and_then(|defaults| defaults.ubatch), BUILTIN_UBATCH, ); - let cache_type_k = resolve_cache_type_k(context, &kv, kv_policy); - let cache_type_v = resolve_cache_type_v(context, &kv, kv_policy); + let cache_type_k = resolve_cache_type_k(context, &kv, kv_policy, family_policy); + let cache_type_v = resolve_cache_type_v(context, &kv, kv_policy, family_policy); let kv_offload = resolve_kv_offload(context, &kv); let flash_attention = context .model_fit @@ -237,11 +237,25 @@ fn resolve_cache_type_k( context: &ResolverContext<'_>, kv: &KvDefaults, kv_policy: KvCachePolicy, + family_policy: &super::super::family_policy::FamilyPolicy, ) -> String { + if let Some(explicit) = context + .model_fit + .and_then(|fit| non_auto_string(fit.cache_type_k.as_deref())) + { + return explicit.to_string(); + } + if let Some(family_default) = family_policy.default_kv_cache_type { + if let Some(explicit) = context + .global_model_fit + .and_then(|fit| non_auto_string(fit.cache_type_k.as_deref())) + { + return explicit.to_string(); + } + return family_default.to_string(); + } resolve_field_string( - context - .model_fit - .and_then(|fit| non_auto_string(fit.cache_type_k.as_deref())), + None, kv.model_macro .as_ref() .and_then(|defaults| defaults.cache_type_k.as_deref()), @@ -259,11 +273,25 @@ fn resolve_cache_type_v( context: &ResolverContext<'_>, kv: &KvDefaults, kv_policy: KvCachePolicy, + family_policy: &super::super::family_policy::FamilyPolicy, ) -> String { + if let Some(explicit) = context + .model_fit + .and_then(|fit| non_auto_string(fit.cache_type_v.as_deref())) + { + return explicit.to_string(); + } + if let Some(family_default) = family_policy.default_kv_cache_type { + if let Some(explicit) = context + .global_model_fit + .and_then(|fit| non_auto_string(fit.cache_type_v.as_deref())) + { + return explicit.to_string(); + } + return family_default.to_string(); + } resolve_field_string( - context - .model_fit - .and_then(|fit| non_auto_string(fit.cache_type_v.as_deref())), + None, kv.model_macro .as_ref() .and_then(|defaults| defaults.cache_type_v.as_deref()), diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs index 0320d219f..dbacb3ee9 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs @@ -563,7 +563,7 @@ fn verify_window_config(policy: &PackageWindowPolicyInfo) -> VerifyWindowConfig VerifyWindowConfig { min_tokens: policy.min_window as usize, max_tokens: policy.max_window as usize, - pipeline_depth: 1, + pipeline_depth: policy.pipeline_depth.unwrap_or(1) as usize, } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs index 17240816b..a557b51de 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs @@ -714,6 +714,102 @@ fn family_policy_beats_builtin_wire_dtype_when_config_is_unset() { assert_eq!(resolved.skippy.activation_wire_dtype, StageWireDType::F32); } +#[test] +fn inkling_family_defaults_to_f32_wire_and_q4_kv() { + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &MeshConfig::default(), + model_id: "meshllm/inkling-UD-Q2_K_XL-layers", + model_path: Path::new("/models/inkling.gguf"), + model_bytes: 316 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + }) + .unwrap(); + + assert_eq!(resolved.skippy.activation_wire_dtype, StageWireDType::F32); + assert_eq!(resolved.model_fit.cache_type_k, "q4_0"); + assert_eq!(resolved.model_fit.cache_type_v, "q4_0"); +} + +#[test] +fn inkling_family_kv_default_beats_generic_saver_macro() { + let mesh_config = parse_config( + r#" +[[models]] +model = "meshllm/inkling-UD-Q2_K_XL-layers" + +[models.model_fit] +kv_cache_policy = "saver" +"#, + ); + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &mesh_config, + model_id: "meshllm/inkling-UD-Q2_K_XL-layers", + model_path: Path::new("/models/inkling.gguf"), + model_bytes: 316 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + }) + .unwrap(); + + assert_eq!(resolved.model_fit.cache_type_k, "q4_0"); + assert_eq!(resolved.model_fit.cache_type_v, "q4_0"); +} + +#[test] +fn explicit_inkling_kv_override_beats_family_default() { + let mesh_config = parse_config( + r#" +[[models]] +model = "meshllm/inkling-UD-Q2_K_XL-layers" +cache_type_k = "q8_0" +cache_type_v = "q8_0" +"#, + ); + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &mesh_config, + model_id: "meshllm/inkling-UD-Q2_K_XL-layers", + model_path: Path::new("/models/inkling.gguf"), + model_bytes: 316 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + }) + .unwrap(); + + assert_eq!(resolved.model_fit.cache_type_k, "q8_0"); + assert_eq!(resolved.model_fit.cache_type_v, "q8_0"); +} + +#[test] +fn explicit_global_inkling_kv_override_beats_family_default() { + let mesh_config = parse_config( + r#" +[defaults.model_fit] +cache_type_k = "q8_0" +cache_type_v = "q8_0" + +[[models]] +model = "meshllm/inkling-UD-Q2_K_XL-layers" +"#, + ); + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &mesh_config, + model_id: "meshllm/inkling-UD-Q2_K_XL-layers", + model_path: Path::new("/models/inkling.gguf"), + model_bytes: 316 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + }) + .unwrap(); + + assert_eq!(resolved.model_fit.cache_type_k, "q8_0"); + assert_eq!(resolved.model_fit.cache_type_v, "q8_0"); +} + #[test] fn family_policy_wires_prefix_cache_by_default_for_supported_models() { let model_file = temp_model_file(); @@ -870,6 +966,43 @@ fn layer_package_translation_does_not_treat_hf_ref_as_direct_gguf() { assert_eq!(options.config.load_mode, LoadMode::LayerPackage); assert_eq!(options.config.model_path.as_deref(), Some(package_ref)); + let kv_cache = options + .config + .kv_cache + .expect("packaged supported family should retain its cache policy"); + assert_eq!(kv_cache.mode, StageKvCacheMode::LookupRecord); + assert_eq!(kv_cache.payload, StageKvCachePayload::ResidentKv); + assert_eq!(kv_cache.max_bytes, 0); +} + +#[test] +fn inkling_layer_package_retains_recurrent_cache_policy_before_materialization() { + let config = MeshConfig::default(); + let package_ref = "hf://meshllm/inkling-UD-Q2_K_XL-layers"; + let resolved = resolve_skippy_config(SkippyConfigResolveRequest { + mesh_config: &config, + model_id: "meshllm/inkling-UD-Q2_K_XL-layers", + model_path: Path::new(package_ref), + model_bytes: 316 * 1024 * 1024 * 1024, + allocatable_memory_bytes: None, + request_defaults: None, + package_generation: None, + }) + .expect("Inkling package config should resolve"); + + let stage = resolved + .to_stage_config(Some(fake_hf_package_identity(66)), LoadMode::LayerPackage) + .expect("Inkling package stage config should build"); + let kv_cache = stage + .kv_cache + .expect("Inkling package should retain its recurrent cache policy"); + + assert_eq!(kv_cache.mode, StageKvCacheMode::LookupRecord); + assert_eq!(kv_cache.payload, StageKvCachePayload::KvRecurrent); + assert!(kv_cache.max_entries > 0); + assert!(kv_cache.max_entries <= 16); + assert_eq!(kv_cache.max_bytes, 0); + assert_eq!(kv_cache.min_tokens, 256); } #[test] diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs index db2da2f8f..8200e6ecf 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs @@ -1,6 +1,7 @@ use std::{ collections::HashMap, net::SocketAddr, + path::Path, sync::{ Arc, atomic::{AtomicBool, Ordering}, @@ -738,7 +739,12 @@ fn stage_config( downstream: load.downstream.as_ref().map(peer_config), }; let family_policy = super::family_policy_for_stage_config(&config); - config.kv_cache = family_policy.stage_kv_cache_config_for_stage(&config); + config.kv_cache = package.map_or_else( + || family_policy.stage_kv_cache_config_for_stage(&config), + |package| { + family_policy.stage_kv_cache_config_for_package(&config, Path::new(&package.local_ref)) + }, + ); Ok(config) } diff --git a/crates/mesh-llm-host-runtime/src/mesh/capacity.rs b/crates/mesh-llm-host-runtime/src/mesh/capacity.rs new file mode 100644 index 000000000..1ace29d8a --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/mesh/capacity.rs @@ -0,0 +1,145 @@ +use crate::system::hardware::HardwareSurvey; + +pub(super) fn mesh_capacity_bytes(hw: &HardwareSurvey) -> u64 { + let unified_memory_only = + hw.is_soc && (hw.gpus.is_empty() || hw.gpus.iter().all(|gpu| gpu.unified_memory)); + if unified_memory_only { + return hw.vram_bytes; + } + + let gpu_capacity = hw + .gpus + .iter() + .map(|gpu| mesh_llm_system::vram::allocatable_bytes(gpu.vram_bytes, gpu.reserved_bytes)) + .sum(); + if gpu_capacity > 0 { + return gpu_capacity; + } + + let legacy_gpu_capacity = hw + .gpu_vram + .iter() + .enumerate() + .map(|(index, &vram)| { + mesh_llm_system::vram::allocatable_bytes( + vram, + hw.gpu_reserved.get(index).copied().flatten(), + ) + }) + .sum(); + if legacy_gpu_capacity > 0 { + legacy_gpu_capacity + } else { + // A non-SoC node without enumerated accelerator memory cannot host a + // GPU stage. Keep the broader RAM/offload budget local-only instead + // of advertising it as accelerator capacity. + 0 + } +} + +pub(super) fn capped_capacity_bytes(capacity_bytes: u64, max_vram_gb: Option) -> u64 { + max_vram_gb + .map(|cap| capacity_bytes.min((cap * 1e9) as u64)) + .unwrap_or(capacity_bytes) +} + +pub(super) fn advertised_capacity_bytes(hw: &HardwareSurvey, max_vram_gb: Option) -> u64 { + let detected = mesh_capacity_bytes(hw); + match (detected, max_vram_gb) { + (0, Some(cap)) => hw.vram_bytes.min((cap * 1e9) as u64), + _ => capped_capacity_bytes(detected, max_vram_gb), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mesh::{NodeRole, node::hardware_snapshot_for_start}; + use crate::system::hardware::GpuFacts; + + fn gpu(vram_bytes: u64, reserved_bytes: Option, unified_memory: bool) -> GpuFacts { + GpuFacts { + vram_bytes, + reserved_bytes, + unified_memory, + ..GpuFacts::default() + } + } + + #[test] + fn discrete_gpu_mesh_capacity_excludes_host_ram_offload_budget() { + let hw = HardwareSurvey { + vram_bytes: 491_000_000_000, + gpu_vram: vec![40_000_000_000], + gpu_reserved: vec![Some(1_000_000_000)], + gpus: vec![gpu(40_000_000_000, Some(1_000_000_000), false)], + ..HardwareSurvey::default() + }; + + let snapshot = hardware_snapshot_for_start(hw, &NodeRole::Worker, None); + + assert_eq!(snapshot.vram_bytes, 39_000_000_000); + assert_eq!(snapshot.local_runtime_capacity_bytes, 491_000_000_000); + } + + #[test] + fn unified_memory_mesh_capacity_keeps_recommended_working_set() { + let hw = HardwareSurvey { + vram_bytes: 96_000_000_000, + is_soc: true, + gpu_vram: vec![128_000_000_000], + gpu_reserved: vec![Some(16_000_000_000)], + gpus: vec![gpu(128_000_000_000, Some(16_000_000_000), true)], + ..HardwareSurvey::default() + }; + + let snapshot = hardware_snapshot_for_start(hw, &NodeRole::Worker, None); + + assert_eq!(snapshot.vram_bytes, 96_000_000_000); + assert_eq!(snapshot.local_runtime_capacity_bytes, 96_000_000_000); + } + + #[test] + fn missing_discrete_gpu_facts_do_not_advertise_host_ram_as_stage_capacity() { + let hw = HardwareSurvey { + vram_bytes: 491_000_000_000, + is_soc: false, + ..HardwareSurvey::default() + }; + + let snapshot = hardware_snapshot_for_start(hw, &NodeRole::Worker, None); + + assert_eq!(snapshot.vram_bytes, 0); + assert_eq!(snapshot.local_runtime_capacity_bytes, 491_000_000_000); + } + + #[test] + fn explicit_cpu_budget_advertises_bounded_stage_capacity() { + let hw = HardwareSurvey { + vram_bytes: 16_000_000_000, + is_soc: false, + ..HardwareSurvey::default() + }; + + let snapshot = hardware_snapshot_for_start(hw, &NodeRole::Worker, Some(1.0)); + + assert_eq!(snapshot.vram_bytes, 1_000_000_000); + assert_eq!(snapshot.local_runtime_capacity_bytes, 1_000_000_000); + } + + #[test] + fn max_vram_caps_mesh_and_local_runtime_capacities() { + let hw = HardwareSurvey { + vram_bytes: 491_000_000_000, + gpu_vram: vec![40_000_000_000], + gpu_reserved: vec![Some(1_000_000_000)], + gpus: vec![gpu(40_000_000_000, Some(1_000_000_000), false)], + ..HardwareSurvey::default() + }; + + let snapshot = hardware_snapshot_for_start(hw, &NodeRole::Worker, Some(32.0)); + + assert_eq!(snapshot.vram_bytes, 32_000_000_000); + assert_eq!(snapshot.local_runtime_capacity_bytes, 32_000_000_000); + } +} diff --git a/crates/mesh-llm-host-runtime/src/mesh/connections.rs b/crates/mesh-llm-host-runtime/src/mesh/connections.rs index 1616662cc..b4c5018e0 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/connections.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/connections.rs @@ -2,7 +2,10 @@ use super::*; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; pub(crate) struct NodeHardwareSnapshot { + /// Accelerator-resident capacity advertised for mesh stage placement. pub(crate) vram_bytes: u64, + /// Broader local fit budget, which may include CPU offload memory. + pub(crate) local_runtime_capacity_bytes: u64, pub(crate) gpu_name: Option, pub(crate) hostname: Option, pub(crate) is_soc: Option, @@ -31,6 +34,10 @@ pub(crate) struct AcceptedMeshStream { pub(crate) const MAX_CONTROL_STREAM_WORK_PER_CONNECTION: usize = 32; const MESH_STREAM_TYPE_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); +/// Grace period for a replaced peer connection to finish in-flight streams +/// before it is explicitly closed. +const REPLACED_CONNECTION_DRAIN_GRACE: std::time::Duration = std::time::Duration::from_secs(5); + pub(crate) fn control_stream_semaphore() -> Arc { Arc::new(tokio::sync::Semaphore::new( MAX_CONTROL_STREAM_WORK_PER_CONNECTION, @@ -695,10 +702,38 @@ impl Node { remote.fmt_short() )); } - state.connections.insert(remote, conn.clone()); + let replaced = state.connections.insert(remote, conn.clone()); + drop(state); + if let Some(replaced) = replaced + && replaced.stable_id() != conn.stable_id() + { + Self::spawn_replaced_connection_drain(remote, replaced); + } (was_dead, admitted) } + /// Retire a connection that a newer one just replaced. + /// + /// Existing streams get a bounded grace period to finish, then the + /// connection is explicitly closed. Without the terminal close a replaced + /// connection is preserved indefinitely; without the grace period in-flight + /// requests are cut off mid-stream. + pub(super) fn spawn_replaced_connection_drain(remote: EndpointId, replaced: Connection) { + tokio::spawn(async move { + tracing::debug!( + peer = %remote.fmt_short(), + drain_ms = REPLACED_CONNECTION_DRAIN_GRACE.as_millis(), + "draining replaced peer connection" + ); + tokio::select! { + _ = replaced.closed() => {} + _ = tokio::time::sleep(REPLACED_CONNECTION_DRAIN_GRACE) => { + replaced.close(0u32.into(), b"connection-replaced"); + } + } + }); + } + pub(crate) fn spawn_reconnect_gossip(&self, conn: Connection, remote: EndpointId) { let node = self.clone(); tokio::spawn(async move { diff --git a/crates/mesh-llm-host-runtime/src/mesh/direct_path.rs b/crates/mesh-llm-host-runtime/src/mesh/direct_path.rs index 44b86b84b..f6d8047a4 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/direct_path.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/direct_path.rs @@ -220,8 +220,26 @@ impl Node { tokio::spawn(async move { node.dispatch_streams(conn_for_dispatch, remote).await; }); - if let Some(existing) = existing { - existing.close(0u32.into(), b"direct-path-replaced"); + if let Some(existing) = existing.filter(|existing| existing.stable_id() != conn.stable_id()) + { + record_draining_replaced_connection(remote, Some(&existing), &conn); + Self::spawn_replaced_connection_drain(remote, existing); } } } + +fn record_draining_replaced_connection( + remote: EndpointId, + existing: Option<&Connection>, + replacement: &Connection, +) { + let Some(existing) = existing else { + return; + }; + tracing::debug!( + peer = %remote.fmt_short(), + replaced_stable_id = existing.stable_id(), + replacement_stable_id = replacement.stable_id(), + "Direct path connection replaced; allowing existing streams to drain" + ); +} diff --git a/crates/mesh-llm-host-runtime/src/mesh/mod.rs b/crates/mesh-llm-host-runtime/src/mesh/mod.rs index 9a4d0681f..67bf98877 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/mod.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/mod.rs @@ -5,10 +5,13 @@ //! mixed-version compatibility. Skippy activation transport remains on the //! latency-sensitive `skippy-stage/2` ALPN. +#[cfg(test)] +pub(crate) use mesh_llm_types::mesh::MAX_SPLIT_RTT_MS; pub use mesh_llm_types::mesh::{ - MAX_SPLIT_RTT_MS, ModelDemand, ModelRuntimeDescriptor, ModelSourceKind, ServedModelDescriptor, + ModelDemand, ModelRuntimeDescriptor, ModelSourceKind, ServedModelDescriptor, ServedModelIdentity, ServedModelMetadata, infer_available_model_descriptors, - infer_local_served_model_descriptor, infer_served_model_descriptors, + infer_local_served_model_descriptor, infer_served_model_descriptors, max_split_rtt_ms, + split_allow_relay_paths, }; use anyhow::{Context, Result}; @@ -76,6 +79,7 @@ pub(crate) fn elapsed_ms_u64(duration: std::time::Duration) -> u64 { } mod artifact_transfer_io; +mod capacity; mod connection_reservation; mod connections; mod direct_path; diff --git a/crates/mesh-llm-host-runtime/src/mesh/node.rs b/crates/mesh-llm-host-runtime/src/mesh/node.rs index 1de6abcb4..f871aeae8 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/node.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/node.rs @@ -1,19 +1,15 @@ use super::*; -use crate::mesh::identity_persistence::load_or_create_key; use mesh_llm_types::mesh::{DEMAND_TTL_SECS, merge_demand}; use serde_json::json; use std::net::SocketAddr; -pub fn detect_vram_bytes_capped(max_vram_gb: Option) -> u64 { - let mut detected = crate::system::hardware::survey().vram_bytes; - if let Some(cap) = max_vram_gb { - let cap_bytes = (cap * 1e9) as u64; - if cap_bytes < detected { - detected = cap_bytes; - } - } - detected -} +mod startup; + +pub use startup::detect_vram_bytes_capped; +use startup::{ + bind_mesh_endpoint, init_owner_runtime, startup_secret_key, wait_for_endpoint_online, +}; +pub(crate) use startup::{default_plugin_event_source, hardware_snapshot_for_start}; /// Lightweight routing table for passive nodes (clients + standby GPU). /// Contains just enough info to route requests to the right host. @@ -33,218 +29,6 @@ pub struct RouteEntry { pub vram_gb: f64, } -pub(crate) async fn startup_secret_key(role: &NodeRole) -> Result { - if matches!(role, NodeRole::Client) || std::env::var("MESH_LLM_EPHEMERAL_KEY").is_ok() { - let key = SecretKey::generate(); - tracing::info!("Using ephemeral key (unique identity)"); - Ok(key) - } else { - load_or_create_key().await - } -} - -pub(crate) fn startup_transport_config() -> iroh::endpoint::QuicTransportConfig { - // We only raise the concurrent bidi-stream ceiling; everything else uses - // iroh's tuned defaults. - // - // History: this function used to override keep-alive (10s) and idle - // timeouts (300s connection + 300s per path). iroh 1.0 clamps per-path idle - // to 15s and already sends keep-alive PINGs every 5s, so those overrides do - // not provide the intended behavior and needlessly diverge from iroh's path - // management defaults. - // Mesh multiplexes many concurrent streams (gossip + heartbeat + inference - // tunnels) over one connection per peer, so we keep a generous bidi ceiling. - iroh::endpoint::QuicTransportConfig::builder() - .max_concurrent_bidi_streams(1024u32.into()) - .build() -} - -pub(crate) fn relay_mode_for_startup(relay: RelayConfig<'_>) -> Result { - let urls = effective_relay_urls(relay.policy, relay.urls); - if relay.policy.uses_relay() { - tracing::info!("Relay: {:?}", urls); - Ok(iroh::endpoint::RelayMode::Custom(relay_map_from_urls( - &urls, - relay.auths, - )?)) - } else { - let reason = match relay.policy { - RelayPolicy::ExplicitlyDisabled => "disabled by embedded config", - RelayPolicy::Disabled => "disabled by LAN-only discovery mode", - RelayPolicy::DefaultPublic => unreachable!("default public uses relays"), - }; - tracing::info!("Relay: {reason}"); - Ok(iroh::endpoint::RelayMode::Disabled) - } -} - -pub(crate) async fn bind_mesh_endpoint( - secret_key: SecretKey, - relay: RelayConfig<'_>, - quic_bind: QuicBindSelection, -) -> Result { - let mut builder = Endpoint::builder(iroh::endpoint::presets::Minimal) - .secret_key(secret_key) - .alpns(vec![ - ALPN_V1.to_vec(), - skippy_protocol::STAGE_ALPN_V2.to_vec(), - ]) - .transport_config(startup_transport_config()) - .relay_mode(relay_mode_for_startup(relay)?); - - if let Some(addr) = quic_bind_addr(quic_bind) { - tracing::info!("Binding QUIC to {addr}"); - if !relay.policy.uses_relay() && addr.is_ipv4() { - // LAN-only (relay-disabled) mode with a specific IPv4 bind: clear the - // pre-configured default sockets first. `bind_addr` only replaces the - // default for the *same* address family, so binding a specific IPv4 - // would otherwise leave the default IPv6 `[::]` socket in place. That - // extra local IPv6 path becomes a second candidate, and with no relay - // iroh's multipath negotiation across the IPv4+IPv6 locals fails with - // `MultipathNotNegotiated`, stalling the connection with no fallback. - // Pinning a single IPv4 socket keeps one local path family so the LAN - // direct path establishes cleanly. In relay (public) mode we keep the - // defaults so relay/IPv6 reachability is unaffected. - builder = builder.clear_ip_transports(); - } - builder = builder.bind_addr(addr)?; - } - - builder.bind().await.map_err(Into::into) -} - -pub(crate) async fn wait_for_endpoint_online( - endpoint: &Endpoint, - connected_log: &str, - timeout_log: &str, -) { - match tokio::time::timeout(std::time::Duration::from_secs(5), endpoint.online()).await { - Ok(()) => tracing::info!("{connected_log}"), - Err(_) => tracing::warn!("{timeout_log}"), - } -} - -pub(crate) fn hardware_snapshot_for_start( - hw: crate::system::hardware::HardwareSurvey, - role: &NodeRole, - max_vram_gb: Option, -) -> NodeHardwareSnapshot { - let mut vram_bytes = hw.vram_bytes; - let gpu_name = if matches!(role, NodeRole::Client) { - None - } else { - hw.gpu_name - }; - let hostname = hw.hostname; - let is_soc = Some(hw.is_soc); - let gpu_vram = (!hw.gpu_vram.is_empty()).then(|| { - hw.gpu_vram - .iter() - .map(|b| b.to_string()) - .collect::>() - .join(",") - }); - let gpu_reserved_bytes = if hw.gpu_reserved.iter().all(Option::is_none) { - None - } else { - Some( - hw.gpu_reserved - .iter() - .map(|value| value.map(|v| v.to_string()).unwrap_or_default()) - .collect::>() - .join(","), - ) - }; - - log_detected_vram(&mut vram_bytes, max_vram_gb); - - NodeHardwareSnapshot { - vram_bytes, - gpu_name, - hostname, - is_soc, - gpu_vram, - gpu_reserved_bytes, - } -} - -pub(crate) fn detected_vram_log(vram_bytes: u64, max_vram_gb: Option) -> DetectedVramLog { - let detected_gb = vram_bytes as f64 / 1e9; - let capped_bytes = max_vram_gb - .map(|max_gb| ((max_gb * 1e9) as u64, max_gb)) - .and_then(|(max_bytes, _)| (max_bytes < vram_bytes).then_some(max_bytes)); - DetectedVramLog { - detected_gb, - max_gb: max_vram_gb, - capped_bytes, - } -} - -pub(crate) fn log_detected_vram(vram_bytes: &mut u64, max_vram_gb: Option) { - let log = detected_vram_log(*vram_bytes, max_vram_gb); - if let Some(max_gb) = log.max_gb { - log_detected_vram_with_cap(vram_bytes, log.detected_gb, max_gb, log.capped_bytes); - } else { - tracing::info!("Detected VRAM: {:.1} GB", log.detected_gb); - } -} - -pub(crate) fn log_detected_vram_with_cap( - vram_bytes: &mut u64, - detected_gb: f64, - max_gb: f64, - capped_bytes: Option, -) { - if let Some(capped_bytes) = capped_bytes { - tracing::info!( - "Detected VRAM: {:.1} GB, capped to {:.1} GB (--max-vram)", - detected_gb, - max_gb - ); - *vram_bytes = capped_bytes; - } else { - tracing::info!( - "Detected VRAM: {:.1} GB (--max-vram {:.1} has no effect)", - detected_gb, - max_gb - ); - } -} - -pub(crate) fn init_owner_runtime( - owner_config: Option<&OwnerRuntimeConfig>, - endpoint_id: EndpointId, - hostname: Option, -) -> Result { - let trust_store = owner_config - .map(|config| config.trust_store.clone()) - .unwrap_or_default(); - let trust_policy = owner_config - .map(|config| config.trust_policy) - .unwrap_or_default(); - let owner_attestation = match owner_config.and_then(|config| config.keypair.as_ref()) { - Some(keypair) => Some(load_or_refresh_owner_attestation( - keypair, - endpoint_id, - owner_config.and_then(|config| config.node_label.clone()), - hostname, - )?), - None => None, - }; - - Ok(OwnerRuntimeInit { - trust_store, - trust_policy, - owner_attestation, - }) -} - -pub(crate) fn default_plugin_event_source(endpoint_id: EndpointId, source_peer_id: &mut String) { - if source_peer_id.is_empty() { - *source_peer_id = endpoint_id_hex(endpoint_id); - } -} - #[derive(Clone)] pub struct Node { pub(crate) endpoint: Endpoint, @@ -280,7 +64,10 @@ pub struct Node { pub(crate) join_targets: Arc>>, pub(crate) first_joined_mesh_ts: Arc>>, pub(crate) accepting: Arc<(tokio::sync::Notify, std::sync::atomic::AtomicBool)>, + /// Accelerator-resident capacity advertised to peers and split placement. pub(crate) vram_bytes: u64, + /// Local fit budget, which may additionally include CPU offload memory. + pub(crate) local_runtime_capacity_bytes: u64, pub(crate) peer_change_tx: watch::Sender, pub peer_change_rx: watch::Receiver, pub(crate) inflight_requests: Arc, @@ -1011,6 +798,7 @@ impl Node { std::sync::atomic::AtomicBool::new(false), )), vram_bytes: hardware.vram_bytes, + local_runtime_capacity_bytes: hardware.local_runtime_capacity_bytes, peer_change_tx, peer_change_rx, inflight_requests: Arc::new(std::sync::atomic::AtomicUsize::new(0)), @@ -1185,6 +973,7 @@ impl Node { std::sync::atomic::AtomicBool::new(false), )), vram_bytes: 0, + local_runtime_capacity_bytes: 0, peer_change_tx, peer_change_rx, inflight_requests: Arc::new(std::sync::atomic::AtomicUsize::new(0)), @@ -1721,12 +1510,12 @@ impl Node { }; if let Some(peer) = updated_peer { tracing::info!("Peer {} RTT: {}ms", id.fmt_short(), rtt_ms); - // If RTT dropped from above the split threshold (80ms) to below it + // If RTT dropped from above the configured split threshold to below it // (e.g. relay → direct), trigger a re-election so the peer can now // be included in split mode. let became_split_eligible = old_rtt - .map(|old| old > MAX_SPLIT_RTT_MS && rtt_ms <= MAX_SPLIT_RTT_MS) - .unwrap_or(rtt_ms <= MAX_SPLIT_RTT_MS); + .map(|old| old > max_split_rtt_ms() && rtt_ms <= max_split_rtt_ms()) + .unwrap_or(rtt_ms <= max_split_rtt_ms()); if became_split_eligible { emit_mesh_info(format!( "📡 Peer {} RTT improved ({}ms → {}ms) — re-electing for split", diff --git a/crates/mesh-llm-host-runtime/src/mesh/node/startup.rs b/crates/mesh-llm-host-runtime/src/mesh/node/startup.rs new file mode 100644 index 000000000..83ebf4d2d --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/mesh/node/startup.rs @@ -0,0 +1,228 @@ +use super::*; +use crate::mesh::identity_persistence::load_or_create_key; + +pub fn detect_vram_bytes_capped(max_vram_gb: Option) -> u64 { + let mut detected = crate::system::hardware::survey().vram_bytes; + if let Some(cap) = max_vram_gb { + let cap_bytes = (cap * 1e9) as u64; + if cap_bytes < detected { + detected = cap_bytes; + } + } + detected +} + +pub(super) async fn startup_secret_key(role: &NodeRole) -> Result { + if matches!(role, NodeRole::Client) || std::env::var("MESH_LLM_EPHEMERAL_KEY").is_ok() { + let key = SecretKey::generate(); + tracing::info!("Using ephemeral key (unique identity)"); + Ok(key) + } else { + load_or_create_key().await + } +} + +fn startup_transport_config() -> iroh::endpoint::QuicTransportConfig { + // We only raise the concurrent bidi-stream ceiling; everything else uses + // iroh's tuned defaults. + // + // History: this function used to override keep-alive (10s) and idle + // timeouts (300s connection + 300s per path). iroh 1.0 clamps per-path idle + // to 15s and already sends keep-alive PINGs every 5s, so those overrides do + // not provide the intended behavior and needlessly diverge from iroh's path + // management defaults. + // Mesh multiplexes many concurrent streams (gossip + heartbeat + inference + // tunnels) over one connection per peer, so we keep a generous bidi ceiling. + iroh::endpoint::QuicTransportConfig::builder() + .max_concurrent_bidi_streams(1024u32.into()) + .build() +} + +fn relay_mode_for_startup(relay: RelayConfig<'_>) -> Result { + let urls = effective_relay_urls(relay.policy, relay.urls); + if relay.policy.uses_relay() { + tracing::info!("Relay: {:?}", urls); + Ok(iroh::endpoint::RelayMode::Custom(relay_map_from_urls( + &urls, + relay.auths, + )?)) + } else { + let reason = match relay.policy { + RelayPolicy::ExplicitlyDisabled => "disabled by embedded config", + RelayPolicy::Disabled => "disabled by LAN-only discovery mode", + RelayPolicy::DefaultPublic => unreachable!("default public uses relays"), + }; + tracing::info!("Relay: {reason}"); + Ok(iroh::endpoint::RelayMode::Disabled) + } +} + +pub(super) async fn bind_mesh_endpoint( + secret_key: SecretKey, + relay: RelayConfig<'_>, + quic_bind: QuicBindSelection, +) -> Result { + let mut builder = Endpoint::builder(iroh::endpoint::presets::Minimal) + .secret_key(secret_key) + .alpns(vec![ + ALPN_V1.to_vec(), + skippy_protocol::STAGE_ALPN_V2.to_vec(), + ]) + .transport_config(startup_transport_config()) + .relay_mode(relay_mode_for_startup(relay)?); + + if let Some(addr) = quic_bind_addr(quic_bind) { + tracing::info!("Binding QUIC to {addr}"); + if !relay.policy.uses_relay() && addr.is_ipv4() { + // LAN-only (relay-disabled) mode with a specific IPv4 bind: clear the + // pre-configured default sockets first. `bind_addr` only replaces the + // default for the *same* address family, so binding a specific IPv4 + // would otherwise leave the default IPv6 `[::]` socket in place. That + // extra local IPv6 path becomes a second candidate, and with no relay + // iroh's multipath negotiation across the IPv4+IPv6 locals fails with + // `MultipathNotNegotiated`, stalling the connection with no fallback. + // Pinning a single IPv4 socket keeps one local path family so the LAN + // direct path establishes cleanly. In relay (public) mode we keep the + // defaults so relay/IPv6 reachability is unaffected. + builder = builder.clear_ip_transports(); + } + builder = builder.bind_addr(addr)?; + } + + builder.bind().await.map_err(Into::into) +} + +pub(super) async fn wait_for_endpoint_online( + endpoint: &Endpoint, + connected_log: &str, + timeout_log: &str, +) { + match tokio::time::timeout(std::time::Duration::from_secs(5), endpoint.online()).await { + Ok(()) => tracing::info!("{connected_log}"), + Err(_) => tracing::warn!("{timeout_log}"), + } +} + +pub(crate) fn hardware_snapshot_for_start( + hw: crate::system::hardware::HardwareSurvey, + role: &NodeRole, + max_vram_gb: Option, +) -> NodeHardwareSnapshot { + let local_runtime_capacity_bytes = + super::super::capacity::capped_capacity_bytes(hw.vram_bytes, max_vram_gb); + let mut vram_bytes = super::super::capacity::advertised_capacity_bytes(&hw, max_vram_gb); + let gpu_name = if matches!(role, NodeRole::Client) { + None + } else { + hw.gpu_name + }; + let hostname = hw.hostname; + let is_soc = Some(hw.is_soc); + let gpu_vram = (!hw.gpu_vram.is_empty()).then(|| { + hw.gpu_vram + .iter() + .map(|b| b.to_string()) + .collect::>() + .join(",") + }); + let gpu_reserved_bytes = if hw.gpu_reserved.iter().all(Option::is_none) { + None + } else { + Some( + hw.gpu_reserved + .iter() + .map(|value| value.map(|v| v.to_string()).unwrap_or_default()) + .collect::>() + .join(","), + ) + }; + + log_detected_vram(&mut vram_bytes, max_vram_gb); + + NodeHardwareSnapshot { + vram_bytes, + local_runtime_capacity_bytes, + gpu_name, + hostname, + is_soc, + gpu_vram, + gpu_reserved_bytes, + } +} + +fn detected_vram_log(vram_bytes: u64, max_vram_gb: Option) -> DetectedVramLog { + let detected_gb = vram_bytes as f64 / 1e9; + let capped_bytes = max_vram_gb + .map(|max_gb| ((max_gb * 1e9) as u64, max_gb)) + .and_then(|(max_bytes, _)| (max_bytes < vram_bytes).then_some(max_bytes)); + DetectedVramLog { + detected_gb, + max_gb: max_vram_gb, + capped_bytes, + } +} + +fn log_detected_vram(vram_bytes: &mut u64, max_vram_gb: Option) { + let log = detected_vram_log(*vram_bytes, max_vram_gb); + if let Some(max_gb) = log.max_gb { + log_detected_vram_with_cap(vram_bytes, log.detected_gb, max_gb, log.capped_bytes); + } else { + tracing::info!("Detected VRAM: {:.1} GB", log.detected_gb); + } +} + +fn log_detected_vram_with_cap( + vram_bytes: &mut u64, + detected_gb: f64, + max_gb: f64, + capped_bytes: Option, +) { + if let Some(capped_bytes) = capped_bytes { + tracing::info!( + "Detected VRAM: {:.1} GB, capped to {:.1} GB (--max-vram)", + detected_gb, + max_gb + ); + *vram_bytes = capped_bytes; + } else { + tracing::info!( + "Detected VRAM: {:.1} GB (--max-vram {:.1} has no effect)", + detected_gb, + max_gb + ); + } +} + +pub(super) fn init_owner_runtime( + owner_config: Option<&OwnerRuntimeConfig>, + endpoint_id: EndpointId, + hostname: Option, +) -> Result { + let trust_store = owner_config + .map(|config| config.trust_store.clone()) + .unwrap_or_default(); + let trust_policy = owner_config + .map(|config| config.trust_policy) + .unwrap_or_default(); + let owner_attestation = match owner_config.and_then(|config| config.keypair.as_ref()) { + Some(keypair) => Some(load_or_refresh_owner_attestation( + keypair, + endpoint_id, + owner_config.and_then(|config| config.node_label.clone()), + hostname, + )?), + None => None, + }; + + Ok(OwnerRuntimeInit { + trust_store, + trust_policy, + owner_attestation, + }) +} + +pub(crate) fn default_plugin_event_source(endpoint_id: EndpointId, source_peer_id: &mut String) { + if source_peer_id.is_empty() { + *source_peer_id = endpoint_id_hex(endpoint_id); + } +} diff --git a/crates/mesh-llm-host-runtime/src/mesh/peer_state.rs b/crates/mesh-llm-host-runtime/src/mesh/peer_state.rs index b07d9018a..8b0300957 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/peer_state.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/peer_state.rs @@ -851,10 +851,16 @@ impl Node { RoutingTable { hosts, mesh_id } } + /// Accelerator-resident capacity used for mesh stage placement. pub fn vram_bytes(&self) -> u64 { self.vram_bytes } + /// Local model-fit budget, including supported CPU offload memory. + pub fn local_runtime_capacity_bytes(&self) -> u64 { + self.local_runtime_capacity_bytes + } + pub async fn peers(&self) -> Vec { self.state .lock() diff --git a/crates/mesh-llm-host-runtime/src/mesh/stage_artifacts.rs b/crates/mesh-llm-host-runtime/src/mesh/stage_artifacts.rs index c9a1bcea6..b740d8c1b 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/stage_artifacts.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/stage_artifacts.rs @@ -10,8 +10,7 @@ use crate::mesh::stage_proto::{ }; use crate::mesh::stage_transport::{ ARTIFACT_TRANSFER_BUFFER_BYTES, ARTIFACT_TRANSFER_INVALID_OFFSET_ERROR, - ARTIFACT_TRANSFER_OPEN_TIMEOUT, ARTIFACT_TRANSFER_READ_IDLE_TIMEOUT, - LOCAL_STAGE_CONTROL_RESPONSE_TIMEOUT, StageTopologyInstance, + ARTIFACT_TRANSFER_OPEN_TIMEOUT, ARTIFACT_TRANSFER_READ_IDLE_TIMEOUT, StageTopologyInstance, artifact_transfer_allowed_by_topology, wait_local_stage_control_response, write_artifact_transfer_response, }; @@ -143,6 +142,10 @@ impl Node { &self, request: crate::inference::skippy::StageControlRequest, ) -> anyhow::Result { + // Load/Prepare can take minutes on large stages; use the same + // per-request budget the remote sender uses instead of the short + // local default, otherwise the executing node rejects its own load. + let timeout = Self::stage_control_request_timeout(&request); let control_tx = self.stage_control_tx.lock().await.clone(); match control_tx { Some(tx) => { @@ -152,8 +155,7 @@ impl Node { resp: resp_tx, }) .map_err(|_| anyhow::anyhow!("stage control loop is unavailable"))?; - wait_local_stage_control_response(resp_rx, LOCAL_STAGE_CONTROL_RESPONSE_TIMEOUT) - .await + wait_local_stage_control_response(resp_rx, timeout).await } None => Ok(stage_control_unavailable_response(request)), } diff --git a/crates/mesh-llm-host-runtime/src/mesh/stage_transport.rs b/crates/mesh-llm-host-runtime/src/mesh/stage_transport.rs index 5e016c665..eceb3bee5 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/stage_transport.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/stage_transport.rs @@ -1,4 +1,4 @@ -use super::{MAX_SPLIT_RTT_MS, Node, PeerInfo, StageTransportBridgeLabel}; +use super::{Node, PeerInfo, StageTransportBridgeLabel}; use crate::mesh::node::LocalRequestMetricsSampler; use crate::mesh::stage_proto::{ stage_control_request_to_proto, stage_control_response_from_proto, @@ -129,14 +129,20 @@ impl SplitStagePathSnapshot { } } - pub(crate) const fn stage_path_rejection(self) -> Option { + pub(crate) fn stage_path_rejection(self) -> Option { match self.kind { SplitStagePathKind::Direct => match self.rtt_ms { - Some(rtt_ms) if rtt_ms <= MAX_SPLIT_RTT_MS => None, + Some(rtt_ms) if rtt_ms <= super::max_split_rtt_ms() => None, Some(_) => Some(SplitStagePathRejection::StagePathTooSlow), None => Some(SplitStagePathRejection::MissingStagePath), }, - SplitStagePathKind::Relay => Some(SplitStagePathRejection::StagePathRelayOnly), + SplitStagePathKind::Relay => { + if super::split_allow_relay_paths() { + None + } else { + Some(SplitStagePathRejection::StagePathRelayOnly) + } + } SplitStagePathKind::Unknown => Some(SplitStagePathRejection::MissingStagePath), } } @@ -915,6 +921,9 @@ impl Node { self.record_stage_topology(stage_topology_from_load(self.endpoint.id(), load)) .await; } + // Load/Prepare can take minutes on large stages; use the same + // per-request budget remote control uses instead of the short default. + let timeout = Self::stage_control_request_timeout(&request); let control_tx = self.stage_control_tx.lock().await.clone(); let Some(tx) = control_tx else { anyhow::bail!("stage control is not available"); @@ -925,9 +934,7 @@ impl Node { resp: resp_tx, }) .map_err(|_| anyhow::anyhow!("stage control loop is unavailable"))?; - let response = - wait_local_stage_control_response(resp_rx, LOCAL_STAGE_CONTROL_RESPONSE_TIMEOUT) - .await?; + let response = wait_local_stage_control_response(resp_rx, timeout).await?; match &response { crate::inference::skippy::StageControlResponse::Ready(ready) => { self.record_stage_status(Some(self.endpoint.id()), ready.status.clone()) diff --git a/crates/mesh-llm-host-runtime/src/mesh/tests/connections.rs b/crates/mesh-llm-host-runtime/src/mesh/tests/connections.rs index 1618b7622..b202dc87d 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/tests/connections.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/tests/connections.rs @@ -441,6 +441,7 @@ async fn make_test_node_with_requirements( std::sync::atomic::AtomicBool::new(false), )), vram_bytes: 64 * 1024 * 1024 * 1024, + local_runtime_capacity_bytes: 64 * 1024 * 1024 * 1024, peer_change_tx, peer_change_rx, inflight_requests: Arc::new(std::sync::atomic::AtomicUsize::new(0)), diff --git a/crates/mesh-llm-host-runtime/src/mesh/tests/direct_path.rs b/crates/mesh-llm-host-runtime/src/mesh/tests/direct_path.rs index 4c01b2912..0d01f54af 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/tests/direct_path.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/tests/direct_path.rs @@ -95,6 +95,56 @@ fn direct_path_reverse_dial_keeps_existing_connection_when_gossip_fails() -> any }) } +#[test] +fn direct_path_reverse_dial_keeps_replaced_connection_open_for_inflight_streams() +-> anyhow::Result<()> { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build()? + .block_on(async { + let node = make_test_node(super::super::NodeRole::Worker).await?; + let remote = make_test_node(super::super::NodeRole::Worker).await?; + remote.start_accepting(); + + let existing = + connect_mesh(&node.endpoint, remote.endpoint_addr_for_advertisement()).await?; + let existing_id = existing.stable_id(); + { + let mut state = node.state.lock().await; + state.connections.insert(remote.id(), existing.clone()); + state + .peers + .insert(remote.id(), super::make_test_peer_info(remote.id())); + } + + let replacement = + connect_mesh(&node.endpoint, remote.endpoint_addr_for_advertisement()).await?; + let replacement_id = replacement.stable_id(); + node.install_direct_path_request_connection(remote.id(), replacement) + .await; + + let tracked_id = node + .state + .lock() + .await + .connections + .get(&remote.id()) + .expect("replacement connection should be tracked") + .stable_id(); + assert_eq!(tracked_id, replacement_id); + assert_ne!(tracked_id, existing_id); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), existing.closed()) + .await + .is_err(), + "replaced connection must remain open so existing streams can drain" + ); + + Ok(()) + }) +} + #[test] fn direct_path_reverse_dial_does_not_publish_during_pending_handshake() -> anyhow::Result<()> { tokio::runtime::Builder::new_multi_thread() diff --git a/crates/mesh-llm-host-runtime/src/models/capabilities.rs b/crates/mesh-llm-host-runtime/src/models/capabilities.rs index e76344e19..2ecf49e3e 100644 --- a/crates/mesh-llm-host-runtime/src/models/capabilities.rs +++ b/crates/mesh-llm-host-runtime/src/models/capabilities.rs @@ -6,11 +6,45 @@ pub use mesh_llm_types::models::capabilities::{ use super::build_hf_tokio_api; use super::remote_catalog; use serde_json::Value; -use std::path::Path; +use std::path::{Path, PathBuf}; #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] pub struct RuntimeMediaCapabilityEvidence { pub vision_projector_loaded: bool, + pub audio_projector_loaded: bool, +} + +pub async fn runtime_media_capability_evidence( + projector_path: Option, +) -> RuntimeMediaCapabilityEvidence { + match tokio::task::spawn_blocking(move || { + scan_runtime_media_capability_evidence(projector_path) + }) + .await + { + Ok(evidence) => evidence, + Err(error) => { + tracing::warn!(%error, "projector metadata scan task failed"); + RuntimeMediaCapabilityEvidence::default() + } + } +} + +fn scan_runtime_media_capability_evidence( + projector_path: Option, +) -> RuntimeMediaCapabilityEvidence { + let Some(projector_path) = projector_path else { + return RuntimeMediaCapabilityEvidence::default(); + }; + let projector_meta = model_artifact::gguf::scan_gguf_projector_meta(Path::new(&projector_path)); + RuntimeMediaCapabilityEvidence { + vision_projector_loaded: projector_meta + .and_then(|meta| meta.has_vision_encoder) + .unwrap_or(true), + audio_projector_loaded: projector_meta + .and_then(|meta| meta.has_audio_encoder) + .unwrap_or(false), + } } pub fn infer_remote_catalog_capabilities( @@ -73,6 +107,10 @@ pub fn runtime_verified_capabilities_from_static( caps.multimodal = false; } } + if evidence.audio_projector_loaded { + caps.audio = CapabilityLevel::Supported; + caps.multimodal = true; + } caps.normalize() } @@ -184,6 +222,7 @@ mod tests { Path::new("/models/Qwen3VL-2B-Instruct-Q4_K_M.gguf"), RuntimeMediaCapabilityEvidence { vision_projector_loaded: false, + audio_projector_loaded: false, }, ); @@ -201,6 +240,7 @@ mod tests { Path::new("/models/Qwen3VL-2B-Instruct-Q4_K_M.gguf"), RuntimeMediaCapabilityEvidence { vision_projector_loaded: true, + audio_projector_loaded: false, }, ); @@ -211,6 +251,23 @@ mod tests { assert!(caps.supports_multimodal_runtime()); } + #[test] + fn runtime_media_verification_promotes_loaded_audio_encoder() { + let caps = runtime_verified_model_capabilities( + "inkling-UD-Q2_K_XL", + Path::new("/models/inkling-UD-Q2_K_XL.gguf"), + RuntimeMediaCapabilityEvidence { + vision_projector_loaded: true, + audio_projector_loaded: true, + }, + ); + + assert_eq!(caps.vision, CapabilityLevel::Supported); + assert_eq!(caps.audio, CapabilityLevel::Supported); + assert!(caps.supports_multimodal_runtime()); + assert!(caps.supports_audio_runtime()); + } + #[test] fn runtime_media_verification_preserves_audio_and_non_media_traits() { let caps = ModelCapabilities { @@ -226,6 +283,7 @@ mod tests { caps, RuntimeMediaCapabilityEvidence { vision_projector_loaded: false, + audio_projector_loaded: false, }, ); diff --git a/crates/mesh-llm-host-runtime/src/models/mod.rs b/crates/mesh-llm-host-runtime/src/models/mod.rs index d848302df..653148318 100644 --- a/crates/mesh-llm-host-runtime/src/models/mod.rs +++ b/crates/mesh-llm-host-runtime/src/models/mod.rs @@ -25,7 +25,7 @@ use hf_hub::{HFClient, HFClientBuilder, HFClientSync}; pub use capabilities::{ CapabilityLevel, ModelCapabilities, RuntimeMediaCapabilityEvidence, - runtime_verified_model_capabilities, + runtime_media_capability_evidence, runtime_verified_model_capabilities, }; pub use download_transfer::DownloadTransferStats; pub(crate) use external_inference::append_external_inference_models; diff --git a/crates/mesh-llm-host-runtime/src/protocol/convert.rs b/crates/mesh-llm-host-runtime/src/protocol/convert.rs index 72a384ad5..4232786e6 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/convert.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/convert.rs @@ -17,7 +17,7 @@ fn skippy_stage_subprotocols( let mut features = vec![skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL.to_string()]; if stage_protocol_generation_supported { features.push( - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V3.to_string(), + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4.to_string(), ); } if artifact_transfer_supported { @@ -50,7 +50,7 @@ fn supports_skippy_status_list(subprotocols: &[crate::proto::node::MeshSubprotoc fn supports_skippy_stage_generation(subprotocols: &[crate::proto::node::MeshSubprotocol]) -> bool { supports_skippy_stage_feature( subprotocols, - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V3, + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, ) && supports_skippy_stage_feature( subprotocols, skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL, diff --git a/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs b/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs index 06dda3a54..691796c9d 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs @@ -82,7 +82,7 @@ fn owner_fields_roundtrip_through_proto_announcement() { .any(|feature| feature == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST) ); assert!(skippy.features.iter().any(|feature| feature - == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V3)); + == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4)); assert_eq!( proto_pa .owner_attestation @@ -310,7 +310,7 @@ fn proto_announcement_without_stage_control_is_not_stage_compatible() { name: skippy_protocol::STAGE_SUBPROTOCOL_NAME.to_string(), major: skippy_protocol::STAGE_SUBPROTOCOL_MAJOR, features: vec![ - skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V3.to_string(), + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4.to_string(), ], }], ..Default::default() diff --git a/crates/mesh-llm-host-runtime/src/runtime/local.rs b/crates/mesh-llm-host-runtime/src/runtime/local.rs index 74ae6b451..2119b8587 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local.rs @@ -531,7 +531,7 @@ pub(super) async fn start_runtime_local_model( let my_vram = spec .capacity_budget_bytes .or_else(|| spec.pinned_gpu.map(|gpu| gpu.allocatable_vram_bytes())) - .unwrap_or_else(|| spec.node.vram_bytes()); + .unwrap_or_else(|| spec.node.local_runtime_capacity_bytes()); // For split/layer-package models, compute the local share of model weights // and the layer fraction so the context planner budgets correctly. @@ -632,9 +632,14 @@ async fn start_runtime_skippy_model( let capabilities = models::runtime_verified_model_capabilities( &model_name, spec.model_path, - models::RuntimeMediaCapabilityEvidence { - vision_projector_loaded: resolved.hardware.projector_path.is_some(), - }, + models::runtime_media_capability_evidence( + resolved + .hardware + .projector_path + .as_deref() + .map(PathBuf::from), + ) + .await, ); let embedded_openai = resolved.to_embedded_openai_args(0, false)?; let mut options = resolved @@ -718,9 +723,14 @@ async fn start_runtime_layer_package_model( let capabilities = models::runtime_verified_model_capabilities( &model_name, spec.model_path, - models::RuntimeMediaCapabilityEvidence { - vision_projector_loaded: resolved.hardware.projector_path.is_some(), - }, + models::runtime_media_capability_evidence( + resolved + .hardware + .projector_path + .as_deref() + .map(PathBuf::from), + ) + .await, ); let activation_width = skippy_stage_activation_width(package.activation_width, &model_name)?; let run_id = format!("mesh-skippy-{}", now_unix_nanos()); diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_package.rs b/crates/mesh-llm-host-runtime/src/runtime/local_package.rs index 52e16393b..8286d97a3 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_package.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_package.rs @@ -9,10 +9,7 @@ use crate::models; use anyhow::{Context, Result}; use sha2::{Digest, Sha256}; use std::path::Path; -use std::time::Duration; -const SPLIT_PARTICIPANT_POLL_INTERVAL: Duration = Duration::from_millis(500); -const SPLIT_PARTICIPANT_STABLE_FOR: Duration = Duration::from_secs(2); pub(super) const SPLIT_DEFAULT_MIN_PARTICIPANTS: usize = 2; /// Try to extract GGUF architecture metadata from a layer package's shared @@ -67,12 +64,14 @@ pub(super) async fn split_runtime_compact_meta( pub(super) fn split_runtime_kv_bytes_per_token( package: &skippy::SkippyPackageIdentity, compact_meta: &models::gguf::GgufCompactMeta, + model_ref: &str, cache_type_k_override: Option<&str>, cache_type_v_override: Option<&str>, ) -> Result { - let split_kv_policy = skippy::KvCachePolicy::for_model_size(package.source_model_bytes); - let kv_cache_quant = split_kv_cache_quant( - &split_kv_policy, + let kv_cache_quant = split_effective_kv_cache_quant( + package, + compact_meta, + model_ref, cache_type_k_override, cache_type_v_override, ); @@ -80,6 +79,36 @@ pub(super) fn split_runtime_kv_bytes_per_token( .kv_cache_bytes_per_token(compact_meta) .context("split topology planning requires KV cache byte metadata") } + +/// Resolve the K/V cache types that split stages will actually load with. +/// +/// Stage loading applies the family default (for example Inkling's Q4_0 K/V) +/// ahead of the size-tiered `KvCachePolicy`. Planning must resolve K/V the same +/// way, or it budgets for a cheaper cache than the stages allocate and +/// over-packs the topology into an out-of-memory load. +pub(super) fn split_effective_kv_cache_quant( + package: &skippy::SkippyPackageIdentity, + compact_meta: &models::gguf::GgufCompactMeta, + model_ref: &str, + cache_type_k_override: Option<&str>, + cache_type_v_override: Option<&str>, +) -> models::gguf::GgufKvCacheQuant { + let size_policy = skippy::KvCachePolicy::for_model_size(package.source_model_bytes); + let family_default = + skippy::family_policy_for_compact_meta(compact_meta, Some(model_ref)).default_kv_cache_type; + + // Explicit user overrides win, then the family default, then model size. + let effective_k = cache_type_k_override + .or(family_default) + .unwrap_or(size_policy.cache_type_k()); + let effective_v = cache_type_v_override + .or(family_default) + .unwrap_or(size_policy.cache_type_v()); + + models::gguf::GgufKvCacheQuant::from_llama_args(effective_k, effective_v).unwrap_or_else(|| { + split_kv_cache_quant(&size_policy, cache_type_k_override, cache_type_v_override) + }) +} pub(super) async fn resolve_split_runtime_package( model_path: &Path, model_ref: &str, @@ -309,101 +338,6 @@ pub(super) fn package_ref_has_independent_prepare_source(package_ref: &str) -> b skippy_runtime::package::is_hf_package_ref(package_ref) } -pub(super) async fn wait_for_split_participants( - node: &mesh::Node, - model_name: &str, - model_ref: &str, - package: &skippy::SkippyPackageIdentity, - local_vram_override: Option, - timeout: Duration, -) -> Result { - let deadline = tokio::time::Instant::now() + timeout; - let mut best: Vec = Vec::new(); - let mut best_excluded: Vec = Vec::new(); - let mut last_signature: SplitParticipantSignature = Vec::new(); - let mut stable_since = tokio::time::Instant::now(); - loop { - let snapshot = - collect_split_participants(node, model_name, model_ref, package, local_vram_override) - .await; - let signature = split_participant_signature(&snapshot.participants); - let now = tokio::time::Instant::now(); - split_participant_signature_changed( - model_ref, - &snapshot, - &signature, - &mut last_signature, - &mut stable_since, - now, - ); - record_best_split_participants(&snapshot, &mut best, &mut best_excluded); - - let stable_for = now.saturating_duration_since(stable_since); - if split_participants_ready(&snapshot, stable_for) { - tracing::info!( - model_ref, - stable_for_ms = stable_for.as_millis(), - participants = ?split_participant_labels(&snapshot.participants), - "split topology participant set accepted" - ); - return Ok(snapshot); - } - - if now >= deadline { - ensure_split_participant_timeout_has_quorum(model_ref, &best, &best_excluded)?; - tracing::warn!( - model_ref, - participants = ?split_participant_labels(&best), - excluded = ?split_participant_exclusion_labels(&best_excluded), - "split topology participant wait timed out; using best observed set" - ); - return Ok(best_split_participant_snapshot(best, best_excluded)); - } - - tokio::time::sleep(SPLIT_PARTICIPANT_POLL_INTERVAL).await; - } -} - -pub(super) fn split_participant_signature_changed( - model_ref: &str, - snapshot: &SplitParticipantSnapshot, - signature: &SplitParticipantSignature, - last_signature: &mut SplitParticipantSignature, - stable_since: &mut tokio::time::Instant, - now: tokio::time::Instant, -) { - if signature == last_signature { - return; - } - *stable_since = now; - *last_signature = signature.clone(); - tracing::info!( - model_ref, - included = ?split_participant_labels(&snapshot.participants), - excluded = ?split_participant_exclusion_labels(&snapshot.excluded), - "split topology participant set changed" - ); -} - -pub(super) fn record_best_split_participants( - snapshot: &SplitParticipantSnapshot, - best: &mut Vec, - best_excluded: &mut Vec, -) { - if snapshot.participants.len() >= best.len() { - *best = snapshot.participants.clone(); - *best_excluded = snapshot.excluded.clone(); - } -} - -pub(super) fn split_participants_ready( - snapshot: &SplitParticipantSnapshot, - stable_for: Duration, -) -> bool { - snapshot.participants.len() >= SPLIT_DEFAULT_MIN_PARTICIPANTS - && stable_for >= SPLIT_PARTICIPANT_STABLE_FOR -} - pub(super) fn ensure_split_participant_timeout_has_quorum( model_ref: &str, best: &[SplitParticipant], @@ -508,10 +442,43 @@ pub(super) fn blocker_reason_rank(reason: &str) -> usize { .unwrap_or(usize::MAX) } -pub(super) fn best_split_participant_snapshot( - participants: Vec, - excluded: Vec, +pub(super) async fn collect_split_participant_membership( + node: &mesh::Node, + model_name: &str, + model_ref: &str, ) -> SplitParticipantSnapshot { + let mut participants = vec![SplitParticipant::new( + node.id(), + node.vram_bytes(), + Some(node.first_joined_mesh_ts().await.unwrap_or(0)), + )]; + let mut excluded = Vec::new(); + for peer in node.peers().await { + if let Some(reason) = split_peer_preflight_exclusion_reason(&peer, model_name, model_ref) { + excluded.push(SplitParticipantExclusion { + node_id: peer.id, + reason, + }); + continue; + } + if let Some(reason) = + split_peer_stage_path_exclusion_reason(node.split_stage_path_snapshot(peer.id).await) + { + excluded.push(SplitParticipantExclusion { + node_id: peer.id, + reason, + }); + continue; + } + participants.push(SplitParticipant::new( + peer.id, + peer.vram_bytes, + peer.first_joined_mesh_ts, + )); + } + sort_split_participants(&mut participants); + excluded.sort_by_key(|exclusion| exclusion.node_id.to_string()); + excluded.dedup_by_key(|exclusion| exclusion.node_id); SplitParticipantSnapshot { participants, excluded, @@ -578,8 +545,7 @@ pub(super) async fn collect_split_participants( } } } - participants.sort_by_key(|participant| participant.node_id.to_string()); - participants.dedup_by_key(|participant| participant.node_id); + sort_split_participants(&mut participants); excluded.sort_by_key(|exclusion| exclusion.node_id.to_string()); excluded.dedup_by_key(|exclusion| exclusion.node_id); SplitParticipantSnapshot { @@ -588,6 +554,11 @@ pub(super) async fn collect_split_participants( } } +fn sort_split_participants(participants: &mut Vec) { + participants.sort_by_key(|participant| participant.node_id.to_string()); + participants.dedup_by_key(|participant| participant.node_id); +} + pub(super) fn split_peer_preflight_exclusion_reason( peer: &mesh::PeerInfo, model_name: &str, diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_split.rs b/crates/mesh-llm-host-runtime/src/runtime/local_split.rs index b45e7a086..cba5ad109 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split.rs @@ -15,13 +15,14 @@ use super::local::{ use super::local_package::{ SplitParticipant, SplitParticipantExclusion, SplitParticipantSnapshot, collect_split_participants, resolve_split_runtime_package, split_runtime_compact_meta, - split_runtime_kv_bytes_per_token, wait_for_split_participants, + split_runtime_kv_bytes_per_token, }; +use super::split_participant_settle::{wait_for_split_membership, wait_for_split_participants}; use super::split_planning::{ PlannedRuntimeSliceTopology, RuntimeSliceStagePlan, SplitTopologyResourceInputs, plan_locked_runtime_slice_topology_with_resources, plan_runtime_slice_topology_with_resources, - split_participant_exclusion_labels, split_participant_labels, split_participants_for_stages, - split_stage_plan_labels, + plan_runtime_slice_topology_with_resources_and_stage0, split_participant_exclusion_labels, + split_participant_labels, split_participants_for_stages, split_stage_plan_labels, }; use super::split_topology_lock::load_locked_split_assignments; use crate::inference::skippy; @@ -51,6 +52,11 @@ pub(super) enum SplitRuntimeStart { Standby { coordinator: iroh::EndpointId }, } +enum CanonicalCoordinatorGate { + Coordinator(T), + Standby { coordinator: iroh::EndpointId }, +} + #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(super) enum StartupRuntimePlan { Local, @@ -143,11 +149,25 @@ pub(super) async fn start_runtime_split_model( spec: LocalRuntimeModelStartSpec<'_>, model_ref: &str, ) -> Result { + let coordinator_start = elect_split_start_coordinator(&spec, model_ref).await?; + let (canonical_coordinator, settled_membership) = match coordinator_start { + CanonicalCoordinatorGate::Coordinator(membership) => (spec.node.id(), membership), + CanonicalCoordinatorGate::Standby { coordinator } => { + return Ok(SplitRuntimeStart::Standby { coordinator }); + } + }; + let run_id = format!("mesh-split-{}", now_unix_nanos()); let topology_id = format!("topology-{run_id}"); - let split_setup = - prepare_split_runtime_start(&spec, model_ref, &topology_id, Duration::from_secs(30)) - .await?; + let split_setup = prepare_split_runtime_start( + &spec, + model_ref, + &topology_id, + &settled_membership, + canonical_coordinator, + Duration::from_secs(30), + ) + .await?; let SplitRuntimeStartPreparation { package, participant_snapshot, @@ -166,6 +186,12 @@ pub(super) async fn start_runtime_split_model( let stage0 = stages .first() .context("split topology did not produce stage 0")?; + anyhow::ensure!( + stage0.node_id == canonical_coordinator, + "split topology stage 0 {} does not match canonical coordinator {}", + stage0.node_id.fmt_short(), + canonical_coordinator.fmt_short() + ); tracing::info!( model_ref, topology_id, @@ -179,21 +205,6 @@ pub(super) async fn start_runtime_split_model( excluded = ?split_participant_exclusion_labels(&participant_snapshot.excluded), "split topology planned; elected coordinator from stage 0" ); - if let Some(standby) = - split_runtime_standby_start(spec.node, model_ref, &topology_id, &run_id, stage0) - { - return Ok(standby); - } - tracing::info!( - model_ref, - topology_id, - run_id, - local_node = %spec.node.id().fmt_short(), - context_length = planned_topology.context_length, - parallel_lanes = planned_topology.slots, - "split topology election selected local node as coordinator" - ); - let ctx_size = planned_topology.context_length; let slots = planned_topology.slots; let projector_path = spec @@ -279,6 +290,8 @@ async fn prepare_split_runtime_start( spec: &LocalRuntimeModelStartSpec<'_>, model_ref: &str, topology_id: &str, + settled_membership: &[SplitParticipant], + canonical_coordinator: iroh::EndpointId, timeout: Duration, ) -> Result { let package = resolve_split_runtime_package(spec.model_path, model_ref).await?; @@ -288,6 +301,7 @@ async fn prepare_split_runtime_start( model_ref, &package, spec.pinned_gpu.map(|gpu| gpu.allocatable_vram_bytes()), + settled_membership, timeout, ) .await?; @@ -295,6 +309,7 @@ async fn prepare_split_runtime_start( let kv_bytes_per_token = split_runtime_kv_bytes_per_token( &package, &compact_meta, + model_ref, spec.cache_type_k_override, spec.cache_type_v_override, )?; @@ -313,6 +328,13 @@ async fn prepare_split_runtime_start( &participant_snapshot.participants, ) .await?; + anyhow::ensure!( + locked_stages + .first() + .is_some_and(|stage| stage.node_id == canonical_coordinator), + "split topology lock stage 0 must be canonical coordinator {}", + canonical_coordinator.fmt_short() + ); plan_locked_runtime_slice_topology_with_resources( topology_id, model_ref, @@ -323,13 +345,14 @@ async fn prepare_split_runtime_start( &locked_stages, )? } else { - plan_runtime_slice_topology_with_resources( + plan_runtime_slice_topology_with_resources_and_stage0( topology_id, model_ref, &package, &participant_snapshot.participants, &participant_snapshot.excluded, resources, + Some(canonical_coordinator), )? }; Ok(SplitRuntimeStartPreparation { @@ -342,27 +365,47 @@ async fn prepare_split_runtime_start( }) } -fn split_runtime_standby_start( - node: &mesh::Node, +async fn elect_split_start_coordinator( + spec: &LocalRuntimeModelStartSpec<'_>, model_ref: &str, - topology_id: &str, - run_id: &str, - stage0: &RuntimeSliceStagePlan, -) -> Option { - if stage0.node_id == node.id() { - return None; +) -> Result>> { + let membership = + wait_for_split_membership(spec.node, model_ref, model_ref, Duration::from_secs(30)).await?; + let gate = canonical_coordinator_gate(spec.node.id(), membership.participants)?; + if let CanonicalCoordinatorGate::Standby { coordinator } = gate { + tracing::info!( + model_ref, + local_node = %spec.node.id().fmt_short(), + elected_coordinator = %coordinator.fmt_short(), + "canonical split coordinator is remote; local node entering standby before package planning" + ); + return Ok(CanonicalCoordinatorGate::Standby { coordinator }); } - tracing::info!( - model_ref, - topology_id, - run_id, - local_node = %node.id().fmt_short(), - elected_coordinator = %stage0.node_id.fmt_short(), - "split topology election selected a remote coordinator; local node entering standby" - ); - Some(SplitRuntimeStart::Standby { - coordinator: stage0.node_id, - }) + Ok(gate) +} + +fn canonical_coordinator_gate( + local_node: iroh::EndpointId, + membership: Vec, +) -> Result>> { + let coordinator = canonical_split_coordinator(&membership) + .context("split membership did not produce a canonical coordinator")?; + if coordinator == local_node { + Ok(CanonicalCoordinatorGate::Coordinator(membership)) + } else { + Ok(CanonicalCoordinatorGate::Standby { coordinator }) + } +} + +fn canonical_split_coordinator(membership: &[SplitParticipant]) -> Option { + membership + .iter() + .max_by(|left, right| { + left.vram_bytes + .cmp(&right.vram_bytes) + .then_with(|| right.node_id.to_string().cmp(&left.node_id.to_string())) + }) + .map(|participant| participant.node_id) } use loading::{SplitGenerationLoadSpec, load_split_runtime_generation}; diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs b/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs index 2e0a9a95f..f3dd625c5 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs @@ -20,6 +20,10 @@ use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::time::Duration; +pub(super) const MIN_STAGE_SOURCE_PREPARE_TIMEOUT: Duration = Duration::from_secs(30 * 60); +const STAGE_SOURCE_PREPARE_ALLOWANCE: Duration = Duration::from_secs(10 * 60); +const STAGE_SOURCE_MIN_BYTES_PER_SEC: u64 = 16 * 1024 * 1024; + pub(super) struct SplitGenerationLoadSpec<'a> { pub(super) node: &'a mesh::Node, pub(super) mesh_config: &'a plugin::MeshConfig, @@ -176,7 +180,14 @@ pub(super) async fn load_split_runtime_generation_inner( stage_index: downstream.stage_index, endpoint: downstream_endpoint, }); - let vision_projector_loaded = runtime_options.config.projector_path.is_some(); + let media_capability_evidence = models::runtime_media_capability_evidence( + runtime_options + .config + .projector_path + .as_deref() + .map(std::path::PathBuf::from), + ) + .await; let node_for_hook = spec.node.clone(); let model_ref = spec.model_ref.to_string(); let reporter_model_ref = model_ref.clone(); @@ -209,9 +220,7 @@ pub(super) async fn load_split_runtime_generation_inner( let capabilities = models::runtime_verified_model_capabilities( spec.model_ref, spec.model_path, - models::RuntimeMediaCapabilityEvidence { - vision_projector_loaded, - }, + media_capability_evidence, ); spec.node @@ -270,7 +279,7 @@ pub(super) async fn load_downstream_split_runtime_stages( spec.node, stage.node_id, &load, - Duration::from_secs(30 * 60), + stage_source_prepare_timeout(spec.package, stage), ) .await .with_context(|| { @@ -322,6 +331,22 @@ pub(super) async fn load_downstream_split_runtime_stages( .context("split topology missing downstream stage") } +pub(super) fn stage_source_prepare_timeout( + package: &skippy::SkippyPackageIdentity, + stage: &RuntimeSliceStagePlan, +) -> Duration { + let package_layers = u64::from(package.layer_count.max(1)); + let stage_layers = u64::from(stage.layer_end.saturating_sub(stage.layer_start).max(1)); + let estimated_stage_bytes = package + .source_model_bytes + .saturating_mul(stage_layers) + .div_ceil(package_layers); + let transfer_secs = estimated_stage_bytes.div_ceil(STAGE_SOURCE_MIN_BYTES_PER_SEC); + Duration::from_secs(transfer_secs) + .saturating_add(STAGE_SOURCE_PREPARE_ALLOWANCE) + .max(MIN_STAGE_SOURCE_PREPARE_TIMEOUT) +} + pub(super) fn split_runtime_stage_load_request( spec: &SplitGenerationLoadSpec<'_>, settings: &SplitGenerationLoadSettings<'_>, diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs b/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs index e132980c1..9bf59e949 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs @@ -36,6 +36,65 @@ fn runtime_local_targets_keep_duplicate_same_model_ports() { ); } +#[test] +fn canonical_coordinator_is_identical_with_divergent_observer_signals() { + let capacities = [ + (1, 24_000_000_000), + (2, 48_000_000_000), + (3, 48_000_000_000), + ]; + let observer_a = capacities + .into_iter() + .map(|(seed, capacity)| { + SplitParticipant::new(make_id(seed), capacity, None).with_package_signals( + SplitParticipantPackageSignal { + cached_slice_bytes: u64::from(seed) * 10_000, + missing_artifact_bytes: u64::from(4 - seed) * 20_000, + availability_score: u32::from(seed), + }, + Some(u32::from(seed) * 40), + true, + ) + }) + .collect::>(); + let observer_b = capacities + .into_iter() + .rev() + .map(|(seed, capacity)| { + SplitParticipant::new(make_id(seed), capacity, None).with_package_signals( + SplitParticipantPackageSignal { + cached_slice_bytes: u64::from(4 - seed) * 100_000, + missing_artifact_bytes: u64::from(seed) * 50_000, + availability_score: u32::from(4 - seed), + }, + Some(u32::from(4 - seed)), + false, + ) + }) + .collect::>(); + + let expected = [make_id(2), make_id(3)].into_iter().min().unwrap(); + assert_eq!(canonical_split_coordinator(&observer_a), Some(expected)); + assert_eq!(canonical_split_coordinator(&observer_b), Some(expected)); +} + +#[test] +fn noncanonical_gate_returns_standby_without_invoking_package_planning() { + let local = SplitParticipant::new(make_id(1), 24_000_000_000, None); + let coordinator = SplitParticipant::new(make_id(2), 48_000_000_000, None); + + let gate = canonical_coordinator_gate(local.node_id, vec![local, coordinator]) + .expect("canonical coordinator gate"); + match gate { + CanonicalCoordinatorGate::Standby { + coordinator: selected, + } => assert_eq!(selected, coordinator.node_id), + CanonicalCoordinatorGate::Coordinator(_) => { + panic!("local node must not be elected coordinator") + } + } +} + #[test] fn split_topology_planner_uses_all_eligible_participants() { let participants = vec![ @@ -67,6 +126,57 @@ fn split_topology_planner_uses_all_eligible_participants() { assert_eq!(stages.last().unwrap().layer_end, 40); } +#[test] +fn resource_planner_keeps_canonical_coordinator_at_stage_zero() { + let canonical = SplitParticipant::new(make_id(1), 48_000_000_000, None).with_package_signals( + SplitParticipantPackageSignal { + cached_slice_bytes: 0, + missing_artifact_bytes: 40_000_000, + availability_score: 0, + }, + Some(200), + true, + ); + let fast_a = SplitParticipant::new(make_id(2), 32_000_000_000, None).with_package_signals( + SplitParticipantPackageSignal { + cached_slice_bytes: 40_000_000, + missing_artifact_bytes: 0, + availability_score: 40, + }, + Some(1), + true, + ); + let fast_b = SplitParticipant::new(make_id(3), 32_000_000_000, None).with_package_signals( + SplitParticipantPackageSignal { + cached_slice_bytes: 40_000_000, + missing_artifact_bytes: 0, + availability_score: 40, + }, + Some(1), + true, + ); + let participants = [canonical, fast_a, fast_b]; + let package = package(40); + + let planned = plan_runtime_slice_topology_with_resources_and_stage0( + "topology-test", + "test-model", + &package, + &participants, + &[], + SplitTopologyResourceInputs { + native_context_length: 65_536, + kv_bytes_per_token: 64 * 1024, + ctx_size_override: Some(65_536), + parallel_override: Some(1), + }, + Some(canonical.node_id), + ) + .expect("canonical stage-zero topology"); + + assert_eq!(planned.stages.first().unwrap().node_id, canonical.node_id); +} + #[test] fn split_topology_planner_prefers_cached_participant_in_runtime_path() { let cold = SplitParticipant::new(make_id(1), 24_000_000_000, None).with_package_signals( @@ -502,6 +612,38 @@ fn split_startup_error_messages_include_specific_blocker_tokens() { assert!(timeout.contains("30s")); } +#[test] +fn stage_source_prepare_timeout_scales_with_assigned_package_bytes() { + let package = skippy::SkippyPackageIdentity { + source_model_bytes: 975_000_000_000, + layer_count: 66, + ..package(66) + }; + let small_stage = RuntimeSliceStagePlan { + stage_id: "stage-2".to_string(), + stage_index: 2, + node_id: make_id(2), + layer_start: 59, + layer_end: 66, + parameter_bytes: 0, + }; + let large_stage = RuntimeSliceStagePlan { + stage_id: "stage-0".to_string(), + stage_index: 0, + node_id: make_id(1), + layer_start: 0, + layer_end: 39, + parameter_bytes: 0, + }; + + let small_timeout = stage_source_prepare_timeout(&package, &small_stage); + let large_timeout = stage_source_prepare_timeout(&package, &large_stage); + + assert!(small_timeout > MIN_STAGE_SOURCE_PREPARE_TIMEOUT); + assert!(large_timeout > small_timeout); + assert!(large_timeout > Duration::from_secs(6 * 60 * 60)); +} + #[test] fn startup_runtime_plan_auto_splits_when_model_exceeds_local_capacity() { assert_eq!( @@ -745,7 +887,7 @@ fn split_participant_signature_includes_vram_for_stability() { } #[test] -fn split_participant_signature_includes_package_signals_for_stability() { +fn split_participant_signature_includes_package_signals_for_claim_identity() { let node_id = make_id(9); let first = vec![SplitParticipant::new(node_id, 24_000_000_000, None)]; let second = vec![ @@ -1424,3 +1566,109 @@ fn split_topology_minimum_rejects_single_stage_split_candidate() { ])); assert!(!split_stages_meet_minimum(&[stage(1, 0, 0, 40)])); } + +#[test] +fn split_planning_uses_family_kv_defaults_for_inkling() { + let mut meta = crate::models::gguf::GgufCompactMeta { + architecture: "inkling".to_string(), + context_length: 65_536, + embedding_size: 4096, + head_count: 32, + kv_head_count: 8, + layer_count: 66, + key_length: 128, + value_length: 128, + ..Default::default() + }; + meta.kv_head_counts = vec![8; 66]; + + // Inkling's reviewed family default keeps both planning and stage loading + // on quantized Q4_0 K/V rather than silently expanding to F16. + let mut identity = package(66); + identity.source_model_bytes = 318 * 1024 * 1024 * 1024; + + let planned = + split_runtime_kv_bytes_per_token(&identity, &meta, "tml/inkling-q2", None, None).unwrap(); + let expected_q4 = crate::models::gguf::GgufKvCacheQuant::from_llama_args("q4_0", "q4_0") + .unwrap() + .kv_cache_bytes_per_token(&meta) + .unwrap(); + assert_eq!(planned, expected_q4); + + // Explicit user overrides still win over the family default. + let overridden = split_runtime_kv_bytes_per_token( + &identity, + &meta, + "tml/inkling-q2", + Some("f16"), + Some("f16"), + ) + .unwrap(); + assert!(overridden > planned); +} + +/// Validates the finding-#1 fix against a real Inkling layer package. +/// +/// Set `INKLING_METADATA_GGUF` to a package's `shared/metadata.gguf` to run it; +/// skipped otherwise so CI stays hermetic. +#[test] +fn real_inkling_metadata_plans_family_kv_not_size_tiered() { + let Ok(path) = std::env::var("INKLING_METADATA_GGUF") else { + eprintln!("skip: INKLING_METADATA_GGUF not set"); + return; + }; + let meta = crate::models::gguf::scan_gguf_compact_meta(std::path::Path::new(&path)) + .expect("scan real inkling metadata gguf"); + eprintln!( + "REAL META arch={} layers={} embed={} kv_heads={} k_len={} v_len={} ctx={}", + meta.architecture, + meta.layer_count, + meta.embedding_size, + meta.kv_head_count, + meta.key_length, + meta.value_length, + meta.context_length + ); + + let model_ref = "unsloth/inkling-GGUF:UD-Q2_K_XL"; + let policy = crate::inference::skippy::family_policy_for_compact_meta(&meta, Some(model_ref)); + eprintln!( + "FAMILY default_kv_cache_type={:?}", + policy.default_kv_cache_type + ); + + let mut identity = package(meta.layer_count); + identity.source_model_bytes = 318 * 1024 * 1024 * 1024; + + let planned = + split_runtime_kv_bytes_per_token(&identity, &meta, model_ref, None, None).unwrap(); + let expected_q4 = crate::models::gguf::GgufKvCacheQuant::from_llama_args("q4_0", "q4_0") + .unwrap() + .kv_cache_bytes_per_token(&meta) + .unwrap(); + let size_tiered = { + let p = + crate::inference::skippy::KvCachePolicy::for_model_size(identity.source_model_bytes); + split_kv_cache_quant(&p, None, None) + .kv_cache_bytes_per_token(&meta) + .unwrap() + }; + let ctx = u64::from(meta.context_length.max(1)); + eprintln!( + "KV/token planned={planned} size_tiered={size_tiered} ratio={:.2}x | @ctx{ctx}: planned={:.1}GiB size_tiered={:.1}GiB under_budget={:.1}GiB", + planned as f64 / size_tiered.max(1) as f64, + (planned * ctx) as f64 / (1024.0 * 1024.0 * 1024.0), + (size_tiered * ctx) as f64 / (1024.0 * 1024.0 * 1024.0), + ((planned - size_tiered.min(planned)) * ctx) as f64 / (1024.0 * 1024.0 * 1024.0), + ); + + assert_eq!( + policy.default_kv_cache_type, + Some("q4_0"), + "inkling must resolve a q4_0 family K/V default" + ); + assert_eq!( + planned, expected_q4, + "family-aware planning must use the Inkling Q4_0 K/V default" + ); +} diff --git a/crates/mesh-llm-host-runtime/src/runtime/mod.rs b/crates/mesh-llm-host-runtime/src/runtime/mod.rs index 4b1482cad..1ad700e6e 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/mod.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/mod.rs @@ -22,11 +22,13 @@ mod release_attestation; mod run_auto; mod runtime_registry; mod serving_surface; +mod split_participant_settle; mod split_planning; mod split_topology_lock; mod startup_handles; mod startup_identity; mod startup_models; +mod startup_retry; mod status; pub(crate) mod survey; #[cfg(test)] diff --git a/crates/mesh-llm-host-runtime/src/runtime/model_lifecycle/load.rs b/crates/mesh-llm-host-runtime/src/runtime/model_lifecycle/load.rs index 3233c6924..8ea8fb828 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/model_lifecycle/load.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/model_lifecycle/load.rs @@ -49,7 +49,7 @@ pub(crate) async fn run_auto_load_runtime_model( &instance_id, &runtime_model_name, None, - ctx.node.vram_bytes(), + ctx.node.local_runtime_capacity_bytes(), model_bytes, )?; add_serving_assignment(ctx.node, ctx.primary_model_name, &runtime_model_name).await; diff --git a/crates/mesh-llm-host-runtime/src/runtime/model_lifecycle/reconciliation.rs b/crates/mesh-llm-host-runtime/src/runtime/model_lifecycle/reconciliation.rs index 54e0cbcb4..50adcd24d 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/model_lifecycle/reconciliation.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/model_lifecycle/reconciliation.rs @@ -66,7 +66,7 @@ pub(crate) async fn reconcile_model_targets_once(ctx: ReconcileModelTargetsConte } let target_lookup = console_state.model_target_lookup().await; - let local_vram_bytes = node.vram_bytes(); + let local_vram_bytes = node.local_runtime_capacity_bytes(); let targets = target_lookup .targets .into_iter() diff --git a/crates/mesh-llm-host-runtime/src/runtime/split_participant_settle.rs b/crates/mesh-llm-host-runtime/src/runtime/split_participant_settle.rs new file mode 100644 index 000000000..925687053 --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/runtime/split_participant_settle.rs @@ -0,0 +1,394 @@ +use super::local_package::{ + SPLIT_DEFAULT_MIN_PARTICIPANTS, SplitParticipant, SplitParticipantSnapshot, + collect_split_participant_membership, collect_split_participants, + ensure_split_participant_timeout_has_quorum, +}; +use super::split_planning::{split_participant_exclusion_labels, split_participant_labels}; +use crate::inference::skippy; +use crate::mesh; +use anyhow::Result; +use std::time::Duration; + +const SPLIT_PARTICIPANT_POLL_INTERVAL: Duration = Duration::from_millis(500); +/// An automatic split cannot know how many nodes the operator intends to start. +/// Wait for additions/capacity changes to stop instead of claiming the first +/// two-node quorum that happens to become visible. +const SPLIT_MEMBERSHIP_SETTLE_DWELL: Duration = Duration::from_secs(8); +/// Even an immediately stable two-node quorum gets a bounded discovery window. +const SPLIT_FIRST_QUORUM_OBSERVATION: Duration = Duration::from_secs(8); + +fn membership_settle_timeout(requested: Duration) -> Duration { + requested.max(SPLIT_FIRST_QUORUM_OBSERVATION.max(SPLIT_MEMBERSHIP_SETTLE_DWELL)) +} + +type SplitMembershipSignature = Vec<(String, u64)>; + +#[derive(Debug, Default)] +struct SplitMembershipSettleBarrier { + signature: SplitMembershipSignature, + first_quorum_observed: Option, + stable_since: Option, +} + +impl SplitMembershipSettleBarrier { + fn observe(&mut self, participants: &[SplitParticipant], now: tokio::time::Instant) -> bool { + let signature = split_membership_signature(participants); + if signature != self.signature { + self.signature = signature; + self.stable_since = Some(now); + } + if participants.len() >= SPLIT_DEFAULT_MIN_PARTICIPANTS + && self.first_quorum_observed.is_none() + { + self.first_quorum_observed = Some(now); + } + self.is_ready(participants.len(), now) + } + + fn is_ready(&self, participant_count: usize, now: tokio::time::Instant) -> bool { + if participant_count < SPLIT_DEFAULT_MIN_PARTICIPANTS { + return false; + } + let Some(first_quorum_observed) = self.first_quorum_observed else { + return false; + }; + let Some(stable_since) = self.stable_since else { + return false; + }; + now.saturating_duration_since(first_quorum_observed) >= SPLIT_FIRST_QUORUM_OBSERVATION + && now.saturating_duration_since(stable_since) >= SPLIT_MEMBERSHIP_SETTLE_DWELL + } + + fn stable_for(&self, now: tokio::time::Instant) -> Duration { + self.stable_since + .map(|stable_since| now.saturating_duration_since(stable_since)) + .unwrap_or_default() + } +} + +struct SplitMembershipWait<'a> { + node: &'a mesh::Node, + model_name: &'a str, + model_ref: &'a str, + deadline: tokio::time::Instant, + barrier: SplitMembershipSettleBarrier, + last_logged_signature: SplitMembershipSignature, +} + +impl<'a> SplitMembershipWait<'a> { + fn new( + node: &'a mesh::Node, + model_name: &'a str, + model_ref: &'a str, + timeout: Duration, + ) -> Self { + Self { + node, + model_name, + model_ref, + deadline: tokio::time::Instant::now() + membership_settle_timeout(timeout), + barrier: SplitMembershipSettleBarrier::default(), + last_logged_signature: Vec::new(), + } + } + + async fn run(mut self) -> Result { + loop { + let snapshot = + collect_split_participant_membership(self.node, self.model_name, self.model_ref) + .await; + let signature = split_membership_signature(&snapshot.participants); + self.log_membership_change(&snapshot, &signature); + let now = tokio::time::Instant::now(); + if self.barrier.observe(&snapshot.participants, now) { + self.log_accepted(&snapshot, now); + return Ok(snapshot); + } + if now >= self.deadline { + return self.finish_at_timeout().await; + } + tokio::time::sleep(SPLIT_PARTICIPANT_POLL_INTERVAL).await; + } + } + + fn log_membership_change( + &mut self, + snapshot: &SplitParticipantSnapshot, + signature: &SplitMembershipSignature, + ) { + if signature == &self.last_logged_signature { + return; + } + tracing::info!( + model_ref = self.model_ref, + members = ?split_participant_labels(&snapshot.participants), + excluded = ?split_participant_exclusion_labels(&snapshot.excluded), + "split topology stable membership changed" + ); + self.last_logged_signature = signature.clone(); + } + + fn log_accepted(&self, snapshot: &SplitParticipantSnapshot, now: tokio::time::Instant) { + tracing::info!( + model_ref = self.model_ref, + stable_for_ms = self.barrier.stable_for(now).as_millis(), + participants = ?split_participant_labels(&snapshot.participants), + "split topology membership accepted for canonical coordinator election" + ); + } + + /// Elect only from a freshly revalidated snapshot. + /// + /// A best-ever set can name peers that have since vanished, which puts + /// dead nodes into the elected topology. Re-collect at the deadline so the + /// final membership reflects peers that are still present. + async fn finish_at_timeout(self) -> Result { + let snapshot = + collect_split_participant_membership(self.node, self.model_name, self.model_ref).await; + ensure_split_participant_timeout_has_quorum( + self.model_ref, + &snapshot.participants, + &snapshot.excluded, + )?; + tracing::warn!( + model_ref = self.model_ref, + participants = ?split_participant_labels(&snapshot.participants), + excluded = ?split_participant_exclusion_labels(&snapshot.excluded), + "split topology membership settle timed out; using revalidated final snapshot" + ); + Ok(snapshot) + } +} + +struct SplitEligibilityWait<'a> { + node: &'a mesh::Node, + model_name: &'a str, + model_ref: &'a str, + package: &'a skippy::SkippyPackageIdentity, + local_vram_override: Option, + expected_node_ids: Vec, + deadline: tokio::time::Instant, +} + +impl<'a> SplitEligibilityWait<'a> { + async fn run(self) -> Result { + loop { + let snapshot = self.collect_snapshot().await; + if self.snapshot_is_complete(&snapshot) { + self.log_complete(&snapshot); + return Ok(snapshot); + } + if tokio::time::Instant::now() >= self.deadline { + return self.finish_at_timeout().await; + } + self.log_pending(&snapshot); + tokio::time::sleep(SPLIT_PARTICIPANT_POLL_INTERVAL).await; + } + } + + async fn collect_snapshot(&self) -> SplitParticipantSnapshot { + collect_split_participants( + self.node, + self.model_name, + self.model_ref, + self.package, + self.local_vram_override, + ) + .await + } + + fn snapshot_is_complete(&self, snapshot: &SplitParticipantSnapshot) -> bool { + split_membership_node_ids(&snapshot.participants) == self.expected_node_ids + } + + fn log_complete(&self, snapshot: &SplitParticipantSnapshot) { + tracing::info!( + model_ref = self.model_ref, + participants = ?split_participant_labels(&snapshot.participants), + "canonical coordinator accepted full split package inventory" + ); + } + + /// Elect only from a freshly revalidated eligible snapshot; see + /// `SplitMembershipWait::finish_at_timeout`. + async fn finish_at_timeout(self) -> Result { + let snapshot = self.collect_snapshot().await; + ensure_split_participant_timeout_has_quorum( + self.model_ref, + &snapshot.participants, + &snapshot.excluded, + )?; + tracing::warn!( + model_ref = self.model_ref, + participants = ?split_participant_labels(&snapshot.participants), + excluded = ?split_participant_exclusion_labels(&snapshot.excluded), + "split package inventory wait timed out; using revalidated final snapshot" + ); + Ok(snapshot) + } + + fn log_pending(&self, snapshot: &SplitParticipantSnapshot) { + tracing::debug!( + model_ref = self.model_ref, + expected_members = ?self.expected_node_ids, + eligible = ?split_participant_labels(&snapshot.participants), + excluded = ?split_participant_exclusion_labels(&snapshot.excluded), + "canonical coordinator waiting for full split package inventory" + ); + } +} + +pub(super) async fn wait_for_split_membership( + node: &mesh::Node, + model_name: &str, + model_ref: &str, + timeout: Duration, +) -> Result { + SplitMembershipWait::new(node, model_name, model_ref, timeout) + .run() + .await +} + +pub(super) async fn wait_for_split_participants( + node: &mesh::Node, + model_name: &str, + model_ref: &str, + package: &skippy::SkippyPackageIdentity, + local_vram_override: Option, + expected_membership: &[SplitParticipant], + timeout: Duration, +) -> Result { + SplitEligibilityWait { + node, + model_name, + model_ref, + package, + local_vram_override, + expected_node_ids: split_membership_node_ids(expected_membership), + deadline: tokio::time::Instant::now() + timeout, + } + .run() + .await +} + +fn split_membership_signature(participants: &[SplitParticipant]) -> SplitMembershipSignature { + participants + .iter() + .map(|participant| (participant.node_id.to_string(), participant.vram_bytes)) + .collect() +} + +fn split_membership_node_ids(participants: &[SplitParticipant]) -> Vec { + participants + .iter() + .map(|participant| participant.node_id.to_string()) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::runtime::local_package::{ + SplitParticipantPackageSignal, split_participant_signature, + }; + + fn make_id(seed: u8) -> iroh::EndpointId { + let secret = iroh::SecretKey::from_bytes(&[seed; 32]); + secret.public() + } + + fn participant(seed: u8) -> SplitParticipant { + SplitParticipant::new(make_id(seed), u64::from(seed) * 1_000_000_000, None) + } + + #[test] + fn membership_settle_waits_for_two_to_six_arrivals() { + let start = tokio::time::Instant::now(); + let mut barrier = SplitMembershipSettleBarrier::default(); + + assert!(!barrier.observe(&[participant(1), participant(2)], start)); + for (seconds, count) in [(2, 3), (4, 4), (6, 5), (7, 6)] { + let participants = (1..=count).map(participant).collect::>(); + assert!(!barrier.observe(&participants, start + Duration::from_secs(seconds))); + } + let participants = (1..=6).map(participant).collect::>(); + assert!(!barrier.observe(&participants, start + Duration::from_secs(14))); + assert!(barrier.observe(&participants, start + Duration::from_secs(15))); + } + + #[test] + fn two_node_cohort_settles_after_bounded_discovery_window() { + let start = tokio::time::Instant::now(); + let participants = vec![participant(1), participant(2)]; + let mut barrier = SplitMembershipSettleBarrier::default(); + + assert!(!barrier.observe(&participants, start)); + assert!(!barrier.observe(&participants, start + Duration::from_secs(7))); + assert!(barrier.observe(&participants, start + Duration::from_secs(8))); + } + + #[test] + fn short_caller_timeout_still_allows_the_settle_barrier() { + assert_eq!( + membership_settle_timeout(Duration::from_secs(1)), + Duration::from_secs(8) + ); + assert_eq!( + membership_settle_timeout(Duration::from_secs(30)), + Duration::from_secs(30) + ); + } + + #[test] + fn volatile_package_and_rtt_signals_do_not_reset_membership_dwell() { + let start = tokio::time::Instant::now(); + let mut participants = vec![participant(1), participant(2)]; + let mut barrier = SplitMembershipSettleBarrier::default(); + assert!(!barrier.observe(&participants, start)); + + participants[0] = participants[0].with_package_signals( + SplitParticipantPackageSignal { + cached_slice_bytes: 10_000, + missing_artifact_bytes: 90_000, + availability_score: 3, + }, + Some(80), + true, + ); + participants[1] = participants[1].with_package_signals( + SplitParticipantPackageSignal { + cached_slice_bytes: 100_000, + missing_artifact_bytes: 0, + availability_score: 30, + }, + Some(4), + true, + ); + + assert!(barrier.observe(&participants, start + Duration::from_secs(8))); + assert_eq!( + split_membership_signature(&participants), + vec![ + (make_id(1).to_string(), 1_000_000_000), + (make_id(2).to_string(), 2_000_000_000), + ] + ); + assert_ne!(split_participant_signature(&participants), Vec::new()); + } + + #[test] + fn capacity_change_restarts_membership_dwell() { + let start = tokio::time::Instant::now(); + let participants = vec![participant(1), participant(2)]; + let mut barrier = SplitMembershipSettleBarrier::default(); + assert!(!barrier.observe(&participants, start)); + + let changed = vec![ + participant(1), + SplitParticipant::new(make_id(2), 9_000, None), + ]; + assert!(!barrier.observe(&changed, start + Duration::from_secs(7))); + assert!(!barrier.observe(&changed, start + Duration::from_secs(14))); + assert!(barrier.observe(&changed, start + Duration::from_secs(15))); + } +} diff --git a/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs b/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs index ade2449ed..7d2af937e 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs @@ -2,7 +2,7 @@ use crate::inference::skippy; use anyhow::{Context, Result}; use skippy_coordinator::topology::{ LockedTopologyStage, TopologyNode, TopologyPlanningInput, TopologyStagePlan, - minimum_valid_context, plan_locked_topology, plan_topology, + minimum_valid_context, plan_locked_topology, plan_topology, plan_topology_with_stage0, }; use std::collections::HashMap; @@ -100,14 +100,28 @@ pub(super) struct PlannedRuntimeSliceTopology { pub(super) fn plan_split_topology(input: SplitTopologyPlanInput) -> Result { let plan = plan_topology(&topology_planning_input(input)).context("plan skippy split topology")?; + Ok(split_topology_plan(plan)) +} - Ok(SplitTopologyPlan { +fn plan_split_topology_with_stage0( + input: SplitTopologyPlanInput, + required_stage0: iroh::EndpointId, +) -> Result { + let input = topology_planning_input(input); + let required_stage0 = required_stage0.to_string(); + let plan = plan_topology_with_stage0(&input, &required_stage0) + .context("plan skippy split topology with canonical stage 0")?; + Ok(split_topology_plan(plan)) +} + +fn split_topology_plan(plan: skippy_coordinator::topology::TopologyPlan) -> SplitTopologyPlan { + SplitTopologyPlan { context_length: plan.context_length, parallel_lanes: plan.parallel_lanes, estimated_decode_network_ms_per_token: plan.estimated_decode_network_ms_per_token, decode_tpot_target_met: plan.decode_tpot_target_met, stages: plan.stages, - }) + } } fn topology_planning_input(input: SplitTopologyPlanInput) -> TopologyPlanningInput { @@ -163,6 +177,26 @@ pub(super) fn plan_runtime_slice_topology_with_resources( participants: &[SplitParticipant], excluded: &[SplitParticipantExclusion], resources: SplitTopologyResourceInputs, +) -> Result { + plan_runtime_slice_topology_with_resources_and_stage0( + topology_id, + model_ref, + package, + participants, + excluded, + resources, + None, + ) +} + +pub(super) fn plan_runtime_slice_topology_with_resources_and_stage0( + topology_id: &str, + model_ref: &str, + package: &skippy::SkippyPackageIdentity, + participants: &[SplitParticipant], + excluded: &[SplitParticipantExclusion], + resources: SplitTopologyResourceInputs, + required_stage0: Option, ) -> Result { tracing::info!( topology_id, @@ -176,13 +210,16 @@ pub(super) fn plan_runtime_slice_topology_with_resources( let participant_by_id = participant_index_by_id(participants); let plan_input = runtime_slice_plan_input(package, participants, resources); let plan = plan_runtime_slice_topology_result( - topology_id, - model_ref, - package, - participants, - excluded, - resources, + SplitPlanAttempt { + topology_id, + model_ref, + package, + participants, + excluded, + resources, + }, plan_input, + required_stage0, )?; let mut stages = map_runtime_slice_stages(plan.stages, &participant_by_id)?; @@ -255,32 +292,41 @@ pub(super) fn plan_locked_runtime_slice_topology_with_resources( }) } -fn plan_runtime_slice_topology_result( - topology_id: &str, - model_ref: &str, - package: &skippy::SkippyPackageIdentity, - participants: &[SplitParticipant], - excluded: &[SplitParticipantExclusion], +struct SplitPlanAttempt<'a> { + topology_id: &'a str, + model_ref: &'a str, + package: &'a skippy::SkippyPackageIdentity, + participants: &'a [SplitParticipant], + excluded: &'a [SplitParticipantExclusion], resources: SplitTopologyResourceInputs, +} + +fn plan_runtime_slice_topology_result( + attempt: SplitPlanAttempt<'_>, plan_input: SplitTopologyPlanInput, + required_stage0: Option, ) -> Result { - match plan_split_topology(plan_input) { + let result = match required_stage0 { + Some(node_id) => plan_split_topology_with_stage0(plan_input, node_id), + None => plan_split_topology(plan_input), + }; + match result { Ok(plan) => Ok(plan), Err(err) => { let reason = split_topology_failure_reason( - model_ref, - package, - participants, - excluded, - resources, + attempt.model_ref, + attempt.package, + attempt.participants, + attempt.excluded, + attempt.resources, ); tracing::warn!( - topology_id, - model_ref, + topology_id = attempt.topology_id, + model_ref = attempt.model_ref, error = %err, reason = %reason, - participants = ?split_participant_labels(participants), - excluded = ?split_participant_exclusion_labels(excluded), + participants = ?split_participant_labels(attempt.participants), + excluded = ?split_participant_exclusion_labels(attempt.excluded), "failed to plan resource-aware split runtime topology" ); Err(err.context(reason)) diff --git a/crates/mesh-llm-host-runtime/src/runtime/startup_handles.rs b/crates/mesh-llm-host-runtime/src/runtime/startup_handles.rs index 40a3659de..e3e540e3b 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/startup_handles.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/startup_handles.rs @@ -1,3 +1,4 @@ +use super::startup_retry::is_retryable_split_start_failure; use super::status::current_time_unix_ms; use super::status::single_quote_shell_arg; use super::{ @@ -406,11 +407,9 @@ where Err(err) => { drop(startup_load_guard); let err_msg = format!("{err:#}"); - let is_participant_shortage = err_msg.contains("at least two participating nodes") - || err_msg.contains("at least two stage participants"); - if is_participant_shortage { + if is_retryable_split_start_failure(&err_msg) { let _ = emit_event(OutputEvent::Info { - message: format!("Split waiting for peers: {err_msg}"), + message: format!("Split waiting to retry: {err_msg}"), context: Some(format!("model={model_name}")), }); } else { @@ -1006,7 +1005,7 @@ pub(super) async fn startup_prepare_launch( let local_capacity = ctx .pinned_gpu .map(|gpu| gpu.allocatable_vram_bytes()) - .unwrap_or_else(|| ctx.node.vram_bytes()); + .unwrap_or_else(|| ctx.node.local_runtime_capacity_bytes()); let model_bytes = startup_planning_model_bytes(&ctx).await?; let runtime_plan = startup_runtime_plan(ctx.split, local_capacity, model_bytes); let launch_kind = startup_launch_kind(runtime_plan, ctx.survey_launch_kind); diff --git a/crates/mesh-llm-host-runtime/src/runtime/startup_retry.rs b/crates/mesh-llm-host-runtime/src/runtime/startup_retry.rs new file mode 100644 index 000000000..ec547fc17 --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/runtime/startup_retry.rs @@ -0,0 +1,107 @@ +pub(super) fn is_retryable_split_start_failure(message: &str) -> bool { + split_participants_are_still_converging(message) + || split_control_transport_failed(message) + || split_stage_source_preparation_timed_out(message) +} + +fn split_participants_are_still_converging(message: &str) -> bool { + message.contains("at least two participating nodes") + || message.contains("at least two stage participants") + || message.contains("split_capacity_shortfall") + || message.contains("canonical coordinator") + || (message.contains("split topology lock stage") + && message.contains("matched 0 eligible nodes")) +} + +fn split_control_transport_failed(message: &str) -> bool { + let is_control_operation = message.contains("load split stage") + || message.contains("prepare split stage") + || message.contains("stage_control_unreachable"); + let is_transport_failure = message.contains("connection lost") + || message.contains("stream finished early") + || message.contains("timeout waiting for stage control response"); + is_control_operation && is_transport_failure +} + +fn split_stage_source_preparation_timed_out(message: &str) -> bool { + message.contains("stage_source_prepare_timeout") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn participant_shortage_is_retryable() { + assert!(is_retryable_split_start_failure( + "split runtime needs at least two participating nodes for model" + )); + assert!(is_retryable_split_start_failure( + "split runtime needs at least two stage participants" + )); + } + + #[test] + fn transient_capacity_shortfall_is_retryable() { + assert!(is_retryable_split_start_failure( + "split_capacity_shortfall: unable to plan split topology: max_placeable_layers_at_evaluated_shape=21/66" + )); + } + + #[test] + fn missing_locked_participant_is_retryable() { + assert!(is_retryable_split_start_failure( + "split topology lock stage 3 selector \"worker\" matched 0 eligible nodes; available: local" + )); + } + + #[test] + fn canonical_coordinator_mismatch_is_retryable() { + assert!(is_retryable_split_start_failure( + "split topology stage 0 node-a does not match canonical coordinator node-b" + )); + assert!(is_retryable_split_start_failure( + "split topology lock stage 0 must be canonical coordinator node-b" + )); + } + + #[test] + fn ambiguous_or_invalid_lock_is_not_retryable() { + assert!(!is_retryable_split_start_failure( + "split topology lock stage 3 selector \"worker\" matched 2 eligible nodes" + )); + assert!(!is_retryable_split_start_failure( + "split topology lock manifest abc does not match resolved package manifest def" + )); + } + + #[test] + fn stage_control_transport_failure_is_retryable() { + assert!(is_retryable_split_start_failure( + "load split stage stage-1: connection lost: closed" + )); + assert!(is_retryable_split_start_failure( + "prepare split stage stage-2: stage_control_unreachable: stream finished early" + )); + assert!(is_retryable_split_start_failure( + "load split stage stage-3: timeout waiting for stage control response" + )); + } + + #[test] + fn stage_source_preparation_timeout_is_retryable() { + assert!(is_retryable_split_start_failure( + "prepare split stage stage-1: stage_source_prepare_timeout: timed out waiting for stage source availability after 30m" + )); + } + + #[test] + fn runtime_or_unrelated_transport_failure_is_not_retryable() { + assert!(!is_retryable_split_start_failure( + "load skippy stage 0 runtime: unable to allocate CUDA0 buffer" + )); + assert!(!is_retryable_split_start_failure( + "artifact transfer connection lost: closed" + )); + } +} diff --git a/crates/mesh-llm-host-runtime/src/sdk.rs b/crates/mesh-llm-host-runtime/src/sdk.rs index 615a49a1f..fc76ea12a 100644 --- a/crates/mesh-llm-host-runtime/src/sdk.rs +++ b/crates/mesh-llm-host-runtime/src/sdk.rs @@ -503,9 +503,7 @@ impl ServingController for EmbeddedServingController { let capabilities = models::runtime_verified_model_capabilities( &model_id, &model_path, - models::RuntimeMediaCapabilityEvidence { - vision_projector_loaded: false, - }, + models::RuntimeMediaCapabilityEvidence::default(), ); let mut state = self.inner.lock().await; diff --git a/crates/mesh-llm-types/src/mesh/mod.rs b/crates/mesh-llm-types/src/mesh/mod.rs index b5c9b5aaf..4dda2bbb0 100644 --- a/crates/mesh-llm-types/src/mesh/mod.rs +++ b/crates/mesh-llm-types/src/mesh/mod.rs @@ -12,6 +12,21 @@ pub const DEMAND_TTL_SECS: u64 = 86400; pub const MAX_SPLIT_RTT_MS: u32 = 80; +/// Split admission RTT ceiling in milliseconds. Defaults to +/// [`MAX_SPLIT_RTT_MS`]; the `MESH_SPLIT_MAX_RTT_MS` environment variable +/// overrides it for long-haul WAN experiments where the fixed ceiling would +/// reject every peer. +pub fn max_split_rtt_ms() -> u32 { + static VALUE: std::sync::OnceLock = std::sync::OnceLock::new(); + *VALUE.get_or_init(|| { + std::env::var("MESH_SPLIT_MAX_RTT_MS") + .ok() + .and_then(|raw| raw.trim().parse().ok()) + .filter(|ms| *ms > 0) + .unwrap_or(MAX_SPLIT_RTT_MS) + }) +} + #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] #[serde(rename_all = "snake_case")] pub enum ModelSourceKind { @@ -418,3 +433,15 @@ fn identity_hash_for(input: &str) -> String { hasher.update(input.as_bytes()); hex::encode(hasher.finalize()) } + +/// Whether relay-only stage paths may participate in splits. Off by default; +/// `MESH_SPLIT_ALLOW_RELAY=1` enables it for WAN experiments where NAT +/// prevents direct QUIC paths between stage peers. +pub fn split_allow_relay_paths() -> bool { + static VALUE: std::sync::OnceLock = std::sync::OnceLock::new(); + *VALUE.get_or_init(|| { + std::env::var("MESH_SPLIT_ALLOW_RELAY") + .map(|raw| raw.trim() == "1") + .unwrap_or(false) + }) +} diff --git a/crates/mesh-llm/src/commands/mod.rs b/crates/mesh-llm/src/commands/mod.rs index cce410694..e6d4f4708 100644 --- a/crates/mesh-llm/src/commands/mod.rs +++ b/crates/mesh-llm/src/commands/mod.rs @@ -164,6 +164,7 @@ async fn dispatch_model_prepare(cmd: &Command) -> Result<()> { flavor, timeout, mesh_llm_ref, + experimental: false, dry_run: *dry_run, confirm: *confirm, follow: *follow, diff --git a/crates/mesh-llm/src/commands/models/mod.rs b/crates/mesh-llm/src/commands/models/mod.rs index 24e59a293..1ff45b7d0 100644 --- a/crates/mesh-llm/src/commands/models/mod.rs +++ b/crates/mesh-llm/src/commands/models/mod.rs @@ -447,6 +447,7 @@ pub async fn dispatch_models_command(command: &ModelsCommand) -> Result<()> { flavor, timeout, mesh_llm_ref, + experimental, dry_run, confirm, follow, @@ -466,6 +467,7 @@ pub async fn dispatch_models_command(command: &ModelsCommand) -> Result<()> { flavor, timeout, mesh_llm_ref, + experimental: *experimental, dry_run: *dry_run, confirm: *confirm, follow: *follow, diff --git a/crates/model-artifact/src/gguf.rs b/crates/model-artifact/src/gguf.rs index 802caf4a5..80e89aecc 100644 --- a/crates/model-artifact/src/gguf.rs +++ b/crates/model-artifact/src/gguf.rs @@ -198,6 +198,29 @@ fn read_gguf_value_as_u32(f: &mut std::fs::File, typ: GgufType) -> std::io::Resu } } +fn read_gguf_value_as_u32_list( + f: &mut std::fs::File, + typ: GgufType, +) -> std::io::Result>> { + if typ != GgufType::Array { + return Ok(read_gguf_value_as_u32(f, typ)?.map(|value| vec![value])); + } + + let elem_type = GgufType::from_u32(read_u32(f)?) + .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "bad array type"))?; + let count = read_bounded_len(f, MAX_GGUF_ARRAY_ELEMENTS, "array")?; + let mut values = Vec::with_capacity(count); + let mut supported = true; + for _ in 0..count { + match read_gguf_value_as_u32(f, elem_type)? { + Some(value) if supported => values.push(value), + Some(_) => {} + None => supported = false, + } + } + Ok(supported.then_some(values)) +} + fn read_gguf_value_as_f32(f: &mut std::fs::File, typ: GgufType) -> std::io::Result> { match typ { GgufType::Float32 => { @@ -212,6 +235,15 @@ fn read_gguf_value_as_f32(f: &mut std::fs::File, typ: GgufType) -> std::io::Resu } } +fn read_gguf_value_as_bool(f: &mut std::fs::File, typ: GgufType) -> std::io::Result> { + if typ == GgufType::Bool { + let mut value = [0u8; 1]; + f.read_exact(&mut value)?; + return Ok(Some(value[0] != 0)); + } + Ok(read_gguf_value_as_u32(f, typ)?.map(|value| value != 0)) +} + fn read_gguf_value_as_string_opt( f: &mut std::fs::File, typ: GgufType, @@ -234,6 +266,7 @@ pub struct GgufCompactMeta { pub embedding_size: u32, pub head_count: u32, pub kv_head_count: u32, + pub kv_head_counts: Vec, pub layer_count: u32, pub feed_forward_length: u32, pub key_length: u32, @@ -247,10 +280,20 @@ pub struct GgufCompactMeta { pub nextn_predict_layers: u32, } +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct GgufProjectorMeta { + pub has_vision_encoder: Option, + pub has_audio_encoder: Option, +} + impl GgufCompactMeta { pub fn effective_kv_head_count(&self) -> Option { if self.kv_head_count > 0 { Some(self.kv_head_count) + } else if let Some(kv_head_count) = self.kv_head_counts.iter().copied().max() + && kv_head_count > 0 + { + Some(kv_head_count) } else if self.head_count > 0 { Some(self.head_count) } else { @@ -348,8 +391,12 @@ pub fn scan_gguf_compact_meta(path: &Path) -> Option { meta.head_count = v; } } else if key.ends_with(".attention.head_count_kv") { - if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) { - meta.kv_head_count = v; + if let Ok(Some(values)) = read_gguf_value_as_u32_list(&mut f, vtype) { + if values.len() == 1 { + meta.kv_head_count = values[0]; + } else { + meta.kv_head_counts = values; + } } } else if key.ends_with(".block_count") { if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) { @@ -419,6 +466,28 @@ pub fn scan_gguf_compact_meta(path: &Path) -> Option { Some(meta) } +/// Scan the modality flags stored in a multimodal projector GGUF. +pub fn scan_gguf_projector_meta(path: &Path) -> Option { + let GgufHeader { + file: mut f, n_kv, .. + } = open_gguf_header(path)?; + let mut meta = GgufProjectorMeta::default(); + for _ in 0..n_kv { + let key = read_gguf_string(&mut f).ok()?; + let value_type = GgufType::from_u32(read_u32(&mut f).ok()?)?; + match key.as_str() { + "clip.has_vision_encoder" => { + meta.has_vision_encoder = read_gguf_value_as_bool(&mut f, value_type).ok()?; + } + "clip.has_audio_encoder" => { + meta.has_audio_encoder = read_gguf_value_as_bool(&mut f, value_type).ok()?; + } + _ => skip_gguf_value(&mut f, value_type).ok()?, + } + } + Some(meta) +} + fn align_offset(value: u64, alignment: u32) -> u64 { let alignment = u64::from(alignment.max(1)); let remainder = value % alignment; @@ -634,6 +703,21 @@ mod tests { bytes.extend_from_slice(&value.to_le_bytes()); } + fn push_bool_kv(bytes: &mut Vec, key: &str, value: bool) { + push_gguf_string(bytes, key); + bytes.extend_from_slice(&(GgufType::Bool as u32).to_le_bytes()); + bytes.push(u8::from(value)); + } + + fn push_u32_array_kv(bytes: &mut Vec, key: &str, values: &[u32]) { + push_gguf_string(bytes, key); + bytes.extend_from_slice(&(GgufType::Array as u32).to_le_bytes()); + push_array_header(bytes, GgufType::Uint32, values.len() as u64); + for value in values { + bytes.extend_from_slice(&value.to_le_bytes()); + } + } + fn push_tensor_info(bytes: &mut Vec, name: &str, offset: u64) { push_gguf_string(bytes, name); bytes.extend_from_slice(&1u32.to_le_bytes()); @@ -750,6 +834,61 @@ mod tests { let _ = std::fs::remove_file(path); } + #[test] + fn scan_gguf_compact_meta_prices_per_layer_kv_head_counts() { + let mut bytes = Vec::new(); + bytes.extend_from_slice(b"GGUF"); + bytes.extend_from_slice(&2u32.to_le_bytes()); + bytes.extend_from_slice(&0i64.to_le_bytes()); + bytes.extend_from_slice(&7i64.to_le_bytes()); + push_gguf_string(&mut bytes, "general.architecture"); + bytes.extend_from_slice(&(GgufType::String as u32).to_le_bytes()); + push_gguf_string(&mut bytes, "inkling"); + push_u32_kv(&mut bytes, "inkling.embedding_length", 6144); + push_u32_kv(&mut bytes, "inkling.attention.head_count", 64); + let kv_head_counts = (0..66) + .map(|layer| if layer % 6 == 5 { 8 } else { 16 }) + .collect::>(); + push_u32_array_kv( + &mut bytes, + "inkling.attention.head_count_kv", + &kv_head_counts, + ); + push_u32_kv(&mut bytes, "inkling.block_count", 66); + push_u32_kv(&mut bytes, "inkling.attention.key_length", 128); + push_u32_kv(&mut bytes, "inkling.attention.value_length", 128); + + let path = write_bytes("model-artifact-gguf-inkling-kv-head-counts", &bytes); + let meta = scan_gguf_compact_meta(&path).expect("should parse GGUF"); + assert_eq!(meta.kv_head_count, 0); + assert_eq!(meta.kv_head_counts, kv_head_counts); + assert_eq!(meta.effective_kv_head_count(), Some(16)); + assert_eq!(meta.k_cache_bytes_per_token_f16(), Some(247_808)); + assert_eq!(meta.v_cache_bytes_per_token_f16(), Some(247_808)); + assert_eq!(meta.kv_cache_bytes_per_token_f16(), Some(495_616)); + let _ = std::fs::remove_file(path); + } + + #[test] + fn scan_gguf_projector_meta_preserves_vision_and_audio_flags() { + let mut bytes = Vec::new(); + bytes.extend_from_slice(b"GGUF"); + bytes.extend_from_slice(&2u32.to_le_bytes()); + bytes.extend_from_slice(&0i64.to_le_bytes()); + bytes.extend_from_slice(&3i64.to_le_bytes()); + push_gguf_string(&mut bytes, "general.architecture"); + bytes.extend_from_slice(&(GgufType::String as u32).to_le_bytes()); + push_gguf_string(&mut bytes, "clip"); + push_bool_kv(&mut bytes, "clip.has_vision_encoder", true); + push_bool_kv(&mut bytes, "clip.has_audio_encoder", true); + + let path = write_bytes("model-artifact-gguf-inkling-projector", &bytes); + let meta = scan_gguf_projector_meta(&path).expect("should parse projector GGUF"); + assert_eq!(meta.has_vision_encoder, Some(true)); + assert_eq!(meta.has_audio_encoder, Some(true)); + let _ = std::fs::remove_file(path); + } + #[test] fn scan_gguf_compact_meta_preserves_nextn_predict_layers() { let mut bytes = Vec::new(); diff --git a/crates/model-artifact/src/gguf/kv_cache.rs b/crates/model-artifact/src/gguf/kv_cache.rs index 17935797a..4c4c7e4ec 100644 --- a/crates/model-artifact/src/gguf/kv_cache.rs +++ b/crates/model-artifact/src/gguf/kv_cache.rs @@ -153,9 +153,24 @@ fn cache_bytes_per_token( vector_length: u32, cache_type: GgufKvCacheType, ) -> Option { - let kv_heads = u64::from(meta.kv_cache_head_count()?); let vector_length = u64::from((vector_length > 0).then_some(vector_length)?); let layers = u64::from((meta.layer_count > 0).then_some(meta.layer_count)?); + // Models with per-layer KV head counts (e.g. inkling's hybrid attention) + // price each layer by its own head count rather than a single global one. + // GLM-DSA keeps its absorbed-MLA special case via kv_cache_head_count. + if meta.architecture != "glm-dsa" + && meta.kv_head_counts.len() == layers as usize + && meta.kv_head_counts.iter().all(|head_count| *head_count > 0) + { + return meta + .kv_head_counts + .iter() + .try_fold(0u64, |total, head_count| { + let elements = u64::from(*head_count).checked_mul(vector_length)?; + total.checked_add(cache_type.bytes_for_elements(elements)?) + }); + } + let kv_heads = u64::from(meta.kv_cache_head_count()?); let elements_per_layer = kv_heads.checked_mul(vector_length)?; cache_type .bytes_for_elements(elements_per_layer)? diff --git a/crates/model-package/src/bin/queue-unsloth-layer-packages.rs b/crates/model-package/src/bin/queue-unsloth-layer-packages.rs index 49978fb75..b7af2dd04 100644 --- a/crates/model-package/src/bin/queue-unsloth-layer-packages.rs +++ b/crates/model-package/src/bin/queue-unsloth-layer-packages.rs @@ -9,7 +9,7 @@ use hf_hub::{ repository::{AddSource, ModelInfo}, }; use model_package::jobs::{CpuJobPlan, HfJobsClient, JobInfo, JobSpec, JobStage, JobVolume}; -use model_package::prepare::{self, DiscoveredQuant}; +use model_package::prepare::{self, DiscoveredProjector, DiscoveredQuant}; use model_package::script; use serde::Serialize; use serde_json::Value; @@ -54,7 +54,10 @@ struct RankedModel { #[derive(Debug, Clone)] struct Candidate { model: RankedModel, + source_revision: String, + source_pipeline_tag: String, quant: DiscoveredQuant, + projectors: Vec, target_repo: String, model_layer_repos: Vec, model_id: String, @@ -88,7 +91,9 @@ struct QueueMarker<'a> { schema_version: u32, queued_at: String, source_repo: &'a str, + source_revision: &'a str, source_file: &'a str, + source_projectors: &'a [DiscoveredProjector], quant: &'a str, target_repo: &'a str, model_id: &'a str, @@ -103,7 +108,9 @@ struct QueueFailureMarker<'a> { schema_version: u32, failed_at: String, source_repo: &'a str, + source_revision: &'a str, source_file: &'a str, + source_projectors: &'a [DiscoveredProjector], quant: &'a str, target_repo: &'a str, model_id: &'a str, @@ -155,10 +162,7 @@ async fn main() -> Result<()> { ranked.len(), args.author ); - println!( - "Preferred 4-bit quants: {}", - args.quant_preference.join(", ") - ); + println!("Preferred quants: {}", args.quant_preference.join(", ")); println!( "Split candidates: selected quant requires more than {} with 10% runtime headroom.", prepare::format_size(args.split_candidate_vram_bytes) @@ -249,11 +253,12 @@ async fn main() -> Result<()> { QueueStatus::Missing | QueueStatus::StaleQueued => {} } + let source_total_bytes = candidate_source_total_bytes(&candidate); let job_plan = model_package::jobs::plan_cpu_job_from_hardware( &hardware, &args.flavor, args.timeout_seconds, - candidate.quant.total_bytes, + source_total_bytes, )?; total_max_cost_usd += job_plan.max_cost_usd; @@ -268,10 +273,8 @@ async fn main() -> Result<()> { candidate.model_id, candidate.target_repo, candidate.quant.name, - prepare::format_size(candidate.quant.total_bytes), - prepare::format_size(estimated_bucket_workspace_bytes( - candidate.quant.total_bytes - )), + prepare::format_size(source_total_bytes), + prepare::format_size(estimated_bucket_workspace_bytes(source_total_bytes)), shard_label(candidate.quant.shard_count), rank_label(&candidate.model), candidate.family, @@ -684,17 +687,28 @@ async fn build_candidate( eprintln!("skip {}: {reason}", model.repo_id); return Ok(None); } - - let quants = match prepare::list_quants(client, &model.repo_id).await { - Ok(quants) => quants, - Err(err) => { - eprintln!( - "skip {}: failed to list GGUF quants: {err:#}", - model.repo_id - ); - return Ok(None); - } + let Some(source_revision) = source_info.sha.clone() else { + eprintln!( + "skip {}: source repo info has no immutable commit SHA", + model.repo_id + ); + return Ok(None); }; + let source_pipeline_tag = + model_pipeline_tag(&source_info).expect("compatible model must have a pipeline tag"); + + let inventory = + match prepare::list_inventory(client, &model.repo_id, Some(&source_revision)).await { + Ok(inventory) => inventory, + Err(err) => { + eprintln!( + "skip {}: failed to list GGUF quants: {err:#}", + model.repo_id + ); + return Ok(None); + } + }; + let quants = inventory.quants; let Some(quant) = select_preferred_quant(&quants, &args.quant_preference) else { let available = quants @@ -703,7 +717,7 @@ async fn build_candidate( .collect::>() .join(", "); eprintln!( - "skip {}: no preferred 4-bit quant ({}); available: {}", + "skip {}: no preferred quant ({}); available: {}", model.repo_id, args.quant_preference.join(", "), if available.is_empty() { @@ -737,7 +751,10 @@ async fn build_candidate( Ok(Some(Candidate { model, + source_revision, + source_pipeline_tag, quant, + projectors: inventory.projectors, target_repo, model_layer_repos, model_id, @@ -772,6 +789,15 @@ fn select_preferred_quant( }) } +fn candidate_source_total_bytes(candidate: &Candidate) -> u64 { + candidate + .projectors + .iter() + .fold(candidate.quant.total_bytes, |total, projector| { + total.saturating_add(projector.total_bytes) + }) +} + async fn candidate_status( client: &HFClient, candidate: &Candidate, @@ -841,12 +867,15 @@ fn model_split_compatibility(info: &ModelInfo) -> SplitCompatibility { "missing text-generation pipeline_tag".to_string(), ); }; - if pipeline_tag.eq_ignore_ascii_case("text-generation") { + if pipeline_tag.eq_ignore_ascii_case("text-generation") + || (pipeline_tag.eq_ignore_ascii_case("image-text-to-text") + && info.id.to_ascii_lowercase().contains("inkling")) + { return SplitCompatibility::Compatible; } SplitCompatibility::Incompatible(format!( - "unsupported pipeline_tag '{pipeline_tag}'; layer packages currently require text-generation" + "unsupported pipeline_tag '{pipeline_tag}'; automated layer packages currently require text-generation or a supported multimodal family" )) } @@ -988,7 +1017,9 @@ async fn write_queue_marker(client: &HFClient, candidate: &Candidate, args: &Arg schema_version: 1, queued_at: Utc::now().to_rfc3339(), source_repo: &candidate.model.repo_id, + source_revision: &candidate.source_revision, source_file: &candidate.quant.first_file, + source_projectors: &candidate.projectors, quant: &candidate.quant.name, target_repo: &candidate.target_repo, model_id: &candidate.model_id, @@ -1019,7 +1050,9 @@ async fn write_queue_failure_marker( schema_version: 1, failed_at: Utc::now().to_rfc3339(), source_repo: &submitted.candidate.model.repo_id, + source_revision: &submitted.candidate.source_revision, source_file: &submitted.candidate.quant.first_file, + source_projectors: &submitted.candidate.projectors, quant: &submitted.candidate.quant.name, target_repo: &submitted.candidate.target_repo, model_id: &submitted.candidate.model_id, @@ -1064,17 +1097,41 @@ fn job_spec_with_token( hf_token: &str, job_plan: &CpuJobPlan, ) -> Result { + let projector_bytes = candidate + .projectors + .iter() + .try_fold(0u64, |total, projector| { + total.checked_add(projector.total_bytes) + }) + .context("source GGUF and projector sizes overflowed u64")?; + let source_total_bytes = candidate + .quant + .total_bytes + .checked_add(projector_bytes) + .context("source GGUF and projector sizes overflowed u64")?; let mut environment = HashMap::new(); environment.insert("SOURCE_REPO".into(), candidate.model.repo_id.clone()); environment.insert("SOURCE_FILE".into(), candidate.quant.first_file.clone()); environment.insert("SOURCE_QUANT".into(), candidate.quant.name.clone()); - environment.insert( - "SOURCE_TOTAL_BYTES".into(), - candidate.quant.total_bytes.to_string(), - ); + environment.insert("SOURCE_TOTAL_BYTES".into(), source_total_bytes.to_string()); environment.insert("TARGET_REPO".into(), candidate.target_repo.clone()); environment.insert("MODEL_ID".into(), candidate.model_id.clone()); - environment.insert("SOURCE_REVISION".into(), "main".into()); + environment.insert("SOURCE_REVISION".into(), candidate.source_revision.clone()); + environment.insert( + "SOURCE_PIPELINE_TAG".into(), + candidate.source_pipeline_tag.clone(), + ); + if !candidate.projectors.is_empty() { + environment.insert( + "SOURCE_PROJECTOR_FILES".into(), + candidate + .projectors + .iter() + .map(|projector| projector.path.as_str()) + .collect::>() + .join("\n"), + ); + } environment.insert("MESH_LLM_REF".into(), args.mesh_llm_ref.clone()); environment.insert( "CATALOG_CREATE_PR".into(), @@ -1098,12 +1155,14 @@ fn job_spec_with_token( source: "meshllm/layer-split-output".into(), mount_path: "/bucket".into(), read_only: None, + revision: None, }, JobVolume { volume_type: "model".into(), source: candidate.model.repo_id.clone(), mount_path: "/source".into(), read_only: Some(true), + revision: Some(candidate.source_revision.clone()), }, ], }) @@ -1316,7 +1375,7 @@ mod tests { use hf_hub::repository::ModelInfo; use super::{ - Args, Candidate, DiscoveredQuant, RankedModel, SplitCompatibility, + Args, Candidate, DiscoveredProjector, DiscoveredQuant, RankedModel, SplitCompatibility, estimated_bucket_workspace_bytes, job_spec_with_token, json_layer_package_repo, model_family_key, model_layer_repos, model_split_compatibility, }; @@ -1365,6 +1424,20 @@ mod tests { ); } + #[test] + fn split_compatibility_accepts_inkling_multimodal_pipeline() { + let info = model_info(serde_json::json!({ + "id": "unsloth/inkling-GGUF", + "pipeline_tag": "image-text-to-text", + "tags": ["gguf", "multimodal"] + })); + + assert_eq!( + model_split_compatibility(&info), + SplitCompatibility::Compatible + ); + } + #[test] fn split_compatibility_rejects_media_generation_pipeline() { let info = model_info(serde_json::json!({ @@ -1421,12 +1494,18 @@ mod tests { recent_rank: None, popular_rank: None, }, + source_revision: "0123456789abcdef".to_string(), + source_pipeline_tag: "text-generation".to_string(), quant: DiscoveredQuant { name: "UD-Q4_K_XL".to_string(), shard_count: 10, total_bytes: 401, first_file: "UD-Q4_K_XL/GLM-5-UD-Q4_K_XL-00001-of-00010.gguf".to_string(), }, + projectors: vec![DiscoveredProjector { + path: "mmproj-BF16.gguf".to_string(), + total_bytes: 183, + }], target_repo: "meshllm/GLM-5-UD-Q4_K_XL-layers".to_string(), model_layer_repos: vec!["meshllm/GLM-5-UD-Q4_K_XL-layers".to_string()], model_id: "unsloth/GLM-5-GGUF:UD-Q4_K_XL".to_string(), @@ -1480,13 +1559,27 @@ mod tests { spec.environment .get("SOURCE_TOTAL_BYTES") .map(String::as_str), - Some("401") + Some("584") + ); + assert_eq!( + spec.environment.get("SOURCE_REVISION").map(String::as_str), + Some("0123456789abcdef") + ); + assert_eq!( + spec.environment + .get("SOURCE_PROJECTOR_FILES") + .map(String::as_str), + Some("mmproj-BF16.gguf") ); assert_eq!(spec.volumes.len(), 2); assert_eq!(spec.volumes[0].volume_type, "bucket"); assert_eq!(spec.volumes[0].mount_path, "/bucket"); assert_eq!(spec.volumes[1].volume_type, "model"); assert_eq!(spec.volumes[1].source, candidate.model.repo_id); + assert_eq!( + spec.volumes[1].revision.as_deref(), + Some("0123456789abcdef") + ); assert_eq!(spec.volumes[1].mount_path, "/source"); assert_eq!(spec.volumes[1].read_only, Some(true)); } diff --git a/crates/model-package/src/jobs.rs b/crates/model-package/src/jobs.rs index 76eafd2b5..301c3069b 100644 --- a/crates/model-package/src/jobs.rs +++ b/crates/model-package/src/jobs.rs @@ -41,6 +41,8 @@ pub struct JobVolume { pub mount_path: String, #[serde(rename = "readOnly", skip_serializing_if = "Option::is_none")] pub read_only: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub revision: Option, } #[derive(Debug, Clone, Deserialize, Serialize)] diff --git a/crates/model-package/src/prepare.rs b/crates/model-package/src/prepare.rs index 4666b475c..32a5e44dd 100644 --- a/crates/model-package/src/prepare.rs +++ b/crates/model-package/src/prepare.rs @@ -16,23 +16,28 @@ use crate::permissions::PermissionCheck; /// Parameters for a model-package job. pub struct PrepareParams { pub source_repo: String, + pub source_revision: Option, pub quant: Option, pub target: Option, pub model_id: Option, pub flavor: String, pub timeout_seconds: u64, pub mesh_llm_ref: String, + pub experimental: bool, pub hf_token: Option, } /// A fully resolved model-package job, ready to submit. pub struct PrepareJob { pub source_repo: String, + pub source_revision: String, pub source_file: String, + pub projectors: Vec, pub target_repo: String, pub model_id: String, pub namespace: String, pub catalog_create_pr: bool, + pub experimental: bool, pub job_plan: CpuJobPlan, pub spec: JobSpec, } @@ -50,13 +55,39 @@ pub struct DiscoveredQuant { pub first_file: String, } +/// A multimodal projector sidecar discovered in a HF model repo. +#[derive(Debug, Clone, Serialize, Eq, PartialEq)] +pub struct DiscoveredProjector { + /// Repo-relative projector path. + pub path: String, + /// Projector size in bytes. + pub total_bytes: u64, +} + +/// GGUF artifacts discovered in a HF model repo, separated by runtime role. +#[derive(Debug, Clone, Default)] +pub struct RepoGgufInventory { + pub quants: Vec, + pub projectors: Vec, +} + /// List all available GGUF quant variants in a HF model repo. pub async fn list_quants(client: &HFClient, repo: &str) -> Result> { + Ok(list_inventory(client, repo, None).await?.quants) +} + +/// List model GGUF quants and multimodal projectors at an optional repo revision. +pub async fn list_inventory( + client: &HFClient, + repo: &str, + revision: Option<&str>, +) -> Result { let (owner, name) = parse_repo(repo)?; let hf_repo = client.model(&owner, &name); let stream = hf_repo .list_tree() + .maybe_revision(revision.map(str::to_string)) .recursive(true) .send() .context("list repo tree")?; @@ -74,7 +105,33 @@ pub async fn list_quants(client: &HFClient, repo: &str) -> Result) -> RepoGgufInventory { + let (projector_files, model_files): (Vec<_>, Vec<_>) = gguf_files + .into_iter() + .partition(|(path, _)| is_projector_path(path)); + let mut projectors = projector_files + .into_iter() + .map(|(path, total_bytes)| DiscoveredProjector { path, total_bytes }) + .collect::>(); + projectors.sort_by(|a, b| a.path.cmp(&b.path)); + RepoGgufInventory { + quants: discover_quants_from_gguf_files(model_files), + projectors, + } +} + +fn is_projector_path(path: &str) -> bool { + std::path::Path::new(path) + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| { + let name = name.to_ascii_lowercase(); + name.starts_with("mmproj") && name.ends_with(".gguf") + }) } /// Group GGUF files into quant variants. @@ -122,7 +179,28 @@ pub async fn resolve( .as_deref() .context("--quant is required when submitting a job")?; - let quants = list_quants(client, ¶ms.source_repo).await?; + let (owner, name) = parse_repo(¶ms.source_repo)?; + let hf_repo = client.model(&owner, &name); + let requested_revision = params.source_revision.as_deref().unwrap_or("main"); + let source_info = hf_repo + .info() + .revision(requested_revision.to_string()) + .send() + .await + .with_context(|| { + format!( + "resolve source revision {}@{requested_revision}", + params.source_repo + ) + })?; + let source_revision = source_info + .sha + .context("source repo info did not include a commit SHA")?; + let source_pipeline_tag = source_info + .pipeline_tag + .unwrap_or_else(|| "text-generation".to_string()); + let inventory = list_inventory(client, ¶ms.source_repo, Some(&source_revision)).await?; + let quants = inventory.quants; if quants.is_empty() { anyhow::bail!("No GGUF files found in {}", params.source_repo); @@ -178,11 +256,22 @@ pub async fn resolve( &matched.name, )?; + let projector_bytes = inventory + .projectors + .iter() + .try_fold(0u64, |total, projector| { + total.checked_add(projector.total_bytes) + }) + .context("source GGUF and projector sizes overflowed u64")?; + let source_total_bytes = matched + .total_bytes + .checked_add(projector_bytes) + .context("source GGUF and projector sizes overflowed u64")?; let job_plan = crate::jobs::plan_cpu_job( &crate::jobs::hf_endpoint(), ¶ms.flavor, params.timeout_seconds, - matched.total_bytes, + source_total_bytes, ) .await?; @@ -191,19 +280,31 @@ pub async fn resolve( environment.insert("SOURCE_REPO".into(), params.source_repo.clone()); environment.insert("SOURCE_FILE".into(), source_file.clone()); environment.insert("SOURCE_QUANT".into(), matched.name.clone()); - environment.insert("SOURCE_TOTAL_BYTES".into(), matched.total_bytes.to_string()); + environment.insert("SOURCE_TOTAL_BYTES".into(), source_total_bytes.to_string()); environment.insert("TARGET_REPO".into(), target_repo.clone()); environment.insert("MODEL_ID".into(), model_id.clone()); - environment.insert("SOURCE_REVISION".into(), "main".into()); + environment.insert("SOURCE_REVISION".into(), source_revision.clone()); + environment.insert("SOURCE_PIPELINE_TAG".into(), source_pipeline_tag); + if !inventory.projectors.is_empty() { + environment.insert( + "SOURCE_PROJECTOR_FILES".into(), + inventory + .projectors + .iter() + .map(|projector| projector.path.as_str()) + .collect::>() + .join("\n"), + ); + } environment.insert("MESH_LLM_REF".into(), params.mesh_llm_ref.clone()); + let catalog_create_pr = params.experimental || permissions.catalog_create_pr; environment.insert( "CATALOG_CREATE_PR".into(), - if permissions.catalog_create_pr { - "true" - } else { - "false" - } - .into(), + if catalog_create_pr { "true" } else { "false" }.into(), + ); + environment.insert( + "PACKAGE_EXPERIMENTAL".into(), + if params.experimental { "true" } else { "false" }.into(), ); // The HF Jobs API passes secrets as env vars inside the container. @@ -219,12 +320,14 @@ pub async fn resolve( source: "meshllm/layer-split-output".into(), mount_path: "/bucket".into(), read_only: None, + revision: None, }, JobVolume { volume_type: "model".into(), source: params.source_repo.clone(), mount_path: "/source".into(), read_only: Some(true), + revision: Some(source_revision.clone()), }, ]; @@ -241,11 +344,14 @@ pub async fn resolve( Ok(PrepareJob { source_repo: params.source_repo, + source_revision, source_file, + projectors: inventory.projectors, target_repo, model_id, namespace: permissions.namespace.clone(), - catalog_create_pr: permissions.catalog_create_pr, + catalog_create_pr, + experimental: params.experimental, job_plan, spec, }) @@ -336,6 +442,27 @@ mod tests { assert_eq!(names, vec!["Q4_K_M", "Q8_0"]); } + #[test] + fn separates_multimodal_projectors_from_model_quants() { + let inventory = discover_inventory_from_gguf_files(vec![ + ( + "UD-Q2_K_XL/Inkling-UD-Q2_K_XL-00001-of-00008.gguf".to_string(), + 317, + ), + ("mmproj-BF16.gguf".to_string(), 183), + ]); + + assert_eq!(inventory.quants.len(), 1); + assert_eq!(inventory.quants[0].name, "UD-Q2_K_XL"); + assert_eq!( + inventory.projectors, + vec![DiscoveredProjector { + path: "mmproj-BF16.gguf".to_string(), + total_bytes: 183, + }] + ); + } + #[test] fn accepts_explicit_model_id_coordinate() { let model_id = resolve_model_id( diff --git a/crates/model-package/src/script.rs b/crates/model-package/src/script.rs index edd1946fe..fd6a2912b 100644 --- a/crates/model-package/src/script.rs +++ b/crates/model-package/src/script.rs @@ -129,7 +129,14 @@ mod tests { #[test] fn embedded_script_writes_rich_model_card() { - assert!(EMBEDDED_SCRIPT.contains("pipeline_tag: text-generation")); + assert!(EMBEDDED_SCRIPT.contains("SOURCE_PIPELINE_TAG")); + assert!(EMBEDDED_SCRIPT.contains("PACKAGE_EXPERIMENTAL")); + assert!(EMBEDDED_SCRIPT.contains("pipeline_tag: {yaml_quote(source_pipeline_tag)}")); + assert!(EMBEDDED_SCRIPT.contains("- experimental")); + assert!(EMBEDDED_SCRIPT.contains("Experimental package:")); + assert!(EMBEDDED_SCRIPT.contains("resolve_upstream_license")); + assert!(EMBEDDED_SCRIPT.contains("license_frontmatter")); + assert!(EMBEDDED_SCRIPT.contains("could not resolve upstream license metadata")); assert!(EMBEDDED_SCRIPT.contains("- openai-compatible")); assert!(EMBEDDED_SCRIPT.contains("## Model Overview")); assert!(EMBEDDED_SCRIPT.contains("## Highlights")); @@ -198,7 +205,20 @@ mod tests { assert!(EMBEDDED_SCRIPT.contains(r#"MOUNTED_SOURCE_PATH="/source/${SOURCE_FILE}""#)); assert!(EMBEDDED_SCRIPT.contains(r#"WRITE_PACKAGE_INPUT="$MOUNTED_SOURCE_PATH""#)); assert!(EMBEDDED_SCRIPT.contains(r#"--source-file "$SOURCE_FILE""#)); + assert!(EMBEDDED_SCRIPT.contains("SOURCE_PROJECTOR_FILES")); + assert!( + EMBEDDED_SCRIPT + .contains(r#"WRITE_PACKAGE_PROJECTOR_ARGS+=(--projector "$PROJECTOR_PATH")"#) + ); + assert!(EMBEDDED_SCRIPT.contains(r#""${WRITE_PACKAGE_PROJECTOR_ARGS[@]}""#)); assert!(EMBEDDED_SCRIPT.contains(r#"time "$SLICER" write-package "$WRITE_PACKAGE_INPUT""#)); assert!(!EMBEDDED_SCRIPT.contains(r#"time $SLICER write-package "$SOURCE_PATH""#)); } + + #[test] + fn embedded_script_preserves_catalog_source_revision() { + assert!(EMBEDDED_SCRIPT.contains(r#""source_revision": source_revision"#)); + assert!(EMBEDDED_SCRIPT.contains(r#"variants[variant_name]["source"] = source_entry"#)); + assert!(EMBEDDED_SCRIPT.contains(r#"existing_variant["source"] = source_entry"#)); + } } diff --git a/crates/model-package/src/scripts/split-model-job.sh b/crates/model-package/src/scripts/split-model-job.sh index 03541fc87..76605a210 100755 --- a/crates/model-package/src/scripts/split-model-job.sh +++ b/crates/model-package/src/scripts/split-model-job.sh @@ -6,8 +6,11 @@ set -euo pipefail # # Environment variables (set by mesh-llm model-package job spec): # SOURCE_REPO, SOURCE_FILE, SOURCE_QUANT, TARGET_REPO, MODEL_ID, SOURCE_REVISION +# SOURCE_PROJECTOR_FILES — optional newline-delimited repo-relative mmproj GGUFs +# SOURCE_PIPELINE_TAG — source model pipeline tag for the published model card # MESH_LLM_REF — git ref to build from (default: main) # CATALOG_CREATE_PR — "true" to open a PR for catalog updates (non-org members) +# PACKAGE_EXPERIMENTAL — "true" to label the public package as not runtime-certified # HF_TOKEN — injected as a secret by HF Jobs # # Volumes: @@ -287,6 +290,34 @@ else WRITE_PACKAGE_IDENTITY_ARGS=() echo " Source mount: not available; falling back to Hugging Face cache download" fi +WRITE_PACKAGE_PROJECTOR_ARGS=() +while IFS= read -r PROJECTOR_FILE; do + if [ -z "$PROJECTOR_FILE" ]; then + continue + fi + MOUNTED_PROJECTOR_PATH="/source/${PROJECTOR_FILE}" + if [ -f "$MOUNTED_PROJECTOR_PATH" ]; then + PROJECTOR_PATH="$MOUNTED_PROJECTOR_PATH" + else + echo " Projector mount missing; downloading ${PROJECTOR_FILE} at ${SOURCE_REVISION}" + PROJECTOR_PATH="$("$VENV_DIR/bin/python3" - "$PROJECTOR_FILE" <<'PYTHON' +from huggingface_hub import hf_hub_download +import os +import sys + +print(hf_hub_download( + repo_id=os.environ["SOURCE_REPO"], + filename=sys.argv[1], + revision=os.environ["SOURCE_REVISION"], + cache_dir=os.environ["HF_HUB_CACHE"], + token=os.environ.get("HF_TOKEN"), +)) +PYTHON +)" + fi + echo " Projector: $PROJECTOR_PATH" + WRITE_PACKAGE_PROJECTOR_ARGS+=(--projector "$PROJECTOR_PATH") +done <<< "${SOURCE_PROJECTOR_FILES:-}" echo " Hugging Face cache: $HF_HUB_CACHE" echo " Package workspace: $PACKAGE_DIR" echo " Temporary workspace: $TMPDIR" @@ -309,6 +340,7 @@ set +e time "$SLICER" write-package "$WRITE_PACKAGE_INPUT" \ --out-dir "$PACKAGE_DIR" \ --after-artifact-command "$ARTIFACT_UPLOAD_HOOK" \ + "${WRITE_PACKAGE_PROJECTOR_ARGS[@]}" \ "${WRITE_PACKAGE_IDENTITY_ARGS[@]}" WRITE_PACKAGE_STATUS=$? set -e @@ -429,6 +461,13 @@ package_entry = { "type": "layer-package", "repo": target_repo, "layer_count": layer_count, + "source_revision": source_revision, +} + +source_entry = { + "repo": source_repo, + "file": source_file, + "revision": source_revision, } # Handle both dict-style and list-style variants @@ -439,14 +478,11 @@ if isinstance(variants, dict): packages = variants[variant_name].get("packages", []) packages = [p for p in packages if p.get("repo") != target_repo] packages.append(package_entry) + variants[variant_name]["source"] = source_entry variants[variant_name]["packages"] = packages else: variants[variant_name] = { - "source": { - "repo": source_repo, - "file": source_file, - "revision": source_revision, - }, + "source": source_entry, "curated": { "name": variant_name, "size": f"{layer_count} layers", @@ -466,14 +502,11 @@ else: packages = existing_variant.get("packages", []) packages = [p for p in packages if p.get("repo") != target_repo] packages.append(package_entry) + existing_variant["source"] = source_entry existing_variant["packages"] = packages else: variants.append({ - "source": { - "repo": source_repo, - "file": source_file, - "revision": source_revision, - }, + "source": source_entry, "curated": { "name": variant_name, "size": f"{layer_count} layers", @@ -522,6 +555,21 @@ source_revision = os.environ.get("SOURCE_REVISION", "main") target_repo = os.environ["TARGET_REPO"] model_id = os.environ.get("MODEL_ID", manifest.get("model_id", target_repo)) mesh_llm_ref = os.environ.get("MESH_LLM_REF", "main") +source_pipeline_tag = os.environ.get("SOURCE_PIPELINE_TAG", "text-generation").strip() +if not source_pipeline_tag: + source_pipeline_tag = "text-generation" +experimental = os.environ.get("PACKAGE_EXPERIMENTAL", "false").lower() == "true" +experimental_tag = "- experimental\n" if experimental else "" +experimental_warning = ( + "> [!WARNING]\n" + "> **Experimental package:** artifact integrity may be validated, but runtime, " + "split-correctness, and multimodal certification are still pending. This package " + "is not discoverable through `meshllm/catalog@main` until its Hugging Face catalog " + "PR is reviewed and merged.\n\n" + if experimental + else "" +) +api = HfApi(token=os.environ["HF_TOKEN"]) def sha256(path: Path) -> str: digest = hashlib.sha256() @@ -555,6 +603,38 @@ def code(value) -> str: def yaml_quote(value: str) -> str: return json.dumps(value) +def card_value(info, key: str): + card_data = getattr(info, "card_data", None) + if card_data is None: + return None + if isinstance(card_data, dict): + return card_data.get(key) + return getattr(card_data, key, None) + +def first_model_id(value): + if isinstance(value, list): + value = value[0] if value else None + if isinstance(value, dict): + return value.get("id") or value.get("modelId") + return value if isinstance(value, str) and value else None + +def resolve_upstream_license(): + try: + source_info = api.model_info(source_repo, revision=source_revision) + source_license = card_value(source_info, "license") + if source_license: + return str(source_license), source_repo + + base_repo = first_model_id(card_value(source_info, "base_model")) + if base_repo: + base_info = api.model_info(base_repo) + base_license = card_value(base_info, "license") + if base_license: + return str(base_license), base_repo + except Exception as error: + print(f" WARNING: could not resolve upstream license metadata: {error}") + return None, None + def infer_model_family(name: str) -> str: lowered = name.lower() for family in ["Qwen3", "Qwen2.5", "DeepSeek", "Kimi", "Gemma", "GLM", "Llama"]: @@ -601,6 +681,10 @@ activation_width = manifest.get("activation_width") or "not recorded" skippy_abi = manifest.get("skippy_abi_version") or "not recorded" source_sha = source_model.get("sha256") or "not recorded" canonical_ref = source_model.get("canonical_ref") or f"{source_repo}@{source_revision}/{source_file}" +upstream_license, license_source_repo = resolve_upstream_license() +license_frontmatter = ( + f"license: {yaml_quote(upstream_license)}\n" if upstream_license else "" +) file_rows = [ ("Manifest", "model-package.json", "Package schema, source identity, checksums", manifest_hash), @@ -647,12 +731,18 @@ rows = [ ("Source file", code(source_file)), ("Package repo", link(target_repo, f"https://huggingface.co/{target_repo}")), ] +if upstream_license and license_source_repo: + rows.append(( + "License", + f"{code(upstream_license)} from " + f"{link(license_source_repo, f'https://huggingface.co/{license_source_repo}')}", + )) readme = f"""--- library_name: mesh-llm -base_model: +{license_frontmatter}base_model: - {yaml_quote(source_repo)} -pipeline_tag: text-generation +pipeline_tag: {yaml_quote(source_pipeline_tag)} tags: - gguf - mesh-llm @@ -661,7 +751,7 @@ tags: - distributed-inference - local-inference - openai-compatible ---- +{experimental_tag}--- -GGUF layer package for running **{display_name}** across a local Mesh LLM cluster. +{experimental_warning}GGUF layer package for running **{display_name}** across a local Mesh LLM cluster. This package is derived from [{source_repo}](https://huggingface.co/{source_repo}) and keeps the original GGUF distribution split into per-layer artifacts for distributed inference. @@ -781,7 +871,6 @@ skippy-model-package write-package "{source_path}" --out-dir "{package_dir}" Path("/tmp/README.md").write_text(readme) -api = HfApi(token=os.environ["HF_TOKEN"]) api.upload_file( path_or_fileobj="/tmp/README.md", path_in_repo="README.md", diff --git a/crates/skippy-coordinator/src/topology.rs b/crates/skippy-coordinator/src/topology.rs index 8f311861d..cd49cffc7 100644 --- a/crates/skippy-coordinator/src/topology.rs +++ b/crates/skippy-coordinator/src/topology.rs @@ -113,6 +113,20 @@ pub enum TopologyPlanError { } pub fn plan_topology(input: &TopologyPlanningInput) -> Result { + plan_topology_with_required_stage0(input, None) +} + +pub fn plan_topology_with_stage0( + input: &TopologyPlanningInput, + stage0_node_id: &str, +) -> Result { + plan_topology_with_required_stage0(input, Some(stage0_node_id)) +} + +fn plan_topology_with_required_stage0( + input: &TopologyPlanningInput, + required_stage0_node_id: Option<&str>, +) -> Result { validate_input(input)?; let minimum_context = minimum_valid_context(input.native_context_length); @@ -137,6 +151,9 @@ pub fn plan_topology(input: &TopologyPlanningInput) -> Result, +) -> bool { + required_stage0_node_id.is_none_or(|required| { + candidate + .plan + .stages + .first() + .is_some_and(|stage| stage.node_id == required) + }) +} + fn candidate_better_for_same_shape(candidate: &CandidatePlan, current: &CandidatePlan) -> bool { let candidate_estimate = candidate .plan diff --git a/crates/skippy-correctness/src/cli.rs b/crates/skippy-correctness/src/cli.rs index b18a57474..63291938f 100644 --- a/crates/skippy-correctness/src/cli.rs +++ b/crates/skippy-correctness/src/cli.rs @@ -19,6 +19,7 @@ pub enum CommandKind { StateHandoff(StateHandoffArgs), NativeMtpOpenAiAb(Box), GlmDsaStage0Trace(Box), + StageFaParity(StageFaParityArgs), } #[derive(Args, Clone)] @@ -324,3 +325,27 @@ pub enum StatePayloadKind { RecurrentOnly, KvRecurrent, } + +#[derive(Args)] +pub struct StageFaParityArgs { + #[arg(long)] + pub model: PathBuf, + #[arg(long, default_value = "unsloth/inkling-GGUF:UD-Q2_K_XL")] + pub model_id: String, + #[arg(long, default_value_t = 0)] + pub layer_start: u32, + #[arg(long)] + pub layer_end: u32, + #[arg(long, default_value_t = 2048)] + pub ctx_size: u32, + #[arg(long, default_value_t = 99)] + pub n_gpu_layers: i32, + #[arg(long, default_value = "Hello")] + pub prompt: String, + #[arg(long, default_value_t = 5e-3)] + pub max_abs: f32, + #[arg(long)] + pub enabled_output: Option, + #[arg(long)] + pub disabled_output: Option, +} diff --git a/crates/skippy-correctness/src/main.rs b/crates/skippy-correctness/src/main.rs index d1fd813ae..269229389 100644 --- a/crates/skippy-correctness/src/main.rs +++ b/crates/skippy-correctness/src/main.rs @@ -12,7 +12,7 @@ use crate::{ cli::{Cli, CommandKind}, glm_dsa_trace::glm_dsa_stage0_trace, native_mtp_openai::native_mtp_openai_ab, - runner::{chain, dtype_matrix, single_step, split_scan, state_handoff}, + runner::{chain, dtype_matrix, single_step, split_scan, stage_fa_parity, state_handoff}, }; fn prepare_model_download_directories() { @@ -43,5 +43,6 @@ fn main() -> Result<()> { CommandKind::StateHandoff(args) => state_handoff(args), CommandKind::NativeMtpOpenAiAb(args) => native_mtp_openai_ab(*args), CommandKind::GlmDsaStage0Trace(args) => glm_dsa_stage0_trace(*args), + CommandKind::StageFaParity(args) => stage_fa_parity(args), } } diff --git a/crates/skippy-correctness/src/runner/mod.rs b/crates/skippy-correctness/src/runner/mod.rs index 9ec2ef3a7..ed3a8301b 100644 --- a/crates/skippy-correctness/src/runner/mod.rs +++ b/crates/skippy-correctness/src/runner/mod.rs @@ -1,9 +1,12 @@ pub(crate) mod native_mtp; +mod prediction_return; mod single_step; mod split_chain; mod stage_execution; +mod stage_fa_parity; mod state_handoff; pub use single_step::single_step; pub use split_chain::{chain, dtype_matrix, split_scan}; +pub use stage_fa_parity::stage_fa_parity; pub use state_handoff::state_handoff; diff --git a/crates/skippy-correctness/src/runner/prediction_return.rs b/crates/skippy-correctness/src/runner/prediction_return.rs new file mode 100644 index 000000000..c58f82236 --- /dev/null +++ b/crates/skippy-correctness/src/runner/prediction_return.rs @@ -0,0 +1,251 @@ +use std::{ + env, io, + io::Write, + net::{Shutdown, SocketAddr, TcpListener, TcpStream}, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + mpsc, + }, + thread::{self, JoinHandle}, + time::Duration, +}; + +use anyhow::{Context, Result, anyhow, bail}; +use skippy_protocol::binary::{ + READY_MAGIC, StageReply, WireMessageKind, read_stage_message, recv_ready, recv_reply, + send_ready, +}; + +const CLIENT_READY_HELLO_ENV: &str = "SKIPPY_STAGE_CLIENT_READY_HELLO"; +const CLIENT_READY_HELLO_PEEK_TIMEOUT: Duration = Duration::from_millis(500); + +pub(super) struct PredictionReturnListener { + bind_addr: SocketAddr, + receiver: mpsc::Receiver>, + shutdown: Arc, + connection: Arc>>, + thread: Option>, +} + +impl PredictionReturnListener { + pub(super) fn start() -> Result { + let listener = TcpListener::bind("127.0.0.1:0") + .context("bind correctness prediction return listener")?; + let bind_addr = listener + .local_addr() + .context("read correctness prediction return listener address")?; + listener + .set_nonblocking(true) + .context("set correctness prediction return listener nonblocking")?; + let shutdown = Arc::new(AtomicBool::new(false)); + let thread_shutdown = shutdown.clone(); + let connection = Arc::new(Mutex::new(None)); + let thread_connection = connection.clone(); + let (sender, receiver) = mpsc::channel(); + let thread = thread::spawn(move || { + let result = + accept_prediction_return(listener, &thread_shutdown, &thread_connection, &sender); + if let Err(error) = result { + let _ = sender.send(Err(format!("{error:#}"))); + } + }); + Ok(Self { + bind_addr, + receiver, + shutdown, + connection, + thread: Some(thread), + }) + } + + pub(super) fn endpoint(&self) -> String { + format!("tcp://{}", self.bind_addr) + } + + pub(super) fn receive(&self, timeout: Duration) -> Result { + match self.receiver.recv_timeout(timeout) { + Ok(Ok(reply)) => Ok(reply), + Ok(Err(error)) => Err(anyhow!(error)), + Err(mpsc::RecvTimeoutError::Timeout) => { + bail!("timed out waiting for direct prediction return") + } + Err(mpsc::RecvTimeoutError::Disconnected) => { + bail!("direct prediction return listener disconnected") + } + } + } +} + +impl Drop for PredictionReturnListener { + fn drop(&mut self) { + self.shutdown.store(true, Ordering::SeqCst); + if let Ok(connection) = self.connection.lock() + && let Some(stream) = connection.as_ref() + { + let _ = stream.shutdown(Shutdown::Both); + } + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} + +fn accept_prediction_return( + listener: TcpListener, + shutdown: &AtomicBool, + connection: &Mutex>, + sender: &mpsc::Sender>, +) -> Result<()> { + let mut stream = loop { + match listener.accept() { + Ok((stream, _)) => break stream, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + if shutdown.load(Ordering::SeqCst) { + return Ok(()); + } + thread::sleep(Duration::from_millis(10)); + } + Err(error) if error.kind() == io::ErrorKind::Interrupted => {} + Err(error) => return Err(error).context("accept direct prediction return"), + } + }; + stream + .set_nonblocking(false) + .context("set direct prediction return stream blocking")?; + let shutdown_stream = stream + .try_clone() + .context("clone direct prediction return stream for shutdown")?; + { + let mut connection = connection + .lock() + .map_err(|_| anyhow!("direct prediction return connection lock poisoned"))?; + if shutdown.load(Ordering::SeqCst) { + return Ok(()); + } + *connection = Some(shutdown_stream); + } + consume_optional_client_ready_hello(&mut stream)?; + send_ready(&mut stream).context("send direct prediction return ready")?; + stream.flush().ok(); + let open = + read_stage_message(&mut stream, 0).context("read direct prediction return open message")?; + if open.kind != WireMessageKind::PredictionReturnOpen { + bail!("expected direct prediction return open message"); + } + loop { + match recv_reply(&mut stream) { + Ok(reply) => { + if sender.send(Ok(reply)).is_err() { + return Ok(()); + } + } + Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(()), + Err(error) => return Err(error).context("read direct prediction return reply"), + } + } +} + +fn consume_optional_client_ready_hello(stream: &mut TcpStream) -> Result<()> { + if !client_ready_hello_enabled() { + return Ok(()); + } + let previous_timeout = stream + .read_timeout() + .context("read direct prediction return timeout")?; + stream + .set_read_timeout(Some(CLIENT_READY_HELLO_PEEK_TIMEOUT)) + .context("set direct prediction return hello timeout")?; + let mut bytes = [0_u8; 4]; + let peek_result = stream.peek(&mut bytes); + stream + .set_read_timeout(previous_timeout) + .context("restore direct prediction return timeout")?; + + match peek_result { + Ok(4) if i32::from_le_bytes(bytes) == READY_MAGIC => { + recv_ready(stream).context("consume direct prediction return client ready hello")?; + } + Ok(_) => {} + Err(error) + if matches!( + error.kind(), + io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut + ) => {} + Err(error) => { + return Err(error).context("peek direct prediction return client ready hello"); + } + } + Ok(()) +} + +fn client_ready_hello_enabled() -> bool { + env::var(CLIENT_READY_HELLO_ENV) + .map(|value| matches!(value.as_str(), "1" | "true" | "TRUE" | "yes" | "on")) + .unwrap_or(false) +} + +#[cfg(test)] +mod tests { + use super::*; + use skippy_protocol::binary::{ + StageStateHeader, StageWireMessage, WireActivationDType, send_reply_predicted, + write_stage_message, + }; + + #[test] + fn receives_prediction_over_direct_return_endpoint() { + let listener = PredictionReturnListener::start().unwrap(); + let endpoint = listener.endpoint(); + let address = endpoint.strip_prefix("tcp://").unwrap().to_string(); + let client = thread::spawn(move || { + let mut stream = TcpStream::connect(address).unwrap(); + if client_ready_hello_enabled() { + send_ready(&mut stream).unwrap(); + } + recv_ready(&mut stream).unwrap(); + let kind = WireMessageKind::PredictionReturnOpen; + let open = StageWireMessage { + kind, + pos_start: 0, + token_count: 0, + state: StageStateHeader::new(kind, WireActivationDType::F32), + request_id: 11, + session_id: 13, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + }; + write_stage_message(&mut stream, &open, WireActivationDType::F32).unwrap(); + send_reply_predicted(&mut stream, 674).unwrap(); + }); + + let reply = listener.receive(Duration::from_secs(1)).unwrap(); + + assert_eq!(reply.predicted, 674); + client.join().unwrap(); + } + + #[test] + fn drop_stops_when_connected_peer_stalls() { + let listener = PredictionReturnListener::start().unwrap(); + let address = listener + .endpoint() + .strip_prefix("tcp://") + .unwrap() + .to_string(); + let mut stream = TcpStream::connect(address).unwrap(); + if client_ready_hello_enabled() { + send_ready(&mut stream).unwrap(); + } + recv_ready(&mut stream).unwrap(); + + let started = std::time::Instant::now(); + drop(listener); + + assert!(started.elapsed() < Duration::from_secs(1)); + } +} diff --git a/crates/skippy-correctness/src/runner/split_chain.rs b/crates/skippy-correctness/src/runner/split_chain.rs index c10667b9f..52f4b9cbd 100644 --- a/crates/skippy-correctness/src/runner/split_chain.rs +++ b/crates/skippy-correctness/src/runner/split_chain.rs @@ -1,9 +1,15 @@ -use std::{fs, net::SocketAddr, path::PathBuf, process::Command, time::Instant}; +use std::{ + fs, + net::SocketAddr, + path::PathBuf, + process::Command, + time::{Duration, Instant}, +}; use anyhow::{Context, Result, bail}; use model_artifact::ModelIdentity; use serde_json::json; -use skippy_protocol::binary::{StageWireMessage, WireReplyKind, recv_reply, write_stage_message}; +use skippy_protocol::binary::{StageWireMessage, WireReplyKind, write_stage_message}; use skippy_runtime::{GGML_TYPE_F16, RuntimeConfig, StageModel}; use crate::{ @@ -24,6 +30,7 @@ use super::{ native_mtp_satisfies_requirement, native_mtp_sideband_report, native_mtp_verification_report, native_mtp_verification_satisfies_requirement, }, + prediction_return::PredictionReturnListener, single_step::{SingleStepCase, run_full_model_decode, run_single_step_with_baseline}, stage_execution::{ BinaryDecodeMessageArgs, CorrectnessTopologyStage, FullModelResult, PackageStageSpec, @@ -347,6 +354,8 @@ fn run_binary_chain(args: BinaryChainConfig) -> Result { bail!("stage 0 produced an empty activation frame"); } let activation_width = activation_width(&boundary)?; + let prediction_return = PredictionReturnListener::start()?; + let stage0_endpoint = prediction_return.endpoint(); let run_id = generate_run_id(); let model_id = args.model_identity.model_id.clone(); @@ -407,7 +416,7 @@ fn run_binary_chain(args: BinaryChainConfig) -> Result { "upstream": { "stage_id": "stage-0", "stage_index": 0, - "endpoint": "driver" + "endpoint": stage0_endpoint }, "downstream": { "stage_id": "stage-2", @@ -422,7 +431,7 @@ fn run_binary_chain(args: BinaryChainConfig) -> Result { CorrectnessTopologyStage { stage_id: "stage-0", stage_index: 0, - endpoint: "driver".to_string(), + endpoint: stage0_endpoint, layer_start: 0, layer_end: args.split_layer_1, load_mode: protocol_load_mode(args.stage_load_mode), @@ -521,7 +530,9 @@ fn run_binary_chain(args: BinaryChainConfig) -> Result { session_id, })?; write_stage_message(&mut stream, &message, wire_dtype).context("send binary chain decode")?; - let reply = recv_reply(&mut stream).context("receive binary chain prediction reply")?; + let reply = prediction_return + .receive(Duration::from_secs(args.startup_timeout_secs)) + .context("receive binary chain direct prediction reply")?; ensure_reply_kind(&reply, WireReplyKind::PredictedToken)?; let native_mtp = native_mtp_sideband_report(&reply); let (second_predicted_token, native_mtp_verification_compute_us) = @@ -542,8 +553,9 @@ fn run_binary_chain(args: BinaryChainConfig) -> Result { })?; write_stage_message(&mut stream, &second_message, wire_dtype) .context("send second binary chain decode")?; - let second_reply = - recv_reply(&mut stream).context("receive second binary chain prediction reply")?; + let second_reply = prediction_return + .receive(Duration::from_secs(args.startup_timeout_secs)) + .context("receive second binary chain direct prediction reply")?; ensure_reply_kind(&second_reply, WireReplyKind::PredictedToken)?; ( Some(second_reply.predicted), diff --git a/crates/skippy-correctness/src/runner/stage_fa_parity.rs b/crates/skippy-correctness/src/runner/stage_fa_parity.rs new file mode 100644 index 000000000..8205a4948 --- /dev/null +++ b/crates/skippy-correctness/src/runner/stage_fa_parity.rs @@ -0,0 +1,134 @@ +use anyhow::{Context, Result, bail}; +use skippy_runtime::{ + FlashAttentionType, GGML_TYPE_F16, RuntimeConfig, RuntimeLoadMode, StageModel, + package::{PackageStageRequest, select_layer_package_parts}, +}; + +use crate::cli::StageFaParityArgs; + +pub fn stage_fa_parity(args: StageFaParityArgs) -> Result<()> { + if args.layer_start >= args.layer_end { + bail!( + "layer_start ({}) must be less than layer_end ({})", + args.layer_start, + args.layer_end + ); + } + let enabled = decode_boundary(&args, FlashAttentionType::Enabled)?; + let disabled = decode_boundary(&args, FlashAttentionType::Disabled)?; + if enabled.desc != disabled.desc { + bail!( + "activation descriptors differ: enabled={:?} disabled={:?}", + enabled.desc, + disabled.desc + ); + } + if let Some(path) = args.enabled_output.as_deref() { + std::fs::write(path, &enabled.payload) + .with_context(|| format!("write enabled activation {}", path.display()))?; + } + if let Some(path) = args.disabled_output.as_deref() { + std::fs::write(path, &disabled.payload) + .with_context(|| format!("write disabled activation {}", path.display()))?; + } + let enabled_values = payload_f32(&enabled.payload)?; + let disabled_values = payload_f32(&disabled.payload)?; + if enabled_values.is_empty() || enabled_values.len() != disabled_values.len() { + bail!( + "activation payload length mismatch or empty: enabled={} disabled={}", + enabled_values.len(), + disabled_values.len() + ); + } + let mut max_abs = 0.0_f32; + let mut sum_sq = 0.0_f64; + for (lhs, rhs) in enabled_values.iter().zip(&disabled_values) { + let delta = (lhs - rhs).abs(); + max_abs = max_abs.max(delta); + sum_sq += f64::from(delta) * f64::from(delta); + } + let rms = (sum_sq / enabled_values.len() as f64).sqrt(); + println!( + "stage_fa_parity elements={} max_abs={max_abs:.8e} rms={rms:.8e}", + enabled_values.len() + ); + if max_abs > args.max_abs { + bail!( + "stage FA parity max_abs {max_abs:.8e} exceeds tolerance {:.8e}", + args.max_abs + ); + } + Ok(()) +} + +fn decode_boundary( + args: &StageFaParityArgs, + flash_attn_type: FlashAttentionType, +) -> Result { + let config = RuntimeConfig { + stage_index: 0, + layer_start: args.layer_start, + layer_end: args.layer_end, + ctx_size: args.ctx_size, + lane_count: 1, + n_batch: None, + n_ubatch: None, + n_threads: None, + n_threads_batch: None, + n_gpu_layers: args.n_gpu_layers, + mmap: None, + mlock: false, + selected_backend_device: None, + cache_type_k: GGML_TYPE_F16, + cache_type_v: GGML_TYPE_F16, + flash_attn_type, + load_mode: RuntimeLoadMode::LayerPackage, + projector_path: None, + include_embeddings: true, + include_output: false, + filter_tensors_on_load: true, + }; + let selection = select_layer_package_parts(&PackageStageRequest { + model_id: args.model_id.clone(), + topology_id: "stage-fa-parity".to_string(), + package_ref: args.model.display().to_string(), + stage_id: "stage-0".to_string(), + layer_start: args.layer_start, + layer_end: args.layer_end, + include_embeddings: true, + include_output: false, + }) + .context("select package parts")?; + let model = StageModel::open_from_parts(&selection.absolute_paths, &config) + .context("open stage model")?; + let tokens = model + .tokenize(&args.prompt, true) + .context("tokenize prompt")?; + if tokens.is_empty() { + bail!("prompt produced no token"); + } + let mut session = model.create_session().context("create stage session")?; + if tokens.len() == 1 { + let (_, frame) = session + .decode_step_frame(tokens[0], None, 0) + .context("decode stage boundary")?; + Ok(frame) + } else { + session + .prefill_chunk_frame(&tokens, None, 0) + .context("prefill stage boundary") + } +} + +fn payload_f32(payload: &[u8]) -> Result> { + if !payload.len().is_multiple_of(4) { + bail!( + "activation payload is not f32-aligned: {} bytes", + payload.len() + ); + } + Ok(payload + .chunks_exact(4) + .map(|bytes| f32::from_le_bytes(bytes.try_into().unwrap())) + .collect()) +} diff --git a/crates/skippy-ffi/src/lib.rs b/crates/skippy-ffi/src/lib.rs index 99ddd56bb..a76464b9a 100644 --- a/crates/skippy-ffi/src/lib.rs +++ b/crates/skippy-ffi/src/lib.rs @@ -1,10 +1,11 @@ pub const ABI_VERSION_MAJOR: u32 = 0; pub const ABI_VERSION_MINOR: u32 = 1; -pub const ABI_VERSION_PATCH: u32 = 33; +pub const ABI_VERSION_PATCH: u32 = 35; pub const FEATURE_BACKEND_DEVICES: u64 = 1 << 23; pub const FEATURE_RUNTIME_EVENTS: u64 = 1 << 24; pub const FEATURE_NATIVE_MTP_N1: u64 = 1 << 25; pub const FEATURE_NGRAM_CACHE_DRAFT: u64 = 1 << 26; +pub const FEATURE_INKLING_MTP_MM: u64 = 1 << 27; #[cfg(feature = "dynamic-runtime")] mod dynamic_library; @@ -18,19 +19,20 @@ pub struct AbiVersion { } /// Whether a native runtime reporting `version` can back this binary's ABI -/// bindings. Required symbol signatures may change between patches (for -/// example `skippy_apply_chat_template_json` gained an argument in 0.1.28), -/// so older runtimes must be rejected at load time. +/// bindings. Required symbol signatures and by-value struct layouts may change +/// between patches, so the loader requires an exact ABI match. pub const fn runtime_abi_supported(version: AbiVersion) -> bool { version.major == ABI_VERSION_MAJOR && version.minor == ABI_VERSION_MINOR - && version.patch >= ABI_VERSION_PATCH + && version.patch == ABI_VERSION_PATCH } use std::ffi::{c_char, c_int, c_void}; pub type LlamaLogCallback = Option; +pub type MtmdProgressCallback = + Option bool>; pub type SkippyRuntimeEventCallback = Option; @@ -349,6 +351,9 @@ pub struct MtmdContextParams { pub image_max_tokens: c_int, pub cb_eval: *mut c_void, pub cb_eval_user_data: *mut c_void, + pub batch_max_tokens: c_int, + pub progress_callback: MtmdProgressCallback, + pub progress_callback_user_data: *mut c_void, } #[repr(C)] @@ -377,6 +382,8 @@ pub struct ActivationDesc { pub flags: u64, } +pub const ACTIVATION_FLAG_INKLING_MTP_EMBD: u64 = 1 << 2; + #[repr(C)] #[derive(Debug, Clone, Copy)] pub struct LogitBias { @@ -422,6 +429,7 @@ pub enum LlamaFileType { MostlyMxfp4Moe = 38, MostlyNvfp4 = 39, MostlyQ1_0 = 40, + MostlyQ2_0 = 41, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -461,7 +469,8 @@ pub enum GgmlType { Mxfp4 = 39, Nvfp4 = 40, Q1_0 = 41, - Count = 42, + Q2_0 = 42, + Count = 43, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -847,6 +856,7 @@ mod dynamic { skippy_session_last_token_signal(session: *mut Session, out_signal: *mut TokenSignal, out_error: *mut *mut Error) -> Status; skippy_session_signal_window(session: *mut Session, window_tokens: u32, out_window: *mut GenerationSignalWindow, out_error: *mut *mut Error) -> Status; skippy_trim_session(session: *mut Session, token_count: u64, out_error: *mut *mut Error) -> Status; + skippy_retire_verify_checkpoint(session: *mut Session, token_start: u64, token_count: u64, out_error: *mut *mut Error) -> Status; skippy_export_state(session: *mut Session, layer_start: i32, layer_end: i32, output: *mut c_void, output_capacity: usize, out_bytes: *mut usize, out_error: *mut *mut Error) -> Status; skippy_import_state(session: *mut Session, layer_start: i32, layer_end: i32, input: *const c_void, input_bytes: usize, out_error: *mut *mut Error) -> Status; skippy_export_full_state(session: *mut Session, layer_start: i32, layer_end: i32, output: *mut c_void, output_capacity: usize, out_bytes: *mut usize, out_error: *mut *mut Error) -> Status; @@ -887,6 +897,7 @@ mod dynamic { mtmd_decode_use_mrope(ctx: *const MtmdContext) -> bool; mtmd_input_chunk_get_type(chunk: *const Opaque) -> MtmdInputChunkType; mtmd_input_chunk_get_n_tokens(chunk: *const Opaque) -> usize; + mtmd_input_chunk_get_tokens_text(chunk: *const Opaque, out_count: *mut usize) -> *const i32; mtmd_input_chunk_get_tokens_image(chunk: *const Opaque) -> *const Opaque; mtmd_helper_image_get_decoder_pos(image: *const Opaque, pos_0: i32, out_pos: *mut MtmdDecoderPos); mtmd_helper_eval_chunks(ctx: *mut MtmdContext, lctx: *mut Opaque, chunks: *const MtmdInputChunks, n_past: i32, seq_id: i32, n_batch: i32, logits_last: bool, new_n_past: *mut i32) -> c_int; @@ -1172,9 +1183,36 @@ pub use dynamic::*; /// Returns the skippy ABI feature bitmask. /// Requires the native runtime to be loaded first (checked by caller). pub fn skippy_abi_features() -> u64 { - let fns = dynamic::skippy_abi_features_optional() - .expect("skippy_abi_features not available in loaded runtime"); - unsafe { fns() } + try_abi_features().expect("skippy_abi_features not available in loaded runtime") +} + +/// Returns the Skippy ABI feature bitmask when the loaded dynamic runtime +/// exports feature probing. +#[cfg(feature = "dynamic-runtime")] +pub fn try_abi_features() -> Option { + dynamic::skippy_abi_features_optional().map(|features| unsafe { features() }) +} + +/// Returns the active Skippy ABI feature bitmask through a safe Rust wrapper. +#[cfg(feature = "dynamic-runtime")] +pub fn abi_features() -> u64 { + skippy_abi_features() +} + +/// Returns the statically linked Skippy ABI feature bitmask. +#[cfg(not(feature = "dynamic-runtime"))] +pub fn try_abi_features() -> Option { + // SAFETY: the statically linked ABI exposes this nullary query with no + // caller-owned pointers or lifetime requirements. + Some(unsafe { skippy_abi_features() }) +} + +/// Returns the statically linked Skippy ABI feature bitmask. +#[cfg(not(feature = "dynamic-runtime"))] +pub fn abi_features() -> u64 { + // SAFETY: the statically linked ABI exposes this nullary query with no + // caller-owned pointers or lifetime requirements. + unsafe { skippy_abi_features() } } #[cfg(not(feature = "dynamic-runtime"))] @@ -1529,6 +1567,13 @@ unsafe extern "C" { out_error: *mut *mut Error, ) -> Status; + pub fn skippy_retire_verify_checkpoint( + session: *mut Session, + token_start: u64, + token_count: u64, + out_error: *mut *mut Error, + ) -> Status; + pub fn skippy_export_state( session: *mut Session, layer_start: i32, @@ -1795,6 +1840,11 @@ unsafe extern "C" { pub fn mtmd_input_chunk_get_n_tokens(chunk: *const Opaque) -> usize; + pub fn mtmd_input_chunk_get_tokens_text( + chunk: *const Opaque, + out_count: *mut usize, + ) -> *const i32; + pub fn mtmd_input_chunk_get_tokens_image(chunk: *const Opaque) -> *const Opaque; pub fn mtmd_helper_image_get_decoder_pos( @@ -1841,6 +1891,7 @@ pub type Opaque = c_void; #[cfg(test)] mod tests { use super::*; + use std::mem::{offset_of, size_of}; const fn version(major: u32, minor: u32, patch: u32) -> AbiVersion { AbiVersion { @@ -1851,21 +1902,21 @@ mod tests { } #[test] - fn accepts_current_and_newer_patch_runtimes() { + fn accepts_current_patch_runtime() { assert!(runtime_abi_supported(version( ABI_VERSION_MAJOR, ABI_VERSION_MINOR, ABI_VERSION_PATCH, ))); - assert!(runtime_abi_supported(version( + } + + #[test] + fn rejects_other_patch_runtimes() { + assert!(!runtime_abi_supported(version( ABI_VERSION_MAJOR, ABI_VERSION_MINOR, ABI_VERSION_PATCH + 1, ))); - } - - #[test] - fn rejects_older_patch_runtimes() { assert!(!runtime_abi_supported(version( ABI_VERSION_MAJOR, ABI_VERSION_MINOR, @@ -1886,4 +1937,26 @@ mod tests { ABI_VERSION_PATCH, ))); } + + #[test] + #[cfg(target_pointer_width = "64")] + fn mtmd_context_params_matches_native_layout() { + assert_eq!(size_of::(), 80); + assert_eq!(offset_of!(MtmdContextParams, batch_max_tokens), 56); + assert_eq!(offset_of!(MtmdContextParams, progress_callback), 64); + assert_eq!( + offset_of!(MtmdContextParams, progress_callback_user_data), + 72 + ); + } + + #[test] + #[cfg(not(feature = "dynamic-runtime"))] + fn native_mtmd_defaults_cross_the_ffi_boundary() { + let params = unsafe { mtmd_context_params_default() }; + + assert_eq!(params.batch_max_tokens, 1024); + assert!(params.progress_callback.is_none()); + assert!(params.progress_callback_user_data.is_null()); + } } diff --git a/crates/skippy-model-package/Cargo.toml b/crates/skippy-model-package/Cargo.toml index fe5c4eb59..feae00841 100644 --- a/crates/skippy-model-package/Cargo.toml +++ b/crates/skippy-model-package/Cargo.toml @@ -8,6 +8,7 @@ version.workspace = true anyhow.workspace = true clap.workspace = true skippy-ffi = { path = "../skippy-ffi", default-features = false } +skippy-protocol = { path = "../skippy-protocol" } skippy-runtime = { path = "../skippy-runtime" } model-artifact = { path = "../model-artifact" } model-hf = { path = "../model-hf" } diff --git a/crates/skippy-model-package/src/package.rs b/crates/skippy-model-package/src/package.rs index 49f954a1f..5b67c5cd3 100644 --- a/crates/skippy-model-package/src/package.rs +++ b/crates/skippy-model-package/src/package.rs @@ -124,6 +124,8 @@ pub(crate) struct PackageWindowPolicy { pub(crate) initial_window: u32, pub(crate) min_window: u32, pub(crate) max_window: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) pipeline_depth: Option, } #[derive(Debug, Deserialize, Serialize)] @@ -663,6 +665,7 @@ pub(crate) fn package_generation(tensors: &[TensorInfo]) -> Option, } #[derive(Debug, Deserialize)] @@ -273,6 +282,8 @@ struct PackageWindowPolicy { initial_window: u32, min_window: u32, max_window: u32, + #[serde(default)] + pipeline_depth: Option, } #[derive(Debug, Deserialize)] @@ -1014,6 +1025,23 @@ fn validate_window_policy( "use positive window sizes", ); } + if window.pipeline_depth.is_some_and(|depth| { + depth == 0 + || usize::try_from(depth) + .map(|depth| depth > MAX_VERIFY_WINDOW_PIPELINE_DEPTH) + .unwrap_or(true) + }) { + report.error( + "invalid_window_policy_pipeline_depth", + format!( + "speculative strategy {name} window_policy.pipeline_depth must be between 1 and {MAX_VERIFY_WINDOW_PIPELINE_DEPTH}" + ), + Some("model-package.json".to_string()), + format!( + "set pipeline_depth to an in-flight verification-window capacity no greater than {MAX_VERIFY_WINDOW_PIPELINE_DEPTH}" + ), + ); + } if window.min_window > window.max_window { report.error( "invalid_window_policy_bounds", @@ -1135,432 +1163,10 @@ fn preflight_window_policy(window: &PackageWindowPolicy) -> PreflightWindowPolic initial_window: window.initial_window, min_window: window.min_window, max_window: window.max_window, + pipeline_depth: window.pipeline_depth, } } -fn collect_artifacts(manifest: &PackageManifest) -> Vec { - let mut artifacts = vec![ - artifact_spec("metadata", None, &manifest.shared.metadata), - artifact_spec("embeddings", None, &manifest.shared.embeddings), - artifact_spec("output", None, &manifest.shared.output), - ]; - artifacts.extend( - manifest - .layers - .iter() - .map(|layer| layer_artifact_spec(layer.layer_index, layer)), - ); - artifacts.extend(manifest.projectors.iter().map(projector_artifact_spec)); - artifacts -} - -fn artifact_spec( - role: &'static str, - layer_index: Option, - artifact: &PackageArtifact, -) -> ArtifactSpec { - ArtifactSpec { - role, - layer_index, - path: artifact.path.clone(), - tensor_count: artifact.tensor_count, - tensor_bytes: artifact.tensor_bytes, - artifact_bytes: artifact.artifact_bytes, - sha256: artifact.sha256.clone(), - } -} - -fn layer_artifact_spec(layer_index: u32, layer: &PackageLayer) -> ArtifactSpec { - ArtifactSpec { - role: "layer", - layer_index: Some(layer_index), - path: layer.path.clone(), - tensor_count: layer.tensor_count, - tensor_bytes: layer.tensor_bytes, - artifact_bytes: layer.artifact_bytes, - sha256: layer.sha256.clone(), - } -} - -fn projector_artifact_spec(projector: &PackageProjector) -> ArtifactSpec { - ArtifactSpec { - role: "projector", - layer_index: None, - path: projector.path.clone(), - tensor_count: projector.tensor_count, - tensor_bytes: projector.tensor_bytes, - artifact_bytes: projector.artifact_bytes, - sha256: projector.sha256.clone(), - } -} - -fn validate_layer_coverage(manifest: &PackageManifest, report: &mut PackagePreflightReport) { - let mut counts = BTreeMap::::new(); - for layer in &manifest.layers { - *counts.entry(layer.layer_index).or_default() += 1; - if layer.layer_index >= manifest.layer_count { - report.error( - "layer_index_out_of_range", - format!( - "package layer index {} exceeds layer_count {}", - layer.layer_index, manifest.layer_count - ), - Some(layer.path.clone()), - "rebuild the package so layer indexes are contiguous and in range", - ); - } - } - for layer_index in 0..manifest.layer_count { - if !counts.contains_key(&layer_index) { - report.error( - "missing_layer", - format!("package manifest is missing layer {layer_index}"), - Some("model-package.json".to_string()), - "rebuild the package so every transformer layer has one artifact", - ); - } - } - for (layer_index, count) in counts { - if count > 1 { - report.error( - "duplicate_layer", - format!("package manifest contains layer {layer_index} {count} times"), - Some("model-package.json".to_string()), - "rebuild the package so each layer appears exactly once", - ); - } - } -} - -fn validate_artifacts( - package: &Path, - artifacts: &[ArtifactSpec], - verify_sha256: bool, - report: &mut PackagePreflightReport, -) { - for artifact in artifacts { - report.artifacts.push(preflight_artifact( - package, - artifact, - verify_sha256, - &mut report.issues, - )); - } -} - -fn preflight_artifact( - package: &Path, - artifact: &ArtifactSpec, - verify_sha256: bool, - issues: &mut Vec, -) -> PreflightArtifact { - let path = match safe_relative_path(&artifact.path) { - Ok(path) => path, - Err(message) => { - push_error( - issues, - "unsafe_artifact_path", - format!( - "package {} artifact path is unsafe: {message}", - artifact.role - ), - Some(artifact.path.clone()), - "rebuild the package so artifact paths stay inside the package directory", - ); - return artifact_output(artifact, false, None, None, None); - } - }; - validate_artifact_manifest(artifact, issues); - let absolute = package.join(&path); - let metadata = match fs::metadata(&absolute) { - Ok(metadata) if metadata.is_file() => metadata, - Ok(_) => { - push_error( - issues, - "artifact_not_file", - format!("package artifact {} is not a file", artifact.path), - Some(artifact.path.clone()), - "replace the artifact path with a regular GGUF file", - ); - return artifact_output(artifact, false, None, None, None); - } - Err(error) => { - push_error( - issues, - "missing_artifact", - format!("package artifact {} is missing: {error}", artifact.path), - Some(artifact.path.clone()), - "download or rebuild the package artifact before starting split serving", - ); - return artifact_output(artifact, false, None, None, None); - } - }; - let actual_len = metadata.len(); - let size_matches = actual_len == artifact.artifact_bytes; - if !size_matches { - push_error( - issues, - "artifact_size_mismatch", - format!( - "package artifact {} has {} bytes, manifest expects {}", - artifact.path, actual_len, artifact.artifact_bytes - ), - Some(artifact.path.clone()), - "redownload or rebuild the package artifact so manifest sizes match", - ); - } - let sha_matches = if verify_sha256 { - Some(validate_artifact_sha(&absolute, artifact, issues)) - } else { - None - }; - artifact_output( - artifact, - true, - Some(actual_len), - Some(size_matches), - sha_matches, - ) -} - -fn validate_artifact_manifest(artifact: &ArtifactSpec, issues: &mut Vec) { - if artifact.artifact_bytes == 0 { - push_error( - issues, - "empty_artifact", - format!("package {} artifact declares zero bytes", artifact.role), - Some(artifact.path.clone()), - "rebuild the package; split artifacts must be non-empty files", - ); - } - if artifact.tensor_count == 0 && artifact.tensor_bytes > 0 { - push_error( - issues, - "invalid_tensor_bytes", - format!( - "package {} artifact declares tensor_bytes without tensors", - artifact.role - ), - Some(artifact.path.clone()), - "rebuild the package manifest so tensor counts and bytes agree", - ); - } - if artifact.tensor_count > 0 && artifact.tensor_bytes == 0 { - push_error( - issues, - "invalid_tensor_bytes", - format!( - "package {} artifact declares tensors but zero tensor_bytes", - artifact.role - ), - Some(artifact.path.clone()), - "rebuild the package manifest so tensor counts and bytes agree", - ); - } - if !is_sha256(&artifact.sha256) { - push_error( - issues, - "invalid_artifact_sha256", - format!( - "package {} artifact sha256 is not a hex digest", - artifact.role - ), - Some(artifact.path.clone()), - "rebuild the package manifest so artifact checksums are valid", - ); - } -} - -fn validate_artifact_sha( - path: &Path, - artifact: &ArtifactSpec, - issues: &mut Vec, -) -> bool { - match file_sha256(path) { - Ok(actual) if actual == artifact.sha256.to_ascii_lowercase() => true, - Ok(actual) => { - push_error( - issues, - "artifact_sha256_mismatch", - format!( - "package artifact {} checksum mismatch: expected {}, got {}", - artifact.path, artifact.sha256, actual - ), - Some(artifact.path.clone()), - "redownload or rebuild the package artifact so checksums match", - ); - false - } - Err(error) => { - push_error( - issues, - "artifact_sha256_unreadable", - format!("cannot hash package artifact {}: {error}", artifact.path), - Some(artifact.path.clone()), - "ensure the artifact is readable before enabling checksum verification", - ); - false - } - } -} - -fn artifact_output( - artifact: &ArtifactSpec, - present: bool, - actual_artifact_bytes: Option, - size_matches_manifest: Option, - sha256_matches_manifest: Option, -) -> PreflightArtifact { - PreflightArtifact { - role: artifact.role.to_string(), - layer_index: artifact.layer_index, - path: artifact.path.clone(), - present, - declared_artifact_bytes: artifact.artifact_bytes, - actual_artifact_bytes, - size_matches_manifest, - sha256_matches_manifest, - } -} - -fn build_stage_reports( - manifest: &PackageManifest, - stages: Option, - report: &mut PackagePreflightReport, -) { - let Some(stage_count) = stages else { - return; - }; - if stage_count == 0 { - report.error( - "invalid_stage_count", - "--stages must be greater than zero", - Some("model-package.json".to_string()), - "choose a positive stage count for split preflight", - ); - return; - } - if stage_count as u32 > manifest.layer_count { - report.error( - "stage_count_exceeds_layer_count", - format!( - "--stages {stage_count} exceeds package layer_count {}", - manifest.layer_count - ), - Some("model-package.json".to_string()), - "use at most one split stage per transformer layer", - ); - return; - } - let artifact_map = stage_artifacts(report); - for (stage_index, (layer_start, layer_end)) in - partition_layers(manifest.layer_count, stage_count) - .into_iter() - .enumerate() - { - report.stages.push(stage_report( - stage_index, - layer_start, - layer_end, - stage_count, - &artifact_map, - )); - } -} - -fn stage_report( - stage_index: usize, - layer_start: u32, - layer_end: u32, - stage_count: usize, - artifact_map: &BTreeMap, -) -> PreflightStage { - let includes_embeddings = stage_index == 0; - let includes_output = stage_index + 1 == stage_count; - let mut parts = vec!["metadata".to_string()]; - if includes_embeddings { - parts.push("embeddings".to_string()); - } - for layer_index in layer_start..layer_end { - parts.push(format!("layer:{layer_index}")); - } - if includes_output { - parts.push("output".to_string()); - } - let artifact_bytes = parts - .iter() - .filter_map(|part| artifact_map.get(part)) - .filter(|artifact| artifact.present) - .map(|artifact| artifact.bytes) - .sum(); - let missing_parts = parts - .iter() - .filter(|part| { - !artifact_map - .get(*part) - .is_some_and(|artifact| artifact.present) - }) - .cloned() - .collect::>(); - PreflightStage { - stage_index, - layer_start, - layer_end, - includes_embeddings, - includes_output, - part_count: parts.len(), - artifact_bytes, - parts, - missing_parts, - } -} - -#[derive(Clone, Copy)] -struct StageArtifact { - present: bool, - bytes: u64, -} - -fn stage_artifacts(report: &PackagePreflightReport) -> BTreeMap { - report - .artifacts - .iter() - .map(|artifact| { - ( - stage_part_key(artifact), - StageArtifact { - present: artifact.present, - bytes: artifact - .actual_artifact_bytes - .unwrap_or(artifact.declared_artifact_bytes), - }, - ) - }) - .collect() -} - -fn stage_part_key(artifact: &PreflightArtifact) -> String { - match (artifact.role.as_str(), artifact.layer_index) { - ("layer", Some(layer)) => format!("layer:{layer}"), - (role, _) => role.to_string(), - } -} - -fn partition_layers(layer_count: u32, stages: usize) -> Vec<(u32, u32)> { - let base = layer_count / stages as u32; - let extra = layer_count % stages as u32; - let mut start = 0; - (0..stages) - .map(|stage_index| { - let width = base + u32::from((stage_index as u32) < extra); - let end = start + width; - let range = (start, end); - start = end; - range - }) - .collect() -} - fn push_error( issues: &mut Vec, code: impl Into, @@ -1735,7 +1341,8 @@ mod tests { "default": "fixed", "initial_window": 1, "min_window": 1, - "max_window": 1 + "max_window": 1, + "pipeline_depth": 2 } } } @@ -1770,6 +1377,7 @@ mod tests { assert_eq!(window_policy.initial_window, 1); assert_eq!(window_policy.min_window, 1); assert_eq!(window_policy.max_window, 1); + assert_eq!(window_policy.pipeline_depth, Some(2)); fs::remove_dir_all(dir).unwrap(); } @@ -2104,6 +1712,96 @@ mod tests { fs::remove_dir_all(dir).unwrap(); } + #[test] + fn preflight_rejects_zero_verify_window_pipeline_depth() { + let dir = unique_test_dir("zero-window-pipeline-depth"); + let package = write_package_fixture(&dir, true); + write_generation_to_manifest( + &package, + serde_json::json!({ + "speculative_decoding": { + "default": "ngram-suffix", + "proposers": { + "suffix": { + "type": "ngram-suffix", + "ngram_min": 5, + "ngram_max": 32, + "max_proposal_tokens": 48, + "history_scope": "request" + } + }, + "strategies": { + "ngram-suffix": { + "type": "ngram-suffix", + "proposer": "suffix", + "window_policy": { + "default": "fixed", + "initial_window": 32, + "min_window": 1, + "max_window": 32, + "pipeline_depth": 0 + } + } + } + } + }), + ); + + let report = preflight_package(&package, &PackagePreflightOptions::default()); + + assert!(!report.valid); + assert_issue(&report, "invalid_window_policy_pipeline_depth"); + fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn preflight_enforces_verify_window_pipeline_depth_maximum() { + for (depth, expected_valid) in [ + (MAX_VERIFY_WINDOW_PIPELINE_DEPTH, true), + (MAX_VERIFY_WINDOW_PIPELINE_DEPTH + 1, false), + ] { + let dir = unique_test_dir(&format!("window-pipeline-depth-{depth}")); + let package = write_package_fixture(&dir, true); + write_generation_to_manifest( + &package, + serde_json::json!({ + "speculative_decoding": { + "default": "ngram-suffix", + "proposers": { + "suffix": { + "type": "ngram-suffix", + "ngram_min": 5, + "ngram_max": 32, + "max_proposal_tokens": 48, + "history_scope": "request" + } + }, + "strategies": { + "ngram-suffix": { + "type": "ngram-suffix", + "proposer": "suffix", + "window_policy": { + "default": "fixed", + "initial_window": 32, + "min_window": 1, + "max_window": 32, + "pipeline_depth": depth + } + } + } + } + }), + ); + + let report = preflight_package(&package, &PackagePreflightOptions::default()); + assert_eq!(report.valid, expected_valid, "pipeline depth {depth}"); + if !expected_valid { + assert_issue(&report, "invalid_window_policy_pipeline_depth"); + } + fs::remove_dir_all(dir).unwrap(); + } + } + fn assert_issue(report: &PackagePreflightReport, code: &str) { assert!( report.issues.iter().any(|issue| issue.code == code), diff --git a/crates/skippy-model-package/src/preflight/artifacts.rs b/crates/skippy-model-package/src/preflight/artifacts.rs new file mode 100644 index 000000000..f3f623456 --- /dev/null +++ b/crates/skippy-model-package/src/preflight/artifacts.rs @@ -0,0 +1,427 @@ +use super::*; + +pub(super) fn collect_artifacts(manifest: &PackageManifest) -> Vec { + let mut artifacts = vec![ + artifact_spec("metadata", None, &manifest.shared.metadata), + artifact_spec("embeddings", None, &manifest.shared.embeddings), + artifact_spec("output", None, &manifest.shared.output), + ]; + artifacts.extend( + manifest + .layers + .iter() + .map(|layer| layer_artifact_spec(layer.layer_index, layer)), + ); + artifacts.extend(manifest.projectors.iter().map(projector_artifact_spec)); + artifacts +} + +fn artifact_spec( + role: &'static str, + layer_index: Option, + artifact: &PackageArtifact, +) -> ArtifactSpec { + ArtifactSpec { + role, + layer_index, + path: artifact.path.clone(), + tensor_count: artifact.tensor_count, + tensor_bytes: artifact.tensor_bytes, + artifact_bytes: artifact.artifact_bytes, + sha256: artifact.sha256.clone(), + } +} + +fn layer_artifact_spec(layer_index: u32, layer: &PackageLayer) -> ArtifactSpec { + ArtifactSpec { + role: "layer", + layer_index: Some(layer_index), + path: layer.path.clone(), + tensor_count: layer.tensor_count, + tensor_bytes: layer.tensor_bytes, + artifact_bytes: layer.artifact_bytes, + sha256: layer.sha256.clone(), + } +} + +fn projector_artifact_spec(projector: &PackageProjector) -> ArtifactSpec { + ArtifactSpec { + role: "projector", + layer_index: None, + path: projector.path.clone(), + tensor_count: projector.tensor_count, + tensor_bytes: projector.tensor_bytes, + artifact_bytes: projector.artifact_bytes, + sha256: projector.sha256.clone(), + } +} + +pub(super) fn validate_layer_coverage( + manifest: &PackageManifest, + report: &mut PackagePreflightReport, +) { + let mut counts = BTreeMap::::new(); + for layer in &manifest.layers { + *counts.entry(layer.layer_index).or_default() += 1; + if layer.layer_index >= manifest.layer_count { + report.error( + "layer_index_out_of_range", + format!( + "package layer index {} exceeds layer_count {}", + layer.layer_index, manifest.layer_count + ), + Some(layer.path.clone()), + "rebuild the package so layer indexes are contiguous and in range", + ); + } + } + for layer_index in 0..manifest.layer_count { + if !counts.contains_key(&layer_index) { + report.error( + "missing_layer", + format!("package manifest is missing layer {layer_index}"), + Some("model-package.json".to_string()), + "rebuild the package so every transformer layer has one artifact", + ); + } + } + for (layer_index, count) in counts { + if count > 1 { + report.error( + "duplicate_layer", + format!("package manifest contains layer {layer_index} {count} times"), + Some("model-package.json".to_string()), + "rebuild the package so each layer appears exactly once", + ); + } + } +} + +pub(super) fn validate_artifacts( + package: &Path, + artifacts: &[ArtifactSpec], + verify_sha256: bool, + report: &mut PackagePreflightReport, +) { + for artifact in artifacts { + report.artifacts.push(preflight_artifact( + package, + artifact, + verify_sha256, + &mut report.issues, + )); + } +} + +fn preflight_artifact( + package: &Path, + artifact: &ArtifactSpec, + verify_sha256: bool, + issues: &mut Vec, +) -> PreflightArtifact { + let path = match safe_relative_path(&artifact.path) { + Ok(path) => path, + Err(message) => { + push_error( + issues, + "unsafe_artifact_path", + format!( + "package {} artifact path is unsafe: {message}", + artifact.role + ), + Some(artifact.path.clone()), + "rebuild the package so artifact paths stay inside the package directory", + ); + return artifact_output(artifact, false, None, None, None); + } + }; + validate_artifact_manifest(artifact, issues); + let absolute = package.join(&path); + let metadata = match fs::metadata(&absolute) { + Ok(metadata) if metadata.is_file() => metadata, + Ok(_) => { + push_error( + issues, + "artifact_not_file", + format!("package artifact {} is not a file", artifact.path), + Some(artifact.path.clone()), + "replace the artifact path with a regular GGUF file", + ); + return artifact_output(artifact, false, None, None, None); + } + Err(error) => { + push_error( + issues, + "missing_artifact", + format!("package artifact {} is missing: {error}", artifact.path), + Some(artifact.path.clone()), + "download or rebuild the package artifact before starting split serving", + ); + return artifact_output(artifact, false, None, None, None); + } + }; + let actual_len = metadata.len(); + let size_matches = actual_len == artifact.artifact_bytes; + if !size_matches { + push_error( + issues, + "artifact_size_mismatch", + format!( + "package artifact {} has {} bytes, manifest expects {}", + artifact.path, actual_len, artifact.artifact_bytes + ), + Some(artifact.path.clone()), + "redownload or rebuild the package artifact so manifest sizes match", + ); + } + let sha_matches = if verify_sha256 { + Some(validate_artifact_sha(&absolute, artifact, issues)) + } else { + None + }; + artifact_output( + artifact, + true, + Some(actual_len), + Some(size_matches), + sha_matches, + ) +} + +fn validate_artifact_manifest(artifact: &ArtifactSpec, issues: &mut Vec) { + if artifact.artifact_bytes == 0 { + push_error( + issues, + "empty_artifact", + format!("package {} artifact declares zero bytes", artifact.role), + Some(artifact.path.clone()), + "rebuild the package; split artifacts must be non-empty files", + ); + } + if artifact.tensor_count == 0 && artifact.tensor_bytes > 0 { + push_error( + issues, + "invalid_tensor_bytes", + format!( + "package {} artifact declares tensor_bytes without tensors", + artifact.role + ), + Some(artifact.path.clone()), + "rebuild the package manifest so tensor counts and bytes agree", + ); + } + if artifact.tensor_count > 0 && artifact.tensor_bytes == 0 { + push_error( + issues, + "invalid_tensor_bytes", + format!( + "package {} artifact declares tensors but zero tensor_bytes", + artifact.role + ), + Some(artifact.path.clone()), + "rebuild the package manifest so tensor counts and bytes agree", + ); + } + if !is_sha256(&artifact.sha256) { + push_error( + issues, + "invalid_artifact_sha256", + format!( + "package {} artifact sha256 is not a hex digest", + artifact.role + ), + Some(artifact.path.clone()), + "rebuild the package manifest so artifact checksums are valid", + ); + } +} + +pub(super) fn validate_artifact_sha( + path: &Path, + artifact: &ArtifactSpec, + issues: &mut Vec, +) -> bool { + match file_sha256(path) { + Ok(actual) if actual == artifact.sha256.to_ascii_lowercase() => true, + Ok(actual) => { + push_error( + issues, + "artifact_sha256_mismatch", + format!( + "package artifact {} checksum mismatch: expected {}, got {}", + artifact.path, artifact.sha256, actual + ), + Some(artifact.path.clone()), + "redownload or rebuild the package artifact so checksums match", + ); + false + } + Err(error) => { + push_error( + issues, + "artifact_sha256_unreadable", + format!("cannot hash package artifact {}: {error}", artifact.path), + Some(artifact.path.clone()), + "ensure the artifact is readable before enabling checksum verification", + ); + false + } + } +} + +fn artifact_output( + artifact: &ArtifactSpec, + present: bool, + actual_artifact_bytes: Option, + size_matches_manifest: Option, + sha256_matches_manifest: Option, +) -> PreflightArtifact { + PreflightArtifact { + role: artifact.role.to_string(), + layer_index: artifact.layer_index, + path: artifact.path.clone(), + present, + declared_artifact_bytes: artifact.artifact_bytes, + actual_artifact_bytes, + size_matches_manifest, + sha256_matches_manifest, + } +} + +pub(super) fn build_stage_reports( + manifest: &PackageManifest, + stages: Option, + report: &mut PackagePreflightReport, +) { + let Some(stage_count) = stages else { + return; + }; + if stage_count == 0 { + report.error( + "invalid_stage_count", + "--stages must be greater than zero", + Some("model-package.json".to_string()), + "choose a positive stage count for split preflight", + ); + return; + } + if stage_count as u32 > manifest.layer_count { + report.error( + "stage_count_exceeds_layer_count", + format!( + "--stages {stage_count} exceeds package layer_count {}", + manifest.layer_count + ), + Some("model-package.json".to_string()), + "use at most one split stage per transformer layer", + ); + return; + } + let artifact_map = stage_artifacts(report); + for (stage_index, (layer_start, layer_end)) in + partition_layers(manifest.layer_count, stage_count) + .into_iter() + .enumerate() + { + report.stages.push(stage_report( + stage_index, + layer_start, + layer_end, + stage_count, + &artifact_map, + )); + } +} + +fn stage_report( + stage_index: usize, + layer_start: u32, + layer_end: u32, + stage_count: usize, + artifact_map: &BTreeMap, +) -> PreflightStage { + let includes_embeddings = stage_index == 0; + let includes_output = stage_index + 1 == stage_count; + let mut parts = vec!["metadata".to_string()]; + if includes_embeddings { + parts.push("embeddings".to_string()); + } + for layer_index in layer_start..layer_end { + parts.push(format!("layer:{layer_index}")); + } + if includes_output { + parts.push("output".to_string()); + } + let artifact_bytes = parts + .iter() + .filter_map(|part| artifact_map.get(part)) + .filter(|artifact| artifact.present) + .map(|artifact| artifact.bytes) + .sum(); + let missing_parts = parts + .iter() + .filter(|part| { + !artifact_map + .get(*part) + .is_some_and(|artifact| artifact.present) + }) + .cloned() + .collect::>(); + PreflightStage { + stage_index, + layer_start, + layer_end, + includes_embeddings, + includes_output, + part_count: parts.len(), + artifact_bytes, + parts, + missing_parts, + } +} + +#[derive(Clone, Copy)] +struct StageArtifact { + present: bool, + bytes: u64, +} + +fn stage_artifacts(report: &PackagePreflightReport) -> BTreeMap { + report + .artifacts + .iter() + .map(|artifact| { + ( + stage_part_key(artifact), + StageArtifact { + present: artifact.present, + bytes: artifact + .actual_artifact_bytes + .unwrap_or(artifact.declared_artifact_bytes), + }, + ) + }) + .collect() +} + +fn stage_part_key(artifact: &PreflightArtifact) -> String { + match (artifact.role.as_str(), artifact.layer_index) { + ("layer", Some(layer)) => format!("layer:{layer}"), + (role, _) => role.to_string(), + } +} + +fn partition_layers(layer_count: u32, stages: usize) -> Vec<(u32, u32)> { + let base = layer_count / stages as u32; + let extra = layer_count % stages as u32; + let mut start = 0; + (0..stages) + .map(|stage_index| { + let width = base + u32::from((stage_index as u32) < extra); + let end = start + width; + let range = (start, end); + start = end; + range + }) + .collect() +} diff --git a/crates/skippy-protocol/src/binary/activation.rs b/crates/skippy-protocol/src/binary/activation.rs index 644bbb117..2e056a931 100644 --- a/crates/skippy-protocol/src/binary/activation.rs +++ b/crates/skippy-protocol/src/binary/activation.rs @@ -102,7 +102,10 @@ pub fn encode_f32_activation_payload_with_state_flags( pub fn activation_payload_multiplier_from_state_flags(state_flag_bits: i32) -> usize { if (state_flag_bits & state_flags::GEMMA3N_ALTUP_SIDEBAND) != 0 { 4 - } else if (state_flag_bits & state_flags::RWKV7_V_FIRST_SIDEBAND) != 0 { + } else if (state_flag_bits + & (state_flags::INKLING_MTP_EMBD_SIDEBAND | state_flags::RWKV7_V_FIRST_SIDEBAND)) + != 0 + { 2 } else { 1 diff --git a/crates/skippy-protocol/src/binary/mod.rs b/crates/skippy-protocol/src/binary/mod.rs index cb8004b66..f7dfbdce9 100644 --- a/crates/skippy-protocol/src/binary/mod.rs +++ b/crates/skippy-protocol/src/binary/mod.rs @@ -15,8 +15,8 @@ pub use codec::{ send_reply_predicted_with_tokens_window_and_stats, write_stage_message, }; pub use types::{ - ACTIVATION_FLAG_GEMMA3N_ALTUP, ACTIVATION_FLAG_RWKV7_V_FIRST, LLAMA_TOKEN_NULL, - MAX_STAGE_ACTIVATION_BYTES, MAX_STAGE_CHAT_SAMPLING_METADATA_BYTES, + ACTIVATION_FLAG_GEMMA3N_ALTUP, ACTIVATION_FLAG_INKLING_MTP_EMBD, ACTIVATION_FLAG_RWKV7_V_FIRST, + LLAMA_TOKEN_NULL, MAX_STAGE_ACTIVATION_BYTES, MAX_STAGE_CHAT_SAMPLING_METADATA_BYTES, MAX_STAGE_DECODED_ACTIVATION_BYTES, MAX_STAGE_LOGIT_BIAS, MAX_STAGE_PREDICTED_TOKENS, MAX_STAGE_SIDEBAND_VALUES, MAX_STAGE_STATE_IMPORT_BYTES, READY_MAGIC, STAGE_LOGIT_BIAS_WIRE_BYTES, STAGE_SAMPLING_CONFIG_BASE_BYTES, STAGE_STATE_HEADER_BYTES, @@ -626,6 +626,33 @@ mod tests { const { assert!(STAGE_WIRE_FIXED_HEADER_BYTES <= 80) }; } + #[test] + fn verify_retirement_round_trips_exact_identity() { + let kind = WireMessageKind::RetireVerifyWindow; + let message = StageWireMessage { + kind, + pos_start: 128, + token_count: 8, + state: StageStateHeader::new(kind, WireActivationDType::F32), + request_id: 23, + session_id: 29, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + }; + let mut bytes = Vec::new(); + write_stage_message(&mut bytes, &message, WireActivationDType::F32).unwrap(); + let decoded = read_stage_message(Cursor::new(bytes), 2048).unwrap(); + + assert_eq!(decoded.kind, kind); + assert_eq!(decoded.pos_start, 128); + assert_eq!(decoded.token_count, 8); + assert!(decoded.state.matches_kind(kind)); + } + #[test] fn session_control_messages_are_fixed_header_only() { let kind = WireMessageKind::TrimSession; @@ -905,6 +932,48 @@ mod tests { ); } + #[test] + fn inkling_mtp_embedding_sideband_activation_round_trips() { + let mut state = + StageStateHeader::new(WireMessageKind::PrefillEmbd, WireActivationDType::F32); + state.source_stage_index = 0; + state.flags |= state_flags::INKLING_MTP_EMBD_SIDEBAND; + let mut activation = Vec::new(); + for value in [1.0_f32, 2.0, 3.0, 4.0] { + activation.extend_from_slice(&value.to_le_bytes()); + } + let message = StageWireMessage { + kind: WireMessageKind::PrefillEmbd, + pos_start: 0, + token_count: 1, + state, + request_id: 7, + session_id: 9, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation, + raw_bytes: Vec::new(), + }; + let mut bytes = Vec::new(); + write_stage_message(&mut bytes, &message, WireActivationDType::F32).unwrap(); + let decoded = read_stage_message(Cursor::new(bytes), 2).unwrap(); + assert_eq!(decoded.activation.len(), 16); + assert_eq!( + activation_frame_flags_from_state_flags(decoded.state.flags), + ACTIVATION_FLAG_INKLING_MTP_EMBD + ); + assert_eq!( + activation_state_flags_from_frame_flags(ACTIVATION_FLAG_INKLING_MTP_EMBD), + state_flags::INKLING_MTP_EMBD_SIDEBAND + ); + assert_eq!( + decoded.activation_f32_payload(2).unwrap(), + message.activation + ); + } + #[test] fn f32_activation_payload_can_be_moved_without_clone() { let state = StageStateHeader::new(WireMessageKind::DecodeEmbd, WireActivationDType::F32); diff --git a/crates/skippy-protocol/src/binary/types.rs b/crates/skippy-protocol/src/binary/types.rs index c89eb8ca5..4541f7f9f 100644 --- a/crates/skippy-protocol/src/binary/types.rs +++ b/crates/skippy-protocol/src/binary/types.rs @@ -5,9 +5,10 @@ use super::{ invalid_data, }; -// v10 makes the coordinator the sole owner of verify-window acceptance and removes the -// redundant tail-stage acceptance/correction fields. Stage peers must be upgraded together. -pub const STAGE_STATE_VERSION: i32 = 10; +// v11 adds the Inkling MTP embedding sideband and makes the coordinator the sole owner of +// verify-window acceptance, removing redundant tail-stage acceptance/correction fields. Stage +// peers must be upgraded together so older readers reject the changed payload contract. +pub const STAGE_STATE_VERSION: i32 = 11; pub const MAX_STAGE_LOGIT_BIAS: usize = 256; pub const MAX_STAGE_PREDICTED_TOKENS: usize = 262_144; pub const MAX_STAGE_SIDEBAND_VALUES: usize = 1_048_576; @@ -56,6 +57,7 @@ pub enum WireMessageKind { DecodeReadout = 8, DecodeLightCtx = 9, VerifyWindow = 21, + RetireVerifyWindow = 22, StateExport = 13, ConfigureGeneration = 14, ProbePrefill = 15, @@ -95,6 +97,10 @@ impl WireMessageKind { matches!(self, Self::TrimSession) } + pub fn is_verify_retirement(self) -> bool { + matches!(self, Self::RetireVerifyWindow) + } + pub fn is_generation_control(self) -> bool { matches!(self, Self::ConfigureGeneration) } @@ -140,6 +146,7 @@ impl TryFrom for WireMessageKind { 19 => Ok(Self::TrimSession), 20 => Ok(Self::PredictionReturnOpen), 21 => Ok(Self::VerifyWindow), + 22 => Ok(Self::RetireVerifyWindow), _ => Err(invalid_data("unknown stage message kind")), } } @@ -183,10 +190,12 @@ pub mod state_flags { pub const CHAT_SAMPLING_METADATA: i32 = 1 << 5; pub const RWKV7_V_FIRST_SIDEBAND: i32 = 1 << 6; pub const GEMMA3N_ALTUP_SIDEBAND: i32 = 1 << 7; + pub const INKLING_MTP_EMBD_SIDEBAND: i32 = 1 << 8; } pub const ACTIVATION_FLAG_RWKV7_V_FIRST: u64 = 1 << 0; pub const ACTIVATION_FLAG_GEMMA3N_ALTUP: u64 = 1 << 1; +pub const ACTIVATION_FLAG_INKLING_MTP_EMBD: u64 = 1 << 2; pub fn activation_frame_flags_from_state_flags(flags: i32) -> u64 { let mut frame_flags = 0; @@ -196,6 +205,9 @@ pub fn activation_frame_flags_from_state_flags(flags: i32) -> u64 { if (flags & state_flags::GEMMA3N_ALTUP_SIDEBAND) != 0 { frame_flags |= ACTIVATION_FLAG_GEMMA3N_ALTUP; } + if (flags & state_flags::INKLING_MTP_EMBD_SIDEBAND) != 0 { + frame_flags |= ACTIVATION_FLAG_INKLING_MTP_EMBD; + } frame_flags } @@ -207,6 +219,9 @@ pub fn activation_state_flags_from_frame_flags(flags: u64) -> i32 { if (flags & ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0 { state |= state_flags::GEMMA3N_ALTUP_SIDEBAND; } + if (flags & ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0 { + state |= state_flags::INKLING_MTP_EMBD_SIDEBAND; + } state } @@ -313,6 +328,7 @@ impl StageStateHeader { kind, WireMessageKind::StateImport | WireMessageKind::StateExport ) || kind.is_session_control() + || kind.is_verify_retirement() || kind.is_generation_control() { return true; @@ -391,7 +407,7 @@ pub struct StageWireMessage { impl StageWireMessage { /// The committed session position that must exist before this message runs. /// - /// Stage-state v10 makes this absolute position authoritative: a worker whose + /// Stage-state v11 makes this absolute position authoritative: a worker whose /// speculative KV is ahead must rewind locally before executing the message. pub fn authoritative_session_position(&self) -> Option { if !matches!( @@ -695,6 +711,7 @@ fn expected_phase(kind: WireMessageKind) -> WireStagePhase { WireMessageKind::StateImport | WireMessageKind::StateExport ) || kind.is_session_control() + || kind.is_verify_retirement() || kind.is_generation_control() || kind.is_prefix_cache_control() { diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index b2aeec0ba..da4626698 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -12,19 +12,21 @@ pub const STAGE_ALPN_V2: &[u8] = b"skippy-stage/2"; pub const STAGE_SUBPROTOCOL_NAME: &str = "skippy-stage"; pub const STAGE_SUBPROTOCOL_MAJOR: u32 = 2; pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL: &str = "stage-control"; -pub const STAGE_PROTOCOL_GENERATION: u32 = 3; +pub const STAGE_PROTOCOL_GENERATION: u32 = 4; /// Generation-scoped stage capability. A peer can advertise `stage-control` /// while still rejecting current-generation frames, so split planning gates on /// this exact token before sending current-generation control requests. -pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V3: &str = "stage-generation-3"; +pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4: &str = "stage-generation-4"; pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION: &str = - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V3; + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4; pub const STAGE_SUBPROTOCOL_FEATURE_ARTIFACT_TRANSFER: &str = "artifact-transfer"; pub const STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST: &str = "status-list"; pub const STAGE_STREAM_CONTROL: u8 = 0x01; pub const STAGE_STREAM_TRANSPORT: u8 = 0x02; pub const STAGE_STREAM_ARTIFACT_TRANSFER: u8 = 0x03; pub const MAX_STAGE_FRAME_BYTES: usize = 8 * 1024 * 1024; +/// Maximum number of unresolved verify windows covered by native checkpoints. +pub const MAX_VERIFY_WINDOW_PIPELINE_DEPTH: usize = 64; #[derive(Debug, Clone, PartialEq, Eq)] pub enum StageFrameError { @@ -662,7 +664,7 @@ mod tests { stage_control_response, }; use super::{ - STAGE_PROTOCOL_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V3, + STAGE_PROTOCOL_GENERATION, STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, StageFrameError, validate_stage_artifact_transfer_request, validate_stage_artifact_transfer_response, validate_stage_control_request, validate_stage_control_response, validate_stage_transport_open, @@ -671,7 +673,7 @@ mod tests { #[test] fn stage_protocol_generation_feature_names_current_generation() { assert_eq!( - STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V3, + STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V4, format!("stage-generation-{STAGE_PROTOCOL_GENERATION}") ); } diff --git a/crates/skippy-quantize/Cargo.toml b/crates/skippy-quantize/Cargo.toml index c75e8b469..f24f51507 100644 --- a/crates/skippy-quantize/Cargo.toml +++ b/crates/skippy-quantize/Cargo.toml @@ -10,8 +10,9 @@ readme = "README.md" publish = false [features] -default = [] +default = ["dynamic-skippy-runtime"] dynamic-llama-quant = ["llama-quant-ffi/dynamic-runtime"] +dynamic-skippy-runtime = ["skippy-ffi/dynamic-runtime", "skippy-runtime/dynamic-native-runtime"] [dependencies] anyhow.workspace = true @@ -20,7 +21,8 @@ libc = "0.2" llama-quant-ffi = { path = "../llama-quant-ffi" } serde.workspace = true serde_json.workspace = true -skippy-ffi = { path = "../skippy-ffi" } +skippy-ffi = { path = "../skippy-ffi", default-features = false } +skippy-runtime = { path = "../skippy-runtime", default-features = false } [dev-dependencies] regex-lite = "0.1" diff --git a/crates/skippy-quantize/src/backend.rs b/crates/skippy-quantize/src/backend.rs index 6907a8eb6..e4ac486dc 100644 --- a/crates/skippy-quantize/src/backend.rs +++ b/crates/skippy-quantize/src/backend.rs @@ -129,7 +129,7 @@ fn skippy_abi_capabilities(skippy_runtime_libraries: &[PathBuf]) -> SkippyAbiCap let load_error = load_skippy_runtime_for_probe(skippy_runtime_libraries); let runtime_loaded = skippy_ffi::native_runtime_loaded(); let feature_mask = if runtime_loaded { - std::panic::catch_unwind(skippy_ffi::skippy_abi_features).ok() + std::panic::catch_unwind(skippy_ffi::abi_features).ok() } else { None }; @@ -170,7 +170,7 @@ fn skippy_abi_reason( return "no Skippy native runtime library was loaded for ABI probing".to_string(); } if feature_mask.is_none() { - return "loaded Skippy runtime does not expose skippy_abi_features".to_string(); + return "loaded Skippy runtime does not expose abi_features".to_string(); } if gguf_slice_write { return "loaded Skippy ABI exposes GGUF slice writing and the linked llama symbols can be used for GGUF quantization, but not HF checkpoint conversion".to_string(); diff --git a/crates/skippy-quantize/src/direct_convert.rs b/crates/skippy-quantize/src/direct_convert.rs index a07d8619c..28cd04507 100644 --- a/crates/skippy-quantize/src/direct_convert.rs +++ b/crates/skippy-quantize/src/direct_convert.rs @@ -89,7 +89,8 @@ pub(crate) fn run_direct_convert(args: DirectConvertArgs) -> Result<()> { window_size: args.window_size, manifest: manifest_path.clone(), }; - let manifest = convert_manifest_from_args(&manifest_args)?; + let mut manifest = convert_manifest_from_args(&manifest_args)?; + crate::native_convert::apply_native_convert_split_max_size(&runner, &mut manifest)?; if args.preflight_only { return run_job_preflight( &manifest_path, diff --git a/crates/skippy-quantize/src/gguf_metadata.rs b/crates/skippy-quantize/src/gguf_metadata.rs new file mode 100644 index 000000000..e960acbee --- /dev/null +++ b/crates/skippy-quantize/src/gguf_metadata.rs @@ -0,0 +1,207 @@ +use std::io::Write; + +use anyhow::{Result, ensure}; + +pub(crate) const GGUF_TYPE_UINT16: u32 = 2; +pub(crate) const GGUF_TYPE_UINT32: u32 = 4; +pub(crate) const GGUF_TYPE_INT32: u32 = 5; +pub(crate) const GGUF_TYPE_FLOAT32: u32 = 6; +pub(crate) const GGUF_TYPE_BOOL: u32 = 7; +pub(crate) const GGUF_TYPE_STRING: u32 = 8; +pub(crate) const GGUF_TYPE_ARRAY: u32 = 9; +pub(crate) const GGUF_TYPE_UINT64: u32 = 10; + +#[derive(Debug, Clone)] +pub(crate) enum GgufKv { + ArrayBool { key: String, value: Vec }, + ArrayF32 { key: String, value: Vec }, + ArrayI32 { key: String, value: Vec }, + ArrayString { key: String, value: Vec }, + ArrayU32 { key: String, value: Vec }, + Bool { key: String, value: bool }, + F32 { key: String, value: f32 }, + I32 { key: String, value: i32 }, + String { key: String, value: String }, + U16 { key: String, value: u16 }, + U32 { key: String, value: u32 }, + U64 { key: String, value: u64 }, +} + +impl GgufKv { + pub(crate) fn array_bool(key: &str, value: Vec) -> Self { + Self::ArrayBool { + key: key.to_string(), + value, + } + } + + pub(crate) fn array_f32(key: &str, value: Vec) -> Self { + Self::ArrayF32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn array_i32(key: &str, value: Vec) -> Self { + Self::ArrayI32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn array_string(key: &str, value: Vec) -> Self { + Self::ArrayString { + key: key.to_string(), + value, + } + } + + pub(crate) fn array_u32(key: &str, value: Vec) -> Self { + Self::ArrayU32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn bool(key: &str, value: bool) -> Self { + Self::Bool { + key: key.to_string(), + value, + } + } + + pub(crate) fn f32(key: &str, value: f32) -> Self { + Self::F32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn i32(key: &str, value: i32) -> Self { + Self::I32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn string(key: &str, value: &str) -> Self { + Self::String { + key: key.to_string(), + value: value.to_string(), + } + } + + pub(crate) fn u16(key: &str, value: u16) -> Self { + Self::U16 { + key: key.to_string(), + value, + } + } + + pub(crate) fn u32(key: &str, value: u32) -> Self { + Self::U32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn u64(key: &str, value: u64) -> Self { + Self::U64 { + key: key.to_string(), + value, + } + } +} + +pub(crate) fn write_kv(writer: &mut W, kv: &GgufKv) -> Result<()> { + match kv { + GgufKv::ArrayBool { key, value } => { + write_array_header(writer, key, GGUF_TYPE_BOOL, value.len())?; + for item in value { + writer.write_all(&[*item as u8])?; + } + } + GgufKv::ArrayF32 { key, value } => { + write_array_header(writer, key, GGUF_TYPE_FLOAT32, value.len())?; + for item in value { + writer.write_all(&item.to_le_bytes())?; + } + } + GgufKv::ArrayI32 { key, value } => { + write_array_header(writer, key, GGUF_TYPE_INT32, value.len())?; + for item in value { + writer.write_all(&item.to_le_bytes())?; + } + } + GgufKv::ArrayString { key, value } => { + write_array_header(writer, key, GGUF_TYPE_STRING, value.len())?; + for item in value { + write_string(writer, item)?; + } + } + GgufKv::ArrayU32 { key, value } => { + write_array_header(writer, key, GGUF_TYPE_UINT32, value.len())?; + for item in value { + writer.write_all(&item.to_le_bytes())?; + } + } + GgufKv::Bool { key, value } => { + write_scalar_header(writer, key, GGUF_TYPE_BOOL)?; + writer.write_all(&[*value as u8])?; + } + GgufKv::F32 { key, value } => { + write_scalar_header(writer, key, GGUF_TYPE_FLOAT32)?; + writer.write_all(&value.to_le_bytes())?; + } + GgufKv::I32 { key, value } => { + write_scalar_header(writer, key, GGUF_TYPE_INT32)?; + writer.write_all(&value.to_le_bytes())?; + } + GgufKv::String { key, value } => { + write_scalar_header(writer, key, GGUF_TYPE_STRING)?; + write_string(writer, value)?; + } + GgufKv::U16 { key, value } => { + write_scalar_header(writer, key, GGUF_TYPE_UINT16)?; + writer.write_all(&value.to_le_bytes())?; + } + GgufKv::U32 { key, value } => { + write_scalar_header(writer, key, GGUF_TYPE_UINT32)?; + writer.write_all(&value.to_le_bytes())?; + } + GgufKv::U64 { key, value } => { + write_scalar_header(writer, key, GGUF_TYPE_UINT64)?; + writer.write_all(&value.to_le_bytes())?; + } + } + Ok(()) +} + +fn write_array_header( + writer: &mut W, + key: &str, + element_type: u32, + len: usize, +) -> Result<()> { + ensure!( + len > 0, + "GGUF array metadata {key:?} cannot be empty because llama.cpp rejects empty arrays" + ); + write_scalar_header(writer, key, GGUF_TYPE_ARRAY)?; + writer.write_all(&element_type.to_le_bytes())?; + writer.write_all(&(len as u64).to_le_bytes())?; + Ok(()) +} + +fn write_scalar_header(writer: &mut W, key: &str, value_type: u32) -> Result<()> { + ensure!(!key.is_empty(), "GGUF metadata key cannot be empty"); + write_string(writer, key)?; + writer.write_all(&value_type.to_le_bytes())?; + Ok(()) +} + +fn write_string(writer: &mut W, value: &str) -> Result<()> { + writer.write_all(&(value.len() as u64).to_le_bytes())?; + writer.write_all(value.as_bytes())?; + Ok(()) +} diff --git a/crates/skippy-quantize/src/gguf_template.rs b/crates/skippy-quantize/src/gguf_template.rs index b871a1f4b..04fdae709 100644 --- a/crates/skippy-quantize/src/gguf_template.rs +++ b/crates/skippy-quantize/src/gguf_template.rs @@ -5,6 +5,7 @@ use anyhow::{Context, Result, ensure}; use serde_json::Value; use crate::gguf_writer::GgufKv; +use crate::inkling_metadata; use crate::tokenizer_metadata::push_tokenizer_metadata; #[derive(Debug, Clone, Copy)] @@ -24,6 +25,9 @@ pub(crate) fn metadata_from_hf_config(source: &Path, tensor_count: usize) -> Res pub(crate) fn mtp_layer_start_from_hf_config(source: &Path) -> Result> { let config = read_hf_config(source)?; + if inkling_metadata::is_inkling_config(&config) { + return inkling_metadata::mtp_layer_start(&config); + } let Some(nextn_layers) = optional_u32(&config, "num_nextn_predict_layers") .or_else(|| optional_u32(&config, "mtp_num_hidden_layers")) else { @@ -41,6 +45,9 @@ pub(crate) fn metadata_from_hf_config_with_options( options: MetadataOptions, ) -> Result> { let config = read_hf_config(source)?; + if inkling_metadata::is_inkling_config(&config) { + return inkling_metadata::metadata(source, tensor_count, &config, options.include_mtp); + } let arch = architecture_name(&config)?; let mut metadata = vec![ GgufKv::string("general.architecture", arch), @@ -1139,6 +1146,75 @@ mod tests { } } + #[test] + fn builds_inkling_multi_depth_mtp_metadata() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("config.json"), + r#"{ + "model_type": "inkling_mm_model", + "eos_token_id": 200006, + "text_config": { + "model_max_length": 1048576, + "hidden_size": 6144, + "num_hidden_layers": 66, + "vocab_size": 201024, + "num_attention_heads": 64, + "num_key_value_heads": 8, + "head_dim": 128, + "d_rel": 16, + "rel_extent": 1024, + "log_scaling_n_floor": 128000, + "log_scaling_alpha": 0.1, + "rms_norm_eps": 1e-6, + "local_layer_ids": [0, 1, 2, 3, 4], + "dense_mlp_idx": 2, + "sconv_kernel_size": 4, + "unpadded_vocab_size": 200058, + "logits_mup_width_multiplier": 24.0, + "swa_num_key_value_heads": 16, + "sliding_window_size": 512, + "n_routed_experts": 256, + "num_experts_per_tok": 6, + "n_shared_experts": 2, + "dense_intermediate_size": 24576, + "intermediate_size": 3072, + "route_scale": 8.0, + "norm_after_topk": true, + "shared_expert_sink": true, + "use_sconv": true, + "use_embed_norm": true, + "use_gate_bias": true, + "use_global_scale": true, + "gate_activation": "sigmoid" + }, + "mtp_config": { + "num_nextn_predict_layers": 8, + "chain_hidden_post_norm": false, + "local_layer_ids": [0, 2, 4, 5, 6, 7] + } + }"#, + ) + .unwrap(); + + assert_eq!(mtp_layer_start_from_hf_config(&root).unwrap(), Some(66)); + let metadata = metadata_from_hf_config(&root, 25).unwrap(); + assert!(metadata.iter().any(|kv| { + matches!(kv, GgufKv::U32 { key, value } if key == "inkling.block_count" && *value == 74) + })); + assert!(metadata.iter().any(|kv| { + matches!(kv, GgufKv::U32 { key, value } if key == "inkling.nextn_predict_layers" && *value == 8) + })); + assert!(metadata.iter().any(|kv| { + matches!(kv, GgufKv::ArrayU32 { key, value } if key == "inkling.attention.head_count_kv" && value.len() == 74 && value[66] == 16 && value[67] == 8) + })); + assert!(metadata.iter().any(|kv| { + matches!(kv, GgufKv::ArrayBool { key, value } if key == "inkling.attention.sliding_window_pattern" && value.len() == 74 && value[66] && !value[67]) + })); + fs::remove_dir_all(root).unwrap(); + } + fn unique_temp_dir() -> PathBuf { static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); let nanos = std::time::SystemTime::now() diff --git a/crates/skippy-quantize/src/gguf_writer.rs b/crates/skippy-quantize/src/gguf_writer.rs index c845eff4e..1c0ffc5ba 100644 --- a/crates/skippy-quantize/src/gguf_writer.rs +++ b/crates/skippy-quantize/src/gguf_writer.rs @@ -7,9 +7,17 @@ use anyhow::{Context, Result, ensure}; use serde::Serialize; use crate::float_convert::{FloatDType, convert_float_chunk, target_dtype_for_tensor}; +pub(crate) use crate::gguf_metadata::GgufKv; +use crate::gguf_metadata::write_kv; +#[cfg(test)] +use crate::gguf_metadata::{ + GGUF_TYPE_ARRAY, GGUF_TYPE_BOOL, GGUF_TYPE_FLOAT32, GGUF_TYPE_INT32, GGUF_TYPE_STRING, + GGUF_TYPE_UINT16, GGUF_TYPE_UINT32, GGUF_TYPE_UINT64, +}; use crate::hf_checkpoint::{SafetensorFile, SafetensorTensorInfo, open_safetensor_files}; use crate::tensor_map::{ - TensorNameMap, hf_layer_id, is_mtp_source_tensor, is_shared_mtp_context_tensor, + TensorNameMap, hf_layer_id, inkling_mtp_depth, is_inkling_fused_w13, is_mtp_source_tensor, + is_shared_mtp_context_tensor, }; use crate::types::ConvertOutputType; @@ -23,14 +31,6 @@ use glm_dsa::{ const GGUF_MAGIC: &[u8; 4] = b"GGUF"; const GGUF_VERSION: u32 = 3; const GGUF_ALIGNMENT: u64 = 32; -const GGUF_TYPE_BOOL: u32 = 7; -const GGUF_TYPE_UINT32: u32 = 4; -const GGUF_TYPE_INT32: u32 = 5; -const GGUF_TYPE_FLOAT32: u32 = 6; -const GGUF_TYPE_STRING: u32 = 8; -const GGUF_TYPE_ARRAY: u32 = 9; -const GGUF_TYPE_UINT16: u32 = 2; -const GGUF_TYPE_UINT64: u32 = 10; const GGML_TYPE_F32: u32 = 0; const GGML_TYPE_F16: u32 = 1; const GGML_TYPE_BF16: u32 = 30; @@ -96,6 +96,58 @@ pub(crate) fn validate_raw_safetensors_gguf( }) } +pub(crate) fn recommended_raw_safetensors_gguf_split_count( + source: &Path, + mut options: RawGgufWriteOptions, + max_tensor_bytes: u64, +) -> Result { + ensure!( + max_tensor_bytes > 0, + "split maximum tensor bytes must be greater than zero" + ); + options.split = None; + let PreparedGgufWrite { tensors, .. } = prepare_raw_safetensors_gguf(source, &options)?; + let largest_tensor_bytes = tensors + .iter() + .map(|tensor| tensor.byte_len) + .max() + .unwrap_or_default(); + ensure!( + largest_tensor_bytes <= max_tensor_bytes, + "largest selected tensor is {largest_tensor_bytes} bytes, exceeding split maximum {max_tensor_bytes} bytes" + ); + let total_tensor_bytes = tensors + .iter() + .try_fold(0_u64, |total, tensor| total.checked_add(tensor.byte_len)); + let total_tensor_bytes = total_tensor_bytes.context("selected tensor byte total overflow")?; + let minimum_count = total_tensor_bytes + .div_ceil(max_tensor_bytes) + .max(1) + .min(tensors.len() as u64); + let minimum_count = u32::try_from(minimum_count).context("split count does not fit u32")?; + let maximum_count = u32::try_from(tensors.len()).context("tensor count does not fit u32")?; + + for split_count in minimum_count..=maximum_count { + let split = GgufSplit { + split_index: 1, + split_count, + }; + let boundaries = byte_balanced_split_boundaries(&tensors, split)?; + let every_split_fits = boundaries.windows(2).all(|range| { + tensors[range[0]..range[1]] + .iter() + .map(|tensor| tensor.byte_len) + .sum::() + <= max_tensor_bytes + }); + if every_split_fits { + return Ok(split_count); + } + } + + anyhow::bail!("could not partition selected tensors within the split maximum") +} + struct PreparedGgufWrite { files: Vec, tensors: Vec, @@ -305,6 +357,15 @@ fn collect_tensor_sources( if !tensor_selection.includes(tensor.name())? { continue; } + if is_inkling_fused_w13(tensor.name()) { + tensors.extend(inkling_w13_tensor_sources( + file_index, + tensor, + tensor_name_map, + output_type, + )?); + continue; + } if matches!( tensor_name_map, TensorNameMap::HfToGguf | TensorNameMap::HfToGgufWithMtp { .. } @@ -394,9 +455,11 @@ impl TensorSource { let source_dtype = FloatDType::from_safetensor(tensor.dtype()).with_context(|| { format!("unsupported dtype {} for {}", tensor.dtype(), tensor.name()) })?; - let target_dtype = target_dtype_for_tensor(source_dtype, output_type, tensor.shape())?; let name = tensor_name_map.map_tensor_name(tensor.name())?; + let target_dtype = + target_dtype_for_mapped_tensor(source_dtype, output_type, tensor.shape(), &name)?; let element_count = tensor_element_count(tensor)?; + let dims = mapped_tensor_dims(tensor.shape(), &name)?; Ok(Self { segments: vec![TensorSegment { file_index, @@ -409,7 +472,7 @@ impl TensorSource { transform: TensorTransform::Identity, }], name, - dims: tensor.shape().iter().rev().copied().collect(), + dims, ggml_type: ggml_type_for_dtype(target_dtype), byte_len: tensor_byte_len(element_count, target_dtype)?, gguf_offset: 0, @@ -417,6 +480,98 @@ impl TensorSource { } } +fn target_dtype_for_mapped_tensor( + source_dtype: FloatDType, + output_type: Option, + shape: &[u64], + mapped_name: &str, +) -> Result { + if mapped_name.ends_with("attn_rel_proj.weight") || mapped_name.contains(".shortconv_") { + return Ok(FloatDType::F32); + } + target_dtype_for_tensor(source_dtype, output_type, shape) +} + +fn mapped_tensor_dims(shape: &[u64], mapped_name: &str) -> Result> { + if mapped_name.contains(".shortconv_") { + ensure!( + shape.len() == 3 && shape[1] == 1, + "Inkling shortconv tensor {mapped_name} must have shape [channels, 1, kernel], got {shape:?}" + ); + return Ok(vec![shape[2], shape[0]]); + } + Ok(shape.iter().rev().copied().collect()) +} + +fn inkling_w13_tensor_sources( + file_index: usize, + tensor: &SafetensorTensorInfo, + tensor_name_map: TensorNameMap, + output_type: Option, +) -> Result> { + let layer = if let Some(depth) = inkling_mtp_depth(tensor.name())? { + let TensorNameMap::HfToGgufWithMtp { layer_start } = tensor_name_map else { + anyhow::bail!("Inkling MTP conversion requires an MTP-aware tensor name map"); + }; + layer_start + .checked_add(depth) + .context("Inkling MTP layer id overflow")? + } else { + ensure!( + matches!( + tensor_name_map, + TensorNameMap::HfToGguf | TensorNameMap::HfToGgufWithMtp { .. } + ), + "Inkling fused w13 conversion requires an HF tensor name map" + ); + hf_layer_id(tensor.name())? + .with_context(|| format!("missing Inkling layer id in {}", tensor.name()))? + }; + ensure!( + tensor.shape().len() == 2, + "Inkling MTP fused w13 tensor {} must be rank 2, got {:?}", + tensor.name(), + tensor.shape() + ); + ensure!( + tensor.shape()[0].is_multiple_of(2), + "Inkling MTP fused w13 tensor {} must have an even row count", + tensor.name() + ); + let source_dtype = FloatDType::from_safetensor(tensor.dtype()) + .with_context(|| format!("unsupported dtype {} for {}", tensor.dtype(), tensor.name()))?; + let output_shape = [tensor.shape()[0] / 2, tensor.shape()[1]]; + let target_dtype = target_dtype_for_tensor(source_dtype, output_type, &output_shape)?; + let element_count = output_shape[0] + .checked_mul(output_shape[1]) + .context("Inkling MTP w13 output element count overflow")?; + let target_byte_len = tensor_byte_len(element_count, target_dtype)?; + let dims = output_shape.iter().rev().copied().collect::>(); + Ok([("ffn_gate", 0_u64), ("ffn_up", 1_u64)] + .into_iter() + .map(|(projection, parity)| TensorSource { + segments: vec![TensorSegment { + file_index, + source_name: tensor.name().to_string(), + source_dtype, + target_dtype, + element_count, + source_byte_len: tensor.byte_len(), + target_byte_len, + transform: TensorTransform::AlternatingRows { + parity, + row_elements: tensor.shape()[1], + }, + }], + name: format!("blk.{layer}.{projection}.weight"), + dims: dims.clone(), + ggml_type: ggml_type_for_dtype(target_dtype), + byte_len: target_byte_len, + gguf_offset: 0, + }) + .collect()) +} + struct TensorSegment { file_index: usize, source_name: String, @@ -653,92 +808,6 @@ fn raw_metadata(source: &Path, tensor_count: usize) -> Vec { ] } -#[derive(Debug, Clone)] -pub(crate) enum GgufKv { - ArrayF32 { key: String, value: Vec }, - ArrayI32 { key: String, value: Vec }, - ArrayString { key: String, value: Vec }, - Bool { key: String, value: bool }, - F32 { key: String, value: f32 }, - I32 { key: String, value: i32 }, - String { key: String, value: String }, - U16 { key: String, value: u16 }, - U32 { key: String, value: u32 }, - U64 { key: String, value: u64 }, -} - -impl GgufKv { - pub(crate) fn array_f32(key: &str, value: Vec) -> Self { - Self::ArrayF32 { - key: key.to_string(), - value, - } - } - - pub(crate) fn array_i32(key: &str, value: Vec) -> Self { - Self::ArrayI32 { - key: key.to_string(), - value, - } - } - - pub(crate) fn array_string(key: &str, value: Vec) -> Self { - Self::ArrayString { - key: key.to_string(), - value, - } - } - - pub(crate) fn bool(key: &str, value: bool) -> Self { - Self::Bool { - key: key.to_string(), - value, - } - } - - pub(crate) fn f32(key: &str, value: f32) -> Self { - Self::F32 { - key: key.to_string(), - value, - } - } - - pub(crate) fn i32(key: &str, value: i32) -> Self { - Self::I32 { - key: key.to_string(), - value, - } - } - - pub(crate) fn string(key: &str, value: &str) -> Self { - Self::String { - key: key.to_string(), - value: value.to_string(), - } - } - - pub(crate) fn u16(key: &str, value: u16) -> Self { - Self::U16 { - key: key.to_string(), - value, - } - } - - pub(crate) fn u32(key: &str, value: u32) -> Self { - Self::U32 { - key: key.to_string(), - value, - } - } - - pub(crate) fn u64(key: &str, value: u64) -> Self { - Self::U64 { - key: key.to_string(), - value, - } - } -} - fn write_header_and_tensor_table( writer: &mut W, metadata: &[GgufKv], @@ -763,82 +832,6 @@ fn write_header_and_tensor_table( Ok(()) } -fn write_kv(writer: &mut W, kv: &GgufKv) -> Result<()> { - match kv { - GgufKv::ArrayF32 { key, value } => { - write_array_header(writer, key, GGUF_TYPE_FLOAT32, value.len())?; - for item in value { - writer.write_all(&item.to_le_bytes())?; - } - } - GgufKv::ArrayI32 { key, value } => { - write_array_header(writer, key, GGUF_TYPE_INT32, value.len())?; - for item in value { - writer.write_all(&item.to_le_bytes())?; - } - } - GgufKv::ArrayString { key, value } => { - write_array_header(writer, key, GGUF_TYPE_STRING, value.len())?; - for item in value { - write_string(writer, item)?; - } - } - GgufKv::Bool { key, value } => { - write_string(writer, key)?; - write_u32(writer, GGUF_TYPE_BOOL)?; - writer.write_all(&[*value as u8])?; - } - GgufKv::F32 { key, value } => { - write_string(writer, key)?; - write_u32(writer, GGUF_TYPE_FLOAT32)?; - writer.write_all(&value.to_le_bytes())?; - } - GgufKv::I32 { key, value } => { - write_string(writer, key)?; - write_u32(writer, GGUF_TYPE_INT32)?; - writer.write_all(&value.to_le_bytes())?; - } - GgufKv::String { key, value } => { - write_string(writer, key)?; - write_u32(writer, GGUF_TYPE_STRING)?; - write_string(writer, value)?; - } - GgufKv::U16 { key, value } => { - write_string(writer, key)?; - write_u32(writer, GGUF_TYPE_UINT16)?; - writer.write_all(&value.to_le_bytes())?; - } - GgufKv::U32 { key, value } => { - write_string(writer, key)?; - write_u32(writer, GGUF_TYPE_UINT32)?; - write_u32(writer, *value)?; - } - GgufKv::U64 { key, value } => { - write_string(writer, key)?; - write_u32(writer, GGUF_TYPE_UINT64)?; - write_u64(writer, *value)?; - } - } - Ok(()) -} - -fn write_array_header( - writer: &mut W, - key: &str, - element_type: u32, - len: usize, -) -> Result<()> { - ensure!(!key.is_empty(), "GGUF metadata key cannot be empty"); - ensure!( - len > 0, - "GGUF array metadata {key:?} cannot be empty because llama.cpp rejects empty arrays" - ); - write_string(writer, key)?; - write_u32(writer, GGUF_TYPE_ARRAY)?; - write_u32(writer, element_type)?; - write_u64(writer, len as u64) -} - fn stream_tensor_data( writer: &mut File, files: &[SafetensorFile], @@ -880,6 +873,13 @@ fn stream_segment( segment: &TensorSegment, buffer_size: usize, ) -> Result { + if let TensorTransform::AlternatingRows { + parity, + row_elements, + } = segment.transform + { + return stream_alternating_rows(writer, file, segment, buffer_size, parity, row_elements); + } if let Some(written) = stream_transformed_segment(writer, file, segment, buffer_size)? { return Ok(written); } @@ -927,6 +927,59 @@ fn stream_segment( Ok(output_bytes) } +/// Deinterleave alternating rows of a fused SwiGLU tensor (Inkling MTP fused +/// w13): parity 0 keeps even rows (gate), parity 1 keeps odd rows (up). +fn stream_alternating_rows( + writer: &mut File, + file: &SafetensorFile, + segment: &TensorSegment, + buffer_size: usize, + parity: u64, + row_elements: u64, +) -> Result { + ensure!(parity < 2, "alternating-row parity must be zero or one"); + ensure!(row_elements > 0, "alternating-row width must be non-zero"); + let row_bytes = row_elements + .checked_mul(segment.source_dtype.byte_size()) + .context("alternating-row byte length overflow")?; + let row_bytes = usize::try_from(row_bytes).context("row byte length does not fit usize")?; + let chunk_size = aligned_chunk_size(buffer_size, row_bytes); + let mut source_bytes = 0_u64; + let mut output_bytes = 0_u64; + let mut row_index = 0_u64; + file.stream_tensor_chunks(&segment.source_name, chunk_size, |chunk| { + ensure!( + chunk.len() % row_bytes == 0, + "chunk for {} split a fused SwiGLU row", + segment.source_name + ); + source_bytes += chunk.len() as u64; + for row in chunk.chunks_exact(row_bytes) { + if row_index % 2 == parity { + output_bytes += + convert_float_chunk(row, segment.source_dtype, segment.target_dtype, writer)?; + } + row_index += 1; + } + Ok(()) + })?; + ensure!( + source_bytes == segment.source_byte_len, + "read {} bytes for {}, expected {}", + source_bytes, + segment.source_name, + segment.source_byte_len + ); + ensure!( + output_bytes == segment.target_byte_len, + "deinterleaved {} bytes for {}, expected {}", + output_bytes, + segment.source_name, + segment.target_byte_len + ); + Ok(output_bytes) +} + fn aligned_chunk_size(buffer_size: usize, element_size: usize) -> usize { let aligned = buffer_size - (buffer_size % element_size); aligned.max(element_size) diff --git a/crates/skippy-quantize/src/gguf_writer/glm_dsa.rs b/crates/skippy-quantize/src/gguf_writer/glm_dsa.rs index fcb31c2bd..a5f63c052 100644 --- a/crates/skippy-quantize/src/gguf_writer/glm_dsa.rs +++ b/crates/skippy-quantize/src/gguf_writer/glm_dsa.rs @@ -368,9 +368,11 @@ fn metadata_u32(metadata: &[GgufKv], key: &str) -> Option { impl GgufKv { fn key(&self) -> &str { match self { - Self::ArrayF32 { key, .. } + Self::ArrayBool { key, .. } + | Self::ArrayF32 { key, .. } | Self::ArrayI32 { key, .. } | Self::ArrayString { key, .. } + | Self::ArrayU32 { key, .. } | Self::Bool { key, .. } | Self::F32 { key, .. } | Self::I32 { key, .. } @@ -385,6 +387,13 @@ impl GgufKv { #[derive(Debug, Clone, Copy)] pub(super) enum TensorTransform { Identity, + /// Deinterleave alternating rows of a fused SwiGLU tensor (Inkling MTP + /// fused w13): parity 0 keeps even rows (gate), parity 1 keeps odd rows + /// (up). + AlternatingRows { + parity: u64, + row_elements: u64, + }, GlmDsaKvB { split: GlmDsaKvBSplitConfig, part: GlmDsaKvBPart, diff --git a/crates/skippy-quantize/src/gguf_writer_tests.rs b/crates/skippy-quantize/src/gguf_writer_tests.rs index ba09b3e6c..8a89f6b4a 100644 --- a/crates/skippy-quantize/src/gguf_writer_tests.rs +++ b/crates/skippy-quantize/src/gguf_writer_tests.rs @@ -260,6 +260,139 @@ fn writes_qwen_style_mtp_only_tensors_with_shared_context() { fs::remove_dir_all(root).unwrap(); } +#[test] +fn writes_inkling_mtp_streaming_transforms() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let w13 = (1_u32..=8) + .flat_map(|value| (value as f32).to_le_bytes()) + .collect::>(); + let bf16_values = [0x80, 0x3f, 0x00, 0x40, 0x40, 0x40, 0x80, 0x40]; + write_safetensor( + &root.join("model.safetensors"), + &[ + ("model.llm.embed.weight", "F32", &[1], &[1, 0, 0, 0]), + ("model.llm.embed_norm.weight", "F32", &[1], &[2, 0, 0, 0]), + ("model.llm.norm.weight", "F32", &[1], &[3, 0, 0, 0]), + ("model.llm.unembed.weight", "F32", &[1], &[4, 0, 0, 0]), + ( + "model.mtp.layers.0.embed_norm.weight", + "F32", + &[1], + &[5, 0, 0, 0], + ), + ( + "model.mtp.layers.0.transformer_block.attn.rel_logits_proj.proj", + "BF16", + &[2, 2], + &bf16_values, + ), + ( + "model.mtp.layers.0.transformer_block.attn.k_sconv.weight", + "BF16", + &[2, 1, 2], + &bf16_values, + ), + ( + "model.mtp.layers.0.transformer_block.mlp.w13_dn.weight", + "F32", + &[4, 2], + &w13, + ), + ], + ); + let output = root.join("inkling-mtp.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 9, + metadata: None, + tensor_name_map: TensorNameMap::HfToGgufWithMtp { layer_start: 66 }, + split: None, + output_type: Some(ConvertOutputType::Bf16), + tensor_selection: TensorSelection::MtpOnly { layer_start: 66 }, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert_eq!(parsed.tensor_count, 9); + let shortconv = parsed.tensor("blk.66.shortconv_k.weight"); + assert_eq!(shortconv.dims, vec![2, 2]); + assert_eq!(shortconv.ggml_type, GGML_TYPE_F32); + let rel_proj = parsed.tensor("blk.66.attn_rel_proj.weight"); + assert_eq!(rel_proj.ggml_type, GGML_TYPE_F32); + let gate = parsed.tensor("blk.66.ffn_gate.weight"); + let up = parsed.tensor("blk.66.ffn_up.weight"); + assert_eq!(gate.dims, vec![2, 2]); + assert_eq!(up.dims, vec![2, 2]); + assert_eq!(gate.ggml_type, GGML_TYPE_BF16); + assert_eq!(up.ggml_type, GGML_TYPE_BF16); + let gate_expected = [0x80, 0x3f, 0x00, 0x40, 0xa0, 0x40, 0xc0, 0x40]; + let up_expected = [0x40, 0x40, 0x80, 0x40, 0xe0, 0x40, 0x00, 0x41]; + assert_eq!( + &bytes[gate.absolute_offset..gate.absolute_offset + gate_expected.len()], + gate_expected + ); + assert_eq!( + &bytes[up.absolute_offset..up.absolute_offset + up_expected.len()], + up_expected + ); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn writes_inkling_trunk_fused_w13_streaming_transforms() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let w13 = (1_u32..=8) + .flat_map(|value| (value as f32).to_le_bytes()) + .collect::>(); + write_safetensor( + &root.join("model.safetensors"), + &[("model.layers.3.mlp.w13_dn.weight", "F32", &[4, 2], &w13)], + ); + let output = root.join("inkling-trunk.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 9, + metadata: None, + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: Some(ConvertOutputType::Bf16), + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert_eq!(parsed.tensor_count, 2); + let gate = parsed.tensor("blk.3.ffn_gate.weight"); + let up = parsed.tensor("blk.3.ffn_up.weight"); + assert_eq!(gate.dims, vec![2, 2]); + assert_eq!(up.dims, vec![2, 2]); + assert_eq!(gate.ggml_type, GGML_TYPE_BF16); + assert_eq!(up.ggml_type, GGML_TYPE_BF16); + let gate_expected = [0x80, 0x3f, 0x00, 0x40, 0xa0, 0x40, 0xc0, 0x40]; + let up_expected = [0x40, 0x40, 0x80, 0x40, 0xe0, 0x40, 0x00, 0x41]; + assert_eq!( + &bytes[gate.absolute_offset..gate.absolute_offset + gate_expected.len()], + gate_expected + ); + assert_eq!( + &bytes[up.absolute_offset..up.absolute_offset + up_expected.len()], + up_expected + ); + fs::remove_dir_all(root).unwrap(); +} + #[test] fn validates_qwen_dense_native_conversion_fixture() { let root = unique_temp_dir(); @@ -1104,6 +1237,64 @@ fn native_splits_are_byte_balanced_not_tensor_count_balanced() { fs::remove_dir_all(root).unwrap(); } +#[test] +fn recommends_enough_byte_balanced_splits_for_the_size_limit() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("a.weight", "F32", &[4], &[1; 16]), + ("b.weight", "F32", &[4], &[2; 16]), + ("c.weight", "F32", &[4], &[3; 16]), + ], + ); + + let split_count = recommended_raw_safetensors_gguf_split_count( + &root, + RawGgufWriteOptions { + buffer_size: 4, + metadata: None, + tensor_name_map: TensorNameMap::Raw, + split: None, + output_type: None, + tensor_selection: TensorSelection::All, + }, + 20, + ) + .unwrap(); + + assert_eq!(split_count, 3); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn rejects_a_size_limit_smaller_than_one_selected_tensor() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[("a.weight", "F32", &[4], &[1; 16])], + ); + + let error = recommended_raw_safetensors_gguf_split_count( + &root, + RawGgufWriteOptions { + buffer_size: 4, + metadata: None, + tensor_name_map: TensorNameMap::Raw, + split: None, + output_type: None, + tensor_selection: TensorSelection::All, + }, + 15, + ) + .unwrap_err(); + + assert!(error.to_string().contains("largest selected tensor")); + fs::remove_dir_all(root).unwrap(); +} + #[test] fn keeps_rank_one_f32_tensor_as_f32_for_bf16_output() { let root = unique_temp_dir(); @@ -1267,6 +1458,13 @@ fn read_string_array_or_skip(cursor: &mut std::io::Cursor<&[u8]>) -> Option, element_type: u32, len: u64) { for _ in 0..len { match element_type { + GGUF_TYPE_BOOL => { + let mut value = [0_u8; 1]; + cursor.read_exact(&mut value).unwrap(); + } + GGUF_TYPE_STRING => { + let _ = read_string(cursor); + } GGUF_TYPE_INT32 | GGUF_TYPE_FLOAT32 | GGUF_TYPE_UINT32 => { let _ = read_u32(cursor); } diff --git a/crates/skippy-quantize/src/hf_checkpoint.rs b/crates/skippy-quantize/src/hf_checkpoint.rs index e898523f9..96361e2c3 100644 --- a/crates/skippy-quantize/src/hf_checkpoint.rs +++ b/crates/skippy-quantize/src/hf_checkpoint.rs @@ -334,6 +334,11 @@ fn discover_safetensors(source: &Path) -> Result> { ); let mut indexed = discover_indexed_safetensors(source)?; if !indexed.is_empty() { + let mtp_sidecar = source.join("mtp.safetensors"); + if mtp_sidecar.is_file() && !indexed.contains(&mtp_sidecar) { + indexed.push(mtp_sidecar); + indexed.sort(); + } return Ok(indexed); } indexed = fs::read_dir(source) @@ -604,6 +609,36 @@ mod tests { fs::remove_dir_all(root).unwrap(); } + #[test] + fn includes_unindexed_mtp_sidecar_with_indexed_checkpoint() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("shard-a.safetensors"), + &[("a.weight", "F32", &[1], &[1, 2, 3, 4])], + ); + write_safetensor( + &root.join("mtp.safetensors"), + &[("model.mtp.layers.0.weight", "F32", &[1], &[5, 6, 7, 8])], + ); + fs::write( + root.join("model.safetensors.index.json"), + r#"{"metadata":{},"weight_map":{"a.weight":"shard-a.safetensors"}}"#, + ) + .unwrap(); + + let files = discover_safetensors(&root).unwrap(); + + assert_eq!( + files, + vec![ + root.join("mtp.safetensors"), + root.join("shard-a.safetensors") + ] + ); + fs::remove_dir_all(root).unwrap(); + } + #[test] fn streams_tensor_bytes_without_reading_neighbor_tensors() { let root = unique_temp_dir(); diff --git a/crates/skippy-quantize/src/inkling_metadata.rs b/crates/skippy-quantize/src/inkling_metadata.rs new file mode 100644 index 000000000..64683de1f --- /dev/null +++ b/crates/skippy-quantize/src/inkling_metadata.rs @@ -0,0 +1,243 @@ +use std::path::Path; + +use anyhow::{Context, Result, ensure}; +use serde_json::Value; + +use crate::gguf_writer::GgufKv; +use crate::tokenizer_metadata::push_tokenizer_metadata; + +const ARCH: &str = "inkling"; + +pub(crate) fn is_inkling_config(config: &Value) -> bool { + config.get("model_type").and_then(Value::as_str) == Some("inkling_mm_model") +} + +pub(crate) fn mtp_layer_start(config: &Value) -> Result> { + if !is_inkling_config(config) { + return Ok(None); + } + let nextn = nested_u32(config, &["mtp_config", "num_nextn_predict_layers"]).unwrap_or_default(); + if nextn == 0 { + return Ok(None); + } + required_nested_u32(config, &["text_config", "num_hidden_layers"]).map(Some) +} + +pub(crate) fn metadata( + source: &Path, + tensor_count: usize, + config: &Value, + include_mtp: bool, +) -> Result> { + let text = config + .get("text_config") + .context("Inkling config missing text_config")?; + validate_supported_variant(text, config)?; + let trunk_layers = required_u32(text, "num_hidden_layers")?; + let mtp_layers = if include_mtp { + nested_u32(config, &["mtp_config", "num_nextn_predict_layers"]).unwrap_or_default() + } else { + 0 + }; + let local_flags = local_layer_flags(text, config, trunk_layers, mtp_layers)?; + let local_kv_heads = required_u32(text, "swa_num_key_value_heads")?; + let global_kv_heads = required_u32(text, "num_key_value_heads")?; + let head_count_kv = local_flags + .iter() + .map(|local| { + if *local { + local_kv_heads + } else { + global_kv_heads + } + }) + .collect(); + let context_length = required_u32(text, "model_max_length")?; + let sliding_window = required_u32(text, "sliding_window_size")?; + let head_dim = required_u32(text, "head_dim")?; + + let mut result = vec![ + GgufKv::string("general.architecture", ARCH), + GgufKv::string("general.name", model_name(source)), + GgufKv::bool("skippy.convert.raw_safetensors", false), + GgufKv::u64("skippy.convert.tensor_count", tensor_count as u64), + GgufKv::u32("inkling.vocab_size", required_u32(text, "vocab_size")?), + GgufKv::u32("inkling.context_length", context_length), + GgufKv::u32( + "inkling.embedding_length", + required_u32(text, "hidden_size")?, + ), + GgufKv::u32("inkling.block_count", trunk_layers + mtp_layers), + GgufKv::u32( + "inkling.feed_forward_length", + required_u32(text, "dense_intermediate_size")?, + ), + GgufKv::f32( + "inkling.attention.layer_norm_rms_epsilon", + required_f32(text, "rms_norm_eps")?, + ), + GgufKv::u32( + "inkling.attention.head_count", + required_u32(text, "num_attention_heads")?, + ), + GgufKv::array_u32("inkling.attention.head_count_kv", head_count_kv), + GgufKv::u32("inkling.attention.key_length", head_dim), + GgufKv::u32("inkling.attention.value_length", head_dim), + GgufKv::u32("inkling.attention.sliding_window", sliding_window), + GgufKv::array_bool("inkling.attention.sliding_window_pattern", local_flags), + GgufKv::u32( + "inkling.expert_count", + required_u32(text, "n_routed_experts")?, + ), + GgufKv::u32( + "inkling.expert_used_count", + required_u32(text, "num_experts_per_tok")?, + ), + GgufKv::u32( + "inkling.expert_shared_count", + required_u32(text, "n_shared_experts")?, + ), + GgufKv::u32( + "inkling.expert_feed_forward_length", + required_u32(text, "intermediate_size")?, + ), + GgufKv::f32( + "inkling.expert_weights_scale", + required_f32(text, "route_scale")?, + ), + GgufKv::u32("inkling.expert_gating_func", 2), + GgufKv::u32("inkling.d_rel", required_u32(text, "d_rel")?), + GgufKv::u32("inkling.rel_extent", required_u32(text, "rel_extent")?), + GgufKv::u32("inkling.rel_extent_swa", sliding_window), + GgufKv::u32( + "inkling.shortconv_kernel", + required_u32(text, "sconv_kernel_size")?, + ), + GgufKv::u32( + "inkling.dense_block_count", + required_u32(text, "dense_mlp_idx")?, + ), + GgufKv::f32( + "inkling.logit_scale_denom", + required_f32(text, "logits_mup_width_multiplier")?, + ), + GgufKv::u32( + "inkling.log_scaling_n_floor", + required_u32(text, "log_scaling_n_floor")?, + ), + GgufKv::f32( + "inkling.log_scaling_alpha", + required_f32(text, "log_scaling_alpha")?, + ), + GgufKv::u32( + "inkling.unpadded_vocab_size", + required_u32(text, "unpadded_vocab_size")?, + ), + ]; + if mtp_layers > 0 { + result.push(GgufKv::u32("inkling.nextn_predict_layers", mtp_layers)); + } + push_tokenizer_metadata(&mut result, source, config)?; + Ok(result) +} + +fn validate_supported_variant(text: &Value, config: &Value) -> Result<()> { + for key in [ + "norm_after_topk", + "shared_expert_sink", + "use_sconv", + "use_embed_norm", + "use_gate_bias", + "use_global_scale", + ] { + ensure!( + text.get(key).and_then(Value::as_bool) == Some(true), + "unsupported Inkling {key}; native conversion requires true" + ); + } + ensure!( + text.get("gate_activation").and_then(Value::as_str) == Some("sigmoid"), + "unsupported Inkling gate_activation; native conversion requires sigmoid" + ); + ensure!( + nested_bool(config, &["mtp_config", "chain_hidden_post_norm"]) != Some(true), + "Inkling chain_hidden_post_norm=true is not supported" + ); + Ok(()) +} + +fn local_layer_flags( + text: &Value, + config: &Value, + trunk_layers: u32, + mtp_layers: u32, +) -> Result> { + let mut result = flags_from_ids(text, "local_layer_ids", trunk_layers)?; + if mtp_layers > 0 { + let mtp = config + .get("mtp_config") + .context("Inkling config missing mtp_config")?; + result.extend(flags_from_ids(mtp, "local_layer_ids", mtp_layers)?); + } + Ok(result) +} + +fn flags_from_ids(config: &Value, key: &str, count: u32) -> Result> { + let values = config + .get(key) + .and_then(Value::as_array) + .with_context(|| format!("Inkling config missing array {key}"))?; + let mut result = vec![false; count as usize]; + for value in values { + let id = value + .as_u64() + .and_then(|id| usize::try_from(id).ok()) + .with_context(|| format!("invalid Inkling {key} entry"))?; + ensure!( + id < result.len(), + "Inkling {key} entry {id} is out of range" + ); + result[id] = true; + } + Ok(result) +} + +fn required_nested_u32(config: &Value, path: &[&str]) -> Result { + nested_u32(config, path).with_context(|| format!("config missing {}", path.join("."))) +} + +fn nested_u32(config: &Value, path: &[&str]) -> Option { + path.iter() + .try_fold(config, |value, key| value.get(key))? + .as_u64() + .and_then(|value| u32::try_from(value).ok()) +} + +fn nested_bool(config: &Value, path: &[&str]) -> Option { + path.iter() + .try_fold(config, |value, key| value.get(key))? + .as_bool() +} + +fn required_u32(config: &Value, key: &str) -> Result { + config + .get(key) + .and_then(Value::as_u64) + .and_then(|value| u32::try_from(value).ok()) + .with_context(|| format!("config missing positive u32 {key}")) +} + +fn required_f32(config: &Value, key: &str) -> Result { + config + .get(key) + .and_then(Value::as_f64) + .map(|value| value as f32) + .with_context(|| format!("config missing f32 {key}")) +} + +fn model_name(source: &Path) -> &str { + source + .file_name() + .and_then(|value| value.to_str()) + .unwrap_or("Inkling") +} diff --git a/crates/skippy-quantize/src/main.rs b/crates/skippy-quantize/src/main.rs index 347d9f139..43874e830 100644 --- a/crates/skippy-quantize/src/main.rs +++ b/crates/skippy-quantize/src/main.rs @@ -12,19 +12,23 @@ mod command_reports; mod direct_convert; mod direct_quantize; mod float_convert; +mod gguf_metadata; mod gguf_template; mod gguf_writer; mod hf_checkpoint; mod imatrix; +mod inkling_metadata; mod llama_load; mod locking; mod manifest; mod memory_budget; +mod mtp_attach; mod native_convert; mod native_quantize; mod output; mod plan_convert; mod preflight; +mod projector_validate; mod quantize; mod records; mod residency; @@ -56,6 +60,7 @@ use memory_budget::{ MemoryBudgetPlanInput, MemoryPolicy, MemorySize, effective_stream_buffer_bytes, native_convert_stream_working_set_bytes, print_memory_budget_plan, }; +use mtp_attach::{ValidateMtpAttachArgs, run_validate_mtp_attach}; use native_convert::{build_native_convert_command, run_native_convert}; use native_quantize::{build_native_quantize_command, run_native_quantize}; use output::{ @@ -64,6 +69,7 @@ use output::{ }; use plan_convert::{PlanConvertArgs, run_plan_convert}; use preflight::run_job_preflight; +use projector_validate::{ValidateProjectorArgs, run_validate_projector}; use records::{WindowRunRecordInput, unix_timestamp_ms, write_window_record}; use residency::remove_dir_if_exists; use splits::{ @@ -109,6 +115,8 @@ enum Command { RunQuantWindow(RunQuantWindowArgs), VerifyJob(VerifyJobArgs), ValidateLlamaLoad(ValidateLlamaLoadArgs), + ValidateMtpAttach(ValidateMtpAttachArgs), + ValidateProjector(ValidateProjectorArgs), ValidateTensorTypes(ValidateTensorTypesArgs), ValidateSplits(ValidateSplitsArgs), } @@ -510,6 +518,8 @@ fn main() -> Result<()> { args.json, ), Command::ValidateLlamaLoad(args) => run_validate_llama_load(args), + Command::ValidateMtpAttach(args) => run_validate_mtp_attach(args), + Command::ValidateProjector(args) => run_validate_projector(args), Command::ValidateTensorTypes(args) => validate_tensor_types_command(&args.file, args.json), Command::ValidateSplits(args) => validate_splits_command( &args.root, @@ -990,9 +1000,10 @@ fn init_convert(args: InitConvertArgs) -> Result<()> { } fn convert_job(args: ConvertJobArgs) -> Result<()> { - let manifest = convert_manifest_from_args(&args.init)?; - let manifest_path = args.init.manifest.clone(); let runner = prepare_convert_runner(args.run.runner)?; + let mut manifest = convert_manifest_from_args(&args.init)?; + native_convert::apply_native_convert_split_max_size(&runner, &mut manifest)?; + let manifest_path = args.init.manifest.clone(); if args.run.preflight_only { return run_job_preflight( &manifest_path, diff --git a/crates/skippy-quantize/src/mtp_attach.rs b/crates/skippy-quantize/src/mtp_attach.rs new file mode 100644 index 000000000..698a6ec6a --- /dev/null +++ b/crates/skippy-quantize/src/mtp_attach.rs @@ -0,0 +1,204 @@ +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, ensure}; +use clap::Parser; +use serde::Serialize; +use skippy_runtime::{ + FlashAttentionType, GGML_TYPE_F16, ModelInfo, RuntimeConfig, RuntimeLoadMode, StageModel, +}; + +use crate::output::{print_json_pretty, print_success}; + +#[derive(Debug, Parser)] +pub(crate) struct ValidateMtpAttachArgs { + #[arg(long = "model", required = true)] + model_parts: Vec, + #[arg(long)] + mtp_draft: PathBuf, + #[arg(long)] + layer_count: u32, + #[arg(long)] + mtp_layer_count: Option, + #[arg(long)] + projector: Option, + #[arg(long, default_value_t = 64)] + ctx_size: u32, + #[arg(long, default_value_t = 0)] + n_gpu_layers: i32, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Serialize)] +struct MtpAttachReport { + model_parts: Vec, + mtp_draft: PathBuf, + projector: Option, + layer_count: u32, + mtp_layer_count: u32, + ctx_size: u32, + native_mtp_multimodal_feature: bool, + session_created: bool, +} + +pub(crate) fn run_validate_mtp_attach(args: ValidateMtpAttachArgs) -> Result<()> { + validate_paths(&args)?; + ensure!( + skippy_ffi::native_runtime_loaded(), + "validate-mtp-attach requires a statically linked standalone build or a loaded native runtime" + ); + let abi_features = skippy_ffi::try_abi_features() + .context("loaded native runtime does not expose Skippy ABI feature probing")?; + let native_mtp_multimodal_feature = abi_features & skippy_ffi::FEATURE_INKLING_MTP_MM != 0; + ensure!( + native_mtp_multimodal_feature, + "native runtime does not advertise Inkling multimodal MTP support" + ); + let target_config = runtime_config( + args.layer_count, + args.ctx_size, + args.n_gpu_layers, + args.projector.as_deref(), + ); + let mut target = if args.model_parts.len() == 1 { + StageModel::open(&args.model_parts[0], &target_config) + } else { + StageModel::open_from_parts(&args.model_parts, &target_config) + } + .context("open target model for MTP attach validation")?; + + let mtp_layer_count = match args.mtp_layer_count { + Some(layer_count) => layer_count, + None => infer_layer_count(&args.mtp_draft)?, + }; + ensure!( + mtp_layer_count > 0, + "--mtp-layer-count must be greater than zero" + ); + let draft_config = runtime_config(mtp_layer_count, args.ctx_size, args.n_gpu_layers, None); + target + .attach_mtp_draft_model(&args.mtp_draft, &draft_config) + .context("attach MTP draft model")?; + let _session = target + .create_session() + .context("create target session after MTP attach")?; + + let report = MtpAttachReport { + model_parts: args.model_parts, + mtp_draft: args.mtp_draft, + projector: args.projector, + layer_count: args.layer_count, + mtp_layer_count, + ctx_size: args.ctx_size, + native_mtp_multimodal_feature, + session_created: true, + }; + if args.json { + print_json_pretty(&report)?; + } else { + print_success(format!( + "MTP attach valid: parts={} layers={} mtp_layers={} projector={} session_created=true", + report.model_parts.len(), + report.layer_count, + report.mtp_layer_count, + report + .projector + .as_deref() + .map_or_else(|| "none".to_string(), |path| path.display().to_string()) + )); + } + Ok(()) +} + +fn validate_paths(args: &ValidateMtpAttachArgs) -> Result<()> { + ensure!( + args.layer_count > 0, + "--layer-count must be greater than zero" + ); + ensure!(args.ctx_size > 0, "--ctx-size must be greater than zero"); + for path in &args.model_parts { + ensure!( + path.is_file(), + "model part does not exist: {}", + path.display() + ); + } + ensure!( + args.mtp_draft.is_file(), + "MTP draft does not exist: {}", + args.mtp_draft.display() + ); + if let Some(projector) = &args.projector { + ensure!( + projector.is_file(), + "projector does not exist: {}", + projector.display() + ); + } + Ok(()) +} + +fn infer_layer_count(path: &Path) -> Result { + ModelInfo::open(path) + .with_context(|| format!("open MTP model info {}", path.display()))? + .tensors()? + .into_iter() + .filter_map(|tensor| tensor.layer_index) + .max() + .and_then(|index| index.checked_add(1)) + .context("MTP draft contains no layer-indexed tensors") +} + +fn runtime_config( + layer_end: u32, + ctx_size: u32, + n_gpu_layers: i32, + projector: Option<&Path>, +) -> RuntimeConfig { + RuntimeConfig { + stage_index: 0, + layer_start: 0, + layer_end, + ctx_size, + lane_count: 1, + n_batch: Some(ctx_size), + n_ubatch: Some(ctx_size), + n_threads: None, + n_threads_batch: None, + n_gpu_layers, + mmap: Some(true), + mlock: false, + selected_backend_device: None, + cache_type_k: GGML_TYPE_F16, + cache_type_v: GGML_TYPE_F16, + flash_attn_type: FlashAttentionType::Auto, + load_mode: RuntimeLoadMode::RuntimeSlice, + projector_path: projector.map(|path| path.display().to_string()), + include_embeddings: true, + include_output: true, + filter_tensors_on_load: false, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn attach_probe_uses_small_mmap_runtime_config() { + let config = runtime_config(66, 64, 0, Some(Path::new("/models/mmproj.gguf"))); + + assert_eq!(config.layer_start, 0); + assert_eq!(config.layer_end, 66); + assert_eq!(config.ctx_size, 64); + assert_eq!(config.n_batch, Some(64)); + assert_eq!(config.n_ubatch, Some(64)); + assert_eq!(config.mmap, Some(true)); + assert_eq!( + config.projector_path.as_deref(), + Some("/models/mmproj.gguf") + ); + assert!(config.include_embeddings); + assert!(config.include_output); + } +} diff --git a/crates/skippy-quantize/src/native_convert.rs b/crates/skippy-quantize/src/native_convert.rs index 454abd334..3b9831259 100644 --- a/crates/skippy-quantize/src/native_convert.rs +++ b/crates/skippy-quantize/src/native_convert.rs @@ -8,12 +8,14 @@ use crate::gguf_template::{ MetadataOptions, metadata_from_hf_config_with_options, mtp_layer_start_from_hf_config, }; use crate::gguf_writer::{ - GgufSplit, RawGgufWriteOptions, TensorSelection, write_raw_safetensors_gguf, + GgufSplit, RawGgufWriteOptions, TensorSelection, recommended_raw_safetensors_gguf_split_count, + write_raw_safetensors_gguf, }; use crate::hf_checkpoint::{inspect_hf_checkpoint, resolve_auto_output_type}; use crate::manifest::Manifest; use crate::memory_budget::{ - effective_stream_buffer_bytes, enforce_memory_budget, native_convert_stream_working_set_bytes, + MemorySize, effective_stream_buffer_bytes, enforce_memory_budget, + native_convert_stream_working_set_bytes, }; use crate::output::{format_bytes, print_info}; use crate::splits::SplitWindow; @@ -51,6 +53,54 @@ pub(crate) fn build_native_convert_command( command } +pub(crate) fn apply_native_convert_split_max_size( + runner: &ConvertRunnerArgs, + manifest: &mut Manifest, +) -> Result<()> { + let max_size = runner + .split_max_size + .parse::() + .map_err(anyhow::Error::msg)?; + if max_size.bytes() == 0 { + return Ok(()); + } + + let output_type = manifest + .output_type + .map(|kind| resolve_auto_output_type(&manifest.source, kind)) + .transpose()?; + let plan = inspect_hf_checkpoint(&manifest.source, runner.max_memory, 0.60)?; + let mtp_layer_start = mtp_layer_start_from_hf_config(&manifest.source)?; + let recommended = recommended_raw_safetensors_gguf_split_count( + &manifest.source, + RawGgufWriteOptions { + buffer_size: runner.stream_buffer_bytes, + metadata: Some(metadata_from_hf_config_with_options( + &manifest.source, + plan.tensor_count, + MetadataOptions { + include_mtp: !runner.no_mtp, + }, + )?), + tensor_name_map: native_tensor_name_map(mtp_layer_start), + split: None, + output_type, + tensor_selection: native_tensor_selection(runner, mtp_layer_start)?, + }, + max_size.bytes(), + )?; + let requested = manifest.expected_splits; + manifest.expected_splits = requested.max(recommended); + if manifest.expected_splits != requested { + print_info(format!( + "Raised native conversion split count from {requested} to {} to honor --split-max-size {}", + manifest.expected_splits, + format_bytes(max_size.bytes()) + )); + } + Ok(()) +} + pub(crate) fn run_native_convert( runner: &ConvertRunnerArgs, manifest: &Manifest, diff --git a/crates/skippy-quantize/src/projector_validate.rs b/crates/skippy-quantize/src/projector_validate.rs new file mode 100644 index 000000000..9e4804585 --- /dev/null +++ b/crates/skippy-quantize/src/projector_validate.rs @@ -0,0 +1,68 @@ +use std::ffi::CString; +use std::path::PathBuf; + +use anyhow::{Context, Result, ensure}; +use clap::Parser; +use serde::Serialize; + +use crate::output::{print_json_pretty, print_success}; + +#[derive(Debug, Parser)] +pub(crate) struct ValidateProjectorArgs { + #[arg(long)] + projector: PathBuf, + #[arg(long = "no-warmup", action = clap::ArgAction::SetFalse, default_value_t = true)] + warmup: bool, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Serialize)] +struct ProjectorReport { + projector: PathBuf, + warmup: bool, + loaded: bool, +} + +pub(crate) fn run_validate_projector(args: ValidateProjectorArgs) -> Result<()> { + ensure!( + args.projector.is_file(), + "projector does not exist: {}", + args.projector.display() + ); + ensure!( + skippy_ffi::native_runtime_loaded(), + "validate-projector requires a statically linked standalone build or a loaded native runtime" + ); + let projector = CString::new(args.projector.to_string_lossy().as_bytes()) + .context("projector path contains an interior NUL byte")?; + let mut params = unsafe { skippy_ffi::mtmd_context_params_default() }; + params.use_gpu = false; + params.warmup = args.warmup; + params.progress_callback = None; + params.progress_callback_user_data = std::ptr::null_mut(); + let raw = + unsafe { skippy_ffi::mtmd_init_from_file(projector.as_ptr(), std::ptr::null(), params) }; + ensure!( + !raw.is_null(), + "failed to load multimodal projector {}", + args.projector.display() + ); + unsafe { skippy_ffi::mtmd_free(raw) }; + + let report = ProjectorReport { + projector: args.projector, + warmup: args.warmup, + loaded: true, + }; + if args.json { + print_json_pretty(&report)?; + } else { + print_success(format!( + "projector valid: path={} warmup={} loaded=true", + report.projector.display(), + report.warmup + )); + } + Ok(()) +} diff --git a/crates/skippy-quantize/src/tensor_map.rs b/crates/skippy-quantize/src/tensor_map.rs index 3fa4b3920..26013add1 100644 --- a/crates/skippy-quantize/src/tensor_map.rs +++ b/crates/skippy-quantize/src/tensor_map.rs @@ -18,23 +18,31 @@ impl TensorNameMap { } fn map_hf_to_gguf(name: &str, mtp_layer_start: Option) -> Result { + if let Some(layer_start) = mtp_layer_start + && let Some(normalized) = normalize_inkling_mtp_source_name(name, layer_start)? + { + return map_hf_to_gguf(&normalized, mtp_layer_start); + } if let Some(layer_start) = mtp_layer_start && let Some(normalized) = normalize_qwen_mtp_source_name(name, layer_start)? { return map_hf_to_gguf(&normalized, mtp_layer_start); } - if name == "model.embed_tokens.weight" { + if matches!(name, "model.embed_tokens.weight" | "model.llm.embed.weight") { return Ok("token_embd.weight".to_string()); } if name == "embed_tokens.weight" { return Ok("token_embd.weight".to_string()); } - if name == "lm_head.weight" { + if matches!(name, "lm_head.weight" | "model.llm.unembed.weight") { return Ok("output.weight".to_string()); } - if name == "model.norm.weight" { + if matches!(name, "model.norm.weight" | "model.llm.norm.weight") { return Ok("output_norm.weight".to_string()); } + if name == "model.llm.embed_norm.weight" { + return Ok("token_embd_norm.weight".to_string()); + } if name == "norm.weight" { return Ok("output_norm.weight".to_string()); } @@ -48,7 +56,8 @@ fn map_hf_to_gguf(name: &str, mtp_layer_start: Option) -> Result { } pub(crate) fn is_mtp_source_tensor(name: &str) -> bool { - is_qwen_mtp_source_tensor(name) + is_inkling_mtp_source_tensor(name) + || is_qwen_mtp_source_tensor(name) || map_mtp_source_tensor(name).is_ok_and(|mapped| mapped.is_some()) } @@ -73,9 +82,57 @@ pub(crate) fn is_shared_mtp_context_tensor(name: &str) -> bool { | "model.norm.weight" | "norm.weight" | "lm_head.weight" + | "model.llm.embed.weight" + | "model.llm.embed_norm.weight" + | "model.llm.norm.weight" + | "model.llm.unembed.weight" ) } +pub(crate) fn is_inkling_fused_w13(name: &str) -> bool { + (name.starts_with("model.layers.") || name.starts_with("model.mtp.layers.")) + && name.ends_with(".mlp.w13_dn.weight") +} + +pub(crate) fn inkling_mtp_depth(name: &str) -> Result> { + let Some(rest) = name.strip_prefix("model.mtp.layers.") else { + return Ok(None); + }; + let Some((depth, _)) = rest.split_once('.') else { + bail!("malformed Inkling MTP tensor name {name}"); + }; + depth + .parse::() + .map(Some) + .map_err(|err| anyhow!("malformed Inkling MTP depth in {name}: {err}")) +} + +fn is_inkling_mtp_source_tensor(name: &str) -> bool { + name.starts_with("model.mtp.layers.") +} + +fn normalize_inkling_mtp_source_name(name: &str, layer_start: u32) -> Result> { + let Some(rest) = name.strip_prefix("model.mtp.layers.") else { + return Ok(None); + }; + let Some((depth, suffix)) = rest.split_once('.') else { + bail!("malformed Inkling MTP tensor name {name}"); + }; + let depth = depth + .parse::() + .map_err(|err| anyhow!("malformed Inkling MTP depth in {name}: {err}"))?; + let layer = layer_start + .checked_add(depth) + .ok_or_else(|| anyhow!("Inkling MTP layer id overflow for {name}"))?; + let suffix = match suffix { + "embed_norm.weight" => "enorm.weight", + "hidden_norm.weight" => "hnorm.weight", + "input_proj.weight" => "eh_proj.weight", + value => value.strip_prefix("transformer_block.").unwrap_or(value), + }; + Ok(Some(format!("model.layers.{layer}.{suffix}"))) +} + fn map_mtp_source_tensor(name: &str) -> Result> { let mapped = match name { "pre_projection" | "pre_projection.weight" => "nextn.pre_projection.weight".to_string(), @@ -182,6 +239,29 @@ impl<'a> HfLayerTensor<'a> { "self_attn.kv_a_proj_with_mqa.weight" => Ok(format!("blk.{bid}.attn_kv_a_mqa.weight")), "self_attn.kv_b_proj.weight" => Ok(format!("blk.{bid}.attn_kv_b.weight")), "self_attn.kv_a_layernorm.weight" => Ok(format!("blk.{bid}.attn_kv_a_norm.weight")), + "attn_norm.weight" => Ok(format!("blk.{bid}.attn_norm.weight")), + "attn.wq_du.weight" => Ok(format!("blk.{bid}.attn_q.weight")), + "attn.wk_dv.weight" => Ok(format!("blk.{bid}.attn_k.weight")), + "attn.wv_dv.weight" => Ok(format!("blk.{bid}.attn_v.weight")), + "attn.wr_du.weight" => Ok(format!("blk.{bid}.attn_r.weight")), + "attn.wo_ud.weight" => Ok(format!("blk.{bid}.attn_output.weight")), + "attn.q_norm.weight" => Ok(format!("blk.{bid}.attn_q_norm.weight")), + "attn.k_norm.weight" => Ok(format!("blk.{bid}.attn_k_norm.weight")), + "attn.rel_logits_proj.proj" | "attn.rel_logits_proj.weight" => { + Ok(format!("blk.{bid}.attn_rel_proj.weight")) + } + "attn.k_sconv.weight" => Ok(format!("blk.{bid}.shortconv_k.weight")), + "attn.v_sconv.weight" => Ok(format!("blk.{bid}.shortconv_v.weight")), + "attn_sconv.weight" => Ok(format!("blk.{bid}.shortconv_attn.weight")), + "mlp_sconv.weight" => Ok(format!("blk.{bid}.shortconv_mlp.weight")), + "mlp_norm.weight" => Ok(format!("blk.{bid}.ffn_norm.weight")), + "mlp.w2_md.weight" => Ok(format!("blk.{bid}.ffn_down.weight")), + "mlp.global_scale" | "mlp.global_scale.weight" => { + Ok(format!("blk.{bid}.ffn_gscale.weight")) + } + "mlp.w13_dn.weight" => { + bail!("Inkling fused w13 tensor requires streaming deinterleave") + } "self_attn.indexer.k_norm.weight" => Ok(format!("blk.{bid}.indexer.k_norm.weight")), "self_attn.indexer.k_norm.bias" => Ok(format!("blk.{bid}.indexer.k_norm.bias")), "self_attn.indexer.weights_proj.weight" => Ok(format!("blk.{bid}.indexer.proj.weight")), @@ -409,4 +489,57 @@ mod tests { "blk.33.attn_q.weight" ); } + + #[test] + fn recognizes_and_maps_inkling_mtp_tensors() { + let map = TensorNameMap::HfToGgufWithMtp { layer_start: 66 }; + for (source, expected) in [ + ( + "model.mtp.layers.0.embed_norm.weight", + "blk.66.nextn.enorm.weight", + ), + ( + "model.mtp.layers.1.hidden_norm.weight", + "blk.67.nextn.hnorm.weight", + ), + ( + "model.mtp.layers.2.input_proj.weight", + "blk.68.nextn.eh_proj.weight", + ), + ( + "model.mtp.layers.3.transformer_block.attn.wq_du.weight", + "blk.69.attn_q.weight", + ), + ( + "model.mtp.layers.4.transformer_block.attn.rel_logits_proj.proj", + "blk.70.attn_rel_proj.weight", + ), + ( + "model.mtp.layers.7.transformer_block.mlp.global_scale", + "blk.73.ffn_gscale.weight", + ), + ] { + assert!(is_mtp_source_tensor(source)); + assert_eq!(map.map_tensor_name(source).unwrap(), expected); + } + assert_eq!( + map.map_tensor_name("model.llm.embed_norm.weight").unwrap(), + "token_embd_norm.weight" + ); + assert!(is_shared_mtp_context_tensor("model.llm.unembed.weight")); + assert!(is_inkling_fused_w13( + "model.mtp.layers.0.transformer_block.mlp.w13_dn.weight" + )); + } + + #[test] + fn rejects_inkling_mtp_layer_overflow() { + let error = TensorNameMap::HfToGgufWithMtp { + layer_start: u32::MAX, + } + .map_tensor_name("model.mtp.layers.1.embed_norm.weight") + .unwrap_err(); + + assert!(error.to_string().contains("MTP layer id overflow")); + } } diff --git a/crates/skippy-quantize/src/tokenizer_metadata.rs b/crates/skippy-quantize/src/tokenizer_metadata.rs index a6377f9cc..05c5877c5 100644 --- a/crates/skippy-quantize/src/tokenizer_metadata.rs +++ b/crates/skippy-quantize/src/tokenizer_metadata.rs @@ -8,6 +8,7 @@ use serde_json::Value; use crate::gguf_writer::GgufKv; const TOKEN_TYPE_NORMAL: i32 = 1; +const TOKEN_TYPE_UNUSED: i32 = 5; const TOKEN_TYPE_CONTROL: i32 = 3; pub(crate) fn push_tokenizer_metadata( @@ -64,7 +65,7 @@ struct BpeVocabMetadata { added_tokens: BTreeMap, } -fn read_byte_level_bpe(tokenizer: &Value, _config: &Value) -> Result { +fn read_byte_level_bpe(tokenizer: &Value, config: &Value) -> Result { let model = tokenizer .get("model") .and_then(Value::as_object) @@ -86,7 +87,16 @@ fn read_byte_level_bpe(tokenizer: &Value, _config: &Value) -> Result= tokenizer_vocab_size, + "configured vocab_size {vocab_size} is smaller than tokenizer vocab size {tokenizer_vocab_size}" + ); let mut tokens = vec![String::new(); vocab_size]; let mut token_types = vec![TOKEN_TYPE_NORMAL; vocab_size]; let mut scores = vec![0.0_f32; vocab_size]; @@ -115,6 +125,21 @@ fn read_byte_level_bpe(tokenizer: &Value, _config: &Value) -> Result Result<&'static str> { if model_type.starts_with("qwen2") || model_type.starts_with("qwen3") { return Ok("qwen2"); } + if model_type == "inkling_mm_model" { + return Ok("inkling"); + } if matches!(model_type, "llama" | "mistral") { return Ok("llama-bpe"); } @@ -256,6 +284,19 @@ fn push_special_token_ids( .get("model_type") .and_then(Value::as_str) .unwrap_or_default(); + if model_type == "inkling_mm_model" { + let mut eos = config + .get("eos_token_id") + .and_then(u32_value) + .unwrap_or(200006); + if eos < 199998 { + eos = 200006; + } + metadata.push(GgufKv::u32("tokenizer.ggml.eos_token_id", eos)); + metadata.push(GgufKv::u32("tokenizer.ggml.bos_token_id", eos)); + metadata.push(GgufKv::bool("tokenizer.ggml.add_bos_token", false)); + return; + } if model_type.starts_with("glm") { push_added_token_id( metadata, @@ -297,6 +338,13 @@ fn push_special_token_ids( push_tokenizer_bool(metadata, tokenizer_config, "add_eos_token"); } +fn inkling_text_u32(config: &Value, key: &str) -> Option { + if config.get("model_type").and_then(Value::as_str) != Some("inkling_mm_model") { + return None; + } + config.get("text_config")?.get(key).and_then(u32_value) +} + fn push_added_token_id( metadata: &mut Vec, key: &str, @@ -470,6 +518,40 @@ mod tests { assert_eq!(metadata.tokens[2], "<|endoftext|>"); } + #[test] + fn fills_inkling_padded_vocab_as_unused_tokens() { + let tokenizer: Value = serde_json::from_str( + r#"{ + "model": { + "type": "BPE", + "vocab": {"a": 0, "b": 1, "<|end|>": 2}, + "merges": ["a b"] + }, + "decoder": {"type": "ByteLevel"}, + "added_tokens": [ + {"id": 2, "content": "<|end|>", "special": true} + ] + }"#, + ) + .unwrap(); + let config: Value = serde_json::from_str( + r#"{ + "model_type": "inkling_mm_model", + "eos_token_id": 2, + "text_config": {"vocab_size": 6, "unpadded_vocab_size": 3} + }"#, + ) + .unwrap(); + + let metadata = read_byte_level_bpe(&tokenizer, &config).unwrap(); + + assert_eq!(metadata.tokens.len(), 6); + assert_eq!(metadata.tokens[3], "[PAD3]"); + assert_eq!(metadata.token_types[2], TOKEN_TYPE_CONTROL); + assert_eq!(metadata.token_types[3], TOKEN_TYPE_UNUSED); + assert_eq!(tokenizer_pre(&config).unwrap(), "inkling"); + } + #[test] fn builds_llama_byte_level_bpe_tokenizer_metadata() { let root = unique_temp_dir(); diff --git a/crates/skippy-quantize/src/types.rs b/crates/skippy-quantize/src/types.rs index 41671fa9e..3dce8433b 100644 --- a/crates/skippy-quantize/src/types.rs +++ b/crates/skippy-quantize/src/types.rs @@ -40,6 +40,7 @@ pub enum JobKind { #[serde(rename_all = "SCREAMING_SNAKE_CASE")] pub enum QuantType { Q1_0, + Q2_0, Q2K, Q2KS, Q3K, @@ -127,6 +128,7 @@ impl FromStr for QuantType { let normalized = normalize_type_name(raw); let quant = match normalized.as_str() { "Q10" => Self::Q1_0, + "Q20" => Self::Q2_0, "Q2K" => Self::Q2K, "Q2KS" => Self::Q2KS, "Q3K" => Self::Q3K, @@ -173,6 +175,7 @@ impl FromStr for QuantType { impl QuantType { pub const ALL: &'static [Self] = &[ Self::Q1_0, + Self::Q2_0, Self::Q2K, Self::Q2KS, Self::Q3K, @@ -215,6 +218,7 @@ impl QuantType { pub fn as_llama_name(self) -> &'static str { match self { Self::Q1_0 => "Q1_0", + Self::Q2_0 => "Q2_0", Self::Q2K => "Q2_K", Self::Q2KS => "Q2_K_S", Self::Q3K => "Q3_K", @@ -291,6 +295,7 @@ impl QuantType { 37 => Some(Self::TQ2_0), 38 => Some(Self::Mxfp4Moe), 40 => Some(Self::Q1_0), + 41 => Some(Self::Q2_0), _ => None, } } @@ -298,6 +303,7 @@ impl QuantType { pub fn as_llama_file_type(self) -> llama_quant_ffi::LlamaFileType { match self { Self::Q1_0 => llama_quant_ffi::LlamaFileType::MostlyQ1_0, + Self::Q2_0 => llama_quant_ffi::LlamaFileType::MostlyQ2_0, Self::Q2K => llama_quant_ffi::LlamaFileType::MostlyQ2K, Self::Q2KS => llama_quant_ffi::LlamaFileType::MostlyQ2KS, Self::Q3K | Self::Q3KM => llama_quant_ffi::LlamaFileType::MostlyQ3KM, @@ -735,7 +741,9 @@ mod tests { } fn pinned_llama_quant_option_names() -> Vec { - const UNSUPPORTED_FFI_QUANT_MODES: &[&str] = &["Q2_0"]; + // Q2_0 gained FFI support with the Inkling Q2 work; no llama-quantize + // modes are currently unsupported by the local catalog. + const UNSUPPORTED_FFI_QUANT_MODES: &[&str] = &[]; let quantize_cpp = repo_root().join(".deps/llama.cpp/tools/quantize/quantize.cpp"); let source = fs::read_to_string(&quantize_cpp) .unwrap_or_else(|err| panic!("read {}: {err}", quantize_cpp.display())); diff --git a/crates/skippy-runtime/src/activation.rs b/crates/skippy-runtime/src/activation.rs index ace1e9571..a50994b76 100644 --- a/crates/skippy-runtime/src/activation.rs +++ b/crates/skippy-runtime/src/activation.rs @@ -594,6 +594,9 @@ impl StageSession { self.verify_tokens_frame_sampled(token_ids, None, input, output_capacity, 0) } + /// Verifies a frame-backed speculative window. + /// + /// Retire the exact checkpoint after full acceptance, or trim after rejection. pub fn verify_tokens_frame_sampled( &mut self, token_ids: &[i32], diff --git a/crates/skippy-runtime/src/media.rs b/crates/skippy-runtime/src/media.rs index a9a3470a3..93b211963 100644 --- a/crates/skippy-runtime/src/media.rs +++ b/crates/skippy-runtime/src/media.rs @@ -385,6 +385,20 @@ impl StageModel { "multimodal chunk {index} has {chunk_tokens} tokens, exceeding n_batch {n_batch}; increase n_batch for staged media prefill" )); } + let chunk_token_ids = if chunk_type == skippy_ffi::MtmdInputChunkType::Text { + let mut text_token_count = 0usize; + let text_tokens = unsafe { + skippy_ffi::mtmd_input_chunk_get_tokens_text(chunk, &mut text_token_count) + }; + if text_tokens.is_null() || text_token_count != chunk_tokens { + return Err(anyhow!( + "multimodal text chunk {index} token view did not match its declared token count" + )); + } + unsafe { std::slice::from_raw_parts(text_tokens, text_token_count) }.to_vec() + } else { + Vec::new() + }; let chunk_positions = if use_mrope { let chunk_positions = match chunk_type { skippy_ffi::MtmdInputChunkType::Image => { @@ -481,6 +495,7 @@ impl StageModel { output_payload.extend_from_slice(&frame.payload); chunk_frames.push(MediaPrefillChunkFrame { token_count: chunk_tokens, + tokens: chunk_token_ids, positions: chunk_positions, output: frame, }); diff --git a/crates/skippy-runtime/src/package.rs b/crates/skippy-runtime/src/package.rs index 69a13be16..efbcc6eea 100644 --- a/crates/skippy-runtime/src/package.rs +++ b/crates/skippy-runtime/src/package.rs @@ -104,6 +104,7 @@ pub struct PackageWindowPolicyInfo { pub initial_window: u32, pub min_window: u32, pub max_window: u32, + pub pipeline_depth: Option, } #[derive(Debug, Clone)] @@ -288,6 +289,8 @@ struct PackageWindowPolicy { initial_window: u32, min_window: u32, max_window: u32, + #[serde(default)] + pipeline_depth: Option, } #[derive(Debug, Deserialize)] @@ -651,6 +654,7 @@ fn package_speculative_strategy_info( initial_window: window.initial_window, min_window: window.min_window, max_window: window.max_window, + pipeline_depth: window.pipeline_depth, }), proposer: strategy.proposer, primary: strategy.primary, diff --git a/crates/skippy-runtime/src/runtime_events.rs b/crates/skippy-runtime/src/runtime_events.rs index d36378125..424cce471 100644 --- a/crates/skippy-runtime/src/runtime_events.rs +++ b/crates/skippy-runtime/src/runtime_events.rs @@ -355,7 +355,7 @@ fn abi_features_bitmask() -> Option { } #[cfg(not(feature = "dynamic-native-runtime"))] { - Some(unsafe { skippy_ffi::skippy_abi_features() }) + Some(skippy_ffi::abi_features()) } } diff --git a/crates/skippy-runtime/src/session.rs b/crates/skippy-runtime/src/session.rs index 550309cc5..5bf689168 100644 --- a/crates/skippy-runtime/src/session.rs +++ b/crates/skippy-runtime/src/session.rs @@ -82,6 +82,20 @@ impl StageSession { ensure_ok(status, error) } + /// Retires the exact recovery checkpoint for a fully accepted verify window. + pub fn retire_verify_checkpoint(&mut self, token_start: u64, token_count: u64) -> Result<()> { + let mut error = ptr::null_mut(); + let status = unsafe { + skippy_ffi::skippy_retire_verify_checkpoint( + self.raw, + token_start, + token_count, + &mut error, + ) + }; + ensure_ok(status, error) + } + pub fn trim_session(&mut self, token_count: u64) -> Result<()> { let mut error = ptr::null_mut(); let status = unsafe { skippy_ffi::skippy_trim_session(self.raw, token_count, &mut error) }; @@ -316,6 +330,10 @@ impl StageSession { Ok(window.into()) } + /// Verifies a speculative window and advances the session. + /// + /// Call [`Self::retire_verify_checkpoint`] when the complete window is accepted, + /// or [`Self::trim_session`] when any suffix is rejected. pub fn verify_tokens(&mut self, token_ids: &[i32]) -> Result> { if token_ids.is_empty() { return Ok(Vec::new()); diff --git a/crates/skippy-runtime/src/types.rs b/crates/skippy-runtime/src/types.rs index 0c0057d97..a49ae5441 100644 --- a/crates/skippy-runtime/src/types.rs +++ b/crates/skippy-runtime/src/types.rs @@ -233,6 +233,7 @@ pub struct MediaPrefill { #[derive(Debug, Clone, PartialEq, Eq)] pub struct MediaPrefillChunkFrame { pub token_count: usize, + pub tokens: Vec, pub positions: Vec, pub output: ActivationFrame, } diff --git a/crates/skippy-server/README.md b/crates/skippy-server/README.md index 6bd26ec87..d3b4eb0dc 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -15,7 +15,7 @@ mesh/openai-frontend; diagnostic and benchmark clients may connect directly to the first stage. The full request/reply path is tip-to-tip: token IDs enter at the driver-facing -tip, and activations flow through the stage chain. Stage protocol generation 3 +tip, and activations flow through the stage chain. Stage protocol generation 4 is a compatibility-breaking contract: prediction-bearing replies return directly from the final/readout tip to the driver-facing stage instead of being relayed back through intermediate stages. Middle-out is the prefill optimization @@ -80,10 +80,11 @@ unload or replan. ## Notes - `serve-binary` is the tuned binary stage-to-stage path. -- `serve-binary` participates in the breaking generation-3 stage protocol. - Stage compatibility requires `stage-generation-3`; direct prediction return is - part of that generation's contract, so older chained-reply peers are rejected - during split planning instead of being mixed into a generation-3 topology. +- `serve-binary` participates in the breaking generation-4 stage protocol. + Stage compatibility requires `stage-generation-4`; direct prediction return and + exact verify-checkpoint retirement are part of that generation's contract, so + older peers are rejected during split planning instead of being mixed into a + generation-4 topology. - `serve-binary` accepts upstream protocol connections concurrently. Model execution remains serialized by the per-process runtime lock, but readiness, abandoned, or broken connections do not monopolize the listener and block the @@ -99,7 +100,7 @@ unload or replan. `/v1/completions` using the shared `openai-frontend` crate for a local final/single-stage config with no downstream peer. Split serving uses embedded stage-0 OpenAI serving from `serve-binary --openai-bind-addr` because - generation-3 prediction returns flow directly from the final stage to stage 0. + generation-4 prediction returns flow directly from the final stage to stage 0. The older standalone `serve-openai --first-stage-addr` adapter is no longer supported. `--model-id` is the exact served model id to advertise and accept, for example `org/repo:Q4_K_M`; it is not parsed as stage topology. diff --git a/crates/skippy-server/src/binary_transport/binary_messaging.rs b/crates/skippy-server/src/binary_transport/binary_messaging.rs index a9f1dcdf4..987cffbdf 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -11,8 +11,8 @@ use std::{ }; use super::stage_execution::{ - consume_optional_client_ready_hello, prepare_binary_stage_connection, - take_warm_or_connect_downstream, warm_downstream_preconnect_enabled, + consume_optional_client_ready_hello, prepare_binary_stage_connection, take_ready_downstream, + warm_downstream_preconnect_enabled, }; use super::{ decode_batcher::DecodeFrameBatcher, @@ -34,8 +34,12 @@ use skippy_protocol::binary::{WireMessageKind, read_stage_message, send_ready}; pub(in crate::binary_transport) mod async_forwarder; mod connection; +mod control_messages; +mod message_receive; mod prefill_recording; pub(in crate::binary_transport) mod reply; +mod session_lifecycle; +mod session_tracker; mod summary; mod telemetry; @@ -162,7 +166,7 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc) -> R speculative: openai_options.speculative.clone(), native_mtp_enabled: native_mtp_enabled && openai_options.speculative.native_mtp.enabled, - native_mtp_draft_model_path: None, + native_mtp_draft_model_path: openai_options.native_mtp_draft_model_path, native_mtp_max_tokens: openai_options.native_mtp_max_tokens, native_mtp_min_tokens: openai_options.native_mtp_min_tokens, activation_width, @@ -243,7 +247,7 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc) -> R } return prediction_return_sinks.insert_opened_sink(first_message, upstream); } - let downstream = take_warm_or_connect_downstream( + let downstream = take_ready_downstream( &config, &warm_downstream, downstream_connect_timeout_secs, diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs b/crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs index d334f6c88..a64b50751 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs @@ -193,7 +193,82 @@ impl AsyncForwardReceipt { #[cfg(test)] mod tests { + use std::net::TcpListener; + + use skippy_protocol::binary::{StageStateHeader, WireMessageKind, read_stage_message}; + use super::*; + use crate::binary_transport::stage_execution::prefix_cache_test_config; + use crate::telemetry::TelemetryLevel; + + fn message(kind: WireMessageKind, pos_start: i32) -> StageWireMessage { + StageWireMessage { + kind, + pos_start, + token_count: if kind == WireMessageKind::RetireVerifyWindow { + 4 + } else { + 0 + }, + state: StageStateHeader::new(kind, WireActivationDType::F32), + request_id: 1, + session_id: 2, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + } + } + + #[test] + fn retirement_receipt_orders_all_prior_verify_writes() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let client = TcpStream::connect(address).unwrap(); + let (mut server, _) = listener.accept().unwrap(); + let telemetry = Telemetry::new(None, 1, prefix_cache_test_config(), TelemetryLevel::Off); + let mut forwarder = AsyncForwarder::new(&client, telemetry, 3).unwrap(); + let condition = WireCondition::new(0.0, None).unwrap(); + + forwarder + .send( + message(WireMessageKind::VerifyWindow, 10), + WireActivationDType::F32, + condition, + BTreeMap::new(), + ) + .unwrap(); + forwarder + .send( + message(WireMessageKind::VerifyWindow, 14), + WireActivationDType::F32, + condition, + BTreeMap::new(), + ) + .unwrap(); + forwarder + .send_tracked( + message(WireMessageKind::RetireVerifyWindow, 10), + WireActivationDType::F32, + condition, + BTreeMap::new(), + ) + .unwrap() + .finish() + .unwrap(); + + let first = read_stage_message(&mut server, 1).unwrap(); + let second = read_stage_message(&mut server, 1).unwrap(); + let retire = read_stage_message(&mut server, 1).unwrap(); + assert_eq!(first.kind, WireMessageKind::VerifyWindow); + assert_eq!(first.pos_start, 10); + assert_eq!(second.kind, WireMessageKind::VerifyWindow); + assert_eq!(second.pos_start, 14); + assert_eq!(retire.kind, WireMessageKind::RetireVerifyWindow); + assert_eq!(retire.pos_start, 10); + } #[test] fn forward_receipt_has_a_terminal_wait_bound() { diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs index f0fc56b4c..3c809768e 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -1,11 +1,21 @@ use super::async_forwarder::AsyncForwarder; -use super::reply::drain_deferred_prefill_replies; +use super::control_messages::{ + handle_generation_control, handle_prefix_cache_control, handle_session_control, handle_stop, + handle_verify_retirement, +}; +use super::message_receive::{next_connection_session_id, receive_next_message}; +use super::reply::reply_window_for_message; use super::reply::send_stage_reply; -use super::reply::{configure_prediction_return_stream, reply_window_for_message}; +use super::session_lifecycle::align_session_to_message; +use super::session_tracker::{ + ConnectionSessionTracker, combine_connection_and_cleanup_results, + release_tracked_connection_sessions, +}; use super::summary::BinaryMessageObservation; use super::summary::BinaryRequestSummary; use super::telemetry::UpstreamReplyWriteSpan; use super::telemetry::{ + BinaryMessageTiming, emit_binary_message_received, emit_binary_message_timing, emit_upstream_reply_write_span, insert_runtime_session_stats, record_prefill_edge_transport, record_verify_window_timing, }; @@ -16,23 +26,19 @@ use crate::binary_transport::binary_kv::accumulate_prefill_tokens; use crate::binary_transport::binary_kv::add_binary_record_stats; use crate::binary_transport::binary_kv::emit_binary_proactive_eviction; use crate::binary_transport::binary_kv::maybe_lookup_binary_prefill; -use crate::binary_transport::binary_kv::maybe_prefix_cache_control; use crate::binary_transport::binary_kv::maybe_record_binary_full_prefill; use crate::binary_transport::direct_return; use crate::binary_transport::direct_return::PredictionReturnSinks; use crate::binary_transport::forwarded_stage_message_timed; use crate::binary_transport::kv_eviction::binary_proactive_eviction_plan; use crate::binary_transport::kv_eviction::evict_binary_resident_prefix_for_decode; -use crate::binary_transport::restore_prefill_decode::handle_binary_restore_prefill_decode_control; use crate::binary_transport::run_binary_stage_message; -use crate::binary_transport::send_client_ready_hello_if_enabled; use crate::binary_transport::stage_execution::binary_message_attrs; use crate::binary_transport::stage_execution::binary_message_session_id; use crate::binary_transport::stage_execution::decode_record_tokens_sideband; use crate::binary_transport::stage_execution::elapsed_ms; use crate::binary_transport::stage_execution::empty_activation_frame; use crate::binary_transport::stage_execution::input_activation_frame; -use crate::binary_transport::stage_execution::insert_optional_unix_nanos; use crate::binary_transport::stage_execution::is_decode_frame_batch_candidate; use crate::binary_transport::stage_execution::nanos_delta_ms; use crate::binary_transport::stage_execution::runtime_sampling_config; @@ -41,39 +47,77 @@ use crate::binary_transport::stage_execution::stage_mask; use crate::binary_transport::stage_execution::token_sideband_or_fill; use crate::binary_transport::stage_output_activation_capacity; use crate::binary_transport::write_stage_message_conditioned; -use crate::kv_integration::{KvStageIntegration, model_requires_recurrent_state}; +use crate::kv_integration::KvStageIntegration; use crate::runtime_state::RuntimeState; use crate::telemetry::Telemetry; use crate::telemetry::now_unix_nanos; -use anyhow::Context; -use anyhow::Result; -use anyhow::bail; +use anyhow::{Context, Result, bail}; use serde_json::json; -use skippy_protocol::StageConfig; -use skippy_protocol::StageTopology; use skippy_protocol::binary::StageReply; use skippy_protocol::binary::StageReplyStats; use skippy_protocol::binary::StageWireMessage; use skippy_protocol::binary::WireActivationDType; use skippy_protocol::binary::WireMessageKind; use skippy_protocol::binary::WireReplyKind; -use skippy_protocol::binary::read_stage_message; use skippy_protocol::binary::recv_reply; use skippy_protocol::binary::send_reply_ack; use skippy_protocol::binary::send_reply_ack_with_stats; +use skippy_protocol::{StageConfig, StageTopology}; use std::collections::BTreeMap; -use std::io; use std::net::TcpStream; -use std::sync::Arc; -use std::sync::Mutex; -use std::sync::atomic::AtomicU64; -use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; use std::time::Instant; -static BINARY_SESSION_COUNTER: AtomicU64 = AtomicU64::new(1); - #[allow(clippy::too_many_arguments)] pub(super) fn handle_binary_connection( + config: &StageConfig, + topology: Option<&StageTopology>, + runtime: &Arc>, + decode_frame_batcher: &DecodeFrameBatcher, + kv: Option<&Arc>, + telemetry: &Telemetry, + upstream: &mut TcpStream, + downstream: Option, + activation_width: i32, + wire_dtype: WireActivationDType, + max_inflight: usize, + reply_credit_limit: Option, + async_prefill_forward: bool, + downstream_wire_condition: WireCondition, + downstream_connect_timeout_secs: u64, + native_mtp_enabled: bool, + prediction_return_sinks: &PredictionReturnSinks, + first_message: StageWireMessage, +) -> Result<()> { + let mut session_tracker = ConnectionSessionTracker::default(); + let result = handle_binary_connection_messages( + config, + topology, + runtime, + decode_frame_batcher, + kv, + telemetry, + upstream, + downstream, + activation_width, + wire_dtype, + max_inflight, + reply_credit_limit, + async_prefill_forward, + downstream_wire_condition, + downstream_connect_timeout_secs, + native_mtp_enabled, + prediction_return_sinks, + first_message, + &mut session_tracker, + ); + let cleanup_result = + release_tracked_connection_sessions(config, runtime, telemetry, &mut session_tracker); + combine_connection_and_cleanup_results(result, cleanup_result) +} + +#[allow(clippy::too_many_arguments)] +fn handle_binary_connection_messages( config: &StageConfig, topology: Option<&StageTopology>, runtime: &Arc>, @@ -92,16 +136,9 @@ pub(super) fn handle_binary_connection( native_mtp_enabled: bool, prediction_return_sinks: &PredictionReturnSinks, first_message: StageWireMessage, + session_tracker: &mut ConnectionSessionTracker, ) -> Result<()> { - if let Some(downstream) = downstream.as_mut() { - send_client_ready_hello_if_enabled(&mut *downstream) - .context("send downstream client ready hello")?; - skippy_protocol::binary::recv_ready(&mut *downstream) - .context("downstream binary stage did not become ready")?; - } - - let connection_session_id = BINARY_SESSION_COUNTER.fetch_add(1, Ordering::Relaxed); - let positional_speculation_supported = !model_requires_recurrent_state(config); + let connection_session_id = next_connection_session_id(); let max_deferred_prefill_replies = reply_credit_limit.unwrap_or_else(|| max_inflight.saturating_sub(1)); let mut pending_prefill_replies = 0usize; @@ -125,20 +162,15 @@ pub(super) fn handle_binary_connection( loop { let recv_start_unix_nanos = now_unix_nanos() as u64; let recv_started = Instant::now(); - let mut message = if let Some(message) = next_message.take() { - message - } else { - match read_stage_message(&mut *upstream, activation_width) { - Ok(message) => message, - Err(error) - if error.kind() == io::ErrorKind::UnexpectedEof - && pending_prefill_replies == 0 - && request_summary.message_count == 0 => - { - return Ok(()); - } - Err(error) => return Err(error).context("read binary stage message"), - } + let Some(mut message) = receive_next_message( + upstream, + activation_width, + next_message.take(), + pending_prefill_replies, + request_summary.message_count, + )? + else { + return Ok(()); }; let recv_end_unix_nanos = now_unix_nanos() as u64; let recv_read_ms = elapsed_ms(recv_started); @@ -146,325 +178,115 @@ pub(super) fn handle_binary_connection( let message_started = Instant::now(); let session_id = binary_message_session_id(connection_session_id, &message); let session_key = session_id.to_string(); - if message.kind == WireMessageKind::VerifyWindow && !positional_speculation_supported { - bail!( - "stage-state v10 positional speculation requires an attention-only stage; {} contains recurrent state", - config.stage_id - ); - } - if telemetry.is_debug_enabled() { - let mut recv_attrs = binary_message_attrs(config, session_id, &message); - recv_attrs.insert( - "llama_stage.recv_start_unix_nanos".to_string(), - json!(recv_start_unix_nanos), - ); - recv_attrs.insert( - "llama_stage.recv_end_unix_nanos".to_string(), - json!(recv_end_unix_nanos), - ); - recv_attrs.insert("llama_stage.recv_read_ms".to_string(), json!(recv_read_ms)); - recv_attrs.insert( - "skippy.upstream_message_wait_ms".to_string(), - json!(recv_read_ms), - ); - recv_attrs.insert( - "llama_stage.source_stage_index".to_string(), - json!(message.state.source_stage_index), - ); - recv_attrs.insert( - "llama_stage.configured_upstream_stage_index".to_string(), - json!(config.upstream.as_ref().map(|peer| peer.stage_index)), - ); - recv_attrs.insert( - "llama_stage.message_wire_bytes".to_string(), - json!(message.estimated_wire_bytes()), - ); - recv_attrs.insert( - "skippy.activation_bytes".to_string(), - json!(message.activation.len()), - ); - telemetry.emit_debug_span( - "stage.binary_recv", - recv_attrs, - recv_start_unix_nanos, - recv_end_unix_nanos, - ); - } + session_tracker.touch(&session_key); + emit_binary_message_received( + telemetry, + config, + session_id, + &message, + recv_start_unix_nanos, + recv_end_unix_nanos, + recv_read_ms, + ); if message.kind == WireMessageKind::Stop { - if pending_prefill_replies != 0 { - bail!("cannot stop with {pending_prefill_replies} deferred prefill replies"); - } - let mut stop_stats = std::mem::take(&mut pending_reply_stats); - request_summary.emit(telemetry, config, session_id); - request_summary = BinaryRequestSummary::default(); - if let Some(downstream) = downstream.as_mut() { - if let Some(forwarder) = async_forwarder.as_mut() { - forwarder - .flush() - .context("flush async forwards before stop")?; - } - write_stage_message_conditioned( - &mut *downstream, - &message, - wire_dtype, - downstream_wire_condition, - ) - .context("forward binary stop")?; - let reply = recv_reply(&mut *downstream).context("stop downstream ACK")?; - if reply.kind != WireReplyKind::Ack { - bail!("stop expected downstream ACK"); - } - stop_stats.merge(reply.stats); - } - let reset_start_unix_nanos = now_unix_nanos() as u64; - let reset_timer = Instant::now(); - let lock_timer = Instant::now(); - let mut runtime = runtime.lock().expect("runtime lock poisoned"); - let runtime_lock_wait_ms = elapsed_ms(lock_timer); - let accumulated = std::mem::take(&mut accumulated_prefill_tokens); - for (prefill_session_key, tokens) in accumulated { - let record = maybe_record_binary_full_prefill( - config, - &mut runtime, - kv, - telemetry, - &prefill_session_key, - &message, - &tokens, - ); - if record.recorded_pages > 0 { - stop_stats.kv_recorded_pages += record.recorded_pages as i64; - stop_stats.kv_record_stage_mask |= stage_mask(config.stage_index); - } - } - let drop_stats = runtime - .drop_session_timed(&session_key) - .context("reset binary stage session")?; - drop(runtime); - let reset_end_unix_nanos = now_unix_nanos() as u64; - let mut reset_attrs = binary_message_attrs(config, session_id, &message); - reset_attrs.insert( - "llama_stage.runtime_lock_wait_ms".to_string(), - json!(runtime_lock_wait_ms), - ); - reset_attrs.insert( - "llama_stage.session_reset_ms".to_string(), - json!(drop_stats.reset_ms), - ); - reset_attrs.insert( - "llama_stage.session_reset".to_string(), - json!(drop_stats.reset_session), - ); - reset_attrs.insert( - "llama_stage.lane_discarded".to_string(), - json!(drop_stats.lane_discarded), - ); - if let Some(reason) = drop_stats.lane_discard_reason.as_deref() { - reset_attrs.insert("llama_stage.lane_discard_reason".to_string(), json!(reason)); - } - reset_attrs.insert( - "llama_stage.elapsed_ms".to_string(), - json!(elapsed_ms(reset_timer)), - ); - insert_runtime_session_stats( - &mut reset_attrs, - "llama_stage.runtime_sessions_after", - &drop_stats.stats_after, - ); - telemetry.emit_debug_span( - "stage.binary_session_stop", - reset_attrs, - reset_start_unix_nanos, - reset_end_unix_nanos, - ); - prediction_return_streams.remove(&(message.request_id, message.session_id)); - prediction_return_sinks.remove(message.request_id, message.session_id); - send_reply_ack_with_stats(&mut *upstream, stop_stats).context("send stop ACK")?; - continue; - } - - if message.kind.is_session_control() { - let mut control_stats = std::mem::take(&mut pending_reply_stats); - if let Some(forwarder) = async_forwarder.as_mut() { - forwarder - .flush() - .context("flush async forwards before session control")?; - } - drain_deferred_prefill_replies( - downstream.as_mut(), - &mut pending_prefill_replies, - &mut control_stats, - ) - .context("drain deferred replies before session control")?; - { - let mut runtime = runtime.lock().expect("runtime lock poisoned"); - match message.kind { - WireMessageKind::TrimSession => runtime - .trim_session(&session_key, message.token_count.max(0) as u64) - .context("trim binary stage session")?, - _ => unreachable!("session control checked above"), - } - } - if let Some(downstream) = downstream.as_mut() { - write_stage_message_conditioned( - &mut *downstream, - &message, - wire_dtype, - downstream_wire_condition, - ) - .context("forward session control")?; - let reply = - recv_reply(&mut *downstream).context("session control downstream ACK")?; - if reply.kind != WireReplyKind::Ack { - bail!("session control expected downstream ACK"); - } - control_stats.merge(reply.stats); - } - send_reply_ack_with_stats(&mut *upstream, control_stats) - .context("session control ack")?; - continue; - } - - if message.kind.is_generation_control() { - let mut generation_stats = std::mem::take(&mut pending_reply_stats); - if let Some(forwarder) = async_forwarder.as_mut() { - forwarder - .flush() - .context("flush async forwards before generation config")?; - } - drain_deferred_prefill_replies( - downstream.as_mut(), - &mut pending_prefill_replies, - &mut generation_stats, - ) - .context("drain deferred replies before generation config")?; - if let Some(downstream) = downstream.as_mut() { - write_stage_message_conditioned( - &mut *downstream, - &message, - wire_dtype, - downstream_wire_condition, - ) - .context("forward generation config")?; - let reply = - recv_reply(&mut *downstream).context("generation config downstream ACK")?; - if reply.kind != WireReplyKind::Ack { - bail!("generation config expected downstream ACK"); - } - generation_stats.merge(reply.stats); - } else { - if let Some(metadata) = message.chat_sampling_metadata.as_deref() { - let sampling = runtime_sampling_config(message.sampling.as_ref()); - let mut runtime = runtime.lock().expect("runtime lock poisoned"); - runtime - .configure_chat_sampling( - &session_key, - metadata, - message.state.prompt_token_count.max(0) as u64, - sampling.as_ref(), - ) - .context("configure binary stage generation")?; - } - configure_prediction_return_stream( - config, - topology, - message.request_id, - message.session_id, - wire_dtype, - downstream_connect_timeout_secs, - prediction_return_sinks, - &mut prediction_return_streams, - ); - } - send_reply_ack_with_stats(&mut *upstream, generation_stats) - .context("generation config ack")?; - continue; - } - - if message.kind.is_prefix_cache_control() { - let control_started = Instant::now(); - let mut control_stats = std::mem::take(&mut pending_reply_stats); - if let Some(forwarder) = async_forwarder.as_mut() { - forwarder - .flush() - .context("flush async forwards before prefix cache control")?; - } - drain_deferred_prefill_replies( - downstream.as_mut(), - &mut pending_prefill_replies, - &mut control_stats, - ) - .context("drain deferred replies before prefix cache control")?; - if message.kind == WireMessageKind::TryRestorePrefillDecode { - handle_binary_restore_prefill_decode_control( - config, - topology, - runtime, - kv, - telemetry, - &session_key, - session_id, - message, - downstream.as_mut(), - wire_dtype, - downstream_wire_condition, - activation_width, - control_started, - control_stats, - prediction_return_sinks, - &mut prediction_return_streams, - downstream_connect_timeout_secs, - native_mtp_enabled, - ) - .context("handle restore-prefill-decode control")?; - continue; - } - let token_ids = token_sideband_or_fill(&message) - .context("read prefix cache control token sideband")?; - let local = maybe_prefix_cache_control( + handle_stop( config, runtime, kv, telemetry, - &session_key, + upstream, + downstream.as_mut(), + wire_dtype, + downstream_wire_condition, &message, - &token_ids, - ); - control_stats.merge(local.stats); - if local.hit - && let Some(downstream) = downstream.as_mut() - { - write_stage_message_conditioned( - &mut *downstream, - &message, - wire_dtype, - downstream_wire_condition, - ) - .context("forward prefix cache control")?; - let reply = recv_reply(&mut *downstream).context("prefix cache downstream ACK")?; - if reply.kind != WireReplyKind::Ack { - bail!("prefix cache control expected downstream ACK"); - } - let downstream_missed = message.kind == WireMessageKind::TryRestorePrefill - && (reply.stats.kv_lookup_misses > 0 - || reply.stats.kv_lookup_errors > 0 - || reply.stats.kv_lookup_hits == 0); - control_stats.merge(reply.stats); - if downstream_missed { - let mut runtime = runtime.lock().expect("runtime lock poisoned"); - let _ = runtime.drop_session_timed(&session_key); - } - } - let mut attrs = binary_message_attrs(config, session_id, &message); - attrs.insert("skippy.kv.control_hit".to_string(), json!(local.hit)); - attrs.insert( - "llama_stage.elapsed_ms".to_string(), - json!(elapsed_ms(control_started)), - ); - telemetry.emit_debug("stage.binary_prefix_cache_control", attrs); - send_reply_ack_with_stats(&mut *upstream, control_stats) - .context("prefix cache control ack")?; + &session_key, + session_id, + pending_prefill_replies, + &mut pending_reply_stats, + &mut request_summary, + &mut accumulated_prefill_tokens, + async_forwarder.as_mut(), + session_tracker, + &mut prediction_return_streams, + prediction_return_sinks, + )?; + continue; + } + + if message.kind.is_verify_retirement() { + handle_verify_retirement( + runtime, + downstream.as_mut(), + wire_dtype, + downstream_wire_condition, + &message, + &session_key, + async_forwarder.as_mut(), + )?; + continue; + } + + if message.kind.is_session_control() { + handle_session_control( + runtime, + upstream, + downstream.as_mut(), + wire_dtype, + downstream_wire_condition, + &message, + &session_key, + &mut pending_prefill_replies, + &mut pending_reply_stats, + async_forwarder.as_mut(), + )?; + continue; + } + + if message.kind.is_generation_control() { + handle_generation_control( + config, + topology, + runtime, + upstream, + downstream.as_mut(), + wire_dtype, + downstream_wire_condition, + downstream_connect_timeout_secs, + &message, + &session_key, + &mut pending_prefill_replies, + &mut pending_reply_stats, + async_forwarder.as_mut(), + prediction_return_sinks, + &mut prediction_return_streams, + )?; + continue; + } + + if message.kind.is_prefix_cache_control() { + handle_prefix_cache_control( + config, + topology, + runtime, + kv, + telemetry, + upstream, + downstream.as_mut(), + wire_dtype, + downstream_wire_condition, + downstream_connect_timeout_secs, + activation_width, + native_mtp_enabled, + message, + &session_key, + session_id, + &mut pending_prefill_replies, + &mut pending_reply_stats, + async_forwarder.as_mut(), + prediction_return_sinks, + &mut prediction_return_streams, + )?; continue; } @@ -495,37 +317,17 @@ pub(super) fn handle_binary_connection( } let token_ids = token_sideband_or_fill(&message)?; - let mut session_auto_align_count = 0usize; - let mut session_auto_align_ms = 0.0; - let mut session_auto_align_trimmed_tokens = 0u64; - if let Some(target_token_count) = message.authoritative_session_position() { - let align_started = Instant::now(); - let align = { - let mut runtime = runtime.lock().expect("runtime lock poisoned"); - runtime - .align_session_to_token_count_if_ahead(&session_key, target_token_count) - .context("auto-align binary stage session")? - }; - if let Some(align) = align { - let align_ms = elapsed_ms(align_started); - session_auto_align_count = 1; - session_auto_align_ms = align_ms; - session_auto_align_trimmed_tokens = align - .before_token_count - .saturating_sub(align.after_token_count); - let mut attrs = binary_message_attrs(config, session_id, &message); - attrs.insert( - "llama_stage.session_auto_align_before_tokens".to_string(), - json!(align.before_token_count), - ); - attrs.insert( - "llama_stage.session_auto_align_after_tokens".to_string(), - json!(align.after_token_count), - ); - attrs.insert("llama_stage.elapsed_ms".to_string(), json!(align_ms)); - telemetry.emit_debug("stage.binary_session_auto_align", attrs); - } - } + let auto_align = align_session_to_message( + config, + runtime, + telemetry, + &session_key, + session_id, + &message, + )?; + let session_auto_align_count = auto_align.count; + let session_auto_align_ms = auto_align.elapsed_ms; + let session_auto_align_trimmed_tokens = auto_align.trimmed_tokens; if message.kind.is_prefill() { accumulate_prefill_tokens( &mut accumulated_prefill_tokens, @@ -1149,154 +951,44 @@ pub(super) fn handle_binary_connection( upstream_message_wait_ms: recv_read_ms, }); - if telemetry.is_debug_enabled() { - let mut timing_attrs = binary_message_attrs(config, session_id, &message); - timing_attrs.insert( - "llama_stage.message_start_unix_nanos".to_string(), - json!(message_start_unix_nanos), - ); - timing_attrs.insert( - "llama_stage.message_end_unix_nanos".to_string(), - json!(message_end_unix_nanos), - ); - timing_attrs.insert( - "llama_stage.compute_start_unix_nanos".to_string(), - json!(compute_start_unix_nanos), - ); - timing_attrs.insert( - "llama_stage.compute_end_unix_nanos".to_string(), - json!(compute_end_unix_nanos), - ); - timing_attrs.insert("llama_stage.compute_ms".to_string(), json!(compute_ms)); - timing_attrs.insert("llama_stage.recv_read_ms".to_string(), json!(recv_read_ms)); - timing_attrs.insert( - "skippy.upstream_message_wait_ms".to_string(), - json!(recv_read_ms), - ); - timing_attrs.insert( - "llama_stage.input_activation_decode_ms".to_string(), - json!(input_activation_decode_ms), - ); - timing_attrs.insert( - "llama_stage.runtime_lock_wait_ms".to_string(), - json!(runtime_lock_wait_ms), - ); - timing_attrs.insert( - "llama_stage.runtime_lock_hold_ms".to_string(), - json!(runtime_lock_hold_ms), - ); - timing_attrs.insert( - "llama_stage.runtime_lock_acquires".to_string(), - json!(runtime_lock_acquires), - ); - if let Some(stats) = runtime_sessions_before.as_ref() { - insert_runtime_session_stats( - &mut timing_attrs, - "llama_stage.runtime_sessions_before", - stats, - ); - } - if let Some(stats) = runtime_sessions_after.as_ref() { - insert_runtime_session_stats( - &mut timing_attrs, - "llama_stage.runtime_sessions_after", - stats, - ); - } - timing_attrs.insert( - "llama_stage.forward_write_ms".to_string(), - json!(forward_write_ms), - ); - timing_attrs.insert( - "llama_stage.activation_encode_ms".to_string(), - json!(forward_activation_encode_ms), - ); - timing_attrs.insert( - "llama_stage.downstream_wait_ms".to_string(), - json!(downstream_wait_ms), - ); - timing_attrs.insert("skippy.compute_ms".to_string(), json!(compute_ms)); - timing_attrs.insert( - "skippy.forward_write_ms".to_string(), - json!(forward_write_ms), - ); - timing_attrs.insert( - "skippy.downstream_wait_ms".to_string(), - json!(downstream_wait_ms), - ); - timing_attrs.insert( - "skippy.upstream_reply_ms".to_string(), - json!(upstream_reply_ms), - ); - timing_attrs.insert("llama_stage.forward_mode".to_string(), json!(forward_mode)); - insert_optional_unix_nanos( - &mut timing_attrs, - "llama_stage.forward_write_start_unix_nanos", - forward_write_start_unix_nanos, - ); - insert_optional_unix_nanos( - &mut timing_attrs, - "llama_stage.forward_write_end_unix_nanos", - forward_write_end_unix_nanos, - ); - insert_optional_unix_nanos( - &mut timing_attrs, - "llama_stage.downstream_wait_start_unix_nanos", - downstream_wait_start_unix_nanos, - ); - insert_optional_unix_nanos( - &mut timing_attrs, - "llama_stage.downstream_wait_end_unix_nanos", - downstream_wait_end_unix_nanos, - ); - insert_optional_unix_nanos( - &mut timing_attrs, - "llama_stage.upstream_reply_start_unix_nanos", - upstream_reply_start_unix_nanos, - ); - insert_optional_unix_nanos( - &mut timing_attrs, - "llama_stage.upstream_reply_end_unix_nanos", - upstream_reply_end_unix_nanos, - ); - timing_attrs.insert( - "skippy.message_elapsed_ms".to_string(), - json!(message_elapsed_ms), - ); - timing_attrs.insert( - "skippy.input_activation_bytes".to_string(), - json!(input_activation_bytes), - ); - timing_attrs.insert( - "skippy.output_activation_bytes".to_string(), - json!(output.payload.len()), - ); - timing_attrs.insert( - "skippy.prefill_credit_limit".to_string(), - json!(max_deferred_prefill_replies), - ); - timing_attrs.insert( - "skippy.prefill_pending_replies_before".to_string(), - json!(pending_prefill_replies_before), - ); - timing_attrs.insert( - "skippy.prefill_pending_replies_after".to_string(), - json!(pending_prefill_replies), - ); - timing_attrs.insert( - "skippy.prefill_credit_wait_count".to_string(), - json!(credit_wait_count), - ); - timing_attrs.insert( - "skippy.prefill_deferred_replies_drained".to_string(), - json!(deferred_prefill_replies_drained), - ); - telemetry.emit_debug_span( - "stage.binary_message_timing", - timing_attrs, + emit_binary_message_timing( + telemetry, + config, + session_id, + &message, + BinaryMessageTiming { message_start_unix_nanos, message_end_unix_nanos, - ); - } + compute_start_unix_nanos, + compute_end_unix_nanos, + forward_write_start_unix_nanos, + forward_write_end_unix_nanos, + downstream_wait_start_unix_nanos, + downstream_wait_end_unix_nanos, + upstream_reply_start_unix_nanos, + upstream_reply_end_unix_nanos, + compute_ms, + recv_read_ms, + input_activation_decode_ms, + runtime_lock_wait_ms, + runtime_lock_hold_ms, + runtime_lock_acquires, + runtime_sessions_before: runtime_sessions_before.as_ref(), + runtime_sessions_after: runtime_sessions_after.as_ref(), + forward_write_ms, + forward_activation_encode_ms, + downstream_wait_ms, + upstream_reply_ms, + forward_mode, + message_elapsed_ms, + input_activation_bytes, + output_activation_bytes: output.payload.len(), + max_deferred_prefill_replies, + pending_prefill_replies_before, + pending_prefill_replies_after: pending_prefill_replies, + credit_wait_count, + deferred_prefill_replies_drained, + }, + ); } } diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/control_messages.rs b/crates/skippy-server/src/binary_transport/binary_messaging/control_messages.rs new file mode 100644 index 000000000..ad69ee5e2 --- /dev/null +++ b/crates/skippy-server/src/binary_transport/binary_messaging/control_messages.rs @@ -0,0 +1,407 @@ +use super::async_forwarder::AsyncForwarder; +use super::reply::{ + configure_prediction_return_stream, drain_deferred_prefill_replies, + normalize_downstream_prefix_restore_reply, +}; +use super::session_tracker::ConnectionSessionTracker; +use super::summary::BinaryRequestSummary; +use crate::binary_transport::WireCondition; +use crate::binary_transport::binary_kv::{ + maybe_prefix_cache_control, maybe_record_binary_full_prefill, +}; +use crate::binary_transport::direct_return::PredictionReturnSinks; +use crate::binary_transport::restore_prefill_decode::handle_binary_restore_prefill_decode_control; +use crate::binary_transport::stage_execution::{ + binary_message_attrs, elapsed_ms, runtime_sampling_config, stage_mask, token_sideband_or_fill, +}; +use crate::binary_transport::write_stage_message_conditioned; +use crate::kv_integration::KvStageIntegration; +use crate::runtime_state::RuntimeState; +use crate::telemetry::{Telemetry, now_unix_nanos}; +use anyhow::{Context, Result, bail}; +use serde_json::json; +use skippy_protocol::binary::{ + StageReplyStats, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, + recv_reply, send_reply_ack_with_stats, +}; +use skippy_protocol::{StageConfig, StageTopology}; +use std::collections::BTreeMap; +use std::net::TcpStream; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +#[allow(clippy::too_many_arguments)] +pub(super) fn handle_stop( + config: &StageConfig, + runtime: &Arc>, + kv: Option<&Arc>, + telemetry: &Telemetry, + upstream: &mut TcpStream, + mut downstream: Option<&mut TcpStream>, + wire_dtype: WireActivationDType, + downstream_wire_condition: WireCondition, + message: &StageWireMessage, + session_key: &str, + session_id: u64, + pending_prefill_replies: usize, + pending_reply_stats: &mut StageReplyStats, + request_summary: &mut BinaryRequestSummary, + accumulated_prefill_tokens: &mut BTreeMap>, + async_forwarder: Option<&mut AsyncForwarder>, + session_tracker: &mut ConnectionSessionTracker, + prediction_return_streams: &mut BTreeMap<(u64, u64), TcpStream>, + prediction_return_sinks: &PredictionReturnSinks, +) -> Result<()> { + if pending_prefill_replies != 0 { + bail!("cannot stop with {pending_prefill_replies} deferred prefill replies"); + } + let mut stop_stats = std::mem::take(pending_reply_stats); + request_summary.emit(telemetry, config, session_id); + *request_summary = BinaryRequestSummary::default(); + if let Some(downstream) = downstream.as_mut() { + if let Some(forwarder) = async_forwarder { + forwarder + .flush() + .context("flush async forwards before stop")?; + } + write_stage_message_conditioned( + &mut **downstream, + message, + wire_dtype, + downstream_wire_condition, + ) + .context("forward binary stop")?; + let reply = recv_reply(&mut **downstream).context("stop downstream ACK")?; + if reply.kind != WireReplyKind::Ack { + bail!("stop expected downstream ACK"); + } + stop_stats.merge(reply.stats); + } + let reset_start_unix_nanos = now_unix_nanos() as u64; + let reset_timer = Instant::now(); + let lock_timer = Instant::now(); + let mut runtime = runtime.lock().expect("runtime lock poisoned"); + let runtime_lock_wait_ms = elapsed_ms(lock_timer); + let accumulated = std::mem::take(accumulated_prefill_tokens); + for (prefill_session_key, tokens) in accumulated { + let record = maybe_record_binary_full_prefill( + config, + &mut runtime, + kv, + telemetry, + &prefill_session_key, + message, + &tokens, + ); + if record.recorded_pages > 0 { + stop_stats.kv_recorded_pages += record.recorded_pages as i64; + stop_stats.kv_record_stage_mask |= stage_mask(config.stage_index); + } + } + let drop_stats = runtime + .drop_session_timed(session_key) + .context("reset binary stage session")?; + drop(runtime); + let reset_end_unix_nanos = now_unix_nanos() as u64; + let mut reset_attrs = binary_message_attrs(config, session_id, message); + reset_attrs.insert( + "llama_stage.runtime_lock_wait_ms".to_string(), + json!(runtime_lock_wait_ms), + ); + reset_attrs.insert( + "llama_stage.session_reset_ms".to_string(), + json!(drop_stats.reset_ms), + ); + reset_attrs.insert( + "llama_stage.session_reset".to_string(), + json!(drop_stats.reset_session), + ); + reset_attrs.insert( + "llama_stage.lane_discarded".to_string(), + json!(drop_stats.lane_discarded), + ); + if let Some(reason) = drop_stats.lane_discard_reason.as_deref() { + reset_attrs.insert("llama_stage.lane_discard_reason".to_string(), json!(reason)); + } + reset_attrs.insert( + "llama_stage.elapsed_ms".to_string(), + json!(elapsed_ms(reset_timer)), + ); + super::telemetry::insert_runtime_session_stats( + &mut reset_attrs, + "llama_stage.runtime_sessions_after", + &drop_stats.stats_after, + ); + telemetry.emit_debug_span( + "stage.binary_session_stop", + reset_attrs, + reset_start_unix_nanos, + reset_end_unix_nanos, + ); + session_tracker.stopped(session_key); + prediction_return_streams.remove(&(message.request_id, message.session_id)); + prediction_return_sinks.remove(message.request_id, message.session_id); + send_reply_ack_with_stats(upstream, stop_stats).context("send stop ACK") +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn handle_verify_retirement( + runtime: &Arc>, + mut downstream: Option<&mut TcpStream>, + wire_dtype: WireActivationDType, + downstream_wire_condition: WireCondition, + message: &StageWireMessage, + session_key: &str, + async_forwarder: Option<&mut AsyncForwarder>, +) -> Result<()> { + if let Some(forwarder) = async_forwarder { + forwarder + .flush() + .context("flush async forwards before verify retirement")?; + } + let token_start = u64::try_from(message.pos_start) + .context("verify retirement position must be non-negative")?; + let token_count = u64::try_from(message.token_count) + .context("verify retirement count must be non-negative")?; + runtime + .lock() + .expect("runtime lock poisoned") + .retire_verify_checkpoint(session_key, token_start, token_count) + .context("retire binary stage verify checkpoint")?; + if let Some(downstream) = downstream.as_mut() { + write_stage_message_conditioned( + &mut **downstream, + message, + wire_dtype, + downstream_wire_condition, + ) + .context("forward verify retirement")?; + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn handle_session_control( + runtime: &Arc>, + upstream: &mut TcpStream, + mut downstream: Option<&mut TcpStream>, + wire_dtype: WireActivationDType, + downstream_wire_condition: WireCondition, + message: &StageWireMessage, + session_key: &str, + pending_prefill_replies: &mut usize, + pending_reply_stats: &mut StageReplyStats, + async_forwarder: Option<&mut AsyncForwarder>, +) -> Result<()> { + let mut control_stats = std::mem::take(pending_reply_stats); + if let Some(forwarder) = async_forwarder { + forwarder + .flush() + .context("flush async forwards before session control")?; + } + drain_deferred_prefill_replies( + downstream.as_deref_mut(), + pending_prefill_replies, + &mut control_stats, + ) + .context("drain deferred replies before session control")?; + match message.kind { + WireMessageKind::TrimSession => runtime + .lock() + .expect("runtime lock poisoned") + .trim_session(session_key, message.token_count.max(0) as u64) + .context("trim binary stage session")?, + _ => unreachable!("session control checked above"), + } + if let Some(downstream) = downstream.as_mut() { + write_stage_message_conditioned( + &mut **downstream, + message, + wire_dtype, + downstream_wire_condition, + ) + .context("forward session control")?; + let reply = recv_reply(&mut **downstream).context("session control downstream ACK")?; + if reply.kind != WireReplyKind::Ack { + bail!("session control expected downstream ACK"); + } + control_stats.merge(reply.stats); + } + send_reply_ack_with_stats(upstream, control_stats).context("session control ack") +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn handle_generation_control( + config: &StageConfig, + topology: Option<&StageTopology>, + runtime: &Arc>, + upstream: &mut TcpStream, + mut downstream: Option<&mut TcpStream>, + wire_dtype: WireActivationDType, + downstream_wire_condition: WireCondition, + downstream_connect_timeout_secs: u64, + message: &StageWireMessage, + session_key: &str, + pending_prefill_replies: &mut usize, + pending_reply_stats: &mut StageReplyStats, + async_forwarder: Option<&mut AsyncForwarder>, + prediction_return_sinks: &PredictionReturnSinks, + prediction_return_streams: &mut BTreeMap<(u64, u64), TcpStream>, +) -> Result<()> { + let mut generation_stats = std::mem::take(pending_reply_stats); + if let Some(forwarder) = async_forwarder { + forwarder + .flush() + .context("flush async forwards before generation config")?; + } + drain_deferred_prefill_replies( + downstream.as_deref_mut(), + pending_prefill_replies, + &mut generation_stats, + ) + .context("drain deferred replies before generation config")?; + if let Some(downstream) = downstream.as_mut() { + write_stage_message_conditioned( + &mut **downstream, + message, + wire_dtype, + downstream_wire_condition, + ) + .context("forward generation config")?; + let reply = recv_reply(&mut **downstream).context("generation config downstream ACK")?; + if reply.kind != WireReplyKind::Ack { + bail!("generation config expected downstream ACK"); + } + generation_stats.merge(reply.stats); + } else { + if let Some(metadata) = message.chat_sampling_metadata.as_deref() { + let sampling = runtime_sampling_config(message.sampling.as_ref()); + runtime + .lock() + .expect("runtime lock poisoned") + .configure_chat_sampling( + session_key, + metadata, + message.state.prompt_token_count.max(0) as u64, + sampling.as_ref(), + ) + .context("configure binary stage generation")?; + } + configure_prediction_return_stream( + config, + topology, + message.request_id, + message.session_id, + wire_dtype, + downstream_connect_timeout_secs, + prediction_return_sinks, + prediction_return_streams, + ); + } + send_reply_ack_with_stats(upstream, generation_stats).context("generation config ack") +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn handle_prefix_cache_control( + config: &StageConfig, + topology: Option<&StageTopology>, + runtime: &Arc>, + kv: Option<&Arc>, + telemetry: &Telemetry, + upstream: &mut TcpStream, + mut downstream: Option<&mut TcpStream>, + wire_dtype: WireActivationDType, + downstream_wire_condition: WireCondition, + downstream_connect_timeout_secs: u64, + activation_width: i32, + native_mtp_enabled: bool, + message: StageWireMessage, + session_key: &str, + session_id: u64, + pending_prefill_replies: &mut usize, + pending_reply_stats: &mut StageReplyStats, + async_forwarder: Option<&mut AsyncForwarder>, + prediction_return_sinks: &PredictionReturnSinks, + prediction_return_streams: &mut BTreeMap<(u64, u64), TcpStream>, +) -> Result<()> { + let control_started = Instant::now(); + let mut control_stats = std::mem::take(pending_reply_stats); + if let Some(forwarder) = async_forwarder { + forwarder + .flush() + .context("flush async forwards before prefix cache control")?; + } + drain_deferred_prefill_replies( + downstream.as_deref_mut(), + pending_prefill_replies, + &mut control_stats, + ) + .context("drain deferred replies before prefix cache control")?; + if message.kind == WireMessageKind::TryRestorePrefillDecode { + return handle_binary_restore_prefill_decode_control( + config, + topology, + runtime, + kv, + telemetry, + session_key, + session_id, + message, + downstream, + wire_dtype, + downstream_wire_condition, + activation_width, + control_started, + control_stats, + prediction_return_sinks, + prediction_return_streams, + downstream_connect_timeout_secs, + native_mtp_enabled, + ) + .context("handle restore-prefill-decode control"); + } + let token_ids = + token_sideband_or_fill(&message).context("read prefix cache control token sideband")?; + let local = maybe_prefix_cache_control( + config, + runtime, + kv, + telemetry, + session_key, + &message, + &token_ids, + ); + control_stats.merge(local.stats); + if local.hit + && let Some(downstream) = downstream.as_mut() + { + write_stage_message_conditioned( + &mut **downstream, + &message, + wire_dtype, + downstream_wire_condition, + ) + .context("forward prefix cache control")?; + let mut reply = recv_reply(&mut **downstream).context("prefix cache downstream ACK")?; + if reply.kind != WireReplyKind::Ack { + bail!("prefix cache control expected downstream ACK"); + } + let downstream_missed = + normalize_downstream_prefix_restore_reply(message.kind, &mut reply.stats); + control_stats.merge(reply.stats); + if downstream_missed { + let _ = runtime + .lock() + .expect("runtime lock poisoned") + .drop_session_timed(session_key); + } + } + let mut attrs = binary_message_attrs(config, session_id, &message); + attrs.insert("skippy.kv.control_hit".to_string(), json!(local.hit)); + attrs.insert( + "llama_stage.elapsed_ms".to_string(), + json!(elapsed_ms(control_started)), + ); + telemetry.emit_debug("stage.binary_prefix_cache_control", attrs); + send_reply_ack_with_stats(upstream, control_stats).context("prefix cache control ack") +} diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs b/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs new file mode 100644 index 000000000..8450c2c9a --- /dev/null +++ b/crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs @@ -0,0 +1,34 @@ +use anyhow::{Context, Result}; +use skippy_protocol::binary::{StageWireMessage, read_stage_message}; +use std::io; +use std::net::TcpStream; +use std::sync::atomic::{AtomicU64, Ordering}; + +static BINARY_SESSION_COUNTER: AtomicU64 = AtomicU64::new(1); + +pub(super) fn next_connection_session_id() -> u64 { + BINARY_SESSION_COUNTER.fetch_add(1, Ordering::Relaxed) +} + +pub(super) fn receive_next_message( + upstream: &mut TcpStream, + activation_width: i32, + first_message: Option, + pending_prefill_replies: usize, + observed_message_count: usize, +) -> Result> { + if first_message.is_some() { + return Ok(first_message); + } + match read_stage_message(upstream, activation_width) { + Ok(message) => Ok(Some(message)), + Err(error) + if error.kind() == io::ErrorKind::UnexpectedEof + && pending_prefill_replies == 0 + && observed_message_count == 0 => + { + Ok(None) + } + Err(error) => Err(error).context("read binary stage message"), + } +} diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/reply.rs b/crates/skippy-server/src/binary_transport/binary_messaging/reply.rs index 2edd6a919..17e7f37f5 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/reply.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/reply.rs @@ -9,6 +9,7 @@ use skippy_protocol::binary::StageReply; use skippy_protocol::binary::StageReplyStats; use skippy_protocol::binary::StageReplyWindow; use skippy_protocol::binary::WireActivationDType; +use skippy_protocol::binary::WireMessageKind; use skippy_protocol::binary::WireReplyKind; use skippy_protocol::binary::recv_reply; use skippy_protocol::binary::send_reply_message; @@ -96,6 +97,29 @@ pub(super) fn reply_window_for_message( } } +/// Normalizes a downstream `TryRestorePrefill` reply before its stats are +/// merged upstream. A stage without cache integration returns neutral stats; +/// after an upstream stage restored successfully, that neutral response means +/// the chain is incomplete and must be reported as a miss. +pub(super) fn normalize_downstream_prefix_restore_reply( + kind: WireMessageKind, + stats: &mut StageReplyStats, +) -> bool { + if kind != WireMessageKind::TryRestorePrefill { + return false; + } + let missed = + stats.kv_lookup_misses > 0 || stats.kv_lookup_errors > 0 || stats.kv_lookup_hits == 0; + if missed + && stats.kv_lookup_hits == 0 + && stats.kv_lookup_misses == 0 + && stats.kv_lookup_errors == 0 + { + stats.kv_lookup_misses = 1; + } + missed +} + #[cfg(test)] mod tests { use super::*; @@ -125,4 +149,34 @@ mod tests { assert_eq!(reply.window_id, 42); } + + #[test] + fn neutral_terminal_restore_reply_becomes_a_chain_miss() { + let mut stats = StageReplyStats::default(); + + let missed = normalize_downstream_prefix_restore_reply( + WireMessageKind::TryRestorePrefill, + &mut stats, + ); + + assert!(missed); + assert_eq!(stats.kv_lookup_misses, 1); + } + + #[test] + fn downstream_restore_hit_remains_a_hit() { + let mut stats = StageReplyStats { + kv_lookup_hits: 1, + ..StageReplyStats::default() + }; + + let missed = normalize_downstream_prefix_restore_reply( + WireMessageKind::TryRestorePrefill, + &mut stats, + ); + + assert!(!missed); + assert_eq!(stats.kv_lookup_hits, 1); + assert_eq!(stats.kv_lookup_misses, 0); + } } diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/session_lifecycle.rs b/crates/skippy-server/src/binary_transport/binary_messaging/session_lifecycle.rs new file mode 100644 index 000000000..75025fc47 --- /dev/null +++ b/crates/skippy-server/src/binary_transport/binary_messaging/session_lifecycle.rs @@ -0,0 +1,57 @@ +use crate::binary_transport::stage_execution::{binary_message_attrs, elapsed_ms}; +use crate::runtime_state::RuntimeState; +use crate::telemetry::Telemetry; +use anyhow::{Context, Result}; +use serde_json::json; +use skippy_protocol::StageConfig; +use skippy_protocol::binary::StageWireMessage; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +#[derive(Default)] +pub(super) struct SessionAutoAlignObservation { + pub(super) count: usize, + pub(super) elapsed_ms: f64, + pub(super) trimmed_tokens: u64, +} + +pub(super) fn align_session_to_message( + config: &StageConfig, + runtime: &Arc>, + telemetry: &Telemetry, + session_key: &str, + session_id: u64, + message: &StageWireMessage, +) -> Result { + let Some(target_token_count) = message.authoritative_session_position() else { + return Ok(SessionAutoAlignObservation::default()); + }; + let started = Instant::now(); + let align = runtime + .lock() + .expect("runtime lock poisoned") + .align_session_to_token_count_if_ahead(session_key, target_token_count) + .context("auto-align binary stage session")?; + let Some(align) = align else { + return Ok(SessionAutoAlignObservation::default()); + }; + let elapsed_ms = elapsed_ms(started); + let mut attrs = binary_message_attrs(config, session_id, message); + attrs.insert( + "llama_stage.session_auto_align_before_tokens".to_string(), + json!(align.before_token_count), + ); + attrs.insert( + "llama_stage.session_auto_align_after_tokens".to_string(), + json!(align.after_token_count), + ); + attrs.insert("llama_stage.elapsed_ms".to_string(), json!(elapsed_ms)); + telemetry.emit_debug("stage.binary_session_auto_align", attrs); + Ok(SessionAutoAlignObservation { + count: 1, + elapsed_ms, + trimmed_tokens: align + .before_token_count + .saturating_sub(align.after_token_count), + }) +} diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/session_tracker.rs b/crates/skippy-server/src/binary_transport/binary_messaging/session_tracker.rs new file mode 100644 index 000000000..77030200e --- /dev/null +++ b/crates/skippy-server/src/binary_transport/binary_messaging/session_tracker.rs @@ -0,0 +1,151 @@ +use super::telemetry::insert_runtime_session_stats; +use crate::{ + runtime_state::RuntimeState, + telemetry::{Telemetry, lifecycle_attrs}, +}; +use anyhow::{Context, Result, bail}; +use serde_json::json; +use skippy_protocol::StageConfig; +use std::{ + collections::BTreeSet, + sync::{Arc, Mutex}, +}; + +/// Runtime session keys created by one binary stage connection. +/// +/// A connection that fails before its graceful `Stop` message would otherwise +/// leave those sessions holding execution lanes indefinitely. +#[derive(Default)] +pub(super) struct ConnectionSessionTracker { + active: BTreeSet, +} + +impl ConnectionSessionTracker { + pub(super) fn touch(&mut self, session_key: &str) { + self.active.insert(session_key.to_string()); + } + + pub(super) fn stopped(&mut self, session_key: &str) { + self.active.remove(session_key); + } + + fn drain(&mut self) -> Vec { + std::mem::take(&mut self.active).into_iter().collect() + } +} + +/// Returns lanes held by sessions that never reached a graceful `Stop`. +pub(super) fn release_tracked_connection_sessions( + config: &StageConfig, + runtime: &Arc>, + telemetry: &Telemetry, + session_tracker: &mut ConnectionSessionTracker, +) -> Result<()> { + let orphaned = session_tracker.drain(); + if orphaned.is_empty() { + return Ok(()); + } + let orphaned_count = orphaned.len(); + let mut runtime = runtime.lock().map_err(|_| { + anyhow::anyhow!( + "failed to reclaim {orphaned_count} orphaned binary stage session(s): runtime lock poisoned" + ) + })?; + let mut failures = Vec::new(); + for session_key in orphaned { + match runtime.drop_session_timed(&session_key) { + Ok(drop_stats) => { + let mut attrs = lifecycle_attrs(config); + attrs.insert("llama_stage.session_key".to_string(), json!(session_key)); + attrs.insert( + "llama_stage.session_reset".to_string(), + json!(drop_stats.reset_session), + ); + attrs.insert( + "llama_stage.lane_discarded".to_string(), + json!(drop_stats.lane_discarded), + ); + insert_runtime_session_stats( + &mut attrs, + "llama_stage.runtime_sessions_after", + &drop_stats.stats_after, + ); + telemetry.emit("stage.binary_session_orphan_reclaimed", attrs); + } + Err(error) => { + failures.push(format!("{session_key}: {error:#}")); + } + } + } + if !failures.is_empty() { + bail!( + "failed to reclaim {}/{} orphaned binary stage session(s): {}", + failures.len(), + orphaned_count, + failures.join("; ") + ); + } + Ok(()) +} + +pub(super) fn combine_connection_and_cleanup_results( + connection_result: Result<()>, + cleanup_result: Result<()>, +) -> Result<()> { + match (connection_result, cleanup_result) { + (Ok(()), Ok(())) => Ok(()), + (Err(connection_error), Ok(())) => Err(connection_error), + (Ok(()), Err(cleanup_error)) => Err(cleanup_error), + (Err(connection_error), Err(cleanup_error)) => Err(connection_error).with_context(|| { + format!("orphaned binary stage session cleanup also failed: {cleanup_error:#}") + }), + } +} + +#[cfg(test)] +mod tests { + use super::{ConnectionSessionTracker, combine_connection_and_cleanup_results}; + use anyhow::anyhow; + + #[test] + fn tracker_drains_sessions_that_never_saw_a_stop() { + let mut tracker = ConnectionSessionTracker::default(); + tracker.touch("session-a"); + tracker.touch("session-a"); + tracker.touch("session-b"); + tracker.stopped("session-b"); + + assert_eq!(tracker.drain(), vec!["session-a"]); + assert!(tracker.drain().is_empty()); + } + + #[test] + fn tracker_reclaims_nothing_after_graceful_stop() { + let mut tracker = ConnectionSessionTracker::default(); + tracker.touch("session-a"); + tracker.stopped("session-a"); + assert!(tracker.drain().is_empty()); + } + + #[test] + fn cleanup_failure_is_returned_when_connection_succeeded() { + let error = + combine_connection_and_cleanup_results(Ok(()), Err(anyhow!("orphan cleanup failed"))) + .expect_err("cleanup failure must reach the connection supervisor"); + + assert!(error.to_string().contains("orphan cleanup failed")); + } + + #[test] + fn connection_and_cleanup_failures_are_both_preserved() { + let error = combine_connection_and_cleanup_results( + Err(anyhow!("connection failed")), + Err(anyhow!("orphan cleanup failed")), + ) + .expect_err("combined lifecycle failures must be returned"); + let message = format!("{error:#}"); + + assert!(message.contains("connection failed")); + assert!(message.contains("orphan cleanup failed")); + } +} diff --git a/crates/skippy-server/src/binary_transport/binary_messaging/telemetry.rs b/crates/skippy-server/src/binary_transport/binary_messaging/telemetry.rs index a3cba4a44..bb6115999 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/telemetry.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/telemetry.rs @@ -1,5 +1,6 @@ use crate::binary_transport::stage_execution::binary_message_attrs; use crate::binary_transport::stage_execution::estimated_reply_wire_bytes; +use crate::binary_transport::stage_execution::insert_optional_unix_nanos; use crate::binary_transport::stage_execution::ms_to_us; use crate::runtime_state::RuntimeSessionStats; use crate::telemetry::Telemetry; @@ -20,6 +21,200 @@ pub(super) struct UpstreamReplyWriteSpan { pub(super) write_ms: f64, } +pub(super) struct BinaryMessageTiming<'a> { + pub(super) message_start_unix_nanos: u64, + pub(super) message_end_unix_nanos: u64, + pub(super) compute_start_unix_nanos: u64, + pub(super) compute_end_unix_nanos: u64, + pub(super) forward_write_start_unix_nanos: Option, + pub(super) forward_write_end_unix_nanos: Option, + pub(super) downstream_wait_start_unix_nanos: Option, + pub(super) downstream_wait_end_unix_nanos: Option, + pub(super) upstream_reply_start_unix_nanos: Option, + pub(super) upstream_reply_end_unix_nanos: Option, + pub(super) compute_ms: f64, + pub(super) recv_read_ms: f64, + pub(super) input_activation_decode_ms: f64, + pub(super) runtime_lock_wait_ms: f64, + pub(super) runtime_lock_hold_ms: f64, + pub(super) runtime_lock_acquires: usize, + pub(super) runtime_sessions_before: Option<&'a RuntimeSessionStats>, + pub(super) runtime_sessions_after: Option<&'a RuntimeSessionStats>, + pub(super) forward_write_ms: f64, + pub(super) forward_activation_encode_ms: f64, + pub(super) downstream_wait_ms: f64, + pub(super) upstream_reply_ms: f64, + pub(super) forward_mode: &'a str, + pub(super) message_elapsed_ms: f64, + pub(super) input_activation_bytes: usize, + pub(super) output_activation_bytes: usize, + pub(super) max_deferred_prefill_replies: usize, + pub(super) pending_prefill_replies_before: usize, + pub(super) pending_prefill_replies_after: usize, + pub(super) credit_wait_count: usize, + pub(super) deferred_prefill_replies_drained: usize, +} + +pub(super) fn emit_binary_message_timing( + telemetry: &Telemetry, + config: &StageConfig, + session_id: u64, + message: &StageWireMessage, + timing: BinaryMessageTiming<'_>, +) { + if !telemetry.is_debug_enabled() { + return; + } + let mut attrs = binary_message_attrs(config, session_id, message); + attrs.insert( + "llama_stage.message_start_unix_nanos".to_string(), + json!(timing.message_start_unix_nanos), + ); + attrs.insert( + "llama_stage.message_end_unix_nanos".to_string(), + json!(timing.message_end_unix_nanos), + ); + attrs.insert( + "llama_stage.compute_start_unix_nanos".to_string(), + json!(timing.compute_start_unix_nanos), + ); + attrs.insert( + "llama_stage.compute_end_unix_nanos".to_string(), + json!(timing.compute_end_unix_nanos), + ); + attrs.insert( + "llama_stage.compute_ms".to_string(), + json!(timing.compute_ms), + ); + attrs.insert( + "llama_stage.recv_read_ms".to_string(), + json!(timing.recv_read_ms), + ); + attrs.insert( + "skippy.upstream_message_wait_ms".to_string(), + json!(timing.recv_read_ms), + ); + attrs.insert( + "llama_stage.input_activation_decode_ms".to_string(), + json!(timing.input_activation_decode_ms), + ); + attrs.insert( + "llama_stage.runtime_lock_wait_ms".to_string(), + json!(timing.runtime_lock_wait_ms), + ); + attrs.insert( + "llama_stage.runtime_lock_hold_ms".to_string(), + json!(timing.runtime_lock_hold_ms), + ); + attrs.insert( + "llama_stage.runtime_lock_acquires".to_string(), + json!(timing.runtime_lock_acquires), + ); + if let Some(stats) = timing.runtime_sessions_before { + insert_runtime_session_stats(&mut attrs, "llama_stage.runtime_sessions_before", stats); + } + if let Some(stats) = timing.runtime_sessions_after { + insert_runtime_session_stats(&mut attrs, "llama_stage.runtime_sessions_after", stats); + } + attrs.insert( + "llama_stage.forward_write_ms".to_string(), + json!(timing.forward_write_ms), + ); + attrs.insert( + "llama_stage.activation_encode_ms".to_string(), + json!(timing.forward_activation_encode_ms), + ); + attrs.insert( + "llama_stage.downstream_wait_ms".to_string(), + json!(timing.downstream_wait_ms), + ); + attrs.insert("skippy.compute_ms".to_string(), json!(timing.compute_ms)); + attrs.insert( + "skippy.forward_write_ms".to_string(), + json!(timing.forward_write_ms), + ); + attrs.insert( + "skippy.downstream_wait_ms".to_string(), + json!(timing.downstream_wait_ms), + ); + attrs.insert( + "skippy.upstream_reply_ms".to_string(), + json!(timing.upstream_reply_ms), + ); + attrs.insert( + "llama_stage.forward_mode".to_string(), + json!(timing.forward_mode), + ); + insert_optional_unix_nanos( + &mut attrs, + "llama_stage.forward_write_start_unix_nanos", + timing.forward_write_start_unix_nanos, + ); + insert_optional_unix_nanos( + &mut attrs, + "llama_stage.forward_write_end_unix_nanos", + timing.forward_write_end_unix_nanos, + ); + insert_optional_unix_nanos( + &mut attrs, + "llama_stage.downstream_wait_start_unix_nanos", + timing.downstream_wait_start_unix_nanos, + ); + insert_optional_unix_nanos( + &mut attrs, + "llama_stage.downstream_wait_end_unix_nanos", + timing.downstream_wait_end_unix_nanos, + ); + insert_optional_unix_nanos( + &mut attrs, + "llama_stage.upstream_reply_start_unix_nanos", + timing.upstream_reply_start_unix_nanos, + ); + insert_optional_unix_nanos( + &mut attrs, + "llama_stage.upstream_reply_end_unix_nanos", + timing.upstream_reply_end_unix_nanos, + ); + attrs.insert( + "skippy.message_elapsed_ms".to_string(), + json!(timing.message_elapsed_ms), + ); + attrs.insert( + "skippy.input_activation_bytes".to_string(), + json!(timing.input_activation_bytes), + ); + attrs.insert( + "skippy.output_activation_bytes".to_string(), + json!(timing.output_activation_bytes), + ); + attrs.insert( + "skippy.prefill_credit_limit".to_string(), + json!(timing.max_deferred_prefill_replies), + ); + attrs.insert( + "skippy.prefill_pending_replies_before".to_string(), + json!(timing.pending_prefill_replies_before), + ); + attrs.insert( + "skippy.prefill_pending_replies_after".to_string(), + json!(timing.pending_prefill_replies_after), + ); + attrs.insert( + "skippy.prefill_credit_wait_count".to_string(), + json!(timing.credit_wait_count), + ); + attrs.insert( + "skippy.prefill_deferred_replies_drained".to_string(), + json!(timing.deferred_prefill_replies_drained), + ); + telemetry.emit_debug_span( + "stage.binary_message_timing", + attrs, + timing.message_start_unix_nanos, + timing.message_end_unix_nanos, + ); +} + pub(super) fn emit_upstream_reply_write_span( telemetry: &Telemetry, config: &StageConfig, @@ -63,6 +258,52 @@ pub(super) fn emit_upstream_reply_write_span( ); } +#[allow(clippy::too_many_arguments)] +pub(super) fn emit_binary_message_received( + telemetry: &Telemetry, + config: &StageConfig, + session_id: u64, + message: &StageWireMessage, + start_unix_nanos: u64, + end_unix_nanos: u64, + read_ms: f64, +) { + if !telemetry.is_debug_enabled() { + return; + } + let mut attrs = binary_message_attrs(config, session_id, message); + attrs.insert( + "llama_stage.recv_start_unix_nanos".to_string(), + json!(start_unix_nanos), + ); + attrs.insert( + "llama_stage.recv_end_unix_nanos".to_string(), + json!(end_unix_nanos), + ); + attrs.insert("llama_stage.recv_read_ms".to_string(), json!(read_ms)); + attrs.insert( + "skippy.upstream_message_wait_ms".to_string(), + json!(read_ms), + ); + attrs.insert( + "llama_stage.source_stage_index".to_string(), + json!(message.state.source_stage_index), + ); + attrs.insert( + "llama_stage.configured_upstream_stage_index".to_string(), + json!(config.upstream.as_ref().map(|peer| peer.stage_index)), + ); + attrs.insert( + "llama_stage.message_wire_bytes".to_string(), + json!(message.estimated_wire_bytes()), + ); + attrs.insert( + "skippy.activation_bytes".to_string(), + json!(message.activation.len()), + ); + telemetry.emit_debug_span("stage.binary_recv", attrs, start_unix_nanos, end_unix_nanos); +} + pub(super) fn insert_runtime_session_stats( attrs: &mut BTreeMap, prefix: &str, diff --git a/crates/skippy-server/src/binary_transport/options.rs b/crates/skippy-server/src/binary_transport/options.rs index c24cdb6b3..b3a1bc20e 100644 --- a/crates/skippy-server/src/binary_transport/options.rs +++ b/crates/skippy-server/src/binary_transport/options.rs @@ -45,6 +45,7 @@ pub struct EmbeddedOpenAiStageOptions { pub speculative_window: usize, pub adaptive_speculative_window: bool, pub draft_n_gpu_layers: Option, + pub native_mtp_draft_model_path: Option, pub native_mtp_max_tokens: usize, pub native_mtp_min_tokens: usize, pub speculative: SpeculativeDecodeConfig, @@ -99,6 +100,7 @@ impl BinaryStageOptions { speculative_window: args.openai_speculative_window, adaptive_speculative_window: args.openai_adaptive_speculative_window, draft_n_gpu_layers: args.openai_draft_n_gpu_layers, + native_mtp_draft_model_path: args.openai_native_mtp_draft_model_path, native_mtp_max_tokens: 3, native_mtp_min_tokens: 0, speculative: openai_speculative, @@ -256,6 +258,47 @@ mod tests { assert_eq!(openai.speculative, expected); } + #[test] + fn native_mtp_sidecar_path_reaches_the_embedded_stage() { + let dir = tempfile::tempdir().expect("create temp directory"); + let stage_path = dir.path().join("stage.json"); + let sidecar_path = dir.path().join("sidecar-mtp.gguf"); + fs::write( + &stage_path, + serde_json::to_vec(&stage_config()).expect("serialize stage config"), + ) + .expect("write stage config"); + fs::write(&sidecar_path, b"gguf-stub").expect("write sidecar stub"); + + let cli = Cli::try_parse_from([ + "skippy-server", + "serve-binary", + "--config", + stage_path.to_str().expect("UTF-8 stage path"), + "--activation-width", + "2048", + "--openai-bind-addr", + "127.0.0.1:9337", + "--openai-native-mtp-draft-model-path", + sidecar_path.to_str().expect("UTF-8 sidecar path"), + ]) + .expect("parse binary stage CLI"); + let Command::ServeBinary(args) = cli.command else { + panic!("expected serve-binary command"); + }; + + let options = BinaryStageOptions::from_cli_args(args).expect("resolve binary stage"); + let openai = options.openai.expect("embedded OpenAI configuration"); + + // The sidecar attaches MTP heads to the served model; it must not be + // opened as a standalone draft model. + assert_eq!( + openai.native_mtp_draft_model_path.as_deref(), + Some(sidecar_path.as_path()) + ); + assert_eq!(openai.draft_model_path, None); + } + #[test] fn cache_composite_plan_is_json_stable_for_stage_handoff() { let plan = cache_composite_plan(); diff --git a/crates/skippy-server/src/binary_transport/preconnect.rs b/crates/skippy-server/src/binary_transport/preconnect.rs index 4ae9728e4..02e0caa09 100644 --- a/crates/skippy-server/src/binary_transport/preconnect.rs +++ b/crates/skippy-server/src/binary_transport/preconnect.rs @@ -14,7 +14,7 @@ use super::stage_execution::connect_binary_downstream; const WARM_DOWNSTREAM_RETRY_SLEEP: Duration = Duration::from_millis(500); const WARM_DOWNSTREAM_SLOT_POLL: Duration = Duration::from_millis(50); -const WARM_DOWNSTREAM_CONNECT_TIMEOUT_SECS: u64 = 2; +const WARM_DOWNSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); pub(super) fn spawn_downstream_preconnector( config: StageConfig, @@ -40,7 +40,7 @@ fn run_downstream_preconnector( thread::sleep(WARM_DOWNSTREAM_SLOT_POLL); continue; } - match connect_binary_downstream(&config, WARM_DOWNSTREAM_CONNECT_TIMEOUT_SECS) { + match connect_binary_downstream(&config, WARM_DOWNSTREAM_CONNECT_TIMEOUT) { Ok(Some(stream)) => { eprintln!( "downstream warm preconnect ready: stage_id={} local={:?} remote={:?}", diff --git a/crates/skippy-server/src/binary_transport/stage_execution.rs b/crates/skippy-server/src/binary_transport/stage_execution.rs index 8b767b024..9afcfad2b 100644 --- a/crates/skippy-server/src/binary_transport/stage_execution.rs +++ b/crates/skippy-server/src/binary_transport/stage_execution.rs @@ -53,7 +53,7 @@ fn warm_downstream_preconnect_enabled_from(value: Option<&str>) -> bool { pub(in crate::binary_transport) fn take_warm_or_connect_downstream( config: &StageConfig, warm_downstream: &Arc>>, - timeout_secs: u64, + timeout: Duration, ) -> Result> { if config.downstream.is_none() { return Ok(None); @@ -64,10 +64,79 @@ pub(in crate::binary_transport) fn take_warm_or_connect_downstream( .take(); match warm { Some(stream) if warm_downstream_is_healthy(&stream)? => Ok(Some(stream)), - Some(_) | None => connect_binary_downstream(config, timeout_secs), + Some(_) | None => connect_binary_downstream(config, timeout), } } +pub(in crate::binary_transport) fn take_ready_downstream( + config: &StageConfig, + warm_downstream: &Arc>>, + timeout_secs: u64, +) -> Result> { + if config.downstream.is_none() { + return Ok(None); + } + let deadline = Instant::now() + Duration::from_secs(timeout_secs.max(1)); + let mut last_error = None; + loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + break; + } + match take_warm_or_connect_downstream(config, warm_downstream, remaining) { + Ok(Some(mut stream)) => { + let handshake_remaining = deadline.saturating_duration_since(Instant::now()); + if handshake_remaining.is_zero() { + last_error = Some(anyhow!("downstream ready deadline expired after connect")); + continue; + } + match complete_downstream_ready( + &mut stream, + handshake_remaining.min(Duration::from_secs(10)), + ) { + Ok(()) => return Ok(Some(stream)), + Err(error) => last_error = Some(error), + } + } + Ok(None) => return Ok(None), + Err(error) => last_error = Some(error), + } + let retry_sleep = deadline + .saturating_duration_since(Instant::now()) + .min(Duration::from_millis(100)); + if !retry_sleep.is_zero() { + thread::sleep(retry_sleep); + } + } + Err(last_error + .unwrap_or_else(|| anyhow!("downstream ready deadline expired")) + .context(format!( + "downstream stage did not become ready within {}s", + timeout_secs.max(1) + ))) +} + +fn complete_downstream_ready(stream: &mut TcpStream, timeout: Duration) -> Result<()> { + stream + .set_write_timeout(Some(timeout)) + .context("set downstream ready write timeout")?; + let hello_result = + send_client_ready_hello_if_enabled(stream).context("send downstream client ready hello"); + stream + .set_write_timeout(None) + .context("clear downstream ready write timeout")?; + hello_result?; + stream + .set_read_timeout(Some(timeout)) + .context("set downstream ready timeout")?; + let result = skippy_protocol::binary::recv_ready(&mut *stream) + .context("downstream binary stage did not become ready"); + stream + .set_read_timeout(None) + .context("clear downstream ready timeout")?; + result +} + pub(in crate::binary_transport) fn warm_downstream_is_healthy(stream: &TcpStream) -> Result { let previous_timeout = stream .read_timeout() @@ -365,7 +434,7 @@ pub(in crate::binary_transport) fn binary_message_request_id(message: &StageWire } pub(crate) fn connect_binary_downstream( config: &StageConfig, - timeout_secs: u64, + timeout: Duration, ) -> Result> { let Some(peer) = config.downstream.as_ref() else { return Ok(None); @@ -376,17 +445,30 @@ pub(crate) fn connect_binary_downstream( .unwrap_or(&peer.endpoint); let downstream_addr = resolve_downstream_endpoint(endpoint)?; let source_ip = downstream_source_ip(config)?; - let attempts = timeout_secs.saturating_mul(2).max(1); + let deadline = Instant::now() + timeout.max(Duration::from_millis(1)); let mut last_error = None; - for _ in 0..attempts { - match connect_downstream_socket(downstream_addr, source_ip, Duration::from_secs(2)) { + loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + break; + } + match connect_downstream_socket( + downstream_addr, + source_ip, + remaining.min(Duration::from_secs(2)), + ) { Ok(stream) => { stream.set_nodelay(true).ok(); return Ok(Some(stream)); } Err(error) => { last_error = Some(anyhow!(error)); - thread::sleep(Duration::from_millis(500)); + let retry_sleep = deadline + .saturating_duration_since(Instant::now()) + .min(Duration::from_millis(500)); + if !retry_sleep.is_zero() { + thread::sleep(retry_sleep); + } } } } @@ -521,6 +603,7 @@ pub(crate) fn run_binary_stage_message( | WireMessageKind::StateExport | WireMessageKind::ConfigureGeneration | WireMessageKind::TrimSession + | WireMessageKind::RetireVerifyWindow | WireMessageKind::ProbePrefill | WireMessageKind::RestorePrefill | WireMessageKind::TryRestorePrefill @@ -794,8 +877,9 @@ mod tests { use super::{ decode_record_tokens_sideband, first_decode_message_with_full_prompt_sideband, is_decode_frame_batch_candidate, prefix_cache_test_config, prepare_binary_stage_connection, - split_native_mtp_reply, take_warm_or_connect_downstream, token_sideband_or_fill, - warm_downstream_is_healthy, warm_downstream_preconnect_enabled_from, + split_native_mtp_reply, take_ready_downstream, take_warm_or_connect_downstream, + token_sideband_or_fill, warm_downstream_is_healthy, + warm_downstream_preconnect_enabled_from, }; use skippy_protocol::binary::{ StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, @@ -805,7 +889,7 @@ mod tests { net::{Shutdown, TcpListener, TcpStream}, os::fd::AsRawFd, thread, - time::Duration, + time::{Duration, Instant}, }; #[test] @@ -848,9 +932,13 @@ mod tests { let (server, _) = listener.accept().unwrap(); let warm = std::sync::Arc::new(std::sync::Mutex::new(Some(server))); - let result = take_warm_or_connect_downstream(&prefix_cache_test_config(), &warm, 1) - .unwrap() - .unwrap(); + let result = take_warm_or_connect_downstream( + &prefix_cache_test_config(), + &warm, + Duration::from_secs(1), + ) + .unwrap() + .unwrap(); assert_eq!(result.peer_addr().unwrap(), client.local_addr().unwrap()); assert!(warm.lock().unwrap().is_none()); @@ -875,7 +963,7 @@ mod tests { let mut config = prefix_cache_test_config(); config.downstream.as_mut().unwrap().endpoint = endpoint; let warm = std::sync::Arc::new(std::sync::Mutex::new(Some(stale_server))); - let replacement = take_warm_or_connect_downstream(&config, &warm, 1) + let replacement = take_warm_or_connect_downstream(&config, &warm, Duration::from_secs(1)) .unwrap() .unwrap(); let (accepted, _) = listener.accept().unwrap(); @@ -886,6 +974,50 @@ mod tests { ); assert!(warm.lock().unwrap().is_none()); } + #[test] + fn downstream_ready_retries_after_failed_handshake() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let endpoint = listener.local_addr().unwrap().to_string(); + let server = thread::spawn(move || { + let (first, _) = listener.accept().unwrap(); + drop(first); + let (mut second, _) = listener.accept().unwrap(); + skippy_protocol::binary::send_ready(&mut second).unwrap(); + }); + let mut config = prefix_cache_test_config(); + config.downstream.as_mut().unwrap().endpoint = endpoint; + let warm = std::sync::Arc::new(std::sync::Mutex::new(None)); + + let ready = take_ready_downstream(&config, &warm, 2) + .unwrap() + .expect("downstream should be present"); + + assert!(ready.peer_addr().is_ok()); + server.join().unwrap(); + } + + #[test] + fn downstream_ready_honors_absolute_deadline() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let endpoint = listener.local_addr().unwrap().to_string(); + let server = thread::spawn(move || { + let (_stream, _) = listener.accept().unwrap(); + thread::sleep(Duration::from_secs(2)); + }); + let mut config = prefix_cache_test_config(); + config.downstream.as_mut().unwrap().endpoint = endpoint; + let warm = std::sync::Arc::new(std::sync::Mutex::new(None)); + + let started = Instant::now(); + let error = take_ready_downstream(&config, &warm, 1).unwrap_err(); + let elapsed = started.elapsed(); + + assert!(error.to_string().contains("did not become ready")); + assert!(elapsed >= Duration::from_millis(800)); + assert!(elapsed < Duration::from_millis(1500)); + server.join().unwrap(); + } + #[test] fn decode_record_tokens_sideband_records_metadata_without_changing_exec_token() { let message = first_decode_message_with_full_prompt_sideband(); diff --git a/crates/skippy-server/src/cli.rs b/crates/skippy-server/src/cli.rs index 3c0d05d24..030f4ab68 100644 --- a/crates/skippy-server/src/cli.rs +++ b/crates/skippy-server/src/cli.rs @@ -131,6 +131,11 @@ pub struct ServeBinaryArgs { help = "Override n_gpu_layers for the embedded OpenAI draft model. Defaults to the stage config n_gpu_layers." )] pub openai_draft_n_gpu_layers: Option, + #[arg( + long, + help = "Native MTP sidecar GGUF to attach to the stage-0 model. Unlike --openai-draft-model-path this is not opened as a standalone draft model; its MTP heads are attached to the served model." + )] + pub openai_native_mtp_draft_model_path: Option, #[arg( long, help = "JSON file containing the complete resolved speculative decode plan." diff --git a/crates/skippy-server/src/frontend/decode_scheduler.rs b/crates/skippy-server/src/frontend/decode_scheduler.rs index 7150ab3ec..ca04c5441 100644 --- a/crates/skippy-server/src/frontend/decode_scheduler.rs +++ b/crates/skippy-server/src/frontend/decode_scheduler.rs @@ -429,6 +429,26 @@ mod tests { assert!(!scheduler.stats().direct_prediction_return); } + #[test] + fn depth_nine_keeps_the_first_window_restorable() { + let mut scheduler = VerifyWindowScheduler::new(VerifyWindowPipelineConfig { depth: 9 }); + let windows: Vec<_> = (0..9) + .map(|step| scheduler.open(10 + step, step).unwrap()) + .collect(); + + assert_eq!(scheduler.in_flight_len(), 9); + assert!(scheduler.open(19, 9).is_err()); + assert_eq!(scheduler.stats().max_in_flight, 9); + + // The first window must still be completable after the ninth opens. + // Native checkpoint retention has to cover every in-flight window, so + // a partially accepted first window stays restorable. + for window in windows { + assert_eq!(scheduler.complete_next(window.id).unwrap(), window); + } + assert_eq!(scheduler.in_flight_len(), 0); + } + #[test] fn discards_stale_windows_after_divergence() { let config = VerifyWindowPipelineConfig { depth: 3 }; diff --git a/crates/skippy-server/src/frontend/embedded_execution.rs b/crates/skippy-server/src/frontend/embedded_execution.rs index 939705b08..c35bc283b 100644 --- a/crates/skippy-server/src/frontend/embedded_execution.rs +++ b/crates/skippy-server/src/frontend/embedded_execution.rs @@ -12,9 +12,11 @@ use crate::frontend::generation::EmbeddedStageExecution; use crate::frontend::generation::EmbeddedStageZeroGeneration; use crate::frontend::generation::PhaseTimer; use crate::frontend::generation::StageOpenAiBackend; +use crate::frontend::generation::stage_reply_timeout; use crate::frontend::util::ms_to_us; use crate::frontend::util::openai_backend_error; use crate::frontend::util::openai_io_error; +use crate::frontend::wire_messages::retire_verify_window_message; use crate::telemetry::now_unix_nanos; use openai_frontend::OpenAiError; use openai_frontend::OpenAiResult; @@ -36,7 +38,13 @@ const DIRECT_RETURN_FALLBACK_POLL: Duration = Duration::from_millis(10); // a generation permit indefinitely. This is deliberately much larger than a // normal WAN verify traversal while remaining shorter than the HTTP client's // request timeout. -const DIRECT_RETURN_FALLBACK_TIMEOUT: Duration = Duration::from_secs(30); + +pub(super) struct VerifyRetirement { + pub(super) request_id: u64, + pub(super) session_id: u64, + pub(super) token_start: usize, + pub(super) token_count: usize, +} pub(super) struct DispatchedEmbeddedStage { started: Instant, @@ -48,6 +56,56 @@ pub(super) struct DispatchedEmbeddedStage { } impl StageOpenAiBackend { + pub(super) fn retire_verify_window( + &self, + request: &EmbeddedStageZeroGeneration<'_>, + downstream: &mut TcpStream, + async_forwarder: Option<&mut AsyncForwarder>, + session_key: &str, + retirement: VerifyRetirement, + ) -> OpenAiResult<()> { + { + let mut runtime = self.runtime.lock().map_err(|_| { + OpenAiError::backend("runtime lock poisoned during verify retirement") + })?; + runtime + .retire_verify_checkpoint( + session_key, + retirement.token_start as u64, + retirement.token_count as u64, + ) + .map_err(openai_backend_error)?; + } + let message = retire_verify_window_message( + request.wire_dtype, + retirement.request_id, + retirement.session_id, + retirement.token_start, + retirement.token_count, + )?; + if let Some(forwarder) = async_forwarder { + forwarder + .send_tracked( + message, + request.wire_dtype, + request.downstream_wire_condition, + self.openai_attrs(request.ids), + ) + .map_err(openai_backend_error)? + .finish() + .map_err(openai_backend_error)?; + } else { + write_stage_message_conditioned( + downstream, + &message, + request.wire_dtype, + request.downstream_wire_condition, + ) + .map_err(openai_io_error)?; + } + Ok(()) + } + pub(super) fn execute_embedded_stage_message( &self, request: &EmbeddedStageZeroGeneration<'_>, @@ -296,7 +354,7 @@ fn receive_direct_prediction_return( OpenAiError::backend("direct prediction return was required but is not configured") })?; prediction_return - .recv_expected_timeout(expected_reply, DIRECT_RETURN_FALLBACK_TIMEOUT) + .recv_expected_timeout(expected_reply, stage_reply_timeout()) .map_err(openai_backend_error)? .ok_or_else(|| { OpenAiError::backend(format!( @@ -340,6 +398,7 @@ fn poll_direct_or_downstream_reply( ) -> OpenAiResult { let mut timeout_restore = DirectReturnFallbackTimeout::install(downstream)?; let started = Instant::now(); + let reply_timeout = stage_reply_timeout(); let result = loop { if let Some(reply) = prediction_return .try_recv_one_of(expected_replies) @@ -353,13 +412,13 @@ fn poll_direct_or_downstream_reply( // while decoding the complete frame turns an ordinary partial // arrival into EWOULDBLOCK. Once downstream wins the race, give // the frame the remainder of the bounded fallback deadline. - let remaining = DIRECT_RETURN_FALLBACK_TIMEOUT.saturating_sub(started.elapsed()); + let remaining = reply_timeout.saturating_sub(started.elapsed()); downstream .set_read_timeout(Some(remaining.max(DIRECT_RETURN_FALLBACK_POLL))) .map_err(openai_io_error)?; break receive_downstream_stage_reply_one_of(downstream, expected_replies); } - if started.elapsed() >= DIRECT_RETURN_FALLBACK_TIMEOUT { + if started.elapsed() >= reply_timeout { break Err(OpenAiError::backend(format!( "timed out waiting for one of {expected_replies:?} from direct return or downstream" ))); diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 2bb2d2d8a..5f5dfe9bb 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -7,10 +7,11 @@ use crate::binary_transport::{ AsyncForwarder, BinaryStageExecutionOptions, forwarded_stage_message, forwarded_stage_message_timed, run_binary_stage_message, write_stage_message_conditioned, }; +use crate::frontend::embedded_execution::VerifyRetirement; use crate::frontend::request::wire_sampling_config; use crate::frontend::speculative::{ OpenAiSpeculativeStats, classify_verify_window, propose_configured_ngram_tokens, - verify_inputs_for_proposals, + verify_checkpoint_no_longer_needed, verify_inputs_for_proposals, }; use crate::frontend::util::{ ms_to_us, openai_backend_error, openai_io_error, saturating_u32, token_is_eog_with_runtime, @@ -34,8 +35,8 @@ use crate::telemetry::now_unix_nanos; use lifecycle::{ DirectPredictionReturnPath, EmbeddedDecodeSummary, PipelinedCompositeWindow, can_seed_pipeline, compose_target_predictions, decode_uses_context_sideband, direct_prediction_return_path, - mark_epoch_stale, pipelined_window_layout, queued_active_tokens, - refill_pipeline_ngram_candidates, + mark_epoch_stale, open_upstream_prediction_return, pipelined_window_layout, + queued_active_tokens, refill_pipeline_ngram_candidates, speculation_after_prefix_restore, }; use openai_frontend::{OpenAiError, OpenAiResult}; use serde_json::json; @@ -60,25 +61,7 @@ impl StageOpenAiBackend { .as_ref() .ok_or_else(|| OpenAiError::backend("embedded stage 0 has no downstream lane pool"))?; let mut lane = lane_pool.checkout(request.ids)?; - let mut direct_prediction_return_opened = false; - if let Some(prediction_return) = request.prediction_return.as_ref() { - match crate::binary_transport::direct_return::open_downstream_prediction_return_stream( - request.config, - request_id, - session_id, - request.wire_dtype, - ) { - Ok(stream) => { - prediction_return.attach_opened_stream(stream); - direct_prediction_return_opened = true; - } - Err(error) => { - eprintln!( - "direct prediction return upstream-opened sink unavailable: {error:#}" - ); - } - } - } + let direct_prediction_return_opened = open_upstream_prediction_return(&request); let mut cache_stats = GenerationCacheStats::default(); let result = (|| { @@ -711,7 +694,19 @@ impl StageOpenAiBackend { )?; let mut fused_reached_stop = false; let mut native_mtp = NativeMtpVerifier::default(); - let native_mtp_options = NativeMtpDecodeOptions::from_config(request.speculative); + let effective_speculative = + speculation_after_prefix_restore(request.speculative, prefill_chain_cache_restored); + if request.speculative.ngram.is_some() && effective_speculative.ngram.is_none() { + let mut attrs = self.openai_attrs(request.ids); + attrs.insert( + "llama_stage.spec.bypass_reason".to_string(), + json!("distributed_prefix_restored"), + ); + self.telemetry + .emit("stage.openai_speculation_bypass", attrs); + } + let effective_speculative = effective_speculative.as_ref(); + let native_mtp_options = NativeMtpDecodeOptions::from_config(effective_speculative); let mut native_mtp_counters = NativeMtpDecodeCounters::default(); let mut native_mtp_reject_cooldown_remaining = 0usize; let mut native_mtp_suppress_cooldown_drafts_remaining = 0usize; @@ -859,7 +854,8 @@ impl StageOpenAiBackend { } } } - let mut cached_ngram_proposer = HistoryNgramProposer::from_config(request.speculative)?; + let mut cached_ngram_proposer = + HistoryNgramProposer::from_config(effective_speculative)?; let max_speculative_window = request.speculative_window.max(1); let mut adaptive_window = if request.adaptive_speculative_window { max_speculative_window.min(4) @@ -871,7 +867,7 @@ impl StageOpenAiBackend { adaptive_window_final: adaptive_window, adaptive_window_max: max_speculative_window, adaptive_window_min: if request.draft.is_some() - || request.speculative.ngram.is_some() + || effective_speculative.ngram.is_some() { adaptive_window } else { @@ -910,7 +906,7 @@ impl StageOpenAiBackend { _ => None, }; let mut verify_window_scheduler = VerifyWindowScheduler::new( - VerifyWindowPipelineConfig::new(request.speculative.verify_window.pipeline_depth), + VerifyWindowPipelineConfig::new(effective_speculative.verify_window.pipeline_depth), ); let composite_sidecar_enabled = native_mtp_options.ngram_hybrid && draft_guard.is_none(); @@ -919,7 +915,7 @@ impl StageOpenAiBackend { // stays composite-only, so standalone drafting still falls back to the // serial block at depth 1. let standalone_ngram_pipelining = !request.native_mtp_enabled - && request.speculative.ngram.is_some() + && effective_speculative.ngram.is_some() && draft_guard.is_none(); let native_mtp_verify_windows_enabled = (request.native_mtp_enabled || composite_sidecar_enabled) && draft_guard.is_none(); @@ -928,12 +924,12 @@ impl StageOpenAiBackend { && verify_window_scheduler.depth() > 1; let mut verify_window_forwarder = None; if let Some(direct_return_path) = direct_prediction_return_path( - native_mtp_verify_windows_enabled, + native_mtp_verify_windows_enabled || pipelined_decode_enabled, request.prediction_return.is_some(), direct_prediction_return_opened, )? { // The final stage first consumes the upstream-opened sink, then - // falls back to opening the v10 direct-return stream back to the + // falls back to opening the v11 direct-return stream back to the // registered stage-0 receiver. A transient failure opening the // preferred sink must not fail an otherwise healthy request. verify_window_scheduler.mark_direct_prediction_return(matches!( @@ -1340,6 +1336,23 @@ impl StageOpenAiBackend { fully_accepted_window, later_active_window || undispatched_candidates, ); + if verify_checkpoint_no_longer_needed( + commit_count, + window.input_tokens.len(), + ) { + self.retire_verify_window( + &request, + downstream, + verify_window_forwarder.as_mut(), + &session_key, + VerifyRetirement { + request_id, + session_id, + token_start: window.window.base_position, + token_count: window.input_tokens.len(), + }, + )?; + } for token in target_predictions.iter().copied().take(commit_count) { current = token; decoded_tokens += 1; @@ -1451,7 +1464,7 @@ impl StageOpenAiBackend { } } if draft_guard.is_some() - || (request.speculative.ngram.is_some() && !pipelined_decode_enabled) + || (effective_speculative.ngram.is_some() && !pipelined_decode_enabled) { let remaining = (request.max_tokens as usize).saturating_sub(decoded_tokens); if remaining == 0 { @@ -1472,9 +1485,9 @@ impl StageOpenAiBackend { proposal_source = "draft-model"; } } - if draft_tokens.is_empty() && request.speculative.ngram.is_some() { + if draft_tokens.is_empty() && effective_speculative.ngram.is_some() { let proposal = propose_configured_ngram_tokens( - request.speculative, + effective_speculative, &mut cached_ngram_proposer, &context_tokens, proposal_limit.min(request.ngram_max), @@ -1557,6 +1570,24 @@ impl StageOpenAiBackend { request.max_tokens as usize, |token| token_is_eog_with_runtime(&self.runtime, token), )?; + let checkpoint_no_longer_needed = verify_checkpoint_no_longer_needed( + decision.commit_count, + verify_inputs.len(), + ); + if checkpoint_no_longer_needed { + self.retire_verify_window( + &request, + downstream, + None, + &session_key, + VerifyRetirement { + request_id, + session_id, + token_start: prefill_token_count + decoded_tokens, + token_count: verify_inputs.len(), + }, + )?; + } speculative_stats.observe_verify_decision( decision, &mut adaptive_window, diff --git a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs index 028a81be3..9396a3a12 100644 --- a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs +++ b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs @@ -1,4 +1,4 @@ -use std::collections::VecDeque; +use std::{borrow::Cow, collections::VecDeque}; use openai_frontend::{OpenAiError, OpenAiResult}; use serde_json::json; @@ -13,10 +13,28 @@ use crate::frontend::{ EmbeddedStageZeroGeneration, GenerationCacheStats, LocalGeneration, PersistentStageLane, PersistentStageLanePool, PhaseTimer, StageOpenAiBackend, TokenControl, }, - speculative::OpenAiSpeculativeStats, + speculative::{OpenAiSpeculativeStats, SpeculativeDecodeConfig}, util::openai_io_error, }; +/// Keeps the configured speculative plan unless a distributed prefix restore +/// has already populated every stage's session. Pure N-gram verification after +/// that restore currently races the restored session lifecycle, so only the +/// history-based N-gram pieces are suppressed for this request. +pub(super) fn speculation_after_prefix_restore( + config: &SpeculativeDecodeConfig, + prefix_restored: bool, +) -> Cow<'_, SpeculativeDecodeConfig> { + if !prefix_restored || config.ngram.is_none() { + return Cow::Borrowed(config); + } + + let mut safe = config.clone(); + safe.ngram = None; + safe.extension = None; + Cow::Owned(safe) +} + pub(super) struct PipelinedCompositeWindow { pub(super) epoch: u64, pub(super) stale: bool, @@ -102,6 +120,44 @@ pub(super) enum DirectPredictionReturnPath { ReverseFallback, } +pub(super) fn should_open_upstream_prediction_return( + native_mtp_enabled: bool, + standalone_ngram_pipelining: bool, +) -> bool { + native_mtp_enabled || standalone_ngram_pipelining +} + +pub(super) fn open_upstream_prediction_return(request: &EmbeddedStageZeroGeneration<'_>) -> bool { + let standalone_ngram_pipelining = !request.native_mtp_enabled + && request.speculative.ngram.is_some() + && request.draft.is_none() + && request.speculative.verify_window.pipeline_depth > 1; + if !should_open_upstream_prediction_return( + request.native_mtp_enabled, + standalone_ngram_pipelining, + ) { + return false; + } + let Some(prediction_return) = request.prediction_return.as_ref() else { + return false; + }; + match crate::binary_transport::direct_return::open_downstream_prediction_return_stream( + request.config, + request.ids.request_id, + request.ids.session_id, + request.wire_dtype, + ) { + Ok(stream) => { + prediction_return.attach_opened_stream(stream); + true + } + Err(error) => { + eprintln!("direct prediction return upstream-opened sink unavailable: {error:#}"); + false + } + } +} + pub(super) fn direct_prediction_return_path( verify_windows_enabled: bool, receiver_registered: bool, @@ -340,9 +396,15 @@ impl StageOpenAiBackend { // The generation error may be the downstream peer disappearing. // A graceful Stop/ACK exchange would then turn the bounded decode // failure into an unbounded teardown wait. Retire the suspect lane - // immediately; replacement uses its own bounded handshake. + // immediately; replacement uses its own bounded handshake. Drop + // the old stream before opening its replacement so the remote + // connection handler can reclaim the request's execution session + // before the replacement admits another request on a one-lane + // stage. self.drop_embedded_runtime_session(request, session_key); - lane_pool.replace_lane(lane.id); + let lane_id = lane.id; + drop(lane); + lane_pool.replace_lane(lane_id); return Ok(()); } @@ -372,7 +434,10 @@ impl StageOpenAiBackend { let stop_result = stop_result.map_err(openai_io_error); match &stop_result { Ok(_) => lane_pool.return_lane(lane), - Err(_) => lane_pool.replace_lane(lane_id), + Err(_) => { + drop(lane); + lane_pool.replace_lane(lane_id); + } } stop_result?; Ok(()) @@ -433,7 +498,45 @@ pub(super) fn decode_uses_context_sideband( #[cfg(test)] mod tests { use super::*; - use crate::frontend::NativeMtpHybridProposal; + use crate::frontend::{ + NativeMtpHybridProposal, + speculative::{NgramProposalConfig, NgramProposerKind}, + }; + + fn ngram_config() -> SpeculativeDecodeConfig { + SpeculativeDecodeConfig { + requested_strategy: "ngram".to_string(), + effective_strategy: "ngram-suffix".to_string(), + ngram: Some(NgramProposalConfig { + kind: NgramProposerKind::Suffix, + min_ngram: 2, + max_ngram: 16, + max_proposal_tokens: 4, + }), + ..SpeculativeDecodeConfig::default() + } + } + + #[test] + fn restored_prefix_bypasses_history_ngram_for_that_request() { + let config = ngram_config(); + + let effective = speculation_after_prefix_restore(&config, true); + + assert!(matches!(effective, Cow::Owned(_))); + assert!(effective.ngram.is_none()); + assert!(config.ngram.is_some()); + } + + #[test] + fn cache_miss_preserves_configured_ngram() { + let config = ngram_config(); + + let effective = speculation_after_prefix_restore(&config, false); + + assert!(matches!(effective, Cow::Borrowed(_))); + assert!(effective.ngram.is_some()); + } #[test] fn direct_return_falls_back_only_with_a_registered_receiver() { @@ -448,6 +551,13 @@ mod tests { ); } + #[test] + fn direct_return_opens_for_native_mtp_or_pipelined_ngram() { + assert!(!should_open_upstream_prediction_return(false, false)); + assert!(should_open_upstream_prediction_return(true, false)); + assert!(should_open_upstream_prediction_return(false, true)); + } + #[test] fn refills_from_an_optimistic_suffix_without_indexing_it() { let committed = vec![1, 2, 3, 1, 2, 3, 1, 2]; diff --git a/crates/skippy-server/src/frontend/generation.rs b/crates/skippy-server/src/frontend/generation.rs index 1378d475b..3d1f569c4 100644 --- a/crates/skippy-server/src/frontend/generation.rs +++ b/crates/skippy-server/src/frontend/generation.rs @@ -5,6 +5,7 @@ mod persistent_lanes; mod queue; mod server; mod streaming; +mod timeouts; mod types; pub use cache_hints::{CONTEXT_BUDGET_MAX_TOKENS, DEFAULT_EMBEDDED_MAX_TOKENS}; @@ -26,4 +27,5 @@ pub(in crate::frontend) use parsing::*; pub(in crate::frontend) use persistent_lanes::*; pub(in crate::frontend) use queue::*; pub(in crate::frontend) use streaming::*; +pub(in crate::frontend) use timeouts::*; pub(in crate::frontend) use types::*; diff --git a/crates/skippy-server/src/frontend/generation/persistent_lanes.rs b/crates/skippy-server/src/frontend/generation/persistent_lanes.rs index 4682a1aec..16e0f54ce 100644 --- a/crates/skippy-server/src/frontend/generation/persistent_lanes.rs +++ b/crates/skippy-server/src/frontend/generation/persistent_lanes.rs @@ -55,7 +55,6 @@ pub(in crate::frontend) struct PrefillTransportEstimate { /// installed on the persistent lane. pub(in crate::frontend) const LANE_READY_READ_TIMEOUT: Duration = Duration::from_secs(20); pub(in crate::frontend) const LANE_STEADY_CONNECT_TIMEOUT: Duration = Duration::from_secs(3); -pub(in crate::frontend) const LANE_STEADY_IO_TIMEOUT: Duration = Duration::from_secs(30); impl PersistentStageLanePool { const PREFILL_TRANSPORT_EWMA_ALPHA: f64 = 0.25; @@ -310,7 +309,7 @@ impl PersistentStageLanePool { connect_timeout: Duration, ready_timeout: Duration, ) -> Result { - let mut stream = connect_binary_downstream(&self.config, connect_timeout.as_secs().max(1))? + let mut stream = connect_binary_downstream(&self.config, connect_timeout)? .ok_or_else(|| anyhow!("embedded stage0 has no downstream"))?; let local_addr = stream.local_addr().ok(); let peer_addr = stream.peer_addr().ok(); @@ -333,11 +332,12 @@ impl PersistentStageLanePool { pub(in crate::frontend) fn configure_persistent_lane_io_deadlines( stream: &TcpStream, ) -> Result<()> { + let timeout = super::stage_reply_timeout(); stream - .set_read_timeout(Some(LANE_STEADY_IO_TIMEOUT)) + .set_read_timeout(Some(timeout)) .context("set persistent downstream lane read timeout")?; stream - .set_write_timeout(Some(LANE_STEADY_IO_TIMEOUT)) + .set_write_timeout(Some(timeout)) .context("set persistent downstream lane write timeout") } diff --git a/crates/skippy-server/src/frontend/generation/server.rs b/crates/skippy-server/src/frontend/generation/server.rs index c802e39d0..58539c0a0 100644 --- a/crates/skippy-server/src/frontend/generation/server.rs +++ b/crates/skippy-server/src/frontend/generation/server.rs @@ -23,7 +23,7 @@ use crate::frontend::prefill::PrefillChunkPolicyArgs; use crate::frontend::speculative::{ SpeculativeDecodeConfig, load_standalone_speculative_config, standalone_ngram_proposal_limit, }; -use crate::kv_integration::{KvStageIntegration, model_requires_recurrent_state}; +use crate::kv_integration::KvStageIntegration; use crate::runtime_state::RuntimeState; use crate::runtime_state::load_runtime; use crate::telemetry::Telemetry; @@ -308,14 +308,9 @@ pub fn embedded_openai_backend(args: EmbeddedOpenAiArgs) -> Result Duration { + stage_reply_timeout_from(std::env::var(STAGE_REPLY_TIMEOUT_ENV).ok().as_deref()) +} + +fn stage_reply_timeout_from(value: Option<&str>) -> Duration { + let seconds = value + .and_then(|value| value.trim().parse::().ok()) + .unwrap_or(DEFAULT_STAGE_REPLY_TIMEOUT_SECS) + .clamp(MIN_STAGE_REPLY_TIMEOUT_SECS, MAX_STAGE_REPLY_TIMEOUT_SECS); + Duration::from_secs(seconds) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn reply_timeout_defaults_to_thirty_seconds() { + assert_eq!(stage_reply_timeout_from(None), Duration::from_secs(30)); + assert_eq!( + stage_reply_timeout_from(Some("invalid")), + Duration::from_secs(30) + ); + } + + #[test] + fn reply_timeout_accepts_and_bounds_override() { + assert_eq!( + stage_reply_timeout_from(Some(" 180 ")), + Duration::from_secs(180) + ); + assert_eq!(stage_reply_timeout_from(Some("0")), Duration::from_secs(1)); + assert_eq!( + stage_reply_timeout_from(Some("99999")), + Duration::from_secs(3600) + ); + } +} diff --git a/crates/skippy-server/src/frontend/generation_flow.rs b/crates/skippy-server/src/frontend/generation_flow.rs index 6f0bc3a90..9ee8a82db 100644 --- a/crates/skippy-server/src/frontend/generation_flow.rs +++ b/crates/skippy-server/src/frontend/generation_flow.rs @@ -740,6 +740,7 @@ impl StageOpenAiBackend { prompt_token_count: prefill.token_count, pos_start: prefill_pos_start, token_count: chunk.token_count, + tokens: chunk.tokens.clone(), positions: chunk.positions.clone(), sampling: is_final_chunk.then_some(wire_sampling.clone()).flatten(), final_chunk: is_final_chunk, diff --git a/crates/skippy-server/src/frontend/linear_proposal.rs b/crates/skippy-server/src/frontend/linear_proposal.rs index c7502a2c6..fbdf37764 100644 --- a/crates/skippy-server/src/frontend/linear_proposal.rs +++ b/crates/skippy-server/src/frontend/linear_proposal.rs @@ -643,6 +643,8 @@ impl StageOpenAiBackend { let repair = finish_linear_proposal_after_repair(callback_error, || { self.trim_branch_suffix_or_retire( params.session_id, + params.base_position, + verify_inputs.len(), canonical_position, position_after_verification, ) @@ -666,10 +668,18 @@ impl StageOpenAiBackend { fn trim_branch_suffix_or_retire( &self, session_id: &str, + checkpoint_start: u64, + checkpoint_count: usize, canonical_position: u64, position_after_verification: u64, ) -> OpenAiResult { if canonical_position >= position_after_verification { + let mut runtime = self.runtime.lock().map_err(|_| { + OpenAiError::backend("runtime lock poisoned during verify retirement") + })?; + runtime + .retire_verify_checkpoint(session_id, checkpoint_start, checkpoint_count as u64) + .map_err(openai_backend_error)?; return Ok(LinearProposalRepairTiming::default()); } diff --git a/crates/skippy-server/src/frontend/native_mtp/verify_window.rs b/crates/skippy-server/src/frontend/native_mtp/verify_window.rs index b6852d3a7..7049364e8 100644 --- a/crates/skippy-server/src/frontend/native_mtp/verify_window.rs +++ b/crates/skippy-server/src/frontend/native_mtp/verify_window.rs @@ -3,6 +3,8 @@ use std::net::TcpStream; use openai_frontend::{OpenAiError, OpenAiResult}; use skippy_protocol::binary::{StageNativeMtpDraft, WireReplyKind}; +use crate::frontend::embedded_execution::VerifyRetirement; + use super::super::{ AdaptiveVerifyWindow, BufferedCompositeProposal, CompositeProposalProvider, EmbeddedStageZeroGeneration, HistoryNgramProposer, NativeMtpDecodeCounters, @@ -10,7 +12,7 @@ use super::super::{ NgramSidecarController, PendingNativeMtpDraft, PhaseTimer, StageOpenAiBackend, TokenControl, VerifyWindowMessageArgs, VerifyWindowScheduler, WireSamplingConfig, classify_native_mtp_verify_window, embedded_verify_window_message, ms_to_us, - token_is_eog_with_runtime, + token_is_eog_with_runtime, verify_checkpoint_no_longer_needed, }; /// Control signal returned after processing a batched native MTP verify step. @@ -207,6 +209,22 @@ impl StageOpenAiBackend { && native_mtp_verify_decision.accepted_proposal_tokens == proposal_tokens.len() && committed_positions == consumed_positions && !reached_stop; + let checkpoint_no_longer_needed = + verify_checkpoint_no_longer_needed(committed_positions, consumed_positions); + if checkpoint_no_longer_needed { + self.retire_verify_window( + request, + downstream, + None, + session_key, + VerifyRetirement { + request_id, + session_id, + token_start: window.base_position, + token_count: verify_inputs.len(), + }, + )?; + } let decision_rejected_native_mtp_prefix = proposal_buffer.as_ref().is_some_and(|buffer| { buffer.native_mtp_prefix_rejected_after( native_mtp_verify_decision.accepted_proposal_tokens, diff --git a/crates/skippy-server/src/frontend/prefix_cache.rs b/crates/skippy-server/src/frontend/prefix_cache.rs index be25393b0..bc982bc0e 100644 --- a/crates/skippy-server/src/frontend/prefix_cache.rs +++ b/crates/skippy-server/src/frontend/prefix_cache.rs @@ -719,7 +719,12 @@ impl StageOpenAiBackend { else { continue; }; - if restore.restored_tokens < checkpoint_tokens.len() { + if exact_replay_restore_is_partial(restore.restored_tokens, checkpoint_tokens.len()) { + // This trial activated the same request session on every stage. + // Retire it before trying the next-shorter replay checkpoint; + // otherwise the next local restore collides with the trial + // session and fails with `session ... already exists`. + self.drop_embedded_split_restore(request, session_key, downstream); continue; } let replay = replay_tokens[..replay_len].to_vec(); @@ -1193,6 +1198,10 @@ impl StageOpenAiBackend { } } +fn exact_replay_restore_is_partial(restored_tokens: usize, checkpoint_tokens: usize) -> bool { + restored_tokens < checkpoint_tokens +} + #[cfg(test)] mod tests { use super::*; @@ -1232,4 +1241,10 @@ mod tests { assert_eq!(f32.stage0_activation_bytes_avoided, 5_242_880); assert_eq!(f32.interstage_activation_bytes_avoided_estimate, 5_242_880); } + + #[test] + fn exact_replay_rejects_a_shorter_restored_checkpoint() { + assert!(exact_replay_restore_is_partial(44_466, 44_467)); + assert!(!exact_replay_restore_is_partial(44_467, 44_467)); + } } diff --git a/crates/skippy-server/src/frontend/request.rs b/crates/skippy-server/src/frontend/request.rs index 77234ff57..fb7c800c7 100644 --- a/crates/skippy-server/src/frontend/request.rs +++ b/crates/skippy-server/src/frontend/request.rs @@ -1,4 +1,6 @@ use crate::frontend::EmbeddedOpenAiRequestDefaults; +use crate::frontend::EmbeddedReasoningBudget; +use crate::frontend::EmbeddedReasoningEnabled; use crate::frontend::EmbeddedReasoningFormat; use base64::Engine; use openai_frontend::ChatCompletionRequest; @@ -285,11 +287,32 @@ pub(super) fn chat_template_options( )?; Ok(ChatTemplateOptions { reasoning_format: Some(chat_reasoning_format(defaults.reasoning_format)), - enable_thinking: reasoning.enable_thinking, + enable_thinking: reasoning + .enable_thinking + .or_else(|| default_reasoning_enabled(defaults.reasoning_enabled)) + .or_else(|| default_reasoning_budget_enabled(defaults.reasoning_budget)), ..ChatTemplateOptions::default() }) } +fn default_reasoning_enabled(value: Option) -> Option { + match value { + Some(EmbeddedReasoningEnabled::Disabled) => Some(false), + Some(EmbeddedReasoningEnabled::Enabled) => Some(true), + Some(EmbeddedReasoningEnabled::Auto) | None => None, + } +} + +fn default_reasoning_budget_enabled(value: Option) -> Option { + match value { + Some(EmbeddedReasoningBudget::Tokens(0)) => Some(false), + Some(EmbeddedReasoningBudget::Tokens(_)) | Some(EmbeddedReasoningBudget::Effort(_)) => { + Some(true) + } + Some(EmbeddedReasoningBudget::Auto) | None => None, + } +} + fn chat_reasoning_format(value: Option) -> ChatReasoningFormat { match value.unwrap_or(EmbeddedReasoningFormat::Hidden) { EmbeddedReasoningFormat::Auto => ChatReasoningFormat::Auto, diff --git a/crates/skippy-server/src/frontend/speculative.rs b/crates/skippy-server/src/frontend/speculative.rs index e85d6b054..6abb8e61e 100644 --- a/crates/skippy-server/src/frontend/speculative.rs +++ b/crates/skippy-server/src/frontend/speculative.rs @@ -4,6 +4,7 @@ use openai_frontend::OpenAiResult; use serde::{Deserialize, Serialize}; use serde_json::Value; use serde_json::json; +use skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH; use std::collections::BTreeMap; use std::path::PathBuf; use std::time::Instant; @@ -166,8 +167,11 @@ impl SpeculativeDecodeConfig { if self.verify_window.min_tokens == 0 || self.verify_window.min_tokens > self.verify_window.max_tokens || self.verify_window.pipeline_depth == 0 + || self.verify_window.pipeline_depth > MAX_VERIFY_WINDOW_PIPELINE_DEPTH { - bail!("verify window requires 0 < min_tokens <= max_tokens and pipeline_depth > 0"); + bail!( + "verify window requires 0 < min_tokens <= max_tokens and 0 < pipeline_depth <= {MAX_VERIFY_WINDOW_PIPELINE_DEPTH}" + ); } Ok(()) } @@ -218,6 +222,31 @@ mod standalone_speculative_config_tests { ); } + #[test] + fn verify_window_depth_is_bounded_by_native_checkpoint_retention() { + let mut config = SpeculativeDecodeConfig::default(); + config.verify_window.pipeline_depth = MAX_VERIFY_WINDOW_PIPELINE_DEPTH; + config + .validate() + .expect("native retention boundary should be accepted"); + + config.verify_window.pipeline_depth = MAX_VERIFY_WINDOW_PIPELINE_DEPTH + 1; + let error = config + .validate() + .expect_err("depth above native retention must fail"); + + assert!(error.to_string().contains(&format!( + "pipeline_depth <= {MAX_VERIFY_WINDOW_PIPELINE_DEPTH}" + ))); + } + + #[test] + fn checkpoint_retires_only_when_no_verified_suffix_remains() { + assert!(verify_checkpoint_no_longer_needed(4, 4)); + assert!(verify_checkpoint_no_longer_needed(5, 4)); + assert!(!verify_checkpoint_no_longer_needed(3, 4)); + } + #[test] fn standalone_speculative_config_round_trips_cache_composite() { let config = SpeculativeDecodeConfig { @@ -936,6 +965,13 @@ where }) } +pub(super) fn verify_checkpoint_no_longer_needed( + committed_positions: usize, + consumed_positions: usize, +) -> bool { + committed_positions >= consumed_positions +} + pub(super) fn nonzero_min(current: usize, candidate: usize) -> usize { if current == 0 { candidate diff --git a/crates/skippy-server/src/frontend/tests/prefill.rs b/crates/skippy-server/src/frontend/tests/prefill.rs index f6c9f51c5..7117b7353 100644 --- a/crates/skippy-server/src/frontend/tests/prefill.rs +++ b/crates/skippy-server/src/frontend/tests/prefill.rs @@ -232,11 +232,8 @@ fn persistent_lane_steady_state_io_is_bounded() { configure_persistent_lane_io_deadlines(&client).unwrap(); - assert_eq!(client.read_timeout().unwrap(), Some(LANE_STEADY_IO_TIMEOUT)); - assert_eq!( - client.write_timeout().unwrap(), - Some(LANE_STEADY_IO_TIMEOUT) - ); + assert_eq!(client.read_timeout().unwrap(), Some(stage_reply_timeout())); + assert_eq!(client.write_timeout().unwrap(), Some(stage_reply_timeout())); drop(peer); } diff --git a/crates/skippy-server/src/frontend/tests/request.rs b/crates/skippy-server/src/frontend/tests/request.rs index 2402681e8..bdaedb938 100644 --- a/crates/skippy-server/src/frontend/tests/request.rs +++ b/crates/skippy-server/src/frontend/tests/request.rs @@ -152,7 +152,7 @@ fn request_defaults_fill_omitted_chat_fields_only() { assert_eq!(sampling.penalty_last_n, 64); assert_eq!(sampling.logit_bias.len(), 2); let template_options = chat_template_options(&request, &test_request_defaults()).unwrap(); - assert_eq!(template_options.enable_thinking, None); + assert_eq!(template_options.enable_thinking, Some(true)); assert_eq!( template_options.reasoning_format, Some(ChatReasoningFormat::Hidden) @@ -345,6 +345,68 @@ fn chat_template_options_default_to_hidden_reasoning_parser() { assert_eq!(options.reasoning_format, Some(ChatReasoningFormat::Hidden)); } +#[test] +fn request_default_reasoning_enabled_controls_chat_template() { + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "test", + "messages": [{"role": "user", "content": "hello"}] + })) + .unwrap(); + + for (configured, expected) in [ + (EmbeddedReasoningEnabled::Disabled, Some(false)), + (EmbeddedReasoningEnabled::Enabled, Some(true)), + (EmbeddedReasoningEnabled::Auto, None), + ] { + let defaults = EmbeddedOpenAiRequestDefaults { + reasoning_enabled: Some(configured), + ..EmbeddedOpenAiRequestDefaults::default() + }; + let options = chat_template_options(&request, &defaults).expect("template options"); + assert_eq!(options.enable_thinking, expected); + } +} + +#[test] +fn explicit_request_reasoning_overrides_configured_default() { + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "test", + "messages": [{"role": "user", "content": "hello"}], + "reasoning": {"enabled": true} + })) + .unwrap(); + let defaults = EmbeddedOpenAiRequestDefaults { + reasoning_enabled: Some(EmbeddedReasoningEnabled::Disabled), + ..EmbeddedOpenAiRequestDefaults::default() + }; + + let options = chat_template_options(&request, &defaults).expect("template options"); + + assert_eq!(options.enable_thinking, Some(true)); +} + +#[test] +fn request_default_reasoning_budget_controls_chat_template() { + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "test", + "messages": [{"role": "user", "content": "hello"}] + })) + .unwrap(); + + for (configured, expected) in [ + (EmbeddedReasoningBudget::Tokens(0), Some(false)), + (EmbeddedReasoningBudget::Tokens(256), Some(true)), + (EmbeddedReasoningBudget::Auto, None), + ] { + let defaults = EmbeddedOpenAiRequestDefaults { + reasoning_budget: Some(configured), + ..EmbeddedOpenAiRequestDefaults::default() + }; + let options = chat_template_options(&request, &defaults).expect("template options"); + assert_eq!(options.enable_thinking, expected); + } +} + #[test] fn request_default_reasoning_format_controls_chat_parser_mode() { let request: ChatCompletionRequest = serde_json::from_value(json!({ diff --git a/crates/skippy-server/src/frontend/tests/wire_messages.rs b/crates/skippy-server/src/frontend/tests/wire_messages.rs index 1680a1490..c4eb56811 100644 --- a/crates/skippy-server/src/frontend/tests/wire_messages.rs +++ b/crates/skippy-server/src/frontend/tests/wire_messages.rs @@ -16,6 +16,7 @@ fn multimodal_final_prefill_message_requests_downstream_prediction() { prompt_token_count: 17, pos_start: 0, token_count: 17, + tokens: vec![3; 17], positions: Vec::new(), sampling: Some(sampling.clone()), final_chunk: true, @@ -27,6 +28,7 @@ fn multimodal_final_prefill_message_requests_downstream_prediction() { assert!(message.kind.requires_predicted_reply()); assert_eq!(message.token_count, 17); assert_eq!(message.state.current_token, LLAMA_TOKEN_NULL); + assert_eq!(message.tokens, vec![3; 17]); assert_eq!(message.sampling, Some(sampling)); } diff --git a/crates/skippy-server/src/frontend/tool_emulation.rs b/crates/skippy-server/src/frontend/tool_emulation.rs index 9ac7930cc..29a276a8e 100644 --- a/crates/skippy-server/src/frontend/tool_emulation.rs +++ b/crates/skippy-server/src/frontend/tool_emulation.rs @@ -42,11 +42,11 @@ pub(super) const TOOL_CALL_MARKER: &str = "TOOL_CALL"; /// true for every tools request) and `chat_parser` is always a non-empty PEG /// structure, so neither field distinguishes a tool-capable template. /// -/// The signal that does distinguish them is `grammar_triggers`: when the jinja -/// template natively describes tool calls, applying it with tools yields a -/// tool-call grammar trigger (e.g. ``). A template with no native -/// tool support (for example SmolLM2-135M) yields an empty `grammar_triggers` -/// list. We treat a non-empty `grammar_triggers` as native tool-call support. +/// Native tool templates expose either a sampling grammar trigger or semantic +/// `tool*` tags in the serialized PEG response parser. The latter matters for +/// formats such as Inkling: its PEG parser extracts tool calls, but it does not +/// install a lazy sampling grammar. A template with no native tool support (for +/// example SmolLM2-135M) exposes neither signal. pub(super) fn template_supports_native_tool_calls(metadata_json: &str) -> bool { let Ok(metadata) = serde_json::from_str::(metadata_json) else { return false; @@ -55,6 +55,28 @@ pub(super) fn template_supports_native_tool_calls(metadata_json: &str) -> bool { .get("grammar_triggers") .and_then(Value::as_array) .is_some_and(|triggers| !triggers.is_empty()) + || chat_parser_has_tool_semantics(&metadata) +} + +fn chat_parser_has_tool_semantics(metadata: &Value) -> bool { + let Some(serialized_parser) = metadata.get("chat_parser").and_then(Value::as_str) else { + return false; + }; + let Ok(parser) = serde_json::from_str::(serialized_parser) else { + return false; + }; + parser + .get("parsers") + .and_then(Value::as_array) + .is_some_and(|nodes| { + nodes.iter().any(|node| { + node.get("type").and_then(Value::as_str) == Some("tag") + && node + .get("tag") + .and_then(Value::as_str) + .is_some_and(|tag| tag == "tool" || tag.starts_with("tool-")) + }) + }) } /// Environment override, mirroring goose's `ToolCallingMode::ForceEmulated`. @@ -532,15 +554,35 @@ mod tests { } #[test] - fn native_detection_uses_grammar_triggers() { + fn native_detection_uses_grammar_triggers_or_parser_semantics() { // Tool-capable template: applying it yields a tool-call grammar trigger. assert!(template_supports_native_tool_calls( r#"{"chat_format": 2, "grammar_triggers": [{"type": 1, "value": ""}]}"# )); + let inkling_parser = serde_json::json!({ + "parsers": [ + {"type": "literal", "literal": "<|content_invoke_tool_json|>"}, + {"type": "tag", "child": 0, "tag": "tool-name"}, + {"type": "tag", "child": 0, "tag": "tool-args"} + ], + "rules": {}, + "root": 0 + }); + let inkling_metadata = serde_json::json!({ + "chat_format": 2, + "grammar_triggers": [], + "chat_parser": inkling_parser.to_string() + }); + assert!(template_supports_native_tool_calls( + &inkling_metadata.to_string() + )); // Non-tool-capable template (e.g. SmolLM2-135M): empty grammar triggers. assert!(!template_supports_native_tool_calls( r#"{"chat_format": 2, "grammar_triggers": []}"# )); + assert!(!template_supports_native_tool_calls( + r#"{"grammar_triggers": [], "chat_parser": "{\"parsers\":[{\"type\":\"tag\",\"tag\":\"content\",\"child\":0}]}"}"# + )); // Missing field or non-array is treated as non-native. assert!(!template_supports_native_tool_calls( r#"{"chat_format": 2}"# diff --git a/crates/skippy-server/src/frontend/wire_messages.rs b/crates/skippy-server/src/frontend/wire_messages.rs index 7d82e7a75..4fc48b42b 100644 --- a/crates/skippy-server/src/frontend/wire_messages.rs +++ b/crates/skippy-server/src/frontend/wire_messages.rs @@ -174,6 +174,32 @@ pub(super) fn embedded_verify_window_message( }) } +pub(super) fn retire_verify_window_message( + wire_dtype: WireActivationDType, + request_id: u64, + session_id: u64, + token_start: usize, + token_count: usize, +) -> OpenAiResult { + let kind = WireMessageKind::RetireVerifyWindow; + Ok(StageWireMessage { + kind, + pos_start: i32::try_from(token_start) + .map_err(|_| OpenAiError::backend("verify retirement position exceeds i32"))?, + token_count: i32::try_from(token_count) + .map_err(|_| OpenAiError::backend("verify retirement count exceeds i32"))?, + state: StageStateHeader::new(kind, wire_dtype), + request_id, + session_id, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + }) +} + pub(super) fn generation_config_message( wire_dtype: WireActivationDType, request_id: u64, @@ -322,6 +348,7 @@ pub(super) struct MultimodalPrefillArgs { pub(super) prompt_token_count: usize, pub(super) pos_start: usize, pub(super) token_count: usize, + pub(super) tokens: Vec, pub(super) positions: Vec, pub(super) sampling: Option, pub(super) final_chunk: bool, @@ -353,7 +380,7 @@ pub(super) fn multimodal_prefill_message( session_id: args.session_id, sampling: args.sampling, chat_sampling_metadata: None, - tokens: Vec::new(), + tokens: args.tokens, positions: args.positions, activation: Vec::new(), raw_bytes: Vec::new(), diff --git a/crates/skippy-server/src/kv_integration/mod.rs b/crates/skippy-server/src/kv_integration/mod.rs index 0c364ad55..15d7158d9 100644 --- a/crates/skippy-server/src/kv_integration/mod.rs +++ b/crates/skippy-server/src/kv_integration/mod.rs @@ -23,8 +23,6 @@ mod identity; mod records; mod resident_prefix; -pub(crate) use config::model_requires_recurrent_state; - pub use records::{ AttachedPage, ExactStateRecord, ExactStateRestore, LookupBatchOutcome, PrefillKvIdentity, RecordPageOutcome, ResidentActivationRecord, ResidentActivationRestore, ResidentPrefixRecord, diff --git a/crates/skippy-server/src/runtime_state.rs b/crates/skippy-server/src/runtime_state.rs index c137a6aaa..97348491b 100644 --- a/crates/skippy-server/src/runtime_state.rs +++ b/crates/skippy-server/src/runtime_state.rs @@ -535,6 +535,16 @@ impl RuntimeState { )) } + pub fn retire_verify_checkpoint( + &mut self, + session_id: &str, + token_start: u64, + token_count: u64, + ) -> Result<()> { + self.active_session(session_id)? + .retire_verify_checkpoint(token_start, token_count) + } + pub fn trim_session(&mut self, session_id: &str, token_count: u64) -> Result<()> { let session = self.session(session_id)?; session.trim_session(token_count)?; @@ -907,10 +917,11 @@ impl RuntimeState { let layer_end = i32::try_from(self.model_layer_end())?; let session = self.session(session_id)?; session.import_state_for_token_count(layer_start, layer_end, bytes, token_count)?; - self.session_token_counts - .entry(session_id.to_string()) - .and_modify(|current| *current = (*current).max(token_count)) - .or_insert(token_count); + record_restored_session_token_count( + &mut self.session_token_counts, + session_id, + token_count, + ); Ok(()) } @@ -938,10 +949,11 @@ impl RuntimeState { let layer_end = i32::try_from(self.model_layer_end())?; let session = self.session(session_id)?; session.import_full_state_for_token_count(layer_start, layer_end, bytes, token_count)?; - self.session_token_counts - .entry(session_id.to_string()) - .and_modify(|current| *current = (*current).max(token_count)) - .or_insert(token_count); + record_restored_session_token_count( + &mut self.session_token_counts, + session_id, + token_count, + ); Ok(()) } @@ -957,10 +969,11 @@ impl RuntimeState { ) -> Result<()> { self.session(session_id)? .import_recurrent_state_for_token_count(bytes, token_count)?; - self.session_token_counts - .entry(session_id.to_string()) - .and_modify(|current| *current = (*current).max(token_count)) - .or_insert(token_count); + record_restored_session_token_count( + &mut self.session_token_counts, + session_id, + token_count, + ); Ok(()) } @@ -1075,6 +1088,18 @@ impl RuntimeState { } } +fn record_restored_session_token_count( + session_token_counts: &mut BTreeMap, + session_id: &str, + token_count: u64, +) { + // A prefix restore can move an existing lane backwards to a shorter + // common prefix. The tracked position must follow the imported native + // state exactly; retaining the previous high-water mark submits the next + // divergent token at the wrong position and makes llama_decode fail. + session_token_counts.insert(session_id.to_string(), token_count); +} + fn split_activation_frame( input: Option<&ActivationFrame>, token_count: usize, @@ -1390,9 +1415,19 @@ mod tests { use super::{ RuntimeLaunchOverrides, create_indexed_lane_resource, load_runtime_with_overrides, - runtime_config_from_stage_config, should_attach_package_projector, + record_restored_session_token_count, runtime_config_from_stage_config, + should_attach_package_projector, }; + #[test] + fn prefix_restore_moves_tracked_position_backwards() { + let mut token_counts = std::collections::BTreeMap::from([("lane-a".to_string(), 3_535)]); + + record_restored_session_token_count(&mut token_counts, "lane-a", 3_530); + + assert_eq!(token_counts.get("lane-a"), Some(&3_530)); + } + #[test] fn create_indexed_lane_resource_keeps_index_available_when_creation_fails() { let mut next_lane_index = 0; diff --git a/crates/skippy-topology/capabilities/reviewed-family-capabilities.json b/crates/skippy-topology/capabilities/reviewed-family-capabilities.json index 35e5cabfa..840640c64 100644 --- a/crates/skippy-topology/capabilities/reviewed-family-capabilities.json +++ b/crates/skippy-topology/capabilities/reviewed-family-capabilities.json @@ -2070,5 +2070,48 @@ "split_constraints": [], "sidebands": [] } + }, + { + "model_id": "meshllm/inkling-UD-Q2_K_XL-layers", + "source_repo": "meshllm/inkling-UD-Q2_K_XL-layers", + "source_revision": "9b4b91a7ddd978dd7a01679bc977f6e53777f2c7", + "distribution_id": "inkling-UD-Q2_K_XL-layers", + "selector": "Q2_K_XL", + "capability": { + "family_id": "inkling", + "layer_count": 66, + "activation_width": 6144, + "default_wire_dtype": "f32", + "q8_wire_validation": "rejected", + "exact_state_mobility": "rejected_too_large", + "recurrent_ranges": [ + { + "start": 0, + "end": 66 + } + ], + "split_constraints": [], + "sidebands": [] + } + }, + { + "model_id": "poolside/Laguna-S-2.1-GGUF:Q4_K_M", + "source_repo": "poolside/Laguna-S-2.1-GGUF", + "source_revision": "edd093522473dc7313b0738d8b4116b7f8b9745f", + "source_file": "laguna-s-2.1-Q4_K_M.gguf", + "canonical_ref": "poolside/Laguna-S-2.1-GGUF@edd093522473dc7313b0738d8b4116b7f8b9745f/laguna-s-2.1-Q4_K_M.gguf", + "distribution_id": "laguna-s-2.1-Q4_K_M", + "selector": "Q4_K_M", + "capability": { + "family_id": "laguna", + "layer_count": 48, + "activation_width": 3072, + "default_wire_dtype": "f16", + "q8_wire_validation": "untested", + "exact_state_mobility": "untested", + "recurrent_ranges": [], + "split_constraints": [], + "sidebands": [] + } } ] diff --git a/crates/skippy-topology/src/family_capability.rs b/crates/skippy-topology/src/family_capability.rs index 05def6beb..045df24f5 100644 --- a/crates/skippy-topology/src/family_capability.rs +++ b/crates/skippy-topology/src/family_capability.rs @@ -125,6 +125,11 @@ pub const STAGE_RUNTIME_LLAMA_FAMILY_EXPECTATIONS: &[StageRuntimeFamilyExpectati family_id: "hunyuan_vl", recurrent_or_hybrid: false, }, + StageRuntimeFamilyExpectation { + llama_architecture: "inkling", + family_id: "inkling", + recurrent_or_hybrid: true, + }, StageRuntimeFamilyExpectation { llama_architecture: "internlm2", family_id: "internlm2", @@ -145,6 +150,11 @@ pub const STAGE_RUNTIME_LLAMA_FAMILY_EXPECTATIONS: &[StageRuntimeFamilyExpectati family_id: "jamba", recurrent_or_hybrid: true, }, + StageRuntimeFamilyExpectation { + llama_architecture: "laguna", + family_id: "laguna", + recurrent_or_hybrid: false, + }, StageRuntimeFamilyExpectation { llama_architecture: "lfm2", family_id: "lfm2", @@ -551,6 +561,16 @@ pub fn qwen3moe_capability(layer_count: u32, activation_width: u32) -> FamilyCap ) } +pub fn laguna_capability(layer_count: u32, activation_width: u32) -> FamilyCapabilityRecord { + dense_family_capability( + "laguna", + layer_count, + activation_width, + WireValidation::Untested, + ExactStateMobility::Untested, + ) +} + pub fn dense_family_capability( family_id: impl Into, layer_count: u32, @@ -753,6 +773,23 @@ pub fn qwen35_series_capability( } } +pub fn inkling_capability(layer_count: u32, activation_width: u32) -> FamilyCapabilityRecord { + FamilyCapabilityRecord { + family_id: "inkling".to_string(), + layer_count, + activation_width, + default_wire_dtype: WireDType::F32, + q8_wire_validation: WireValidation::Rejected, + exact_state_mobility: ExactStateMobility::RejectedTooLarge, + recurrent_ranges: vec![LayerRange { + start: 0, + end: layer_count, + }], + split_constraints: Vec::new(), + sidebands: Vec::new(), + } +} + pub fn qwen3next_capability( layer_count: u32, activation_width: u32, @@ -1081,6 +1118,9 @@ fn infer_mistral_olmo_llama_capability( if compact.contains("olmo") { return Some(olmo_capability(layer_count, activation_width)); } + if compact.contains("laguna") { + return Some(laguna_capability(layer_count, activation_width)); + } if compact.contains("llama") { return Some(llama_capability(layer_count, activation_width)); } @@ -1112,6 +1152,9 @@ fn infer_recurrent_capability( layer_count: u32, activation_width: u32, ) -> Option { + if compact.contains("inkling") { + return Some(inkling_capability(layer_count, activation_width)); + } if compact.contains("kimilinear") { return Some(kimi_linear_capability(layer_count, activation_width)); } diff --git a/crates/skippy-topology/src/lib.rs b/crates/skippy-topology/src/lib.rs index 2015b6ff7..d06b39c2b 100644 --- a/crates/skippy-topology/src/lib.rs +++ b/crates/skippy-topology/src/lib.rs @@ -12,11 +12,11 @@ pub use family_capability::{ deepseek3_capability, dense_attention_layers, dense_family_capability, falcon_h1_capability, falcon_h1_layers, gemma2_capability, gemma3_capability, gemma3n_capability, gemma4_a4b_capability, gemma4_e4b_capability, glm4_capability, glm47_flash_capability, - infer_family_capability, kimi_linear_capability, llama_capability, minimax_m27_capability, - olmo_capability, qwen2moe_capability, qwen3_dense_capability, qwen3moe_capability, - qwen3next_capability, qwen3next_layers, qwen35_series_capability, recurrent_family_capability, - reviewed_capability_for_identity, reviewed_capability_records, rwkv6_capability, - rwkv7_capability, + infer_family_capability, kimi_linear_capability, laguna_capability, llama_capability, + minimax_m27_capability, olmo_capability, qwen2moe_capability, qwen3_dense_capability, + qwen3moe_capability, qwen3next_capability, qwen3next_layers, qwen35_series_capability, + recurrent_family_capability, reviewed_capability_for_identity, reviewed_capability_records, + rwkv6_capability, rwkv7_capability, }; pub use planning::{ classify_layers, plan_contiguous_with_splits, plan_even_contiguous, diff --git a/crates/skippy-topology/src/tests.rs b/crates/skippy-topology/src/tests.rs index ff731c259..a044d44c8 100644 --- a/crates/skippy-topology/src/tests.rs +++ b/crates/skippy-topology/src/tests.rs @@ -921,6 +921,16 @@ fn infers_known_family_capabilities_from_model_identity() { assert_eq!(llama.q8_wire_validation, WireValidation::Validated); assert_eq!(llama.exact_state_mobility, ExactStateMobility::Accepted); + let laguna = infer_family_capability( + "poolside/Laguna-S-2.1-GGUF@edd093522473dc7313b0738d8b4116b7f8b9745f/laguna-s-2.1-Q4_K_M.gguf", + 48, + 3072, + ) + .expect("reviewed Poolside Laguna S 2.1 Q4_K_M"); + assert_eq!(laguna.family_id, "laguna"); + assert_eq!(laguna.default_wire_dtype, WireDType::F16); + assert_eq!(laguna.q8_wire_validation, WireValidation::Untested); + let gemma4_e4b = infer_family_capability( "unsloth/gemma-4-E4B-it-GGUF@315e03409eb1cdde302488d66e586dea1e82aad1/gemma-4-E4B-it-Q4_K_M.gguf", 42, @@ -950,6 +960,20 @@ fn infers_known_family_capabilities_from_model_identity() { .family_id, "qwen3next" ); + let inkling = + infer_family_capability("meshllm/inkling-UD-Q2_K_XL-layers", 66, 6144).expect("inkling"); + assert_eq!(inkling.family_id, "inkling"); + assert_eq!(inkling.default_wire_dtype, WireDType::F32); + assert_eq!(inkling.q8_wire_validation, WireValidation::Rejected); + assert_eq!( + inkling.exact_state_mobility, + ExactStateMobility::RejectedTooLarge + ); + assert_eq!( + inkling.recurrent_ranges, + vec![LayerRange { start: 0, end: 66 }] + ); + let rwkv6 = infer_family_capability("latestissue/rwkv-6-finch-1b6-gguf:Q4_K", 24, 2048).expect("rwkv6"); assert_eq!(rwkv6.family_id, "rwkv6"); @@ -1019,6 +1043,20 @@ fn infers_known_family_capabilities_from_model_identity() { .expect("qwen3moe"); assert_eq!(qwen3moe.family_id, "qwen3moe"); assert_eq!(qwen3moe.q8_wire_validation, WireValidation::Validated); + for identity in [ + "laguna", + "poolside/Laguna-XS-2.1-GGUF:Q4_K_M", + "poolside/Laguna-S-2.1-GGUF:Q4_K_M", + "poolside/Laguna-XS.2-GGUF:Q4_K_M", + "poolside/Laguna-M.1-GGUF:Q4_K_M", + ] { + let laguna = infer_family_capability(identity, 48, 3072) + .unwrap_or_else(|| panic!("failed to infer {identity}")); + assert_eq!(laguna.family_id, "laguna", "{identity}"); + assert_eq!(laguna.q8_wire_validation, WireValidation::Untested); + assert_eq!(laguna.exact_state_mobility, ExactStateMobility::Untested); + assert!(laguna.recurrent_ranges.is_empty()); + } let openai_moe = infer_family_capability("ggml-org/gpt-oss-20b-GGUF:gpt-oss-20b-mxfp4", 24, 2880) .expect("openai_moe/gpt-oss"); diff --git a/docs/LAYER_PACKAGE_REPOS.md b/docs/LAYER_PACKAGE_REPOS.md index 73411f264..30f5980dc 100644 --- a/docs/LAYER_PACKAGE_REPOS.md +++ b/docs/LAYER_PACKAGE_REPOS.md @@ -96,7 +96,8 @@ Minimal GLM-DSA shape: "default": "fixed", "initial_window": 1, "min_window": 1, - "max_window": 1 + "max_window": 1, + "pipeline_depth": 1 } } } @@ -112,7 +113,7 @@ Authoring rule of thumb: | `generation.policy` | Stable semantic execution choices validated for the package. | `profile`, `decode`, `short_prefill`, `long_prefill`, `verify`, `indexshare` | | `generation.policy.experimental` | Named opt-in paths that need package/backend evidence before becoming defaults. | `selected_row_flash`, `moe_weighted_down`, `moe_merged_shared_gate_up` | | `generation.thresholds` | Numeric resolver inputs used to accept, reject, or fall back from a policy. | `short_prefill_max_tokens`, `compact_flash_min_kv`, `dense_mask_max_bytes` | -| `generation.speculative_decoding` | Package-owned native or draft speculation strategy defaults. | `native-mtp-n1`, `prediction_depth`, `layer_indices`, `window_policy` | +| `generation.speculative_decoding` | Package-owned native, N-gram, or draft speculation strategy defaults. | strategy id, proposer bounds, `window_policy`, optional positive `pipeline_depth` | | GGUF metadata | Architecture correctness and tensor layout requirements. | GLM-DSA q/k/v split dimensions, IndexShare roles, MTP tensor presence | Writers should emit a profile only after the artifact actually matches that @@ -260,6 +261,8 @@ dry-run by default and must be confirmed explicitly before submitting jobs: ```bash mesh-llm models package unsloth/Qwen3-8B-GGUF:Q4_K_M --dry-run mesh-llm models package unsloth/Qwen3-8B-GGUF:Q4_K_M --confirm --follow +mesh-llm models package unsloth/inkling-GGUF:UD-Q2_K_XL --dry-run +mesh-llm models package unsloth/inkling-GGUF:UD-Q2_K_XL --experimental --confirm --follow ``` The hidden compatibility alias is `mesh-llm model-package`; prefer @@ -274,13 +277,26 @@ Important options: - `--dry-run`: print the resolved package plan and maximum cost without side effects. - `--confirm`: submit the job. - `--follow`: wait and stream job progress. +- `--experimental`: publish the package publicly with an experimental warning + and tag, and open an unmerged Hugging Face `meshllm/catalog` PR instead of + committing it to the catalog's `main` revision. Mesh discovery does not see + the package until that HF PR is reviewed and merged. - `--status `, `--logs `, `--cancel `, `--list`: inspect or manage submitted jobs. - `--update-script`: refresh the bucket script when needed. The source model should stay in colon-selector form, for example -`unsloth/Qwen3-8B-GGUF:Q4_K_M`. Do not split the quant into a separate `--quant` -argument for generated job inputs. +`unsloth/Qwen3-8B-GGUF:Q4_K_M`. A source revision may be requested as +`org/repo@revision:quant`. The job resolves that revision to an immutable commit +SHA before planning and mounts the source model volume at that SHA. Do not split +the quant into a separate `--quant` argument for generated job inputs. + +Repository GGUFs whose basenames start with `mmproj` are discovered as +multimodal projector sidecars, not model quants. The job passes them to +`skippy-model-package write-package`, publishes them under `projectors/`, and +preserves the source pipeline tag in the package model card. This is how a +combined vision/audio projector such as Inkling's `mmproj-BF16.gguf` travels +with its Q2 layer package. ## Publishing flow diff --git a/docs/SKIPPY_SPLITS.md b/docs/SKIPPY_SPLITS.md index 13bb69c17..228624be8 100644 --- a/docs/SKIPPY_SPLITS.md +++ b/docs/SKIPPY_SPLITS.md @@ -80,6 +80,83 @@ curl -sS http://127.0.0.1:9447/v1/chat/completions \ -d '{"model":"meshllm/Qwen3-8B-Q4_K_M-layers","messages":[{"role":"user","content":"Reply with OK"}],"max_tokens":16}' ``` +## Try Inkling Q2 text splits (experimental) + +Inkling is available as an immutable layer package for operators who want to +evaluate the text path before it is promoted to the customer support matrix: + +```text +meshllm/inkling-UD-Q2_K_XL-layers@9b4b91a7ddd978dd7a01679bc977f6e53777f2c7 +``` + +The package is about 296.5 GiB and contains 66 model layers plus shared and +projector artifacts. Each node materializes only its assigned layer range, but +the participating nodes still need enough aggregate GPU memory and per-host +system memory for the model, KV cache, runtime workspaces, and headroom. Start +every node with the same pinned package and context allocation: + +```bash +# first node +mesh-llm serve \ + --model meshllm/inkling-UD-Q2_K_XL-layers@9b4b91a7ddd978dd7a01679bc977f6e53777f2c7 \ + --split \ + --ctx-size 131072 \ + --bind-port 7842 + +# each additional node +mesh-llm serve \ + --model meshllm/inkling-UD-Q2_K_XL-layers@9b4b91a7ddd978dd7a01679bc977f6e53777f2c7 \ + --split \ + --ctx-size 131072 \ + --bind-port 7842 \ + --join +``` + +Use directly reachable, low-latency UDP paths. If a cloud provider remaps the +container UDP port to a different public port, confirm that the invite advertises +the reachable public endpoint and that runtime diagnostics report a direct path +before paying the model-load cost. A relay-only peer is deliberately excluded +from the Inkling split plan. + +Current Inkling policy uses an F32 activation wire and Q4_0 K/V cache. F16 and +Q8 activation wires are not interchangeable shortcuts: both failed the current +correctness policy. The published package has no default speculative strategy, +and live native MTP and multimodal serving are not yet operator claims. + +PR #1118 has exercised ordinary all-CUDA Mesh planning on a direct roughly 5 ms +Iroh/QUIC path using one 4 x 96 GB node and one 48 GB node. Automatic placement +produced ranges `0..65 / 65..66`, four lanes, and a 131,072-token allocation; +a short exact-answer request completed at 14.98 generated tokens/s. This is a +runnable research topology, not a recommendation that the highly imbalanced +range is optimal: the 65-layer head reserved about 589 GiB of CUDA host compute +workspace. Do not size a host from package bytes and VRAM alone. Prefer multiple +nearby nodes with enough system-memory headroom for a more balanced plan, and +inspect `GET /api/runtime/stages` before inference. + +The same run completed two sequential OpenAI tool loops, each with two native +structured tool calls, two intervening pressure turns, and final recall. Exact +prompt replay restored all 3,531 prompt tokens and the native log scan found no +fatal KV/decode/slot/eviction error. Both overlapping phases missed the full +harness bar: one failed while opening a direct prediction-return sink; the other +completed its tool behavior but reported zero changed-tail cached tokens. The +separate same-prefix phase reported the same cache miss. Treat concurrent +admission and suffix cache reuse as active validation gaps; the sequential tool +and exact-cache results do not waive them. + +Use streaming for a cold long-context request. Inkling Q2 prefill at this scale +can exceed the OpenAI frontend's 300-second non-streaming backend deadline; a +non-streaming request then returns HTTP 504 even though native prefill is still +healthy. Streaming establishes the response before prefill and also propagates +client cancellation to the generation worker. + +Do not treat the 131,072-token allocation as a completed 128K workload proof. +A 480,000-character repository prompt kept native prefill active for 3,429 +seconds without a fatal native-log pattern, but the two SSH-launched Mesh +processes ended together before an SSE data event was delivered. The client saw +an empty HTTP 200 stream with no content or usage. That probe is inconclusive; +an operator-facing long-context claim still requires a completed response with +reported prompt-token usage and correct far-prefix recall. + ## Lock node order and layer ranges Maintainer and benchmark runs can replace automatic placement with an exact, diff --git a/docs/design/TESTING.md b/docs/design/TESTING.md index 7affe6812..aa0867f5a 100644 --- a/docs/design/TESTING.md +++ b/docs/design/TESTING.md @@ -750,8 +750,8 @@ cached and a worker does not: to open `skippy-stage/2`, then Skippy artifact-transfer stream 0x03, to fetch only its assigned package files before the normal HF fallback path. - Current/released mixed mesh: a released coordinator without advertised - `skippy-stage/2` `artifact-transfer`, `stage-generation-3`, and - `direct-prediction-return` support must not be selected for a generation-3 + `skippy-stage/2` `artifact-transfer`, `stage-generation-4`, and + `direct-prediction-return` support must not be selected for a generation-4 split topology; the worker must fall back to local/HF package resolution. - Default public-mesh safety: with `MESH_LLM_ARTIFACT_TRANSFER` unset, the node must advertise no `artifact-transfer` feature, reject inbound artifact diff --git a/docs/design/message_protocol.md b/docs/design/message_protocol.md index 6867c337b..34c5a03b1 100644 --- a/docs/design/message_protocol.md +++ b/docs/design/message_protocol.md @@ -448,7 +448,7 @@ message MeshSubprotocolOpen { } ``` -Outbound transfer uses mesh stream `0x0d` (`STREAM_SUBPROTOCOL`), followed by a length-prefixed `MeshSubprotocolOpen { name: "skippy-stage", major: 2 }`, the Skippy-owned stream kind `0x03`, a length-prefixed `StageArtifactTransferRequest`, a length-prefixed `StageArtifactTransferResponse`, and raw artifact bytes when accepted. Skippy stage major 2 is a compatibility break for generation-3 direct prediction return. +Outbound transfer uses mesh stream `0x0d` (`STREAM_SUBPROTOCOL`), followed by a length-prefixed `MeshSubprotocolOpen { name: "skippy-stage", major: 2 }`, the Skippy-owned stream kind `0x03`, a length-prefixed `StageArtifactTransferRequest`, a length-prefixed `StageArtifactTransferResponse`, and raw artifact bytes when accepted. Skippy stage major 2 is a compatibility break for generation-4 direct prediction return and verify retirement. **Request:** ```proto diff --git a/docs/skippy/DATA_FLOW.md b/docs/skippy/DATA_FLOW.md index 58810c655..9f78d34a1 100644 --- a/docs/skippy/DATA_FLOW.md +++ b/docs/skippy/DATA_FLOW.md @@ -40,11 +40,11 @@ activation links and then crossed three reply links before stage 0 could emit the token. On a topology with a fixed 10 ms delay per inter-stage hop, the reply chain alone makes the hot path six hops, or about 60 ms before compute. -## Generation 3 Direct Prediction Return +## Generation 4 Direct Prediction Return and Verify Retirement -Stage protocol generation 3 is a compatibility-breaking change. A peer is stage -compatible only when it advertises the `skippy-stage` major version for -generation 3 plus `stage-generation-3`. Prediction-bearing messages return +Stage protocol generation 4 is a compatibility-breaking change. A peer is stage +compatible only when it advertises both `skippy-stage/2` and +`stage-generation-4`. Prediction-bearing messages return directly from the final/readout stage to the driver-facing stage. Intermediate stages continue to forward activations and may handle cold-path control acknowledgments, but they are not part of the decode-token prediction return path. diff --git a/docs/skippy/FAMILY_CERTIFY.md b/docs/skippy/FAMILY_CERTIFY.md index c645f6435..454bba491 100644 --- a/docs/skippy/FAMILY_CERTIFY.md +++ b/docs/skippy/FAMILY_CERTIFY.md @@ -18,12 +18,22 @@ A family is certified when the current recommended artifact has evidence for: | --- | --- | | `single-step` | A two-stage split produces the same next token as full-model execution. | | `chain` | The recommended multi-stage split produces the same next token as full-model execution, unless the family has a documented two-stage-only split. | -| `dtype-matrix` | `f32` and `f16` activation transfer are exact; `q8` is marked validated or rejected for that family/split. | +| `dtype-matrix` | `f32` is exact; `f16` and `q8` are each marked validated or rejected for that family/split, and the selected family default is exact. | | `state-handoff` | Exact live state mobility is accepted or explicitly rejected by the Qwen3 baseline rule. | +| `context-capacity` | The staged serving path allocates the required context and completes a near-limit prefill plus continuation. For models whose native context is at least 131,072 tokens, this lane must use at least 131,072; smaller-context models must use their native limit. | | `llama-spec-bench` | Optional target/draft speculative compatibility checks. | -The default shipping wire dtype is `f16`. `q8` is opt-in only when the -dtype-matrix lane proves exactness for that family and split. +The usual shipping wire dtype is `f16`. A family whose F16 lane is rejected +must select `f32` explicitly in its capability record. `q8` is opt-in only when +the dtype-matrix lane proves exactness for that family and split. + +The small `--ctx-size 256` correctness run below proves graph and split parity; +it does not prove usable context capacity. Context support is a separate live +serving result. For a 131,072-token lane, require the server to report +`ctx_size >= 131072`, successfully prefill a request with at least 120,000 +prompt tokens, and generate a continuation without a stage failure, KV-slot +failure, or context truncation. Record the KV types, lane count, stage ranges, +and per-stage cache allocation with that result. ## Certification Command diff --git a/docs/skippy/FAMILY_STATUS.md b/docs/skippy/FAMILY_STATUS.md index 0c2dc326f..90b95daec 100644 --- a/docs/skippy/FAMILY_STATUS.md +++ b/docs/skippy/FAMILY_STATUS.md @@ -8,7 +8,18 @@ Certification process lives in `docs/FAMILY_CERTIFY.md`. Payload measurements and topology constraints are summarized here so this file stays the only customer-facing source of truth. -Last updated: 2026-05-07. +Last updated: 2026-07-31. + +## Context Capacity Contract + +The parity-sized family certification lanes do not establish a production +context window. A reviewed context claim requires a separate staged-serving +capacity result at `min(native_context, 131072)` tokens. For models with native +context of at least 131,072, the minimum evidence is a 131,072-token allocation, +a request reporting at least 120,000 prompt tokens, and a successful +continuation. The support record must name the tested artifact, KV cache types, +lane count, and stage plan. Do not infer native-context support from a small +correctness smoke or from GGUF metadata alone. ## Customer Support Matrix @@ -44,6 +55,7 @@ Last updated: 2026-05-07. | LFM2 | Supported | `meshllm/lfm2-350m-parity-q4_k_m-gguf:q4_k_m` | `layer_end=16`, `splits=5,10`, activation width `1024` | `f16`; q8 validated | `baseline,ngram,ngram-adaptive` | Keep recurrent range `0..16` sticky for normal decode. | Use `KvRecurrent` for exact prefix cache restore; native sequence remap cache smoke passed. | | Jamba | Supported | `bartowski/ai21labs_AI21-Jamba2-3B-GGUF:Q4_K_M` | `layer_end=28`, `splits=9,18`, activation width `2560` | `f16`; q8 validated | `baseline,ngram,ngram-adaptive` | Keep recurrent range `0..28` sticky for normal decode. | Use `KvRecurrent` for exact prefix cache restore; middle-stage recurrent-only slices are valid. | | Kimi Linear | Supported | `bartowski/moonshotai_Kimi-Linear-48B-A3B-Instruct-GGUF:IQ2_XXS` | `layer_end=27`, `splits=9,18`, activation width `2304` | `f16`; q8 validated | `baseline,ngram,ngram-adaptive` | Keep recurrent KDA ranges `0..3`, `4..7`, `8..11`, `12..15`, `16..19`, `20..23`, and `24..26` sticky for normal decode. | Use `KvRecurrent`; sparse K-only MLA KV pages plus recurrent state are required for exact prefix restore. | +| Laguna S 2.1 | Supported for the pinned Q4_K_M package | `meshllm/laguna-s-2.1-Q4_K_M-layers@0c467ad441ee94cb5a76f626294d963c4048507d` | `layer_end=48`, activation width `3072`; observed CUDA plan `0..25 / 25..39 / 39..48` | `f16`; q8 untested | package-default suffix N-gram depth 2 | Keep the hybrid recurrent/attention state on its owning stage; use direct inter-stage paths. | Use one lane with Q4_0 K/V for the reviewed 131,072-token allocation. A 128,080-token cold prompt completed twice; 44,416-token exact-prefix restore passed, while the 128K repeats ran with zero cached tokens. Structured OpenAI tool call and tool-result turns passed. | | Mamba | Supported | `mradermacher/mamba-130m-hf-GGUF:Q4_K_M` | `layer_end=24`, `splits=8,16`, activation width `768` | `f16`; q8 validated | `baseline,ngram,ngram-adaptive` | Keep recurrent range `0..24` sticky for normal decode. | Use `KvRecurrent`; cache restore can have zero native KV bytes. | | Mamba2 | Supported | `mradermacher/mamba-2.8b-hf-GGUF:Q4_K_M` | `layer_end=64`, `splits=21,42`, activation width `2560` | `f16`; q8 validated | `baseline,ngram,ngram-adaptive` | Keep recurrent range `0..64` sticky for normal decode. | Use `KvRecurrent`; full-state mobility is rejected as too large. | | RWKV6 | Supported | `latestissue/rwkv-6-finch-1b6-gguf:Q4_K` | `layer_end=24`, `splits=8,16`, activation width `2048` | `f16`; q8 rejected | `baseline,ngram,ngram-adaptive` | Keep recurrent range `0..24` sticky for normal decode. | Use `KvRecurrent`; cache restore can have zero native KV bytes. | @@ -88,13 +100,22 @@ Last updated: 2026-05-07. ## Text-Split Candidates -These families now pass the cheap runtime-slice text lane, but are not promoted -to the customer support matrix until the remaining cache smoke, reviewed -topology records, and family-specific policy notes are updated. +These families now pass a cheap runtime-slice or package-backed text lane, but +are not promoted to the customer support matrix until the remaining cache +smoke, reviewed topology records, and family-specific policy notes are updated. -```text -Gemma text -``` +- **Gemma text:** the sampled `gemma` artifact currently requires an F32 + activation wire; see the exception below. +- **Inkling text:** use pinned package + `meshllm/inkling-UD-Q2_K_XL-layers@9b4b91a7ddd978dd7a01679bc977f6e53777f2c7`. + It has 66 layers, activation width 6144, F32 wire, and Q4_0 K/V policy. PR + #1118 has proven ordinary all-CUDA Mesh planning and real split generation at + a 131,072-token allocation over a direct approximately 5 ms path, two + sequential native OpenAI tool loops, and exact-prefix cache restoration. + Promotion still requires a completed long-context continuation plus reliable + overlapping-request admission and changed-tail same-prefix cache reuse. + F16/Q8 wire, native MTP, and multimodal inference remain outside the current + claim. ## Exceptions diff --git a/docs/skippy/LLAMA_PARITY.md b/docs/skippy/LLAMA_PARITY.md index 4fdf1d0b7..4c300ba6e 100644 --- a/docs/skippy/LLAMA_PARITY.md +++ b/docs/skippy/LLAMA_PARITY.md @@ -360,9 +360,10 @@ implementation. ## Current Local Evidence -These rows were collected on the local Mac Studio against the Metal stage ABI. -They are cheap text-split and cache-smoke evidence, not full promotion by -themselves until the reviewed topology records are updated. +These rows were collected primarily on the local Mac Studio against the Metal +stage ABI. They are cheap text-split and cache-smoke evidence, not full +promotion by themselves until the reviewed topology records are updated. +Rows with distributed evidence call out the second backend explicitly. | Family | Artifact | Text Split | q8 Wire | Exact State | Cache | | --- | --- | --- | --- | --- | --- | @@ -370,6 +371,8 @@ themselves until the reviewed topology records are updated. | `deepseek` | `Morgen0052/deepseek-llm-7b-chat-Q4_K_M-GGUF` | `single-step`, `chain`, and f16 dtype matrix passed | rejected | accepted | `ResidentKv` borrowed-hit smoke passed, 64-token prefix, 1.58x cache-hit speedup | | `openai_moe` | `ggml-org/gpt-oss-20b-GGUF:gpt-oss-20b-mxfp4` | `single-step`, `chain`, and dtype matrix passed | rejected | accepted | `ResidentKv` state handoff passed; llama.cpp model file is `openai-moe`, GGUF architecture is `gpt-oss` | | `ernie4_5_moe` | `lmstudio-community/ERNIE-4.5-21B-A3B-PT-GGUF:Q4_K_M` | `single-step`, `chain`, and dtype matrix passed | validated | accepted | `ResidentKv` state handoff passed | +| `laguna` | `meshllm/laguna-s-2.1-Q4_K_M-layers@0c467ad441ee94cb5a76f626294d963c4048507d` | package-backed `single-step` and three-stage `chain` parity passed; ordinary three-stage CUDA Mesh serving completed two 128,080-token prompts | untested | accepted through bounded recurrent checkpoint/replay | Three-stage CUDA placement `0..25 / 25..39 / 39..48` used a 131,072-token allocation, Q4_0 K/V, F16 activation wire, and package-default suffix N-gram depth 2. The two 128K cold prompts ran at 956.12/958.05 prompt tok/s and 12.45/12.59 generation tok/s. A 44,460-token exact replay restored 44,416 tokens; structured OpenAI tool call and result turns passed; 12/12 subsequent stability requests passed. This evidence promotes only the pinned Q4_K_M package. The 128K repeats reported zero cached tokens; 256K, Q8 wire, other quants, and Metal at 128K remain unproven. | +| `qwen3next` (Inkling) | `meshllm/inkling-UD-Q2_K_XL-layers@9b4b91a7ddd978dd7a01679bc977f6e53777f2c7` | experimental two-stage CUDA text serving passed at a 131,072-token allocation | rejected; F32 required | exact replay accepted; recurrent verification recovery exercised | Automatic placement `0..65 / 65..66` on a direct approximately 5 ms path generated at 14.98 tok/s and completed two sequential native OpenAI tool loops. Exact replay restored 3,531/3,531 prompt tokens and the native fatal-pattern scan passed. This lopsided research topology reserved about 589 GiB of head host workspace and is not a deployment recommendation. Overlapping admission and changed-tail same-prefix cache reuse still fail the formal harness; native MTP, multimodal serving, and clean Metal teardown remain unproven. | | `llama4` | `ggml-org/Llama-4-Scout-17B-16E-Instruct-GGUF:Q4_K_M` | package validated | untested | untested | package-only validation passed: 48 layers, 627 owned tensors, 51 artifacts, no missing/duplicate tensors | | `mistral4` | `bartowski/mistralai_Mistral-Small-4-119B-2603-GGUF:IQ2_XXS` | package validated | untested | untested | package-only validation passed: 36 layers, 579 tensors, 39 artifacts, no missing/duplicate tensors | | `nemotron_h_moe` | `lmstudio-community/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-GGUF:Q4_K_M` | package validated | untested | rejected-too-large | package-only validation passed: 52 layers, 401 tensors, 55 artifacts; `KvRecurrent` target | diff --git a/docs/skippy/NEW_MODEL_ONBOARDING.md b/docs/skippy/NEW_MODEL_ONBOARDING.md index f50ccb42c..c9106ad70 100644 --- a/docs/skippy/NEW_MODEL_ONBOARDING.md +++ b/docs/skippy/NEW_MODEL_ONBOARDING.md @@ -16,7 +16,7 @@ Issue #630, Cohere Command A+, is the first model tracked with this flow. | Artifact inspected | GGUF metadata or package manifest gives architecture, layer count, activation width, quant, shard layout, and tokenizer sidecars. | May plan a package job. | | Package validated | `skippy-model-package` writes, validates, and preflights a package with no unresolved manifest, artifact, sidecar, materialization, missing, duplicate, or checksum diagnostics. | May test staged serving. | | Runtime smoke passed | A package-backed model starts and answers through the OpenAI-compatible surface. | May collect serving evidence. | -| Family certified | Split correctness, dtype matrix, state handoff/cache policy, and required multimodal sidebands pass. | May promote to reviewed support. | +| Family certified | Split correctness, dtype matrix, state handoff/cache policy, required context capacity, and required multimodal sidebands pass. | May promote to reviewed support. | | Reviewed support | `docs/skippy/FAMILY_STATUS.md` and reviewed topology records are updated from evidence. | User-visible support claim. | Do not skip from candidate to reviewed support. A big or popular model is still @@ -36,7 +36,11 @@ only a candidate until the evidence exists. load and inspect it. 5. Run package validation before runtime certification for models too large for a local full-model parity pass. -6. Promote only after the evidence maps cleanly to both +6. Run the context-capacity lane at `min(native_context, 131072)`. For models + whose native context is at least 131,072 tokens, a tiny-context runtime smoke + is not promotion evidence: the live split must allocate 131,072 and complete + a request with at least 120,000 prompt tokens plus a continuation. +7. Promote only after the evidence maps cleanly to both `docs/skippy/FAMILY_STATUS.md` and `crates/skippy-topology/capabilities/reviewed-family-capabilities.json`. diff --git a/docs/skippy/PIPELINED_VERIFY_WINDOW.md b/docs/skippy/PIPELINED_VERIFY_WINDOW.md index 3f7b090b8..f40d78ba5 100644 --- a/docs/skippy/PIPELINED_VERIFY_WINDOW.md +++ b/docs/skippy/PIPELINED_VERIFY_WINDOW.md @@ -41,9 +41,9 @@ The request-local N-gram index contains committed target history only. Optimistic suffixes may be queried to extend the current branch, but they are never inserted into the index before target acceptance. -## Stage Protocol v10 +## Stage Protocol v11 -`STAGE_STATE_VERSION` is `10`, and `VerifyWindow` is message kind `21`. +`STAGE_STATE_VERSION` is `11`, and `VerifyWindow` is message kind `21`. Every verification request carries: - a FIFO window ID; diff --git a/docs/skippy/SUFFIX_NGRAM_PROPOSER.md b/docs/skippy/SUFFIX_NGRAM_PROPOSER.md index ba1bbd7fc..396d84d96 100644 --- a/docs/skippy/SUFFIX_NGRAM_PROPOSER.md +++ b/docs/skippy/SUFFIX_NGRAM_PROPOSER.md @@ -165,7 +165,14 @@ strategy types. A suffix package proposer must declare request-local history: "strategies": { "ngram-suffix": { "type": "ngram-suffix", - "proposer": "suffix" + "proposer": "suffix", + "window_policy": { + "default": "fixed", + "initial_window": 32, + "min_window": 1, + "max_window": 32, + "pipeline_depth": 2 + } } } } @@ -173,6 +180,11 @@ strategy types. A suffix package proposer must declare request-local history: } ``` +`window_policy.pipeline_depth` is optional and defaults to `1` for older +packages. A package may select a deeper verification pipeline only after the +specific artifact and serving topology have passed workload and stability +gates at that depth. + The cache and suffix limits intentionally differ. Cache uses llama.cpp's stateful lookup with a match window no larger than four tokens. Suffix may use a much longer exact match. For both proposers, proposal output length is diff --git a/docs/skippy/WAN_SPLIT_PERF.md b/docs/skippy/WAN_SPLIT_PERF.md index 9a5631dec..f40f06590 100644 --- a/docs/skippy/WAN_SPLIT_PERF.md +++ b/docs/skippy/WAN_SPLIT_PERF.md @@ -155,7 +155,7 @@ The smoking gun is `window_shrinks 0`: the adaptive policy never narrowed the window under a sustained reject storm, so it kept proposing deep, kept rejecting, and kept paying the 3× round-trip recovery. -Stage-state v10 deletes that recovery path. Decode and verify messages carry an +Stage-state v11 removes that recovery path. Decode and verify messages carry an authoritative absolute position, while non-overlapping continuation chunks make the fully accepted path advance monotonically without rewinding. Only a real divergence trims an invalid target suffix locally. No checkpoint, restore ACK, diff --git a/docs/skippy/llama-parity-candidates.json b/docs/skippy/llama-parity-candidates.json index 22796aefd..6b1613ca5 100644 --- a/docs/skippy/llama-parity-candidates.json +++ b/docs/skippy/llama-parity-candidates.json @@ -185,6 +185,14 @@ "include": "Qwen3-MOE-4x0.6B-2.4B-Writing-Thunder.Q4_K_M.gguf", "notes": "text lane passed in llama-parity-qwen3moe-runtime-slice-2; q8 activation wire validated; ResidentKv native-sequence remap cache smoke passed; MoE expert-stage smoke already passed" }, + { + "llama_model": "laguna", + "family": "laguna", + "status": "certified_package_only", + "repo": "poolside/Laguna-S-2.1-GGUF", + "include": "laguna-s-2.1-Q4_K_M.gguf", + "notes": "package-only certification status retained while broader gates remain pending; source revision edd093522473dc7313b0738d8b4116b7f8b9745f (SHA-256 a34c74e46688122bef83122f4133031bababbefcf57436dde97048c91e2cc6ff) and package revision 0c467ad441ee94cb5a76f626294d963c4048507d (manifest SHA-256 0250cfb54ceeb94a9c71e48df447f780e32fc625553844d6403770f315be0237) passed M5 Max package-backed single-step parity at split 24 and three-stage chain parity at splits 16,32 with f16 activation transport: baseline token 674 matched both staged predictions, activation width 3072, payload 12288 bytes, wire payload 6144 bytes; an ordinary M5 Metal plus Australian Vast RTX 6000 Ada CUDA private Mesh run passed normal 0..36/36..48 placement, direct Iroh transport, configured context 262144, a 44460-token prompt, exact 65-token completion agreement, and suffix N-gram verification at depth 2 with 32/64 speculative tokens accepted per request; the final three-request recovery probe decoded at 19.74, 18.24, and 15.42 tok/s and an exact-prefix hit reused 44416 prompt tokens; 63/63 stability requests returned HTTP 200; suffix N-gram min 5, max 32, proposal cap 48, and verify pipeline depth 2 are now the package default and no-config resolver coverage consumes that policy, while a fresh live M5 plus Vast no-override confirmation remains pending; 24/24 full-context placement exhausted the 48 GB CUDA worker, the structured tool harness passed 6/13, and Q8 wire, state handoff, structured tool use, and full native-context saturation remain unproven" + }, { "llama_model": "mistral3", "family": "mistral", @@ -523,6 +531,18 @@ ], "notes": "HunyuanOCR split multimodal serving passed with ggml-org/HunyuanOCR-GGUF:Q8_0 plus mmproj-HunyuanOCR-Q8_0.gguf using the shared Hunyuan-Dense/VL graph filter and the media activation position sideband; default f16 activation wire is the support target" }, + { + "llama_model": "inkling", + "family": "inkling", + "status": "package_or_remote_only", + "repo": "unsloth/inkling-GGUF", + "include": [ + "UD-Q2_K_XL/*.gguf", + "mmproj-BF16.gguf" + ], + "recurrent": "all", + "notes": "975B multimodal hybrid attention with per-layer short convolution; Q2 layer-package target; certify text, image, audio, two-stage and multi-stage serving with KvRecurrent before promotion" + }, { "llama_model": "internlm2", "family": "internlm2", diff --git a/docs/specs/layer-package-repos.md b/docs/specs/layer-package-repos.md index 4f4b0beaf..03e6a92cd 100644 --- a/docs/specs/layer-package-repos.md +++ b/docs/specs/layer-package-repos.md @@ -214,7 +214,8 @@ Minimal shape: "default": "fixed", "initial_window": 1, "min_window": 1, - "max_window": 1 + "max_window": 1, + "pipeline_depth": 1 } } } @@ -688,7 +689,8 @@ The current native MTP strategy shape is: "default": "fixed", "initial_window": 1, "min_window": 1, - "max_window": 1 + "max_window": 1, + "pipeline_depth": 1 } } ``` @@ -746,6 +748,7 @@ The package schema separates a proposer match length from its output budget: | `ngram_min` / `ngram_max` | N-gram proposers | Define the historical token match range. Both are required and `ngram_min <= ngram_max`. | | `max_proposal_tokens` | N-gram proposers | Caps how many continuation tokens the proposer may return. It is independent of `ngram_max`. | | `history_scope` | `ngram-cache`, `ngram-suffix` | Must be `"request"`; a history proposer never observes another request's tokens. | +| `window_policy.pipeline_depth` | All strategies | Optional positive per-request capacity for in-flight verification windows. Omission preserves the legacy depth of `1`; package defaults above `1` require topology/workload-specific evidence. | | `initial_tokens` / `max_tokens` | composite extension policy | Bound the adaptive N-gram tail after an MTP prefix. | | `tail_backoff_proposals` | composite extension policy | Sets how many proposals to back off after an unhelpful tail. | diff --git a/scripts/hf-skippy-convert-job.py b/scripts/hf-skippy-convert-job.py new file mode 100644 index 000000000..784787e3b --- /dev/null +++ b/scripts/hf-skippy-convert-job.py @@ -0,0 +1,241 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = ["huggingface_hub[hf_xet]>=0.34"] +# /// + +"""Run a native skippy-quantize HF checkpoint conversion on Hugging Face Jobs.""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import subprocess +from pathlib import Path + + +def run(*command: str, cwd: Path | None = None) -> None: + print("+", " ".join(command), flush=True) + subprocess.run(command, cwd=cwd, check=True) + + +def ensure_build_tools() -> None: + required = ("git", "curl", "cmake", "c++", "ld.lld") + if any(shutil.which(tool) is None for tool in required): + if shutil.which("apt-get") is None: + raise RuntimeError(f"missing build tools: {required}") + run("apt-get", "update") + run( + "apt-get", + "install", + "-y", + "build-essential", + "cmake", + "curl", + "git", + "lld", + "pkg-config", + ) + if shutil.which("cargo") is None: + run( + "sh", + "-c", + "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y", + ) + os.environ["PATH"] = f"{Path.home() / '.cargo' / 'bin'}:{os.environ['PATH']}" + if shutil.which("just") is None: + run("cargo", "install", "just", "--locked") + + +def checkout_mesh(repo: str, revision: str, root: Path) -> None: + if root.exists(): + shutil.rmtree(root) + run("git", "clone", "--filter=blob:none", repo, str(root)) + run("git", "checkout", revision, cwd=root) + + +def write_beta_card(artifact_dir: Path, source_repo: str, revision: str) -> None: + card = f"""--- +license: apache-2.0 +base_model: {source_repo} +tags: +- gguf +- beta +- skippy +- mtp +--- + +# Inkling MTP sidecar (beta) + +This is a public beta artifact for Skippy compatibility testing. It contains +Inkling's multi-token-prediction depths plus the shared embedding/output +context needed by distributed final stages. It is not a standalone chat model +and is not a promoted mesh-llm catalog entry. + +Built with native `skippy-quantize` from mesh-llm revision `{revision}`. +""" + (artifact_dir / "README.md").write_text(card, encoding="utf-8") + + +def convert(args: argparse.Namespace, root: Path) -> Path: + binary = root / "target" / "release" / "skippy-quantize" + run("just", "skippy-quantize-standalone-release-build", "cpu", cwd=root) + work = Path(args.work_dir) + target = work / "target" + artifact_dir = target / args.target_prefix + manifest = work / "convert-manifest.json" + records = work / "records" + spool = work / "spool" + status = work / "status.json" + work.mkdir(parents=True, exist_ok=True) + run( + str(binary), + "convert-job", + "--source", + args.source, + "--target", + str(target), + "--target-prefix", + args.target_prefix, + "--output-basename", + args.output_basename, + "--output-type", + "bf16", + "--expected-splits", + str(args.expected_splits), + "--window-size", + "1", + "--manifest", + str(manifest), + "--mtp", + "--split-max-size", + args.split_max_size, + "--max-memory", + args.max_memory, + "--stream-buffer-bytes", + "8388608", + "--spool-dir", + str(spool), + "--record-dir", + str(records), + "--json-event-file", + str(status), + "--json-event-interval-seconds", + "60", + "--json-event-window", + "8", + "--watchdog-seconds", + "300", + ) + run(str(binary), "verify-job", "--manifest", str(manifest), "--json") + if not artifact_dir.is_dir(): + raise FileNotFoundError(f"convert-job did not produce {artifact_dir}") + shutil.copy2(manifest, artifact_dir / "skippy-convert-manifest.json") + if status.exists(): + shutil.copy2(status, artifact_dir / "skippy-convert-status.json") + write_beta_card(artifact_dir, args.source_repo, args.mesh_revision) + validate_converted_artifact(artifact_dir) + return artifact_dir + + +def upload(args: argparse.Namespace, artifact_dir: Path) -> None: + from huggingface_hub import HfApi + + api = HfApi(token=os.environ["HF_TOKEN"]) + api.create_repo(args.target_repo, repo_type="model", private=False, exist_ok=True) + api.upload_folder( + repo_id=args.target_repo, + repo_type="model", + folder_path=str(artifact_dir), + ) + + +def validate_converted_artifact(artifact_dir: Path) -> None: + required_files = ("README.md", "skippy-convert-manifest.json") + missing = [name for name in required_files if not (artifact_dir / name).is_file()] + if not artifact_dir.is_dir() or missing: + details = f"; missing {', '.join(missing)}" if missing else "" + raise FileNotFoundError( + f"complete converted artifact not found: {artifact_dir}{details}" + ) + + manifest_path = artifact_dir / "skippy-convert-manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + expected_splits = manifest.get("expected_splits") + basename = manifest.get("output_basename") + if not isinstance(expected_splits, int) or expected_splits < 1: + raise ValueError(f"invalid expected_splits in {manifest_path}") + if not isinstance(basename, str) or not basename: + raise ValueError(f"invalid output_basename in {manifest_path}") + + if expected_splits == 1: + expected_names = [f"{basename}.gguf"] + else: + expected_names = [ + f"{basename}-{index:05}-of-{expected_splits:05}.gguf" + for index in range(1, expected_splits + 1) + ] + missing_shards = [name for name in expected_names if not (artifact_dir / name).is_file()] + if missing_shards: + raise FileNotFoundError( + f"converted artifact is incomplete: missing {', '.join(missing_shards)}" + ) + + +def converted_artifact_dir(args: argparse.Namespace) -> Path: + artifact_dir = Path(args.work_dir) / "target" / args.target_prefix + if args.upload_only: + validate_converted_artifact(artifact_dir) + return artifact_dir + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--source", default="/mnt/checkpoint") + parser.add_argument("--source-repo", required=True) + parser.add_argument("--target-repo", required=True) + parser.add_argument("--target-prefix", default="BF16") + parser.add_argument("--output-basename", required=True) + parser.add_argument("--expected-splits", type=int, default=1) + parser.add_argument("--split-max-size", default="50G") + parser.add_argument("--max-memory", default="24G") + parser.add_argument("--work-dir", default="/data/skippy-convert") + parser.add_argument("--mesh-repo", default="https://github.com/Mesh-LLM/mesh-llm.git") + parser.add_argument("--mesh-revision", required=True) + parser.add_argument( + "--upload-only", + action="store_true", + help="publish an existing converted artifact without rebuilding or reconverting it", + ) + parser.add_argument( + "--xet-high-performance", + action="store_true", + help="enable Hugging Face Xet high-performance upload mode (use on hosts with at least 64 GB RAM)", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + os.environ.setdefault("HF_HOME", str(Path(args.work_dir) / "hf-home")) + # The work directory can be a mounted bucket. Xet's shard cache performs + # poorly on network filesystems, so keep it on the Job's local SSD. + os.environ.setdefault("HF_XET_CACHE", "/tmp/hf-xet") + if args.xet_high_performance: + os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1") + if args.upload_only: + artifact_dir = converted_artifact_dir(args) + upload(args, artifact_dir) + print(f"published https://huggingface.co/{args.target_repo}", flush=True) + return + ensure_build_tools() + mesh_root = Path("/tmp/mesh-llm") + checkout_mesh(args.mesh_repo, args.mesh_revision, mesh_root) + artifact_dir = convert(args, mesh_root) + upload(args, artifact_dir) + print(f"published https://huggingface.co/{args.target_repo}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/hf-skippy-mtp-certify-job.py b/scripts/hf-skippy-mtp-certify-job.py new file mode 100644 index 000000000..d9f0b3d3a --- /dev/null +++ b/scripts/hf-skippy-mtp-certify-job.py @@ -0,0 +1,229 @@ +# /// script +# requires-python = ">=3.11" +# /// + +"""Validate a mounted target GGUF, projector, and external MTP sidecar on HF Jobs.""" + +from __future__ import annotations + +import argparse +import ipaddress +import os +import shutil +import socket +import subprocess +import urllib.parse +import urllib.request +from pathlib import Path + + +PROJECTOR_DOWNLOAD_TIMEOUT_SECONDS = 60 +PROJECTOR_DOWNLOAD_MAX_BYTES = 64 * 1024 * 1024 * 1024 +PROJECTOR_DOWNLOAD_CHUNK_BYTES = 8 * 1024 * 1024 +TRUSTED_PROJECTOR_HOST_SUFFIXES = ("huggingface.co", "hf.co", "xethub.hf.co") + + +def run(*command: str, cwd: Path | None = None) -> None: + print("+", " ".join(command), flush=True) + subprocess.run(command, cwd=cwd, check=True) + + +def ensure_build_tools() -> None: + required = ("git", "curl", "cmake", "c++", "ld.lld") + if any(shutil.which(tool) is None for tool in required): + if shutil.which("apt-get") is None: + raise RuntimeError(f"missing build tools: {required}") + run("apt-get", "update") + run( + "apt-get", + "install", + "-y", + "build-essential", + "cmake", + "curl", + "git", + "lld", + "pkg-config", + ) + if shutil.which("cargo") is None: + run( + "sh", + "-c", + "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y", + ) + os.environ["PATH"] = f"{Path.home() / '.cargo' / 'bin'}:{os.environ['PATH']}" + if shutil.which("just") is None: + run("cargo", "install", "just", "--locked") + + +def checkout_mesh(repo: str, revision: str, root: Path) -> None: + if root.exists(): + shutil.rmtree(root) + run("git", "clone", "--filter=blob:none", repo, str(root)) + run("git", "checkout", revision, cwd=root) + + +def model_parts(args: argparse.Namespace) -> list[Path]: + parts = sorted(Path(args.model_root).glob(args.model_pattern)) + if len(parts) != args.expected_parts: + raise RuntimeError( + f"expected {args.expected_parts} target parts matching {args.model_pattern!r}, " + f"found {len(parts)}" + ) + return parts + + +def require_gguf_magic(path: Path) -> Path: + with path.open("rb") as handle: + magic = handle.read(4) + if magic != b"GGUF": + raise RuntimeError(f"invalid GGUF magic for {path}: {magic!r}") + return path + + +def validate_projector_url(url: str) -> urllib.parse.ParseResult: + parsed = urllib.parse.urlparse(url) + if parsed.scheme != "https": + raise RuntimeError(f"unsupported projector URL scheme: {parsed.scheme!r}") + hostname = parsed.hostname + if hostname is None or parsed.username is not None or parsed.password is not None: + raise RuntimeError("projector URL must contain a trusted HTTPS host without credentials") + hostname = hostname.rstrip(".").lower() + if not any( + hostname == suffix or hostname.endswith(f".{suffix}") + for suffix in TRUSTED_PROJECTOR_HOST_SUFFIXES + ): + raise RuntimeError(f"untrusted projector URL host: {hostname!r}") + if parsed.port not in (None, 443): + raise RuntimeError(f"unsupported projector URL port: {parsed.port}") + addresses = socket.getaddrinfo(hostname, 443, type=socket.SOCK_STREAM) + if not addresses: + raise RuntimeError(f"projector URL host did not resolve: {hostname!r}") + for address in addresses: + resolved = ipaddress.ip_address(address[4][0]) + if not resolved.is_global: + raise RuntimeError( + f"projector URL host resolved to a non-public address: {hostname!r}" + ) + return parsed + + +class TrustedProjectorRedirectHandler(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): # noqa: ANN001 + validate_projector_url(newurl) + return super().redirect_request(req, fp, code, msg, headers, newurl) + + +def copy_projector_response(response, output) -> None: # noqa: ANN001 + content_length = response.headers.get("Content-Length") + if content_length is not None and int(content_length) > PROJECTOR_DOWNLOAD_MAX_BYTES: + raise RuntimeError("projector download exceeds the maximum supported size") + copied = 0 + while chunk := response.read(PROJECTOR_DOWNLOAD_CHUNK_BYTES): + copied += len(chunk) + if copied > PROJECTOR_DOWNLOAD_MAX_BYTES: + raise RuntimeError("projector download exceeds the maximum supported size") + output.write(chunk) + + +def projector_path(args: argparse.Namespace) -> Path: + if not args.projector_url: + return require_gguf_magic(Path(args.projector)) + parsed = validate_projector_url(args.projector_url) + target = Path(args.projector_local_path) + target.parent.mkdir(parents=True, exist_ok=True) + temporary = target.with_suffix(f"{target.suffix}.part") + redacted_url = urllib.parse.urlunparse(parsed._replace(query="", fragment="")) + print(f"+ download {redacted_url} -> {target}", flush=True) + opener = urllib.request.build_opener(TrustedProjectorRedirectHandler()) + try: + with opener.open( + args.projector_url, + timeout=PROJECTOR_DOWNLOAD_TIMEOUT_SECONDS, + ) as response, temporary.open("wb") as output: + copy_projector_response(response, output) + require_gguf_magic(temporary) + temporary.replace(target) + except Exception: + temporary.unlink(missing_ok=True) + raise + return require_gguf_magic(target) + + +def run_report(command: list[str], report_out: str) -> None: + print("+", " ".join(command), flush=True) + completed = subprocess.run(command, text=True, capture_output=True) + if completed.stderr: + print(completed.stderr, end="", flush=True) + if completed.stdout: + print(completed.stdout, end="", flush=True) + if completed.returncode != 0: + raise subprocess.CalledProcessError(completed.returncode, command) + report = Path(report_out) + report.parent.mkdir(parents=True, exist_ok=True) + report.write_text(completed.stdout, encoding="utf-8") + print(f"certification report: {report}", flush=True) + + +def certify(args: argparse.Namespace, mesh_root: Path) -> None: + binary = mesh_root / "target" / "release" / "skippy-quantize" + run("just", "skippy-quantize-standalone-release-build", "cpu", cwd=mesh_root) + projector = projector_path(args) + if args.projector_only: + run_report( + [str(binary), "validate-projector", "--projector", str(projector), "--json"], + args.report_out, + ) + return + target_parts = [require_gguf_magic(path) for path in model_parts(args)] + mtp_draft = require_gguf_magic(Path(args.mtp_draft)) + command = [str(binary), "validate-mtp-attach"] + for part in target_parts: + command.extend(("--model", str(part))) + command.extend( + ( + "--mtp-draft", + str(mtp_draft), + "--layer-count", + str(args.layer_count), + "--ctx-size", + str(args.ctx_size), + "--projector", + str(projector), + "--json", + ) + ) + if args.mtp_layer_count is not None: + command.extend(("--mtp-layer-count", str(args.mtp_layer_count))) + run_report(command, args.report_out) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--model-root", default="/target") + parser.add_argument("--model-pattern", required=True) + parser.add_argument("--expected-parts", type=int, default=1) + parser.add_argument("--mtp-draft", required=True) + parser.add_argument("--projector", required=True) + parser.add_argument("--projector-url") + parser.add_argument("--projector-local-path", default="/tmp/mmproj.gguf") + parser.add_argument("--projector-only", action="store_true") + parser.add_argument("--layer-count", type=int, required=True) + parser.add_argument("--mtp-layer-count", type=int) + parser.add_argument("--ctx-size", type=int, default=64) + parser.add_argument("--report-out", default="/results/mtp-attach-certification.json") + parser.add_argument("--mesh-repo", default="https://github.com/Mesh-LLM/mesh-llm.git") + parser.add_argument("--mesh-revision", required=True) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + ensure_build_tools() + mesh_root = Path("/tmp/mesh-llm") + checkout_mesh(args.mesh_repo, args.mesh_revision, mesh_root) + certify(args, mesh_root) + + +if __name__ == "__main__": + main() diff --git a/scripts/tests/test_hf_skippy_convert_job.py b/scripts/tests/test_hf_skippy_convert_job.py new file mode 100644 index 000000000..958483a5b --- /dev/null +++ b/scripts/tests/test_hf_skippy_convert_job.py @@ -0,0 +1,56 @@ +import importlib.util +import json +import tempfile +import unittest +from pathlib import Path + + +SCRIPT = Path(__file__).parents[1] / "hf-skippy-convert-job.py" +SPEC = importlib.util.spec_from_file_location("hf_skippy_convert_job", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +MODULE = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(MODULE) + + +class ConvertedArtifactValidationTests(unittest.TestCase): + def write_artifact(self, root: Path, expected_splits: int) -> None: + (root / "README.md").write_text("beta", encoding="utf-8") + (root / "skippy-convert-manifest.json").write_text( + json.dumps( + { + "expected_splits": expected_splits, + "output_basename": "Inkling-BF16", + } + ), + encoding="utf-8", + ) + + def test_accepts_every_declared_split(self): + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + self.write_artifact(root, 2) + (root / "Inkling-BF16-00001-of-00002.gguf").write_bytes(b"one") + (root / "Inkling-BF16-00002-of-00002.gguf").write_bytes(b"two") + + MODULE.validate_converted_artifact(root) + + def test_rejects_interrupted_multi_split_conversion(self): + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + self.write_artifact(root, 2) + (root / "Inkling-BF16-00001-of-00002.gguf").write_bytes(b"one") + + with self.assertRaisesRegex(FileNotFoundError, "00002-of-00002"): + MODULE.validate_converted_artifact(root) + + def test_accepts_declared_unsplit_output(self): + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + self.write_artifact(root, 1) + (root / "Inkling-BF16.gguf").write_bytes(b"one") + + MODULE.validate_converted_artifact(root) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/tests/test_hf_skippy_mtp_certify_job.py b/scripts/tests/test_hf_skippy_mtp_certify_job.py new file mode 100644 index 000000000..8461ce02b --- /dev/null +++ b/scripts/tests/test_hf_skippy_mtp_certify_job.py @@ -0,0 +1,55 @@ +import importlib.util +import io +import unittest +from pathlib import Path +from unittest import mock + + +SCRIPT = Path(__file__).parents[1] / "hf-skippy-mtp-certify-job.py" +SPEC = importlib.util.spec_from_file_location("hf_skippy_mtp_certify_job", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +MODULE = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(MODULE) + + +class ProjectorDownloadSafetyTests(unittest.TestCase): + @mock.patch.object(MODULE.socket, "getaddrinfo") + def test_accepts_public_hugging_face_https_url(self, getaddrinfo): + getaddrinfo.return_value = [ + (MODULE.socket.AF_INET, MODULE.socket.SOCK_STREAM, 6, "", ("13.33.88.1", 443)) + ] + + parsed = MODULE.validate_projector_url( + "https://huggingface.co/org/model/resolve/main/mmproj.gguf?download=true" + ) + + self.assertEqual(parsed.hostname, "huggingface.co") + + def test_rejects_untrusted_projector_host(self): + with self.assertRaisesRegex(RuntimeError, "untrusted projector URL host"): + MODULE.validate_projector_url("https://example.com/mmproj.gguf") + + @mock.patch.object(MODULE.socket, "getaddrinfo") + def test_rejects_private_resolution_for_trusted_host(self, getaddrinfo): + getaddrinfo.return_value = [ + (MODULE.socket.AF_INET, MODULE.socket.SOCK_STREAM, 6, "", ("169.254.169.254", 443)) + ] + + with self.assertRaisesRegex(RuntimeError, "non-public address"): + MODULE.validate_projector_url("https://huggingface.co/mmproj.gguf") + + def test_copy_rejects_body_over_limit(self): + response = mock.Mock() + response.headers = {} + response.read = mock.Mock( + side_effect=[b"GGUF", b"x", b""], + ) + output = io.BytesIO() + + with mock.patch.object(MODULE, "PROJECTOR_DOWNLOAD_MAX_BYTES", 4): + with self.assertRaisesRegex(RuntimeError, "maximum supported size"): + MODULE.copy_projector_response(response, output) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/tests/test_windows_native_runtime_deps.py b/scripts/tests/test_windows_native_runtime_deps.py index 7ee63fbf3..637695945 100644 --- a/scripts/tests/test_windows_native_runtime_deps.py +++ b/scripts/tests/test_windows_native_runtime_deps.py @@ -129,7 +129,7 @@ class WindowsNativeRuntimeDepsTests(unittest.TestCase): "runtime": { "id": artifact.name, "mesh_version": "0.72.1", - "skippy_abi": "0.1.32", + "skippy_abi": "0.1.35", "platform": { "os": "windows", "arch": "x86_64", diff --git a/third_party/llama.cpp/patches/0048-Support-Laguna-staged-execution.patch b/third_party/llama.cpp/patches/0048-Support-Laguna-staged-execution.patch new file mode 100644 index 000000000..94dc4f1a0 --- /dev/null +++ b/third_party/llama.cpp/patches/0048-Support-Laguna-staged-execution.patch @@ -0,0 +1,72 @@ +From 3140a1e0aa272ccaf89d280fa14f293e9234ed08 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Mon, 27 Jul 2026 17:19:02 +1000 +Subject: [PATCH 48/61] Support Laguna staged execution + +--- + src/models/laguna.cpp | 19 ++++++++++++++++--- + src/skippy.cpp | 1 + + 2 files changed, 17 insertions(+), 3 deletions(-) + +diff --git a/src/models/laguna.cpp b/src/models/laguna.cpp +index fb55ec12f..7144a389c 100644 +--- a/src/models/laguna.cpp ++++ b/src/models/laguna.cpp +@@ -157,7 +157,12 @@ llama_model_laguna::graph::graph(const llama_model & model, const llm_graph_para + ggml_tensor * cur; + ggml_tensor * inpL; + +- inpL = build_inp_embd(model.tok_embd); ++ const skippy_graph_filter & stage_filter = skippy_graph_get_filter(); ++ const bool stage_filtered = stage_filter.enabled; ++ const int il_start = stage_filtered ? stage_filter.layer_start : 0; ++ const int il_end = stage_filtered ? stage_filter.layer_end : n_layer; ++ ++ inpL = build_inp_embd(stage_filtered && il_start > 0 ? nullptr : model.tok_embd); + // No MuP embedding scale (laguna omits this; afmoe scales by sqrt(hidden)). + + ggml_tensor * inp_pos = build_inp_pos(); +@@ -166,11 +171,11 @@ llama_model_laguna::graph::graph(const llama_model & model, const llm_graph_para + const bool has_swa = hparams.swa_type != LLAMA_SWA_TYPE_NONE; + llm_graph_input_attn_kv * inp_attn_kv = has_swa ? nullptr : build_attn_inp_kv(); + llm_graph_input_attn_kv_iswa * inp_attn_iswa = has_swa ? build_attn_inp_kv_iswa() : nullptr; +- ggml_tensor * inp_out_ids = build_inp_out_ids(); ++ ggml_tensor * inp_out_ids = (!stage_filtered || stage_filter.include_output) ? build_inp_out_ids() : nullptr; + + const float kq_scale = 1.0f / sqrtf(float(n_embd_head)); + +- for (int il = 0; il < n_layer; ++il) { ++ for (int il = il_start; il < il_end; ++il) { + const bool is_swa_il = hparams.is_swa(il); + const int64_t n_head_il = hparams.n_head(il); + const int64_t n_head_kv_il = hparams.n_head_kv(il); +@@ -320,6 +325,14 @@ llama_model_laguna::graph::graph(const llama_model & model, const llm_graph_para + } + + cur = inpL; ++ ++ if (stage_filtered && !stage_filter.include_output) { ++ cb(cur, "stage_boundary", il_end - 1); ++ res->t_embd = cur; ++ ggml_build_forward_expand(gf, cur); ++ return; ++ } ++ + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 21db2b77d..3965e0252 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -4191,6 +4191,7 @@ static enum skippy_status skippy_finish_model_open( + model->arch != LLM_ARCH_JAIS && + model->arch != LLM_ARCH_JAIS2 && + model->arch != LLM_ARCH_JAMBA && ++ model->arch != LLM_ARCH_LAGUNA && + model->arch != LLM_ARCH_LFM2 && + model->arch != LLM_ARCH_LLADA && + model->arch != LLM_ARCH_LLADA_MOE && +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0049-Filter-staged-runtime-memory-to-layer-range.patch b/third_party/llama.cpp/patches/0049-Filter-staged-runtime-memory-to-layer-range.patch new file mode 100644 index 000000000..ae62f3f39 --- /dev/null +++ b/third_party/llama.cpp/patches/0049-Filter-staged-runtime-memory-to-layer-range.patch @@ -0,0 +1,72 @@ +From 426444c8c8ee62b6322ac6794ada658df0f077f4 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Tue, 28 Jul 2026 14:23:09 +1000 +Subject: [PATCH 49/61] Filter staged runtime memory to layer range + +Skippy graph execution already filters to the assigned layer range, but +llama_model::create_memory still allocated cache rows for the complete model. +Compose the active stage filter with each memory implementation's architecture +filter so split capacity estimates match the native allocation. Keep MTP +sidecar contexts on their existing next-token layer filter. +--- + src/llama-model.cpp | 23 ++++++++++++++++++++++- + 1 file changed, 22 insertions(+), 1 deletion(-) + +diff --git a/src/llama-model.cpp b/src/llama-model.cpp +index 442387e6a..b23a7e8d8 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -2047,6 +2047,22 @@ ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int i + return layers[il].rope_short; + } + ++static llama_memory_i::layer_filter_cb skippy_stage_memory_filter( ++ llama_memory_i::layer_filter_cb filter, ++ enum llama_context_type ctx_type) { ++ const skippy_graph_filter stage_filter = skippy_graph_get_filter(); ++ if (!stage_filter.enabled || ctx_type == LLAMA_CONTEXT_TYPE_MTP) { ++ return filter; ++ } ++ ++ const int32_t layer_start = stage_filter.layer_start; ++ const int32_t layer_end = stage_filter.layer_end; ++ return [filter = std::move(filter), layer_start, layer_end](int32_t il) { ++ const bool in_stage = il >= layer_start && il < layer_end; ++ return in_stage && (!filter || filter(il)); ++ }; ++} ++ + llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const llama_cparams & cparams) const { + llama_memory_i * res; + +@@ -2107,7 +2123,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, + std::max((uint32_t) 1, cparams.n_seq_max), + cparams.n_seq_max, + cparams.n_rs_seq, +- nullptr); ++ skippy_stage_memory_filter(nullptr, params.ctx_type)); + } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen35) { + // The main difference between hybrid architectures is the + // layer filters, so pick the right one here +@@ -2132,6 +2148,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, + }; + } + ++ filter_attn = skippy_stage_memory_filter(std::move(filter_attn), params.ctx_type); ++ filter_recr = skippy_stage_memory_filter(std::move(filter_recr), params.ctx_type); ++ + if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { + // Use hybrid-iswa for hybrid models with SWA + res = new llama_memory_hybrid_iswa( +@@ -2209,6 +2228,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, + } + } + ++ filter = skippy_stage_memory_filter(std::move(filter), params.ctx_type); ++ + if (arch == LLM_ARCH_DEEPSEEK4) { + GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); + +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0050-Add-TML-Inkling-architecture.patch b/third_party/llama.cpp/patches/0050-Add-TML-Inkling-architecture.patch new file mode 100644 index 000000000..40d1adc6d --- /dev/null +++ b/third_party/llama.cpp/patches/0050-Add-TML-Inkling-architecture.patch @@ -0,0 +1,5828 @@ +From 4508ca657d370fe18b558ec12bcc3e4294925bcc Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Wed, 15 Jul 2026 18:14:40 +0000 +Subject: [PATCH 50/61] Add TML Inkling architecture + +Hybrid attention model: 55 sliding-window plus 11 global layers, banded +content-dependent relative position bias instead of RoPE, per-layer short +convolution state, fine-grained MoE (256 experts top-6 plus 2 shared), +attention log-scaling past 128K, 1M context. + +Includes the GGML_OP_FLASH_ATTN_EXT_BANDED operator (CPU and CUDA, fused +into the MMA flash attention kernel with an fp16 accumulator overflow +guard), HF to GGUF conversion, chat template with typed content block +parsing (interleaved thinking, narration and tool calls), mmproj vision +and audio support, and backend op tests at production shapes. +--- + common/chat-peg-parser.cpp | 8 +- + common/chat.cpp | 121 +++++ + conversion/__init__.py | 2 + + conversion/base.py | 31 ++ + conversion/inkling.py | 350 +++++++++++++ + ggml/include/ggml-rpc.h | 2 +- + ggml/include/ggml.h | 20 +- + ggml/src/ggml-backend-meta.cpp | 13 + + ggml/src/ggml-cpu/ggml-cpu.c | 3 + + ggml/src/ggml-cpu/ops.cpp | 44 +- + ggml/src/ggml-cuda/argsort.cu | 28 +- + ggml/src/ggml-cuda/fattn-banded.cu | 247 +++++++++ + ggml/src/ggml-cuda/fattn-banded.cuh | 7 + + ggml/src/ggml-cuda/fattn-common.cuh | 17 +- + ggml/src/ggml-cuda/fattn-mma-f16.cuh | 111 +++- + ggml/src/ggml-cuda/fattn.cu | 10 +- + ggml/src/ggml-cuda/ggml-cuda.cu | 68 ++- + ggml/src/ggml-cuda/mmf.cuh | 8 +- + ggml/src/ggml-cuda/mmq.cuh | 14 +- + ggml/src/ggml-cuda/mmvf.cu | 4 +- + ggml/src/ggml-cuda/mmvq.cu | 4 +- + ggml/src/ggml-cuda/pad.cu | 2 +- + ggml/src/ggml-cuda/ssm-conv.cu | 28 +- + ggml/src/ggml-rpc/ggml-rpc.cpp | 1 + + ggml/src/ggml.c | 58 ++- + gguf-py/gguf/constants.py | 50 ++ + gguf-py/gguf/tensor_mapping.py | 49 ++ + models/templates/Inkling.jinja | 514 ++++++++++++++++++ + src/llama-arch.cpp | 35 ++ + src/llama-arch.h | 24 + + src/llama-graph.cpp | 30 +- + src/llama-graph.h | 3 +- + src/llama-hparams.cpp | 5 + + src/llama-hparams.h | 11 + + src/llama-kv-cache.cpp | 97 ++++ + src/llama-kv-cache.h | 13 +- + src/llama-model-saver.cpp | 1 + + src/llama-model.cpp | 6 +- + src/llama-model.h | 9 + + src/llama-quant.cpp | 10 + + src/llama-vocab.cpp | 10 + + src/llama-vocab.h | 1 + + src/models/inkling.cpp | 695 +++++++++++++++++++++++++ + src/models/models.h | 13 + + tests/CMakeLists.txt | 8 + + tests/test-backend-ops.cpp | 116 ++++- + tests/test-chat.cpp | 79 +++ + tests/test-flash-attn-bias.cpp | 525 +++++++++++++++++++ + tests/test-flash-attn-generic-hash.cpp | 142 +++++ + tests/test-llama-archs.cpp | 3 + + tools/mtmd/CMakeLists.txt | 1 + + tools/mtmd/clip-impl.h | 9 + + tools/mtmd/clip-model.h | 13 +- + tools/mtmd/clip.cpp | 62 ++- + tools/mtmd/models/inkling.cpp | 106 ++++ + tools/mtmd/models/models.h | 11 + + tools/mtmd/mtmd-audio.cpp | 71 +++ + tools/mtmd/mtmd-audio.h | 10 + + tools/mtmd/mtmd-image.cpp | 203 +++++++- + tools/mtmd/mtmd-image.h | 20 + + tools/mtmd/mtmd.cpp | 37 +- + 61 files changed, 4068 insertions(+), 125 deletions(-) + create mode 100644 conversion/inkling.py + create mode 100644 ggml/src/ggml-cuda/fattn-banded.cu + create mode 100644 ggml/src/ggml-cuda/fattn-banded.cuh + create mode 100644 models/templates/Inkling.jinja + create mode 100644 src/models/inkling.cpp + create mode 100644 tests/test-flash-attn-bias.cpp + create mode 100644 tests/test-flash-attn-generic-hash.cpp + create mode 100644 tools/mtmd/models/inkling.cpp + +diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp +index f786f5ff2..60403d374 100644 +--- a/common/chat-peg-parser.cpp ++++ b/common/chat-peg-parser.cpp +@@ -631,10 +631,10 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key( + // Arguments — either wrapped in args_key or parsed directly + common_peg_parser args_parser = eps(); + if (args_key.empty()) { +- args_parser = tool_args(schema(json(), "tool-" + name + "-schema", params)); ++ args_parser = tool_args(schema(json_object(), "tool-" + name + "-schema", params)); + } else { + args_parser = literal("\"" + effective_args_key + "\"") + space() + literal(":") + space() + +- tool_args(schema(json(), "tool-" + name + "-schema", params)); ++ tool_args(schema(json_object(), "tool-" + name + "-schema", params)); + } + inner_fields.push_back(args_parser); + +@@ -695,7 +695,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys( + auto nested_name = literal("\"" + nested_name_field + "\"") + space() + literal(":") + space() + + atomic(literal("\"") + tool_name(literal(name)) + literal("\"")); + auto nested_args = literal("\"" + nested_args_field + "\"") + space() + literal(":") + space() + +- tool_args(schema(json(), "tool-" + name + "-schema", params)); ++ tool_args(schema(json_object(), "tool-" + name + "-schema", params)); + + auto nested_object = literal("{") + space() + + nested_name + space() + literal(",") + space() + +@@ -764,7 +764,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys( + auto tool_name_ = name_key_parser + space() + literal(":") + space() + + atomic(literal("\"") + tool_name(literal(name)) + literal("\"")); + auto tool_args_ = args_key_parser + space() + literal(":") + space() + +- tool_args(schema(json(), "tool-" + name + "-schema", params)); ++ tool_args(schema(json_object(), "tool-" + name + "-schema", params)); + + // Build ID parsers if keys are provided + common_peg_parser id_parser = eps(); +diff --git a/common/chat.cpp b/common/chat.cpp +index 71871b3df..ae1569ab0 100644 +--- a/common/chat.cpp ++++ b/common/chat.cpp +@@ -2641,6 +2641,119 @@ static common_chat_params common_chat_params_init_minimax_m3(const common_chat_t + return data; + } + ++// Inkling / TML typed-content-block parser: <|end_message|> separates blocks within a turn, ++// <|content_model_end_sampling|> is the sole end-of-generation token (mirrors sglang TmlDetector). ++static common_chat_params common_chat_params_init_inkling(const common_chat_template & tmpl, ++ const autoparser::generation_params & inputs) { ++ common_chat_params data; ++ ++ const std::string MSG_MODEL = "<|message_model|>"; ++ const std::string MSG_USER = "<|message_user|>"; ++ const std::string MSG_SYSTEM = "<|message_system|>"; ++ const std::string MSG_TOOL = "<|message_tool|>"; ++ const std::string THINK = "<|content_thinking|>"; ++ const std::string TEXT = "<|content_text|>"; ++ const std::string END_MESSAGE = "<|end_message|>"; ++ const std::string END_SAMPLING = "<|content_model_end_sampling|>"; ++ const std::string INVOKE_TOOL = "<|content_invoke_tool_json|>"; ++ ++ data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); ++ data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); ++ data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; ++ data.supports_thinking = true; ++ data.thinking_start_tag = THINK; ++ data.thinking_end_tags = {END_MESSAGE}; ++ data.preserved_tokens = { ++ MSG_MODEL, MSG_USER, MSG_SYSTEM, MSG_TOOL, ++ THINK, TEXT, END_MESSAGE, END_SAMPLING, INVOKE_TOOL, ++ }; ++ ++ auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); ++ data.message_delimiters = { ++ { COMMON_CHAT_ROLE_ASSISTANT, MSG_MODEL }, ++ { COMMON_CHAT_ROLE_USER, MSG_USER }, ++ { COMMON_CHAT_ROLE_SYSTEM, MSG_SYSTEM }, ++ { COMMON_CHAT_ROLE_TOOL, MSG_TOOL }, ++ }; ++ ++ auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; ++ ++ if (inputs.has_continuation()) { ++ const auto & msg = inputs.continue_msg; ++ ++ data.generation_prompt = MSG_MODEL + THINK + msg.reasoning_content; ++ if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { ++ data.generation_prompt += END_MESSAGE + TEXT + msg.render_content(); ++ } ++ ++ data.prompt += data.generation_prompt; ++ } ++ ++ auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { ++ auto generation_prompt = p.literal(MSG_MODEL); ++ auto end = p.end(); ++ ++ // thinking block; may also reappear mid-turn (after content), so it is both an optional ++ // prefix and a choice inside the block loops. With reasoning_format=NONE keep it ++ // (markers included) inline as content ++ common_peg_parser reasoning_block = p.eps(); ++ if (extract_reasoning) { ++ reasoning_block = p.literal(THINK) + ++ p.reasoning(p.until_one_of({ END_MESSAGE, TEXT, END_SAMPLING })) + ++ p.optional(p.literal(END_MESSAGE)); ++ } else { ++ reasoning_block = p.content(p.literal(THINK) + ++ p.until_one_of({ END_MESSAGE, TEXT, END_SAMPLING }) + ++ p.optional(p.literal(END_MESSAGE))); ++ } ++ auto reasoning = p.optional(reasoning_block); ++ ++ // TML re-emits <|message_model|> before each content block; a turn may contain several ++ // text blocks (one per content part), so the block repeats and bodies concatenate. ++ // THINK stops the content scan so a mid-turn thinking block is never leaked as text ++ auto text_block = p.optional(p.literal(MSG_MODEL)) + ++ p.optional(p.literal(TEXT)) + ++ p.content(p.until_one_of({ THINK, END_MESSAGE, END_SAMPLING })) + ++ p.optional(p.literal(END_MESSAGE)); ++ auto text_content = p.one_or_more(p.choice({ reasoning_block, text_block })); ++ ++ if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { ++ return generation_prompt + reasoning + text_content + ++ p.optional(p.literal(END_SAMPLING)) + end; ++ } ++ ++ // each call is its own block (role opener + bare name echo + JSON section); ++ // force_tool_calls=true makes the JSON section required so a pure-text answer fails the ++ // block cleanly; parallel calls are separate blocks, hence repeat + parallel=false ++ auto tool_section = p.standard_json_tools( ++ INVOKE_TOOL, END_MESSAGE, inputs.tools, /* parallel_tool_calls = */ false, ++ /* force_tool_calls = */ true, ++ /* name_key = */ "name", ++ /* args_key = */ "args", ++ /* array_wrapped = */ false); ++ // the name-echo scan must stop at any block marker: a greedy until(INVOKE_TOOL) returns ++ // NEED_MORE_INPUT mid-stream, which choice() treats as a match and shadows the text branch ++ auto tool_block = p.optional(p.literal(MSG_MODEL)) + ++ p.until_one_of({ INVOKE_TOOL, TEXT, THINK, END_MESSAGE, END_SAMPLING }) + ++ tool_section; ++ auto tool_calls = inputs.parallel_tool_calls ? p.one_or_more(tool_block) : tool_block; ++ // turns may interleave narration, thinking and calls; parse block-by-block (tool block ++ // first) since a whole-body choice would let the text branch swallow tool blocks into ++ // visible content ++ auto mixed_body = p.one_or_more(p.choice({ tool_block, reasoning_block, text_block })); ++ auto body = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ++ ? tool_calls ++ : mixed_body; ++ ++ return generation_prompt + reasoning + body + ++ p.optional(p.literal(END_SAMPLING)) + end; ++ }); ++ ++ data.parser = parser.save(); ++ ++ return data; ++} ++ + namespace workaround { + + static void map_developer_role_to_system(json & messages) { +@@ -3058,6 +3171,14 @@ std::optional common_chat_try_specialized_template( + return common_chat_params_init_cohere2moe(tmpl, params); + } + ++ // Inkling / TML: this marker combination is unique to the template ++ if (src.find("<|content_thinking|>") != std::string::npos && ++ src.find("<|content_text|>") != std::string::npos && ++ src.find("<|message_model|>") != std::string::npos) { ++ LOG_DBG("Using specialized template: Inkling\n"); ++ return common_chat_params_init_inkling(tmpl, params); ++ } ++ + if (is_lfm2_template(src)) { + LOG_DBG("Using specialized template: LFM2\n"); + return common_chat_params_init_lfm2(tmpl, params, /* tool_list_tokens = */ true); +diff --git a/conversion/__init__.py b/conversion/__init__.py +index 1a47b851a..abdf76cdb 100644 +--- a/conversion/__init__.py ++++ b/conversion/__init__.py +@@ -111,6 +111,7 @@ TEXT_MODEL_MAP: dict[str, str] = { + "HunYuanVLForConditionalGeneration": "hunyuan", + "HYV3ForCausalLM": "hunyuan", + "IQuestCoderForCausalLM": "llama", ++ "InklingForConditionalGeneration": "inkling", + "InternLM2ForCausalLM": "internlm", + "InternLM3ForCausalLM": "internlm", + "JAISLMHeadModel": "jais", +@@ -279,6 +280,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = { + "GraniteSpeechPlusForConditionalGeneration": "granite", + "HunYuanVLForConditionalGeneration": "hunyuan", + "Idefics3ForConditionalGeneration": "smolvlm", ++ "InklingForConditionalGeneration": "inkling", + "InternVisionModel": "internvl", + "JanusForConditionalGeneration": "januspro", + "KimiK25ForConditionalGeneration": "kimivl", +diff --git a/conversion/base.py b/conversion/base.py +index a7cd3fd90..720cfdca3 100644 +--- a/conversion/base.py ++++ b/conversion/base.py +@@ -268,8 +268,17 @@ class ModelBase: + data_gen = lambda data=data_torch: LazyTorchTensor.from_eager(data) # noqa: E731 + else: + data_gen = lambda data=data_torch: data # noqa: E731 ++ # the index maps each tensor to one shard; a duplicate would silently overwrite it ++ if weight_map and name in weight_map and weight_map[name] != part_name: ++ raise ValueError( ++ f"tensor '{name}' found in '{part_name}' but the index assigns " ++ f"it to '{weight_map[name]}'; refusing to load a wrong-shard copy") + if titem := self.filter_tensors((name, data_gen)): + tname, tgen = titem ++ if tname in tensors: ++ raise ValueError( ++ f"duplicate tensor '{tname}' found in multiple model parts; " ++ f"refusing to silently overwrite") + tensors[tname] = tgen + + # verify tensor name presence and identify potentially missing files +@@ -537,6 +546,20 @@ class ModelBase: + else: + raise NotImplementedError(f"Quant format {quant_format!r} for method {quant_method!r} is not yet supported") + elif quant_method == "modelopt": ++ # Stacked-expert NVFP4 checkpoints (e.g. Inkling-NVFP4) store experts as ++ # w13_weight/w2_weight with .scale/.scale2/.input_amax/.original_shape ++ # auxiliaries; the main tensors do not end in .weight, so the NVFP4 path ++ # below would silently skip them. Reject them with a clear error. ++ stacked_expert_aux = [ ++ n for n in self.model_tensors ++ if n.endswith((".scale2", ".input_amax", ".original_shape")) ++ ] ++ if stacked_expert_aux: ++ raise NotImplementedError( ++ "This checkpoint stores quantized experts in the stacked ModelOpt NVFP4 layout " ++ f"({len(stacked_expert_aux)} auxiliary tensors like {stacked_expert_aux[0]!r}), " ++ "which is not supported yet. Convert from the unquantized (BF16) checkpoint instead." ++ ) + # Mixed-precision ModelOpt models: NVFP4 tensors are handled by + # _generate_nvfp4_tensors; FP8 tensors have 1D weight_scale and + # are dequantized here. k/v scale tensors are unused. +@@ -1023,6 +1046,14 @@ class ModelBase: + + def write(self): + self.prepare_tensors() ++ # zero tensors means the shards were never discovered, yet a metadata-only GGUF logs success ++ n_written = sum(len(shard) for shard in self.gguf_writer.tensors) ++ if n_written == 0: ++ raise ValueError( ++ "no tensors were written: the model shards could not be found. " ++ "Check that the safetensors filenames are discoverable (they must " ++ "start with 'model') or that model.safetensors exists." ++ ) + self.prepare_metadata(vocab_only=False) + self.gguf_writer.write_header_to_file(path=self.fname_out) + self.gguf_writer.write_kv_data_to_file() +diff --git a/conversion/inkling.py b/conversion/inkling.py +new file mode 100644 +index 000000000..90ef65c3b +--- /dev/null ++++ b/conversion/inkling.py +@@ -0,0 +1,350 @@ ++from __future__ import annotations ++ ++from typing import Callable, Iterable, TYPE_CHECKING ++ ++if TYPE_CHECKING: ++ from torch import Tensor ++ ++from .base import MmprojModel, ModelBase, TextModel, gguf, logger ++ ++ ++@ModelBase.register("InklingForConditionalGeneration") ++class InklingModel(TextModel): ++ model_arch = gguf.MODEL_ARCH.INKLING ++ undo_permute = False ++ ++ _SKIP_PREFIXES = ("model.visual.", "model.audio.", "model.mtp.") ++ ++ def __init__(self, *args, **kwargs): ++ super().__init__(*args, **kwargs) ++ # explicit raises (not assert, stripped by python -O) guard the single supported variant ++ hp = self.hparams ++ ++ # normalize keys renamed by HF-port re-saved configs back to checkpoint names ++ if "dense_intermediate_size" not in hp and "moe_intermediate_size" in hp: ++ hp["dense_intermediate_size"] = hp["intermediate_size"] ++ hp["intermediate_size"] = hp["moe_intermediate_size"] ++ if "sconv_kernel_size" not in hp and "conv_kernel_size" in hp: ++ hp["sconv_kernel_size"] = hp["conv_kernel_size"] ++ if "dense_mlp_idx" not in hp and hp.get("mlp_layer_types"): ++ types = hp["mlp_layer_types"] ++ hp["dense_mlp_idx"] = next((i for i, t in enumerate(types) if t != "dense"), len(types)) ++ ++ if hp.get("gate_activation", "sigmoid") != "sigmoid": ++ raise NotImplementedError( ++ f"unsupported gate_activation {hp.get('gate_activation')!r}; only 'sigmoid' is implemented" ++ ) ++ for flag, want in ( ++ ("norm_after_topk", True), ++ ("shared_expert_sink", True), ++ ("use_sconv", True), ++ ("use_embed_norm", True), ++ ("use_gate_bias", True), ++ ("use_global_scale", True), ++ ): ++ if hp.get(flag, want) is not want: ++ raise NotImplementedError(f"unsupported {flag}={hp.get(flag)!r}; only {want} is implemented") ++ if hp.get("q_bias", False) is not False or hp.get("o_bias", False) is not False: ++ raise NotImplementedError("attention q_bias / o_bias are not supported") ++ if hp.get("final_logit_softcapping") not in (None, 0, 0.0): ++ raise NotImplementedError( ++ f"final_logit_softcapping={hp.get('final_logit_softcapping')!r} is not supported" ++ ) ++ if hp["swa_head_dim"] != hp["head_dim"]: ++ raise ValueError(f"swa_head_dim {hp['swa_head_dim']} must equal head_dim {hp['head_dim']}") ++ if hp["swa_num_attention_heads"] != hp["num_attention_heads"]: ++ raise ValueError( ++ f"swa_num_attention_heads {hp['swa_num_attention_heads']} must equal " ++ f"num_attention_heads {hp['num_attention_heads']}" ++ ) ++ ++ # context length comes from model_max_length per the design contract ++ if (mml := hp.get("model_max_length")) is not None: ++ self.hparams["max_position_embeddings"] = mml ++ ++ # checked by the base find_hparam list before the MoE "intermediate_size" ++ self.hparams["prefix_dense_intermediate_size"] = hp["dense_intermediate_size"] ++ ++ self._local_layer_flags = self._get_local_layer_flags() ++ self.hparams["num_key_value_heads"] = [ ++ hp["swa_num_key_value_heads"] if is_local else hp["num_key_value_heads"] ++ for is_local in self._local_layer_flags ++ ] ++ ++ def _get_local_layer_flags(self) -> list[bool]: ++ # local_layer_ids is authoritative; a round-tripped layer_types may be stale and must not override it ++ n_layer = self.hparams["num_hidden_layers"] ++ local_ids = self.hparams.get("local_layer_ids") ++ if local_ids is None: ++ # default: global at id % 6 == 5; omitted/null must not collapse to all-global (explicit [] does) ++ local_ids = [i for i in range(n_layer) if i % 6 != 5] ++ local_ids = set(local_ids) ++ return [i in local_ids for i in range(n_layer)] ++ ++ def get_vocab_base(self) -> tuple[list[str], list[int], str]: ++ tokens, toktypes, tokpre = super().get_vocab_base() ++ # dedicated pre-type: o200k-family regex that keeps combining marks attached to base letters ++ tokpre = "inkling" ++ import gguf as _gguf ++ n_vocab = self.hparams["vocab_size"] ++ n_unpadded = self.hparams.get("unpadded_vocab_size") or n_vocab ++ if len(tokens) != n_vocab: ++ raise ValueError(f"Inkling tokenizer produced {len(tokens)} entries, expected {n_vocab}") ++ # force-CONTROL special ids from added_tokens_decoder, else the trailing-60 convention ++ try: ++ import json as _json ++ import pathlib as _pl ++ tc = _json.loads((_pl.Path(self.dir_model) / "tokenizer_config.json").read_text()) ++ special_ids = sorted(int(i) for i, d in tc.get("added_tokens_decoder", {}).items() if d.get("special")) ++ except Exception: ++ special_ids = list(range(n_unpadded - 60, n_unpadded)) ++ for tid in special_ids: ++ if 0 <= tid < n_vocab: ++ toktypes[tid] = _gguf.TokenType.CONTROL ++ if any(t != _gguf.TokenType.UNUSED for t in toktypes[n_unpadded:]): ++ raise ValueError("real tokens found at/above unpadded_vocab_size; padded-vocab mask would hide them") ++ return tokens, toktypes, tokpre ++ ++ def set_vocab(self): ++ self._set_vocab_gpt2() ++ eos_id = int(self.hparams.get("eos_token_id", 200006)) ++ if eos_id < 199998: ++ # HF-port configs re-save generic bos/eos defaults; the real EOS lives at 199998+ ++ eos_id = 200006 ++ # 200006 is the SOLE end-of-generation token; <|end_message|> (200010) is an ++ # intra-turn block separator and must NOT be registered eot/eog ++ self.gguf_writer.add_eos_token_id(eos_id) ++ # no BOS is ever prepended; pin bos to EOS so a stale base-tokenizer bos id never surfaces ++ self.gguf_writer.add_bos_token_id(eos_id) ++ self.gguf_writer.add_add_bos_token(False) ++ ++ def set_gguf_parameters(self): ++ super().set_gguf_parameters() ++ hp = self.hparams ++ ++ self.gguf_writer.add_vocab_size(hp["vocab_size"]) ++ self.gguf_writer.add_expert_feed_forward_length(hp["intermediate_size"]) ++ self.gguf_writer.add_expert_shared_count(hp["n_shared_experts"]) ++ self.gguf_writer.add_expert_weights_scale(hp["route_scale"]) ++ self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) ++ ++ # sliding_window_size is canonical; explicit is-None fallback so a serialized 0 cannot bypass the mismatch check ++ canonical_window = hp["sliding_window_size"] ++ sliding_window = hp.get("sliding_window") ++ if sliding_window is None: ++ sliding_window = canonical_window ++ elif sliding_window != canonical_window: ++ raise ValueError( ++ f"sliding_window {sliding_window} disagrees with sliding_window_size " ++ f"{canonical_window!r}" ++ ) ++ if sliding_window <= 0: ++ raise ValueError(f"sliding_window must be positive, got {sliding_window}") ++ self.gguf_writer.add_sliding_window(sliding_window) ++ # true = local (swa) layer ++ self.gguf_writer.add_sliding_window_pattern(self._local_layer_flags) ++ ++ # no RoPE (arch-determined NONE); custom inkling.* keys per INKLING_DESIGN.md ++ arch = gguf.MODEL_ARCH_NAMES[self.model_arch] ++ self.gguf_writer.add_uint32(f"{arch}.d_rel", hp["d_rel"]) ++ self.gguf_writer.add_uint32(f"{arch}.rel_extent", hp["rel_extent"]) ++ self.gguf_writer.add_uint32(f"{arch}.rel_extent_swa", sliding_window) ++ self.gguf_writer.add_uint32(f"{arch}.shortconv_kernel", hp["sconv_kernel_size"]) ++ self.gguf_writer.add_uint32(f"{arch}.dense_block_count", hp["dense_mlp_idx"]) ++ self.gguf_writer.add_float32(f"{arch}.logit_scale_denom", hp["logits_mup_width_multiplier"]) ++ self.gguf_writer.add_uint32(f"{arch}.log_scaling_n_floor", int(hp.get("log_scaling_n_floor") or 0)) ++ self.gguf_writer.add_float32(f"{arch}.log_scaling_alpha", hp.get("log_scaling_alpha", 0.0)) ++ self.gguf_writer.add_uint32(f"{arch}.unpadded_vocab_size", hp["unpadded_vocab_size"]) ++ ++ logger.info(f"gguf: (inkling) swa pattern (true=local) = {self._local_layer_flags}") ++ logger.info(f"gguf: (inkling) unpadded_vocab_size = {hp['unpadded_vocab_size']}") ++ ++ @classmethod ++ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: ++ name, gen = item ++ ++ if name.startswith(cls._SKIP_PREFIXES): ++ return None ++ ++ name = name.replace("model.llm.", "model.") ++ # parameter has no ".weight"-style suffix in the checkpoint ++ name = name.replace("rel_logits_proj.proj", "rel_logits_proj.weight") ++ ++ return super().filter_tensors((name, gen)) ++ ++ @staticmethod ++ def _deinterleave_w13(w13: Tensor) -> tuple[Tensor, Tensor]: ++ # interleaved SwiGLU along the output rows: silu(z[..., ::2]) * z[..., 1::2] ++ gate = w13[..., 0::2, :].contiguous() ++ up = w13[..., 1::2, :].contiguous() ++ return gate, up ++ ++ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: ++ # short convs: [C, 1, K] -> [C, K] (same layout as LFM2 shortconv.conv) ++ if name.endswith("_sconv.weight"): ++ data_torch = data_torch.squeeze(1) ++ return [(self.map_tensor_name(name), data_torch)] ++ ++ if name.endswith(".mlp.w13_dn.weight"): ++ assert bid is not None ++ gate, up = self._deinterleave_w13(data_torch) ++ return [ ++ (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE, bid), gate), ++ (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP, bid), up), ++ ] ++ ++ if name.endswith(".mlp.global_scale") or name.endswith(".mlp.gate.global_scale"): ++ assert bid is not None ++ return [(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GSCALE, bid), data_torch.float())] ++ ++ if name.endswith(".mlp.gate.bias"): ++ assert bid is not None ++ return [(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_EXP_PROBS_B, bid, ".bias"), data_torch.float())] ++ ++ if name.endswith(".mlp.experts.w13_weight"): ++ assert bid is not None ++ gate, up = self._deinterleave_w13(data_torch) ++ return [ ++ (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), gate), ++ (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), up), ++ ] ++ if name.endswith(".mlp.experts.w2_weight"): ++ assert bid is not None ++ return [(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), data_torch)] ++ ++ # shared experts stored stacked for mul_mat_id ++ if name.endswith(".mlp.shared_experts.shared_w13_weight"): ++ assert bid is not None ++ gate, up = self._deinterleave_w13(data_torch) ++ return [ ++ (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_SHEXP, bid), gate), ++ (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_SHEXP, bid), up), ++ ] ++ if name.endswith(".mlp.shared_experts.shared_w2_weight"): ++ assert bid is not None ++ return [(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, bid), data_torch)] ++ ++ return [(self.map_tensor_name(name), data_torch)] ++ ++ def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int): ++ # used in fp32 rel-bias math; keep full precision ++ if new_name.endswith("attn_rel_proj.weight"): ++ return gguf.GGMLQuantizationType.F32 ++ # ggml_ssm_conv kernels are F32-only ++ if ".shortconv_" in new_name: ++ return gguf.GGMLQuantizationType.F32 ++ return super().tensor_force_quant(name, new_name, bid, n_dims) ++ ++ ++@ModelBase.register("InklingForConditionalGeneration") ++class InklingMmprojModel(MmprojModel): ++ """Export Inkling's hMLP and dMel towers as one mtmd projector.""" ++ ++ has_vision_encoder = True ++ has_audio_encoder = True ++ ++ _IMAGE_MEAN = [0.48145466, 0.4578275, 0.40821073] ++ _IMAGE_STD = [0.26862954, 0.2613026, 0.2757771] ++ ++ def __init__(self, *args, **kwargs): ++ super().__init__(*args, **kwargs) ++ assert self.hparams_vision is not None ++ hp = self.hparams_vision ++ expected = { ++ "vision_encoder_type": "hmlp", ++ "patch_size": 40, ++ "temporal_patch_size": 2, ++ "n_channels": 3, ++ "n_layers": 4, ++ "decoder_dmodel": 6144, ++ "use_vision_norm": True, ++ } ++ for key, want in expected.items(): ++ got = hp.get(key, want) ++ if got != want: ++ raise NotImplementedError( ++ f"Inkling mmproj requires vision_config.{key}={want!r}, got {got!r}" ++ ) ++ ++ assert self.hparams_audio is not None ++ ahp = self.hparams_audio ++ audio_expected = { ++ "audio_mode": "dmel", ++ "decoder_dmodel": 6144, ++ "n_mel_bins": 80, ++ "mel_vocab_size": 16, ++ "use_audio_norm": True, ++ } ++ for key, want in audio_expected.items(): ++ got = ahp.get(key, want) ++ if got != want: ++ raise NotImplementedError( ++ f"Inkling mmproj requires audio_config.{key}={want!r}, got {got!r}" ++ ) ++ ++ def set_gguf_parameters(self): ++ hp = self.hparams_vision ++ assert hp is not None ++ ++ self.gguf_writer.add_file_type(self.ftype) ++ self.gguf_writer.add_clip_has_vision_encoder(True) ++ self.gguf_writer.add_clip_has_audio_encoder(True) ++ self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.INKLING) ++ self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.INKLING) ++ self.gguf_writer.add_vision_projection_dim(hp["decoder_dmodel"]) ++ ++ # clip.cpp requires these common fields even though hMLP is not a ViT. ++ self.gguf_writer.add_vision_image_size(hp["patch_size"]) ++ self.gguf_writer.add_vision_patch_size(hp["patch_size"]) ++ self.gguf_writer.add_vision_embedding_length(hp["n_channels"]) ++ self.gguf_writer.add_vision_feed_forward_length(0) ++ self.gguf_writer.add_vision_block_count(hp["n_layers"]) ++ self.gguf_writer.add_vision_head_count(1) ++ self.gguf_writer.add_vision_attention_layernorm_eps(1e-6) ++ self.gguf_writer.add_vision_image_mean(self._IMAGE_MEAN) ++ self.gguf_writer.add_vision_image_std(self._IMAGE_STD) ++ ++ ahp = self.hparams_audio ++ assert ahp is not None ++ self.gguf_writer.add_audio_projection_dim(ahp["decoder_dmodel"]) ++ self.gguf_writer.add_audio_embedding_length(ahp["decoder_dmodel"]) ++ self.gguf_writer.add_audio_feed_forward_length(0) ++ self.gguf_writer.add_audio_block_count(0) ++ self.gguf_writer.add_audio_head_count(1) ++ self.gguf_writer.add_audio_attention_layernorm_eps(1e-6) ++ self.gguf_writer.add_audio_num_mel_bins(ahp["n_mel_bins"]) ++ ++ @classmethod ++ def filter_tensors(cls, item): ++ name, gen = item ++ if not name.startswith(("model.visual.", "visual.", "model.audio.", "audio.")): ++ return None ++ return name, gen ++ ++ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): ++ del bid ++ if name.startswith(("model.audio.", "audio.")): ++ prefix = "model.audio." if name.startswith("model.audio.") else "audio." ++ local = name.removeprefix(prefix) ++ if local == "encoder.weight": ++ yield "a.dmel.embedding.weight", data_torch ++ return ++ if local == "final_norm.weight": ++ yield "a.dmel.final_norm.weight", data_torch ++ return ++ raise ValueError(f"unexpected Inkling audio tensor {name!r}") ++ ++ prefix = "model.visual." if name.startswith("model.visual.") else "visual." ++ local = name.removeprefix(prefix) ++ if local == "final_norm.weight": ++ yield "v.hmlp.final_norm.weight", data_torch ++ return ++ ++ parts = local.split(".") ++ if len(parts) == 3 and parts[0] == "layers" and parts[2] == "weight": ++ kind, sep, layer_s = parts[1].partition("_") ++ if sep and kind in ("linear", "norm") and layer_s.isdigit(): ++ yield f"v.hmlp.{int(layer_s)}.{kind}.weight", data_torch ++ return ++ ++ raise ValueError(f"unexpected Inkling vision tensor {name!r}") +diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h +index 839363731..086dddd96 100644 +--- a/ggml/include/ggml-rpc.h ++++ b/ggml/include/ggml-rpc.h +@@ -11,7 +11,7 @@ extern "C" { + #define RPC_PROTO_PATCH_VERSION 4 + + #ifdef __cplusplus +-static_assert(GGML_OP_COUNT == 107, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); ++static_assert(GGML_OP_COUNT == 108, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); + #endif + + #define GGML_RPC_MAX_SERVERS 16 +diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h +index 7ef0ec88f..6a5925ce4 100644 +--- a/ggml/include/ggml.h ++++ b/ggml/include/ggml.h +@@ -435,8 +435,9 @@ extern "C" { + + // precision + enum ggml_prec { +- GGML_PREC_DEFAULT = 0, // stored as ggml_tensor.op_params, 0 by default +- GGML_PREC_F32 = 10, ++ GGML_PREC_DEFAULT = 0, // stored as ggml_tensor.op_params, 0 by default ++ GGML_PREC_F32 = 10, ++ GGML_PREC_F32_PEDANTIC = 11, + }; + + // op hint +@@ -558,6 +559,7 @@ extern "C" { + GGML_OP_FILL, + + GGML_OP_FLASH_ATTN_EXT, ++ GGML_OP_FLASH_ATTN_EXT_BANDED, + GGML_OP_FLASH_ATTN_BACK, + GGML_OP_SSM_CONV, + GGML_OP_SSM_SCAN, +@@ -1440,6 +1442,7 @@ extern "C" { + + // change the precision of a matrix multiplication + // set to GGML_PREC_F32 for higher precision (useful for phi-2) ++ // or GGML_PREC_F32_PEDANTIC to require true F32 arithmetic + GGML_API void ggml_mul_mat_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec); +@@ -2439,6 +2442,19 @@ extern "C" { + float max_bias, + float logit_softcap); + ++ // flash attention with an additive banded relative-position bias, applied after scale, no dense bias tensor: ++ // rel_logits: [rel_extent, n_head, n_batch, ne3]; rel_dist = q_idx + (n_kv - n_batch) - kv_idx ++ // score += rel_logits[rel_dist, head, q_idx, batch] iff 0 <= rel_dist < rel_extent ++ GGML_API struct ggml_tensor * ggml_flash_attn_ext_banded( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * k, ++ struct ggml_tensor * v, ++ struct ggml_tensor * mask, ++ struct ggml_tensor * rel_logits, ++ float scale, ++ int64_t rel_extent); ++ + GGML_API void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec); +diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp +index a5a3a58ad..b397fb6ba 100644 +--- a/ggml/src/ggml-backend-meta.cpp ++++ b/ggml/src/ggml-backend-meta.cpp +@@ -753,6 +753,16 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; + }; + ++ auto handle_flash_attn_ext_banded = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { ++ GGML_ASSERT( src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); ++ GGML_ASSERT( src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2); ++ GGML_ASSERT( src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); ++ GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); ++ // rel_logits is [E, H, Q, B], so its head shard is axis 1. ++ GGML_ASSERT( src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_1); ++ return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; ++ }; ++ + auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis == src_ss[1].axis) { + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { +@@ -964,6 +974,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( + case GGML_OP_FLASH_ATTN_EXT: { + split_state = handle_flash_attn_ext(src_ss); + } break; ++ case GGML_OP_FLASH_ATTN_EXT_BANDED: { ++ split_state = handle_flash_attn_ext_banded(src_ss); ++ } break; + case GGML_OP_FLASH_ATTN_BACK: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; +diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c +index eed46cfe2..6f7cc39d1 100644 +--- a/ggml/src/ggml-cpu/ggml-cpu.c ++++ b/ggml/src/ggml-cpu/ggml-cpu.c +@@ -2016,6 +2016,7 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm + ggml_compute_forward_fill(params, tensor); + } break; + case GGML_OP_FLASH_ATTN_EXT: ++ case GGML_OP_FLASH_ATTN_EXT_BANDED: + { + ggml_compute_forward_flash_attn_ext(params, tensor); + } break; +@@ -2438,6 +2439,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { + case GGML_OP_ARGSORT: + case GGML_OP_TOP_K: + case GGML_OP_FLASH_ATTN_EXT: ++ case GGML_OP_FLASH_ATTN_EXT_BANDED: + case GGML_OP_FLASH_ATTN_BACK: + case GGML_OP_SSM_CONV: + case GGML_OP_SSM_SCAN: +@@ -2995,6 +2997,7 @@ struct ggml_cplan ggml_graph_plan( + cur += sizeof(int32_t)*node->src[0]->ne[0]*n_tasks; + } break; + case GGML_OP_FLASH_ATTN_EXT: ++ case GGML_OP_FLASH_ATTN_EXT_BANDED: + { + const int64_t neq2 = node->src[0]->ne[2]; // number of query heads + const int64_t DK = node->src[1]->ne[0]; +diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp +index 61237ed06..9e9b1227b 100644 +--- a/ggml/src/ggml-cpu/ops.cpp ++++ b/ggml/src/ggml-cpu/ops.cpp +@@ -8340,10 +8340,11 @@ template + struct cmp_argsort { + const float * data; + bool operator()(int32_t a, int32_t b) const { ++ // ties must resolve to the lower id (MoE routers); std::sort is unstable + if constexpr (order == GGML_SORT_ORDER_ASC) { +- return data[a] < data[b]; ++ return data[a] < data[b] || (data[a] == data[b] && a < b); + } else { +- return data[a] > data[b]; ++ return data[a] > data[b] || (data[a] == data[b] && a < b); + } + } + }; +@@ -8412,7 +8413,8 @@ void ggml_compute_forward_argsort( + struct cmp_top_k { + const float * data; + bool operator()(int32_t a, int32_t b) const { +- return data[a] > data[b]; ++ // ties must resolve to the lower id so the selected set matches the CUDA backend ++ return data[a] > data[b] || (data[a] == data[b] && a < b); + } + }; + +@@ -8473,6 +8475,30 @@ void ggml_compute_forward_top_k( + } + } + ++static inline float ggml_flash_attn_ext_banded_load( ++ const ggml_tensor * rel, ++ int64_t iq1, ++ int64_t iq2, ++ int64_t iq3, ++ int64_t rel_idx) { ++ const char * ptr = (const char *) rel->data + ++ (size_t) rel_idx * rel->nb[0] + ++ (size_t) iq2 * rel->nb[1] + ++ (size_t) iq1 * rel->nb[2] + ++ (size_t) (iq3 % rel->ne[3]) * rel->nb[3]; ++ ++ switch (rel->type) { ++ case GGML_TYPE_F32: ++ return *(const float *) ptr; ++ case GGML_TYPE_F16: ++ return GGML_CPU_FP16_TO_FP32(*(const ggml_fp16_t *) ptr); ++ case GGML_TYPE_BF16: ++ return GGML_BF16_TO_FP32(*(const ggml_bf16_t *) ptr); ++ default: ++ GGML_ABORT("banded flash attention: unsupported rel_logits type"); ++ } ++} ++ + static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( + const ggml_compute_params * params, + ggml_tensor * dst, +@@ -8486,6 +8512,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * sinks = dst->src[4]; ++ const ggml_tensor * rel = dst->src[5]; + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne) + GGML_TENSOR_LOCALS(size_t, nbq, q, nb) +@@ -8614,6 +8641,14 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( + s = logit_softcap*tanhf(s); + } + ++ if (rel) { ++ // the offset aligns a short decode Q block to the tail of K (FA4 seqlen_k - seqlen_q convention) ++ const int64_t rel_dist = iq1 + (nek1 - neq1) - ic; ++ if (rel_dist >= 0 && rel_dist < rel->ne[0]) { ++ s += ggml_flash_attn_ext_banded_load(rel, iq1, iq2, iq3, rel_dist); ++ } ++ } ++ + s += mv; // apply mask + + const float Mold = M; +@@ -9078,6 +9113,7 @@ static void ggml_compute_forward_flash_attn_ext_f16( + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; ++ const ggml_tensor * rel = dst->src[5]; + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne) + GGML_TENSOR_LOCALS(size_t, nbq, q, nb) +@@ -9177,7 +9213,7 @@ static void ggml_compute_forward_flash_attn_ext_f16( + const int64_t dr = (nr + nchunk - 1) / nchunk; + + static constexpr int64_t Q_TILE_SZ = ggml_fa_tile_config::Q; +- bool use_tiled = !use_ref && ++ bool use_tiled = !use_ref && rel == nullptr && + (q->type == GGML_TYPE_F32 && + kv_is_f32_or_f16 && + k->type == v->type && +diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu +index 26af90025..0352e6bac 100644 +--- a/ggml/src/ggml-cuda/argsort.cu ++++ b/ggml/src/ggml-cuda/argsort.cu +@@ -163,6 +163,22 @@ static inline __device__ void ggml_cuda_swap(T & a, T & b) { + b = tmp; + } + ++// true if ia sorts after ib; padded indices sink to the end, ties break to the lower index (matches CPU cmp_argsort) ++template ++static inline __device__ bool argsort_ranks_after(const float * x_row, int ia, int ib, int ncols) { ++ const bool a_pad = ia >= ncols; ++ const bool b_pad = ib >= ncols; ++ if (a_pad || b_pad) { ++ return a_pad && (!b_pad || ia > ib); ++ } ++ const float xa = x_row[ia]; ++ const float xb = x_row[ib]; ++ if (xa != xb) { ++ return order == GGML_SORT_ORDER_ASC ? (xa > xb) : (xa < xb); ++ } ++ return ia > ib; ++} ++ + template + static __global__ void k_argsort_f32_i32(const float * x, int * dst, const int ncols, int ncols_pad) { + // bitonic sort +@@ -186,19 +202,11 @@ static __global__ void k_argsort_f32_i32(const float * x, int * dst, const int n + int ixj = col ^ j; + if (ixj > col) { + if ((col & k) == 0) { +- if (dst_row[col] >= ncols || +- (dst_row[ixj] < ncols && (order == GGML_SORT_ORDER_ASC ? +- x_row[dst_row[col]] > x_row[dst_row[ixj]] : +- x_row[dst_row[col]] < x_row[dst_row[ixj]])) +- ) { ++ if (argsort_ranks_after(x_row, dst_row[col], dst_row[ixj], ncols)) { + ggml_cuda_swap(dst_row[col], dst_row[ixj]); + } + } else { +- if (dst_row[ixj] >= ncols || +- (dst_row[col] < ncols && (order == GGML_SORT_ORDER_ASC ? +- x_row[dst_row[col]] < x_row[dst_row[ixj]] : +- x_row[dst_row[col]] > x_row[dst_row[ixj]])) +- ) { ++ if (argsort_ranks_after(x_row, dst_row[ixj], dst_row[col], ncols)) { + ggml_cuda_swap(dst_row[col], dst_row[ixj]); + } + } +diff --git a/ggml/src/ggml-cuda/fattn-banded.cu b/ggml/src/ggml-cuda/fattn-banded.cu +new file mode 100644 +index 000000000..98458ae86 +--- /dev/null ++++ b/ggml/src/ggml-cuda/fattn-banded.cu +@@ -0,0 +1,247 @@ ++#include "common.cuh" ++#include "fattn-banded.cuh" ++#include "fattn.cuh" ++ ++#include ++ ++static __device__ __forceinline__ float fattn_banded_load( ++ const char * ptr, const int type) { ++ switch (type) { ++ case GGML_TYPE_F32: ++ return *(const float *) ptr; ++ case GGML_TYPE_F16: ++ return __half2float(*(const half *) ptr); ++ case GGML_TYPE_BF16: { ++ // Read BF16 as raw bits so this kernel needs no native BF16 support; all math stays FP32. ++ const uint32_t bits = uint32_t(*(const uint16_t *) ptr) << 16; ++ return __uint_as_float(bits); ++ } ++ default: ++ return 0.0f; ++ } ++} ++ ++template ++static __global__ void flash_attn_ext_banded_f32( ++ const char * __restrict__ q, ++ const char * __restrict__ k, ++ const char * __restrict__ v, ++ const char * __restrict__ mask, ++ const char * __restrict__ rel, ++ float * __restrict__ dst, ++ float scale, ++ int type_k, ++ int type_v, ++ int type_rel, ++ int64_t n_q, ++ int64_t n_kv, ++ int64_t n_head_q, ++ int64_t n_head_kv, ++ int64_t n_batch, ++ int64_t rel_extent, ++ int64_t mask_ne2, ++ int64_t mask_ne3, ++ uint64_t q_nb1, ++ uint64_t q_nb2, ++ uint64_t q_nb3, ++ uint64_t k_nb0, ++ uint64_t k_nb1, ++ uint64_t k_nb2, ++ uint64_t k_nb3, ++ uint64_t v_nb0, ++ uint64_t v_nb1, ++ uint64_t v_nb2, ++ uint64_t v_nb3, ++ uint64_t m_nb1, ++ uint64_t m_nb2, ++ uint64_t m_nb3, ++ uint64_t r_nb0, ++ uint64_t r_nb1, ++ uint64_t r_nb2, ++ uint64_t r_nb3, ++ int64_t rel_ne3) { ++ constexpr int values_per_lane = D / WARP_SIZE; ++ static_assert(D == 64 || D == 128, "banded FA supports head dimensions 64 and 128"); ++ static_assert(D % WARP_SIZE == 0, "head dimension must be divisible by warp size"); ++ ++ const int lane = threadIdx.x % WARP_SIZE; ++ const int warp = threadIdx.x / WARP_SIZE; ++ const int64_t iq = int64_t(blockIdx.x) * WARPS_PER_BLOCK + warp; ++ const int64_t ih = blockIdx.y; ++ const int64_t ib = blockIdx.z; ++ ++ if (iq >= n_q || ih >= n_head_q || ib >= n_batch) { ++ return; ++ } ++ ++ const int64_t ih_kv = ih / (n_head_q / n_head_kv); ++ const char * q_row = q + uint64_t(iq)*q_nb1 + uint64_t(ih)*q_nb2 + uint64_t(ib)*q_nb3; ++ ++ float q_reg[values_per_lane]; ++ float out[values_per_lane]; ++#pragma unroll ++ for (int j = 0; j < values_per_lane; ++j) { ++ const int d = lane + j*WARP_SIZE; ++ q_reg[j] = *(const float *)(q_row + uint64_t(d)*sizeof(float)); ++ out[j] = 0.0f; ++ } ++ ++ float row_max = -INFINITY; ++ float row_sum = 0.0f; ++ const int64_t q_offset = n_kv - n_q; ++ ++ for (int64_t ik = 0; ik < n_kv; ++ik) { ++ const char * k_row = k + uint64_t(ik)*k_nb1 + uint64_t(ih_kv)*k_nb2 + uint64_t(ib)*k_nb3; ++ float dot = 0.0f; ++#pragma unroll ++ for (int j = 0; j < values_per_lane; ++j) { ++ const int d = lane + j*WARP_SIZE; ++ dot += q_reg[j] * fattn_banded_load(k_row + uint64_t(d)*k_nb0, type_k); ++ } ++ dot = warp_reduce_sum(dot); ++ ++ float score = dot * scale; ++ if (lane == 0) { ++ const int64_t rel_dist = iq + q_offset - ik; ++ if (rel_dist >= 0 && rel_dist < rel_extent) { ++ const char * rel_value = rel + ++ uint64_t(rel_dist)*r_nb0 + uint64_t(ih)*r_nb1 + ++ uint64_t(iq)*r_nb2 + uint64_t(ib % rel_ne3)*r_nb3; ++ score += fattn_banded_load(rel_value, type_rel); ++ } ++ if (mask) { ++ const char * mask_value = mask + uint64_t(ik)*sizeof(half) + ++ uint64_t(iq)*m_nb1 + uint64_t(ih % mask_ne2)*m_nb2 + ++ uint64_t(ib % mask_ne3)*m_nb3; ++ score += __half2float(*(const half *) mask_value); ++ } ++ } ++ score = __shfl_sync(0xffffffff, score, 0, WARP_SIZE); ++ ++ if (score == -INFINITY) { ++ continue; ++ } ++ ++ float old_scale = 1.0f; ++ float value_scale = 1.0f; ++ if (score > row_max) { ++ old_scale = expf(row_max - score); ++ row_max = score; ++ } else { ++ value_scale = expf(score - row_max); ++ } ++ ++ const char * v_row = v + uint64_t(ik)*v_nb1 + uint64_t(ih_kv)*v_nb2 + uint64_t(ib)*v_nb3; ++#pragma unroll ++ for (int j = 0; j < values_per_lane; ++j) { ++ const int d = lane + j*WARP_SIZE; ++ const float vv = fattn_banded_load(v_row + uint64_t(d)*v_nb0, type_v); ++ out[j] = out[j]*old_scale + vv*value_scale; ++ } ++ row_sum = row_sum*old_scale + value_scale; ++ } ++ ++ const float inv_sum = row_sum == 0.0f ? 0.0f : 1.0f/row_sum; ++ float * dst_row = dst + ((ib*n_q + iq)*n_head_q + ih)*D; ++#pragma unroll ++ for (int j = 0; j < values_per_lane; ++j) { ++ const int d = lane + j*WARP_SIZE; ++ dst_row[d] = out[j]*inv_sum; ++ } ++} ++ ++static bool fattn_banded_type_supported(ggml_type type) { ++ return type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16; ++} ++ ++bool ggml_cuda_flash_attn_ext_banded_supported(int device, const ggml_tensor * dst) { ++ GGML_UNUSED(device); ++#if defined(GGML_USE_MUSA) ++ GGML_UNUSED(dst); ++ return false; ++#else ++ if (dst->op != GGML_OP_FLASH_ATTN_EXT_BANDED) { ++ return false; ++ } ++ ++ const ggml_tensor * q = dst->src[0]; ++ const ggml_tensor * k = dst->src[1]; ++ const ggml_tensor * v = dst->src[2]; ++ const ggml_tensor * m = dst->src[3]; ++ const ggml_tensor * rel = dst->src[5]; ++ if (!q || !k || !v || !rel || q->type != GGML_TYPE_F32) { ++ return false; ++ } ++ if (!fattn_banded_type_supported(k->type) || ++ !fattn_banded_type_supported(v->type) || ++ !fattn_banded_type_supported(rel->type)) { ++ return false; ++ } ++ if ((q->ne[0] != 64 && q->ne[0] != 128) || v->ne[0] != q->ne[0] || k->ne[0] != q->ne[0]) { ++ return false; ++ } ++ if (q->ne[2] % k->ne[2] != 0 || q->ne[2] % v->ne[2] != 0 || k->ne[2] != v->ne[2]) { ++ return false; ++ } ++ if (q->ne[3] != k->ne[3] || q->ne[3] != v->ne[3]) { ++ return false; ++ } ++ if (q->nb[0] != sizeof(float) || k->nb[0] != ggml_type_size(k->type) || ++ v->nb[0] != ggml_type_size(v->type) || rel->nb[0] != ggml_type_size(rel->type)) { ++ return false; ++ } ++ if (rel->ne[0] <= 0 || rel->ne[1] != q->ne[2] || rel->ne[2] != q->ne[1] || ++ (rel->ne[3] != 1 && rel->ne[3] != q->ne[3])) { ++ return false; ++ } ++ return !m || (m->type == GGML_TYPE_F16 && ggml_is_contiguous(m) && ++ q->ne[2] % m->ne[2] == 0 && q->ne[3] % m->ne[3] == 0); ++#endif ++} ++ ++void ggml_cuda_flash_attn_ext_banded(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ GGML_ASSERT(ggml_cuda_flash_attn_ext_banded_supported(ctx.device, dst)); ++ ++ const ggml_tensor * q = dst->src[0]; ++ const ggml_tensor * k = dst->src[1]; ++ const ggml_tensor * v = dst->src[2]; ++ const ggml_tensor * m = dst->src[3]; ++ const ggml_tensor * rel = dst->src[5]; ++ ++ // route F16/BF16 K/V to the MMA kernel; keep this FP32 kernel for mixed types and strided rel ++ if (k->type != GGML_TYPE_F32 && v->type != GGML_TYPE_F32 && ++ rel->type == GGML_TYPE_F32 && ggml_is_contiguous(rel) && ++ // MMA ABI indexes rel by Q's batch: a singleton rel batch must take the stride-aware fallback ++ rel->ne[3] == q->ne[3] && rel->ne[0] <= (1 << 20)) { ++ ggml_cuda_flash_attn_ext(ctx, dst); ++ return; ++ } ++ ++ float scale; ++ memcpy(&scale, dst->op_params, sizeof(scale)); ++ // the tensor extent (not op_params) is authoritative after graph cloning ++ const int64_t rel_extent = rel->ne[0]; ++ ++ constexpr int warps_per_block = 4; ++ const dim3 blocks((q->ne[1] + warps_per_block - 1) / warps_per_block, q->ne[2], q->ne[3]); ++ const dim3 threads(warps_per_block * WARP_SIZE, 1, 1); ++ cudaStream_t stream = ctx.stream(); ++ ++#define LAUNCH_BANDED(D) \ ++ flash_attn_ext_banded_f32<<>>( \ ++ (const char *) q->data, (const char *) k->data, (const char *) v->data, \ ++ m ? (const char *) m->data : nullptr, (const char *) rel->data, (float *) dst->data, \ ++ scale, k->type, v->type, rel->type, q->ne[1], k->ne[1], q->ne[2], k->ne[2], q->ne[3], \ ++ rel_extent, m ? m->ne[2] : 1, m ? m->ne[3] : 1, \ ++ q->nb[1], q->nb[2], q->nb[3], k->nb[0], k->nb[1], k->nb[2], k->nb[3], \ ++ v->nb[0], v->nb[1], v->nb[2], v->nb[3], \ ++ m ? m->nb[1] : 0, m ? m->nb[2] : 0, m ? m->nb[3] : 0, \ ++ rel->nb[0], rel->nb[1], rel->nb[2], rel->nb[3], rel->ne[3]) ++ ++ if (q->ne[0] == 64) { ++ LAUNCH_BANDED(64); ++ } else { ++ LAUNCH_BANDED(128); ++ } ++#undef LAUNCH_BANDED ++} +diff --git a/ggml/src/ggml-cuda/fattn-banded.cuh b/ggml/src/ggml-cuda/fattn-banded.cuh +new file mode 100644 +index 000000000..bb56b8986 +--- /dev/null ++++ b/ggml/src/ggml-cuda/fattn-banded.cuh +@@ -0,0 +1,7 @@ ++#pragma once ++ ++#include "common.cuh" ++ ++void ggml_cuda_flash_attn_ext_banded(ggml_backend_cuda_context & ctx, ggml_tensor * dst); ++ ++bool ggml_cuda_flash_attn_ext_banded_supported(int device, const ggml_tensor * dst); +diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh +index e67cc7fdf..b07a0122d 100644 +--- a/ggml/src/ggml-cuda/fattn-common.cuh ++++ b/ggml/src/ggml-cuda/fattn-common.cuh +@@ -52,7 +52,7 @@ struct ggml_cuda_flash_attn_ext_f16_extra_data { + + static inline ggml_cuda_flash_attn_ext_f16_extra_data ggml_cuda_flash_attn_ext_get_f16_extra_data( + const ggml_tensor * dst, const bool need_f16_K, const bool need_f16_V) { +- GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT); ++ GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT || dst->op == GGML_OP_FLASH_ATTN_EXT_BANDED); + + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; +@@ -983,7 +983,8 @@ void launch_fattn( + const bool V_is_K_view = V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)); + + const ggml_tensor * mask = dst->src[3]; +- const ggml_tensor * sinks = dst->src[4]; ++ const ggml_tensor * rel = dst->op == GGML_OP_FLASH_ATTN_EXT_BANDED ? dst->src[5] : nullptr; ++ const ggml_tensor * sinks = rel ? rel : dst->src[4]; + + ggml_tensor * KQV = dst; + +@@ -1192,6 +1193,13 @@ void launch_fattn( + memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float)); + ++ // banded op reuses the MMA ABI: negative max_bias tags the branch, sinks_ptr carries rel_logits, -max_bias is E ++ if (rel) { ++ GGML_ASSERT(rel->type == GGML_TYPE_F32 && ggml_is_contiguous(rel)); ++ GGML_ASSERT(rel->ne[0] <= (1 << 20)); // exactly representable in float ++ max_bias = -float(rel->ne[0]); ++ } ++ + if (logit_softcap != 0.0f) { + scale /= logit_softcap; + } +@@ -1199,8 +1207,9 @@ void launch_fattn( + const uint32_t n_head = Q->ne[2]; + const uint32_t n_head_log2 = 1u << uint32_t(floorf(log2f(float(n_head)))); + +- const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); +- const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); ++ // m0/m1 are unused for banded bias; the negative tag would otherwise blow up the exponent ++ const float m0 = rel ? 1.0f : powf(2.0f, -(max_bias ) / n_head_log2); ++ const float m1 = rel ? 1.0f : powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + + // TODO other tensor dimensions after removal of WMMA kernel: + const uint3 ne01 = init_fastdiv_values(Q->ne[1]); +diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh +index 7f4cfd551..a29f9910b 100644 +--- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh ++++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh +@@ -535,13 +535,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + const half2 * const __restrict__ K_h2, + const half2 * const __restrict__ V_h2, + const half * const __restrict__ mask_h, ++ const float * const __restrict__ rel_f, + float2 * const __restrict__ dstk, + float2 * const __restrict__ dstk_fixup, + const float scale, + const float slope, + const float logit_softcap, ++ const int rel_extent, ++ const int head_q0, + const uint3 ne01, + const int ne02, ++ const int ne11, + const int stride_K, + const int stride_V, + const int stride_mask, +@@ -688,10 +692,41 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + #pragma unroll + for (int col = 0; col < cols_per_thread; ++col) { + KQ_max_new[col] = KQ_max[col]; ++ // The fp16 VKQ accumulator holds the still-unnormalized sum of softmax weights times V. ++ // With large-magnitude V (~1e3) it overflows after a few thousand KV positions even though ++ // all inputs are finite. When the per-thread partial row sum of weights gets large, force ++ // the running maximum up: the regular rescale below then shrinks the accumulator, row sum, ++ // and all further weights by 2^-8, which the final normalization cancels exactly. The bump ++ // reaches all threads sharing the column via the KQ_max_new warp reduction, and re-triggers ++ // only after the row sum regrows 256x, so the number of bumps is logarithmic in n_kv. ++ if (KQ_rowsum[col] > 4.0f) { ++ KQ_max_new[col] += 8.0f*0.6931f; ++ } + } + float KQ_rowsum_add[cols_per_thread] = {0.0f}; + + if constexpr (cols_per_warp == 8) { ++ if (rel_f) { ++#pragma unroll ++ for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::I) { ++ const int i0 = i00 + (threadIdx.y % np)*T_C_KQ::I; ++#pragma unroll ++ for (int l = 0; l < T_C_KQ::ne; ++l) { ++ const int i = i0 + T_C_KQ::get_i(l); ++ const int jc = (threadIdx.y / np)*T_C_KQ::J + T_C_KQ::get_j(l); ++ const int j = jc / ncols2; ++ const int c = jc % ncols2; ++ const int64_t q_idx = int64_t(jt)*ncols1 + j; ++ const int64_t kv_idx = int64_t(k_VKQ_0) + i; ++ const int64_t dist = q_idx + (int64_t(ne11) - ne01.z) - kv_idx; ++ if (q_idx < ne01.z && head_q0 + c < ne02 && kv_idx < ne11 && ++ dist >= 0 && dist < rel_extent) { ++ KQ_C[i00/(np*T_C_KQ::I)].x[l] += ++ rel_f[(q_idx*ne02 + head_q0 + c)*rel_extent + dist]; ++ } ++ } ++ } ++ } + if (ncols2 > 1 || mask_h) { + #pragma unroll + for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::I) { +@@ -754,6 +789,27 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + } + } + } else { // not Turing mma or T_B_KQ::I > 8 ++ if (rel_f) { ++#pragma unroll ++ for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::J) { ++ const int i0 = i00 + (threadIdx.y % np)*T_C_KQ::J; ++#pragma unroll ++ for (int l = 0; l < T_C_KQ::ne; ++l) { ++ const int i = i0 + T_C_KQ::get_j(l); ++ const int jc = (threadIdx.y / np)*cols_per_warp + T_C_KQ::get_i(l); ++ const int j = jc / ncols2; ++ const int c = jc % ncols2; ++ const int64_t q_idx = int64_t(jt)*ncols1 + j; ++ const int64_t kv_idx = int64_t(k_VKQ_0) + i; ++ const int64_t dist = q_idx + (int64_t(ne11) - ne01.z) - kv_idx; ++ if (q_idx < ne01.z && head_q0 + c < ne02 && kv_idx < ne11 && ++ dist >= 0 && dist < rel_extent) { ++ KQ_C[i00/(np*T_C_KQ::J)].x[l] += ++ rel_f[(q_idx*ne02 + head_q0 + c)*rel_extent + dist]; ++ } ++ } ++ } ++ } + if (ncols2 > 1 || mask_h) { + #pragma unroll + for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::J) { +@@ -1015,8 +1071,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + } + } + #else +- GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, +- scale, slope, logit_softcap, ne01, ne02, ++ GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, rel_f, dstk, dstk_fixup, ++ scale, slope, logit_softcap, rel_extent, head_q0, ne01, ne02, ne11, + stride_K, stride_V, stride_mask, + tile_Q, tile_K, tile_V, tile_mask, + Q_B, VKQ_C, KQ_max, KQ_rowsum, kb0); +@@ -1120,11 +1176,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + const half2 * const __restrict__ V_h2, + const half * const __restrict__ mask_h, + const float * const __restrict__ sinks_f, ++ const float * const __restrict__ rel_f, + float2 * const __restrict__ dstk, + float2 * const __restrict__ dstk_fixup, + const float scale, + const float slope, + const float logit_softcap, ++ const int rel_extent, ++ const int head_q0, + const uint3 ne01, + const int ne02, + const int gqa_ratio, +@@ -1278,8 +1337,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + flash_attn_ext_f16_iter + +- (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, +- ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, ++ (Q_f2, K_h2, V_h2, mask_h, rel_f, dstk, dstk_fixup, scale, slope, logit_softcap, ++ rel_extent, head_q0, ne01, ne02, ne11, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); + } + constexpr bool last_iter = true; +@@ -1287,8 +1346,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + flash_attn_ext_f16_iter + +- (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, +- ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, ++ (Q_f2, K_h2, V_h2, mask_h, rel_f, dstk, dstk_fixup, scale, slope, logit_softcap, ++ rel_extent, head_q0, ne01, ne02, ne11, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); + } else { + constexpr bool oob_check = false; +@@ -1298,8 +1357,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + flash_attn_ext_f16_iter + +- (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, +- ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, ++ (Q_f2, K_h2, V_h2, mask_h, rel_f, dstk, dstk_fixup, scale, slope, logit_softcap, ++ rel_extent, head_q0, ne01, ne02, ne11, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); + } + constexpr bool last_iter = true; +@@ -1307,8 +1366,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + flash_attn_ext_f16_iter + +- (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, +- ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, ++ (Q_f2, K_h2, V_h2, mask_h, rel_f, dstk, dstk_fixup, scale, slope, logit_softcap, ++ rel_extent, head_q0, ne01, ne02, ne11, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); + } + +@@ -1692,8 +1751,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + } + } + #else +- GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dstk_fixup, +- scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ++ GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, sinks_f, rel_f, dstk, dstk_fixup, ++ scale, slope, logit_softcap, rel_extent, head_q0, ne01, ne02, gqa_ratio, ne11, + stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, + jt, kb0_start, kb0_stop); + NO_DEVICE_CODE; +@@ -1734,6 +1793,8 @@ static __global__ void flash_attn_ext_f16( + const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr; + float * GGML_CUDA_RESTRICT dst = dst_ptr; + float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr; ++ const bool banded_bias = max_bias < 0.0f; ++ const int rel_extent = banded_bias ? int(-max_bias) : 0; + + // Skip unused kernel variants for faster compilation: + if (use_logit_softcap && !(DKQ == 128 || DKQ == 256 || DKQ == 512)) { +@@ -1819,9 +1880,12 @@ static __global__ void flash_attn_ext_f16( + float2 * dstk = ((float2 *) dst) + (sequence*ne01.z*ne02 + zt_Q) * (DV/2); + + const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); +- const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; ++ const float * sinks_f = sinks && !banded_bias ? (const float *) sinks + zt_Q : nullptr; ++ const float * rel_f = banded_bias ? ++ (const float *) sinks + int64_t(sequence)*ne01.z*ne02*rel_extent : nullptr; + +- const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; ++ const float slope = banded_bias ? 1.0f : ++ (ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f); + + if (KV_max) { + kb0_stop = min(kb0_stop, KV_max[sequence*iter_j + jt] / nbatch_fa); +@@ -1830,13 +1894,13 @@ static __global__ void flash_attn_ext_f16( + if (kb0_start == 0) { + constexpr bool needs_fixup = false; // CUDA block is working on an entire tile. + flash_attn_ext_f16_process_tile +- (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, +- ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); ++ (Q_f2, K_h2, V_h2, mask_h, sinks_f, rel_f, dstk, dst_meta, scale, slope, logit_softcap, ++ rel_extent, zt_Q, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); + } else { + constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. + flash_attn_ext_f16_process_tile +- (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, +- ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); ++ (Q_f2, K_h2, V_h2, mask_h, sinks_f, rel_f, dstk, dst_meta, scale, slope, logit_softcap, ++ rel_extent, zt_Q, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); + } + + kbc += iter_k; +@@ -1865,9 +1929,12 @@ static __global__ void flash_attn_ext_f16( + float2 * dstk = ((float2 *) dst) + (sequence*ne01.z*ne02 + zt_Q) * (DV/2); + + const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); +- const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; ++ const float * sinks_f = sinks && !banded_bias ? (const float *) sinks + zt_Q : nullptr; ++ const float * rel_f = banded_bias ? ++ (const float *) sinks + int64_t(sequence)*ne01.z*ne02*rel_extent : nullptr; + +- const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; ++ const float slope = banded_bias ? 1.0f : ++ (ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f); + + if (KV_max) { + kb0_stop = min(kb0_stop, KV_max[sequence*iter_j + jt] / nbatch_fa); +@@ -1876,8 +1943,8 @@ static __global__ void flash_attn_ext_f16( + constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. + constexpr bool needs_fixup = false; + flash_attn_ext_f16_process_tile +- (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, +- ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); ++ (Q_f2, K_h2, V_h2, mask_h, sinks_f, rel_f, dstk, dst_meta, scale, slope, logit_softcap, ++ rel_extent, zt_Q, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); + #else + GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale, + max_bias, m0, m1, n_head_log2, logit_softcap, +diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu +index ab7a3b297..43849b9ce 100644 +--- a/ggml/src/ggml-cuda/fattn.cu ++++ b/ggml/src/ggml-cuda/fattn.cu +@@ -390,6 +390,14 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const + + const int cc = ggml_cuda_info().devices[device].cc; + ++ // banded bias lives in the F16 MMA loop; F32 K/V keeps the dedicated FP32 warp kernel ++ if (dst->op == GGML_OP_FLASH_ATTN_EXT_BANDED && ++ dst->src[5]->type == GGML_TYPE_F32 && K->type != GGML_TYPE_F32 && V->type != GGML_TYPE_F32 && ++ dst->src[5]->ne[3] == Q->ne[3] && ++ turing_mma_available(cc) && (Q->ne[0] == 64 || Q->ne[0] == 128) && V->ne[0] == Q->ne[0]) { ++ return BEST_FATTN_KERNEL_MMA_F16; ++ } ++ + switch (K->ne[0]) { + case 40: + case 64: +@@ -534,7 +542,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const + } + + size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * dst) { +- GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT); ++ GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT || dst->op == GGML_OP_FLASH_ATTN_EXT_BANDED); + + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; +diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu +index e73a7b890..5373b64d9 100644 +--- a/ggml/src/ggml-cuda/ggml-cuda.cu ++++ b/ggml/src/ggml-cuda/ggml-cuda.cu +@@ -25,6 +25,7 @@ + #include "ggml-cuda/diagmask.cuh" + #include "ggml-cuda/diag.cuh" + #include "ggml-cuda/fattn.cuh" ++#include "ggml-cuda/fattn-banded.cuh" + #include "ggml-cuda/fwht.cuh" + #include "ggml-cuda/getrows.cuh" + #include "ggml-cuda/im2col.cuh" +@@ -906,7 +907,7 @@ static size_t ggml_backend_cuda_buffer_type_get_alignment(ggml_backend_buffer_ty + static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *) buft->context; + +- size_t size = tensor->op == GGML_OP_FLASH_ATTN_EXT ++ size_t size = (tensor->op == GGML_OP_FLASH_ATTN_EXT || tensor->op == GGML_OP_FLASH_ATTN_EXT_BANDED) + ? ggml_cuda_flash_attn_ext_get_alloc_size(buft_ctx->device, tensor) + : ggml_nbytes(tensor); + int64_t ne0 = tensor->ne[0]; +@@ -1496,6 +1497,9 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const + size_t nbd2 = dst->nb[2]; + size_t nbd3 = dst->nb[3]; + ++ const bool f32_pedantic = compute_type == GGML_TYPE_F32 && ++ src0->type == GGML_TYPE_F32 && ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC; ++ + cublasComputeType_t cu_compute_type = traits::compute_type; + cudaDataType_t cu_data_type = traits::data_type; + cudaDataType_t cu_data_type_a = traits::data_type; +@@ -1527,6 +1531,18 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const + } + } + ++ if (f32_pedantic) { ++#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11020 ++ cu_compute_type = CUBLAS_COMPUTE_32F_PEDANTIC; ++#else ++ // no pedantic compute enum here; ordinary F32 is the strongest available contract ++ cu_compute_type = CUBLAS_COMPUTE_32F; ++#endif ++ } ++ ++ const auto cu_gemm_algo = f32_pedantic ? ++ CUBLAS_GEMM_DEFAULT : CUBLAS_GEMM_DEFAULT_TENSOR_OP; ++ + GGML_ASSERT(ne12 % ne02 == 0); + GGML_ASSERT(ne13 % ne03 == 0); + +@@ -1538,12 +1554,23 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const + // However, for some old NVIDIA and AMD GPUs the strided/Ex GEMM is much slower, + // probably because the internal kernel selection logic is suboptimal. + if (compute_type == GGML_TYPE_F32 && ne12 == 1 && ne13 == 1) { +- CUBLAS_CHECK( +- cublasSgemm(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, +- ne01, ne11, ne10, +- (const float *) alpha, (const float *) src0_ptr, s01, +- (const float *) src1_ptr, s11, +- (const float *) beta, (float *) dst_ptr, ne0)); ++ if (f32_pedantic) { ++ CUBLAS_CHECK( ++ cublasGemmEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, ++ ne01, ne11, ne10, ++ alpha, src0_ptr, CUDA_R_32F, s01, ++ src1_ptr, CUDA_R_32F, s11, ++ beta, dst_ptr, CUDA_R_32F, ne0, ++ cu_compute_type, ++ cu_gemm_algo)); ++ } else { ++ CUBLAS_CHECK( ++ cublasSgemm(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, ++ ne01, ne11, ne10, ++ (const float *) alpha, (const float *) src0_ptr, s01, ++ (const float *) src1_ptr, s11, ++ (const float *) beta, (float *) dst_ptr, ne0)); ++ } + } else if (ne12 == 1 && ne13 == 1) { + CUBLAS_CHECK( + cublasGemmEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, +@@ -1552,7 +1579,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const + src1_ptr, cu_data_type_b, s11, + beta, dst_ptr, cu_data_type, ne0, + cu_compute_type, +- CUBLAS_GEMM_DEFAULT_TENSOR_OP)); ++ cu_gemm_algo)); + } else if (r2 == 1 && r3 == 1 && is_src0_cont_2 && is_src1_cont_2) { + // with a [0, 2, 1, 3] perm. and ne02==1 the matrix strides need to be determined from dim 3: + const int64_t sma = ne02 == 1 ? s03 : s02; +@@ -1568,7 +1595,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const + beta, dst_ptr, cu_data_type, ne0, ne1*ne0, // strideC + ne12*ne13, + cu_compute_type, +- CUBLAS_GEMM_DEFAULT_TENSOR_OP)); ++ cu_gemm_algo)); + } else { + // use cublasGemmBatchedEx + const int64_t ne23 = ne12*ne13; +@@ -1606,7 +1633,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const + beta, ( void **) (ptrs_dst.get() + 0*ne23), cu_data_type, ne0, + ne23, + cu_compute_type, +- CUBLAS_GEMM_DEFAULT_TENSOR_OP)); ++ cu_gemm_algo)); + } + + // Convert output back to F32 if needed +@@ -1644,6 +1671,11 @@ static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml + } + } + ++ // a scoped pedantic request overrides the process-wide compute type, for F32 weights only ++ if (src0->type == GGML_TYPE_F32 && ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC) { ++ compute_type = GGML_TYPE_F32; ++ } ++ + switch (compute_type) { + case GGML_TYPE_F32: + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst); +@@ -1829,6 +1861,8 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor + + const int cc = ggml_cuda_info().devices[ctx.device].cc; + const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; ++ const bool f32_pedantic = src0->type == GGML_TYPE_F32 && ++ ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC; + + if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11)) { + // The custom F16 vector kernel can be used over batched cuBLAS GEMM. +@@ -1836,7 +1870,8 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + return; + } +- if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { ++ if (!f32_pedantic && ggml_cuda_should_use_mmf( ++ src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { + ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); + return; + } +@@ -1862,6 +1897,8 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * + GGML_TENSOR_BINARY_OP_LOCALS + + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; ++ const bool f32_pedantic = src0->type == GGML_TYPE_F32 && ++ ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC; + + // [TAG_MUL_MAT_ID_CUDA_GRAPHS] + if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { +@@ -1886,7 +1923,8 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * + return; + } + +- if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { ++ if (!f32_pedantic && ggml_cuda_should_use_mmf( ++ src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst); + return; + } +@@ -1994,6 +2032,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * + dst_slice.nb[2] = dst_slice.ne[1] * dst_slice.nb[1]; + dst_slice.nb[3] = dst_slice.ne[2] * dst_slice.nb[2]; + dst_slice.data = dst_data_cur; ++ dst_slice.op_params[0] = dst->op_params[0]; + + ggml_cuda_mul_mat(ctx, &src0_slice, &src1_slice, &dst_slice); + CUDA_CHECK(cudaGetLastError()); +@@ -2305,6 +2344,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg + case GGML_OP_FLASH_ATTN_EXT: + ggml_cuda_flash_attn_ext(ctx, dst); + break; ++ case GGML_OP_FLASH_ATTN_EXT_BANDED: ++ ggml_cuda_flash_attn_ext_banded(ctx, dst); ++ break; + case GGML_OP_CROSS_ENTROPY_LOSS: + ggml_cuda_cross_entropy_loss(ctx, dst); + break; +@@ -5142,6 +5184,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g + op->type == GGML_TYPE_F32; + case GGML_OP_FLASH_ATTN_EXT: + return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); ++ case GGML_OP_FLASH_ATTN_EXT_BANDED: ++ return ggml_cuda_flash_attn_ext_banded_supported(dev_ctx->device, op); + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + case GGML_OP_OPT_STEP_ADAMW: +diff --git a/ggml/src/ggml-cuda/mmf.cuh b/ggml/src/ggml-cuda/mmf.cuh +index d55cc1ec7..8b2be232e 100644 +--- a/ggml/src/ggml-cuda/mmf.cuh ++++ b/ggml/src/ggml-cuda/mmf.cuh +@@ -110,9 +110,9 @@ static __global__ void mul_mat_f( + const int sample_x = sample_dst / sample_ratio; + const int sample_y = sample_dst; + +- x += int64_t(sample_x) *stride_sample_x + channel_x *stride_channel_x + row0*stride_row ; +- y += int64_t(sample_y) *stride_sample_y + (has_ids ? 0 : channel_y *stride_channel_y); +- dst += int64_t(sample_dst)*stride_sample_dst + (has_ids ? 0 : channel_dst*stride_channel_dst); ++ x += int64_t(sample_x) *stride_sample_x + int64_t(channel_x)*stride_channel_x + int64_t(row0)*stride_row; ++ y += int64_t(sample_y) *stride_sample_y + (has_ids ? 0 : int64_t(channel_y) *stride_channel_y); ++ dst += int64_t(sample_dst)*stride_sample_dst + (has_ids ? 0 : int64_t(channel_dst)*stride_channel_dst); + + if constexpr (has_ids) { + constexpr int y_stride_scale = std::is_same_v ? 1 : 2; +@@ -362,7 +362,7 @@ static __global__ void mul_mat_f_ids( + const int sample_x = sample_dst / sample_ratio; + const int sample_y = sample_dst; + +- x += int64_t(sample_x) *stride_sample_x + channel_x *stride_channel_x + row0*stride_row; ++ x += int64_t(sample_x) *stride_sample_x + int64_t(channel_x)*stride_channel_x + int64_t(row0)*stride_row; + y += int64_t(sample_y) *stride_sample_y; + dst += int64_t(sample_dst)*stride_sample_dst; + +diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh +index 71e3b2647..e8f2c0605 100644 +--- a/ggml/src/ggml-cuda/mmq.cuh ++++ b/ggml/src/ggml-cuda/mmq.cuh +@@ -438,12 +438,12 @@ template static __device__ __forceinline_ + + if constexpr (type == GGML_TYPE_NVFP4) { + if (y_scale_used) { +- dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; ++ dst[(int64_t) ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; + } else { +- dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; ++ dst[(int64_t) ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; + } + } else { +- dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; ++ dst[(int64_t) ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; + GGML_UNUSED(y_scale_used); + } + } +@@ -969,7 +969,7 @@ static __global__ void mul_mat_q( + int col_high = ncols_dst; + int col_diff = ncols_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; +- int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; ++ int64_t offset_dst = (int64_t) wt*stride_sample_dst + (int64_t) zt*stride_channel_dst + (int64_t) jt*J*stride_col_dst; + int offset_y_scale; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; +@@ -1057,7 +1057,7 @@ static __global__ void mul_mat_q( + int col_high = ncols_dst; + int col_diff = ncols_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; +- int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; ++ int64_t offset_dst = (int64_t) wt*stride_sample_dst + (int64_t) zt*stride_channel_dst + (int64_t) jt*J*stride_col_dst; + int offset_y_scale; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; +@@ -1146,7 +1146,7 @@ static __global__ void mul_mat_q( + int col_high = ncols_dst; + int col_diff = ncols_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; +- int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; ++ int64_t offset_dst = (int64_t) wt*stride_sample_dst + (int64_t) zt*stride_channel_dst + (int64_t) jt*J*stride_col_dst; + int offset_y_scale; + if constexpr (type == GGML_TYPE_NVFP4) { + offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; +@@ -1289,7 +1289,7 @@ static __global__ void mul_mat_q_stream_k_fixup( + const int it = tmp2.x; + + if (!ids_dst) { +- const int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst + it*I; ++ const int64_t offset_dst = (int64_t) wt*stride_sample_dst + (int64_t) zt*stride_channel_dst + (int64_t) jt*J*stride_col_dst + (int64_t) it*I; + dst += offset_dst; + + const int i_max = nrows_x - it*I - 1; +diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu +index d7dbc8b99..bbf862b3f 100644 +--- a/ggml/src/ggml-cuda/mmvf.cu ++++ b/ggml/src/ggml-cuda/mmvf.cu +@@ -44,7 +44,7 @@ static __global__ void mul_mat_vec_f( + + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +- x += int64_t(sample_x) *stride_sample_x + channel_x *stride_channel_x + row*stride_row; ++ x += int64_t(sample_x) *stride_sample_x + int64_t(channel_x)*stride_channel_x + int64_t(row)*stride_row; + y += int64_t(sample_y) *stride_sample_y + channel_y *stride_channel_y; + dst += int64_t(sample_dst)*stride_sample_dst + channel_dst*stride_channel_dst; + if constexpr (is_multi_token_id) { +@@ -81,7 +81,7 @@ static __global__ void mul_mat_vec_f( + } + + if (use_gate) { +- gate_x += int64_t(sample_x) *stride_sample_x + channel_x *stride_channel_x + row*stride_row; ++ gate_x += int64_t(sample_x) *stride_sample_x + int64_t(channel_x)*stride_channel_x + int64_t(row)*stride_row; + } + + if constexpr (has_fusion) { +diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu +index e18ada537..f62c4ebb2 100644 +--- a/ggml/src/ggml-cuda/mmvq.cu ++++ b/ggml/src/ggml-cuda/mmvq.cu +@@ -587,7 +587,7 @@ static __global__ void mul_mat_vec_q( + float tmp_gate[ncols_dst][rows_per_cuda_block] = {{0.0f}}; + + const block_q8_1 * y = ((const block_q8_1 *) vy) + sample_y*stride_sample_y + channel_y*stride_channel_y; +- const int kbx_offset = sample_x*stride_sample_x + channel_x*stride_channel_x + row0*stride_row_x; ++ const int64_t kbx_offset = int64_t(sample_x)*stride_sample_x + int64_t(channel_x)*stride_channel_x + int64_t(row0)*stride_row_x; + + for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); // y block index that aligns with kbx +@@ -739,7 +739,7 @@ static __global__ void mul_mat_vec_q_moe( + const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y); + + const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y; +- const int kbx_offset = channel_x*stride_channel_x + row0*stride_row_x; ++ const int64_t kbx_offset = int64_t(channel_x)*stride_channel_x + int64_t(row0)*stride_row_x; + + // partial sum for each thread + float tmp[c_rows_per_block] = {0.0f}; +diff --git a/ggml/src/ggml-cuda/pad.cu b/ggml/src/ggml-cuda/pad.cu +index 31cd00f77..013bd8386 100644 +--- a/ggml/src/ggml-cuda/pad.cu ++++ b/ggml/src/ggml-cuda/pad.cu +@@ -25,7 +25,7 @@ static __global__ void pad_f32(const float * src, size_t s00, size_t s01, size_t + return; + } + +- const int64_t dst_idx = i3 * (ne0 * ne1 * ne2) + i2 * (ne0 * ne1) + i1 * ne0 + i0; ++ const int64_t dst_idx = (int64_t) i3 * ne0 * ne1 * ne2 + (int64_t) i2 * ne0 * ne1 + (int64_t) i1 * ne0 + i0; + + if (!circular) { + if ((i0 >= lp0 && i0 < ne0 - rp0) && (i1 >= lp1 && i1 < ne1 - rp1) && (i2 >= lp2 && i2 < ne2 - rp2) && +diff --git a/ggml/src/ggml-cuda/ssm-conv.cu b/ggml/src/ggml-cuda/ssm-conv.cu +index 1463169cf..1e09177a4 100644 +--- a/ggml/src/ggml-cuda/ssm-conv.cu ++++ b/ggml/src/ggml-cuda/ssm-conv.cu +@@ -5,8 +5,8 @@ + template + static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_ptr, + const float * bias_ptr, +- const int src0_nb0, const int src0_nb1, const int src0_nb2, const int src1_nb1, +- float * dst_ptr, const int dst_nb0, const int dst_nb1, const int dst_nb2, ++ const int64_t src0_nb0, const int64_t src0_nb1, const int64_t src0_nb2, const int64_t src1_nb1, ++ float * dst_ptr, const int64_t dst_nb0, const int64_t dst_nb1, const int64_t dst_nb2, + const int64_t n_t) { + ggml_cuda_pdl_lc(); + const float * GGML_CUDA_RESTRICT src0 = src0_ptr; +@@ -22,9 +22,9 @@ static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_p + const float * w_block = (const float *) ((const char *) src1 + bidy * split_d_inner * src1_nb1); + float * y_block = (float *) ((char *) dst + bidx * dst_nb2 + bidy * split_d_inner * dst_nb0); + +- const int stride_x = src0_nb1 / sizeof(float); +- const int stride_w = src1_nb1 / sizeof(float); +- const int stride_y = dst_nb1 / sizeof(float); ++ const int64_t stride_x = src0_nb1 / sizeof(float); ++ const int64_t stride_w = src1_nb1 / sizeof(float); ++ const int64_t stride_y = dst_nb1 / sizeof(float); + + float x[d_conv] = { 0.0f }; + float w[d_conv] = { 0.0f }; +@@ -60,9 +60,9 @@ static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_p + template + static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, const float * __restrict__ src1, + const float * __restrict__ bias, +- const int src0_nb0, const int src0_nb1, const int src0_nb2, +- const int src1_nb1, float * __restrict__ dst, const int dst_nb0, +- const int dst_nb1, const int dst_nb2, const int64_t n_t) { ++ const int64_t src0_nb0, const int64_t src0_nb1, const int64_t src0_nb2, ++ const int64_t src1_nb1, float * __restrict__ dst, const int64_t dst_nb0, ++ const int64_t dst_nb1, const int64_t dst_nb2, const int64_t n_t) { + const int tid = threadIdx.x; + const int bidx = blockIdx.x; + const int bidy = blockIdx.y; +@@ -74,9 +74,9 @@ static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, + float * y_block = + (float *) ((char *) dst + bidx * dst_nb2 + bidz * split_n_t * dst_nb1 + bidy * split_d_inner * dst_nb0); + +- const int stride_x = src0_nb1 / sizeof(float); +- const int stride_w = src1_nb1 / sizeof(float); +- const int stride_y = dst_nb1 / sizeof(float); ++ const int64_t stride_x = src0_nb1 / sizeof(float); ++ const int64_t stride_w = src1_nb1 / sizeof(float); ++ const int64_t stride_y = dst_nb1 / sizeof(float); + + const int64_t local_n_t = min(split_n_t, n_t - bidz * split_n_t); + const int n_cols = d_conv - 1 + split_n_t; +@@ -124,9 +124,9 @@ static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, + } + + template +-static void ssm_conv_f32_cuda(const float * src0, const float * src1, const float * bias, const int src0_nb0, const int src0_nb1, +- const int src0_nb2, const int src1_nb1, float * dst, const int dst_nb0, const int dst_nb1, +- const int dst_nb2, const int64_t nc, const int64_t nr, const int64_t n_t, ++static void ssm_conv_f32_cuda(const float * src0, const float * src1, const float * bias, const int64_t src0_nb0, const int64_t src0_nb1, ++ const int64_t src0_nb2, const int64_t src1_nb1, float * dst, const int64_t dst_nb0, const int64_t dst_nb1, ++ const int64_t dst_nb2, const int64_t nc, const int64_t nr, const int64_t n_t, + const int64_t n_s, cudaStream_t stream) { + const int threads = 128; + GGML_ASSERT(nr % threads == 0); +diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp +index d38057721..1166178b7 100644 +--- a/ggml/src/ggml-rpc/ggml-rpc.cpp ++++ b/ggml/src/ggml-rpc/ggml-rpc.cpp +@@ -600,6 +600,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty + // ops that require additional memory for fleeting data on certain backends + // ref: https://github.com/ggml-org/llama.cpp/pull/15966 + rpc_get |= tensor->op == GGML_OP_FLASH_ATTN_EXT; ++ rpc_get |= tensor->op == GGML_OP_FLASH_ATTN_EXT_BANDED; + rpc_get |= tensor->op == GGML_OP_MUL_MAT_ID; + + if (rpc_get) { +diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c +index 4b452037b..0b720fff9 100644 +--- a/ggml/src/ggml.c ++++ b/ggml/src/ggml.c +@@ -1067,6 +1067,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + "FILL", + + "FLASH_ATTN_EXT", ++ "FLASH_ATTN_EXT_BANDED", + "FLASH_ATTN_BACK", + "SSM_CONV", + "SSM_SCAN", +@@ -1107,7 +1108,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + "MOE_MUL_MAT_ID", + }; + +-static_assert(GGML_OP_COUNT == 107, "GGML_OP_COUNT != 107"); ++static_assert(GGML_OP_COUNT == 108, "GGML_OP_COUNT != 108"); + + static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "none", +@@ -1189,6 +1190,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "fill(x, c)", + + "flash_attn_ext(x)", ++ "flash_attn_ext_banded(x)", + "flash_attn_back(x)", + "ssm_conv(x)", + "ssm_scan(x)", +@@ -1229,7 +1231,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "moe_mul_mat_id(experts, input, ids, weights)", + }; + +-static_assert(GGML_OP_COUNT == 107, "GGML_OP_COUNT != 107"); ++static_assert(GGML_OP_COUNT == 108, "GGML_OP_COUNT != 108"); + + static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); + +@@ -5480,11 +5482,59 @@ struct ggml_tensor * ggml_flash_attn_ext( + return result; + } + ++struct ggml_tensor * ggml_flash_attn_ext_banded( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * k, ++ struct ggml_tensor * v, ++ struct ggml_tensor * mask, ++ struct ggml_tensor * rel_logits, ++ float scale, ++ int64_t rel_extent) { ++ GGML_ASSERT(ggml_can_mul_mat(k, q)); ++ GGML_ASSERT(q->type == GGML_TYPE_F32); ++ GGML_ASSERT(q->ne[3] == k->ne[3]); ++ GGML_ASSERT(q->ne[3] == v->ne[3]); ++ GGML_ASSERT(q->ne[2] % k->ne[2] == 0); ++ GGML_ASSERT(q->ne[2] % v->ne[2] == 0); ++ ++ GGML_ASSERT(rel_logits != NULL); ++ GGML_ASSERT(rel_logits->type == GGML_TYPE_F32 || ++ rel_logits->type == GGML_TYPE_F16 || ++ rel_logits->type == GGML_TYPE_BF16); ++ GGML_ASSERT(rel_extent > 0); ++ GGML_ASSERT(rel_logits->ne[0] == rel_extent); ++ GGML_ASSERT(rel_logits->ne[1] == q->ne[2]); ++ GGML_ASSERT(rel_logits->ne[2] == q->ne[1]); ++ GGML_ASSERT(rel_logits->ne[3] == 1 || rel_logits->ne[3] == q->ne[3]); ++ ++ if (mask) { ++ GGML_ASSERT(mask->type == GGML_TYPE_F16); ++ GGML_ASSERT(ggml_is_contiguous(mask)); ++ GGML_ASSERT(q->ne[2] % mask->ne[2] == 0); ++ GGML_ASSERT(q->ne[3] % mask->ne[3] == 0); ++ } ++ ++ int64_t ne[4] = { v->ne[0], q->ne[2], q->ne[1], q->ne[3] }; ++ struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); ++ ++ float params[] = { scale, 0.0f, 0.0f }; ++ ggml_set_op_params(result, params, sizeof(params)); ++ ++ result->op = GGML_OP_FLASH_ATTN_EXT_BANDED; ++ result->src[0] = q; ++ result->src[1] = k; ++ result->src[2] = v; ++ result->src[3] = mask; ++ result->src[5] = rel_logits; ++ ++ return result; ++} + + void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec) { +- GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); ++ GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT || a->op == GGML_OP_FLASH_ATTN_EXT_BANDED); + + const int32_t prec_i32 = (int32_t) prec; + +@@ -5493,7 +5543,7 @@ void ggml_flash_attn_ext_set_prec( + + enum ggml_prec ggml_flash_attn_ext_get_prec( + const struct ggml_tensor * a) { +- GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); ++ GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT || a->op == GGML_OP_FLASH_ATTN_EXT_BANDED); + + const int32_t prec_i32 = ggml_get_op_params_i32(a, 3); + +diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py +index 124ea28b0..643063e66 100644 +--- a/gguf-py/gguf/constants.py ++++ b/gguf-py/gguf/constants.py +@@ -556,6 +556,7 @@ class MODEL_ARCH(IntEnum): + TALKIE = auto() + MELLUM = auto() + NANBEIGE = auto() ++ INKLING = auto() + + + class VISION_PROJECTOR_TYPE(IntEnum): +@@ -779,6 +780,14 @@ class MODEL_TENSOR(IntEnum): + SHORTCONV_CONV = auto() + SHORTCONV_INPROJ = auto() + SHORTCONV_OUTPROJ = auto() ++ # inkling ++ ATTN_R = auto() ++ ATTN_REL_PROJ = auto() ++ SHORTCONV_K = auto() ++ SHORTCONV_V = auto() ++ SHORTCONV_ATTN = auto() ++ SHORTCONV_MLP = auto() ++ FFN_GSCALE = auto() + VISEXP_ATTN_QKV = auto() + VISEXP_ATTN_OUT = auto() + VISEXP_GATE = auto() +@@ -1168,6 +1177,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { + MODEL_ARCH.TALKIE: "talkie", + MODEL_ARCH.MELLUM: "mellum", + MODEL_ARCH.NANBEIGE: "nanbeige", ++ MODEL_ARCH.INKLING: "inkling", + } + + VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { +@@ -1389,6 +1399,13 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { + MODEL_TENSOR.SHORTCONV_CONV: "blk.{bid}.shortconv.conv", + MODEL_TENSOR.SHORTCONV_INPROJ: "blk.{bid}.shortconv.in_proj", + MODEL_TENSOR.SHORTCONV_OUTPROJ: "blk.{bid}.shortconv.out_proj", ++ MODEL_TENSOR.ATTN_R: "blk.{bid}.attn_r", # inkling ++ MODEL_TENSOR.ATTN_REL_PROJ: "blk.{bid}.attn_rel_proj", # inkling ++ MODEL_TENSOR.SHORTCONV_K: "blk.{bid}.shortconv_k", # inkling ++ MODEL_TENSOR.SHORTCONV_V: "blk.{bid}.shortconv_v", # inkling ++ MODEL_TENSOR.SHORTCONV_ATTN: "blk.{bid}.shortconv_attn", # inkling ++ MODEL_TENSOR.SHORTCONV_MLP: "blk.{bid}.shortconv_mlp", # inkling ++ MODEL_TENSOR.FFN_GSCALE: "blk.{bid}.ffn_gscale", # inkling + MODEL_TENSOR.VISEXP_ATTN_QKV: "blk.{bid}.vis_attn_qkv", + MODEL_TENSOR.VISEXP_ATTN_OUT: "blk.{bid}.vis_attn_output", + MODEL_TENSOR.VISEXP_GATE: "blk.{bid}.vis_gate", +@@ -4161,6 +4178,38 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + ], ++ MODEL_ARCH.INKLING: [ ++ MODEL_TENSOR.TOKEN_EMBD, ++ MODEL_TENSOR.TOKEN_EMBD_NORM, ++ MODEL_TENSOR.OUTPUT_NORM, ++ MODEL_TENSOR.OUTPUT, ++ MODEL_TENSOR.ATTN_NORM, ++ MODEL_TENSOR.ATTN_Q, ++ MODEL_TENSOR.ATTN_K, ++ MODEL_TENSOR.ATTN_V, ++ MODEL_TENSOR.ATTN_R, ++ MODEL_TENSOR.ATTN_OUT, ++ MODEL_TENSOR.ATTN_Q_NORM, ++ MODEL_TENSOR.ATTN_K_NORM, ++ MODEL_TENSOR.ATTN_REL_PROJ, ++ MODEL_TENSOR.SHORTCONV_K, ++ MODEL_TENSOR.SHORTCONV_V, ++ MODEL_TENSOR.SHORTCONV_ATTN, ++ MODEL_TENSOR.SHORTCONV_MLP, ++ MODEL_TENSOR.FFN_NORM, ++ MODEL_TENSOR.FFN_GATE, ++ MODEL_TENSOR.FFN_UP, ++ MODEL_TENSOR.FFN_DOWN, ++ MODEL_TENSOR.FFN_GSCALE, ++ MODEL_TENSOR.FFN_GATE_INP, ++ MODEL_TENSOR.FFN_EXP_PROBS_B, ++ MODEL_TENSOR.FFN_GATE_EXP, ++ MODEL_TENSOR.FFN_UP_EXP, ++ MODEL_TENSOR.FFN_DOWN_EXP, ++ MODEL_TENSOR.FFN_GATE_SHEXP, ++ MODEL_TENSOR.FFN_UP_SHEXP, ++ MODEL_TENSOR.FFN_DOWN_SHEXP, ++ ], + MODEL_ARCH.SMALLTHINKER: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, +@@ -4834,6 +4883,7 @@ class GGUFValueType(IntEnum): + + + class VisionProjectorType: ++ INKLING = "inkling" + GEMMA3 = "gemma3" + GEMMA3NV = "gemma3nv" + GEMMA3NA = "gemma3na" +diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py +index 1e991b873..5a0175a56 100644 +--- a/gguf-py/gguf/tensor_mapping.py ++++ b/gguf-py/gguf/tensor_mapping.py +@@ -63,6 +63,7 @@ class TensorNameMap: + "model.layers.0.pre_norm", # rwkv7 + "backbone.norm", # wavtokenizer + "model.embedding_norm", # lfm2 ++ "model.embed_norm", # inkling + ), + + # Position embeddings +@@ -86,6 +87,7 @@ class TensorNameMap: + "lm_head", # llama4 + "model.transformer.ff_out", # llada + "head.decoder", # modern-bert ++ "model.unembed", # inkling + ), + MODEL_TENSOR.DENSE_2_OUT: ( + "dense_2_out", # embeddinggemma +@@ -2548,6 +2550,53 @@ class TensorNameMap: + + # architecture-specific block mappings + arch_block_mappings_cfg: dict[MODEL_ARCH, dict[MODEL_TENSOR, tuple[str, ...]]] = { ++ MODEL_ARCH.INKLING: { ++ MODEL_TENSOR.ATTN_NORM: ( ++ "model.layers.{bid}.attn_norm", ++ ), ++ MODEL_TENSOR.ATTN_Q: ( ++ "model.layers.{bid}.attn.wq_du", ++ ), ++ MODEL_TENSOR.ATTN_K: ( ++ "model.layers.{bid}.attn.wk_dv", ++ ), ++ MODEL_TENSOR.ATTN_V: ( ++ "model.layers.{bid}.attn.wv_dv", ++ ), ++ MODEL_TENSOR.ATTN_R: ( ++ "model.layers.{bid}.attn.wr_du", ++ ), ++ MODEL_TENSOR.ATTN_OUT: ( ++ "model.layers.{bid}.attn.wo_ud", ++ ), ++ MODEL_TENSOR.ATTN_Q_NORM: ( ++ "model.layers.{bid}.attn.q_norm", ++ ), ++ MODEL_TENSOR.ATTN_K_NORM: ( ++ "model.layers.{bid}.attn.k_norm", ++ ), ++ MODEL_TENSOR.ATTN_REL_PROJ: ( ++ "model.layers.{bid}.attn.rel_logits_proj", ++ ), ++ MODEL_TENSOR.SHORTCONV_K: ( ++ "model.layers.{bid}.attn.k_sconv", ++ ), ++ MODEL_TENSOR.SHORTCONV_V: ( ++ "model.layers.{bid}.attn.v_sconv", ++ ), ++ MODEL_TENSOR.SHORTCONV_ATTN: ( ++ "model.layers.{bid}.attn_sconv", ++ ), ++ MODEL_TENSOR.SHORTCONV_MLP: ( ++ "model.layers.{bid}.mlp_sconv", ++ ), ++ MODEL_TENSOR.FFN_NORM: ( ++ "model.layers.{bid}.mlp_norm", ++ ), ++ MODEL_TENSOR.FFN_DOWN: ( ++ "model.layers.{bid}.mlp.w2_md", ++ ), ++ }, + MODEL_ARCH.ARCTIC: { + MODEL_TENSOR.FFN_NORM: ( + "model.layers.{bid}.residual_layernorm", +diff --git a/models/templates/Inkling.jinja b/models/templates/Inkling.jinja +new file mode 100644 +index 000000000..95874d23d +--- /dev/null ++++ b/models/templates/Inkling.jinja +@@ -0,0 +1,514 @@ ++{#- Keep Python's floating type spelling when the jinja engine compacts 1.0 to 1. -#} ++{%- macro json_scalar(value) -%} ++ {%- set serialized = value | tojson(ensure_ascii=false, separators=(',', ':')) -%} ++ {%- if value is float and '.' not in serialized and 'e' not in (serialized | lower) -%} ++ {{- serialized -}}{{- '.0' -}} ++ {%- else -%} ++ {{- serialized -}} ++ {%- endif -%} ++{%- endmacro -%} ++ ++{%- macro canonical_json(value) -%} ++ {%- if value is mapping -%} ++ {{- '{' -}} ++ {%- for key, item in value | dictsort(case_sensitive=true) -%} ++ {{- (key | string) | tojson(ensure_ascii=false, separators=(',', ':')) -}}{{- ':' -}} ++ {{- canonical_json(item) -}} ++ {%- if not loop.last -%}{{- ',' -}}{%- endif -%} ++ {%- endfor -%} ++ {{- '}' -}} ++ {%- elif value is sequence and value is not string -%} ++ {{- '[' -}} ++ {%- for item in value -%} ++ {{- canonical_json(item) -}} ++ {%- if not loop.last -%}{{- ',' -}}{%- endif -%} ++ {%- endfor -%} ++ {{- ']' -}} ++ {%- else -%} ++ {{- json_scalar(value) -}} ++ {%- endif -%} ++{%- endmacro -%} ++ ++{#- Advance across insignificant JSON whitespace. -#} ++{%- macro json_skip_ws(source, state) -%} ++ {%- set scan = namespace(done=false) -%} ++ {%- for ignored in range(source | length) -%} ++ {%- if not scan.done and state.pos < source | length -%} ++ {%- set ch = source[state.pos] -%} ++ {%- if ch == ' ' or ch == '\t' or ch == '\r' or ch == '\n' -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- else -%} ++ {%- set scan.done = true -%} ++ {%- endif -%} ++ {%- endif -%} ++ {%- endfor -%} ++{%- endmacro -%} ++ ++{#- ++ Decode a JSON string under Transformers, then re-encode it like json.dumps. ++ llama.cpp converts string arguments to objects after capability detection; ++ the fallback keeps engine-only parsing safe without runtime-only filters. ++-#} ++{%- macro json_string(source, state) -%} ++ {%- if lipsum is defined -%} ++ {%- set out = namespace(decoded='', done=false, code=0, low=0) -%} ++ {%- set hex_values = {'0': 0, '1': 1, '2': 2, '3': 3, '4': 4, '5': 5, '6': 6, '7': 7, '8': 8, '9': 9, 'a': 10, 'b': 11, 'c': 12, 'd': 13, 'e': 14, 'f': 15} -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- for ignored in range(source | length) -%} ++ {%- if not out.done and state.pos < source | length -%} ++ {%- set ch = source[state.pos] -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- if ch == '"' -%} ++ {%- set out.done = true -%} ++ {%- elif ch != '\\' -%} ++ {%- set out.decoded = out.decoded + ch -%} ++ {%- elif state.pos < source | length -%} ++ {%- set escape = source[state.pos] -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- if escape == '"' -%}{%- set out.decoded = out.decoded + '"' -%} ++ {%- elif escape == '\\' -%}{%- set out.decoded = out.decoded + '\\' -%} ++ {%- elif escape == '/' -%}{%- set out.decoded = out.decoded + '/' -%} ++ {%- elif escape == 'b' -%}{%- set out.decoded = out.decoded + '\b' -%} ++ {%- elif escape == 'f' -%}{%- set out.decoded = out.decoded + '\f' -%} ++ {%- elif escape == 'n' -%}{%- set out.decoded = out.decoded + '\n' -%} ++ {%- elif escape == 'r' -%}{%- set out.decoded = out.decoded + '\r' -%} ++ {%- elif escape == 't' -%}{%- set out.decoded = out.decoded + '\t' -%} ++ {%- elif escape == 'u' -%} ++ {%- set out.code = 0 -%} ++ {%- for offset in range(4) -%} ++ {%- set digit = source[state.pos + offset] | lower -%} ++ {%- set out.code = out.code * 16 + hex_values[digit] -%} ++ {%- endfor -%} ++ {%- set state.pos = state.pos + 4 -%} ++ {%- if out.code >= 55296 and out.code <= 56319 and source[state.pos:state.pos + 2] == '\\u' -%} ++ {%- set out.low = 0 -%} ++ {%- for offset in range(4) -%} ++ {%- set digit = source[state.pos + 2 + offset] | lower -%} ++ {%- set out.low = out.low * 16 + hex_values[digit] -%} ++ {%- endfor -%} ++ {%- if out.low >= 56320 and out.low <= 57343 -%} ++ {%- set out.code = 65536 + (out.code - 55296) * 1024 + out.low - 56320 -%} ++ {%- set state.pos = state.pos + 6 -%} ++ {%- endif -%} ++ {%- endif -%} ++ {%- set out.decoded = out.decoded + ('%c' % out.code) -%} ++ {%- else -%} ++ {%- set out.decoded = out.decoded + escape -%} ++ {%- endif -%} ++ {%- endif -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {%- set state.string_value = out.decoded -%} ++ {{- out.decoded | tojson(ensure_ascii=false, separators=(',', ':')) -}} ++ {%- else -%} ++ {%- set out = namespace(start=state.pos, escaped=false, done=false) -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- for ignored in range(source | length) -%} ++ {%- if not out.done and state.pos < source | length -%} ++ {%- set ch = source[state.pos] -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- if out.escaped -%} ++ {%- set out.escaped = false -%} ++ {%- elif ch == '\\' -%} ++ {%- set out.escaped = true -%} ++ {%- elif ch == '"' -%} ++ {%- set out.done = true -%} ++ {%- endif -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {%- set state.string_value = source[out.start:state.pos] -%} ++ {{- source[out.start:state.pos] -}} ++ {%- endif -%} ++{%- endmacro -%} ++ ++{#- ++ Move the cursor across one complete JSON value without rendering it. Object ++ parsing uses the raw slice to sort members before recursively rendering them. ++-#} ++{%- macro json_scan_value(source, state) -%} ++ {{- json_skip_ws(source, state) -}} ++ {%- set scan = namespace(depth=0, quoted=false, escaped=false, done=false) -%} ++ {%- for ignored in range(source | length) -%} ++ {%- if not scan.done and state.pos < source | length -%} ++ {%- set ch = source[state.pos] -%} ++ {%- if scan.quoted -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- if ch == '"' and not scan.escaped -%} ++ {%- set scan.quoted = false -%} ++ {%- elif ch == '\\' and not scan.escaped -%} ++ {%- set scan.escaped = true -%} ++ {%- else -%} ++ {%- set scan.escaped = false -%} ++ {%- endif -%} ++ {%- elif ch == '"' -%} ++ {%- set scan.quoted = true -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- elif ch == '{' or ch == '[' -%} ++ {%- set scan.depth = scan.depth + 1 -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- elif ch == '}' or ch == ']' -%} ++ {%- if scan.depth > 0 -%} ++ {%- set scan.depth = scan.depth - 1 -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- else -%} ++ {%- set scan.done = true -%} ++ {%- endif -%} ++ {%- elif ch == ',' and scan.depth == 0 -%} ++ {%- set scan.done = true -%} ++ {%- else -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- endif -%} ++ {%- endif -%} ++ {%- endfor -%} ++{%- endmacro -%} ++ ++{#- Parse and canonically render one JSON value from source at state.pos. -#} ++{%- macro canonical_json_text_value(source, state) -%} ++ {{- json_skip_ws(source, state) -}} ++ {%- if state.pos >= source | length -%} ++ {{- '{}' -}} ++ {%- elif source[state.pos] == '"' -%} ++ {{- json_string(source, state) -}} ++ {%- elif source[state.pos] == '{' -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {{- json_skip_ws(source, state) -}} ++ {%- set object_state = namespace(pairs=[], done=false) -%} ++ {%- if state.pos < source | length and source[state.pos] == '}' -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- set object_state.done = true -%} ++ {%- endif -%} ++ {%- for ignored in range(source | length) -%} ++ {%- if not object_state.done -%} ++ {{- json_skip_ws(source, state) -}} ++ {%- set key = json_string(source, state) -%} ++ {%- set sort_key = state.string_value -%} ++ {{- json_skip_ws(source, state) -}} ++ {%- if state.pos < source | length and source[state.pos] == ':' -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- endif -%} ++ {{- json_skip_ws(source, state) -}} ++ {%- set value_start = state.pos -%} ++ {{- json_scan_value(source, state) -}} ++ {#- json.loads keeps the final member when a key is duplicated. -#} ++ {%- set unique = namespace(pairs=[]) -%} ++ {%- for previous in object_state.pairs -%} ++ {%- if previous[0] != sort_key -%} ++ {%- set unique.pairs = unique.pairs + [previous] -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {%- set object_state.pairs = unique.pairs + [[sort_key, key, source[value_start:state.pos]]] -%} ++ {{- json_skip_ws(source, state) -}} ++ {%- if state.pos < source | length and source[state.pos] == ',' -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- else -%} ++ {%- if state.pos < source | length and source[state.pos] == '}' -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- endif -%} ++ {%- set object_state.done = true -%} ++ {%- endif -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {{- '{' -}} ++ {%- for pair in object_state.pairs | sort(case_sensitive=true, attribute=0) -%} ++ {{- pair[1] -}}{{- ':' -}} ++ {%- set child_state = namespace(pos=0, string_value='') -%} ++ {{- canonical_json_text_value(pair[2], child_state) -}} ++ {%- if not loop.last -%}{{- ',' -}}{%- endif -%} ++ {%- endfor -%} ++ {{- '}' -}} ++ {%- elif source[state.pos] == '[' -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {{- '[' -}} ++ {{- json_skip_ws(source, state) -}} ++ {%- set array_state = namespace(done=false, first=true) -%} ++ {%- if state.pos < source | length and source[state.pos] == ']' -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- set array_state.done = true -%} ++ {%- endif -%} ++ {%- for ignored in range(source | length) -%} ++ {%- if not array_state.done -%} ++ {%- if not array_state.first -%}{{- ',' -}}{%- endif -%} ++ {%- set array_state.first = false -%} ++ {{- canonical_json_text_value(source, state) -}} ++ {{- json_skip_ws(source, state) -}} ++ {%- if state.pos < source | length and source[state.pos] == ',' -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- else -%} ++ {%- if state.pos < source | length and source[state.pos] == ']' -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- endif -%} ++ {%- set array_state.done = true -%} ++ {%- endif -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {{- ']' -}} ++ {%- else -%} ++ {%- set scalar = namespace(start=state.pos, done=false) -%} ++ {%- for ignored in range(source | length) -%} ++ {%- if not scalar.done and state.pos < source | length -%} ++ {%- set ch = source[state.pos] -%} ++ {%- if ch == ',' or ch == '}' or ch == ']' or ch == ' ' or ch == '\t' or ch == '\r' or ch == '\n' -%} ++ {%- set scalar.done = true -%} ++ {%- else -%} ++ {%- set state.pos = state.pos + 1 -%} ++ {%- endif -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {%- set token = source[scalar.start:state.pos] -%} ++ {%- if token == 'true' or token == 'false' or token == 'null' -%} ++ {{- token -}} ++ {%- elif '.' in token or 'e' in token or 'E' in token -%} ++ {{- json_scalar(token | float) -}} ++ {%- else -%} ++ {{- token | int | tojson(ensure_ascii=false, separators=(',', ':')) -}} ++ {%- endif -%} ++ {%- endif -%} ++{%- endmacro -%} ++ ++{%- macro canonical_json_text(source) -%} ++ {%- set text = source | trim -%} ++ {%- if not text -%} ++ {{- '{}' -}} ++ {%- else -%} ++ {%- set state = namespace(pos=0, string_value='') -%} ++ {{- canonical_json_text_value(text, state) -}} ++ {%- endif -%} ++{%- endmacro -%} ++ ++{#- ++ OpenAI clients use either an argument mapping or a JSON-encoded object. ++ Non-object/empty oddities degrade to {} instead of raising. Object parsing ++ recursively sorts keys, keeps the last duplicate key (json.loads behavior), ++ decodes JSON escapes under Transformers, and preserves array order. ++-#} ++{%- macro canonical_arguments(arguments) -%} ++ {%- if arguments is mapping -%} ++ {{- canonical_json(arguments) -}} ++ {%- elif arguments is string -%} ++ {%- set source = arguments | trim -%} ++ {%- if source and source[0] == '{' -%} ++ {{- canonical_json_text(source) -}} ++ {%- else -%} ++ {{- '{}' -}} ++ {%- endif -%} ++ {%- else -%} ++ {{- '{}' -}} ++ {%- endif -%} ++{%- endmacro -%} ++ ++{#- ++ Match Python f"{float(effort):.2f}" followed by trailing-zero removal. ++ llama.cpp's jinja engine has no round filter; the bit table handles binary64 midpoint equality. ++-#} ++{%- macro reasoning_effort_text(effort) -%} ++ {%- set eff = effort -%} ++ {%- if eff is string -%} ++ {%- set e = eff | trim | lower -%} ++ {%- if e == 'none' -%}{%- set eff = 0.0 -%} ++ {%- elif e == 'minimal' -%}{%- set eff = 0.1 -%} ++ {%- elif e == 'low' -%}{%- set eff = 0.2 -%} ++ {%- elif e == 'medium' -%}{%- set eff = 0.7 -%} ++ {%- elif e == 'high' -%}{%- set eff = 0.9 -%} ++ {%- elif e == 'xhigh' -%}{%- set eff = 0.99 -%} ++ {%- elif e == 'max' -%}{%- set eff = 0.99 -%} ++ {%- else -%}{%- set eff = e | float(-1.0) -%} ++ {%- endif -%} ++ {%- endif -%} ++ {%- set value = eff | float -%} ++ {%- if value < 0 or value > 0.99 -%} ++ {{- raise_exception('Invalid reasoning_effort: ' + (effort | string) + '; expected none/minimal/low/medium/high/xhigh/max or a number in [0.0, 0.99]') -}} ++ {%- endif -%} ++ {%- if value == value | int and value >= 0 and value <= 1 -%} ++ {{- value | int -}} ++ {%- else -%} ++ {%- set midpoint_rounds_up = '1011011011010100100100100111000111000111100011100011111100000001111110000001111110000001111111000000' -%} ++ {%- set rounded = namespace(hundredths=0) -%} ++ {%- for lower_hundredth in range(100) -%} ++ {%- set boundary = (lower_hundredth + 0.5) / 100 -%} ++ {%- if value > boundary or (value == boundary and midpoint_rounds_up[lower_hundredth] == '1') -%} ++ {%- set rounded.hundredths = lower_hundredth + 1 -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {%- if rounded.hundredths == 100 -%} ++ {{- '1' -}} ++ {%- elif rounded.hundredths == 0 -%} ++ {{- '0' -}} ++ {%- elif rounded.hundredths % 10 == 0 -%} ++ {{- '0.' -}}{{- (rounded.hundredths / 10) | int -}} ++ {%- elif rounded.hundredths < 10 -%} ++ {{- '0.0' -}}{{- rounded.hundredths -}} ++ {%- else -%} ++ {{- '0.' -}}{{- rounded.hundredths -}} ++ {%- endif -%} ++ {%- endif -%} ++{%- endmacro -%} ++ ++{%- macro role_token(role) -%} ++ {%- if role == 'user' -%}{{- '<|message_user|>' -}} ++ {%- elif role == 'assistant' -%}{{- '<|message_model|>' -}} ++ {%- elif role == 'system' or role == 'developer' -%}{{- '<|message_system|>' -}} ++ {%- elif role == 'tool' -%}{{- '<|message_tool|>' -}} ++ {%- endif -%} ++{%- endmacro -%} ++ ++{%- macro emit_message(role, kind, content='', author_name='') -%} ++ {{- role_token(role) -}} ++ {%- if author_name -%}{{- author_name -}}{%- endif -%} ++ {%- if kind == 'text' -%} ++ {{- '<|content_text|>' -}}{{- content -}} ++ {%- elif kind == 'thinking' -%} ++ {{- '<|content_thinking|>' -}}{{- content -}} ++ {%- elif kind == 'xml' -%} ++ {{- '<|content_xml|>' -}}{{- content -}} ++ {%- elif kind == 'invoke_tool_json' -%} ++ {{- '<|content_invoke_tool_json|>' -}}{{- content -}} ++ {%- elif kind == 'image' -%} ++ {{- '<|content_image|><|image|>' -}} ++ {%- elif kind == 'audio' -%} ++ {{- '<|content_audio_input|><|audio|><|audio_end|>' -}} ++ {%- endif -%} ++ {{- '<|end_message|>' -}} ++{%- endmacro -%} ++ ++{%- set effort_value = 0.0 if (enable_thinking is defined and enable_thinking is not none and not enable_thinking) else (reasoning_effort if (reasoning_effort is defined and reasoning_effort is not none) else 0.9) -%} ++{%- set eff_ns = namespace(emitted=false) -%} ++{%- set first_ns = namespace(idx=-1) -%} ++{%- for m in messages -%} ++ {%- if first_ns.idx == -1 and m.get('role') not in ['system', 'developer'] -%} ++ {%- set first_ns.idx = loop.index0 -%} ++ {%- endif -%} ++{%- endfor -%} ++ ++{%- if tools is defined and tools -%} ++ {{- '<|message_system|>tool_declare<|content_xml|>[' -}} ++ {%- for tool in tools -%} ++ {%- set function = tool.get('function', {}) if tool.get('function', {}) is mapping else {} -%} ++ {%- set description = function.get('description') or '' -%} ++ {%- set parameters = function.get('parameters') or {} -%} ++ {%- set tool_type = tool.get('type', 'function') -%} ++ {{- '{"description":' -}}{{- canonical_json(description) -}} ++ {{- ',"name":' -}}{{- canonical_json(function.get('name')) -}} ++ {{- ',"parameters":' -}}{{- canonical_json(parameters) -}} ++ {{- ',"type":' -}}{{- canonical_json(tool_type) -}}{{- '}' -}} ++ {%- if not loop.last -%}{{- ',' -}}{%- endif -%} ++ {%- endfor -%} ++ {{- ']<|end_message|>' -}} ++{%- endif -%} ++ ++{#- Last-user boundary used only by the opt-in preserve_thinking=false mode. -#} ++{%- set thinking_state = namespace(last_user_index=messages | length - 1, found_user=false) -%} ++{%- for index in range(messages | length - 1, -1, -1) -%} ++ {%- if not thinking_state.found_user and messages[index].get('role') == 'user' -%} ++ {%- set thinking_state.last_user_index = index -%} ++ {%- set thinking_state.found_user = true -%} ++ {%- endif -%} ++{%- endfor -%} ++ ++{%- for message in messages -%} ++ {%- set message_index = loop.index0 -%} ++ {%- set role = message.get('role') -%} ++ {%- if not eff_ns.emitted and loop.index0 == first_ns.idx -%} ++ {{- emit_message('system', 'text', 'Thinking effort level: ' + reasoning_effort_text(effort_value)) -}} ++ {%- set eff_ns.emitted = true -%} ++ {%- endif -%} ++ {%- if role == 'tool' -%} ++ {%- set resolved = namespace(name=message.get('name') or '') -%} ++ {%- if not resolved.name and message.get('tool_call_id') -%} ++ {%- for prior in messages[:message_index] -%} ++ {%- if prior.get('role') == 'assistant' -%} ++ {%- for call in prior.get('tool_calls') or [] -%} ++ {%- if call.get('id') and (call.get('id') | string) == message.get('tool_call_id') -%} ++ {%- set prior_function = call.get('function', {}) if call.get('function', {}) is mapping else {} -%} ++ {%- set resolved.name = prior_function.get('name') or '' -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {%- endif -%} ++ {%- set tool_content = message.get('content', '') -%} ++ {%- if tool_content is none -%} ++ {%- set tool_content = '' -%} ++ {%- elif tool_content is mapping or (tool_content is sequence and tool_content is not string) -%} ++ {%- set tool_content = canonical_json(tool_content) -%} ++ {%- elif tool_content is not string -%} ++ {%- set tool_content = canonical_json(tool_content) -%} ++ {%- endif -%} ++ {{- emit_message('tool', 'text', tool_content, resolved.name | string) -}} ++ {%- elif role == 'user' or role == 'assistant' or role == 'system' or role == 'developer' -%} ++ {%- set turn_out -%} ++ {%- if role == 'assistant' and message.get('reasoning_content') is string and message.get('reasoning_content') and ++ ((preserve_thinking is not defined) or preserve_thinking is not false or message_index > thinking_state.last_user_index) -%} ++ {{- emit_message('assistant', 'thinking', message.get('reasoning_content')) -}} ++ {%- endif -%} ++ ++ {%- set content = message.get('content', '') -%} ++ {#- Makes llama.cpp retain typed arrays rather than flattening them. -#} ++ {%- set content_probe = content[0] if content is sequence and content | length > 0 else none -%} ++ {%- if content is string -%} ++ {%- if '<__media_' in content -%} ++ {#- Flattened media markers: each part becomes its own message block. The runtime may ++ randomize the marker suffix, so split on the stable prefix and re-emit the exact ++ marker text; the runtime then expands it into the typed content sentinel plus the ++ media embedding rows. -#} ++ {%- for segment in content.split('<__media_') -%} ++ {%- if loop.first -%} ++ {%- if segment -%}{{- emit_message(role, 'text', segment) -}}{%- endif -%} ++ {%- else -%} ++ {%- set mparts = segment.split('>') -%} ++ {%- set rest = mparts[1:] | join('>') -%} ++ {{- role_token(role) -}}{{- '<__media_' + mparts[0] + '>' -}}{{- '<|end_message|>' -}} ++ {%- if rest -%}{{- emit_message(role, 'text', rest) -}}{%- endif -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {%- elif content -%} ++ {{- emit_message(role, 'text', content) -}} ++ {%- endif -%} ++ {%- elif content is sequence -%} ++ {%- for part in content -%} ++ {%- if part is string -%} ++ {{- emit_message(role, 'text', part) -}} ++ {%- elif part is mapping -%} ++ {%- set part_type = part.get('type') -%} ++ {%- if part_type is none or part_type == 'text' or part_type == 'input_text' -%} ++ {%- set part_text = part.get('text', '') -%} ++ {{- emit_message(role, 'text', part_text if part_text is string else '') -}} ++ {%- elif part_type == 'image' or part_type == 'input_image' or part_type == 'image_url' -%} ++ {{- emit_message(role, 'image') -}} ++ {%- elif part_type == 'audio' or part_type == 'input_audio' or part_type == 'audio_url' -%} ++ {{- emit_message(role, 'audio') -}} ++ {%- endif -%} ++ {%- endif -%} ++ {%- endfor -%} ++ {%- endif -%} ++ ++ {%- if role == 'assistant' -%} ++ {%- for call in message.get('tool_calls') or [] -%} ++ {%- set function = call.get('function', {}) if call.get('function', {}) is mapping else {} -%} ++ {%- if function.get('name') is string -%} ++ {%- set raw_arguments = function.get('arguments') or {} -%} ++ {%- set arguments_json = canonical_arguments(raw_arguments) -%} ++ {%- set invocation = '{"name":' + canonical_json(function.get('name')) + ',"args":' + arguments_json + '}' -%} ++ {{- emit_message('assistant', 'invoke_tool_json', invocation, function.get('name')) -}} ++ {%- endif -%} ++ {%- endfor -%} ++ {%- endif -%} ++ {%- endset -%} ++ {{- turn_out -}} ++ {#- Close each historical model turn, but never emit a bare terminator for an ++ assistant message that rendered no blocks. -#} ++ {%- if role == 'assistant' and turn_out -%} ++ {{- '<|content_model_end_sampling|>' -}} ++ {%- endif -%} ++ {%- else -%} ++ {{- raise_exception('Unknown message role: ' + (role | string)) -}} ++ {%- endif -%} ++{%- endfor -%} ++ ++{%- if not eff_ns.emitted -%} ++ {{- emit_message('system', 'text', 'Thinking effort level: ' + reasoning_effort_text(effort_value)) -}} ++{%- endif -%} ++ ++{%- if add_generation_prompt is defined and add_generation_prompt -%} ++ {{- '<|message_model|>' -}} ++{%- endif -%} ++{#- Unsloth translation to jinja from TML's parser #} +diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp +index e81ff647e..99fad9738 100644 +--- a/src/llama-arch.cpp ++++ b/src/llama-arch.cpp +@@ -144,6 +144,7 @@ static const std::map LLM_ARCH_NAMES = { + { LLM_ARCH_TALKIE, "talkie" }, + { LLM_ARCH_MELLUM, "mellum" }, + { LLM_ARCH_NANBEIGE, "nanbeige" }, ++ { LLM_ARCH_INKLING, "inkling" }, + { LLM_ARCH_UNKNOWN, "(unknown)" }, + }; + +@@ -320,6 +321,18 @@ static const std::map LLM_KV_NAMES = { + { LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" }, + + { LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, ++ ++ // inkling (private arch) ++ { LLM_KV_INKLING_D_REL, "%s.d_rel" }, ++ { LLM_KV_INKLING_REL_EXTENT, "%s.rel_extent" }, ++ { LLM_KV_INKLING_REL_EXTENT_SWA, "%s.rel_extent_swa" }, ++ { LLM_KV_INKLING_SHORTCONV_KERNEL, "%s.shortconv_kernel" }, ++ { LLM_KV_INKLING_DENSE_BLOCK_COUNT, "%s.dense_block_count" }, ++ { LLM_KV_INKLING_LOGIT_SCALE_DENOM, "%s.logit_scale_denom" }, ++ { LLM_KV_INKLING_LOG_SCALING_N_FLOOR, "%s.log_scaling_n_floor" }, ++ { LLM_KV_INKLING_LOG_SCALING_ALPHA, "%s.log_scaling_alpha" }, ++ { LLM_KV_INKLING_UNPADDED_VOCAB_SIZE, "%s.unpadded_vocab_size" }, ++ + // sentence-transformers dense modules feature dims + { LLM_KV_DENSE_2_FEAT_IN, "%s.dense_2_feat_in" }, + { LLM_KV_DENSE_2_FEAT_OUT, "%s.dense_2_feat_out" }, +@@ -592,9 +605,19 @@ static const std::map LLM_TENSOR_NAMES = { + { LLM_TENSOR_SHORTCONV_CONV, "blk.%d.shortconv.conv" }, + { LLM_TENSOR_SHORTCONV_INPROJ, "blk.%d.shortconv.in_proj" }, + { LLM_TENSOR_SHORTCONV_OUTPROJ, "blk.%d.shortconv.out_proj" }, ++ { LLM_TENSOR_ATTN_R, "blk.%d.attn_r" }, ++ { LLM_TENSOR_ATTN_REL_PROJ, "blk.%d.attn_rel_proj" }, ++ { LLM_TENSOR_SHORTCONV_K, "blk.%d.shortconv_k" }, ++ { LLM_TENSOR_SHORTCONV_V, "blk.%d.shortconv_v" }, ++ { LLM_TENSOR_SHORTCONV_ATTN, "blk.%d.shortconv_attn" }, ++ { LLM_TENSOR_SHORTCONV_MLP, "blk.%d.shortconv_mlp" }, ++ { LLM_TENSOR_FFN_GSCALE, "blk.%d.ffn_gscale" }, + { LLM_TENSOR_FFN_GATE_CHEXPS, "blk.%d.ffn_gate_chexps" }, + { LLM_TENSOR_FFN_DOWN_CHEXPS, "blk.%d.ffn_down_chexps" }, + { LLM_TENSOR_FFN_UP_CHEXPS, "blk.%d.ffn_up_chexps" }, ++ { LLM_TENSOR_FFN_GATE_SHEXPS, "blk.%d.ffn_gate_shexp" }, ++ { LLM_TENSOR_FFN_DOWN_SHEXPS, "blk.%d.ffn_down_shexp" }, ++ { LLM_TENSOR_FFN_UP_SHEXPS, "blk.%d.ffn_up_shexp" }, + { LLM_TENSOR_VISEXP_ATTN_QKV, "blk.%d.vis_attn_qkv" }, + { LLM_TENSOR_VISEXP_ATTN_OUT, "blk.%d.vis_attn_output" }, + { LLM_TENSOR_VISEXP_FFN_GATE, "blk.%d.vis_gate" }, +@@ -795,6 +818,9 @@ static const std::map LLM_TENSOR_INFOS = { + {LLM_TENSOR_FFN_UP_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_GATE_UP_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_DOWN_CHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, ++ {LLM_TENSOR_FFN_DOWN_SHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, ++ {LLM_TENSOR_FFN_GATE_SHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, ++ {LLM_TENSOR_FFN_UP_SHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_GATE_CHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_UP_CHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_EXP_PROBS_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, +@@ -836,6 +862,13 @@ static const std::map LLM_TENSOR_INFOS = { + {LLM_TENSOR_SHORTCONV_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, + {LLM_TENSOR_SHORTCONV_INPROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_SHORTCONV_OUTPROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, ++ {LLM_TENSOR_ATTN_R, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, ++ {LLM_TENSOR_ATTN_REL_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, ++ {LLM_TENSOR_SHORTCONV_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, ++ {LLM_TENSOR_SHORTCONV_V, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, ++ {LLM_TENSOR_SHORTCONV_ATTN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, ++ {LLM_TENSOR_SHORTCONV_MLP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, ++ {LLM_TENSOR_FFN_GSCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_VISEXP_ATTN_QKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_VISEXP_ATTN_OUT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_VISEXP_FFN_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, +@@ -968,6 +1001,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { + case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_QWEN35: + case LLM_ARCH_QWEN35MOE: ++ case LLM_ARCH_INKLING: + return true; + default: + return false; +@@ -1024,6 +1058,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { + case LLM_ARCH_MINIMAX_M3: + case LLM_ARCH_MISTRAL4: + case LLM_ARCH_KIMI_LINEAR: ++ case LLM_ARCH_INKLING: + return false; + default: + return true; +diff --git a/src/llama-arch.h b/src/llama-arch.h +index cbc97085e..6a283eb6b 100644 +--- a/src/llama-arch.h ++++ b/src/llama-arch.h +@@ -149,6 +149,7 @@ enum llm_arch { + LLM_ARCH_MINIMAX_M3, + LLM_ARCH_DFLASH, + LLM_ARCH_NANBEIGE, ++ LLM_ARCH_INKLING, + LLM_ARCH_UNKNOWN, + }; + +@@ -367,6 +368,17 @@ enum llm_kv { + + LLM_KV_SHORTCONV_L_CACHE, + ++ // inkling (private arch) ++ LLM_KV_INKLING_D_REL, ++ LLM_KV_INKLING_REL_EXTENT, ++ LLM_KV_INKLING_REL_EXTENT_SWA, ++ LLM_KV_INKLING_SHORTCONV_KERNEL, ++ LLM_KV_INKLING_DENSE_BLOCK_COUNT, ++ LLM_KV_INKLING_LOGIT_SCALE_DENOM, ++ LLM_KV_INKLING_LOG_SCALING_N_FLOOR, ++ LLM_KV_INKLING_LOG_SCALING_ALPHA, ++ LLM_KV_INKLING_UNPADDED_VOCAB_SIZE, ++ + LLM_KV_XIELU_ALPHA_N, + LLM_KV_XIELU_ALPHA_P, + LLM_KV_XIELU_BETA, +@@ -434,6 +446,10 @@ enum llm_tensor { + LLM_TENSOR_FFN_DOWN_CHEXPS, + LLM_TENSOR_FFN_GATE_CHEXPS, + LLM_TENSOR_FFN_UP_CHEXPS, ++ // Inkling: same GGUF names as *_SHEXP but registered as MUL_MAT_ID (3D shared-expert bank via ggml_mul_mat_id) ++ LLM_TENSOR_FFN_DOWN_SHEXPS, ++ LLM_TENSOR_FFN_GATE_SHEXPS, ++ LLM_TENSOR_FFN_UP_SHEXPS, + LLM_TENSOR_FFN_EXP_PROBS_B, + LLM_TENSOR_FFN_LATENT_DOWN, + LLM_TENSOR_FFN_LATENT_UP, +@@ -595,6 +611,14 @@ enum llm_tensor { + LLM_TENSOR_SHORTCONV_CONV, + LLM_TENSOR_SHORTCONV_INPROJ, + LLM_TENSOR_SHORTCONV_OUTPROJ, ++ // inkling (private arch) ++ LLM_TENSOR_ATTN_R, ++ LLM_TENSOR_ATTN_REL_PROJ, ++ LLM_TENSOR_SHORTCONV_K, ++ LLM_TENSOR_SHORTCONV_V, ++ LLM_TENSOR_SHORTCONV_ATTN, ++ LLM_TENSOR_SHORTCONV_MLP, ++ LLM_TENSOR_FFN_GSCALE, + LLM_TENSOR_VISEXP_ATTN_QKV, + LLM_TENSOR_VISEXP_ATTN_OUT, + LLM_TENSOR_VISEXP_FFN_GATE, +diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp +index a2467e660..36bf32193 100644 +--- a/src/llama-graph.cpp ++++ b/src/llama-graph.cpp +@@ -1574,9 +1574,15 @@ ggml_tensor * llm_graph_context::build_cvec( + ggml_tensor * llm_graph_context::build_lora_mm( + ggml_tensor * w, + ggml_tensor * cur, +- ggml_tensor * w_s) const { ++ ggml_tensor * w_s, ++ enum ggml_prec prec) const { + ggml_tensor * res = ggml_mul_mat(ctx0, w, cur); + ++ if (prec != GGML_PREC_DEFAULT) { ++ // Set precision on the base MUL_MAT before an optional scale/LoRA attachment changes the root op. ++ ggml_mul_mat_set_prec(res, prec); ++ } ++ + if (w_s) { + res = ggml_mul(ctx0, res, w_s); + } +@@ -2622,9 +2628,25 @@ ggml_tensor * llm_graph_context::build_attn_mha( + + ggml_tensor * cur; + +- const bool use_flash_attn = cparams.flash_attn && kq_b == nullptr; ++ // backends without a fused-bias flash-attention kernel (e.g. Metal for the ++ // banded op) can still run FA by folding the additive KQ bias into the mask ++ const bool fold_kq_b_into_mask = cparams.flash_attn && kq_b != nullptr && kq_mask != nullptr && ++ arch == LLM_ARCH_INKLING; ++ const bool use_flash_attn = cparams.flash_attn && (kq_b == nullptr || fold_kq_b_into_mask); + if (use_flash_attn) { +- GGML_ASSERT(kq_b == nullptr && "Flash attention does not support KQ bias yet"); ++ ggml_tensor * fa_mask = kq_mask; ++ if (fold_kq_b_into_mask) { ++ ggml_tensor * bias = ggml_cont(ctx0, kq_b); ++ if (kq_mask->ne[1] != bias->ne[1]) { ++ // the FA mask rows are padded; pad the bias with zeros to match ++ bias = ggml_pad(ctx0, bias, 0, (int)(kq_mask->ne[1] - bias->ne[1]), 0, 0); ++ } ++ // Metal adds F32 tensors, while FA consumes an F16 mask. Keep the ++ // dense bias construction on-device, then cast the combined mask. ++ ggml_tensor * mask_f32 = ggml_cast(ctx0, kq_mask, GGML_TYPE_F32); ++ fa_mask = ggml_cast(ctx0, ggml_add(ctx0, bias, mask_f32), GGML_TYPE_F16); ++ cb(fa_mask, "kq_mask_biased", il); ++ } + + if (v_trans) { + v = ggml_transpose(ctx0, v); +@@ -2639,7 +2661,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( + v = ggml_cast(ctx0, v, GGML_TYPE_F16); + } + +- cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias, ++ cur = ggml_flash_attn_ext(ctx0, q, k, v, fa_mask, kq_scale, hparams.f_max_alibi_bias, + hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); + res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); + +diff --git a/src/llama-graph.h b/src/llama-graph.h +index ab6e111e1..e74a58a56 100644 +--- a/src/llama-graph.h ++++ b/src/llama-graph.h +@@ -1083,7 +1083,8 @@ struct llm_graph_context { + ggml_tensor * build_lora_mm( + ggml_tensor * w, + ggml_tensor * cur, +- ggml_tensor * w_s = nullptr) const; ++ ggml_tensor * w_s = nullptr, ++ enum ggml_prec prec = GGML_PREC_DEFAULT) const; + + // do mat_mul_id, while optionally apply lora and per-expert scale + ggml_tensor * build_lora_mm_id( +diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp +index 50af97f35..36568582e 100644 +--- a/src/llama-hparams.cpp ++++ b/src/llama-hparams.cpp +@@ -191,6 +191,11 @@ uint32_t llama_hparams::n_embd_k_idx(uint32_t il) const { + } + + uint32_t llama_hparams::n_embd_r() const { ++ if (n_embd_r_impl != 0) { ++ // explicit override (e.g. inkling: 4 packed shortconv streams per layer) ++ return n_embd_r_impl; ++ } ++ + if (wkv_head_size != 0) { + // for RWKV models + return token_shift_count * n_embd; +diff --git a/src/llama-hparams.h b/src/llama-hparams.h +index 0b3626dc8..4d74e3527 100644 +--- a/src/llama-hparams.h ++++ b/src/llama-hparams.h +@@ -80,6 +80,17 @@ struct llama_hparams { + + uint32_t n_shortconv_l_cache = 0; + ++ // explicit override for the rolling state size per layer (see n_embd_r()) ++ uint32_t n_embd_r_impl = 0; ++ ++ // inkling (private arch) ++ uint32_t inkling_d_rel = 0; ++ uint32_t inkling_rel_extent = 0; // global (non-SWA) layers ++ uint32_t inkling_rel_extent_swa = 0; // local (SWA) layers ++ uint32_t inkling_log_n_floor = 0; // 0 = log-N scaling disabled ++ float inkling_log_alpha = 0.0f; ++ uint32_t inkling_unpadded_n_vocab = 0; // 0 = no padded-vocab masking ++ + std::array n_head_arr; + std::array n_head_kv_arr; + std::array n_ff_arr; +diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp +index 033e172d8..eebed9e04 100644 +--- a/src/llama-kv-cache.cpp ++++ b/src/llama-kv-cache.cpp +@@ -1863,6 +1863,51 @@ uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { + return result; + } + ++uint32_t llama_kv_cache::get_n_kv_pos_contiguous(const slot_info & sinfo, const llama_ubatch & ubatch) const { ++ if (sinfo.n_stream() != 1 || ubatch.n_seqs_unq != 1 || ubatch.n_tokens == 0 || ++ ubatch.pos == nullptr || ubatch.n_seq_id == nullptr || ++ ubatch.seq_id == nullptr || ubatch.seq_id[0] == nullptr) { ++ return 0; ++ } ++ ++ const llama_seq_id seq_id = ubatch.seq_id[0][0]; ++ if (seq_id < 0 || (size_t) seq_id >= seq_to_stream.size()) { ++ return 0; ++ } ++ ++ const uint32_t stream = seq_to_stream[seq_id]; ++ if (sinfo.strm[0] < 0 || (uint32_t) sinfo.strm[0] != stream || stream >= v_cells.size()) { ++ return 0; ++ } ++ ++ const auto & cells = v_cells[stream]; ++ const llama_pos pos_max = cells.seq_pos_max(seq_id); ++ ++ if (pos_max < 0 || pos_max >= (llama_pos) cells.size()) { ++ return 0; ++ } ++ ++ // the banded op aligns Q to the tail of K: the ubatch must be that monotonic tail, else dense bias ++ if ((uint32_t) pos_max + 1 < ubatch.n_tokens) { ++ return 0; ++ } ++ const llama_pos pos_start = pos_max + 1 - ubatch.n_tokens; ++ for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { ++ if (ubatch.pos[i] != pos_start + (llama_pos) i || ++ ubatch.n_seq_id[i] < 1 || ubatch.seq_id[i] == nullptr || ubatch.seq_id[i][0] != seq_id) { ++ return 0; ++ } ++ } ++ ++ for (llama_pos pos = 0; pos <= pos_max; ++pos) { ++ if (cells.is_empty(pos) || cells.pos_get(pos) != pos || !cells.seq_has(pos, seq_id)) { ++ return 0; ++ } ++ } ++ ++ return pos_max + 1; ++} ++ + ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { + const int32_t ikv = map_layer_ids.at(il); + +@@ -2441,6 +2486,39 @@ void llama_kv_cache::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch + } + } + ++void llama_kv_cache::set_input_pos_rel_flat(ggml_tensor * dst, const llama_ubatch * ubatch, uint32_t extent) const { ++ const int64_t n_tokens = ubatch->n_tokens; ++ ++ GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); ++ ++ int32_t * data = (int32_t *) dst->data; ++ ++ const int64_t n_kv = dst->ne[0]; ++ GGML_ASSERT(dst->ne[1] == n_tokens); ++ ++ // [n_kv, n_tokens] in GLOBAL token order (stream-major, same as the KQ mask) ++ for (int64_t i = 0; i < n_tokens; ++i) { ++ const llama_seq_id seq_id = ubatch->seq_id[i][0]; ++ ++ const auto & cells = v_cells[seq_to_stream[seq_id]]; ++ ++ const llama_pos p1 = ubatch->pos[i]; ++ ++ for (int64_t j = 0; j < n_kv; ++j) { ++ // use the ACTUAL absolute position in the KV cell; physical slot order is not monotonic ++ int32_t rel = (int32_t) extent; // zero-bias column ++ if (!cells.is_empty(j)) { ++ const llama_pos d = p1 - cells.pos_get(j); ++ if (d >= 0 && d < (llama_pos) extent) { ++ rel = (int32_t) d; ++ } ++ } ++ ++ data[i*n_kv + j] = (int32_t) (i*(extent + 1)) + rel; ++ } ++ } ++} ++ + void llama_kv_cache::set_input_k_rot(ggml_tensor * dst) const { + GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); + +@@ -3339,6 +3417,21 @@ uint32_t llama_kv_cache_context::get_n_kv() const { + return n_kv; + } + ++uint32_t llama_kv_cache_context::get_n_kv_pos_contiguous() const { ++ // Full-cache and update contexts do not carry a concrete ubatch/slot pair. ++ if (ubatches.empty() || sinfos.empty() || i_cur >= ubatches.size() || i_cur >= sinfos.size()) { ++ // reserve context: report the whole cache as position-contiguous so the worst-case graph ++ // is the banded path; reserving the dense fallback is unallocatable at large n_ctx ++ if (kv != nullptr && lctx == nullptr && kv->get_n_stream() == 1) { ++ return n_kv; ++ } ++ return 0; ++ } ++ ++ const uint32_t result = kv->get_n_kv_pos_contiguous(sinfos[i_cur], ubatches[i_cur]); ++ return result <= (uint32_t) n_kv ? result : 0; ++} ++ + ggml_type llama_kv_cache_context::type_k() const { + return kv->type_k(); + } +@@ -3407,6 +3500,10 @@ void llama_kv_cache_context::set_input_pos_bucket(ggml_tensor * dst, const llama + kv->set_input_pos_bucket(dst, ubatch); + } + ++void llama_kv_cache_context::set_input_pos_rel_flat(ggml_tensor * dst, const llama_ubatch * ubatch, uint32_t extent) const { ++ kv->set_input_pos_rel_flat(dst, ubatch, extent); ++} ++ + void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const { + kv->set_input_k_rot(dst); + } +diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h +index b253d9158..3b1d9b36f 100644 +--- a/src/llama-kv-cache.h ++++ b/src/llama-kv-cache.h +@@ -197,6 +197,9 @@ public: + + uint32_t get_n_kv(const slot_info & sinfo) const; + ++ // active cell count when position p lives in physical cell p; 0 for any non-contiguous layout ++ uint32_t get_n_kv_pos_contiguous(const slot_info & sinfo, const llama_ubatch & ubatch) const; ++ + // get views of the current state of the cache + ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; + ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; +@@ -247,6 +250,10 @@ public: + void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; + void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; + ++ // inkling: fill dst I32 [n_kv, n_tokens] with flat rel-bias gather indices, ++ // idx(i, j) = i*(extent + 1) + rel; empty/out-of-band cells map to the zero-bias column `extent` ++ void set_input_pos_rel_flat(ggml_tensor * dst, const llama_ubatch * ubatch, uint32_t extent) const; ++ + void set_input_k_rot(ggml_tensor * dst) const; + void set_input_v_rot(ggml_tensor * dst) const; + +@@ -403,6 +410,7 @@ public: + // + + uint32_t get_n_kv() const; ++ uint32_t get_n_kv_pos_contiguous() const; + + ggml_type type_k() const; + ggml_type type_v() const; +@@ -437,6 +445,7 @@ public: + void set_input_k_shift (ggml_tensor * dst) const; + void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; + void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; ++ void set_input_pos_rel_flat(ggml_tensor * dst, const llama_ubatch * ubatch, uint32_t extent) const; // inkling + + void set_input_k_rot(ggml_tensor * dst) const; + void set_input_v_rot(ggml_tensor * dst) const; +@@ -444,8 +453,8 @@ public: + private: + llama_memory_status status; + +- llama_kv_cache * kv; +- llama_context * lctx; ++ llama_kv_cache * kv = nullptr; ++ llama_context * lctx = nullptr; + + // + // update context +diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp +index 3812c594e..264b0b5b1 100644 +--- a/src/llama-model-saver.cpp ++++ b/src/llama-model-saver.cpp +@@ -29,6 +29,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { + case LLM_ARCH_STEP35: + case LLM_ARCH_MELLUM: + case LLM_ARCH_LAGUNA: ++ case LLM_ARCH_INKLING: + return false; + default: + return true; +diff --git a/src/llama-model.cpp b/src/llama-model.cpp +index b23a7e8d8..f95591860 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -311,6 +311,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params + return new llama_model_kimi_linear(params); + case LLM_ARCH_STEP35: + return new llama_model_step35(params); ++ case LLM_ARCH_INKLING: ++ return new llama_model_inkling(params); + default: + throw std::runtime_error(std::string("unsupported model architecture: '") + llm_arch_name(arch) + "'"); + } +@@ -2129,7 +2131,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, + // layer filters, so pick the right one here + llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; + llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; +- if (arch == LLM_ARCH_FALCON_H1) { ++ if (arch == LLM_ARCH_FALCON_H1 || arch == LLM_ARCH_INKLING) { ++ // all layers have both an attention KV cache and a recurrent (conv) state + filter_attn = [&](uint32_t) { return true; }; + filter_recr = [&](uint32_t) { return true; }; + } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { +@@ -2482,6 +2485,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { + case LLM_ARCH_NEMOTRON_H: + case LLM_ARCH_NEMOTRON_H_MOE: + case LLM_ARCH_KIMI_LINEAR: ++ case LLM_ARCH_INKLING: + return LLAMA_ROPE_TYPE_NONE; + + // use what we call a normal RoPE, operating on pairs of consecutive head values +diff --git a/src/llama-model.h b/src/llama-model.h +index d6a40fa30..8d00d5fa1 100644 +--- a/src/llama-model.h ++++ b/src/llama-model.h +@@ -532,6 +532,15 @@ struct llama_layer { + struct llama_layer_shortconv shortconv; + + struct llama_layer_nextn nextn; ++ ++ // inkling (private arch) ++ struct ggml_tensor * wr = nullptr; // attn_r [n_embd, n_head*d_rel] ++ struct ggml_tensor * attn_rel_proj = nullptr; // [rel_extent, d_rel] (checkpoint [d_rel, E] orientation) ++ struct ggml_tensor * shortconv_k = nullptr; // [K, kvw] ++ struct ggml_tensor * shortconv_v = nullptr; // [K, kvw] ++ struct ggml_tensor * shortconv_attn = nullptr; // [K, n_embd] ++ struct ggml_tensor * shortconv_mlp = nullptr; // [K, n_embd] ++ struct ggml_tensor * ffn_gscale = nullptr; // F32 [1] + }; + + struct llama_device { +diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp +index 03a6c34fe..8ab4e403c 100644 +--- a/src/llama-quant.cpp ++++ b/src/llama-quant.cpp +@@ -326,6 +326,16 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param + quantize &= name.find("ssm_conv1d") == std::string::npos; + quantize &= name.find("shortconv.conv.weight") == std::string::npos; + ++ // keep Inkling's shortconv kernels and rel-proj table unquantized; arch-gated so the ++ // name substrings cannot hit another architecture ++ if (arch == LLM_ARCH_INKLING) { ++ quantize &= name.find("shortconv_k.weight") == std::string::npos; ++ quantize &= name.find("shortconv_v.weight") == std::string::npos; ++ quantize &= name.find("shortconv_attn.weight") == std::string::npos; ++ quantize &= name.find("shortconv_mlp.weight") == std::string::npos; ++ quantize &= name.find("attn_rel_proj.weight") == std::string::npos; ++ } ++ + // do not quantize MiniMax's indexer projection weights, they are tiny + quantize &= name.find("indexer.k_proj.weight") == std::string::npos; + quantize &= name.find("indexer.q_proj.weight") == std::string::npos; +diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp +index 9164a4dd8..74942e812 100644 +--- a/src/llama-vocab.cpp ++++ b/src/llama-vocab.cpp +@@ -433,6 +433,12 @@ struct llm_tokenizer_bpe : llm_tokenizer { + "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))*((?=[\\p{L}])([^A-Z]))+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))+((?=[\\p{L}])([^A-Z]))*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", + }; + break; ++ case LLAMA_VOCAB_PRE_TYPE_INKLING: ++ // o200k-family with \p{M} in the letter classes; own pre-type so GPT4O / MINIMAX_M2 stay unchanged ++ regex_exprs = { ++ "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))*((?=[\\p{L}\\p{M}])([^A-Z]))+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))+((?=[\\p{L}\\p{M}])([^A-Z]))*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", ++ }; ++ break; + case LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI: + // Same lookaheads as GPT4O but with \p{M} added so combining marks + // (diacritics) attach to their base letters. Avoids excessive +@@ -2295,6 +2301,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { + tokenizer_pre == "talkie") { + pre_type = LLAMA_VOCAB_PRE_TYPE_GPT4O; + clean_spaces = false; ++ } else if ( ++ tokenizer_pre == "inkling") { ++ pre_type = LLAMA_VOCAB_PRE_TYPE_INKLING; ++ clean_spaces = false; + } else if ( + tokenizer_pre == "granite-embed-multi-97m") { + pre_type = LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI; +diff --git a/src/llama-vocab.h b/src/llama-vocab.h +index b7c289263..65e43f671 100644 +--- a/src/llama-vocab.h ++++ b/src/llama-vocab.h +@@ -65,6 +65,7 @@ enum llama_vocab_pre_type { + LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI = 54, + LLAMA_VOCAB_PRE_TYPE_MELLUM2 = 55, + LLAMA_VOCAB_PRE_TYPE_LAGUNA = 56, ++ LLAMA_VOCAB_PRE_TYPE_INKLING = 57, + }; + + struct LLM_KV; +diff --git a/src/models/inkling.cpp b/src/models/inkling.cpp +new file mode 100644 +index 000000000..6c9c1aef8 +--- /dev/null ++++ b/src/models/inkling.cpp +@@ -0,0 +1,695 @@ ++// Inkling (PRIVATE arch): hybrid iSWA attention + per-layer packed shortconv state; see INKLING_DESIGN.md. ++ ++#include "models.h" ++ ++#include "../llama-kv-cache-iswa.h" ++#include "../llama-kv-cache.h" ++#include "../llama-memory-hybrid-iswa.h" ++#include "../llama-memory-recurrent.h" ++ ++#include ++#include ++ ++void llama_model_inkling::load_arch_hparams(llama_model_loader & ml) { ++ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); ++ ++ ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); ++ ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); ++ ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); ++ ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); ++ ++ ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); ++ hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; // visible iff pos_q - pos_k < n_swa (includes self) ++ ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer()); ++ ++ for (uint32_t il = 0; il < hparams.n_layer(); ++il) { ++ hparams.is_recr_impl[il] = 1; ++ } ++ ++ ml.get_key(LLM_KV_INKLING_D_REL, hparams.inkling_d_rel); ++ ml.get_key(LLM_KV_INKLING_REL_EXTENT, hparams.inkling_rel_extent); ++ ml.get_key(LLM_KV_INKLING_REL_EXTENT_SWA, hparams.inkling_rel_extent_swa); ++ ml.get_key(LLM_KV_INKLING_SHORTCONV_KERNEL, hparams.n_shortconv_l_cache); ++ ml.get_key(LLM_KV_INKLING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); ++ ++ float logit_scale_denom = 0.0f; ++ ml.get_key(LLM_KV_INKLING_LOGIT_SCALE_DENOM, logit_scale_denom); ++ GGML_ASSERT(logit_scale_denom != 0.0f); ++ hparams.f_logit_scale = 1.0f / logit_scale_denom; ++ ++ ml.get_key(LLM_KV_INKLING_LOG_SCALING_N_FLOOR, hparams.inkling_log_n_floor, false); ++ ml.get_key(LLM_KV_INKLING_LOG_SCALING_ALPHA, hparams.inkling_log_alpha, false); ++ ml.get_key(LLM_KV_INKLING_UNPADDED_VOCAB_SIZE, hparams.inkling_unpadded_n_vocab, false); ++ ++ GGML_ASSERT(hparams.n_shortconv_l_cache > 1); ++ GGML_ASSERT(hparams.inkling_d_rel > 0); ++ GGML_ASSERT(hparams.inkling_rel_extent > 0 && hparams.inkling_rel_extent_swa > 0); ++ ++ // uniform state per cell: 4 packed streams [k | v | attn | mlp] of last K-1 columns, k/v sized for the widest layer ++ const uint32_t d_conv = hparams.n_shortconv_l_cache - 1; ++ hparams.n_embd_r_impl = d_conv * (hparams.n_embd_k_gqa_max() + hparams.n_embd_v_gqa_max() + 2*hparams.n_embd); ++ ++ type = LLM_TYPE_UNKNOWN; ++} ++ ++void llama_model_inkling::load_arch_tensors(llama_model_loader &) { ++ LLAMA_LOAD_LOCALS; ++ ++ const int64_t head_dim = hparams.n_embd_head_k(); ++ const int64_t d_rel = hparams.inkling_d_rel; ++ const int64_t K = hparams.n_shortconv_l_cache; ++ const int64_t n_ff_exp = hparams.n_ff_exp; ++ const int64_t n_shexp = hparams.n_expert_shared; ++ ++ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); ++ tok_norm = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD_NORM, "weight", 0), {n_embd}, 0); // bid 0: compute on the first layer's device ++ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); ++ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); ++ ++ for (int i = 0; i < n_layer; ++i) { ++ auto & layer = layers[i]; ++ ++ const int64_t n_head_kv_i = hparams.n_head_kv(i); ++ const int64_t kvw = n_head_kv_i * head_dim; ++ const int64_t rel_extent = hparams.is_swa(i) ? hparams.inkling_rel_extent_swa : hparams.inkling_rel_extent; ++ ++ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); ++ ++ layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head*head_dim}, 0); ++ layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, kvw}, 0); ++ layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, kvw}, 0); ++ layer.wr = create_tensor(tn(LLM_TENSOR_ATTN_R, "weight", i), {n_embd, n_head*d_rel}, 0); ++ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head*head_dim, n_embd}, 0); ++ ++ layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {head_dim}, 0); ++ layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {head_dim}, 0); ++ ++ // stored in checkpoint orientation [d_rel, E] -> gguf ne = [E, d_rel] ++ layer.attn_rel_proj = create_tensor(tn(LLM_TENSOR_ATTN_REL_PROJ, "weight", i), {rel_extent, d_rel}, 0); ++ ++ layer.shortconv_k = create_tensor(tn(LLM_TENSOR_SHORTCONV_K, "weight", i), {K, kvw}, 0); ++ layer.shortconv_v = create_tensor(tn(LLM_TENSOR_SHORTCONV_V, "weight", i), {K, kvw}, 0); ++ layer.shortconv_attn = create_tensor(tn(LLM_TENSOR_SHORTCONV_ATTN, "weight", i), {K, n_embd}, 0); ++ layer.shortconv_mlp = create_tensor(tn(LLM_TENSOR_SHORTCONV_MLP, "weight", i), {K, n_embd}, 0); ++ ++ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); ++ layer.ffn_gscale = create_tensor(tn(LLM_TENSOR_FFN_GSCALE, "weight", i), {1}, 0); ++ ++ if (i < (int) hparams.n_layer_dense_lead) { ++ const int64_t n_ff_i = hparams.n_ff(i); ++ ++ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff_i}, 0); ++ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff_i}, 0); ++ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff_i, n_embd}, 0); ++ } else { ++ GGML_ASSERT(n_expert > 0 && n_expert_used > 0 && n_shexp > 0); ++ ++ // gate holds n_expert + n_shexp rows (incl. shared-expert sink logits) ++ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert + n_shexp}, 0); ++ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); ++ ++ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); ++ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); ++ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); ++ ++ // shared experts stacked as an n_shexp bank, registered MUL_MAT_ID so the loader picks a mul_mat_id-capable buffer ++ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXPS, "weight", i), {n_embd, n_ff_exp, n_shexp}, 0); ++ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXPS, "weight", i), {n_embd, n_ff_exp, n_shexp}, 0); ++ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXPS, "weight", i), {n_ff_exp, n_embd, n_shexp}, 0); ++ } ++ } ++} ++ ++class llm_graph_input_inkling : public llm_graph_input_i { ++public: ++ llm_graph_input_inkling( ++ const llama_hparams & hparams, ++ const llama_memory_hybrid_iswa_context * mctx) : ++ hparams(hparams), ++ mctx(mctx) {} ++ virtual ~llm_graph_input_inkling() = default; ++ ++ void set_input(const llama_ubatch * ubatch) override { ++ if (tau) { ++ GGML_ASSERT(ggml_backend_buffer_is_host(tau->buffer)); ++ float * data = (float *) tau->data; ++ ++ const float n_floor = (float) hparams.inkling_log_n_floor; ++ const float alpha = hparams.inkling_log_alpha; ++ ++ for (int64_t i = 0; i < (int64_t) ubatch->n_tokens; ++i) { ++ const float eff = (float) (ubatch->pos[i] + 1) / n_floor; ++ data[i] = 1.0f + alpha*logf(std::max(eff, 1.0f)); ++ } ++ } ++ ++ if (rel_idx) { ++ mctx->get_attn()->get_base()->set_input_pos_rel_flat(rel_idx, ubatch, hparams.inkling_rel_extent); ++ } ++ ++ if (rel_idx_swa) { ++ mctx->get_attn()->get_swa()->set_input_pos_rel_flat(rel_idx_swa, ubatch, hparams.inkling_rel_extent_swa); ++ } ++ ++ if (vocab_mask) { ++ GGML_ASSERT(ggml_backend_buffer_is_host(vocab_mask->buffer)); ++ float * data = (float *) vocab_mask->data; ++ ++ const int64_t n_vocab = vocab_mask->ne[0]; ++ const int64_t n_unpadded = hparams.inkling_unpadded_n_vocab; ++ ++ for (int64_t id = 0; id < n_vocab; ++id) { ++ data[id] = id < n_unpadded ? 0.0f : -INFINITY; ++ } ++ } ++ ++ if (shexp_idx) { ++ GGML_ASSERT(ggml_backend_buffer_is_host(shexp_idx->buffer)); ++ int32_t * data = (int32_t *) shexp_idx->data; ++ ++ const int64_t n_shexp = shexp_idx->ne[0]; ++ const int64_t n_tokens = shexp_idx->ne[1]; ++ ++ for (int64_t j = 0; j < n_tokens; ++j) { ++ for (int64_t s = 0; s < n_shexp; ++s) { ++ data[j*n_shexp + s] = (int32_t) s; ++ } ++ } ++ } ++ } ++ ++ ggml_tensor * tau = nullptr; // F32 [1, 1, n_tokens] ++ ggml_tensor * rel_idx = nullptr; // I32 [n_kv_base, n_tokens] ++ ggml_tensor * rel_idx_swa = nullptr; // I32 [n_kv_swa, n_tokens] ++ ggml_tensor * vocab_mask = nullptr; // F32 [n_vocab] ++ ggml_tensor * shexp_idx = nullptr; // I32 [n_shexp, n_tokens], constant 0..n_shexp-1 ++ ++ const llama_hparams hparams; ++ ++ const llama_memory_hybrid_iswa_context * mctx; ++}; ++ ++std::unique_ptr llama_model_inkling::build_arch_graph(const llm_graph_params & params) const { ++ return std::make_unique(*this, params); ++} ++ ++llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_params & params) : ++ llm_graph_context(params) { ++ ++ const int64_t head_dim = hparams.n_embd_head_k(); ++ const int64_t d_rel = hparams.inkling_d_rel; ++ const int64_t d_conv = hparams.n_shortconv_l_cache - 1; ++ const int64_t n_embd_r = hparams.n_embd_r(); ++ const int64_t kw_max = hparams.n_embd_k_gqa_max(); ++ const int64_t vw_max = hparams.n_embd_v_gqa_max(); ++ ++ // packed conv-state stream offsets within one cell: [k | v | attn | mlp] ++ const int64_t off_k = 0; ++ const int64_t off_v = d_conv*kw_max; ++ const int64_t off_attn = d_conv*(kw_max + vw_max); ++ const int64_t off_mlp = d_conv*(kw_max + vw_max + n_embd); ++ ++ const auto * mctx_hyb = static_cast(mctx); ++ const auto * mctx_recr = mctx_hyb->get_recr(); ++ const auto * mctx_attn = mctx_hyb->get_attn(); ++ ++ const uint32_t kv_head = mctx_recr->get_head(); ++ ++ const int64_t n_seq_tokens = ubatch.n_seq_tokens; ++ const int64_t n_seqs = ubatch.n_seqs; ++ ++ GGML_ASSERT(n_seqs != 0); ++ GGML_ASSERT(ubatch.equal_seqs()); ++ GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); ++ ++ const uint32_t n_kv_flash_base = cparams.flash_attn ? mctx_attn->get_base()->get_n_kv_pos_contiguous() : 0; ++ const uint32_t n_kv_flash_swa = cparams.flash_attn ? mctx_attn->get_swa ()->get_n_kv_pos_contiguous() : 0; ++ ++ bool has_global = false; ++ bool needs_rel_idx_local = false; ++ bool needs_rel_idx_global = false; ++ ++ const auto banded_cache_type_supported = [](ggml_type type) { ++ return type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16; ++ }; ++ ++ // probe once per device whether the fused banded FA op is supported there ++ // (e.g. CUDA/CPU yes, Metal no). layers on devices without it use the ++ // dense-bias branch, which build_attn_mha folds into the FA mask. ++ std::unordered_map banded_dev_support; ++ const auto dev_supports_banded_flash = [&](int il) { ++ ggml_backend_dev_t dev = model.dev_layer(il); ++ if (dev == nullptr) { ++ return true; ++ } ++ const auto it = banded_dev_support.find(dev); ++ if (it != banded_dev_support.end()) { ++ return it->second; ++ } ++ // representative probe op; tensors are metadata-only in the graph ctx ++ // and never expanded into the compute graph ++ const int64_t ext = std::max(1, hparams.inkling_rel_extent); ++ ggml_tensor * pq = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, head_dim, 1, n_head, 1); ++ ggml_tensor * pk = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, head_dim, 256, hparams.n_head_kv(il), 1); ++ ggml_tensor * pv = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, head_dim, 256, hparams.n_head_kv(il), 1); ++ ggml_tensor * pm = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, 256, 1, 1, 1); ++ ggml_tensor * pr = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, ext, n_head, 1, 1); ++ ggml_tensor * op = ggml_flash_attn_ext_banded(ctx0, pq, pk, pv, pm, pr, 1.0f/float(head_dim), ext); ++ const bool ok = ggml_backend_dev_supports_op(dev, op); ++ banded_dev_support.emplace(dev, ok); ++ return ok; ++ }; ++ ++ const auto use_banded_flash = [&](int il) { ++ const auto * cache = hparams.is_swa(il) ? mctx_attn->get_swa() : mctx_attn->get_base(); ++ const uint32_t n_kv_flash = hparams.is_swa(il) ? n_kv_flash_swa : n_kv_flash_base; ++ ++ // get_n_kv_pos_contiguous() is 0 for multi-sequence ubatches; the reserve context reports full n_kv ++ return cparams.flash_attn && ++ n_kv_flash > 0 && ++ (head_dim == 64 || head_dim == 128) && ++ hparams.n_embd_head_v(il) == head_dim && ++ hparams.n_head(il) % hparams.n_head_kv(il) == 0 && ++ banded_cache_type_supported(cache->type_k()) && ++ banded_cache_type_supported(cache->type_v()) && ++ dev_supports_banded_flash(il); ++ }; ++ ++ for (int il = 0; il < n_layer; ++il) { ++ if (hparams.is_swa(il)) { ++ needs_rel_idx_local |= !use_banded_flash(il); ++ } else { ++ has_global = true; ++ needs_rel_idx_global |= !use_banded_flash(il); ++ } ++ } ++ ++ auto inp = std::make_unique(hparams, mctx_hyb); ++ ++ if (hparams.inkling_log_n_floor > 0 && has_global) { ++ inp->tau = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens); ++ ggml_set_input(inp->tau); ++ ggml_set_name(inp->tau, "inkling_tau"); ++ } ++ ++ if (needs_rel_idx_global) { ++ const int64_t n_kv = mctx_attn->get_base()->get_n_kv(); ++ inp->rel_idx = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_tokens); ++ ggml_set_input(inp->rel_idx); ++ ggml_set_name(inp->rel_idx, "inkling_rel_idx"); ++ } ++ ++ if (needs_rel_idx_local) { ++ const int64_t n_kv_swa = mctx_attn->get_swa()->get_n_kv(); ++ inp->rel_idx_swa = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv_swa, n_tokens); ++ ggml_set_input(inp->rel_idx_swa); ++ ggml_set_name(inp->rel_idx_swa, "inkling_rel_idx_swa"); ++ } ++ ++ const int64_t n_vocab = model.vocab.n_tokens(); ++ if (!cparams.embeddings && hparams.inkling_unpadded_n_vocab > 0 && (int64_t) hparams.inkling_unpadded_n_vocab < n_vocab) { ++ inp->vocab_mask = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_vocab); ++ ggml_set_input(inp->vocab_mask); ++ ggml_set_name(inp->vocab_mask, "inkling_vocab_mask"); ++ } ++ ++ // shared experts go through mul_mat_id: 2D views into a repacked/quantized 3D bank are invalid ++ if (hparams.n_expert_shared > 0 && (uint32_t) n_layer > hparams.n_layer_dense_lead) { ++ inp->shexp_idx = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, hparams.n_expert_shared, n_tokens); ++ ggml_set_input(inp->shexp_idx); ++ ggml_set_name(inp->shexp_idx, "inkling_shexp_idx"); ++ } ++ ++ ggml_tensor * tau = inp->tau; ++ ggml_tensor * rel_idx = inp->rel_idx; ++ ggml_tensor * rel_idx_swa = inp->rel_idx_swa; ++ ggml_tensor * vocab_mask = inp->vocab_mask; ++ ggml_tensor * shexp_idx = inp->shexp_idx; ++ ++ res->add_input(std::move(inp)); ++ ++ auto * inp_hybrid = build_inp_mem_hybrid_iswa(); ++ ++ // shared by the 4 stream sub-views; build_rs must run exactly once per layer (it zero-inits fresh states) ++ ggml_tensor * conv_rs_cur = nullptr; ++ ++ // sconv(x) = x + causal_depthwise_conv1d(x); rolling state = last K-1 inputs ++ auto build_sconv = [&](ggml_tensor * x2d, ggml_tensor * kernel, int64_t off, int il) -> ggml_tensor * { ++ const int64_t w = x2d->ne[0]; ++ ++ ggml_tensor * x3 = ggml_reshape_3d(ctx0, x2d, w, n_seq_tokens, n_seqs); ++ ggml_tensor * xt = ggml_transpose(ctx0, x3); // time-major for the conv ++ ++ ggml_tensor * conv_state = mctx_recr->get_r_l(il); ++ ggml_tensor * conv_rs = conv_rs_cur; // {n_embd_r, n_seqs} ++ GGML_ASSERT(conv_rs != nullptr); ++ ++ const size_t sz = ggml_element_size(conv_rs); ++ ++ // this stream's slice of the packed state ++ ggml_tensor * state = ggml_view_3d(ctx0, conv_rs, d_conv, w, n_seqs, ++ d_conv*sz, conv_rs->nb[1], off*sz); ++ ++ ggml_tensor * sx = ggml_concat(ctx0, state, xt, 0); // {d_conv + n_seq_tokens, w, n_seqs} ++ ++ // write the last d_conv time columns back into the cache ++ ggml_tensor * new_state = ggml_view_3d(ctx0, sx, d_conv, w, n_seqs, ++ sx->nb[1], sx->nb[2], (sx->ne[0] - d_conv)*sx->nb[0]); ++ ggml_tensor * state_dst = ggml_view_3d(ctx0, conv_state, d_conv, w, n_seqs, ++ d_conv*sz, n_embd_r*sz, (kv_head*n_embd_r + off)*sz); ++ ggml_build_forward_expand(gf, ggml_cpy(ctx0, new_state, state_dst)); ++ ++ ggml_tensor * conv_out = ggml_ssm_conv(ctx0, sx, kernel); // {w, n_seq_tokens, n_seqs} ++ ++ ggml_tensor * y = ggml_add(ctx0, x3, conv_out); // built-in residual, no activation ++ ++ return ggml_reshape_2d(ctx0, y, w, n_seq_tokens*n_seqs); ++ }; ++ ++ auto build_attn_block = [&](ggml_tensor * cur, int il) -> ggml_tensor * { ++ const auto & layer = model.layers[il]; ++ ++ const bool is_swa = hparams.is_swa(il); ++ const int64_t n_head_kv = hparams.n_head_kv(il); ++ const int64_t rel_extent = is_swa ? hparams.inkling_rel_extent_swa : hparams.inkling_rel_extent; ++ ++ ggml_tensor * q = build_lora_mm(layer.wq, cur); ++ ggml_tensor * k = build_lora_mm(layer.wk, cur); ++ ggml_tensor * v = build_lora_mm(layer.wv, cur); ++ ggml_tensor * r = build_lora_mm(layer.wr, cur); ++ cb(q, "inkling_attn_q", il); ++ cb(k, "inkling_attn_k", il); ++ cb(v, "inkling_attn_v", il); ++ cb(r, "inkling_attn_r", il); ++ ++ // k/v short convs on the flat projections, before the head reshape ++ k = build_sconv(k, layer.shortconv_k, off_k, il); ++ v = build_sconv(v, layer.shortconv_v, off_v, il); ++ cb(k, "inkling_attn_k_sconv", il); ++ cb(v, "inkling_attn_v_sconv", il); ++ ++ q = ggml_reshape_3d(ctx0, q, head_dim, n_head, n_tokens); ++ k = ggml_reshape_3d(ctx0, k, head_dim, n_head_kv, n_tokens); ++ v = ggml_reshape_3d(ctx0, v, head_dim, n_head_kv, n_tokens); ++ ++ q = build_norm(q, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); ++ k = build_norm(k, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); ++ cb(q, "inkling_attn_q_norm", il); ++ cb(k, "inkling_attn_k_norm", il); ++ ++ // log-N tau on global layers only, after q_norm ++ if (tau && !is_swa) { ++ q = ggml_mul(ctx0, q, tau); ++ } ++ ++ // relative position bias ++ ggml_tensor * r2 = ggml_reshape_2d(ctx0, r, d_rel, n_head*n_tokens); ++ ++ // proj stored [E, d_rel]; transpose so ggml_mul_mat contracts over d_rel ++ ggml_tensor * proj = ggml_cont(ctx0, ggml_transpose(ctx0, layer.attn_rel_proj)); // {d_rel, E} ++ ++ ggml_tensor * rel = ggml_mul_mat(ctx0, proj, r2); // {E, n_head*n_tokens} ++ ggml_mul_mat_set_prec(rel, GGML_PREC_F32_PEDANTIC); ++ rel = ggml_reshape_3d(ctx0, rel, rel_extent, n_head, n_tokens); ++ ++ if (tau && !is_swa) { ++ rel = ggml_mul(ctx0, rel, tau); ++ } ++ cb(rel, "inkling_rel_logits", il); ++ ++ auto * inp_attn = inp_hybrid->get_attn(); ++ const int64_t n_stream = (is_swa ? inp_attn->get_kq_mask_swa() : inp_attn->get_kq_mask())->ne[3]; ++ GGML_ASSERT(n_tokens % n_stream == 0); ++ ++ if (use_banded_flash(il)) { ++ GGML_ASSERT(q->type == GGML_TYPE_F32); ++ auto * k_rot = is_swa ? inp_attn->self_k_rot_swa : inp_attn->self_k_rot; ++ auto * v_rot = is_swa ? inp_attn->self_v_rot_swa : inp_attn->self_v_rot; ++ ++ if (k_rot) { ++ q = llama_mul_mat_hadamard(ctx0, q, k_rot); ++ k = llama_mul_mat_hadamard(ctx0, k, k_rot); ++ } ++ if (v_rot) { ++ v = llama_mul_mat_hadamard(ctx0, v, v_rot); ++ } ++ ++ ggml_build_forward_expand(gf, q); ++ ggml_build_forward_expand(gf, k); ++ ggml_build_forward_expand(gf, v); ++ ++ const auto * cache = is_swa ? inp_attn->mctx->get_swa() : inp_attn->mctx->get_base(); ++ const auto & k_idxs = is_swa ? inp_attn->get_k_idxs_swa() : inp_attn->get_k_idxs(); ++ const auto & v_idxs = is_swa ? inp_attn->get_v_idxs_swa() : inp_attn->get_v_idxs(); ++ ++ ggml_build_forward_expand(gf, cache->cpy_k(ctx0, k, k_idxs, il)); ++ ggml_build_forward_expand(gf, cache->cpy_v(ctx0, v, v_idxs, il)); ++ ++ ggml_tensor * q_fa = ggml_view_4d(ctx0, q, ++ q->ne[0], q->ne[1], q->ne[2]/n_stream, n_stream, ++ q->nb[1], q->nb[2], q->nb[3]/n_stream, 0); ++ ggml_tensor * k_fa = cache->get_k(ctx0, il); ++ ggml_tensor * v_fa = cache->get_v(ctx0, il); ++ ++ const int64_t n_kv_flash = is_swa ? n_kv_flash_swa : n_kv_flash_base; ++ GGML_ASSERT(n_stream == 1 && n_kv_flash <= k_fa->ne[2]); ++ ++ k_fa = ggml_view_4d(ctx0, k_fa, ++ k_fa->ne[0], k_fa->ne[1], n_kv_flash, k_fa->ne[3], ++ k_fa->nb[1], k_fa->nb[2], k_fa->nb[3], 0); ++ ++ const bool v_trans = v_fa->nb[1] > v_fa->nb[2]; ++ if (v_trans) { ++ GGML_ASSERT(n_kv_flash <= v_fa->ne[0]); ++ v_fa = ggml_view_4d(ctx0, v_fa, ++ n_kv_flash, v_fa->ne[1], v_fa->ne[2], v_fa->ne[3], ++ v_fa->nb[1], v_fa->nb[2], v_fa->nb[3], 0); ++ } else { ++ GGML_ASSERT(n_kv_flash <= v_fa->ne[2]); ++ v_fa = ggml_view_4d(ctx0, v_fa, ++ v_fa->ne[0], v_fa->ne[1], n_kv_flash, v_fa->ne[3], ++ v_fa->nb[1], v_fa->nb[2], v_fa->nb[3], 0); ++ } ++ ++ q_fa = ggml_permute(ctx0, q_fa, 0, 2, 1, 3); ++ k_fa = ggml_permute(ctx0, k_fa, 0, 2, 1, 3); ++ v_fa = ggml_permute(ctx0, v_fa, 0, 2, 1, 3); ++ ++ if (v_trans) { ++ v_fa = ggml_transpose(ctx0, v_fa); ++ } ++ if (k_fa->type == GGML_TYPE_F32) { ++ k_fa = ggml_cast(ctx0, k_fa, GGML_TYPE_F16); ++ } ++ if (v_fa->type == GGML_TYPE_F32) { ++ v_fa = ggml_cast(ctx0, v_fa, GGML_TYPE_F16); ++ } ++ ++ ggml_tensor * rel_fa = ggml_reshape_4d(ctx0, rel, ++ rel_extent, n_head, n_tokens/n_stream, n_stream); ++ ggml_tensor * mask = is_swa ? inp_attn->get_kq_mask_swa() : inp_attn->get_kq_mask(); ++ mask = ggml_cont(ctx0, ggml_view_4d(ctx0, mask, ++ n_kv_flash, mask->ne[1], mask->ne[2], mask->ne[3], ++ mask->nb[1], mask->nb[2], mask->nb[3], 0)); ++ ++ cur = ggml_flash_attn_ext_banded(ctx0, q_fa, k_fa, v_fa, mask, rel_fa, ++ 1.0f/float(head_dim), rel_extent); ++ ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32); ++ res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); ++ ++ cur = ggml_reshape_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]); ++ ggml_build_forward_expand(gf, cur); ++ cb(cur, "kqv_out", il); ++ ++ if (v_rot) { ++ cur = llama_mul_mat_hadamard(ctx0, cur, v_rot); ++ } ++ cur = build_lora_mm(layer.wo, cur); ++ } else { ++ // soft_max_ext scales kq + kq_b jointly: fold 1/head_dim into q to keep the bias unscaled ++ q = ggml_scale(ctx0, q, 1.0f/float(head_dim)); ++ ++ // zero column at index E is gathered by out-of-band / empty-cell indices ++ rel = ggml_pad(ctx0, rel, 1, 0, 0, 0); // {E+1, n_head, n_tokens} ++ rel = ggml_cont(ctx0, ggml_permute(ctx0, rel, 1, 0, 2, 3)); // {n_head, E+1, n_tokens} ++ rel = ggml_reshape_2d(ctx0, rel, n_head, (rel_extent + 1)*n_tokens); ++ ++ ggml_tensor * idx = is_swa ? rel_idx_swa : rel_idx; // {n_kv, n_tokens} ++ GGML_ASSERT(idx != nullptr); ++ const int64_t n_kv = idx->ne[0]; ++ ++ ggml_tensor * idx1 = ggml_reshape_1d(ctx0, idx, n_kv*n_tokens); ++ ++ ggml_tensor * kq_b = ggml_get_rows(ctx0, rel, idx1); // {n_head, n_kv*n_tokens} ++ kq_b = ggml_reshape_3d(ctx0, kq_b, n_head, n_kv, n_tokens); ++ kq_b = ggml_cont(ctx0, ggml_permute(ctx0, kq_b, 2, 0, 1, 3)); // {n_kv, n_tokens, n_head} ++ cb(kq_b, "inkling_kq_b", il); ++ ++ // streamed kq is [n_kv, n_tokens/n_stream, n_head, n_stream], tokens stream-major: view kq_b to match (same trick as the KQ mask) ++ if (n_stream > 1) { ++ kq_b = ggml_view_4d(ctx0, kq_b, n_kv, n_tokens/n_stream, n_head, n_stream, ++ kq_b->nb[1], ++ kq_b->nb[2], ++ (n_tokens/n_stream)*kq_b->nb[1], ++ 0); ++ } ++ ++ cur = build_attn(inp_attn, ++ layer.wo, NULL, NULL, ++ q, k, v, kq_b, nullptr, nullptr, 1.0f, il); ++ } ++ cb(cur, "inkling_attn_o", il); ++ ++ return cur; ++ }; ++ ++ auto build_dense_ffn = [&](ggml_tensor * cur, int il) -> ggml_tensor * { ++ cur = build_ffn(cur, ++ model.layers[il].ffn_up, NULL, NULL, ++ model.layers[il].ffn_gate, NULL, NULL, ++ model.layers[il].ffn_down, NULL, NULL, ++ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); ++ cur = ggml_mul(ctx0, cur, model.layers[il].ffn_gscale); ++ cb(cur, "inkling_dense_ffn_out", il); ++ return cur; ++ }; ++ ++ // custom MoE routing (not expressible via build_moe_ffn): select by top-k(sigmoid(logits) + bias), weight by softmax(logsigmoid(raw logits)) * scales ++ auto build_moe = [&](ggml_tensor * cur, int il) -> ggml_tensor * { ++ const auto & layer = model.layers[il]; ++ ++ const int64_t n_shexp = hparams.n_expert_shared; ++ ++ ggml_tensor * logits = build_lora_mm( ++ layer.ffn_gate_inp, cur, nullptr, GGML_PREC_F32_PEDANTIC); // {n_expert + n_shexp, n_tokens} ++ cb(logits, "inkling_moe_logits", il); ++ ++ const size_t lsz = ggml_element_size(logits); ++ ++ ggml_tensor * routed = ggml_cont(ctx0, ggml_view_2d(ctx0, logits, n_expert, n_tokens, logits->nb[1], 0)); ++ ggml_tensor * shared_logits = ggml_view_2d(ctx0, logits, n_shexp, n_tokens, logits->nb[1], n_expert*lsz); ++ ++ // bias affects selection only, not the weights ++ ggml_tensor * scores = ggml_sigmoid(ctx0, routed); ++ scores = ggml_add(ctx0, scores, layer.ffn_exp_probs_b); ++ cb(scores, "inkling_moe_scores", il); ++ ++ ggml_tensor * selected = ggml_argsort_top_k(ctx0, scores, n_expert_used); // I32 {n_expert_used, n_tokens} ++ cb(selected, "inkling_moe_topk", il); ++ ++ // weights use the raw top-k logits, not the biased scores ++ ggml_tensor * routed3 = ggml_reshape_3d(ctx0, routed, 1, n_expert, n_tokens); ++ ggml_tensor * topk_logits = ggml_get_rows(ctx0, routed3, selected); // {1, n_expert_used, n_tokens} ++ topk_logits = ggml_reshape_2d(ctx0, topk_logits, n_expert_used, n_tokens); ++ ++ ggml_tensor * all_logits = ggml_concat(ctx0, topk_logits, shared_logits, 0); // {n_expert_used + n_shexp, n_tokens} ++ ++ // logsigmoid(x) = -softplus(-x) ++ ggml_tensor * w = ggml_neg(ctx0, ggml_softplus(ctx0, ggml_neg(ctx0, all_logits))); ++ w = ggml_soft_max(ctx0, w); ++ w = ggml_scale(ctx0, w, hparams.expert_weights_scale); ++ w = ggml_mul(ctx0, w, layer.ffn_gscale); // gate global_scale (F32 [1]) ++ cb(w, "inkling_moe_weights", il); ++ ++ const size_t wsz = ggml_element_size(w); ++ ++ ggml_tensor * weights = ggml_cont(ctx0, ggml_view_2d(ctx0, w, n_expert_used, n_tokens, w->nb[1], 0)); ++ weights = ggml_reshape_3d(ctx0, weights, 1, n_expert_used, n_tokens); ++ ++ ggml_tensor * xr = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens); ++ ggml_tensor * gate = build_lora_mm_id(layer.ffn_gate_exps, xr, selected); // {n_ff_exp, n_expert_used, n_tokens} ++ ggml_tensor * up = build_lora_mm_id(layer.ffn_up_exps, xr, selected); ++ ggml_tensor * h = ggml_swiglu_split(ctx0, gate, up); ++ ++ ggml_tensor * experts = build_lora_mm_id(layer.ffn_down_exps, h, selected); // {n_embd, n_expert_used, n_tokens} ++ experts = ggml_mul(ctx0, experts, weights); ++ ++ ggml_tensor * moe_out = nullptr; ++ for (int64_t i = 0; i < n_expert_used; ++i) { ++ ggml_tensor * e = ggml_view_2d(ctx0, experts, n_embd, n_tokens, experts->nb[2], i*experts->nb[1]); ++ moe_out = moe_out ? ggml_add(ctx0, moe_out, e) : e; ++ } ++ ++ // shared experts: mul_mat_id with constant ids (never 2D-view a quantized/repacked weight) ++ GGML_ASSERT(shexp_idx != nullptr); ++ ggml_tensor * gs = build_lora_mm_id(layer.ffn_gate_shexp, xr, shexp_idx); // {n_ff_exp, n_shexp, n_tokens} ++ ggml_tensor * us = build_lora_mm_id(layer.ffn_up_shexp, xr, shexp_idx); ++ ggml_tensor * hs = ggml_swiglu_split(ctx0, gs, us); ++ ++ // gammas (last n_shexp weight rows) must scale hs BEFORE the down-proj to match reference rounding in bf16/quant ++ ggml_tensor * gammas = ggml_cont(ctx0, ggml_view_2d(ctx0, w, n_shexp, n_tokens, w->nb[1], n_expert_used*wsz)); ++ hs = ggml_mul(ctx0, hs, ggml_reshape_3d(ctx0, gammas, 1, n_shexp, n_tokens)); ++ ggml_tensor * ds = build_lora_mm_id(layer.ffn_down_shexp, hs, shexp_idx); // {n_embd, n_shexp, n_tokens} ++ ++ for (int64_t s = 0; s < n_shexp; ++s) { ++ ggml_tensor * e = ggml_view_2d(ctx0, ds, n_embd, n_tokens, ds->nb[2], s*ds->nb[1]); ++ moe_out = ggml_add(ctx0, moe_out, e); ++ } ++ cb(moe_out, "inkling_moe_out", il); ++ ++ return moe_out; ++ }; ++ ++ ggml_tensor * cur = build_inp_embd(model.tok_embd); ++ // mtmd embd rows arrive pre-normalized; embed_norm applies to text token lookups only ++ if (ubatch.token) { ++ cur = build_norm(cur, model.tok_norm, NULL, LLM_NORM_RMS, -1); ++ cb(cur, "inkling_embd_norm", -1); ++ } else { ++ cb(cur, "inkling_mm_embd", -1); ++ } ++ ++ ggml_build_forward_expand(gf, cur); ++ ++ for (int il = 0; il < n_layer; ++il) { ++ conv_rs_cur = build_rs(inp_hybrid->get_recr(), mctx_recr->get_r_l(il), n_embd_r, n_seqs); ++ ++ // h += attn_sconv(attn(attn_norm(h))) ++ ggml_tensor * attn_in = build_norm(cur, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); ++ cb(attn_in, "inkling_attn_norm", il); ++ ggml_tensor * attn_out = build_attn_block(attn_in, il); ++ attn_out = build_sconv(attn_out, model.layers[il].shortconv_attn, off_attn, il); ++ cb(attn_out, "inkling_attn_sconv", il); ++ ++ cur = ggml_add(ctx0, cur, attn_out); ++ ++ // h += mlp_sconv(mlp(mlp_norm(h))) ++ ggml_tensor * ffn_in = build_norm(cur, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); ++ cb(ffn_in, "inkling_ffn_norm", il); ++ ggml_tensor * ffn_out = il < (int) hparams.n_layer_dense_lead ? ++ build_dense_ffn(ffn_in, il) : build_moe(ffn_in, il); ++ ffn_out = build_sconv(ffn_out, model.layers[il].shortconv_mlp, off_mlp, il); ++ cb(ffn_out, "inkling_ffn_sconv", il); ++ ++ cur = ggml_add(ctx0, cur, ffn_out); ++ ++ cur = build_cvec(cur, il); ++ cb(cur, "l_out", il); ++ } ++ ++ // conv states need every layer to see ALL tokens, so trim outputs only after the full stack ++ ggml_tensor * inp_out_ids = build_inp_out_ids(); ++ if (inp_out_ids) { ++ cur = ggml_get_rows(ctx0, cur, inp_out_ids); ++ } ++ ++ cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); ++ cb(cur, "result_norm", -1); ++ res->t_embd = cur; ++ ++ if (!cparams.embeddings) { ++ cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); ++ cur = build_lora_mm( ++ model.output, cur, nullptr, ++ model.output->type == GGML_TYPE_F32 ? GGML_PREC_F32_PEDANTIC : GGML_PREC_DEFAULT); ++ ++ // padded vocab rows get -inf so samplers never emit a padded id ++ if (vocab_mask) { ++ cur = ggml_add(ctx0, cur, vocab_mask); ++ } ++ cb(cur, "result_output", -1); ++ res->t_logits = cur; ++ } ++ ++ ggml_build_forward_expand(gf, cur); ++} +diff --git a/src/models/models.h b/src/models/models.h +index aeeb7b222..8a14222b2 100644 +--- a/src/models/models.h ++++ b/src/models/models.h +@@ -1874,6 +1874,19 @@ struct llama_model_lfm2moe : public llama_model_base { + }; + + ++struct llama_model_inkling : public llama_model_base { ++ llama_model_inkling(const struct llama_model_params & params) : llama_model_base(params) {} ++ void load_arch_hparams(llama_model_loader & ml) override; ++ void load_arch_tensors(llama_model_loader & ml) override; ++ ++ struct graph : public llm_graph_context { ++ graph(const llama_model & model, const llm_graph_params & params); ++ }; ++ ++ std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; ++}; ++ ++ + struct llama_model_smallthinker : public llama_model_base { + llama_model_smallthinker(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; +diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt +index 158ab28df..794c0f047 100644 +--- a/tests/CMakeLists.txt ++++ b/tests/CMakeLists.txt +@@ -266,6 +266,14 @@ if (NOT LLAMA_SANITIZE_ADDRESS AND NOT GGML_SCHED_NO_REALLOC) + endif() + llama_build_and_test(test-backend-ops.cpp) + ++add_executable(test-flash-attn-bias test-flash-attn-bias.cpp) ++target_link_libraries(test-flash-attn-bias PRIVATE ggml) ++add_test(NAME test-flash-attn-bias COMMAND test-flash-attn-bias) ++ ++add_executable(test-flash-attn-generic-hash test-flash-attn-generic-hash.cpp) ++target_link_libraries(test-flash-attn-generic-hash PRIVATE ggml) ++add_test(NAME test-flash-attn-generic-hash COMMAND test-flash-attn-generic-hash) ++ + llama_build_and_test(test-model-load-cancel.cpp LABEL "model") + llama_build_and_test(test-autorelease.cpp LABEL "model") + llama_build_and_test(test-backend-sampler.cpp LABEL "model") +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index e9e926e87..9031e71e9 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -12407,6 +12407,103 @@ struct test_glm_dsa_compact_k_gather : public test_case { + } + }; + ++ ++// GGML_OP_FLASH_ATTN_EXT_BANDED ++struct test_flash_attn_ext_banded : public test_case { ++ const int64_t d; ++ const int64_t n_head; ++ const int64_t n_head_kv; ++ const int64_t n_q; ++ const int64_t n_kv; ++ const int64_t rel_extent; ++ const int mask_kind; // 0: none, 1: causal, 2: causal sliding window with dist < rel_extent ++ const ggml_type kv_type; ++ const ggml_type rel_type; ++ const bool strided; ++ ++ std::string vars() override { ++ return VARS_TO_STR10(d, n_head, n_head_kv, n_q, n_kv, rel_extent, mask_kind, kv_type, rel_type, strided); ++ } ++ ++ double max_nmse_err() override { ++ if (kv_type == GGML_TYPE_F32 && rel_type == GGML_TYPE_F32) { ++ return 2e-6; ++ } ++ // fp16 VKQ accumulation error grows with the KV length (plus periodic accumulator ++ // rescales guarding against fp16 overflow) ++ return n_kv > 8192 ? 1e-3 : 5e-4; ++ } ++ ++ uint64_t op_flops(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return 4*n_head*n_q*n_kv*d; ++ } ++ ++ test_flash_attn_ext_banded( ++ int64_t d, int64_t n_head, int64_t n_head_kv, ++ int64_t n_q, int64_t n_kv, int64_t rel_extent, ++ int mask_kind, ggml_type kv_type, ggml_type rel_type, bool strided = false) ++ : d(d), n_head(n_head), n_head_kv(n_head_kv), n_q(n_q), n_kv(n_kv), ++ rel_extent(rel_extent), mask_kind(mask_kind), kv_type(kv_type), rel_type(rel_type), strided(strided) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d, n_q, n_head, 1); ++ ggml_tensor * k; ++ ggml_tensor * v; ++ ggml_tensor * r; ++ if (strided) { ++ // gaps between rows/heads force the kernels to use the 64-bit byte strides ++ ggml_tensor * kb = ggml_new_tensor_4d(ctx, kv_type, 2*d, n_kv, n_head_kv, 1); ++ ggml_tensor * vb = ggml_new_tensor_4d(ctx, kv_type, 2*d, n_kv, n_head_kv, 1); ++ ggml_tensor * rb = ggml_new_tensor_4d(ctx, rel_type, 2*rel_extent, n_head, n_q, 1); ++ k = ggml_view_4d(ctx, kb, d, n_kv, n_head_kv, 1, kb->nb[1], kb->nb[2], kb->nb[3], 0); ++ v = ggml_view_4d(ctx, vb, d, n_kv, n_head_kv, 1, vb->nb[1], vb->nb[2], vb->nb[3], 0); ++ r = ggml_view_4d(ctx, rb, rel_extent, n_head, n_q, 1, rb->nb[1], rb->nb[2], rb->nb[3], 0); ++ } else { ++ k = ggml_new_tensor_4d(ctx, kv_type, d, n_kv, n_head_kv, 1); ++ v = ggml_new_tensor_4d(ctx, kv_type, d, n_kv, n_head_kv, 1); ++ r = ggml_new_tensor_4d(ctx, rel_type, rel_extent, n_head, n_q, 1); ++ } ++ ggml_set_name(q, "q"); ++ ggml_set_name(k, "k"); ++ ggml_set_name(v, "v"); ++ ggml_set_name(r, "rel_logits"); ++ ++ ggml_tensor * m = nullptr; ++ if (mask_kind != 0) { ++ m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_q, 1, 1); ++ ggml_set_name(m, "m"); ++ } ++ ++ ggml_tensor * out = ggml_flash_attn_ext_banded(ctx, q, k, v, m, r, 1.0f/float(d), rel_extent); ++ ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); ++ ggml_set_name(out, "out"); ++ return out; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "m") == 0) { ++ std::vector data(n_q*n_kv); ++ for (int64_t iq = 0; iq < n_q; ++iq) { ++ for (int64_t ik = 0; ik < n_kv; ++ik) { ++ const int64_t rel_dist = iq + (n_kv - n_q) - ik; ++ const bool visible = rel_dist >= 0 && (mask_kind == 1 || rel_dist < rel_extent); ++ data[iq*n_kv + ik] = ggml_fp32_to_fp16(visible ? 0.0f : -INFINITY); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(data[0])); ++ } else if (strcmp(t->name, "rel_logits") == 0) { ++ // A larger range makes rel_dist = E versus E-1 mistakes immediately visible. ++ init_tensor_uniform(t, -1.0f, 1.0f); ++ } else { ++ init_tensor_uniform(t, -0.25f, 0.25f); ++ } ++ } ++ } ++}; ++ ++ + // GGML_OP_CROSS_ENTROPY_LOSS + struct test_cross_entropy_loss : public test_case { + const ggml_type type; +@@ -12986,6 +13083,7 @@ struct test_generic_op : public test_case { + case GGML_OP_RWKV_WKV7: + return 5e-3; + case GGML_OP_FLASH_ATTN_EXT: ++ case GGML_OP_FLASH_ATTN_EXT_BANDED: + { + // Scale error with kv length to account for accumulating floating point error + const int64_t kv = sources[1].ne[1]; +@@ -13009,7 +13107,7 @@ struct test_generic_op : public test_case { + } + + // FLASH_ATTN_EXT: src[3] is the KQ mask +- if (op == GGML_OP_FLASH_ATTN_EXT && i == 3) { ++ if ((op == GGML_OP_FLASH_ATTN_EXT || op == GGML_OP_FLASH_ATTN_EXT_BANDED) && i == 3) { + init_tensor_kq_mask(t); + continue; + } +@@ -15062,6 +15160,22 @@ static std::vector> make_test_cases_eval() { + test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q1_0)); + test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q1_0, GGML_TYPE_F16)); + ++ // banded score-bias coverage: band edges, masks, decode offset, GQA, head sizes, table types ++ test_cases.emplace_back(new test_flash_attn_ext_banded( 64, 2, 1, 8, 8, 8, 1, GGML_TYPE_F32, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 2, 16, 16, 8, 1, GGML_TYPE_F16, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded( 64, 8, 2, 1, 64, 8, 1, GGML_TYPE_F32, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 2, 64, 64, 8, 2, GGML_TYPE_BF16, GGML_TYPE_BF16)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded( 64, 8, 1, 64, 64, 512, 1, GGML_TYPE_F16, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 2, 17, 33, 8, 1, GGML_TYPE_F16, GGML_TYPE_F16, true)); ++ // production-scale n_kv straddling the observed ~16.4-16.9K garbage threshold ++ test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 8192, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 16384, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 16403, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 16896, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 17024, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 1, 17024, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 32768, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); ++ + // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix + // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). + for (int64_t kv : { 4096, 16384 }) { +diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp +index 01b07953a..d0f985049 100644 +--- a/tests/test-chat.cpp ++++ b/tests/test-chat.cpp +@@ -2996,6 +2996,85 @@ static void test_template_output_peg_parsers(bool detailed_debug) { + .run(); + } + ++ { ++ // Inkling / TML typed content blocks: <|end_message|> separates blocks, <|content_model_end_sampling|> ends the turn. ++ auto tst = peg_tester("models/templates/Inkling.jinja", detailed_debug); ++ ++ // reasoning + visible answer ++ tst.test("<|content_thinking|>I'm\nthinking<|end_message|>" ++ "<|message_model|><|content_text|>Hello, world!\nWhat's up?<|end_message|>" ++ "<|content_model_end_sampling|>") ++ .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) ++ .expect(message_assist_thoughts) ++ .run(); ++ ++ // Visible answer only (reasoning_effort=0 -> no thinking block). ++ tst.test("<|content_text|>Hello, world!\nWhat's up?<|end_message|>" ++ "<|content_model_end_sampling|>") ++ .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) ++ .expect(message_assist) ++ .run(); ++ ++ // Empty thinking block, then the answer. ++ tst.test("<|content_thinking|><|end_message|>" ++ "<|message_model|><|content_text|>Hello, world!\nWhat's up?<|end_message|>" ++ "<|content_model_end_sampling|>") ++ .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) ++ .expect(message_assist) ++ .run(); ++ ++ // single tool call with reasoning; the bare tool-name echo and role opener are dropped ++ tst.test("<|content_thinking|>I'm\nthinking<|end_message|>" ++ "<|message_model|>special_function<|content_invoke_tool_json|>" ++ "{\"name\":\"special_function\",\"args\":{\"arg1\":1}}<|end_message|>" ++ "<|content_model_end_sampling|>") ++ .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) ++ .tools({ special_function_tool }) ++ .expect(message_with_reasoning_and_tool_call("I'm\nthinking", "special_function", "{\"arg1\": 1}")) ++ .run(); ++ ++ // Tool call, no reasoning. ++ tst.test("<|message_model|>special_function<|content_invoke_tool_json|>" ++ "{\"name\":\"special_function\",\"args\":{\"arg1\":1}}<|end_message|>" ++ "<|content_model_end_sampling|>") ++ .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) ++ .tools({ special_function_tool }) ++ .expect(message_assist_call) ++ .run(); ++ ++ // regression: the tool branch must not swallow a pure-text answer ++ tst.test("<|content_thinking|>I'm\nthinking<|end_message|>" ++ "<|message_model|><|content_text|>Hello, world!\nWhat's up?<|end_message|>" ++ "<|content_model_end_sampling|>") ++ .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) ++ .tools({ special_function_tool }) ++ .expect(message_assist_thoughts) ++ .run(); ++ ++ // tools available, content only, no thinking (reasoning_effort=0 leak scenario) ++ tst.test("<|content_text|>Hello, world!\nWhat's up?<|end_message|>" ++ "<|content_model_end_sampling|>") ++ .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) ++ .tools({ special_function_tool }) ++ .expect(message_assist) ++ .run(); ++ ++ // parallel tool calls are separate marker-wrapped blocks ++ tst.test("<|message_model|>special_function<|content_invoke_tool_json|>" ++ "{\"name\":\"special_function\",\"args\":{\"arg1\":1}}<|end_message|>" ++ "<|message_model|>python<|content_invoke_tool_json|>" ++ "{\"name\":\"python\",\"args\":{\"code\":\"print('hey')\"}}<|end_message|>" ++ "<|content_model_end_sampling|>") ++ .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) ++ .parallel_tool_calls(true) ++ .tools({ special_function_tool, python_tool }) ++ .expect_tool_calls({ ++ { "special_function", R"({"arg1": 1})", "" }, ++ { "python", "{\"code\": \"print('hey')\"}", "" }, ++ }) ++ .run(); ++ } ++ + { + // Google Gemma 2 2B - does not support tool calling + auto tst = peg_tester("models/templates/google-gemma-2-2b-it.jinja"); +diff --git a/tests/test-flash-attn-bias.cpp b/tests/test-flash-attn-bias.cpp +new file mode 100644 +index 000000000..94232b32d +--- /dev/null ++++ b/tests/test-flash-attn-bias.cpp +@@ -0,0 +1,525 @@ ++#include "ggml.h" ++#include "ggml-alloc.h" ++#include "ggml-backend.h" ++#include "ggml-cpp.h" ++ ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++ ++struct bias_test_config { ++ const char * name; ++ int64_t d; ++ int64_t nq; ++ int64_t nkv; ++ int64_t hq; ++ int64_t hkv; ++ int64_t extent; ++ ggml_type type; ++ bool use_mask; ++ bool sliding; ++ int64_t n_batch = 1; ++ int64_t rel_batch = 1; ++ bool strided_rel = false; ++}; ++ ++struct test_data { ++ std::vector q; ++ std::vector k; ++ std::vector v; ++ std::vector rel; ++ std::vector mask; ++ std::vector k_typed; ++ std::vector v_typed; ++ std::vector rel_typed; ++ std::vector k_rounded; ++ std::vector v_rounded; ++ std::vector rel_rounded; ++ std::vector dense_bias; ++}; ++ ++struct run_result { ++ std::vector output; ++ size_t allocated_bytes; ++ double ms; ++}; ++ ++static std::vector convert_type(ggml_type type, const std::vector & src, std::vector & rounded) { ++ rounded.resize(src.size()); ++ if (type == GGML_TYPE_F32) { ++ rounded = src; ++ std::vector bytes(src.size()*sizeof(float)); ++ memcpy(bytes.data(), src.data(), bytes.size()); ++ return bytes; ++ } ++ if (type == GGML_TYPE_F16) { ++ std::vector tmp(src.size()); ++ ggml_fp32_to_fp16_row(src.data(), tmp.data(), src.size()); ++ ggml_fp16_to_fp32_row(tmp.data(), rounded.data(), src.size()); ++ std::vector bytes(tmp.size()*sizeof(tmp[0])); ++ memcpy(bytes.data(), tmp.data(), bytes.size()); ++ return bytes; ++ } ++ GGML_ASSERT(type == GGML_TYPE_BF16); ++ std::vector tmp(src.size()); ++ ggml_fp32_to_bf16_row_ref(src.data(), tmp.data(), src.size()); ++ ggml_bf16_to_fp32_row(tmp.data(), rounded.data(), src.size()); ++ std::vector bytes(tmp.size()*sizeof(tmp[0])); ++ memcpy(bytes.data(), tmp.data(), bytes.size()); ++ return bytes; ++} ++ ++static test_data make_data(const bias_test_config & c) { ++ test_data data; ++ data.q.resize(c.d*c.nq*c.hq*c.n_batch); ++ data.k.resize(c.d*c.nkv*c.hkv*c.n_batch); ++ data.v.resize(c.d*c.nkv*c.hkv*c.n_batch); ++ data.rel.resize(c.extent*c.hq*c.nq*c.rel_batch); ++ data.mask.resize(c.nkv*c.nq); ++ ++ for (size_t i = 0; i < data.q.size(); ++i) { ++ data.q[i] = 0.20f*std::sin(float(i)*0.017f + 0.13f); ++ } ++ for (size_t i = 0; i < data.k.size(); ++i) { ++ data.k[i] = 0.25f*std::cos(float(i)*0.013f - 0.29f); ++ data.v[i] = 0.30f*std::sin(float(i)*0.019f + 0.71f); ++ } ++ for (int64_t ib = 0; ib < c.rel_batch; ++ib) { ++ for (int64_t iq = 0; iq < c.nq; ++iq) { ++ for (int64_t ih = 0; ih < c.hq; ++ih) { ++ for (int64_t ie = 0; ie < c.extent; ++ie) { ++ const size_t idx = ((ib*c.nq + iq)*c.hq + ih)*c.extent + ie; ++ data.rel[idx] = 0.75f*std::sin(float(idx)*0.007f + float(ie)*0.021f + 0.31f); ++ } ++ } ++ } ++ } ++ for (int64_t iq = 0; iq < c.nq; ++iq) { ++ for (int64_t ik = 0; ik < c.nkv; ++ik) { ++ const int64_t dist = iq + (c.nkv - c.nq) - ik; ++ const bool visible = !c.use_mask || (dist >= 0 && (!c.sliding || dist < c.extent)); ++ data.mask[iq*c.nkv + ik] = ggml_fp32_to_fp16(visible ? 0.0f : -INFINITY); ++ } ++ } ++ ++ data.k_typed = convert_type(c.type, data.k, data.k_rounded); ++ data.v_typed = convert_type(c.type, data.v, data.v_rounded); ++ data.rel_typed = convert_type(GGML_TYPE_F32, data.rel, data.rel_rounded); ++ ++ data.dense_bias.assign(c.nkv*c.nq*c.hq*c.n_batch, 0.0f); ++ for (int64_t ib = 0; ib < c.n_batch; ++ib) { ++ const int64_t irb = ib % c.rel_batch; ++ for (int64_t ih = 0; ih < c.hq; ++ih) { ++ for (int64_t iq = 0; iq < c.nq; ++iq) { ++ for (int64_t ik = 0; ik < c.nkv; ++ik) { ++ const int64_t dist = iq + (c.nkv - c.nq) - ik; ++ if (dist >= 0 && dist < c.extent) { ++ data.dense_bias[((ib*c.hq + ih)*c.nq + iq)*c.nkv + ik] = ++ data.rel_rounded[((irb*c.nq + iq)*c.hq + ih)*c.extent + dist]; ++ } ++ } ++ } ++ } ++ } ++ return data; ++} ++ ++static run_result run_graph( ++ ggml_backend_t backend, ++ const bias_test_config & c, ++ const test_data & data, ++ bool dense, ++ int repeats) { ++ ggml_init_params params = { ++ /* .mem_size = */ ggml_tensor_overhead()*64 + ggml_graph_overhead_custom(64, false), ++ /* .mem_base = */ nullptr, ++ /* .no_alloc = */ true, ++ }; ++ ggml_context_ptr ctx(ggml_init(params)); ++ GGML_ASSERT(ctx); ++ ++ ggml_tensor * q = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, c.d, c.nq, c.hq, c.n_batch); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx.get(), c.type, c.d, c.nkv, c.hkv, c.n_batch); ++ ggml_tensor * v = ggml_new_tensor_4d(ctx.get(), c.type, c.d, c.nkv, c.hkv, c.n_batch); ++ ggml_tensor * r_storage = nullptr; ++ ggml_tensor * r; ++ if (c.strided_rel) { ++ r_storage = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, ++ 2*c.extent, c.hq, c.nq, c.rel_batch); ++ r = ggml_view_4d(ctx.get(), r_storage, c.extent, c.hq, c.nq, c.rel_batch, ++ r_storage->nb[1], r_storage->nb[2], r_storage->nb[3], 0); ++ } else { ++ r = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, ++ c.extent, c.hq, c.nq, c.rel_batch); ++ } ++ ggml_tensor * m = c.use_mask ? ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, c.nkv, c.nq, 1, 1) : nullptr; ++ ggml_set_name(q, "q"); ++ ggml_set_name(k, "k"); ++ ggml_set_name(v, "v"); ++ ggml_set_name(r, "rel_logits"); ++ if (m) { ++ ggml_set_name(m, "mask"); ++ } ++ ++ ggml_tensor * out; ++ ggml_tensor * bias = nullptr; ++ if (!dense) { ++ out = ggml_flash_attn_ext_banded(ctx.get(), q, k, v, m, r, 1.0f/float(c.d), c.extent); ++ } else { ++ bias = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, c.nkv, c.nq, c.hq, c.n_batch); ++ ggml_set_name(bias, "dense_bias"); ++ ggml_tensor * scores = ggml_mul_mat(ctx.get(), k, q); ++ ggml_mul_mat_set_prec(scores, GGML_PREC_F32); ++ scores = ggml_scale(ctx.get(), scores, 1.0f/float(c.d)); ++ scores = ggml_add(ctx.get(), scores, bias); ++ scores = ggml_soft_max_ext(ctx.get(), scores, m, 1.0f, 0.0f); ++ ggml_tensor * vt = ggml_cont(ctx.get(), ggml_transpose(ctx.get(), v)); ++ out = ggml_mul_mat(ctx.get(), vt, scores); ++ ggml_mul_mat_set_prec(out, GGML_PREC_F32); ++ out = ggml_cont(ctx.get(), ggml_permute(ctx.get(), out, 0, 2, 1, 3)); ++ } ++ ggml_set_name(out, dense ? "out_dense" : "out_flash"); ++ ++ GGML_ASSERT(ggml_backend_supports_op(backend, out)); ++ ggml_backend_buffer_ptr buffer(ggml_backend_alloc_ctx_tensors(ctx.get(), backend)); ++ GGML_ASSERT(buffer); ++ ++ ggml_backend_tensor_set(q, data.q.data(), 0, data.q.size()*sizeof(float)); ++ ggml_backend_tensor_set(k, data.k_typed.data(), 0, data.k_typed.size()); ++ ggml_backend_tensor_set(v, data.v_typed.data(), 0, data.v_typed.size()); ++ if (r_storage) { ++ std::vector physical(2*c.extent*c.hq*c.nq*c.rel_batch, 0.0f); ++ for (int64_t ib = 0; ib < c.rel_batch; ++ib) { ++ for (int64_t iq = 0; iq < c.nq; ++iq) { ++ for (int64_t ih = 0; ih < c.hq; ++ih) { ++ const size_t logical = ((ib*c.nq + iq)*c.hq + ih)*c.extent; ++ const size_t storage = ((ib*c.nq + iq)*c.hq + ih)*(2*c.extent); ++ memcpy(physical.data() + storage, data.rel_rounded.data() + logical, ++ c.extent*sizeof(float)); ++ } ++ } ++ } ++ ggml_backend_tensor_set(r_storage, physical.data(), 0, physical.size()*sizeof(float)); ++ } else { ++ ggml_backend_tensor_set(r, data.rel_typed.data(), 0, data.rel_typed.size()); ++ } ++ if (m) { ++ ggml_backend_tensor_set(m, data.mask.data(), 0, data.mask.size()*sizeof(data.mask[0])); ++ } ++ if (bias) { ++ ggml_backend_tensor_set(bias, data.dense_bias.data(), 0, data.dense_bias.size()*sizeof(float)); ++ } ++ ++ ggml_cgraph * graph = ggml_new_graph_custom(ctx.get(), 64, false); ++ ggml_build_forward_expand(graph, out); ++ GGML_ASSERT(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); ++ ggml_backend_synchronize(backend); ++ ++ const int64_t start = ggml_time_us(); ++ for (int i = 0; i < repeats; ++i) { ++ GGML_ASSERT(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); ++ } ++ ggml_backend_synchronize(backend); ++ const int64_t elapsed = ggml_time_us() - start; ++ ++ run_result result; ++ result.output.resize(ggml_nelements(out)); ++ ggml_backend_tensor_get(out, result.output.data(), 0, result.output.size()*sizeof(float)); ++ result.allocated_bytes = ggml_backend_buffer_get_size(buffer.get()); ++ result.ms = double(elapsed)/1000.0/repeats; ++ return result; ++} ++ ++static std::vector naive_materialized(const bias_test_config & c, const test_data & data) { ++ const int64_t nrows = c.n_batch*c.hq*c.nq; ++ std::vector scores(nrows*c.nkv); ++ std::vector output(c.d*c.hq*c.nq*c.n_batch, 0.0f); ++ std::atomic next_row(0); ++ const unsigned nt = std::max(1u, std::thread::hardware_concurrency()); ++ std::vector workers; ++ workers.reserve(nt); ++ ++ for (unsigned it = 0; it < nt; ++it) { ++ workers.emplace_back([&]() { ++ while (true) { ++ const int64_t row = next_row.fetch_add(1); ++ if (row >= nrows) { ++ break; ++ } ++ const int64_t ib = row / (c.hq*c.nq); ++ const int64_t ih = (row / c.nq) % c.hq; ++ const int64_t iq = row % c.nq; ++ const int64_t ihkv = ih / (c.hq/c.hkv); ++ float row_max = -INFINITY; ++ for (int64_t ik = 0; ik < c.nkv; ++ik) { ++ float dot = 0.0f; ++ for (int64_t id = 0; id < c.d; ++id) { ++ dot += data.q[((ib*c.hq + ih)*c.nq + iq)*c.d + id] * ++ data.k_rounded[((ib*c.hkv + ihkv)*c.nkv + ik)*c.d + id]; ++ } ++ const float mask = ggml_fp16_to_fp32(data.mask[iq*c.nkv + ik]); ++ const float score = dot/float(c.d) + data.dense_bias[row*c.nkv + ik] + mask; ++ scores[row*c.nkv + ik] = score; ++ row_max = std::max(row_max, score); ++ } ++ float sum = 0.0f; ++ for (int64_t ik = 0; ik < c.nkv; ++ik) { ++ const float p = std::exp(scores[row*c.nkv + ik] - row_max); ++ scores[row*c.nkv + ik] = p; ++ sum += p; ++ } ++ for (int64_t ik = 0; ik < c.nkv; ++ik) { ++ const float p = scores[row*c.nkv + ik]/sum; ++ for (int64_t id = 0; id < c.d; ++id) { ++ output[((ib*c.nq + iq)*c.hq + ih)*c.d + id] += ++ p*data.v_rounded[((ib*c.hkv + ihkv)*c.nkv + ik)*c.d + id]; ++ } ++ } ++ } ++ }); ++ } ++ for (auto & worker : workers) { ++ worker.join(); ++ } ++ return output; ++} ++ ++static void error_stats(const std::vector & got, const std::vector & ref, ++ double & max_abs, double & mean_abs, double & max_rel, double & mean_rel, double & rmse) { ++ double sq = 0.0; ++ double abs_sum = 0.0; ++ double rel_sum = 0.0; ++ max_abs = 0.0; ++ max_rel = 0.0; ++ for (size_t i = 0; i < got.size(); ++i) { ++ const double ae = std::abs(double(got[i]) - ref[i]); ++ const double re = ae/std::max(1e-5, std::abs(double(ref[i]))); ++ max_abs = std::max(max_abs, ae); ++ max_rel = std::max(max_rel, re); ++ abs_sum += ae; ++ rel_sum += re; ++ sq += ae*ae; ++ } ++ mean_abs = abs_sum/got.size(); ++ mean_rel = rel_sum/got.size(); ++ rmse = std::sqrt(sq/got.size()); ++} ++ ++static void overflow_arithmetic_self_test() { ++ // mirrors the scalar-path offset math; exact offset checked at 128 bits, lands past 2^31 ++ const uint64_t nb0 = sizeof(float); ++ const uint64_t nb1 = 1024*nb0; ++ const uint64_t nb2 = 64*nb1; ++ const uint64_t nb3 = 131072*nb2; ++ const uint64_t dist = 1023, head = 63, query = 131071, batch = 3; ++ const uint64_t offset = dist*nb0 + head*nb1 + query*nb2 + batch*nb3; ++ __extension__ typedef unsigned __int128 uint128_t; ++ const uint128_t exact = uint128_t(dist)*nb0 + uint128_t(head)*nb1 + ++ uint128_t(query)*nb2 + uint128_t(batch)*nb3; ++ GGML_ASSERT(exact <= UINT64_MAX && offset == (uint64_t) exact && offset > INT32_MAX); ++ ++ const int64_t nq = int64_t(1) << 40; ++ const int64_t nkv = nq + 8192; ++ const int64_t iq = nq - 1; ++ const int64_t ik = nkv - 1024; ++ const int64_t rel_dist = iq + (nkv - nq) - ik; ++ GGML_ASSERT(rel_dist == 1023); ++ printf("overflow_check offset=%llu (>INT32_MAX) large_T=%lld rel_dist=%lld PASS\n", ++ (unsigned long long) offset, (long long) nq, (long long) rel_dist); ++} ++ ++static bool overflow_kernel_test(ggml_backend_t backend, const char * backend_kind) { ++ // rel-logits row for query 1 sits beyond INT32_MAX; only two small logical rows are touched ++ const bias_test_config c = { ++ "overflow_kernel_stride", 64, 2, 2, 2, 1, 8, GGML_TYPE_F32, true, false, ++ }; ++ test_data data = make_data(c); ++ const uint64_t rel_nb2 = (UINT64_C(1) << 31) + 4096; ++ const size_t rel_row_bytes = c.extent*c.hq*sizeof(float); ++ const uint64_t storage_bytes = rel_nb2 + rel_row_bytes; ++ ++ ggml_init_params params = { ++ /* .mem_size = */ ggml_tensor_overhead()*32 + ggml_graph_overhead_custom(32, false), ++ /* .mem_base = */ nullptr, ++ /* .no_alloc = */ true, ++ }; ++ ggml_context_ptr ctx(ggml_init(params)); ++ GGML_ASSERT(ctx); ++ ggml_tensor * q = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, c.d, c.nq, c.hq, 1); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, c.d, c.nkv, c.hkv, 1); ++ ggml_tensor * v = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, c.d, c.nkv, c.hkv, 1); ++ ggml_tensor * m = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, c.nkv, c.nq, 1, 1); ++ ggml_tensor * r_storage = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, ++ (storage_bytes + sizeof(float) - 1)/sizeof(float)); ++ ggml_tensor * r = ggml_view_4d(ctx.get(), r_storage, c.extent, c.hq, c.nq, 1, ++ c.extent*sizeof(float), rel_nb2, rel_nb2*c.nq, 0); ++ ggml_tensor * out = ggml_flash_attn_ext_banded( ++ ctx.get(), q, k, v, m, r, 1.0f/float(c.d), c.extent); ++ ggml_backend_buffer_ptr buffer(ggml_backend_alloc_ctx_tensors(ctx.get(), backend)); ++ GGML_ASSERT(buffer); ++ ++ ggml_backend_tensor_set(q, data.q.data(), 0, data.q.size()*sizeof(float)); ++ ggml_backend_tensor_set(k, data.k_typed.data(), 0, data.k_typed.size()); ++ ggml_backend_tensor_set(v, data.v_typed.data(), 0, data.v_typed.size()); ++ ggml_backend_tensor_set(m, data.mask.data(), 0, data.mask.size()*sizeof(data.mask[0])); ++ ggml_backend_tensor_set(r_storage, data.rel_typed.data(), 0, rel_row_bytes); ++ ggml_backend_tensor_set(r_storage, data.rel_typed.data() + rel_row_bytes, rel_nb2, rel_row_bytes); ++ ++ ggml_cgraph * graph = ggml_new_graph_custom(ctx.get(), 32, false); ++ ggml_build_forward_expand(graph, out); ++ GGML_ASSERT(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); ++ ggml_backend_synchronize(backend); ++ std::vector got(ggml_nelements(out)); ++ ggml_backend_tensor_get(out, got.data(), 0, got.size()*sizeof(float)); ++ const std::vector ref = naive_materialized(c, data); ++ double max_abs, mean_abs, max_rel, mean_rel, rmse; ++ error_stats(got, ref, max_abs, mean_abs, max_rel, mean_rel, rmse); ++ const bool pass = max_abs <= 2e-5; ++ printf("overflow_kernel backend=%s rel_query_stride=%llu allocated_bytes=%zu " ++ "naive_max_abs=%.9g naive_mean_abs=%.9g naive_max_rel=%.9g naive_mean_rel=%.9g naive_rmse=%.9g %s\n", ++ backend_kind, (unsigned long long) rel_nb2, ggml_backend_buffer_get_size(buffer.get()), ++ max_abs, mean_abs, max_rel, mean_rel, rmse, pass ? "PASS" : "FAIL"); ++ return pass; ++} ++ ++int main(int argc, char ** argv) { ++ std::string backend_kind = argc > 1 ? argv[1] : "cpu"; ++ std::string suite = argc > 2 ? argv[2] : "small"; ++ int repeats = argc > 3 ? std::max(1, atoi(argv[3])) : 1; ++ ++ overflow_arithmetic_self_test(); ++ ggml_backend_load_all(); ++ ggml_backend_dev_t chosen = nullptr; ++ const enum ggml_backend_dev_type wanted = backend_kind == "cuda" ? ++ GGML_BACKEND_DEVICE_TYPE_GPU : GGML_BACKEND_DEVICE_TYPE_CPU; ++ for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { ++ ggml_backend_dev_t dev = ggml_backend_dev_get(i); ++ if (ggml_backend_dev_type(dev) == wanted) { ++ chosen = dev; ++ break; ++ } ++ } ++ GGML_ASSERT(chosen); ++ ggml_backend_ptr backend(ggml_backend_dev_init(chosen, nullptr)); ++ GGML_ASSERT(backend); ++ ++ if (suite == "overflow") { ++ return overflow_kernel_test(backend.get(), backend_kind.c_str()) ? 0 : 1; ++ } ++ ++ if (suite == "perf") { ++ const std::vector perf_cases = { ++ {"prefill_t1024", 64, 1024, 1024, 8, 2, 512, GGML_TYPE_F16, true, false}, ++ {"prefill_t2048", 64, 2048, 2048, 8, 2, 512, GGML_TYPE_F16, true, false}, ++ {"prefill_t4096", 64, 4096, 4096, 8, 2, 512, GGML_TYPE_F16, true, false}, ++ {"decode_8k", 128, 1, 8192, 8, 1, 1024, GGML_TYPE_F16, true, false}, ++ {"heads64_gqa", 64, 1024, 1024, 64, 8, 512, GGML_TYPE_F16, true, false}, ++ }; ++ printf("backend=%s suite=perf device=%s repeats=%d\n", backend_kind.c_str(), ++ ggml_backend_dev_description(chosen), repeats); ++ for (const bias_test_config & c : perf_cases) { ++ test_data data = make_data(c); ++ const run_result flash = run_graph(backend.get(), c, data, false, repeats); ++ const run_result dense = run_graph(backend.get(), c, data, true, repeats); ++ printf("%s type=%s D=%lld nq=%lld nkv=%lld hq=%lld hkv=%lld E=%lld " ++ "flash_ms=%.6f dense_ms=%.6f speedup=%.6f flash_bytes=%zu dense_bytes=%zu memory_ratio=%.6f\n", ++ c.name, ggml_type_name(c.type), (long long)c.d, (long long)c.nq, ++ (long long)c.nkv, (long long)c.hq, (long long)c.hkv, (long long)c.extent, ++ flash.ms, dense.ms, dense.ms/flash.ms, flash.allocated_bytes, ++ dense.allocated_bytes, double(dense.allocated_bytes)/flash.allocated_bytes); ++ } ++ return 0; ++ } ++ ++ if (suite == "memory") { ++ const std::vector memory_cases = { ++ {"memory_t1024", 64, 1024, 1024, 8, 2, 512, GGML_TYPE_F16, false, false}, ++ {"memory_t2048", 64, 2048, 2048, 8, 2, 512, GGML_TYPE_F16, false, false}, ++ {"memory_t4096", 64, 4096, 4096, 8, 2, 512, GGML_TYPE_F16, false, false}, ++ }; ++ printf("backend=%s suite=memory device=%s\n", backend_kind.c_str(), ++ ggml_backend_dev_description(chosen)); ++ for (const bias_test_config & c : memory_cases) { ++ test_data data = make_data(c); ++ const run_result flash = run_graph(backend.get(), c, data, false, 1); ++ const run_result dense = run_graph(backend.get(), c, data, true, 1); ++ double max_abs, mean_abs, max_rel, mean_rel, rmse; ++ error_stats(flash.output, dense.output, max_abs, mean_abs, max_rel, mean_rel, rmse); ++ printf("%s T=%lld flash_bytes=%zu dense_bytes=%zu memory_ratio=%.6f " ++ "dense_max_abs=%.9g dense_mean_abs=%.9g PASS\n", ++ c.name, (long long)c.nq, flash.allocated_bytes, dense.allocated_bytes, ++ double(dense.allocated_bytes)/flash.allocated_bytes, max_abs, mean_abs); ++ } ++ return 0; ++ } ++ ++ std::vector configs; ++ if (suite == "small") { ++ configs = { ++ {"edge_e8_f32", 64, 16, 16, 2, 1, 8, GGML_TYPE_F32, true, false}, ++ {"gqa_d128_f16", 128, 16, 16, 8, 2, 8, GGML_TYPE_F16, true, false}, ++ {"sliding_bf16", 64, 64, 64, 8, 2, 8, GGML_TYPE_BF16, true, true }, ++ {"decode_offset", 64, 1, 513, 8, 2, 8, GGML_TYPE_F32, true, false}, ++ {"extent_512_edge", 64, 64, 64, 8, 1, 512, GGML_TYPE_F16, true, false}, ++ {"strided_rel_f16", 64, 17, 33, 8, 2, 8, GGML_TYPE_F16, true, false, 1, 1, true}, ++ {"batch_distinct", 64, 16, 16, 8, 2, 8, GGML_TYPE_F16, true, false, 2, 2, false}, ++ {"batch_broadcast", 64, 16, 16, 8, 2, 8, GGML_TYPE_F16, true, false, 2, 1, false}, ++ {"heads64_gqa4", 64, 16, 16, 64, 16, 8, GGML_TYPE_F16, true, false}, ++ {"heads64_gqa8", 64, 16, 16, 64, 8, 8, GGML_TYPE_F16, true, false}, ++ }; ++ } else if (suite == "medium") { ++ configs = { ++ {"medium_f32_e512", 64, 1024, 1024, 8, 2, 512, GGML_TYPE_F32, true, false}, ++ {"medium_f16_e1024", 128, 1024, 1024, 8, 2, 1024, GGML_TYPE_F16, true, false}, ++ {"medium_bf16_local", 64, 2048, 2048, 8, 2, 512, GGML_TYPE_BF16, true, true }, ++ {"decode_8k_e1024", 128, 1, 8192, 8, 1, 1024, GGML_TYPE_F16, true, false}, ++ }; ++ } else if (suite == "hard") { ++ configs = { ++ {"heads64_gqa", 64, 1024, 1024, 64, 8, 512, GGML_TYPE_F16, true, false}, ++ }; ++ } else { ++ fprintf(stderr, "unknown suite: %s (expected small, medium, hard, perf, memory, or overflow)\n", ++ suite.c_str()); ++ return 2; ++ } ++ ++ bool ok = true; ++ printf("backend=%s suite=%s device=%s\n", backend_kind.c_str(), suite.c_str(), ggml_backend_dev_description(chosen)); ++ for (const bias_test_config & c : configs) { ++ test_data data = make_data(c); ++ run_result flash = run_graph(backend.get(), c, data, false, repeats); ++ run_result dense = run_graph(backend.get(), c, data, true, repeats); ++ // always compare to an independently materialized oracle (O(T^2), test-only) ++ const std::vector naive = naive_materialized(c, data); ++ ++ double abs_dense, mean_abs_dense, rel_dense, mean_rel_dense, rmse_dense; ++ error_stats(flash.output, dense.output, abs_dense, mean_abs_dense, rel_dense, mean_rel_dense, rmse_dense); ++ double abs_naive = 0.0, mean_abs_naive = 0.0, rel_naive = 0.0, mean_rel_naive = 0.0, rmse_naive = 0.0; ++ error_stats(flash.output, naive, abs_naive, mean_abs_naive, rel_naive, mean_rel_naive, rmse_naive); ++ const double tol = c.type == GGML_TYPE_F32 ? 2e-5 : 2e-3; ++ const bool pass = abs_naive <= tol; ++ ok = ok && pass; ++ printf("%s type=%s D=%lld nq=%lld nkv=%lld hq=%lld hkv=%lld E=%lld mask=%d sliding=%d " ++ "batch=%lld rel_batch=%lld strided_rel=%d " ++ "dense_max_abs=%.9g dense_mean_abs=%.9g dense_max_rel=%.9g dense_mean_rel=%.9g dense_rmse=%.9g " ++ "naive_max_abs=%.9g naive_mean_abs=%.9g naive_max_rel=%.9g naive_mean_rel=%.9g naive_rmse=%.9g " ++ "flash_ms=%.4f dense_ms=%.4f speedup=%.4f " ++ "flash_bytes=%zu dense_bytes=%zu memory_ratio=%.4f %s\n", ++ c.name, ggml_type_name(c.type), (long long)c.d, (long long)c.nq, (long long)c.nkv, ++ (long long)c.hq, (long long)c.hkv, (long long)c.extent, c.use_mask, c.sliding, ++ (long long)c.n_batch, (long long)c.rel_batch, c.strided_rel, ++ abs_dense, mean_abs_dense, rel_dense, mean_rel_dense, rmse_dense, ++ abs_naive, mean_abs_naive, rel_naive, mean_rel_naive, rmse_naive, ++ flash.ms, dense.ms, dense.ms/flash.ms, ++ flash.allocated_bytes, dense.allocated_bytes, double(dense.allocated_bytes)/flash.allocated_bytes, ++ pass ? "PASS" : "FAIL"); ++ } ++ return ok ? 0 : 1; ++} +diff --git a/tests/test-flash-attn-generic-hash.cpp b/tests/test-flash-attn-generic-hash.cpp +new file mode 100644 +index 000000000..4e3aca44f +--- /dev/null ++++ b/tests/test-flash-attn-generic-hash.cpp +@@ -0,0 +1,142 @@ ++#include "ggml.h" ++#include "ggml-alloc.h" ++#include "ggml-backend.h" ++ ++#include ++#include ++#include ++#include ++#include ++#include ++ ++// Deterministic FLASH_ATTN_EXT probe; banded-API-free so the same source builds on base and final trees. ++int main(int argc, char ** argv) { ++ const std::string backend_kind = argc > 1 ? argv[1] : "cpu"; ++ const std::string output_path = argc > 2 ? argv[2] : ""; ++ ++ ggml_backend_load_all(); ++ const enum ggml_backend_dev_type wanted = backend_kind == "cuda" ? ++ GGML_BACKEND_DEVICE_TYPE_GPU : GGML_BACKEND_DEVICE_TYPE_CPU; ++ ggml_backend_dev_t chosen = nullptr; ++ for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { ++ ggml_backend_dev_t dev = ggml_backend_dev_get(i); ++ if (ggml_backend_dev_type(dev) == wanted) { ++ chosen = dev; ++ break; ++ } ++ } ++ if (!chosen) { ++ fprintf(stderr, "requested backend is unavailable: %s\n", backend_kind.c_str()); ++ return 2; ++ } ++ ggml_backend_t backend = ggml_backend_dev_init(chosen, nullptr); ++ if (!backend) { ++ return 2; ++ } ++ ++ constexpr int64_t d = 64; ++ constexpr int64_t nq = 33; ++ constexpr int64_t nkv = 47; ++ constexpr int64_t hq = 8; ++ constexpr int64_t hkv = 2; ++ ggml_init_params params = { ++ /* .mem_size = */ ggml_tensor_overhead()*16 + ggml_graph_overhead_custom(16, false), ++ /* .mem_base = */ nullptr, ++ /* .no_alloc = */ true, ++ }; ++ ggml_context * ctx = ggml_init(params); ++ if (!ctx) { ++ ggml_backend_free(backend); ++ return 2; ++ } ++ ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d, nq, hq, 1); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, nkv, hkv, 1); ++ ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, nkv, hkv, 1); ++ ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, nkv, nq, 1, 1); ++ ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/float(d), 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); ++ if (!ggml_backend_supports_op(backend, out)) { ++ fprintf(stderr, "ordinary flash attention is unsupported on %s\n", ++ ggml_backend_dev_description(chosen)); ++ ggml_free(ctx); ++ ggml_backend_free(backend); ++ return 2; ++ } ++ ++ ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); ++ if (!buffer) { ++ ggml_free(ctx); ++ ggml_backend_free(backend); ++ return 2; ++ } ++ ++ std::vector q_data(ggml_nelements(q)); ++ std::vector k_f32(ggml_nelements(k)); ++ std::vector v_f32(ggml_nelements(v)); ++ std::vector k_data(k_f32.size()); ++ std::vector v_data(v_f32.size()); ++ std::vector mask(ggml_nelements(m)); ++ for (size_t i = 0; i < q_data.size(); ++i) { ++ q_data[i] = 0.20f*std::sin(0.017f*float(i) + 0.11f); ++ } ++ for (size_t i = 0; i < k_f32.size(); ++i) { ++ k_f32[i] = 0.23f*std::cos(0.013f*float(i) - 0.29f); ++ v_f32[i] = 0.31f*std::sin(0.019f*float(i) + 0.71f); ++ } ++ ggml_fp32_to_fp16_row(k_f32.data(), k_data.data(), k_data.size()); ++ ggml_fp32_to_fp16_row(v_f32.data(), v_data.data(), v_data.size()); ++ for (int64_t iq = 0; iq < nq; ++iq) { ++ for (int64_t ik = 0; ik < nkv; ++ik) { ++ const int64_t dist = iq + (nkv - nq) - ik; ++ mask[iq*nkv + ik] = ggml_fp32_to_fp16(dist >= 0 && dist < 29 ? 0.0f : -INFINITY); ++ } ++ } ++ ggml_backend_tensor_set(q, q_data.data(), 0, q_data.size()*sizeof(q_data[0])); ++ ggml_backend_tensor_set(k, k_data.data(), 0, k_data.size()*sizeof(k_data[0])); ++ ggml_backend_tensor_set(v, v_data.data(), 0, v_data.size()*sizeof(v_data[0])); ++ ggml_backend_tensor_set(m, mask.data(), 0, mask.size()*sizeof(mask[0])); ++ ++ ggml_cgraph * graph = ggml_new_graph_custom(ctx, 16, false); ++ ggml_build_forward_expand(graph, out); ++ const ggml_status status = ggml_backend_graph_compute(backend, graph); ++ ggml_backend_synchronize(backend); ++ if (status != GGML_STATUS_SUCCESS) { ++ fprintf(stderr, "graph failed with status %d\n", int(status)); ++ ggml_backend_buffer_free(buffer); ++ ggml_free(ctx); ++ ggml_backend_free(backend); ++ return 1; ++ } ++ ++ std::vector result(ggml_nelements(out)); ++ ggml_backend_tensor_get(out, result.data(), 0, result.size()*sizeof(result[0])); ++ uint64_t fnv = UINT64_C(1469598103934665603); ++ const uint8_t * bytes = reinterpret_cast(result.data()); ++ for (size_t i = 0; i < result.size()*sizeof(result[0]); ++i) { ++ fnv ^= bytes[i]; ++ fnv *= UINT64_C(1099511628211); ++ } ++ if (!output_path.empty()) { ++ FILE * fp = fopen(output_path.c_str(), "wb"); ++ if (!fp || fwrite(result.data(), sizeof(result[0]), result.size(), fp) != result.size()) { ++ fprintf(stderr, "failed to write %s\n", output_path.c_str()); ++ if (fp) { ++ fclose(fp); ++ } ++ ggml_backend_buffer_free(buffer); ++ ggml_free(ctx); ++ ggml_backend_free(backend); ++ return 1; ++ } ++ fclose(fp); ++ } ++ printf("generic_hash backend=%s device=%s bytes=%zu fnv1a64=%016llx\n", ++ backend_kind.c_str(), ggml_backend_dev_description(chosen), ++ result.size()*sizeof(result[0]), (unsigned long long) fnv); ++ ++ ggml_backend_buffer_free(buffer); ++ ggml_free(ctx); ++ ggml_backend_free(backend); ++ return 0; ++} +diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp +index 659ea9e9c..c7efc95a4 100644 +--- a/tests/test-llama-archs.cpp ++++ b/tests/test-llama-archs.cpp +@@ -1358,6 +1358,9 @@ static bool arch_supported(const llm_arch arch) { + if (arch == LLM_ARCH_DEEPSEEK4) { + return false; + } ++ if (arch == LLM_ARCH_INKLING) { ++ return false; // TODO fixture params for the arch-specific hparams (d_rel, rel_extent, shortconv, logit_scale_denom) ++ } + + // FIXME some models are segfaulting with WebGPU: + #ifdef GGML_USE_WEBGPU +diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt +index 15040e4af..461c4f429 100644 +--- a/tools/mtmd/CMakeLists.txt ++++ b/tools/mtmd/CMakeLists.txt +@@ -38,6 +38,7 @@ add_library(mtmd + models/granite4-vision.cpp + models/hunyuanvl.cpp + models/internvl.cpp ++ models/inkling.cpp + models/kimivl.cpp + models/kimik25.cpp + models/nemotron-v2-vl.cpp +diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h +index 589fc724e..fc011f714 100644 +--- a/tools/mtmd/clip-impl.h ++++ b/tools/mtmd/clip-impl.h +@@ -285,6 +285,13 @@ + #define TN_A_FFN_POST_NORM "%s.blk.%d.ffn_post_norm.%s" + #define TN_A_FFN_POST_NORM_1 "%s.blk.%d.ffn_post_norm_1.%s" + ++// Inkling hMLP vision and dMel audio towers. ++#define TN_INKLING_HMLP_LINEAR "v.hmlp.%d.linear.weight" ++#define TN_INKLING_HMLP_NORM "v.hmlp.%d.norm.weight" ++#define TN_INKLING_HMLP_FINAL_NORM "v.hmlp.final_norm.weight" ++#define TN_INKLING_DMEL_EMBD "a.dmel.embedding.weight" ++#define TN_INKLING_DMEL_FINAL_NORM "a.dmel.final_norm.weight" ++ + // mobilenetv5 (gemma3n) definitions + #define TN_MNV5_STEM_CONV "v.conv_stem.conv.weight" + #define TN_MNV5_STEM_BIAS "v.conv_stem.conv.bias" +@@ -353,6 +360,7 @@ + struct clip_ctx; + + enum projector_type { ++ PROJECTOR_TYPE_INKLING, + PROJECTOR_TYPE_MLP, + PROJECTOR_TYPE_MLP_NORM, + PROJECTOR_TYPE_LDP, +@@ -411,6 +419,7 @@ enum projector_type { + }; + + static std::map PROJECTOR_TYPE_NAMES = { ++ { PROJECTOR_TYPE_INKLING, "inkling" }, + { PROJECTOR_TYPE_MLP, "mlp" }, + { PROJECTOR_TYPE_LDP, "ldp" }, + { PROJECTOR_TYPE_LDPV2, "ldpv2"}, +diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h +index 146eabce2..bfe021922 100644 +--- a/tools/mtmd/clip-model.h ++++ b/tools/mtmd/clip-model.h +@@ -33,7 +33,7 @@ enum resize_algo { + RESIZE_ALGO_BILINEAR, // stretch to target resolution + RESIZE_ALGO_BICUBIC, // center-crop when aspect ratio doesn't match + RESIZE_ALGO_BICUBIC_PILLOW, +- // RESIZE_ALGO_LANCZOS, // TODO ++ RESIZE_ALGO_LANCZOS, + }; + + // Padding style for img_tool::resize +@@ -363,6 +363,11 @@ struct qf_block { + std::vector qf_proj_layers; + }; + ++struct inkling_hmlp_layer { ++ ggml_tensor * linear_w = nullptr; ++ ggml_tensor * norm_w = nullptr; ++}; ++ + struct clip_model { + clip_modality modality = CLIP_MODALITY_VISION; + projector_type proj_type = PROJECTOR_TYPE_MLP; +@@ -377,6 +382,12 @@ struct clip_model { + ggml_tensor * norm_embd_w = nullptr; + ggml_tensor * norm_embd_b = nullptr; + ++ // Inkling towers (neither uses standard transformer blocks). ++ std::vector inkling_hmlp_layers; ++ ggml_tensor * inkling_hmlp_final_norm_w = nullptr; ++ ggml_tensor * inkling_dmel_embd_w = nullptr; ++ ggml_tensor * inkling_dmel_final_norm_w = nullptr; ++ + // "indexed" patch embedding norms + ggml_tensor * patch_norm_1_w = nullptr; + ggml_tensor * patch_norm_1_b = nullptr; +diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp +index 11f9820ed..f1ec237a4 100644 +--- a/tools/mtmd/clip.cpp ++++ b/tools/mtmd/clip.cpp +@@ -878,6 +878,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const + std::unique_ptr builder; + + switch (ctx->proj_type()) { ++ case PROJECTOR_TYPE_INKLING: ++ { ++ builder = std::make_unique(ctx, img); ++ } break; + case PROJECTOR_TYPE_GEMMA3: + case PROJECTOR_TYPE_IDEFICS3: + case PROJECTOR_TYPE_LFM2: +@@ -1298,6 +1302,21 @@ struct clip_model_loader { + + // model-specific params + switch (model.proj_type) { ++ case PROJECTOR_TYPE_INKLING: ++ { ++ if (is_vision) { ++ hparams.image_resize_algo = RESIZE_ALGO_BICUBIC_PILLOW; ++ hparams.image_resize_pad = PAD_NONE; ++ hparams.warmup_image_size = hparams.patch_size; ++ } else { ++ hparams.audio_chunk_len = 0; ++ hparams.audio_sample_rate = 16000; ++ hparams.audio_n_fft = 1600; ++ hparams.audio_window_len = 1600; ++ hparams.audio_hop_len = 800; ++ hparams.warmup_audio_size = 4; ++ } ++ } break; + case PROJECTOR_TYPE_MLP: + case PROJECTOR_TYPE_MLP_NORM: + case PROJECTOR_TYPE_LDP: +@@ -1971,7 +1990,8 @@ struct clip_model_loader { + model.position_embeddings = get_tensor(string_format(TN_POS_EMBD, prefix), false); + + const bool has_standard_layers = ( +- model.proj_type != PROJECTOR_TYPE_GEMMA3NV); ++ model.proj_type != PROJECTOR_TYPE_GEMMA3NV && ++ model.proj_type != PROJECTOR_TYPE_INKLING); + + // layers + const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; +@@ -2057,6 +2077,23 @@ struct clip_model_loader { + + + switch (model.proj_type) { ++ case PROJECTOR_TYPE_INKLING: ++ { ++ if (model.modality == CLIP_MODALITY_VISION) { ++ model.inkling_hmlp_layers.resize(hparams.n_layer); ++ for (int il = 0; il < hparams.n_layer; ++il) { ++ auto & layer = model.inkling_hmlp_layers[il]; ++ layer.linear_w = get_tensor(string_format(TN_INKLING_HMLP_LINEAR, il)); ++ layer.norm_w = il + 1 < hparams.n_layer ++ ? get_tensor(string_format(TN_INKLING_HMLP_NORM, il)) ++ : nullptr; ++ } ++ model.inkling_hmlp_final_norm_w = get_tensor(TN_INKLING_HMLP_FINAL_NORM); ++ } else { ++ model.inkling_dmel_embd_w = get_tensor(TN_INKLING_DMEL_EMBD); ++ model.inkling_dmel_final_norm_w = get_tensor(TN_INKLING_DMEL_FINAL_NORM); ++ } ++ } break; + case PROJECTOR_TYPE_MLP: + case PROJECTOR_TYPE_MLP_NORM: + { +@@ -3144,6 +3181,11 @@ struct clip_model_loader { + LOG_INF("%s: warmup with audio size = %d\n", __func__, hparams.warmup_audio_size); + } + batch.entries.push_back(img); ++ // One logical Inkling vision patch always contains two adjacent temporal slices (including warmup reserve). ++ if (ctx_clip.model.modality == CLIP_MODALITY_VISION && ++ ctx_clip.model.proj_type == PROJECTOR_TYPE_INKLING) { ++ batch.entries.push_back(img); ++ } + return batch; + } + +@@ -3541,6 +3583,16 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { + projector_type proj = ctx->proj_type(); + + switch (proj) { ++ case PROJECTOR_TYPE_INKLING: ++ { ++ if (ctx->model.modality == CLIP_MODALITY_AUDIO) { ++ // One dMel row is one soft audio token. ++ n_patches = img->nx(); ++ } else { ++ // Batch entries are temporal slices; each pair emits one token. ++ n_patches = 1; ++ } ++ } break; + case PROJECTOR_TYPE_MLP: + case PROJECTOR_TYPE_MLP_NORM: + case PROJECTOR_TYPE_JANUS_PRO: +@@ -3927,6 +3979,10 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 + + // set input per projector + switch (ctx->model.proj_type) { ++ case PROJECTOR_TYPE_INKLING: ++ { ++ // inp_raw is the only graph input. ++ } break; + case PROJECTOR_TYPE_MINICPMV: + { + // inspired from siglip: +@@ -4955,6 +5011,10 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 + + int clip_n_mmproj_embd(const struct clip_ctx * ctx) { + switch (ctx->model.proj_type) { ++ case PROJECTOR_TYPE_INKLING: ++ return ctx->model.modality == CLIP_MODALITY_AUDIO ++ ? ctx->model.inkling_dmel_final_norm_w->ne[0] ++ : ctx->model.inkling_hmlp_final_norm_w->ne[0]; + case PROJECTOR_TYPE_LDP: + return ctx->model.mm_model_block_1_block_2_1_b->ne[0]; + case PROJECTOR_TYPE_LDPV2: +diff --git a/tools/mtmd/models/inkling.cpp b/tools/mtmd/models/inkling.cpp +new file mode 100644 +index 000000000..c072f66a2 +--- /dev/null ++++ b/tools/mtmd/models/inkling.cpp +@@ -0,0 +1,106 @@ ++#include "models.h" ++ ++ggml_tensor * clip_graph_inkling::build_mm(ggml_tensor * w, ggml_tensor * x) const { ++ ggml_tensor * cur = ggml_mul_mat(ctx0, w, x); ++ ggml_mul_mat_set_prec(cur, GGML_PREC_F32); ++ return cur; ++} ++ ++// fold square neighborhoods from W/H into channels; folded order is [h_fold, w_fold, C] ++static ggml_tensor * inkling_fold_spatial( ++ ggml_context * ctx0, ++ ggml_tensor * cur, ++ int scale) { ++ GGML_ASSERT(scale > 0); ++ GGML_ASSERT(cur->ne[1] % scale == 0); ++ GGML_ASSERT(cur->ne[2] % scale == 0); ++ ++ const int64_t c = cur->ne[0]; ++ const int64_t w = cur->ne[1]; ++ const int64_t h = cur->ne[2]; ++ const int64_t b = cur->ne[3]; ++ ++ cur = ggml_reshape_4d(ctx0, cur, c * scale, w / scale, h, b); ++ cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 2, 1, 3)); ++ cur = ggml_reshape_4d(ctx0, cur, c * scale * scale, h / scale, w / scale, b); ++ cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 2, 1, 3)); ++ return cur; ++} ++ ++ggml_cgraph * clip_graph_inkling::build() { ++ return model.modality == CLIP_MODALITY_AUDIO ? build_audio() : build_vision(); ++} ++ ++ggml_cgraph * clip_graph_inkling::build_vision() { ++ static constexpr int temporal_patch_size = 2; ++ static constexpr int spatial_folds[] = {5, 2, 4}; ++ ++ GGML_ASSERT(img.nx() == 40 && img.ny() == 40); ++ GGML_ASSERT(n_batch > 0 && n_batch % temporal_patch_size == 0); ++ GGML_ASSERT(model.inkling_hmlp_layers.size() == 4); ++ GGML_ASSERT(model.inkling_hmlp_final_norm_w); ++ ++ // Raw input is [W,H,RGB,temporal*patch]. Put RGB on ne[0]. ++ ggml_tensor * cur = build_inp_raw(3); ++ cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 2, 0, 3)); ++ ++ for (int il = 0; il < 3; ++il) { ++ cur = inkling_fold_spatial(ctx0, cur, spatial_folds[il]); ++ const int64_t w = cur->ne[1]; ++ const int64_t h = cur->ne[2]; ++ const int64_t b = cur->ne[3]; ++ ++ cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], w * h * b); ++ cur = build_mm(model.inkling_hmlp_layers[il].linear_w, cur); ++ cur = build_norm(cur, model.inkling_hmlp_layers[il].norm_w, ++ nullptr, NORM_TYPE_RMS, eps, il); ++ cur = ggml_gelu_erf(ctx0, cur); ++ cur = ggml_reshape_4d(ctx0, cur, cur->ne[0], w, h, b); ++ } ++ ++ GGML_ASSERT(cur->ne[1] == 1 && cur->ne[2] == 1); ++ const int64_t n_patches = n_batch / temporal_patch_size; ++ cur = ggml_reshape_2d(ctx0, cur, cur->ne[0] * temporal_patch_size, n_patches); ++ cur = build_mm(model.inkling_hmlp_layers[3].linear_w, cur); ++ cur = build_norm(cur, model.inkling_hmlp_final_norm_w, ++ nullptr, NORM_TYPE_RMS, eps, 3); ++ ++ // Batched mtmd convention: one token in ne[1], patch count in ne[2]. ++ cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], 1, n_patches); ++ ggml_build_forward_expand(gf, cur); ++ return gf; ++} ++ ++ggml_cgraph * clip_graph_inkling::build_audio() { ++ static constexpr int n_mels = 80; ++ static constexpr int mel_vocab_size = 16; ++ static constexpr int n_embd = 6144; ++ ++ GGML_ASSERT(img.ny() == n_mels); ++ GGML_ASSERT(model.inkling_dmel_embd_w); ++ GGML_ASSERT(model.inkling_dmel_final_norm_w); ++ GGML_ASSERT(model.inkling_dmel_embd_w->ne[0] == n_embd); ++ GGML_ASSERT(model.inkling_dmel_embd_w->ne[1] == n_mels * mel_vocab_size); ++ ++ const int64_t n_tokens = img.nx(); ++ // mtmd audio storage is mel-major and represented as [token, mel]. ++ ggml_tensor * bins = build_inp_raw(1); ++ bins = ggml_cont(ctx0, ggml_transpose(ctx0, bins)); ++ ggml_tensor * offsets = ggml_arange(ctx0, 0, n_mels, 1); ++ offsets = ggml_scale(ctx0, offsets, mel_vocab_size); ++ offsets = ggml_reshape_2d(ctx0, offsets, n_mels, 1); ++ ggml_tensor * indices = ggml_cast(ctx0, ggml_add(ctx0, bins, offsets), GGML_TYPE_I32); ++ ++ indices = ggml_reshape_1d(ctx0, indices, n_mels * n_tokens); ++ ggml_tensor * cur = ggml_get_rows(ctx0, model.inkling_dmel_embd_w, indices); ++ cur = ggml_reshape_3d(ctx0, cur, n_embd, n_mels, n_tokens); ++ cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 0, 2, 3)); ++ cur = ggml_sum_rows(ctx0, cur); ++ cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 0, 2, 3)); ++ cur = ggml_reshape_2d(ctx0, cur, n_embd, n_tokens); ++ cur = build_norm(cur, model.inkling_dmel_final_norm_w, ++ nullptr, NORM_TYPE_RMS, eps, -1); ++ ++ ggml_build_forward_expand(gf, cur); ++ return gf; ++} +diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h +index e54366a08..6735c9cde 100644 +--- a/tools/mtmd/models/models.h ++++ b/tools/mtmd/models/models.h +@@ -12,6 +12,17 @@ struct clip_graph_siglip : clip_graph { + ggml_cgraph * build() override; + }; + ++struct clip_graph_inkling : clip_graph { ++ clip_graph_inkling(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ++ ggml_cgraph * build() override; ++ ggml_tensor * build_mm(ggml_tensor * w, ggml_tensor * x) const override; ++ bool support_batch() const override { return true; } ++ ++private: ++ ggml_cgraph * build_vision(); ++ ggml_cgraph * build_audio(); ++}; ++ + struct clip_graph_gemma4v : clip_graph { + clip_graph_gemma4v(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; +diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp +index fea03557d..13e31ff54 100644 +--- a/tools/mtmd/mtmd-audio.cpp ++++ b/tools/mtmd/mtmd-audio.cpp +@@ -283,6 +283,7 @@ struct filter_params { + bool norm_per_feature = false; + bool use_magnitude = false; // |X| instead of |X|^2 + float mel_floor = 5.960464477539063e-08f; ++ float power_floor = 0.0f; // clamp |X|^2 to this before sqrt (0 = disabled) + }; + + static void log_mel_spectrogram_worker_thread(int ith, +@@ -327,6 +328,7 @@ static void log_mel_spectrogram_worker_thread(int ith, + // Calculate modulus^2 (power) or modulus (magnitude) + for (int j = 0; j < n_fft_bins; j++) { + float power = (fft_out[2 * j + 0] * fft_out[2 * j + 0] + fft_out[2 * j + 1] * fft_out[2 * j + 1]); ++ power = std::max(power, params.power_floor); + fft_out[j] = params.use_magnitude ? sqrtf(power) : power; + } + +@@ -537,6 +539,75 @@ static bool log_mel_spectrogram( + return true; + } + ++void mtmd_audio_preprocessor_inkling::initialize() { ++ GGML_ASSERT(hparams.n_mel_bins == 80); ++ GGML_ASSERT(hparams.audio_n_fft == 1600); ++ cache.fill_sin_cos_table(hparams.audio_n_fft); ++ cache.fill_hann_window(hparams.audio_window_len, true); ++ cache.fill_mel_filterbank_matrix( ++ hparams.n_mel_bins, hparams.audio_n_fft, hparams.audio_sample_rate, ++ 0.0f, -1.0f, true, 1.0f, false); ++} ++ ++bool mtmd_audio_preprocessor_inkling::preprocess( ++ const float * samples, ++ size_t n_samples, ++ std::vector & output) { ++ if (n_samples == 0) { ++ return false; ++ } ++ GGML_ASSERT(hparams.audio_hop_len == 800); ++ GGML_ASSERT(hparams.audio_window_len == 1600); ++ GGML_ASSERT(!cache.filters.data.empty()); ++ ++ // Reference behavior: left pad by n_fft-hop, then right pad to a whole hop. ++ const size_t left_pad = static_cast(hparams.audio_n_fft - hparams.audio_hop_len); ++ const size_t right_pad = ++ (static_cast(hparams.audio_hop_len) - n_samples % hparams.audio_hop_len) ++ % hparams.audio_hop_len; ++ std::vector padded(left_pad + n_samples + right_pad, 0.0f); ++ std::copy(samples, samples + n_samples, padded.begin() + left_pad); ++ ++ filter_params params; ++ params.n_mel = hparams.n_mel_bins; ++ params.n_fft_bins = 1 + hparams.audio_n_fft / 2; ++ params.hann_window_size = hparams.audio_window_len; ++ params.hop_length = hparams.audio_hop_len; ++ params.sample_rate = hparams.audio_sample_rate; ++ params.no_padding = true; ++ params.use_natural_log = false; ++ params.use_magnitude = true; ++ params.mel_floor = 1e-10f; ++ params.power_floor = 1e-10f; // reference clamps |X|^2 before sqrt ++ ++ mtmd_audio_mel dmel; ++ if (!log_mel_spectrogram( ++ padded.data(), static_cast(padded.size()), 4, ++ params, cache, dmel)) { ++ return false; ++ } ++ dmel.n_len_org = static_cast(n_samples); ++ ++ // nearest of 16 float64 centers over [-7,2]; strict '<' keeps the lower bin on midpoints (torch.argmin) ++ for (float & value_f32 : dmel.data) { ++ const double value = std::max(-7.0, std::min(2.0, static_cast(value_f32))); ++ int best = 0; ++ double best_dist = INFINITY; ++ for (int bin = 0; bin < 16; ++bin) { ++ const double center = -7.0 + 9.0 * static_cast(bin) / 15.0; ++ const double dist = std::abs(value - center); ++ if (dist < best_dist) { ++ best = bin; ++ best_dist = dist; ++ } ++ } ++ value_f32 = static_cast(best); ++ } ++ ++ output.push_back(std::move(dmel)); ++ return true; ++} ++ + // + // mtmd_audio_preprocessor_whisper + // +diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h +index f65f282d9..e79519553 100644 +--- a/tools/mtmd/mtmd-audio.h ++++ b/tools/mtmd/mtmd-audio.h +@@ -69,6 +69,16 @@ struct mtmd_audio_preprocessor_whisper : mtmd_audio_preprocessor { + mtmd_audio_cache cache; + }; + ++// Inkling dMel: 100 ms Slaney-mel magnitude windows at a 50 ms hop, quantized to 16 bins. ++struct mtmd_audio_preprocessor_inkling : mtmd_audio_preprocessor { ++ mtmd_audio_preprocessor_inkling(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} ++ void initialize() override; ++ bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; ++ ++private: ++ mtmd_audio_cache cache; ++}; ++ + struct mtmd_audio_preprocessor_conformer : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_conformer(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} + void initialize() override; +diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp +index 36cd463b2..01cb9f592 100644 +--- a/tools/mtmd/mtmd-image.cpp ++++ b/tools/mtmd/mtmd-image.cpp +@@ -68,6 +68,9 @@ struct img_tool { + case RESIZE_ALGO_BICUBIC_PILLOW: + resize_bicubic_pillow(src, dst, target_resolution.width, target_resolution.height); + break; ++ case RESIZE_ALGO_LANCZOS: ++ resize_lanczos_pillow(src, dst, target_resolution.width, target_resolution.height); ++ break; + default: + throw std::runtime_error("Unsupported resize algorithm"); + } +@@ -97,6 +100,9 @@ struct img_tool { + case RESIZE_ALGO_BICUBIC_PILLOW: + resize_bicubic_pillow(src, resized_image, new_width, new_height); + break; ++ case RESIZE_ALGO_LANCZOS: ++ resize_lanczos_pillow(src, resized_image, new_width, new_height); ++ break; + default: + throw std::runtime_error("Unsupported resize algorithm"); + } +@@ -345,6 +351,19 @@ private: + // 2. Pre-computes normalized filter coefficients for each output pixel + // 3. Applies convolution using fixed-point integer arithmetic for performance + static bool resize_bicubic_pillow(const clip_image_u8 & img, clip_image_u8 & dst, int target_width, int target_height) { ++ return resize_pillow(img, dst, target_width, target_height, false); ++ } ++ ++ static bool resize_lanczos_pillow(const clip_image_u8 & img, clip_image_u8 & dst, int target_width, int target_height) { ++ return resize_pillow(img, dst, target_width, target_height, true); ++ } ++ ++ static bool resize_pillow( ++ const clip_image_u8 & img, ++ clip_image_u8 & dst, ++ int target_width, ++ int target_height, ++ bool use_lanczos) { + // Fixed-point precision: 22 bits = 32 (int32_t) - 8 (uint8_t pixels) - 2 (headroom for accumulation) + // This allows encoding fractional weights as integers: weight * 2^22 + const int PRECISION_BITS = 32 - 8 - 2; +@@ -352,7 +371,21 @@ private: + // Bicubic filter function with a = -0.5 (Note that GGML/PyTorch takes a = -0.75) + // Returns filter weight for distance x from pixel center + // Support: [-2, 2], meaning the filter influences pixels within 2 units of distance +- auto bicubic_filter = [](double x) -> double { ++ auto resample_filter = [use_lanczos](double x) -> double { ++ if (use_lanczos) { ++ if (-3.0 <= x && x < 3.0) { ++ auto sinc = [](double value) { ++ if (value == 0.0) { ++ return 1.0; ++ } ++ const double pix = value * 3.141592653589793238462643383279502884; ++ return std::sin(pix) / pix; ++ }; ++ return sinc(x) * sinc(x / 3.0); ++ } ++ return 0.0; ++ } ++ + constexpr double a = -0.5; + if (x < 0.0) { + x = -x; +@@ -367,7 +400,7 @@ private: + }; + + // Filter support radius: bicubic extends 2 pixels in each direction +- constexpr double filter_support = 2.0; ++ const double filter_support = use_lanczos ? 3.0 : 2.0; + + // Clipping function for 8-bit values + auto clip8 = [](int val) -> uint8_t { +@@ -434,7 +467,7 @@ private: + // Compute filter weights for each contributing input pixel + for (x = 0; x < xmax; x++) { + // Distance from input pixel center to output pixel center in input space +- double w = bicubic_filter((x + xmin - center + 0.5) * ss); ++ double w = resample_filter((x + xmin - center + 0.5) * ss); + pre_weights[xx * ksize + x] = w; + ww += w; // Accumulate for normalization + } +@@ -463,17 +496,9 @@ private: + const double fxp_scale = std::ldexp(1.0, PRECISION_BITS); // 1.0 * 2^PRECISION_BITS + + for (int i = 0; i < outSize * ksize; i++) { +- double tmp_val = pre_weights[i] * fxp_scale; +- if (pre_weights[i] < 0) { +- tmp_val -= 0.5; +- } else { +- tmp_val += 0.5; +- } +- tmp_val = std::round(tmp_val); +- tmp_val = std::clamp(tmp_val, +- static_cast(std::numeric_limits::min()), +- static_cast(std::numeric_limits::max())); +- weights[i] = static_cast(tmp_val); ++ // Pillow adds +/- 0.5 then truncates toward zero; std::round would round twice ++ const double rounded = pre_weights[i] * fxp_scale + (pre_weights[i] < 0 ? -0.5 : 0.5); ++ weights[i] = static_cast(rounded); + } + + return ksize; +@@ -606,6 +631,82 @@ private: + } + }; + ++mtmd_inkling_image_preproc_out mtmd_image_preprocess_inkling( ++ const clip_image_u8 & img, ++ resize_algo algo) { ++ constexpr int patch_size = 40; ++ constexpr int temporal_patch_size = 2; ++ constexpr int max_upscaled_long_edge = 2048; ++ constexpr double rescale_image_frac = 2.0; ++ constexpr float image_mean[3] = { 0.48145466f, 0.4578275f, 0.40821073f }; ++ constexpr float image_std[3] = { 0.26862954f, 0.26130258f, 0.27577711f }; ++ ++ mtmd_inkling_image_preproc_out output; ++ output.source_size = img.get_size(); ++ GGML_ASSERT(output.source_size.width > 0 && output.source_size.height > 0); ++ GGML_ASSERT(!img.is_placeholder()); ++ ++ const int long_edge = std::max(output.source_size.width, output.source_size.height); ++ const double target_long_edge = std::min( ++ static_cast(long_edge) * rescale_image_frac, ++ static_cast(std::max(max_upscaled_long_edge, long_edge))); ++ const double ratio = target_long_edge / long_edge; ++ const auto scaled = [ratio](int size) { ++ // The reference explicitly requests half-up rounding for positive sizes. ++ return std::max(1, static_cast(std::floor(size * ratio + 0.5))); ++ }; ++ output.resized_size = { ++ scaled(output.source_size.width), ++ scaled(output.source_size.height), ++ }; ++ ++ clip_image_u8 resized; ++ img_tool::resize(img, resized, output.resized_size, algo, PAD_NONE); ++ output.resized_rgb = resized.get_ro_buf(); ++ ++ output.patch_rows = (output.resized_size.height + patch_size - 1) / patch_size; ++ // Inkling always appends a right-hand patch, even at exact multiples of 40 ++ output.patch_cols = output.resized_size.width / patch_size + 1; ++ const size_t n_patches = static_cast(output.patch_rows) * output.patch_cols; ++ const size_t values_per_patch = static_cast(temporal_patch_size) * ++ patch_size * patch_size * 3; ++ output.pixel_values_bthwc.resize(n_patches * values_per_patch); ++ ++ // preserve torchvision's fused (raw - mean*255)/(std*255) order (raw=-1 for padding) for ulp parity ++ float mean_255[3]; ++ float std_255[3]; ++ for (int c = 0; c < 3; ++c) { ++ mean_255[c] = image_mean[c] * 255.0f; ++ std_255[c] = image_std[c] * 255.0f; ++ } ++ ++ for (int py = 0; py < output.patch_rows; ++py) { ++ for (int px = 0; px < output.patch_cols; ++px) { ++ const size_t patch_index = static_cast(py) * output.patch_cols + px; ++ for (int t = 0; t < temporal_patch_size; ++t) { ++ for (int y = 0; y < patch_size; ++y) { ++ const int iy = py * patch_size + y; ++ for (int x = 0; x < patch_size; ++x) { ++ const int ix = px * patch_size + x; ++ const bool in_image = iy < output.resized_size.height && ix < output.resized_size.width; ++ const std::array rgb = in_image ++ ? resized.get_pixel(ix, iy) ++ : std::array{ 0, 0, 0 }; ++ for (int c = 0; c < 3; ++c) { ++ const float raw = in_image ? static_cast(rgb[c]) : -1.0f; ++ const size_t offset = (((((patch_index * temporal_patch_size + t) * ++ patch_size + y) * patch_size + x) * 3) + c); ++ output.pixel_values_bthwc[offset] = (raw - mean_255[c]) / std_255[c]; ++ } ++ } ++ } ++ } ++ } ++ } ++ ++ return output; ++} ++ + + // + // mtmd_image_preprocessor_llava_uhd +@@ -885,6 +986,80 @@ mtmd_image_preproc_out mtmd_image_preprocessor_fixed_size::preprocess(const clip + return output; + } + ++// Inkling hMLP patches ++mtmd_image_preproc_out mtmd_image_preprocessor_inkling::preprocess(const clip_image_u8 & img) { ++ GGML_ASSERT(hparams.patch_size == 40); ++ constexpr int temporal_patch_size = 2; ++ constexpr float rescale_frac = 2.0f; ++ constexpr int rescale_max_long_edge = 2048; ++ ++ const auto original = img.get_size(); ++ const int long_edge = std::max(original.width, original.height); ++ const float target_long = std::min( ++ static_cast(long_edge) * rescale_frac, ++ static_cast(std::max(long_edge, rescale_max_long_edge))); ++ const float ratio = long_edge > 0 ? target_long / long_edge : 1.0f; ++ const clip_image_size scaled_size { ++ std::max(1, static_cast(std::floor(original.width * ratio + 0.5f))), ++ std::max(1, static_cast(std::floor(original.height * ratio + 0.5f))), ++ }; ++ ++ clip_image_u8 scaled; ++ // The reference processor resizes with Pillow Lanczos (radius 3). ++ img_tool::resize(img, scaled, scaled_size, RESIZE_ALGO_LANCZOS, PAD_NONE); ++ ++ const int p = hparams.patch_size; ++ const int nph = (scaled_size.height + p - 1) / p; ++ // The extra right column is intentional, including exact multiples of 40. ++ const int npw = scaled_size.width / p + 1; ++ ++ const float pad_raw = -1.0f / 255.0f; ++ // the reference materializes normalized patches as bf16; round through bf16 for bit-parity ++ const auto bf16_round = [](float v) { ++ return ggml_bf16_to_fp32(ggml_fp32_to_bf16(v)); ++ }; ++ float pad_norm[3]; ++ for (int c = 0; c < 3; ++c) { ++ pad_norm[c] = bf16_round((pad_raw - hparams.image_mean[c]) / hparams.image_std[c]); ++ } ++ ++ mtmd_image_preproc_out output; ++ output.entries.reserve(static_cast(nph) * npw * temporal_patch_size); ++ for (int py = 0; py < nph; ++py) { ++ for (int px = 0; px < npw; ++px) { ++ std::vector data(static_cast(p) * p * 3); ++ for (int y = 0; y < p; ++y) { ++ const int iy = py * p + y; ++ for (int x = 0; x < p; ++x) { ++ const int ix = px * p + x; ++ const size_t off = static_cast(y * p + x) * 3; ++ if (iy < scaled_size.height && ix < scaled_size.width && !scaled.is_placeholder()) { ++ const auto rgb = scaled.get_pixel(ix, iy); ++ for (int c = 0; c < 3; ++c) { ++ const float raw = static_cast(rgb[c]) / 255.0f; ++ data[off + c] = bf16_round((raw - hparams.image_mean[c]) / hparams.image_std[c]); ++ } ++ } else { ++ for (int c = 0; c < 3; ++c) { ++ data[off + c] = pad_norm[c]; ++ } ++ } ++ } ++ } ++ ++ for (int t = 0; t < temporal_patch_size; ++t) { ++ clip_image_f32 patch; ++ patch.set_size({p, p}, scaled.is_placeholder(), false); ++ if (!scaled.is_placeholder()) { ++ patch.cpy_buf(data); ++ } ++ output.entries.push_back(std::move(patch)); ++ } ++ } ++ } ++ return output; ++} ++ + // + // mtmd_image_preprocessor_dyn_size + // +diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h +index 115cba51e..044d72f61 100644 +--- a/tools/mtmd/mtmd-image.h ++++ b/tools/mtmd/mtmd-image.h +@@ -26,6 +26,20 @@ struct mtmd_image_preproc_out { + } + }; + ++// Inkling hMLP preprocessing; pixel_values_bthwc is row-major BTHWC [n_patches, 2, 40, 40, 3]. ++struct mtmd_inkling_image_preproc_out { ++ clip_image_size source_size; ++ clip_image_size resized_size; ++ int patch_rows = 0; ++ int patch_cols = 0; ++ std::vector resized_rgb; ++ std::vector pixel_values_bthwc; ++}; ++ ++mtmd_inkling_image_preproc_out mtmd_image_preprocess_inkling( ++ const clip_image_u8 & img, ++ resize_algo algo = RESIZE_ALGO_LANCZOS); ++ + // base class, models must inherit from this class + struct mtmd_image_preprocessor { + const clip_hparams & hparams; +@@ -115,6 +129,12 @@ struct mtmd_image_preprocessor_fixed_size : mtmd_image_preprocessor { + mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; + }; + ++// Inkling: split an image into 40x40 hMLP patches, each duplicated across a fixed temporal dimension of two. ++struct mtmd_image_preprocessor_inkling : mtmd_image_preprocessor { ++ mtmd_image_preprocessor_inkling(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} ++ mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; ++}; ++ + // resize image to multiple of patch_size*n_merge, while preserving aspect ratio + // if image_resize_pad is true, the resized image will be padded, otherwise it will be either stretched or center-cropped depending on image_resize_pad + // this is used by models with native support for dynamic image size, for example: Qwen-VL, Pixtral, Kimi-VL, etc +diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp +index 93ca8cbcf..a70f326a9 100644 +--- a/tools/mtmd/mtmd.cpp ++++ b/tools/mtmd/mtmd.cpp +@@ -393,6 +393,12 @@ struct mtmd_context { + projector_type proj = clip_get_projector_type(ctx_v); + + switch (proj) { ++ case PROJECTOR_TYPE_INKLING: ++ { ++ // renderer opens each image part with <|content_image|>; block framing comes from the template ++ img_beg = "<|content_image|>"; ++ image_preproc = std::make_unique(ctx_v); ++ } break; + case PROJECTOR_TYPE_MLP: + case PROJECTOR_TYPE_MLP_NORM: + case PROJECTOR_TYPE_LDP: +@@ -678,6 +684,13 @@ struct mtmd_context { + + // set preprocessor + switch (proj) { ++ case PROJECTOR_TYPE_INKLING: ++ { ++ // <|content_audio_input|> ... <|audio_end|>, matching the renderer's audio framing ++ aud_beg = "<|content_audio_input|>"; ++ aud_end = "<|audio_end|>"; ++ audio_preproc = std::make_unique(ctx_a); ++ } break; + case PROJECTOR_TYPE_QWEN2A: + case PROJECTOR_TYPE_QWEN25O: + { +@@ -1204,23 +1217,35 @@ struct mtmd_tokenizer { + } + + size_t n_tokens = 0; +- for (auto & e : preproc_out.entries) { +- n_tokens += clip_n_output_tokens(ctx->ctx_v, &e); +- if (clip_model_n_temporal_merge(ctx->ctx_v) == 2) { +- // [QWEN_VIDEO] pair input is merged to the same embd, so only count as one image +- break; ++ if (ctx->proj_type_v() == PROJECTOR_TYPE_INKLING) { ++ GGML_ASSERT(preproc_out.entries.size() % 2 == 0); ++ n_tokens = preproc_out.entries.size() / 2; ++ } else { ++ for (auto & e : preproc_out.entries) { ++ n_tokens += clip_n_output_tokens(ctx->ctx_v, &e); ++ if (clip_model_n_temporal_merge(ctx->ctx_v) == 2) { ++ // [QWEN_VIDEO] pair input is merged to the same embd, so only count as one image ++ break; ++ } + } + } + + mtmd_image_tokens_ptr image_tokens(new mtmd_image_tokens); + + // [QWEN_VIDEO] improve this in the future +- image_tokens->n_temporal_merge = clip_model_n_temporal_merge(ctx->ctx_v); ++ // Inkling's pair is internal to each patch; it must not merge consecutive user bitmaps as video frames. ++ image_tokens->n_temporal_merge = ctx->proj_type_v() == PROJECTOR_TYPE_INKLING ++ ? 2 ++ : clip_model_n_temporal_merge(ctx->ctx_v); + + if (mtmd_decode_use_mrope(ctx)) { + // for Qwen2VL, we need this information for M-RoPE decoding positions + image_tokens->nx = clip_n_output_tokens_x(ctx->ctx_v, &preproc_out.entries[0]); + image_tokens->ny = clip_n_output_tokens_y(ctx->ctx_v, &preproc_out.entries[0]); ++ } else if (ctx->proj_type_v() == PROJECTOR_TYPE_INKLING) { ++ // n_tokens() multiplies 1x1 by the number of temporal groups. ++ image_tokens->nx = 1; ++ image_tokens->ny = 1; + } else { + // other models, we only need the total number of tokens + image_tokens->nx = n_tokens; +-- +2.50.1 (Apple Git-155) diff --git a/third_party/llama.cpp/patches/0051-Add-staged-execution-support-for-Inkling.patch b/third_party/llama.cpp/patches/0051-Add-staged-execution-support-for-Inkling.patch new file mode 100644 index 000000000..029727fc7 --- /dev/null +++ b/third_party/llama.cpp/patches/0051-Add-staged-execution-support-for-Inkling.patch @@ -0,0 +1,108 @@ +From de940c16a6bf7999e9a9b164a3716d6d9f6c7468 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 13:47:13 +1000 +Subject: [PATCH 51/61] Add staged execution support for Inkling + +--- + src/models/inkling.cpp | 30 +++++++++++++++++++++++------- + src/skippy.cpp | 1 + + 2 files changed, 24 insertions(+), 7 deletions(-) + +diff --git a/src/models/inkling.cpp b/src/models/inkling.cpp +index 6c9c1aef8..3d5f8ab56 100644 +--- a/src/models/inkling.cpp ++++ b/src/models/inkling.cpp +@@ -196,6 +196,11 @@ std::unique_ptr llama_model_inkling::build_arch_graph(const l + llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + ++ const skippy_graph_filter & stage_filter = skippy_graph_get_filter(); ++ const bool stage_filtered = stage_filter.enabled; ++ const int il_start = stage_filtered ? stage_filter.layer_start : 0; ++ const int il_end = stage_filtered ? stage_filter.layer_end : n_layer; ++ + const int64_t head_dim = hparams.n_embd_head_k(); + const int64_t d_rel = hparams.inkling_d_rel; + const int64_t d_conv = hparams.n_shortconv_l_cache - 1; +@@ -275,7 +280,7 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + dev_supports_banded_flash(il); + }; + +- for (int il = 0; il < n_layer; ++il) { ++ for (int il = il_start; il < il_end; ++il) { + if (hparams.is_swa(il)) { + needs_rel_idx_local |= !use_banded_flash(il); + } else { +@@ -307,14 +312,17 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + } + + const int64_t n_vocab = model.vocab.n_tokens(); +- if (!cparams.embeddings && hparams.inkling_unpadded_n_vocab > 0 && (int64_t) hparams.inkling_unpadded_n_vocab < n_vocab) { ++ if ((!stage_filtered || stage_filter.include_output) && ++ !cparams.embeddings && ++ hparams.inkling_unpadded_n_vocab > 0 && ++ (int64_t) hparams.inkling_unpadded_n_vocab < n_vocab) { + inp->vocab_mask = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_vocab); + ggml_set_input(inp->vocab_mask); + ggml_set_name(inp->vocab_mask, "inkling_vocab_mask"); + } + + // shared experts go through mul_mat_id: 2D views into a repacked/quantized 3D bank are invalid +- if (hparams.n_expert_shared > 0 && (uint32_t) n_layer > hparams.n_layer_dense_lead) { ++ if (hparams.n_expert_shared > 0 && (uint32_t) il_end > hparams.n_layer_dense_lead) { + inp->shexp_idx = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, hparams.n_expert_shared, n_tokens); + ggml_set_input(inp->shexp_idx); + ggml_set_name(inp->shexp_idx, "inkling_shexp_idx"); +@@ -630,9 +638,9 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + return moe_out; + }; + +- ggml_tensor * cur = build_inp_embd(model.tok_embd); ++ ggml_tensor * cur = build_inp_embd(stage_filtered && il_start > 0 ? nullptr : model.tok_embd); + // mtmd embd rows arrive pre-normalized; embed_norm applies to text token lookups only +- if (ubatch.token) { ++ if (ubatch.token && (!stage_filtered || stage_filter.include_embeddings)) { + cur = build_norm(cur, model.tok_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "inkling_embd_norm", -1); + } else { +@@ -641,7 +649,7 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + + ggml_build_forward_expand(gf, cur); + +- for (int il = 0; il < n_layer; ++il) { ++ for (int il = il_start; il < il_end; ++il) { + conv_rs_cur = build_rs(inp_hybrid->get_recr(), mctx_recr->get_r_l(il), n_embd_r, n_seqs); + + // h += attn_sconv(attn(attn_norm(h))) +@@ -667,7 +675,15 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + cb(cur, "l_out", il); + } + +- // conv states need every layer to see ALL tokens, so trim outputs only after the full stack ++ if (stage_filtered && !stage_filter.include_output) { ++ cb(cur, "stage_boundary", il_end - 1); ++ res->t_embd = cur; ++ ggml_build_forward_expand(gf, cur); ++ return; ++ } ++ ++ // Conv states need every layer to see all tokens. Intermediate stages hand off the full ++ // activation sequence, and the final stage trims outputs only after its last local layer. + ggml_tensor * inp_out_ids = build_inp_out_ids(); + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 3965e0252..f6ccad78c 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -4187,6 +4187,7 @@ static enum skippy_status skippy_finish_model_open( + model->arch != LLM_ARCH_HUNYUAN_DENSE && + model->arch != LLM_ARCH_HUNYUAN_MOE && + model->arch != LLM_ARCH_HUNYUAN_VL && ++ model->arch != LLM_ARCH_INKLING && + model->arch != LLM_ARCH_INTERNLM2 && + model->arch != LLM_ARCH_JAIS && + model->arch != LLM_ARCH_JAIS2 && +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0052-Advance-native-MTP-depth-per-draft-step.patch b/third_party/llama.cpp/patches/0052-Advance-native-MTP-depth-per-draft-step.patch new file mode 100644 index 000000000..4efacbbb7 --- /dev/null +++ b/third_party/llama.cpp/patches/0052-Advance-native-MTP-depth-per-draft-step.patch @@ -0,0 +1,81 @@ +From 1d78b6bfc85213cc7bcacfb60c691f7424c4bd64 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 14:46:51 +1000 +Subject: [PATCH 52/61] Advance native MTP depth per draft step + +--- + src/skippy.cpp | 37 +++++++++++++++++++++++++++++++++++++ + 1 file changed, 37 insertions(+) + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index f6ccad78c..1484f4f3b 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -1792,6 +1792,40 @@ static bool skippy_mtp_available(const skippy_session * session) { + session->stage_model->config.include_output; + } + ++static constexpr int32_t skippy_mtp_depth_for_step(size_t draft_step, uint32_t depth_count) { ++ return depth_count > 0 ? static_cast(draft_step % depth_count) : 0; ++} ++ ++static_assert(skippy_mtp_depth_for_step(0, 8) == 0); ++static_assert(skippy_mtp_depth_for_step(7, 8) == 7); ++static_assert(skippy_mtp_depth_for_step(8, 8) == 0); ++static_assert(skippy_mtp_depth_for_step(5, 1) == 0); ++static_assert(skippy_mtp_depth_for_step(5, 0) == 0); ++ ++struct skippy_mtp_depth_scope { ++ explicit skippy_mtp_depth_scope(llama_context * ctx) ++ : ctx(ctx), depth_count(ctx != nullptr ? llama_get_model(ctx)->hparams.n_layer_nextn : 0) { ++ if (ctx != nullptr) { ++ llama_set_nextn_layer_offset(ctx, 0); ++ } ++ } ++ ++ ~skippy_mtp_depth_scope() { ++ if (ctx != nullptr) { ++ llama_set_nextn_layer_offset(ctx, 0); ++ } ++ } ++ ++ void select(size_t draft_step) const { ++ if (ctx != nullptr && depth_count > 0) { ++ llama_set_nextn_layer_offset(ctx, skippy_mtp_depth_for_step(draft_step, depth_count)); ++ } ++ } ++ ++ llama_context * ctx; ++ uint32_t depth_count; ++}; ++ + static bool skippy_env_enabled(const char * name) { + const char * value = std::getenv(name); + if (value == nullptr || value[0] == '\0') { +@@ -2812,6 +2846,7 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + } + + llama_context * mtp_ctx = session->stage_model->mtp_ctx; ++ skippy_mtp_depth_scope mtp_depth_scope(mtp_ctx); + const bool mtp_shares_target_memory = llama_get_ctx_other(mtp_ctx) == session->ctx; + const int32_t n_embd = llama_model_n_embd(session->stage_model->model); + const size_t row_bytes = static_cast(n_embd)*sizeof(float); +@@ -3587,6 +3622,7 @@ static enum skippy_status skippy_mtp_propose_next( + } + + llama_context * mtp_ctx = session->stage_model->mtp_ctx; ++ skippy_mtp_depth_scope mtp_depth_scope(mtp_ctx); + const bool mtp_shares_target_memory = llama_get_ctx_other(mtp_ctx) == session->ctx; + const int32_t n_embd = llama_model_n_embd(session->stage_model->model); + if (session->mtp_pending_h.size() != static_cast(n_embd)) { +@@ -3606,6 +3642,7 @@ static enum skippy_status skippy_mtp_propose_next( + std::vector proposal_h = session->mtp_pending_h; + int32_t token_count = 0; + for (size_t i = 0; i < draft_limit; ++i) { ++ mtp_depth_scope.select(i); + llama_batch batch = { + /*n_tokens =*/ 1, + /*token =*/ &token, +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0053-Replay-multi-depth-MTP-caches-correctly.patch b/third_party/llama.cpp/patches/0053-Replay-multi-depth-MTP-caches-correctly.patch new file mode 100644 index 000000000..3515bb5b4 --- /dev/null +++ b/third_party/llama.cpp/patches/0053-Replay-multi-depth-MTP-caches-correctly.patch @@ -0,0 +1,368 @@ +From 8e80340a6f85a49e53320c79e3d353b7285b84db Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 14:58:19 +1000 +Subject: [PATCH 53/61] Replay multi-depth MTP caches correctly + +--- + src/models/step35.cpp | 12 ++- + src/skippy.cpp | 226 +++++++++++++++++++++++++++++++++--------- + 2 files changed, 189 insertions(+), 49 deletions(-) + +diff --git a/src/models/step35.cpp b/src/models/step35.cpp +index 4df4605bf..2d546c20f 100644 +--- a/src/models/step35.cpp ++++ b/src/models/step35.cpp +@@ -544,12 +544,16 @@ llama_model_step35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "mtp_post_ffn", il); + ++ // Pre-norm hidden state: used by the AR draft loop to seed the next MTP step. + ggml_tensor * inp_out_ids = build_inp_out_ids(); +- cur = ggml_get_rows(ctx0, cur, inp_out_ids); ++ ggml_tensor * h_nextn = cur; ++ if (cparams.embeddings_nextn_masked && inp_out_ids) { ++ h_nextn = ggml_get_rows(ctx0, h_nextn, inp_out_ids); ++ } ++ cb(h_nextn, "h_nextn", -1); ++ res->t_h_nextn = h_nextn; + +- // Pre-norm hidden state: used by the AR draft loop to seed the next MTP step. +- cb(cur, "h_nextn", -1); +- res->t_h_nextn = cur; ++ cur = ggml_get_rows(ctx0, cur, inp_out_ids); + + ggml_tensor * head_norm_w = layer.nextn.shared_head_norm + ? layer.nextn.shared_head_norm +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 1484f4f3b..52670efe8 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -1822,6 +1822,14 @@ struct skippy_mtp_depth_scope { + } + } + ++ bool chains_heads() const { ++ return depth_count > 1 && llama_get_ctx_other(ctx) == nullptr; ++ } ++ ++ bool chains_at_same_position() const { ++ return chains_heads() && llama_get_model(ctx)->arch == LLM_ARCH_INKLING; ++ } ++ + llama_context * ctx; + uint32_t depth_count; + }; +@@ -2831,6 +2839,11 @@ static void skippy_mtp_clear_session_state(skippy_session * session) { + } + } + ++static llama_token skippy_greedy_sample_context( ++ const llama_model * model, ++ llama_context * ctx, ++ int32_t index); ++ + static enum skippy_status skippy_mtp_sync_target_tokens( + skippy_session * session, + const llama_token * token_ids, +@@ -2848,6 +2861,8 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + llama_context * mtp_ctx = session->stage_model->mtp_ctx; + skippy_mtp_depth_scope mtp_depth_scope(mtp_ctx); + const bool mtp_shares_target_memory = llama_get_ctx_other(mtp_ctx) == session->ctx; ++ const bool chain_heads = mtp_depth_scope.chains_heads(); ++ const bool chain_at_same_position = mtp_depth_scope.chains_at_same_position(); + const int32_t n_embd = llama_model_n_embd(session->stage_model->model); + const size_t row_bytes = static_cast(n_embd)*sizeof(float); + if (session->mtp_pending_h.size() != static_cast(n_embd)) { +@@ -2872,7 +2887,15 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + + const bool reprime_prefix = !session->mtp_prefix_valid; + +- if (session->mtp_has_pending_draft) { ++ if (chain_heads && session->mtp_has_pending_draft) { ++ if (llama_memory_t memory = mtp_ctx->get_memory()) { ++ llama_memory_seq_rm(memory, session->seq_id, token_start, -1); ++ } ++ session->mtp_next_pos = std::min(session->mtp_next_pos, static_cast(token_start)); ++ session->mtp_has_pending_draft = false; ++ session->mtp_pending_draft_pos = 0; ++ session->mtp_pending_draft_token = -1; ++ } else if (session->mtp_has_pending_draft) { + const llama_pos token_end = token_start + static_cast(token_count); + if (session->mtp_pending_draft_pos >= token_start && + session->mtp_pending_draft_pos < token_end) { +@@ -2892,7 +2915,7 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + } + + size_t first_index = 0; +- if (session->mtp_next_pos > token_start) { ++ if (!chain_heads && session->mtp_next_pos > token_start) { + first_index = static_cast(std::min( + static_cast(token_count), + session->mtp_next_pos - token_start)); +@@ -2915,7 +2938,7 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + batch.pos[out_i] = token_start + static_cast(src_i); + batch.n_seq_id[out_i] = 1; + batch.seq_id[out_i][0] = session->seq_id; +- batch.logits[out_i] = 0; ++ batch.logits[out_i] = chain_at_same_position && out_i == n_decode - 1 ? 1 : 0; + + float * dst = batch.embd + static_cast(out_i)*n_embd; + if (src_i == 0) { +@@ -2933,8 +2956,46 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + } + } + +- skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config); +- const int32_t rc = llama_decode(mtp_ctx, batch); ++ int32_t rc = 0; ++ const uint32_t sync_depth_count = chain_heads ? mtp_depth_scope.depth_count : 1; ++ for (uint32_t depth = 0; depth < sync_depth_count; ++depth) { ++ if (chain_heads) { ++ if (llama_memory_t memory = mtp_ctx->get_memory()) { ++ llama_memory_seq_rm(memory, session->seq_id, batch.pos[0], -1); ++ } ++ mtp_depth_scope.select(depth); ++ } ++ ++ { ++ skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config); ++ rc = llama_decode(mtp_ctx, batch); ++ } ++ if (rc != 0 || !chain_heads || depth + 1 == sync_depth_count) { ++ break; ++ } ++ ++ if (!chain_at_same_position) { ++ continue; ++ } ++ ++ for (int32_t out_i = 0; out_i < n_decode; ++out_i) { ++ const float * h_next = llama_get_embeddings_nextn_ith(mtp_ctx, out_i); ++ if (h_next == nullptr) { ++ rc = -1; ++ break; ++ } ++ std::memcpy(batch.embd + static_cast(out_i)*n_embd, h_next, row_bytes); ++ } ++ if (rc != 0) { ++ break; ++ } ++ ++ for (int32_t out_i = 0; out_i + 1 < n_decode; ++out_i) { ++ batch.token[out_i] = batch.token[out_i + 1]; ++ } ++ batch.token[n_decode - 1] = skippy_greedy_sample_context( ++ llama_get_model(mtp_ctx), mtp_ctx, -1); ++ } + std::free(batch.token); + batch.token = nullptr; + llama_batch_free(batch); +@@ -3069,6 +3130,7 @@ static enum skippy_status skippy_verify_token_batch( + batch.logits[i] = 1; + } + ++ const llama_pos token_start = session->n_past; + enum skippy_status status = skippy_decode_batch( + session, + batch, +@@ -3076,6 +3138,9 @@ static enum skippy_status skippy_verify_token_batch( + SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, + out_error); + llama_batch_free(batch); ++ if (status == SKIPPY_STATUS_OK) { ++ status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); ++ } + if (status == SKIPPY_STATUS_OK) { + skippy_record_tokens(session, token_ids, token_count); + } +@@ -3083,7 +3148,7 @@ static enum skippy_status skippy_verify_token_batch( + } + + static llama_token skippy_greedy_sample_context( +- llama_model * model, ++ const llama_model * model, + llama_context * ctx, + int32_t index) { + if (model == nullptr || ctx == nullptr) { +@@ -3629,52 +3694,123 @@ static enum skippy_status skippy_mtp_propose_next( + return skippy_success(out_error); + } + +- const size_t draft_limit = std::min( ++ size_t draft_limit = std::min( + max_draft_tokens, + SKIPPY_NATIVE_MTP_MAX_DRAFT_TOKENS); ++ const bool chain_heads = mtp_depth_scope.chains_heads(); ++ const bool chain_at_same_position = mtp_depth_scope.chains_at_same_position(); ++ if (chain_heads) { ++ draft_limit = std::min(draft_limit, mtp_depth_scope.depth_count); ++ } + llama_token token = predicted_token; +- llama_pos pos = session->n_past; +- int32_t n_seq_id = 1; +- llama_seq_id seq_id = session->seq_id; +- llama_seq_id * seq_ids = &seq_id; +- int8_t logits = 1; + int64_t elapsed_us = 0; +- std::vector proposal_h = session->mtp_pending_h; + int32_t token_count = 0; +- for (size_t i = 0; i < draft_limit; ++i) { +- mtp_depth_scope.select(i); +- llama_batch batch = { +- /*n_tokens =*/ 1, +- /*token =*/ &token, +- /*embd =*/ proposal_h.data(), +- /*pos =*/ &pos, +- /*n_seq_id =*/ &n_seq_id, +- /*seq_id =*/ &seq_ids, +- /*logits =*/ &logits, +- }; ++ if (chain_heads && !chain_at_same_position) { ++ std::vector prefix_tokens = { predicted_token }; ++ prefix_tokens.reserve(draft_limit + 1); ++ std::vector prefix_hidden = session->mtp_pending_h; ++ prefix_hidden.reserve((draft_limit + 1)*static_cast(n_embd)); ++ ++ for (size_t depth = 0; depth < draft_limit; ++depth) { ++ const int32_t n_rows = static_cast(prefix_tokens.size()); ++ llama_batch batch = llama_batch_init(n_rows, n_embd, 1); ++ batch.token = static_cast(std::malloc(sizeof(llama_token)*static_cast(n_rows))); ++ if (batch.token == nullptr) { ++ llama_batch_free(batch); ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to allocate chained MTP proposal batch"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ batch.n_tokens = n_rows; ++ for (int32_t row = 0; row < n_rows; ++row) { ++ batch.token[row] = prefix_tokens[static_cast(row)]; ++ batch.pos[row] = session->n_past + row; ++ batch.n_seq_id[row] = 1; ++ batch.seq_id[row][0] = session->seq_id; ++ batch.logits[row] = row == n_rows - 1 ? 1 : 0; ++ std::memcpy( ++ batch.embd + static_cast(row)*n_embd, ++ prefix_hidden.data() + static_cast(row)*n_embd, ++ static_cast(n_embd)*sizeof(float)); ++ } + +- const int64_t t_start_us = ggml_time_us(); +- skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config); +- const int32_t rc = llama_decode(mtp_ctx, batch); +- elapsed_us += ggml_time_us() - t_start_us; +- if (rc != 0) { +- skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed for MTP sidecar proposal"); +- return SKIPPY_STATUS_RUNTIME_ERROR; +- } ++ if (llama_memory_t memory = mtp_ctx->get_memory()) { ++ llama_memory_seq_rm(memory, session->seq_id, session->n_past, -1); ++ } ++ mtp_depth_scope.select(depth); ++ const int64_t t_start_us = ggml_time_us(); ++ int32_t rc = 0; ++ { ++ skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config); ++ rc = llama_decode(mtp_ctx, batch); ++ } ++ elapsed_us += ggml_time_us() - t_start_us; ++ if (rc != 0) { ++ std::free(batch.token); ++ batch.token = nullptr; ++ llama_batch_free(batch); ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed for chained MTP proposal"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } + +- token = skippy_greedy_sample_context(session->stage_model->model, mtp_ctx, -1); +- if (out_mtp_draft != nullptr) { +- out_mtp_draft->token_ids[token_count] = token; +- } +- ++token_count; +- if (!mtp_shares_target_memory) { +- ++pos; ++ token = skippy_greedy_sample_context(llama_get_model(mtp_ctx), mtp_ctx, -1); ++ if (out_mtp_draft != nullptr) { ++ out_mtp_draft->token_ids[token_count] = token; ++ } ++ ++token_count; ++ const float * h_next = llama_get_embeddings_nextn_ith(mtp_ctx, n_rows - 1); ++ if (h_next != nullptr && depth + 1 < draft_limit) { ++ prefix_tokens.push_back(token); ++ prefix_hidden.insert(prefix_hidden.end(), h_next, h_next + n_embd); ++ } ++ std::free(batch.token); ++ batch.token = nullptr; ++ llama_batch_free(batch); ++ if (h_next == nullptr) { ++ break; ++ } + } +- const float * h_next = llama_get_embeddings_nextn_ith(mtp_ctx, -1); +- if (h_next == nullptr) { +- break; ++ } else { ++ llama_pos pos = session->n_past; ++ int32_t n_seq_id = 1; ++ llama_seq_id seq_id = session->seq_id; ++ llama_seq_id * seq_ids = &seq_id; ++ int8_t logits = 1; ++ std::vector proposal_h = session->mtp_pending_h; ++ for (size_t depth = 0; depth < draft_limit; ++depth) { ++ mtp_depth_scope.select(depth); ++ llama_batch batch = { ++ /*n_tokens =*/ 1, ++ /*token =*/ &token, ++ /*embd =*/ proposal_h.data(), ++ /*pos =*/ &pos, ++ /*n_seq_id =*/ &n_seq_id, ++ /*seq_id =*/ &seq_ids, ++ /*logits =*/ &logits, ++ }; ++ ++ const int64_t t_start_us = ggml_time_us(); ++ skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config); ++ const int32_t rc = llama_decode(mtp_ctx, batch); ++ elapsed_us += ggml_time_us() - t_start_us; ++ if (rc != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed for MTP sidecar proposal"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ ++ token = skippy_greedy_sample_context(llama_get_model(mtp_ctx), mtp_ctx, -1); ++ if (out_mtp_draft != nullptr) { ++ out_mtp_draft->token_ids[token_count] = token; ++ } ++ ++token_count; ++ if (!mtp_shares_target_memory && !chain_at_same_position) { ++ ++pos; ++ } ++ const float * h_next = llama_get_embeddings_nextn_ith(mtp_ctx, 0); ++ if (h_next == nullptr) { ++ break; ++ } ++ std::memcpy(proposal_h.data(), h_next, static_cast(n_embd)*sizeof(float)); + } +- std::memcpy(proposal_h.data(), h_next, static_cast(n_embd)*sizeof(float)); + } + + if (out_mtp_draft != nullptr && token_count > 0) { +@@ -4414,7 +4550,7 @@ static enum skippy_status skippy_finish_model_open( + } + if (stage_model->mtp_ctx != nullptr) { + llama_set_embeddings_nextn(stage_model->ctx, true, false); +- llama_set_embeddings_nextn(stage_model->mtp_ctx, true, true); ++ llama_set_embeddings_nextn(stage_model->mtp_ctx, true, false); + } else { + fprintf(stderr, "skippy: native MTP sidecar unavailable for this final stage; continuing without drafts\n"); + } +@@ -4679,7 +4815,7 @@ enum skippy_status skippy_model_attach_mtp_draft_model( + + target_model->mtp_model = draft_model; + llama_set_embeddings_nextn(target_model->ctx, true, false); +- llama_set_embeddings_nextn(target_model->mtp_ctx, true, true); ++ llama_set_embeddings_nextn(target_model->mtp_ctx, true, false); + return skippy_success(out_error); + } + +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0054-Add-Inkling-multi-depth-MTP-sidecars.patch b/third_party/llama.cpp/patches/0054-Add-Inkling-multi-depth-MTP-sidecars.patch new file mode 100644 index 000000000..d8f60501a --- /dev/null +++ b/third_party/llama.cpp/patches/0054-Add-Inkling-multi-depth-MTP-sidecars.patch @@ -0,0 +1,545 @@ +From a7a426fc1ca1c0fd28933bcd2bf49191e5041580 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 15:09:39 +1000 +Subject: [PATCH 54/61] Add Inkling multi-depth MTP sidecars + +--- + conversion/inkling.py | 78 ++++++++++++++++- + src/llama-context.cpp | 6 ++ + src/llama-model.cpp | 12 ++- + src/models/inkling.cpp | 188 ++++++++++++++++++++++++++++++----------- + src/skippy.cpp | 17 +++- + 5 files changed, 247 insertions(+), 54 deletions(-) + +diff --git a/conversion/inkling.py b/conversion/inkling.py +index 90ef65c3b..079842903 100644 +--- a/conversion/inkling.py ++++ b/conversion/inkling.py +@@ -1,5 +1,7 @@ + from __future__ import annotations + ++import re ++ + from typing import Callable, Iterable, TYPE_CHECKING + + if TYPE_CHECKING: +@@ -12,8 +14,15 @@ from .base import MmprojModel, ModelBase, TextModel, gguf, logger + class InklingModel(TextModel): + model_arch = gguf.MODEL_ARCH.INKLING + undo_permute = False ++ supports_mtp_export = True + +- _SKIP_PREFIXES = ("model.visual.", "model.audio.", "model.mtp.") ++ _SKIP_PREFIXES = ("model.visual.", "model.audio.") ++ _MTP_SHARED_NAMES = { ++ "model.llm.embed.weight", ++ "model.llm.embed_norm.weight", ++ "model.llm.norm.weight", ++ "model.llm.unembed.weight", ++ } + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) +@@ -66,9 +75,25 @@ class InklingModel(TextModel): + self.hparams["prefix_dense_intermediate_size"] = hp["dense_intermediate_size"] + + self._local_layer_flags = self._get_local_layer_flags() ++ mtp = hp.get("mtp_config") or {} ++ self._n_nextn = int( ++ mtp.get("num_nextn_predict_layers", hp.get("num_mtp_layers", 0)) or 0 ++ ) ++ if mtp.get("chain_hidden_post_norm", hp.get("chain_hidden_post_norm", False)): ++ raise NotImplementedError("Inkling MTP chain_hidden_post_norm=true is not supported") ++ mtp_local_ids = set(mtp.get("local_layer_ids", hp.get("mtp_local_layer_ids", [])) or []) ++ self._mtp_local_layer_flags = [i in mtp_local_ids for i in range(self._n_nextn)] ++ ++ if self.no_mtp: ++ self._n_nextn = 0 ++ self._mtp_local_layer_flags = [] ++ self.block_count = hp["num_hidden_layers"] + self._n_nextn ++ self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) ++ ++ all_local_layer_flags = self._local_layer_flags + self._mtp_local_layer_flags + self.hparams["num_key_value_heads"] = [ + hp["swa_num_key_value_heads"] if is_local else hp["num_key_value_heads"] +- for is_local in self._local_layer_flags ++ for is_local in all_local_layer_flags + ] + + def _get_local_layer_flags(self) -> list[bool]: +@@ -142,7 +167,10 @@ class InklingModel(TextModel): + raise ValueError(f"sliding_window must be positive, got {sliding_window}") + self.gguf_writer.add_sliding_window(sliding_window) + # true = local (swa) layer +- self.gguf_writer.add_sliding_window_pattern(self._local_layer_flags) ++ all_local_layer_flags = self._local_layer_flags + self._mtp_local_layer_flags ++ self.gguf_writer.add_sliding_window_pattern(all_local_layer_flags) ++ if self._n_nextn > 0: ++ self.gguf_writer.add_nextn_predict_layers(self._n_nextn) + + # no RoPE (arch-determined NONE); custom inkling.* keys per INKLING_DESIGN.md + arch = gguf.MODEL_ARCH_NAMES[self.model_arch] +@@ -157,6 +185,12 @@ class InklingModel(TextModel): + self.gguf_writer.add_uint32(f"{arch}.unpadded_vocab_size", hp["unpadded_vocab_size"]) + + logger.info(f"gguf: (inkling) swa pattern (true=local) = {self._local_layer_flags}") ++ if self._n_nextn > 0: ++ logger.info( ++ "gguf: (inkling) MTP depths = %d, local pattern = %s", ++ self._n_nextn, ++ self._mtp_local_layer_flags, ++ ) + logger.info(f"gguf: (inkling) unpadded_vocab_size = {hp['unpadded_vocab_size']}") + + @classmethod +@@ -165,6 +199,11 @@ class InklingModel(TextModel): + + if name.startswith(cls._SKIP_PREFIXES): + return None ++ is_mtp = name.startswith("model.mtp.") ++ if is_mtp and cls.no_mtp: ++ return None ++ if cls.mtp_only and not is_mtp and name not in cls._MTP_SHARED_NAMES: ++ return None + + name = name.replace("model.llm.", "model.") + # parameter has no ".weight"-style suffix in the checkpoint +@@ -180,6 +219,22 @@ class InklingModel(TextModel): + return gate, up + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: ++ match = re.match(r"^model\.mtp\.layers\.(\d+)\.(.*)$", name) ++ if match is not None: ++ mtp_idx = int(match.group(1)) ++ if mtp_idx >= self._n_nextn: ++ raise ValueError(f"MTP layer index {mtp_idx} >= declared count {self._n_nextn}") ++ rest = match.group(2) ++ rest = { ++ "embed_norm.weight": "enorm.weight", ++ "hidden_norm.weight": "hnorm.weight", ++ "input_proj.weight": "eh_proj.weight", ++ }.get(rest, rest.removeprefix("transformer_block.")) ++ bid = self.hparams["num_hidden_layers"] + mtp_idx ++ name = f"model.layers.{bid}.{rest}" ++ elif name.startswith("model.mtp."): ++ raise ValueError(f"unexpected Inkling MTP tensor {name!r}") ++ + # short convs: [C, 1, K] -> [C, K] (same layout as LFM2 shortconv.conv) + if name.endswith("_sconv.weight"): + data_torch = data_torch.squeeze(1) +@@ -235,6 +290,23 @@ class InklingModel(TextModel): + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + ++ def prepare_metadata(self, vocab_only: bool): ++ from_dir = self.fname_out.is_dir() ++ super().prepare_metadata(vocab_only=vocab_only) ++ if not self.mtp_only or not from_dir: ++ return ++ output_type: str = self.ftype.name.partition("_")[2] ++ fname_default: str = gguf.naming_convention( ++ self.metadata.name, ++ self.metadata.basename, ++ self.metadata.finetune, ++ self.metadata.version, ++ size_label=None, ++ output_type=output_type, ++ model_type=None, ++ ) ++ self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf" ++ + + @ModelBase.register("InklingForConditionalGeneration") + class InklingMmprojModel(MmprojModel): +diff --git a/src/llama-context.cpp b/src/llama-context.cpp +index c4e057aee..4b47903d3 100644 +--- a/src/llama-context.cpp ++++ b/src/llama-context.cpp +@@ -158,6 +158,12 @@ llama_context::llama_context( + } + } + ++ if (model.arch == LLM_ARCH_INKLING && ++ params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && ++ params.ctx_other != nullptr) { ++ cparams.ctx_other = params.ctx_other; ++ } ++ + // Initialize backend samplers here so they are part of the sampling graph + // before the reserve passes run later in this function. This avoids a later + // re-reserve when graph nodes change. +diff --git a/src/llama-model.cpp b/src/llama-model.cpp +index f95591860..4ffe8df0f 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -2131,10 +2131,20 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, + // layer filters, so pick the right one here + llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; + llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; +- if (arch == LLM_ARCH_FALCON_H1 || arch == LLM_ARCH_INKLING) { ++ if (arch == LLM_ARCH_FALCON_H1) { + // all layers have both an attention KV cache and a recurrent (conv) state + filter_attn = [&](uint32_t) { return true; }; + filter_recr = [&](uint32_t) { return true; }; ++ } else if (arch == LLM_ARCH_INKLING) { ++ // Inkling's MTP depths are appended to the trunk and each ++ // owns both attention KV and packed short-conv state. ++ if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_nextn > 0) { ++ filter_attn = [&](uint32_t il) { return il >= hparams.n_layer(); }; ++ filter_recr = [&](uint32_t il) { return il >= hparams.n_layer(); }; ++ } else { ++ filter_attn = [&](uint32_t il) { return il < hparams.n_layer(); }; ++ filter_recr = [&](uint32_t il) { return il < hparams.n_layer(); }; ++ } + } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { + filter_attn = [&](uint32_t il) { + return !hparams.is_recr(il) && hparams.n_ff(il) == 0; +diff --git a/src/models/inkling.cpp b/src/models/inkling.cpp +index 3d5f8ab56..d435a1f5d 100644 +--- a/src/models/inkling.cpp ++++ b/src/models/inkling.cpp +@@ -13,6 +13,9 @@ + void llama_model_inkling::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ++ ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); ++ GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all); ++ + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); +@@ -20,9 +23,9 @@ void llama_model_inkling::load_arch_hparams(llama_model_loader & ml) { + + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; // visible iff pos_q - pos_k < n_swa (includes self) +- ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer()); ++ ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer_all); + +- for (uint32_t il = 0; il < hparams.n_layer(); ++il) { ++ for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { + hparams.is_recr_impl[il] = 1; + } + +@@ -52,71 +55,91 @@ void llama_model_inkling::load_arch_hparams(llama_model_loader & ml) { + type = LLM_TYPE_UNKNOWN; + } + +-void llama_model_inkling::load_arch_tensors(llama_model_loader &) { ++void llama_model_inkling::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + ++ const bool mtp_only = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr; ++ const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; ++ const bool trunk_only = hparams.n_layer_nextn > 0 && ml.get_weight(mtp_probe.c_str()) == nullptr; ++ const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; ++ const int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; ++ + const int64_t head_dim = hparams.n_embd_head_k(); + const int64_t d_rel = hparams.inkling_d_rel; + const int64_t K = hparams.n_shortconv_l_cache; + const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_shexp = hparams.n_expert_shared; + +- tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); +- tok_norm = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD_NORM, "weight", 0), {n_embd}, 0); // bid 0: compute on the first layer's device +- output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); +- output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); ++ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, mtp_only ? TENSOR_NOT_REQUIRED : 0); ++ tok_norm = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD_NORM, "weight", 0), {n_embd}, mtp_only ? TENSOR_NOT_REQUIRED : 0); ++ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, mtp_only ? TENSOR_NOT_REQUIRED : 0); ++ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, mtp_only ? TENSOR_NOT_REQUIRED : 0); + +- for (int i = 0; i < n_layer; ++i) { ++ auto load_block = [&](int i, int flags, bool force_dense) { + auto & layer = layers[i]; + + const int64_t n_head_kv_i = hparams.n_head_kv(i); + const int64_t kvw = n_head_kv_i * head_dim; + const int64_t rel_extent = hparams.is_swa(i) ? hparams.inkling_rel_extent_swa : hparams.inkling_rel_extent; + +- layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); ++ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); + +- layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head*head_dim}, 0); +- layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, kvw}, 0); +- layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, kvw}, 0); +- layer.wr = create_tensor(tn(LLM_TENSOR_ATTN_R, "weight", i), {n_embd, n_head*d_rel}, 0); +- layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head*head_dim, n_embd}, 0); ++ layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head*head_dim}, flags); ++ layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, kvw}, flags); ++ layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, kvw}, flags); ++ layer.wr = create_tensor(tn(LLM_TENSOR_ATTN_R, "weight", i), {n_embd, n_head*d_rel}, flags); ++ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head*head_dim, n_embd}, flags); + +- layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {head_dim}, 0); +- layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {head_dim}, 0); ++ layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {head_dim}, flags); ++ layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {head_dim}, flags); + + // stored in checkpoint orientation [d_rel, E] -> gguf ne = [E, d_rel] +- layer.attn_rel_proj = create_tensor(tn(LLM_TENSOR_ATTN_REL_PROJ, "weight", i), {rel_extent, d_rel}, 0); ++ layer.attn_rel_proj = create_tensor(tn(LLM_TENSOR_ATTN_REL_PROJ, "weight", i), {rel_extent, d_rel}, flags); + +- layer.shortconv_k = create_tensor(tn(LLM_TENSOR_SHORTCONV_K, "weight", i), {K, kvw}, 0); +- layer.shortconv_v = create_tensor(tn(LLM_TENSOR_SHORTCONV_V, "weight", i), {K, kvw}, 0); +- layer.shortconv_attn = create_tensor(tn(LLM_TENSOR_SHORTCONV_ATTN, "weight", i), {K, n_embd}, 0); +- layer.shortconv_mlp = create_tensor(tn(LLM_TENSOR_SHORTCONV_MLP, "weight", i), {K, n_embd}, 0); ++ layer.shortconv_k = create_tensor(tn(LLM_TENSOR_SHORTCONV_K, "weight", i), {K, kvw}, flags); ++ layer.shortconv_v = create_tensor(tn(LLM_TENSOR_SHORTCONV_V, "weight", i), {K, kvw}, flags); ++ layer.shortconv_attn = create_tensor(tn(LLM_TENSOR_SHORTCONV_ATTN, "weight", i), {K, n_embd}, flags); ++ layer.shortconv_mlp = create_tensor(tn(LLM_TENSOR_SHORTCONV_MLP, "weight", i), {K, n_embd}, flags); + +- layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); +- layer.ffn_gscale = create_tensor(tn(LLM_TENSOR_FFN_GSCALE, "weight", i), {1}, 0); ++ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); ++ layer.ffn_gscale = create_tensor(tn(LLM_TENSOR_FFN_GSCALE, "weight", i), {1}, flags); + +- if (i < (int) hparams.n_layer_dense_lead) { ++ if (force_dense || i < (int) hparams.n_layer_dense_lead) { + const int64_t n_ff_i = hparams.n_ff(i); + +- layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff_i}, 0); +- layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff_i}, 0); +- layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff_i, n_embd}, 0); ++ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff_i}, flags); ++ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff_i}, flags); ++ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff_i, n_embd}, flags); + } else { + GGML_ASSERT(n_expert > 0 && n_expert_used > 0 && n_shexp > 0); + + // gate holds n_expert + n_shexp rows (incl. shared-expert sink logits) +- layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert + n_shexp}, 0); +- layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); ++ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert + n_shexp}, flags); ++ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags); + +- layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); +- layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); +- layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); ++ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, flags); ++ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, flags); ++ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); + + // shared experts stacked as an n_shexp bank, registered MUL_MAT_ID so the loader picks a mul_mat_id-capable buffer +- layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXPS, "weight", i), {n_embd, n_ff_exp, n_shexp}, 0); +- layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXPS, "weight", i), {n_embd, n_ff_exp, n_shexp}, 0); +- layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXPS, "weight", i), {n_ff_exp, n_embd, n_shexp}, 0); ++ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXPS, "weight", i), {n_embd, n_ff_exp, n_shexp}, flags); ++ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXPS, "weight", i), {n_embd, n_ff_exp, n_shexp}, flags); ++ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXPS, "weight", i), {n_ff_exp, n_embd, n_shexp}, flags); + } ++ }; ++ ++ for (int i = 0; i < n_layer; ++i) { ++ load_block(i, trunk_flags, false); ++ } ++ for (int i = n_layer; i < n_layer_all; ++i) { ++ load_block(i, mtp_flags, true); ++ auto & layer = layers[i]; ++ layer.nextn.eh_proj = create_tensor( ++ tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2*n_embd, n_embd}, mtp_flags); ++ layer.nextn.enorm = create_tensor( ++ tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags); ++ layer.nextn.hnorm = create_tensor( ++ tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags); + } + } + +@@ -198,8 +221,12 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + + const skippy_graph_filter & stage_filter = skippy_graph_get_filter(); + const bool stage_filtered = stage_filter.enabled; +- const int il_start = stage_filtered ? stage_filter.layer_start : 0; +- const int il_end = stage_filtered ? stage_filter.layer_end : n_layer; ++ const bool is_mtp = params.gtype == LLM_GRAPH_TYPE_DECODER_MTP; ++ const int mtp_il = n_layer + cparams.nextn_layer_offset; ++ GGML_ASSERT(!is_mtp || (cparams.nextn_layer_offset >= 0 && ++ cparams.nextn_layer_offset < (int) hparams.n_layer_nextn)); ++ const int il_start = is_mtp ? mtp_il : (stage_filtered ? stage_filter.layer_start : 0); ++ const int il_end = is_mtp ? mtp_il + 1 : (stage_filtered ? stage_filter.layer_end : n_layer); + + const int64_t head_dim = hparams.n_embd_head_k(); + const int64_t d_rel = hparams.inkling_d_rel; +@@ -322,7 +349,7 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + } + + // shared experts go through mul_mat_id: 2D views into a repacked/quantized 3D bank are invalid +- if (hparams.n_expert_shared > 0 && (uint32_t) il_end > hparams.n_layer_dense_lead) { ++ if (!is_mtp && hparams.n_expert_shared > 0 && (uint32_t) il_end > hparams.n_layer_dense_lead) { + inp->shexp_idx = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, hparams.n_expert_shared, n_tokens); + ggml_set_input(inp->shexp_idx); + ggml_set_name(inp->shexp_idx, "inkling_shexp_idx"); +@@ -638,18 +665,52 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + return moe_out; + }; + +- ggml_tensor * cur = build_inp_embd(stage_filtered && il_start > 0 ? nullptr : model.tok_embd); +- // mtmd embd rows arrive pre-normalized; embed_norm applies to text token lookups only +- if (ubatch.token && (!stage_filtered || stage_filter.include_embeddings)) { +- cur = build_norm(cur, model.tok_norm, NULL, LLM_NORM_RMS, -1); +- cb(cur, "inkling_embd_norm", -1); ++ const llama_model * shared_model = &model; ++ if (is_mtp && cparams.ctx_other != nullptr) { ++ const llama_model * target_model = llama_get_model(cparams.ctx_other); ++ if (target_model != nullptr && target_model->tok_embd != nullptr && ++ target_model->tok_norm != nullptr && target_model->output != nullptr) { ++ shared_model = target_model; ++ } ++ } ++ ++ ggml_tensor * cur = nullptr; ++ if (is_mtp) { ++ const auto & layer = model.layers[mtp_il]; ++ GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm); ++ GGML_ASSERT(shared_model->tok_embd && shared_model->tok_norm && shared_model->output); ++ ++ auto mtp_inp = std::make_unique(hparams.n_embd); ++ mtp_inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); ++ ggml_set_input(mtp_inp->tokens); ++ mtp_inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); ++ ggml_set_input(mtp_inp->embd); ++ ggml_set_name(mtp_inp->embd, "inkling_mtp_h_input"); ++ ++ ggml_tensor * tok_embd = ggml_get_rows(ctx0, shared_model->tok_embd, mtp_inp->tokens); ++ tok_embd = build_norm(tok_embd, shared_model->tok_norm, nullptr, LLM_NORM_RMS, -1); ++ tok_embd = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, mtp_il); ++ ggml_tensor * hidden = build_norm(mtp_inp->embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, mtp_il); ++ cur = build_lora_mm(layer.nextn.eh_proj, ggml_concat(ctx0, hidden, tok_embd, 0)); ++ cb(cur, "inkling_mtp_input", mtp_il); ++ res->add_input(std::move(mtp_inp)); + } else { +- cb(cur, "inkling_mm_embd", -1); ++ cur = build_inp_embd(stage_filtered && il_start > 0 ? nullptr : model.tok_embd); ++ // mtmd embd rows arrive pre-normalized; embed_norm applies to text token lookups only ++ if (ubatch.token && (!stage_filtered || stage_filter.include_embeddings)) { ++ cur = build_norm(cur, model.tok_norm, NULL, LLM_NORM_RMS, -1); ++ cb(cur, "inkling_embd_norm", -1); ++ } else { ++ cb(cur, "inkling_mm_embd", -1); ++ } + } + + ggml_build_forward_expand(gf, cur); + + for (int il = il_start; il < il_end; ++il) { ++ if (!is_mtp) { ++ res->t_layer_inp[il] = cur; ++ } + conv_rs_cur = build_rs(inp_hybrid->get_recr(), mctx_recr->get_r_l(il), n_embd_r, n_seqs); + + // h += attn_sconv(attn(attn_norm(h))) +@@ -664,7 +725,7 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + // h += mlp_sconv(mlp(mlp_norm(h))) + ggml_tensor * ffn_in = build_norm(cur, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); + cb(ffn_in, "inkling_ffn_norm", il); +- ggml_tensor * ffn_out = il < (int) hparams.n_layer_dense_lead ? ++ ggml_tensor * ffn_out = is_mtp || il < (int) hparams.n_layer_dense_lead ? + build_dense_ffn(ffn_in, il) : build_moe(ffn_in, il); + ffn_out = build_sconv(ffn_out, model.layers[il].shortconv_mlp, off_mlp, il); + cb(ffn_out, "inkling_ffn_sconv", il); +@@ -675,6 +736,31 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + cb(cur, "l_out", il); + } + ++ if (is_mtp) { ++ ggml_tensor * inp_out_ids = build_inp_out_ids(); ++ ggml_tensor * h_nextn = cur; ++ if (cparams.embeddings_nextn_masked && inp_out_ids) { ++ h_nextn = ggml_get_rows(ctx0, h_nextn, inp_out_ids); ++ } ++ cb(h_nextn, "h_nextn", -1); ++ res->t_h_nextn = h_nextn; ++ ++ if (inp_out_ids) { ++ cur = ggml_get_rows(ctx0, cur, inp_out_ids); ++ } ++ cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); ++ cur = build_lora_mm( ++ shared_model->output, cur, nullptr, ++ shared_model->output->type == GGML_TYPE_F32 ? GGML_PREC_F32_PEDANTIC : GGML_PREC_DEFAULT); ++ if (vocab_mask) { ++ cur = ggml_add(ctx0, cur, vocab_mask); ++ } ++ cb(cur, "result_output", -1); ++ res->t_logits = cur; ++ ggml_build_forward_expand(gf, cur); ++ return; ++ } ++ + if (stage_filtered && !stage_filter.include_output) { + cb(cur, "stage_boundary", il_end - 1); + res->t_embd = cur; +@@ -685,11 +771,19 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + // Conv states need every layer to see all tokens. Intermediate stages hand off the full + // activation sequence, and the final stage trims outputs only after its last local layer. + ggml_tensor * inp_out_ids = build_inp_out_ids(); ++ cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); ++ cb(cur, "result_norm_all", -1); ++ if (cparams.embeddings_nextn) { ++ ggml_tensor * h_nextn = cur; ++ if (cparams.embeddings_nextn_masked && inp_out_ids) { ++ h_nextn = ggml_get_rows(ctx0, h_nextn, inp_out_ids); ++ } ++ cb(h_nextn, "h_nextn", -1); ++ res->t_h_nextn = h_nextn; ++ } + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } +- +- cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 52670efe8..913c4d5ee 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -1823,7 +1823,7 @@ struct skippy_mtp_depth_scope { + } + + bool chains_heads() const { +- return depth_count > 1 && llama_get_ctx_other(ctx) == nullptr; ++ return depth_count > 1; + } + + bool chains_at_same_position() const { +@@ -2844,6 +2844,17 @@ static llama_token skippy_greedy_sample_context( + llama_context * ctx, + int32_t index); + ++static bool skippy_mtp_shares_target_memory( ++ const llama_context * mtp_ctx, ++ const llama_context * target_ctx) { ++ if (mtp_ctx == nullptr || target_ctx == nullptr || llama_get_ctx_other(const_cast(mtp_ctx)) != target_ctx) { ++ return false; ++ } ++ // Inkling uses ctx_other only to borrow shared target weights. Its eight ++ // attention/short-conv depth caches are independent of the target cache. ++ return llama_get_model(mtp_ctx)->arch != LLM_ARCH_INKLING; ++} ++ + static enum skippy_status skippy_mtp_sync_target_tokens( + skippy_session * session, + const llama_token * token_ids, +@@ -2860,7 +2871,7 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + + llama_context * mtp_ctx = session->stage_model->mtp_ctx; + skippy_mtp_depth_scope mtp_depth_scope(mtp_ctx); +- const bool mtp_shares_target_memory = llama_get_ctx_other(mtp_ctx) == session->ctx; ++ const bool mtp_shares_target_memory = skippy_mtp_shares_target_memory(mtp_ctx, session->ctx); + const bool chain_heads = mtp_depth_scope.chains_heads(); + const bool chain_at_same_position = mtp_depth_scope.chains_at_same_position(); + const int32_t n_embd = llama_model_n_embd(session->stage_model->model); +@@ -3688,7 +3699,7 @@ static enum skippy_status skippy_mtp_propose_next( + + llama_context * mtp_ctx = session->stage_model->mtp_ctx; + skippy_mtp_depth_scope mtp_depth_scope(mtp_ctx); +- const bool mtp_shares_target_memory = llama_get_ctx_other(mtp_ctx) == session->ctx; ++ const bool mtp_shares_target_memory = skippy_mtp_shares_target_memory(mtp_ctx, session->ctx); + const int32_t n_embd = llama_model_n_embd(session->stage_model->model); + if (session->mtp_pending_h.size() != static_cast(n_embd)) { + return skippy_success(out_error); +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0055-Carry-multimodal-Inkling-embeddings-into-MTP.patch b/third_party/llama.cpp/patches/0055-Carry-multimodal-Inkling-embeddings-into-MTP.patch new file mode 100644 index 000000000..842265636 --- /dev/null +++ b/third_party/llama.cpp/patches/0055-Carry-multimodal-Inkling-embeddings-into-MTP.patch @@ -0,0 +1,567 @@ +From 7412e7ae58b83de1864433828e4a786693e58b60 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 16:04:48 +1000 +Subject: [PATCH 55/61] Carry multimodal Inkling embeddings into MTP + +--- + include/skippy.h | 3 +- + include/skippy/common.h | 1 + + src/llama-context.cpp | 5 +- + src/llama-context.h | 2 + + src/llama-graph.cpp | 25 +++++- + src/llama-graph.h | 9 ++ + src/models/inkling.cpp | 20 +++-- + src/skippy.cpp | 183 +++++++++++++++++++++++++++++++++++----- + 8 files changed, 218 insertions(+), 30 deletions(-) + +diff --git a/include/skippy.h b/include/skippy.h +index e9d31f559..cc2382d37 100644 +--- a/include/skippy.h ++++ b/include/skippy.h +@@ -43,7 +43,8 @@ enum skippy_activation_layout { + + #define SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST (UINT64_C(1) << 0) + #define SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP (UINT64_C(1) << 1) +-#define SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K (UINT64_C(1) << 3) ++#define SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD (UINT64_C(1) << 2) ++#define SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K (UINT64_C(1) << 3) + + #define SKIPPY_GLM_DSA_POLICY_PROFILE_NONE 0 + #define SKIPPY_GLM_DSA_POLICY_PROFILE_V1 1 +diff --git a/include/skippy/common.h b/include/skippy/common.h +index 0145705b2..11f06bf3e 100644 +--- a/include/skippy/common.h ++++ b/include/skippy/common.h +@@ -55,6 +55,7 @@ enum skippy_feature { + SKIPPY_FEATURE_RUNTIME_EVENTS = 1 << 24, + SKIPPY_FEATURE_NATIVE_MTP_N1 = 1 << 25, + SKIPPY_FEATURE_NGRAM_CACHE_DRAFT = 1 << 26, ++ SKIPPY_FEATURE_INKLING_MTP_MM = 1 << 27, + }; + + enum skippy_status { +diff --git a/src/llama-context.cpp b/src/llama-context.cpp +index 4b47903d3..9d528761b 100644 +--- a/src/llama-context.cpp ++++ b/src/llama-context.cpp +@@ -4113,7 +4113,10 @@ int32_t llama_encode( + int32_t llama_decode( + llama_context * ctx, + llama_batch batch) { +- const int ret = ctx->decode(batch); ++ int ret = ctx->decode(batch); ++ if (ret == 0) { ++ ret = skippy_external_decode_observe(ctx, batch); ++ } + if (ret != 0 && ret != 1) { + LLAMA_LOG_ERROR("%s: failed to decode, ret = %d\n", __func__, ret); + } +diff --git a/src/llama-context.h b/src/llama-context.h +index 42442879e..20c121b24 100644 +--- a/src/llama-context.h ++++ b/src/llama-context.h +@@ -20,6 +20,8 @@ class llama_batch_allocr; + class llama_io_read_i; + class llama_io_write_i; + ++int skippy_external_decode_observe(llama_context * ctx, const llama_batch & batch); ++ + // "memory" as in abstract memory for the context + struct llama_memory_i; + struct llama_memory_context_i; +diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp +index 36bf32193..27c5129be 100644 +--- a/src/llama-graph.cpp ++++ b/src/llama-graph.cpp +@@ -26,6 +26,7 @@ static thread_local skippy_activation_tokens g_skippy_activation_tokens; + static thread_local skippy_activation_rwkv7_v_first g_skippy_rwkv7_v_first; + static thread_local skippy_activation_gemma3n_altup g_skippy_gemma3n_altup; + static thread_local skippy_activation_glm_dsa_top_k g_skippy_glm_dsa_top_k; ++static thread_local skippy_activation_mtp_embeddings g_skippy_mtp_embeddings; + + void skippy_graph_set_filter(const skippy_graph_filter & filter) { + g_skippy_graph_filter = filter; +@@ -87,6 +88,18 @@ const skippy_activation_glm_dsa_top_k & skippy_graph_get_glm_dsa_top_k() { + return g_skippy_glm_dsa_top_k; + } + ++void skippy_graph_set_mtp_embeddings(const skippy_activation_mtp_embeddings & values) { ++ g_skippy_mtp_embeddings = values; ++} ++ ++void skippy_graph_clear_mtp_embeddings() { ++ g_skippy_mtp_embeddings = {}; ++} ++ ++const skippy_activation_mtp_embeddings & skippy_graph_get_mtp_embeddings() { ++ return g_skippy_mtp_embeddings; ++} ++ + // dedup helpers + + static ggml_tensor * build_attn_inp_kq_mask( +@@ -156,15 +169,21 @@ bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) { + + void llm_graph_input_embd_h::set_input(const llama_ubatch * ubatch) { + const int64_t n_tokens = ubatch->n_tokens; ++ const skippy_activation_mtp_embeddings & mtp_embeddings = skippy_graph_get_mtp_embeddings(); + + if (ubatch->token) { + ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens)); + } else { +- // note: mtmd embedding input goes through here + GGML_ASSERT(ubatch->embd); + GGML_ASSERT(n_embd == embd->ne[0]); ++ const float * values = ubatch->embd; ++ if (mtp_embeddings.values != nullptr) { ++ GGML_ASSERT(mtp_embeddings.token_count == ubatch->n_tokens); ++ GGML_ASSERT(mtp_embeddings.n_embd == n_embd); ++ values = mtp_embeddings.values; ++ } + +- ggml_backend_tensor_set(embd, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(h)); ++ ggml_backend_tensor_set(embd, values, 0, n_tokens*n_embd*ggml_element_size(embd)); + } + + // TODO: extend llama_ubatch to differentiate between token embeddings and hidden states +@@ -181,7 +200,7 @@ bool llm_graph_input_embd_h::can_reuse(const llm_graph_params & params) { + bool res = true; + + res &= (!params.ubatch.token) || (tokens && tokens->buffer && tokens->ne[0] == params.ubatch.n_tokens); +- res &= (!params.ubatch.embd) || (embd && embd->buffer && embd->ne[1] == params.ubatch.n_tokens); ++ res &= (params.ubatch.token) || (embd && embd->buffer && embd->ne[1] == params.ubatch.n_tokens); + res &= (!params.ubatch.embd) || (h && h->buffer && h->ne[1] == params.ubatch.n_tokens); + + return res; +diff --git a/src/llama-graph.h b/src/llama-graph.h +index e74a58a56..6e66f7fef 100644 +--- a/src/llama-graph.h ++++ b/src/llama-graph.h +@@ -78,6 +78,12 @@ struct skippy_activation_glm_dsa_top_k { + llama_pos pos_start = 0; + }; + ++struct skippy_activation_mtp_embeddings { ++ const float * values = nullptr; ++ uint32_t token_count = 0; ++ uint32_t n_embd = 0; ++}; ++ + void skippy_graph_set_filter(const skippy_graph_filter & filter); + void skippy_graph_clear_filter(); + const skippy_graph_filter & skippy_graph_get_filter(); +@@ -93,6 +99,9 @@ const skippy_activation_gemma3n_altup & skippy_graph_get_gemma3n_altup(); + void skippy_graph_set_glm_dsa_top_k(const skippy_activation_glm_dsa_top_k & values); + void skippy_graph_clear_glm_dsa_top_k(); + const skippy_activation_glm_dsa_top_k & skippy_graph_get_glm_dsa_top_k(); ++void skippy_graph_set_mtp_embeddings(const skippy_activation_mtp_embeddings & values); ++void skippy_graph_clear_mtp_embeddings(); ++const skippy_activation_mtp_embeddings & skippy_graph_get_mtp_embeddings(); + + // certain models (typically multi-modal) can produce different types of graphs + enum llm_graph_type { +diff --git a/src/models/inkling.cpp b/src/models/inkling.cpp +index d435a1f5d..c054f1d98 100644 +--- a/src/models/inkling.cpp ++++ b/src/models/inkling.cpp +@@ -680,17 +680,25 @@ llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_par + GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm); + GGML_ASSERT(shared_model->tok_embd && shared_model->tok_norm && shared_model->output); + +- auto mtp_inp = std::make_unique(hparams.n_embd); ++ auto mtp_inp = std::make_unique(hparams.n_embd); + mtp_inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(mtp_inp->tokens); + mtp_inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(mtp_inp->embd); +- ggml_set_name(mtp_inp->embd, "inkling_mtp_h_input"); +- +- ggml_tensor * tok_embd = ggml_get_rows(ctx0, shared_model->tok_embd, mtp_inp->tokens); +- tok_embd = build_norm(tok_embd, shared_model->tok_norm, nullptr, LLM_NORM_RMS, -1); ++ ggml_set_name(mtp_inp->embd, "inkling_mtp_embd_input"); ++ mtp_inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); ++ ggml_set_input(mtp_inp->h); ++ ggml_set_name(mtp_inp->h, "inkling_mtp_h_input"); ++ ++ ggml_tensor * tok_embd = nullptr; ++ if (ubatch.token) { ++ tok_embd = ggml_get_rows(ctx0, shared_model->tok_embd, mtp_inp->tokens); ++ tok_embd = build_norm(tok_embd, shared_model->tok_norm, nullptr, LLM_NORM_RMS, -1); ++ } else { ++ tok_embd = mtp_inp->embd; ++ } + tok_embd = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, mtp_il); +- ggml_tensor * hidden = build_norm(mtp_inp->embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, mtp_il); ++ ggml_tensor * hidden = build_norm(mtp_inp->h, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, mtp_il); + cur = build_lora_mm(layer.nextn.eh_proj, ggml_concat(ctx0, hidden, tok_embd, 0)); + cb(cur, "inkling_mtp_input", mtp_il); + res->add_input(std::move(mtp_inp)); +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 913c4d5ee..cc50cc501 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -149,8 +149,12 @@ struct skippy_session { + bool mtp_has_pending_draft = false; + llama_pos mtp_pending_draft_pos = 0; + llama_token mtp_pending_draft_token = -1; ++ bool external_mtp_embeddings_available = false; ++ std::vector external_mtp_embeddings; + }; + ++static thread_local skippy_session * g_skippy_external_decode_session = nullptr; ++ + struct skippy_tensor_meta { + std::string name; + int32_t layer_index = -1; +@@ -1287,6 +1291,7 @@ static bool skippy_is_gemma3n_activation_model(const skippy_session * session) { + return session->stage_model->model->arch == LLM_ARCH_GEMMA3N; + } + ++ + static bool skippy_is_glm_dsa_activation_model(const skippy_session * session) { + if (session == nullptr || session->stage_model == nullptr || session->stage_model->model == nullptr) { + return false; +@@ -1405,6 +1410,14 @@ static bool skippy_glm_dsa_stage_starts_in_consumer_group(const skippy_session * + return skippy_glm_dsa_layer_starts_consumer_group(session, session->stage_model->config.layer_start); + } + ++static bool skippy_is_inkling_activation_model(const skippy_session * session) { ++ if (session == nullptr || session->stage_model == nullptr || session->stage_model->model == nullptr) { ++ return false; ++ } ++ return session->stage_model->model->arch == LLM_ARCH_INKLING; ++} ++ ++ + static size_t skippy_activation_hidden_bytes(const skippy_session * session, size_t token_count) { + if (session == nullptr || session->stage_model == nullptr || session->stage_model->model == nullptr) { + return 0; +@@ -1530,6 +1543,14 @@ static uint64_t skippy_output_activation_flags( + skippy_glm_dsa_layer_starts_consumer_group(session, session->stage_model->config.layer_end)) { + return SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K; + } ++ if (skippy_emits_activation_frame(session) && skippy_is_inkling_activation_model(session)) { ++ if (input_desc != nullptr && (input_desc->flags & SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0) { ++ return SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD; ++ } ++ if (session->stage_model->config.layer_start == 0 && session->external_mtp_embeddings_available) { ++ return SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD; ++ } ++ } + if (!skippy_emits_activation_frame(session) || !skippy_is_rwkv7_activation_model(session)) { + return 0; + } +@@ -1559,6 +1580,9 @@ static size_t skippy_activation_payload_bytes( + if ((flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0) { + payload_bytes += hidden_bytes; + } ++ if ((flags & SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0) { ++ payload_bytes += hidden_bytes; ++ } + return payload_bytes; + } + +@@ -1609,6 +1633,7 @@ static enum skippy_status skippy_validate_frame_input_sequences( + const uint64_t supported_flags = + SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST | + SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP | ++ SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD | + SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K; + if ((input_desc->flags & ~supported_flags) != 0) { + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "activation frame has unsupported sideband flags"); +@@ -1649,6 +1674,15 @@ static enum skippy_status skippy_validate_frame_input_sequences( + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "non-first RWKV7 runtime slices require v_first activation sideband"); + return SKIPPY_STATUS_INVALID_ARGUMENT; + } ++ if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0 && !skippy_is_inkling_activation_model(session)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "Inkling MTP embedding sideband is only valid for Inkling stages"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0 && ++ (input_desc->flags & (SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP | SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST)) != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "Inkling MTP embedding sideband cannot be combined with other activation sidebands"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } + + const size_t hidden_bytes = skippy_activation_hidden_bytes(session, expected_token_count); + if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0) { +@@ -2855,13 +2889,33 @@ static bool skippy_mtp_shares_target_memory( + return llama_get_model(mtp_ctx)->arch != LLM_ARCH_INKLING; + } + +-static enum skippy_status skippy_mtp_sync_target_tokens( ++struct skippy_mtp_embeddings_scope { ++ skippy_mtp_embeddings_scope(const float * values, uint32_t token_count, uint32_t n_embd) { ++ if (values != nullptr) { ++ skippy_graph_set_mtp_embeddings({ values, token_count, n_embd }); ++ enabled = true; ++ } ++ } ++ ++ ~skippy_mtp_embeddings_scope() { ++ if (enabled) { ++ skippy_graph_clear_mtp_embeddings(); ++ } ++ } ++ ++ bool enabled = false; ++}; ++ ++static enum skippy_status skippy_mtp_sync_target_inputs( + skippy_session * session, + const llama_token * token_ids, ++ const float * input_embeddings, + size_t token_count, + llama_pos token_start, + struct skippy_error ** out_error) { +- if (!skippy_mtp_available(session) || token_ids == nullptr || token_count == 0) { ++ if (!skippy_mtp_available(session) || ++ (token_ids == nullptr && input_embeddings == nullptr) || ++ token_count == 0) { + return skippy_success(out_error); + } + if (token_count > static_cast(std::numeric_limits::max())) { +@@ -2911,7 +2965,7 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + if (session->mtp_pending_draft_pos >= token_start && + session->mtp_pending_draft_pos < token_end) { + const size_t draft_index = static_cast(session->mtp_pending_draft_pos - token_start); +- if (token_ids[draft_index] == session->mtp_pending_draft_token) { ++ if (token_ids != nullptr && token_ids[draft_index] == session->mtp_pending_draft_token) { + session->mtp_next_pos = std::max(session->mtp_next_pos, session->mtp_pending_draft_pos + 1); + } else { + if (llama_memory_t memory = mtp_ctx->get_memory()) { +@@ -2935,17 +2989,21 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + const int32_t n_decode = static_cast(token_count - first_index); + if (!mtp_shares_target_memory && n_decode > 0) { + llama_batch batch = llama_batch_init(n_decode, n_embd, 1); +- batch.token = static_cast(std::malloc(sizeof(llama_token)*static_cast(n_decode))); +- if (batch.token == nullptr) { +- llama_batch_free(batch); +- skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to allocate MTP token batch"); +- return SKIPPY_STATUS_RUNTIME_ERROR; ++ if (token_ids != nullptr) { ++ batch.token = static_cast(std::malloc(sizeof(llama_token)*static_cast(n_decode))); ++ if (batch.token == nullptr) { ++ llama_batch_free(batch); ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to allocate MTP token batch"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } + } + batch.n_tokens = n_decode; + + for (int32_t out_i = 0; out_i < n_decode; ++out_i) { + const size_t src_i = first_index + static_cast(out_i); +- batch.token[out_i] = token_ids[src_i]; ++ if (batch.token != nullptr) { ++ batch.token[out_i] = token_ids[src_i]; ++ } + batch.pos[out_i] = token_start + static_cast(src_i); + batch.n_seq_id[out_i] = 1; + batch.seq_id[out_i][0] = session->seq_id; +@@ -2969,6 +3027,12 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + + int32_t rc = 0; + const uint32_t sync_depth_count = chain_heads ? mtp_depth_scope.depth_count : 1; ++ const float * mtp_embeddings = input_embeddings != nullptr ? ++ input_embeddings + first_index*static_cast(n_embd) : nullptr; ++ skippy_mtp_embeddings_scope mtp_embeddings_scope( ++ mtp_embeddings, ++ static_cast(n_decode), ++ static_cast(n_embd)); + for (uint32_t depth = 0; depth < sync_depth_count; ++depth) { + if (chain_heads) { + if (llama_memory_t memory = mtp_ctx->get_memory()) { +@@ -3001,14 +3065,19 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + break; + } + ++ if (batch.token == nullptr) { ++ continue; ++ } + for (int32_t out_i = 0; out_i + 1 < n_decode; ++out_i) { + batch.token[out_i] = batch.token[out_i + 1]; + } + batch.token[n_decode - 1] = skippy_greedy_sample_context( + llama_get_model(mtp_ctx), mtp_ctx, -1); + } +- std::free(batch.token); +- batch.token = nullptr; ++ if (batch.token != nullptr) { ++ std::free(batch.token); ++ batch.token = nullptr; ++ } + llama_batch_free(batch); + if (rc != 0) { + skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed for MTP sidecar sync"); +@@ -3030,6 +3099,41 @@ static enum skippy_status skippy_mtp_sync_target_tokens( + return skippy_success(out_error); + } + ++int skippy_external_decode_observe(llama_context * ctx, const llama_batch & batch) { ++ skippy_session * session = g_skippy_external_decode_session; ++ if (session == nullptr || ctx != session->ctx || batch.n_tokens <= 0 || ++ session->stage_model == nullptr || session->stage_model->model == nullptr || ++ session->stage_model->model->arch != LLM_ARCH_INKLING) { ++ return 0; ++ } ++ session->external_mtp_embeddings_available = batch.embd != nullptr; ++ if (batch.embd != nullptr) { ++ const size_t value_count = static_cast(batch.n_tokens)* ++ static_cast(llama_model_n_embd(session->stage_model->model)); ++ session->external_mtp_embeddings.assign(batch.embd, batch.embd + value_count); ++ } else { ++ session->external_mtp_embeddings.clear(); ++ } ++ ++ const llama_pos token_start = batch.pos != nullptr ? batch.pos[0] : session->n_past; ++ skippy_error * error = nullptr; ++ const skippy_status status = skippy_mtp_sync_target_inputs( ++ session, ++ batch.token, ++ batch.embd, ++ static_cast(batch.n_tokens), ++ token_start, ++ &error); ++ if (status == SKIPPY_STATUS_OK) { ++ return 0; ++ } ++ ++ fprintf(stderr, "skippy: external Inkling MTP sync failed: %s\n", ++ error != nullptr && error->message != nullptr ? error->message : "unknown error"); ++ skippy_error_free(error); ++ return -1; ++} ++ + static void skippy_record_tokens( + skippy_session * session, + const llama_token * token_ids, +@@ -3082,7 +3186,7 @@ static enum skippy_status skippy_decode_tokens( + }; + enum skippy_status status = skippy_decode_batch(session, batch, 1, glm_dsa_phase_hint, out_error); + if (status == SKIPPY_STATUS_OK) { +- status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, pos, out_error); ++ status = skippy_mtp_sync_target_inputs(session, token_ids, nullptr, token_count, pos, out_error); + } + if (status == SKIPPY_STATUS_OK) { + skippy_record_tokens(session, token_ids, token_count); +@@ -3104,7 +3208,7 @@ static enum skippy_status skippy_decode_tokens( + enum skippy_status status = skippy_decode_batch(session, batch, token_count, glm_dsa_phase_hint, out_error); + llama_batch_free(batch); + if (status == SKIPPY_STATUS_OK) { +- status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); ++ status = skippy_mtp_sync_target_inputs(session, token_ids, nullptr, token_count, token_start, out_error); + } + if (status == SKIPPY_STATUS_OK) { + skippy_record_tokens(session, token_ids, token_count); +@@ -3150,7 +3254,7 @@ static enum skippy_status skippy_verify_token_batch( + out_error); + llama_batch_free(batch); + if (status == SKIPPY_STATUS_OK) { +- status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); ++ status = skippy_mtp_sync_target_inputs(session, token_ids, nullptr, token_count, token_start, out_error); + } + if (status == SKIPPY_STATUS_OK) { + skippy_record_tokens(session, token_ids, token_count); +@@ -4010,6 +4114,28 @@ static enum skippy_status skippy_copy_output_activation_frame( + if (output_desc != nullptr) { + output_desc->payload_bytes = actual_payload_bytes; + } ++ if ((output_flags & SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0) { ++ uint8_t * sideband_output = static_cast(output_payload) + hidden_bytes; ++ const skippy_runtime_config & config = session->stage_model->config; ++ if (config.layer_start == 0) { ++ if (session->external_mtp_embeddings.size()*sizeof(float) < hidden_bytes) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "Inkling MTP input embedding sideband was not available"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ std::memcpy(sideband_output, session->external_mtp_embeddings.data(), hidden_bytes); ++ } else { ++ if (input_desc == nullptr || ++ input_payload == nullptr || ++ (input_desc->flags & SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD) == 0 || ++ input_desc->payload_bytes < hidden_bytes * 2) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "Inkling downstream slice cannot forward a missing MTP input embedding sideband"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ std::memcpy(sideband_output, static_cast(input_payload) + hidden_bytes, hidden_bytes); ++ } ++ session->external_mtp_embeddings_available = false; ++ session->external_mtp_embeddings.clear(); ++ } + return skippy_success(out_error); + } + +@@ -4129,8 +4255,12 @@ static enum skippy_status skippy_decode_activation_frame( + token_count, + glm_dsa_phase_hint, + out_error); +- if (status == SKIPPY_STATUS_OK && token_ids != nullptr) { +- status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); ++ const float * mtp_input_embeddings = ++ (input_desc->flags & SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0 ? ++ reinterpret_cast(static_cast(input_payload) + hidden_bytes) : nullptr; ++ if (status == SKIPPY_STATUS_OK && (token_ids != nullptr || mtp_input_embeddings != nullptr)) { ++ status = skippy_mtp_sync_target_inputs( ++ session, token_ids, mtp_input_embeddings, token_count, token_start, out_error); + } + if (status == SKIPPY_STATUS_OK && token_ids != nullptr) { + skippy_record_tokens(session, token_ids, token_count); +@@ -4213,8 +4343,12 @@ static enum skippy_status skippy_verify_activation_frame( + if (!alias_input_payload) { + llama_batch_free(batch); + } +- if (status == SKIPPY_STATUS_OK && token_ids != nullptr) { +- status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); ++ const float * mtp_input_embeddings = ++ (input_desc->flags & SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0 ? ++ reinterpret_cast(static_cast(input_payload) + hidden_bytes) : nullptr; ++ if (status == SKIPPY_STATUS_OK && (token_ids != nullptr || mtp_input_embeddings != nullptr)) { ++ status = skippy_mtp_sync_target_inputs( ++ session, token_ids, mtp_input_embeddings, token_count, token_start, out_error); + } + return status; + } +@@ -4251,7 +4385,8 @@ uint64_t skippy_abi_features(void) { + SKIPPY_FEATURE_RUNTIME_EVENTS | + SKIPPY_FEATURE_BACKEND_DEVICES | + SKIPPY_FEATURE_NATIVE_MTP_N1 | +- SKIPPY_FEATURE_NGRAM_CACHE_DRAFT; ++ SKIPPY_FEATURE_NGRAM_CACHE_DRAFT | ++ SKIPPY_FEATURE_INKLING_MTP_MM; + } + + void skippy_error_free(struct skippy_error * error) { +@@ -4940,6 +5075,12 @@ enum skippy_status skippy_session_begin_external_decode( + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session is required"); + return SKIPPY_STATUS_INVALID_ARGUMENT; + } ++ if (g_skippy_external_decode_session != nullptr && g_skippy_external_decode_session != session) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "another external decode session is already active on this thread"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ session->external_mtp_embeddings_available = false; ++ session->external_mtp_embeddings.clear(); + + const skippy_runtime_config & config = session->stage_model->config; + if (config.filter_tensors_on_load) { +@@ -4951,6 +5092,7 @@ enum skippy_status skippy_session_begin_external_decode( + filter.include_output = config.include_output; + skippy_graph_set_filter(filter); + } ++ g_skippy_external_decode_session = session; + + return skippy_success(out_error); + } +@@ -4962,6 +5104,9 @@ enum skippy_status skippy_session_end_external_decode( + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session is required"); + return SKIPPY_STATUS_INVALID_ARGUMENT; + } ++ if (g_skippy_external_decode_session == session) { ++ g_skippy_external_decode_session = nullptr; ++ } + skippy_graph_clear_filter(); + return skippy_success(out_error); + } +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0056-Size-recurrent-memory-for-appended-MTP-layers.patch b/third_party/llama.cpp/patches/0056-Size-recurrent-memory-for-appended-MTP-layers.patch new file mode 100644 index 000000000..afbc226ac --- /dev/null +++ b/third_party/llama.cpp/patches/0056-Size-recurrent-memory-for-appended-MTP-layers.patch @@ -0,0 +1,46 @@ +From 326edb5713500457c0a82fbb450d84fcf3549cb0 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 20:20:45 +1000 +Subject: [PATCH 56/61] Size recurrent memory for appended MTP layers + +--- + src/llama-memory-recurrent.cpp | 9 +++++---- + 1 file changed, 5 insertions(+), 4 deletions(-) + +diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp +index ef82eb976..4f1f1df8d 100644 +--- a/src/llama-memory-recurrent.cpp ++++ b/src/llama-memory-recurrent.cpp +@@ -26,7 +26,8 @@ llama_memory_recurrent::llama_memory_recurrent( + uint32_t n_seq_max, + uint32_t n_rs_seq, + const layer_filter_cb & filter) : hparams(model.hparams), n_seq_max(n_seq_max) { +- const int32_t n_layer = hparams.n_layer(); ++ // MTP layers append recurrent state beyond the trunk layer count. ++ const int32_t n_layer = hparams.n_layer_all; + + head = 0; + size = mem_size; +@@ -865,7 +866,7 @@ void llama_memory_recurrent::state_write_meta(llama_io_write_i & io, const std:: + + void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std::vector> & cell_ranges) const { + const uint32_t s_trans = 0; +- const uint32_t n_layer = hparams.n_layer(); ++ const uint32_t n_layer = hparams.n_layer_all; + + io.write(&s_trans, sizeof(s_trans)); + io.write(&n_layer, sizeof(n_layer)); +@@ -1049,8 +1050,8 @@ bool llama_memory_recurrent::state_read_data(llama_io_read_i & io, uint32_t cell + io.read(&s_trans, sizeof(s_trans)); + io.read(&n_layer, sizeof(n_layer)); + +- if (n_layer != hparams.n_layer()) { +- LLAMA_LOG_ERROR("%s: mismatched layer count (%u instead of %u)\n", __func__, n_layer, hparams.n_layer()); ++ if (n_layer != hparams.n_layer_all) { ++ LLAMA_LOG_ERROR("%s: mismatched layer count (%u instead of %u)\n", __func__, n_layer, hparams.n_layer_all); + return false; + } + if (cell_count > size) { +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0057-Complete-Inkling-and-GLM-metadata-integration.patch b/third_party/llama.cpp/patches/0057-Complete-Inkling-and-GLM-metadata-integration.patch new file mode 100644 index 000000000..95ee48520 --- /dev/null +++ b/third_party/llama.cpp/patches/0057-Complete-Inkling-and-GLM-metadata-integration.patch @@ -0,0 +1,143 @@ +From b01316df2de211495a1754050a1643509ae03b7a Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Mon, 27 Jul 2026 12:01:26 +1000 +Subject: [PATCH 57/61] Complete Inkling and GLM metadata integration + +--- + common/stage-chat.cpp | 6 ++++++ + src/llama-arch.cpp | 2 ++ + src/llama-arch.h | 2 ++ + src/llama-model.cpp | 1 + + src/models/glm-dsa.cpp | 46 +++++++++++++++++++++++++++++++++++------- + 5 files changed, 50 insertions(+), 7 deletions(-) + +diff --git a/common/stage-chat.cpp b/common/stage-chat.cpp +index 3298ed669..16ce255af 100644 +--- a/common/stage-chat.cpp ++++ b/common/stage-chat.cpp +@@ -178,6 +178,9 @@ enum skippy_status skippy_apply_chat_template( + inputs.chat_template_kwargs["thinking_enabled"] = value; + inputs.chat_template_kwargs["enable_think"] = value; + inputs.chat_template_kwargs["think_enabled"] = value; ++ if (!enable_thinking) { ++ inputs.chat_template_kwargs["reasoning_effort"] = "0"; ++ } + } + inputs.messages.reserve(message_count); + +@@ -282,6 +285,9 @@ enum skippy_status skippy_apply_chat_template_json( + inputs.chat_template_kwargs["thinking_enabled"] = value; + inputs.chat_template_kwargs["enable_think"] = value; + inputs.chat_template_kwargs["think_enabled"] = value; ++ if (!enable_thinking) { ++ inputs.chat_template_kwargs["reasoning_effort"] = "0"; ++ } + } + if (reasoning_format_name != nullptr && reasoning_format_name[0] != '\0') { + inputs.reasoning_format = common_reasoning_format_from_name(reasoning_format_name); +diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp +index a489e1ad5..836df068d 100644 +--- a/src/llama-arch.cpp ++++ b/src/llama-arch.cpp +@@ -258,6 +258,8 @@ static const std::map LLM_KV_NAMES = { + { LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" }, + { LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" }, + { LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" }, ++ { LLM_KV_ATTENTION_INDEXER_TOP_K_FREQUENCY, "%s.attention.indexer.top_k_frequency" }, ++ { LLM_KV_ATTENTION_INDEXER_SKIP_TOP_K_OFFSET, "%s.attention.indexer.skip_top_k_offset" }, + { LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" }, + { LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" }, + { LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" }, +diff --git a/src/llama-arch.h b/src/llama-arch.h +index 74446fdaf..0033c7612 100644 +--- a/src/llama-arch.h ++++ b/src/llama-arch.h +@@ -263,6 +263,8 @@ enum llm_kv { + LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, + LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, + LLM_KV_ATTENTION_INDEXER_TOP_K, ++ LLM_KV_ATTENTION_INDEXER_TOP_K_FREQUENCY, ++ LLM_KV_ATTENTION_INDEXER_SKIP_TOP_K_OFFSET, + LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, + LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, + LLM_KV_ATTENTION_INDEXER_TYPES, +diff --git a/src/llama-model.cpp b/src/llama-model.cpp +index 4ffe8df0f..70c105e6a 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -1137,6 +1137,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { + std::fill(hparams.is_swa_impl.begin(), hparams.is_swa_impl.end(), 0); + std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), llm_arch_is_recurrent(ml.get_arch()) ? 1 : 0); + std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 0); ++ std::fill(hparams.indexer_types.begin(), hparams.indexer_types.end(), -1); + + std::fill(hparams.xielu_alpha_n.begin(), hparams.xielu_alpha_n.end(), 0.0f); + std::fill(hparams.xielu_alpha_p.begin(), hparams.xielu_alpha_p.end(), 0.0f); +diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp +index a88aae2cf..57c457725 100644 +--- a/src/models/glm-dsa.cpp ++++ b/src/models/glm-dsa.cpp +@@ -52,6 +52,36 @@ void llama_model_glm_dsa::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + ++ const bool has_indexer_top_k_freq = ++ ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K_FREQUENCY, hparams.indexer_top_k_freq, false); ++ const bool has_indexer_skip_top_k_offset = ++ ml.get_key(LLM_KV_ATTENTION_INDEXER_SKIP_TOP_K_OFFSET, hparams.indexer_skip_top_k_offset, false); ++ if (has_indexer_top_k_freq && hparams.indexer_top_k_freq == 0) { ++ throw std::runtime_error("GLM_DSA attention.indexer.top_k_frequency must be positive when present"); ++ } ++ if (has_indexer_top_k_freq && !has_indexer_skip_top_k_offset) { ++ throw std::runtime_error("GLM_DSA attention.indexer.skip_top_k_offset is required when top_k_frequency is present"); ++ } ++ ++ std::vector indexer_types; ++ if (ml.get_arr(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types, false)) { ++ if (indexer_types.size() != hparams.n_layer()) { ++ throw std::runtime_error("GLM_DSA attention.indexer.types length must match effective decoder layer count"); ++ } ++ hparams.indexer_types_present = true; ++ for (size_t il = 0; il < indexer_types.size(); ++il) { ++ if (indexer_types[il] == "full") { ++ hparams.indexer_types[il] = 1; ++ hparams.is_indexer_full_impl[il] = 1; ++ } else if (indexer_types[il] == "shared") { ++ hparams.indexer_types[il] = 0; ++ hparams.is_indexer_full_impl[il] = 0; ++ } else { ++ throw std::runtime_error("GLM_DSA attention.indexer.types values must be \"full\" or \"shared\""); ++ } ++ } ++ } ++ + // Expert gating function (GLM-4.5 uses sigmoid) + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); + if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) { +@@ -62,14 +92,16 @@ void llama_model_glm_dsa::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all"); + +- // BC for GLM 5, 5.1 (full indexers) without indexer_types metadata +- const bool is_pre_5_2 = hparams.n_ctx_train < 1048576; +- if (is_pre_5_2) { +- std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1); +- } else { +- hparams.is_indexer_full_impl = GLM_5_2_DEFAULT_INDEXER_TYPES; ++ if (!hparams.indexer_types_present) { ++ // BC for GLM 5, 5.1 and numeric role metadata. ++ const bool is_pre_5_2 = hparams.n_ctx_train < 1048576; ++ if (is_pre_5_2) { ++ std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1); ++ } else { ++ hparams.is_indexer_full_impl = GLM_5_2_DEFAULT_INDEXER_TYPES; ++ } ++ ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false); + } +- ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false); + + switch (hparams.n_layer()) { + case 78: type = LLM_TYPE_744B_A40B; break; +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0058-skippy-balance-filtered-stages-across-devices.patch b/third_party/llama.cpp/patches/0058-skippy-balance-filtered-stages-across-devices.patch new file mode 100644 index 000000000..45de95528 --- /dev/null +++ b/third_party/llama.cpp/patches/0058-skippy-balance-filtered-stages-across-devices.patch @@ -0,0 +1,76 @@ +From 77251842c0f5c2f288ac0a2f9039a41478ca6370 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 29 Jul 2026 00:43:18 +1000 +Subject: [PATCH 58/61] skippy: balance filtered stages across devices + +--- + src/llama-model-loader.cpp | 4 ++++ + src/llama-model-loader.h | 1 + + src/llama-model.cpp | 18 +++++++++++++++++- + 3 files changed, 22 insertions(+), 1 deletion(-) + +diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp +index ec530d62e..1ac8a9f93 100644 +--- a/src/llama-model-loader.cpp ++++ b/src/llama-model-loader.cpp +@@ -29,6 +29,10 @@ void llama_model_loader_clear_stage_filter() { + g_skippy_filter = {}; + } + ++llama_model_loader_stage_filter llama_model_loader_get_stage_filter() { ++ return g_skippy_filter; ++} ++ + bool llama_model_loader_last_tensor_filtered() { + return g_skippy_last_tensor_filtered; + } +diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h +index da58661f5..988fcf944 100644 +--- a/src/llama-model-loader.h ++++ b/src/llama-model-loader.h +@@ -31,6 +31,7 @@ struct llama_model_loader_stage_filter { + + void llama_model_loader_set_stage_filter(const llama_model_loader_stage_filter & filter); + void llama_model_loader_clear_stage_filter(); ++llama_model_loader_stage_filter llama_model_loader_get_stage_filter(); + bool llama_model_loader_last_tensor_filtered(); + + enum llama_fver { +diff --git a/src/llama-model.cpp b/src/llama-model.cpp +index 70c105e6a..0e45dc460 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -1316,13 +1316,29 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { + + const int i_gpu_start = std::max(n_layer_all + 1 - n_gpu_layers, 0); + const int act_gpu_layers = devices.empty() ? 0 : std::min(n_gpu_layers, n_layer_all + 1); ++ const auto stage_filter = llama_model_loader_get_stage_filter(); ++ const int stage_gpu_start = std::max(stage_filter.layer_start, i_gpu_start); ++ const int stage_gpu_end = std::min(stage_filter.layer_end, i_gpu_start + act_gpu_layers); ++ const int stage_gpu_layers = std::max(stage_gpu_end - stage_gpu_start, 0); ++ if (stage_filter.enabled && stage_gpu_layers > 0 && n_devices() > 1) { ++ LLAMA_LOG_INFO( ++ "load_tensors: distributing filtered stage layers %d..%d across %zu devices\n", ++ stage_gpu_start, ++ stage_gpu_end, ++ n_devices()); ++ } + auto get_layer_buft_list = [&](int il) -> llama_model::impl::layer_dev { + const bool is_swa = il < n_layer_all && hparams.is_swa(il); + if (il < i_gpu_start || (il - i_gpu_start) >= act_gpu_layers) { + LLAMA_LOG_DEBUG("load_tensors: layer %3d assigned to device %s, is_swa = %d\n", il, ggml_backend_dev_name(cpu_dev), is_swa); + return {cpu_dev, &pimpl->cpu_buft_list}; + } +- const int layer_gpu = std::upper_bound(splits.begin(), splits.begin() + n_devices(), float(il - i_gpu_start)/act_gpu_layers) - splits.begin(); ++ const bool is_filtered_stage_layer = stage_filter.enabled && ++ il >= stage_gpu_start && il < stage_gpu_end && stage_gpu_layers > 0; ++ const float layer_fraction = is_filtered_stage_layer ? ++ float(il - stage_gpu_start)/stage_gpu_layers : ++ float(il - i_gpu_start)/act_gpu_layers; ++ const int layer_gpu = std::upper_bound(splits.begin(), splits.begin() + n_devices(), layer_fraction) - splits.begin(); + auto * dev = devices.at(layer_gpu).dev; + LLAMA_LOG_DEBUG("load_tensors: layer %3d assigned to device %s, is_swa = %d\n", il, ggml_backend_dev_name(dev), is_swa); + return {dev, &pimpl->gpu_buft_list.at(dev)}; +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0059-Preserve-hybrid-verify-window-state-across-trim.patch b/third_party/llama.cpp/patches/0059-Preserve-hybrid-verify-window-state-across-trim.patch new file mode 100644 index 000000000..351b1ce7e --- /dev/null +++ b/third_party/llama.cpp/patches/0059-Preserve-hybrid-verify-window-state-across-trim.patch @@ -0,0 +1,501 @@ +From 57280d642a5806d9464577bfbbc276d5afcab79b Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Tue, 28 Jul 2026 19:12:14 +1000 +Subject: [PATCH 59/61] Preserve hybrid verify-window state across trim + +Checkpoint only partial recurrent state, trim hybrid attention directly, and replay the accepted prefix so pipelined speculative windows can recover without copying the full long-context attention cache. + +Based on the recurrent checkpoint work from the Inkling branch. + +Assisted-by: codex +--- + src/skippy.cpp | 360 +++++++++++++++++++++++++++++++++++++++++++++---- + 1 file changed, 332 insertions(+), 28 deletions(-) + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index cc50cc501..9da45701e 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -9,6 +9,7 @@ + #include "llama-ext.h" + #include "llama-graph.h" + #include "llama-kv-cache.h" ++#include "llama-kv-cache-iswa.h" + #include "llama-kv-cache-dsa.h" + #include "llama-memory-hybrid.h" + #include "llama-memory-hybrid-iswa.h" +@@ -124,6 +125,17 @@ struct skippy_model { + std::vector> lane_resident_prefix_tokens; + }; + ++struct skippy_verify_checkpoint { ++ bool valid = false; ++ int32_t token_start = 0; ++ size_t token_count = 0; ++ std::vector state; ++ std::vector token_ids; ++ bool has_activation_input = false; ++ skippy_activation_desc input_desc = {}; ++ std::vector input_payload; ++}; ++ + struct skippy_session { + skippy_model * stage_model = nullptr; + llama_context * ctx = nullptr; +@@ -151,6 +163,8 @@ struct skippy_session { + llama_token mtp_pending_draft_token = -1; + bool external_mtp_embeddings_available = false; + std::vector external_mtp_embeddings; ++ bool reset_required = false; ++ std::vector verify_checkpoints; + }; + + static thread_local skippy_session * g_skippy_external_decode_session = nullptr; +@@ -250,7 +264,15 @@ static enum skippy_status skippy_success(skippy_error ** out_error) { + } + return SKIPPY_STATUS_OK; + } +- ++static enum skippy_status skippy_require_usable_session( ++ const skippy_session * session, ++ skippy_error ** out_error) { ++ if (session != nullptr && session->reset_required) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "session requires reset after failed verify-window recovery"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ return SKIPPY_STATUS_OK; ++} + static int32_t skippy_stage_layer_count(const llama_model * model) { + if (model == nullptr) { + return 0; +@@ -1590,6 +1612,76 @@ static bool skippy_has_activation_payload(const skippy_activation_desc * desc, c + return desc != nullptr && desc->payload_bytes > 0 && payload != nullptr; + } + ++static bool skippy_memory_needs_verify_checkpoint(llama_memory_t memory) { ++ return dynamic_cast(memory) != nullptr || ++ dynamic_cast(memory) != nullptr || ++ dynamic_cast(memory) != nullptr; ++} ++ ++static enum skippy_status skippy_checkpoint_verify_window( ++ skippy_session * session, ++ const llama_token * token_ids, ++ size_t token_count, ++ const skippy_activation_desc * input_desc, ++ const void * input_payload, ++ skippy_error ** out_error) { ++ enum skippy_status usable_status = skippy_require_usable_session(session, out_error); ++ if (usable_status != SKIPPY_STATUS_OK) { ++ return usable_status; ++ } ++ llama_memory_t memory = session->ctx->get_memory(); ++ if (!skippy_memory_needs_verify_checkpoint(memory)) { ++ return skippy_success(out_error); ++ } ++ ++ session->ctx->synchronize(); ++ const size_t state_size = llama_state_seq_get_size_ext( ++ session->ctx, ++ session->seq_id, ++ LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); ++ if (state_size == 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to size verify-window state checkpoint"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ ++ skippy_verify_checkpoint checkpoint; ++ checkpoint.state.resize(state_size); ++ const size_t written = llama_state_seq_get_data_ext( ++ session->ctx, ++ checkpoint.state.data(), ++ checkpoint.state.size(), ++ session->seq_id, ++ LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); ++ if (written != checkpoint.state.size()) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to checkpoint verify-window sequence state"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ ++ checkpoint.valid = true; ++ checkpoint.token_start = session->n_past; ++ checkpoint.token_count = token_count; ++ if (token_ids != nullptr) { ++ checkpoint.token_ids.assign(token_ids, token_ids + token_count); ++ } ++ checkpoint.has_activation_input = skippy_has_activation_payload(input_desc, input_payload); ++ if (checkpoint.has_activation_input) { ++ checkpoint.input_desc = *input_desc; ++ const uint8_t * bytes = static_cast(input_payload); ++ checkpoint.input_payload.assign(bytes, bytes + input_desc->payload_bytes); ++ } ++ ++ constexpr size_t max_verify_checkpoints = 64; ++ if (session->verify_checkpoints.size() >= max_verify_checkpoints) { ++ skippy_set_error( ++ out_error, ++ SKIPPY_STATUS_UNSUPPORTED, ++ "verify-window pipeline depth exceeds supported checkpoint retention"); ++ return SKIPPY_STATUS_UNSUPPORTED; ++ } ++ session->verify_checkpoints.push_back(std::move(checkpoint)); ++ return skippy_success(out_error); ++} ++ + static enum skippy_status skippy_validate_frame_input_sequences( + skippy_session * session, + const skippy_activation_desc * input_desc, +@@ -1797,7 +1889,10 @@ static enum skippy_status skippy_decode_batch( + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session and at least one token are required"); + return SKIPPY_STATUS_INVALID_ARGUMENT; + } +- ++ enum skippy_status usable_status = skippy_require_usable_session(session, out_error); ++ if (usable_status != SKIPPY_STATUS_OK) { ++ return usable_status; ++ } + skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config, glm_dsa_phase_hint); + const char * phase_log = getenv("SKIPPY_GLM_DSA_LOG_DIRECT_SPARSE_DECISIONS"); + if (phase_log != nullptr && atoi(phase_log) != 0) { +@@ -4353,6 +4448,169 @@ static enum skippy_status skippy_verify_activation_frame( + return status; + } + ++static std::vector skippy_verify_prefix_activation_payload( ++ const skippy_session * session, ++ const skippy_verify_checkpoint & checkpoint, ++ size_t token_count) { ++ const uint64_t flags = checkpoint.input_desc.flags; ++ if ((flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0) { ++ const size_t prefix_bytes = skippy_activation_payload_bytes(session, token_count, flags); ++ return std::vector( ++ checkpoint.input_payload.begin(), ++ checkpoint.input_payload.begin() + static_cast(prefix_bytes)); ++ } ++ ++ const size_t source_hidden_bytes = skippy_activation_hidden_bytes(session, checkpoint.token_count); ++ const size_t prefix_hidden_bytes = skippy_activation_hidden_bytes(session, token_count); ++ const uint32_t glm_dsa_top_k = skippy_glm_dsa_top_k_count_from_desc(session, &checkpoint.input_desc); ++ const size_t prefix_glm_dsa_bytes = (flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0 ? ++ skippy_glm_dsa_top_k_bytes_for_count(token_count, glm_dsa_top_k) : 0; ++ const size_t prefix_rwkv7_bytes = (flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0 ? ++ prefix_hidden_bytes : 0; ++ std::vector payload(prefix_hidden_bytes + prefix_glm_dsa_bytes + prefix_rwkv7_bytes); ++ std::memcpy(payload.data(), checkpoint.input_payload.data(), prefix_hidden_bytes); ++ ++ size_t source_offset = source_hidden_bytes; ++ size_t prefix_offset = prefix_hidden_bytes; ++ if (prefix_glm_dsa_bytes > 0) { ++ std::memcpy( ++ payload.data() + prefix_offset, ++ checkpoint.input_payload.data() + source_offset, ++ prefix_glm_dsa_bytes); ++ source_offset += skippy_glm_dsa_top_k_bytes_for_count(checkpoint.token_count, glm_dsa_top_k); ++ prefix_offset += prefix_glm_dsa_bytes; ++ } ++ if (prefix_rwkv7_bytes > 0) { ++ std::memcpy( ++ payload.data() + prefix_offset, ++ checkpoint.input_payload.data() + source_offset, ++ prefix_rwkv7_bytes); ++ } ++ return payload; ++} ++ ++static bool skippy_trim_verify_attention( ++ llama_memory_t memory, ++ llama_seq_id seq_id, ++ llama_pos token_start) { ++ if (auto * iswa = dynamic_cast(memory)) { ++ return iswa->seq_rm(seq_id, token_start, -1); ++ } ++ if (auto * hybrid = dynamic_cast(memory)) { ++ return hybrid->get_mem_attn()->seq_rm(seq_id, token_start, -1); ++ } ++ if (auto * hybrid_iswa = dynamic_cast(memory)) { ++ return hybrid_iswa->get_mem_attn()->seq_rm(seq_id, token_start, -1); ++ } ++ return true; ++} ++ ++static void skippy_mark_verify_restore_failed(skippy_session * session) { ++ session->reset_required = true; ++ session->verify_checkpoints.clear(); ++ skippy_mtp_clear_session_state(session); ++} ++ ++static enum skippy_status skippy_fail_verify_restore( ++ skippy_session * session, ++ enum skippy_status status, ++ const char * message, ++ skippy_error ** out_error) { ++ skippy_mark_verify_restore_failed(session); ++ skippy_set_error(out_error, status, message); ++ return status; ++} ++ ++static enum skippy_status skippy_restore_verify_prefix( ++ skippy_session * session, ++ uint64_t token_count, ++ bool * out_restored, ++ skippy_error ** out_error) { ++ *out_restored = false; ++ const skippy_verify_checkpoint * checkpoint = nullptr; ++ for (auto it = session->verify_checkpoints.rbegin(); it != session->verify_checkpoints.rend(); ++it) { ++ const uint64_t token_start = static_cast(it->token_start); ++ const uint64_t token_end = token_start + it->token_count; ++ if (it->valid && token_count >= token_start && token_count < token_end) { ++ checkpoint = &*it; ++ break; ++ } ++ } ++ if (checkpoint == nullptr) { ++ return skippy_success(out_error); ++ } ++ ++ const size_t read = llama_state_seq_set_data_ext( ++ session->ctx, ++ checkpoint->state.data(), ++ checkpoint->state.size(), ++ session->seq_id, ++ LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); ++ if (read != checkpoint->state.size()) { ++ return skippy_fail_verify_restore( ++ session, ++ SKIPPY_STATUS_RUNTIME_ERROR, ++ "failed to restore verify-window sequence state", ++ out_error); ++ } ++ ++ llama_memory_t memory = session->ctx->get_memory(); ++ if (!skippy_trim_verify_attention(memory, session->seq_id, checkpoint->token_start)) { ++ return skippy_fail_verify_restore( ++ session, ++ SKIPPY_STATUS_RUNTIME_ERROR, ++ "failed to trim verify-window attention suffix", ++ out_error); ++ } ++ ++ const uint64_t token_start = static_cast(checkpoint->token_start); ++ session->n_past = checkpoint->token_start; ++ const size_t accepted_count = static_cast(token_count - token_start); ++ if (accepted_count > 0 && checkpoint->token_ids.empty()) { ++ return skippy_fail_verify_restore( ++ session, ++ SKIPPY_STATUS_RUNTIME_ERROR, ++ "verify-window checkpoint is missing token ids", ++ out_error); ++ } ++ ++ enum skippy_status status = SKIPPY_STATUS_OK; ++ if (accepted_count > 0 && checkpoint->has_activation_input) { ++ skippy_activation_desc input_desc = checkpoint->input_desc; ++ std::vector input_payload = skippy_verify_prefix_activation_payload( ++ session, ++ *checkpoint, ++ accepted_count); ++ input_desc.token_count = static_cast(accepted_count); ++ input_desc.payload_bytes = input_payload.size(); ++ status = skippy_decode_activation_frame( ++ session, ++ &input_desc, ++ input_payload.data(), ++ checkpoint->token_ids.data(), ++ nullptr, ++ 0, ++ accepted_count, ++ false, ++ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, ++ out_error); ++ } else if (accepted_count > 0) { ++ status = skippy_decode_tokens( ++ session, ++ checkpoint->token_ids.data(), ++ accepted_count, ++ false, ++ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, ++ out_error); ++ } ++ if (status != SKIPPY_STATUS_OK) { ++ skippy_mark_verify_restore_failed(session); ++ return status; ++ } ++ *out_restored = true; ++ return SKIPPY_STATUS_OK; ++} ++ + extern "C" { + + struct skippy_abi_version skippy_abi_version(void) { +@@ -5132,6 +5390,7 @@ enum skippy_status skippy_session_set_position( + session->signal_history.resize(static_cast(n_past)); + } + skippy_mtp_clear_session_state(session); ++ session->verify_checkpoints.clear(); + return skippy_success(out_error); + } + +@@ -5225,6 +5484,8 @@ enum skippy_status skippy_session_reset( + session->signal_history.clear(); + skippy_clear_chat_sampling(session); + skippy_mtp_clear_session_state(session); ++ session->reset_required = false; ++ session->verify_checkpoints.clear(); + session->ctx->synchronize(); + return skippy_success(out_error); + } +@@ -5552,7 +5813,17 @@ enum skippy_status skippy_verify_tokens( + return SKIPPY_STATUS_INVALID_ARGUMENT; + } + +- enum skippy_status status = skippy_verify_token_batch(session, token_ids, token_count, out_error); ++ enum skippy_status status = skippy_checkpoint_verify_window( ++ session, ++ token_ids, ++ token_count, ++ nullptr, ++ nullptr, ++ out_error); ++ if (status != SKIPPY_STATUS_OK) { ++ return status; ++ } ++ status = skippy_verify_token_batch(session, token_ids, token_count, out_error); + if (status == SKIPPY_STATUS_OK) { + const int32_t n_tokens = static_cast(token_count); + for (int32_t i = 0; i < n_tokens; ++i) { +@@ -5613,7 +5884,8 @@ static enum skippy_status skippy_prefill_chunk_frame_impl( + return status; + } + +- if (skippy_is_filtered(session) && session->stage_model->config.layer_start > 0) { ++ const bool activation_input = skippy_is_filtered(session) && session->stage_model->config.layer_start > 0; ++ if (activation_input) { + status = skippy_decode_activation_frame( + session, + input_desc, +@@ -6261,7 +6533,19 @@ enum skippy_status skippy_verify_tokens_frame_sampled( + } + } + +- if (skippy_is_filtered(session) && session->stage_model->config.layer_start > 0) { ++ const bool activation_input = skippy_is_filtered(session) && session->stage_model->config.layer_start > 0; ++ status = skippy_checkpoint_verify_window( ++ session, ++ token_ids, ++ token_count, ++ activation_input ? input_desc : nullptr, ++ activation_input ? input_payload : nullptr, ++ out_error); ++ if (status != SKIPPY_STATUS_OK) { ++ return status; ++ } ++ ++ if (activation_input) { + status = session->stage_model->config.include_output ? + skippy_verify_activation_frame(session, input_desc, input_payload, token_ids, token_count, out_error) : + skippy_decode_activation_frame( +@@ -6380,6 +6664,7 @@ static void skippy_update_session_state_after_import( + session->signal_history.resize(static_cast(session->n_past)); + } + skippy_mtp_clear_session_state(session); ++ session->verify_checkpoints.clear(); + } + + enum skippy_status skippy_export_state( +@@ -6798,6 +7083,7 @@ enum skippy_status skippy_import_kv_page( + desc->token_start + desc->token_count, + static_cast(std::numeric_limits::max())))); + skippy_mtp_clear_session_state(session); ++ session->verify_checkpoints.clear(); + session->ctx->synchronize(); + + return skippy_success(out_error); +@@ -6826,30 +7112,46 @@ enum skippy_status skippy_trim_session( + skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "runtime memory is unavailable"); + return SKIPPY_STATUS_RUNTIME_ERROR; + } +- const llama_pos p0 = static_cast(token_count); +- if (auto * hybrid = dynamic_cast(memory)) { +- if (!hybrid->seq_rm(session->seq_id, p0, -1)) { +- skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim hybrid memory suffix"); +- return SKIPPY_STATUS_RUNTIME_ERROR; +- } +- } else if (auto * hybrid_iswa = dynamic_cast(memory)) { +- if (!hybrid_iswa->seq_rm(session->seq_id, p0, -1)) { +- skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim hybrid ISWA memory suffix"); +- return SKIPPY_STATUS_RUNTIME_ERROR; +- } +- } else if (auto * dsa = dynamic_cast(memory)) { +- if (!dsa->seq_rm(session->seq_id, p0, -1)) { +- skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim GLM-DSA KV suffix"); +- return SKIPPY_STATUS_RUNTIME_ERROR; +- } +- } else if (auto * kv = dynamic_cast(memory)) { +- if (!kv->seq_rm(session->seq_id, p0, -1)) { +- skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim native KV suffix"); +- return SKIPPY_STATUS_RUNTIME_ERROR; ++ bool restored_verify_prefix = false; ++ enum skippy_status status = skippy_restore_verify_prefix( ++ session, ++ token_count, ++ &restored_verify_prefix, ++ out_error); ++ if (status != SKIPPY_STATUS_OK) { ++ return status; ++ } ++ if (!restored_verify_prefix) { ++ const llama_pos p0 = static_cast(token_count); ++ if (auto * hybrid = dynamic_cast(memory)) { ++ if (!hybrid->seq_rm(session->seq_id, p0, -1)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim hybrid memory suffix"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ } else if (auto * hybrid_iswa = dynamic_cast(memory)) { ++ if (!hybrid_iswa->seq_rm(session->seq_id, p0, -1)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim hybrid ISWA memory suffix"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ } else if (auto * iswa = dynamic_cast(memory)) { ++ if (!iswa->seq_rm(session->seq_id, p0, -1)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim interleaved-SWA KV suffix"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ } else if (auto * dsa = dynamic_cast(memory)) { ++ if (!dsa->seq_rm(session->seq_id, p0, -1)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim GLM-DSA KV suffix"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ } else if (auto * kv = dynamic_cast(memory)) { ++ if (!kv->seq_rm(session->seq_id, p0, -1)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim native KV suffix"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ } else if (dynamic_cast(memory) == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "runtime memory type is not supported for trim"); ++ return SKIPPY_STATUS_UNSUPPORTED; + } +- } else if (dynamic_cast(memory) == nullptr) { +- skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "runtime memory type is not supported for trim"); +- return SKIPPY_STATUS_UNSUPPORTED; + } + session->n_past = static_cast(token_count); + if (session->token_history.size() > token_count) { +@@ -6860,6 +7162,7 @@ enum skippy_status skippy_trim_session( + session->signal_history.resize(static_cast(token_count)); + } + skippy_mtp_clear_session_state(session); ++ session->verify_checkpoints.clear(); + session->ctx->synchronize(); + + return skippy_success(out_error); +@@ -7012,6 +7315,7 @@ enum skippy_status skippy_session_restore_prefix( + } + session->signal_history.clear(); + skippy_mtp_clear_session_state(session); ++ session->verify_checkpoints.clear(); + session->ctx->synchronize(); + return skippy_success(out_error); + } +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0060-skippy-retire-accepted-verify-checkpoints-exactly.patch b/third_party/llama.cpp/patches/0060-skippy-retire-accepted-verify-checkpoints-exactly.patch new file mode 100644 index 000000000..be34a3662 --- /dev/null +++ b/third_party/llama.cpp/patches/0060-skippy-retire-accepted-verify-checkpoints-exactly.patch @@ -0,0 +1,160 @@ +From 60f5917b3718cc74b2298571b4fac924e9ecf081 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 29 Jul 2026 11:12:07 +1000 +Subject: [PATCH 60/61] skippy: retire accepted verify checkpoints exactly + +--- + include/skippy.h | 6 +++++ + include/skippy/common.h | 2 +- + src/skippy-verify-checkpoint.h | 24 +++++++++++++++++ + src/skippy.cpp | 21 +++++++++++++++ + tests/CMakeLists.txt | 1 + + ...st-skippy-verify-checkpoint-retirement.cpp | 26 +++++++++++++++++++ + 6 files changed, 79 insertions(+), 1 deletion(-) + create mode 100644 src/skippy-verify-checkpoint.h + create mode 100644 tests/test-skippy-verify-checkpoint-retirement.cpp + +diff --git a/include/skippy.h b/include/skippy.h +index cc2382d37..f60e7c2ee 100644 +--- a/include/skippy.h ++++ b/include/skippy.h +@@ -525,6 +525,12 @@ LLAMA_API enum skippy_status skippy_trim_session( + uint64_t token_count, + struct skippy_error ** out_error); + ++LLAMA_API enum skippy_status skippy_retire_verify_checkpoint( ++ struct skippy_session * session, ++ uint64_t token_start, ++ uint64_t token_count, ++ struct skippy_error ** out_error); ++ + LLAMA_API enum skippy_status skippy_export_kv_page( + struct skippy_session * session, + int32_t layer_start, +diff --git a/include/skippy/common.h b/include/skippy/common.h +index 11f06bf3e..d1e45b44d 100644 +--- a/include/skippy/common.h ++++ b/include/skippy/common.h +@@ -26,7 +26,7 @@ extern "C" { + + #define SKIPPY_ABI_VERSION_MAJOR 0 + #define SKIPPY_ABI_VERSION_MINOR 1 +-#define SKIPPY_ABI_VERSION_PATCH 33 ++#define SKIPPY_ABI_VERSION_PATCH 34 + + enum skippy_feature { + SKIPPY_FEATURE_RUNTIME_SLICE = 1 << 0, +diff --git a/src/skippy-verify-checkpoint.h b/src/skippy-verify-checkpoint.h +new file mode 100644 +index 000000000..8793000c1 +--- /dev/null ++++ b/src/skippy-verify-checkpoint.h +@@ -0,0 +1,24 @@ ++#pragma once ++ ++#include ++#include ++#include ++ ++template ++bool skippy_retire_verify_checkpoint_exact( ++ std::vector & checkpoints, ++ uint64_t token_start, ++ uint64_t token_count) { ++ const auto it = std::find_if( ++ checkpoints.begin(), checkpoints.end(), ++ [token_start, token_count](const T & checkpoint) { ++ return checkpoint.valid && ++ static_cast(checkpoint.token_start) == token_start && ++ checkpoint.token_count == token_count; ++ }); ++ if (it == checkpoints.end()) { ++ return false; ++ } ++ checkpoints.erase(it); ++ return true; ++} +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 9da45701e..cb20422f9 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -1,4 +1,5 @@ + #include "skippy.h" ++#include "skippy-verify-checkpoint.h" + #include "skippy/devices.h" + #include "skippy-signals.h" + +@@ -7089,6 +7090,26 @@ enum skippy_status skippy_import_kv_page( + return skippy_success(out_error); + } + ++enum skippy_status skippy_retire_verify_checkpoint( ++ skippy_session * session, ++ uint64_t token_start, ++ uint64_t token_count, ++ skippy_error ** out_error) { ++ if (session == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session is null"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (!skippy_memory_needs_verify_checkpoint(session->ctx->get_memory())) { ++ return skippy_success(out_error); ++ } ++ if (!skippy_retire_verify_checkpoint_exact( ++ session->verify_checkpoints, token_start, token_count)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "verify-window checkpoint was not found"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ return skippy_success(out_error); ++} ++ + enum skippy_status skippy_trim_session( + struct skippy_session * session, + uint64_t token_count, +diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt +index 794c0f047..03cad21d5 100644 +--- a/tests/CMakeLists.txt ++++ b/tests/CMakeLists.txt +@@ -152,6 +152,7 @@ llama_build(test-recurrent-state-rollback.cpp get-model.cpp) + + if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) + llama_build_and_test(test-skippy-kv-page-export.cpp) ++ llama_build_and_test(test-skippy-verify-checkpoint-retirement.cpp) + # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries) + llama_build_and_test(test-sampling.cpp) + llama_build_and_test(test-reasoning-budget.cpp) +diff --git a/tests/test-skippy-verify-checkpoint-retirement.cpp b/tests/test-skippy-verify-checkpoint-retirement.cpp +new file mode 100644 +index 000000000..d09acccb4 +--- /dev/null ++++ b/tests/test-skippy-verify-checkpoint-retirement.cpp +@@ -0,0 +1,26 @@ ++#include "../src/skippy-verify-checkpoint.h" ++ ++#include ++#include ++ ++struct checkpoint { ++ bool valid; ++ int token_start; ++ size_t token_count; ++}; ++ ++int main() { ++ std::vector checkpoints; ++ for (int window = 0; window < 129; ++window) { ++ checkpoints.push_back({true, window * 4, 4}); ++ if (!skippy_retire_verify_checkpoint_exact(checkpoints, window * 4, 4)) return 1; ++ if (!checkpoints.empty()) return 2; ++ } ++ ++ checkpoints = {{true, 10, 4}, {true, 14, 4}, {true, 18, 4}}; ++ if (!skippy_retire_verify_checkpoint_exact(checkpoints, 14, 4)) return 3; ++ if (checkpoints.size() != 2 || checkpoints[0].token_start != 10 || checkpoints[1].token_start != 18) return 4; ++ if (skippy_retire_verify_checkpoint_exact(checkpoints, 14, 4)) return 5; ++ if (skippy_retire_verify_checkpoint_exact(checkpoints, 10, 3)) return 6; ++ return 0; ++} +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0061-Reject-trims-after-failed-verify-recovery.patch b/third_party/llama.cpp/patches/0061-Reject-trims-after-failed-verify-recovery.patch new file mode 100644 index 000000000..ec4cfcedd --- /dev/null +++ b/third_party/llama.cpp/patches/0061-Reject-trims-after-failed-verify-recovery.patch @@ -0,0 +1,27 @@ +From 8bde3043235f5b795f19dab1e1ae6d08d422979c Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Thu, 30 Jul 2026 16:40:46 +1000 +Subject: [PATCH 61/61] Reject trims after failed verify recovery + +--- + src/skippy.cpp | 4 ++++ + 1 file changed, 4 insertions(+) + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index cb20422f9..e72a8e82e 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -7118,6 +7118,10 @@ enum skippy_status skippy_trim_session( + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session is required"); + return SKIPPY_STATUS_INVALID_ARGUMENT; + } ++ enum skippy_status usable_status = skippy_require_usable_session(session, out_error); ++ if (usable_status != SKIPPY_STATUS_OK) { ++ return usable_status; ++ } + if (token_count > static_cast(std::numeric_limits::max())) { + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "token_count exceeds int32_t range"); + return SKIPPY_STATUS_INVALID_ARGUMENT; +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0062-Harden-Inkling-MTP-and-KV-contiguity-state.patch b/third_party/llama.cpp/patches/0062-Harden-Inkling-MTP-and-KV-contiguity-state.patch new file mode 100644 index 000000000..a308bd296 --- /dev/null +++ b/third_party/llama.cpp/patches/0062-Harden-Inkling-MTP-and-KV-contiguity-state.patch @@ -0,0 +1,269 @@ +From 5c6836dcf71e9b6a1c29ab412cb3ec1bd3c127bf Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 31 Jul 2026 00:15:15 +1000 +Subject: [PATCH 62/62] Harden Inkling MTP and KV contiguity state + +--- + include/skippy/common.h | 2 +- + src/llama-context.cpp | 4 +-- + src/llama-kv-cache.cpp | 6 ++-- + src/llama-kv-cells.h | 44 +++++++++++++++++++---- + src/skippy.cpp | 20 ++++++++++- + tests/CMakeLists.txt | 1 + + tests/test-skippy-kv-cells-contiguous.cpp | 20 +++++++++++ + 7 files changed, 82 insertions(+), 15 deletions(-) + create mode 100644 tests/test-skippy-kv-cells-contiguous.cpp + +diff --git a/include/skippy/common.h b/include/skippy/common.h +index d1e45b44d..979f1c48e 100644 +--- a/include/skippy/common.h ++++ b/include/skippy/common.h +@@ -26,7 +26,7 @@ extern "C" { + + #define SKIPPY_ABI_VERSION_MAJOR 0 + #define SKIPPY_ABI_VERSION_MINOR 1 +-#define SKIPPY_ABI_VERSION_PATCH 34 ++#define SKIPPY_ABI_VERSION_PATCH 35 + + enum skippy_feature { + SKIPPY_FEATURE_RUNTIME_SLICE = 1 << 0, +diff --git a/src/llama-context.cpp b/src/llama-context.cpp +index 9d528761b..e7c065c91 100644 +--- a/src/llama-context.cpp ++++ b/src/llama-context.cpp +@@ -4113,9 +4113,9 @@ int32_t llama_encode( + int32_t llama_decode( + llama_context * ctx, + llama_batch batch) { +- int ret = ctx->decode(batch); ++ const int ret = ctx->decode(batch); + if (ret == 0) { +- ret = skippy_external_decode_observe(ctx, batch); ++ (void) skippy_external_decode_observe(ctx, batch); + } + if (ret != 0 && ret != 1) { + LLAMA_LOG_ERROR("%s: failed to decode, ret = %d\n", __func__, ret); +diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp +index eebed9e04..b58a2bf65 100644 +--- a/src/llama-kv-cache.cpp ++++ b/src/llama-kv-cache.cpp +@@ -1899,10 +1899,8 @@ uint32_t llama_kv_cache::get_n_kv_pos_contiguous(const slot_info & sinfo, const + } + } + +- for (llama_pos pos = 0; pos <= pos_max; ++pos) { +- if (cells.is_empty(pos) || cells.pos_get(pos) != pos || !cells.seq_has(pos, seq_id)) { +- return 0; +- } ++ if (!cells.seq_pos_is_cell_contiguous(seq_id)) { ++ return 0; + } + + return pos_max + 1; +diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h +index c1c8dad93..aab32c9e3 100644 +--- a/src/llama-kv-cells.h ++++ b/src/llama-kv-cells.h +@@ -45,6 +45,7 @@ public: + + for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) { + seq_pos[s].clear(); ++ seq_pos_index_mismatch[s] = 0; + } + } + +@@ -250,7 +251,7 @@ public: + assert(seq_id >= 0); + + seq[i].reset(seq_id); +- seq_pos_dec(seq_id, pos[i]); ++ seq_pos_dec(seq_id, pos[i], i); + + if (seq[i].none()) { + pos[i] = -1; +@@ -274,7 +275,7 @@ public: + seq[i].reset(); + + seq[i].set(seq_id); +- seq_pos_inc(seq_id, pos[i]); ++ seq_pos_inc(seq_id, pos[i], i); + + return false; + } +@@ -320,7 +321,7 @@ public: + assert(!seq[i].test(seq_id)); + + seq[i].set(seq_id); +- seq_pos_inc(seq_id, pos[i]); ++ seq_pos_inc(seq_id, pos[i], i); + } + + // return the sequence id of this cell +@@ -367,6 +368,23 @@ public: + return seq_pos[seq_id].rbegin()->first; + } + ++ // Whether every position for seq_id is stored in the cell with the same ++ // index and the represented positions form the dense range [0, max]. ++ // The mismatch count is maintained with the existing seq_pos index so this ++ // check remains O(1) on the graph-build path. ++ bool seq_pos_is_cell_contiguous(llama_seq_id seq_id) const { ++ assert(seq_id >= 0); ++ assert(seq_id < LLAMA_MAX_SEQ); ++ ++ const auto & positions = seq_pos[seq_id]; ++ if (positions.empty() || seq_pos_index_mismatch[seq_id] != 0) { ++ return false; ++ } ++ const llama_pos pos_max = positions.rbegin()->first; ++ return positions.begin()->first == 0 && pos_max >= 0 && ++ positions.size() == static_cast(pos_max) + 1; ++ } ++ + // note: call only if the cell is not empty + llama_pos pos_get(uint32_t i) const { + assert(i < pos.size()); +@@ -506,26 +524,38 @@ private: + // + std::map seq_pos[LLAMA_MAX_SEQ]; + ++ // Number of cells for each sequence whose physical cell index differs ++ // from its logical position. ++ uint32_t seq_pos_index_mismatch[LLAMA_MAX_SEQ] = {}; ++ + // helper functions for updating `seq_pos`, once cell at a time: + +- void seq_pos_dec(llama_seq_id s, llama_pos p) { ++ void seq_pos_dec(llama_seq_id s, llama_pos p, uint32_t i) { + auto it = seq_pos[s].find(p); + assert(it != seq_pos[s].end()); + ++ if (p != static_cast(i)) { ++ assert(seq_pos_index_mismatch[s] > 0); ++ --seq_pos_index_mismatch[s]; ++ } ++ + if (--it->second == 0) { + seq_pos[s].erase(it); + } + } + +- void seq_pos_inc(llama_seq_id s, llama_pos p) { ++ void seq_pos_inc(llama_seq_id s, llama_pos p, uint32_t i) { + seq_pos[s][p]++; ++ if (p != static_cast(i)) { ++ ++seq_pos_index_mismatch[s]; ++ } + } + + // remove cell i + void seq_pos_rm(uint32_t i) { + for (int s = 0; s < LLAMA_MAX_SEQ; ++s) { + if (seq[i].test(s)) { +- seq_pos_dec(s, pos[i]); ++ seq_pos_dec(s, pos[i], i); + } + } + } +@@ -534,7 +564,7 @@ private: + void seq_pos_add(uint32_t i) { + for (int s = 0; s < LLAMA_MAX_SEQ; ++s) { + if (seq[i].test(s)) { +- seq_pos_inc(s, pos[i]); ++ seq_pos_inc(s, pos[i], i); + } + } + } +diff --git a/src/skippy.cpp b/src/skippy.cpp +index e72a8e82e..26226b55d 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -3227,6 +3227,9 @@ int skippy_external_decode_observe(llama_context * ctx, const llama_batch & batc + fprintf(stderr, "skippy: external Inkling MTP sync failed: %s\n", + error != nullptr && error->message != nullptr ? error->message : "unknown error"); + skippy_error_free(error); ++ // The target decode has already committed. Degrade only the MTP sidecar ++ // until a later target decode can re-prime it. ++ skippy_mtp_clear_session_state(session); + return -1; + } + +@@ -4456,6 +4459,9 @@ static std::vector skippy_verify_prefix_activation_payload( + const uint64_t flags = checkpoint.input_desc.flags; + if ((flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0) { + const size_t prefix_bytes = skippy_activation_payload_bytes(session, token_count, flags); ++ if (prefix_bytes > checkpoint.input_payload.size()) { ++ return {}; ++ } + return std::vector( + checkpoint.input_payload.begin(), + checkpoint.input_payload.begin() + static_cast(prefix_bytes)); +@@ -4468,6 +4474,18 @@ static std::vector skippy_verify_prefix_activation_payload( + skippy_glm_dsa_top_k_bytes_for_count(token_count, glm_dsa_top_k) : 0; + const size_t prefix_rwkv7_bytes = (flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0 ? + prefix_hidden_bytes : 0; ++ const size_t source_glm_dsa_bytes = (flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0 ? ++ skippy_glm_dsa_top_k_bytes_for_count(checkpoint.token_count, glm_dsa_top_k) : 0; ++ const size_t source_rwkv7_bytes = (flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0 ? ++ source_hidden_bytes : 0; ++ if (prefix_hidden_bytes > source_hidden_bytes || ++ source_hidden_bytes > checkpoint.input_payload.size() || ++ source_glm_dsa_bytes > checkpoint.input_payload.size() - source_hidden_bytes || ++ source_rwkv7_bytes > checkpoint.input_payload.size() - source_hidden_bytes - source_glm_dsa_bytes || ++ prefix_glm_dsa_bytes > source_glm_dsa_bytes || ++ prefix_rwkv7_bytes > source_rwkv7_bytes) { ++ return {}; ++ } + std::vector payload(prefix_hidden_bytes + prefix_glm_dsa_bytes + prefix_rwkv7_bytes); + std::memcpy(payload.data(), checkpoint.input_payload.data(), prefix_hidden_bytes); + +@@ -4478,7 +4496,7 @@ static std::vector skippy_verify_prefix_activation_payload( + payload.data() + prefix_offset, + checkpoint.input_payload.data() + source_offset, + prefix_glm_dsa_bytes); +- source_offset += skippy_glm_dsa_top_k_bytes_for_count(checkpoint.token_count, glm_dsa_top_k); ++ source_offset += source_glm_dsa_bytes; + prefix_offset += prefix_glm_dsa_bytes; + } + if (prefix_rwkv7_bytes > 0) { +diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt +index 03cad21d5..4a1db624c 100644 +--- a/tests/CMakeLists.txt ++++ b/tests/CMakeLists.txt +@@ -153,6 +153,7 @@ llama_build(test-recurrent-state-rollback.cpp get-model.cpp) + if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) + llama_build_and_test(test-skippy-kv-page-export.cpp) + llama_build_and_test(test-skippy-verify-checkpoint-retirement.cpp) ++ llama_build_and_test(test-skippy-kv-cells-contiguous.cpp) + # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries) + llama_build_and_test(test-sampling.cpp) + llama_build_and_test(test-reasoning-budget.cpp) +diff --git a/tests/test-skippy-kv-cells-contiguous.cpp b/tests/test-skippy-kv-cells-contiguous.cpp +new file mode 100644 +index 000000000..0e95c35b2 +--- /dev/null ++++ b/tests/test-skippy-kv-cells-contiguous.cpp +@@ -0,0 +1,20 @@ ++#include "../src/llama-kv-cells.h" ++ ++int main() { ++ llama_kv_cells cells; ++ cells.resize(8); ++ for (uint32_t i = 0; i < 4; ++i) { ++ cells.pos_set(i, static_cast(i)); ++ cells.seq_add(i, 0); ++ } ++ if (!cells.seq_pos_is_cell_contiguous(0)) return 1; ++ ++ cells.mv(1, 4); ++ if (cells.seq_pos_is_cell_contiguous(0)) return 2; ++ cells.mv(4, 1); ++ if (!cells.seq_pos_is_cell_contiguous(0)) return 3; ++ ++ if (!cells.seq_rm(2, 0)) return 4; ++ if (cells.seq_pos_is_cell_contiguous(0)) return 5; ++ return 0; ++} +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0063-Harden-Inkling-MTP-sidecar-and-verify-cleanup.patch b/third_party/llama.cpp/patches/0063-Harden-Inkling-MTP-sidecar-and-verify-cleanup.patch new file mode 100644 index 000000000..76af7c730 --- /dev/null +++ b/third_party/llama.cpp/patches/0063-Harden-Inkling-MTP-sidecar-and-verify-cleanup.patch @@ -0,0 +1,140 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: mesh-llm maintainers +Date: Thu, 30 Jul 2026 15:10:00 +0000 +Subject: [PATCH] Harden Inkling MTP sidecar and verify cleanup + +Keep MTP embedding overrides aligned when llama.cpp splits a decode batch, +invalidate partial sidecar state after decode failure, and reject empty +activation prefixes during verify recovery. Also clear the external-decode +session pointer during teardown and advance the RPC patch version for the +added GGML operation. +--- + ggml/include/ggml-rpc.h | 2 +- + src/llama-graph.cpp | 14 ++++++++++---- + src/llama-graph.h | 1 + + src/skippy.cpp | 24 +++++++++++++++++++++--- + 4 files changed, 33 insertions(+), 8 deletions(-) + +diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h +index 086dddd96..2d103664f 100644 +--- a/ggml/include/ggml-rpc.h ++++ b/ggml/include/ggml-rpc.h +@@ -8,7 +8,7 @@ extern "C" { + + #define RPC_PROTO_MAJOR_VERSION 4 + #define RPC_PROTO_MINOR_VERSION 0 +-#define RPC_PROTO_PATCH_VERSION 4 ++#define RPC_PROTO_PATCH_VERSION 5 + + #ifdef __cplusplus + static_assert(GGML_OP_COUNT == 108, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp +index 27c5129be..2d1cff2a3 100644 +--- a/src/llama-graph.cpp ++++ b/src/llama-graph.cpp +@@ -178,9 +178,14 @@ void llm_graph_input_embd_h::set_input(const llama_ubatch * ubatch) { + GGML_ASSERT(n_embd == embd->ne[0]); + const float * values = ubatch->embd; + if (mtp_embeddings.values != nullptr) { +- GGML_ASSERT(mtp_embeddings.token_count == ubatch->n_tokens); + GGML_ASSERT(mtp_embeddings.n_embd == n_embd); +- values = mtp_embeddings.values; ++ GGML_ASSERT(ubatch->pos != nullptr); ++ const llama_pos token_offset = ubatch->pos[0] - mtp_embeddings.pos_start; ++ GGML_ASSERT(token_offset >= 0); ++ GGML_ASSERT( ++ static_cast(token_offset) + ubatch->n_tokens <= ++ mtp_embeddings.token_count); ++ values = mtp_embeddings.values + static_cast(token_offset)*n_embd; + } + + ggml_backend_tensor_set(embd, values, 0, n_tokens*n_embd*ggml_element_size(embd)); +@@ -2656,7 +2661,8 @@ ggml_tensor * llm_graph_context::build_attn_mha( + ggml_tensor * fa_mask = kq_mask; + if (fold_kq_b_into_mask) { + ggml_tensor * bias = ggml_cont(ctx0, kq_b); +- if (kq_mask->ne[1] != bias->ne[1]) { ++ GGML_ASSERT(kq_mask->ne[1] >= bias->ne[1]); ++ if (kq_mask->ne[1] > bias->ne[1]) { + // the FA mask rows are padded; pad the bias with zeros to match + bias = ggml_pad(ctx0, bias, 0, (int)(kq_mask->ne[1] - bias->ne[1]), 0, 0); + } +diff --git a/src/llama-graph.h b/src/llama-graph.h +index 6e66f7fef..a02f65b26 100644 +--- a/src/llama-graph.h ++++ b/src/llama-graph.h +@@ -82,6 +82,7 @@ struct skippy_activation_mtp_embeddings { + const float * values = nullptr; + uint32_t token_count = 0; + uint32_t n_embd = 0; ++ llama_pos pos_start = 0; + }; + + void skippy_graph_set_filter(const skippy_graph_filter & filter); +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 26226b55d..2e5e6e8f3 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -2986,9 +2986,13 @@ static bool skippy_mtp_shares_target_memory( + } + + struct skippy_mtp_embeddings_scope { +- skippy_mtp_embeddings_scope(const float * values, uint32_t token_count, uint32_t n_embd) { ++ skippy_mtp_embeddings_scope( ++ const float * values, ++ uint32_t token_count, ++ uint32_t n_embd, ++ llama_pos pos_start) { + if (values != nullptr) { +- skippy_graph_set_mtp_embeddings({ values, token_count, n_embd }); ++ skippy_graph_set_mtp_embeddings({ values, token_count, n_embd, pos_start }); + enabled = true; + } + } +@@ -3128,7 +3132,8 @@ static enum skippy_status skippy_mtp_sync_target_inputs( + skippy_mtp_embeddings_scope mtp_embeddings_scope( + mtp_embeddings, + static_cast(n_decode), +- static_cast(n_embd)); ++ static_cast(n_embd), ++ token_start + static_cast(first_index)); + for (uint32_t depth = 0; depth < sync_depth_count; ++depth) { + if (chain_heads) { + if (llama_memory_t memory = mtp_ctx->get_memory()) { +@@ -3141,6 +3146,9 @@ static enum skippy_status skippy_mtp_sync_target_inputs( + skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config); + rc = llama_decode(mtp_ctx, batch); + } ++ if (rc != 0) { ++ session->mtp_prefix_valid = false; ++ } + if (rc != 0 || !chain_heads || depth + 1 == sync_depth_count) { + break; + } +@@ -4600,6 +4608,13 @@ static enum skippy_status skippy_restore_verify_prefix( + session, + *checkpoint, + accepted_count); ++ if (input_payload.empty()) { ++ return skippy_fail_verify_restore( ++ session, ++ SKIPPY_STATUS_RUNTIME_ERROR, ++ "verify-window checkpoint has an invalid activation prefix payload", ++ out_error); ++ } + input_desc.token_count = static_cast(accepted_count); + input_desc.payload_bytes = input_payload.size(); + status = skippy_decode_activation_frame( +@@ -5513,6 +5528,9 @@ enum skippy_status skippy_session_free( + struct skippy_session * session, + struct skippy_error ** out_error) { + if (session != nullptr) { ++ if (g_skippy_external_decode_session == session) { ++ g_skippy_external_decode_session = nullptr; ++ } + if (session->ctx != nullptr) { + if (session->preserve_prefix_on_free && !session->borrowed_sequence && session->stage_model != nullptr) { + session->ctx->synchronize(); +-- +2.50.1 + diff --git a/third_party/llama.cpp/patches/0064-Accept-Inkling-arguments-tool-call-key.patch b/third_party/llama.cpp/patches/0064-Accept-Inkling-arguments-tool-call-key.patch new file mode 100644 index 000000000..0831afcdb --- /dev/null +++ b/third_party/llama.cpp/patches/0064-Accept-Inkling-arguments-tool-call-key.patch @@ -0,0 +1,94 @@ +From 1c496ba2c99581768a457d1e87ceeb5a41db7bc0 Mon Sep 17 00:00:00 2001 +From: Mesh LLM +Date: Sat, 1 Aug 2026 08:44:30 +1000 +Subject: [PATCH] Accept Inkling arguments tool-call key + +--- + common/chat-peg-parser.cpp | 4 +++- + common/chat.cpp | 5 ++++- + tests/test-chat.cpp | 27 +++++++++++++++++++++++++++ + 3 files changed, 34 insertions(+), 2 deletions(-) + +diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp +index db9004d84..d2e97bea7 100644 +--- a/common/chat-peg-parser.cpp ++++ b/common/chat-peg-parser.cpp +@@ -751,7 +751,9 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys( + + auto tool_choices = choice(); + auto name_key_parser = literal("\"" + effective_name_key + "\""); +- auto args_key_parser = literal("\"" + effective_args_key + "\""); ++ auto args_key_parser = effective_args_key == "arguments" ++ ? choice({ literal("\"arguments\""), literal("\"args\"") }) ++ : literal("\"" + effective_args_key + "\""); + + for (const auto & tool_def : tools) { + if (!tool_def.contains("function")) { +diff --git a/common/chat.cpp b/common/chat.cpp +index 89f518716..3b71427d9 100644 +--- a/common/chat.cpp ++++ b/common/chat.cpp +@@ -2465,11 +2465,14 @@ static common_chat_params common_chat_params_init_inkling(const common_chat_temp + // each call is its own block (role opener + bare name echo + JSON section); + // force_tool_calls=true makes the JSON section required so a pure-text answer fails the + // block cleanly; parallel calls are separate blocks, hence repeat + parallel=false ++ // The canonical TML history form uses "args", but Inkling can emit the ++ // OpenAI spelling "arguments". The standard parser's "arguments" alias ++ // accepts both generated forms. + auto tool_section = p.standard_json_tools( + INVOKE_TOOL, END_MESSAGE, inputs.tools, /* parallel_tool_calls = */ false, + /* force_tool_calls = */ true, + /* name_key = */ "name", +- /* args_key = */ "args", ++ /* args_key = */ "arguments", + /* array_wrapped = */ false); + // the name-echo scan must stop at any block marker: a greedy until(INVOKE_TOOL) returns + // NEED_MORE_INPUT mid-stream, which choice() treats as a match and shadows the text branch +diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp +index fcf025d07..f57f6292a 100644 +--- a/tests/test-chat.cpp ++++ b/tests/test-chat.cpp +@@ -436,6 +436,21 @@ static common_chat_tool special_function_tool{ + "required": ["arg1"] + })", + }; ++static common_chat_tool lookup_fixture_fact_tool{ ++ /* .name = */ "lookup_fixture_fact", ++ /* .description = */ "Return one deterministic fact from the agent reliability fixture.", ++ /* .parameters = */ R"({ ++ "type": "object", ++ "properties": { ++ "key": { ++ "type": "string", ++ "enum": ["codeword"] ++ } ++ }, ++ "required": ["key"], ++ "additionalProperties": false ++ })", ++}; + static common_chat_tool special_function_tool_with_optional_param{ + /* .name = */ "special_function_with_opt", + /* .description = */ "I'm special but have optional stuff", +@@ -2977,6 +2992,18 @@ static void test_template_output_peg_parsers(bool detailed_debug) { + .expect(message_assist_call) + .run(); + ++ // Inkling also emits the OpenAI key spelling despite the canonical ++ // history template using "args". ++ tst.test("lookup_fixture_fact<|content_invoke_tool_json|>" ++ "{\"name\":\"lookup_fixture_fact\",\"arguments\":{\"key\":\"codeword\"}}<|end_message|>") ++ .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) ++ .tools({ lookup_fixture_fact_tool }) ++ .tool_choice(COMMON_CHAT_TOOL_CHOICE_REQUIRED) ++ .enable_thinking(false) ++ .parallel_tool_calls(false) ++ .expect(simple_assist_msg("", "", "lookup_fixture_fact", "{\"key\":\"codeword\"}")) ++ .run(); ++ + // regression: the tool branch must not swallow a pure-text answer + tst.test("<|content_thinking|>I'm\nthinking<|end_message|>" + "<|message_model|><|content_text|>Hello, world!\nWhat's up?<|end_message|>" +-- +2.50.1 (Apple Git-155) + diff --git a/tools/relay-fly-legacy/README.md b/tools/relay-fly-legacy/README.md index 8c01621c4..412a1fdfe 100644 --- a/tools/relay-fly-legacy/README.md +++ b/tools/relay-fly-legacy/README.md @@ -1,14 +1,19 @@ # Fly relay reference mesh-llm production relay operations use managed iroh relay infrastructure via -[services.iroh.computer](https://services.iroh.computer): +[services.iroh.computer](https://services.iroh.computer). The service operates +and maintains the relay fleet, providing managed regional coverage instead of +requiring mesh-llm to deploy and operate its own relay servers: | Relay | Region | URL | |-------|--------|-----| | USW1-2 | US West | `https://usw1-2.relay.michaelneale.mesh-llm.iroh.link./` | | APS1-1 | Asia-Pacific South | `https://aps1-1.relay.michaelneale.mesh-llm.iroh.link./` | +| EUC1-1 | Europe Central | `https://euc1-1.relay.michaelneale.mesh-llm.iroh.link./` | +| USE1-1 | US East | `https://use1-1.relay.michaelneale.mesh-llm.iroh.link./` | -These are configured as defaults in `crates/mesh-llm/src/mesh/mod.rs`. +These are configured as defaults in +`crates/mesh-llm-host-runtime/src/mesh/connections.rs`. `mesh-llm-relay.fly.dev` is retained here as a Fly.io deployment reference. diff --git a/website/src/docs/pages/CLI.md b/website/src/docs/pages/CLI.md index 50eafac66..842c1d35b 100644 --- a/website/src/docs/pages/CLI.md +++ b/website/src/docs/pages/CLI.md @@ -443,6 +443,10 @@ mesh-llm models package unsloth/Qwen3-8B-GGUF:Q4_K_M --confirm --follow mesh-llm models package --status ``` +Pass `--experimental` to publish a public package marked experimental: the +package README carries an experimental warning and the Hugging Face catalog PR +is opened but left unmerged until the package is certified. + Use `--help` for the full planning, status, logs, cancel, and publishing options.