From 91cd35e241c9cfd8f1c5861a3b7f21644a1edbbf Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:35:04 +1000 Subject: [PATCH 01/26] feat(skippy): add experimental Inkling text split serving --- .../queue-unsloth-layer-packages.yml | 2 +- Cargo.lock | 1 + Justfile | 8 +- crates/llama-quant-ffi/src/lib.rs | 1 + crates/mesh-llm-cli/src/models.rs | 3 + crates/mesh-llm-cli/src/parser.rs | 28 +- crates/mesh-llm-commands/src/model_package.rs | 145 +- .../src/inference/skippy/family_policy.rs | 35 +- .../src/inference/skippy/mod.rs | 4 +- .../inference/skippy/resolver/resolution.rs | 36 +- .../src/inference/skippy/resolver/tests.rs | 106 + .../src/mesh/capacity.rs | 145 + .../src/mesh/connections.rs | 37 +- .../src/mesh/direct_path.rs | 22 +- crates/mesh-llm-host-runtime/src/mesh/mod.rs | 4 +- crates/mesh-llm-host-runtime/src/mesh/node.rs | 14 +- .../src/mesh/peer_state.rs | 6 + .../src/mesh/stage_artifacts.rs | 10 +- .../src/mesh/stage_transport.rs | 21 +- .../src/mesh/tests/connections.rs | 1 + .../src/mesh/tests/direct_path.rs | 50 + .../src/models/capabilities.rs | 60 +- .../mesh-llm-host-runtime/src/models/mod.rs | 2 +- .../src/runtime/local.rs | 24 +- .../src/runtime/local_package.rs | 183 +- .../src/runtime/local_split.rs | 115 +- .../src/runtime/local_split/loading.rs | 39 +- .../src/runtime/local_split/tests.rs | 247 +- .../mesh-llm-host-runtime/src/runtime/mod.rs | 2 + .../src/runtime/model_lifecycle/load.rs | 2 +- .../runtime/model_lifecycle/reconciliation.rs | 2 +- .../src/runtime/split_participant_settle.rs | 378 ++ .../src/runtime/split_planning.rs | 96 +- .../src/runtime/startup_handles.rs | 9 +- .../src/runtime/startup_retry.rs | 96 + crates/mesh-llm-host-runtime/src/sdk.rs | 4 +- crates/mesh-llm-types/src/mesh/mod.rs | 27 + crates/mesh-llm/src/commands/mod.rs | 1 + crates/mesh-llm/src/commands/models/mod.rs | 2 + crates/model-artifact/src/gguf.rs | 143 +- crates/model-artifact/src/gguf/kv_cache.rs | 17 +- .../src/bin/queue-unsloth-layer-packages.rs | 153 +- crates/model-package/src/jobs.rs | 2 + crates/model-package/src/prepare.rs | 151 +- crates/model-package/src/script.rs | 12 +- .../src/scripts/split-model-job.sh | 52 +- crates/skippy-coordinator/src/topology.rs | 30 + crates/skippy-correctness/src/cli.rs | 23 + crates/skippy-correctness/src/main.rs | 3 +- crates/skippy-correctness/src/runner/mod.rs | 2 + .../src/runner/stage_fa_parity.rs | 134 + crates/skippy-ffi/src/lib.rs | 74 +- .../skippy-protocol/src/binary/activation.rs | 5 +- crates/skippy-protocol/src/binary/mod.rs | 46 +- crates/skippy-protocol/src/binary/types.rs | 15 +- crates/skippy-quantize/Cargo.toml | 6 +- crates/skippy-quantize/src/backend.rs | 2 +- crates/skippy-quantize/src/gguf_metadata.rs | 207 + crates/skippy-quantize/src/gguf_template.rs | 76 + crates/skippy-quantize/src/gguf_writer.rs | 337 +- .../src/gguf_writer/glm_dsa.rs | 11 +- .../skippy-quantize/src/gguf_writer_tests.rs | 91 + crates/skippy-quantize/src/hf_checkpoint.rs | 35 + .../skippy-quantize/src/inkling_metadata.rs | 243 + crates/skippy-quantize/src/main.rs | 10 + crates/skippy-quantize/src/mtp_attach.rs | 203 + .../skippy-quantize/src/projector_validate.rs | 68 + crates/skippy-quantize/src/tensor_map.rs | 130 +- .../skippy-quantize/src/tokenizer_metadata.rs | 86 +- crates/skippy-quantize/src/types.rs | 10 +- crates/skippy-runtime/src/media.rs | 15 + crates/skippy-runtime/src/runtime_events.rs | 2 +- crates/skippy-runtime/src/types.rs | 1 + .../src/binary_transport/binary_messaging.rs | 8 +- .../binary_messaging/connection.rs | 173 +- .../binary_messaging/reply.rs | 54 + .../src/binary_transport/options.rs | 43 + .../src/binary_transport/preconnect.rs | 4 +- .../src/binary_transport/stage_execution.rs | 151 +- crates/skippy-server/src/cli.rs | 5 + .../src/frontend/decode_scheduler.rs | 20 + .../src/frontend/embedded_execution.rs | 9 +- .../src/frontend/embedded_generation.rs | 36 +- .../frontend/embedded_generation/lifecycle.rs | 72 +- .../skippy-server/src/frontend/generation.rs | 2 + .../frontend/generation/persistent_lanes.rs | 8 +- .../src/frontend/generation/timeouts.rs | 45 + .../src/frontend/generation_flow.rs | 1 + crates/skippy-server/src/frontend/request.rs | 25 +- .../src/frontend/tests/prefill.rs | 7 +- .../src/frontend/tests/request.rs | 64 +- .../src/frontend/tests/wire_messages.rs | 2 + .../src/frontend/wire_messages.rs | 3 +- crates/skippy-server/src/runtime_state.rs | 33 +- .../reviewed-family-capabilities.json | 23 + .../skippy-topology/src/family_capability.rs | 25 + crates/skippy-topology/src/tests.rs | 14 + docs/LAYER_PACKAGE_REPOS.md | 19 +- docs/skippy/FAMILY_CERTIFY.md | 16 +- docs/skippy/FAMILY_STATUS.md | 11 + docs/skippy/NEW_MODEL_ONBOARDING.md | 8 +- docs/skippy/llama-parity-candidates.json | 12 + scripts/hf-skippy-convert-job.py | 211 + scripts/hf-skippy-mtp-certify-job.py | 168 + .../0048-Add-TML-Inkling-architecture.patch | 5829 +++++++++++++++++ ...staged-execution-support-for-Inkling.patch | 108 + ...ance-native-MTP-depth-per-draft-step.patch | 81 + ...lay-multi-depth-MTP-caches-correctly.patch | 368 ++ ...Add-Inkling-multi-depth-MTP-sidecars.patch | 545 ++ ...ltimodal-Inkling-embeddings-into-MTP.patch | 567 ++ ...rrent-memory-for-appended-MTP-layers.patch | 46 + ...estore-recurrent-verify-window-state.patch | 317 + ...reserve-verify-checkpoint-until-trim.patch | 84 + ...-retain-pipelined-verify-checkpoints.patch | 170 + ...-Inkling-stage-memory-by-layer-range.patch | 72 + ...Inkling-and-GLM-metadata-integration.patch | 143 + ...lance-filtered-stages-across-devices.patch | 76 + tools/relay-fly-legacy/README.md | 9 +- website/src/docs/pages/CLI.md | 4 + 119 files changed, 13447 insertions(+), 607 deletions(-) create mode 100644 crates/mesh-llm-host-runtime/src/mesh/capacity.rs create mode 100644 crates/mesh-llm-host-runtime/src/runtime/split_participant_settle.rs create mode 100644 crates/mesh-llm-host-runtime/src/runtime/startup_retry.rs create mode 100644 crates/skippy-correctness/src/runner/stage_fa_parity.rs create mode 100644 crates/skippy-quantize/src/gguf_metadata.rs create mode 100644 crates/skippy-quantize/src/inkling_metadata.rs create mode 100644 crates/skippy-quantize/src/mtp_attach.rs create mode 100644 crates/skippy-quantize/src/projector_validate.rs create mode 100644 crates/skippy-server/src/frontend/generation/timeouts.rs create mode 100644 scripts/hf-skippy-convert-job.py create mode 100644 scripts/hf-skippy-mtp-certify-job.py create mode 100644 third_party/llama.cpp/patches/0048-Add-TML-Inkling-architecture.patch create mode 100644 third_party/llama.cpp/patches/0049-Add-staged-execution-support-for-Inkling.patch create mode 100644 third_party/llama.cpp/patches/0050-Advance-native-MTP-depth-per-draft-step.patch create mode 100644 third_party/llama.cpp/patches/0051-Replay-multi-depth-MTP-caches-correctly.patch create mode 100644 third_party/llama.cpp/patches/0052-Add-Inkling-multi-depth-MTP-sidecars.patch create mode 100644 third_party/llama.cpp/patches/0053-Carry-multimodal-Inkling-embeddings-into-MTP.patch create mode 100644 third_party/llama.cpp/patches/0054-Size-recurrent-memory-for-appended-MTP-layers.patch create mode 100644 third_party/llama.cpp/patches/0055-Restore-recurrent-verify-window-state.patch create mode 100644 third_party/llama.cpp/patches/0056-skippy-preserve-verify-checkpoint-until-trim.patch create mode 100644 third_party/llama.cpp/patches/0057-skippy-retain-pipelined-verify-checkpoints.patch create mode 100644 third_party/llama.cpp/patches/0058-skippy-filter-Inkling-stage-memory-by-layer-range.patch create mode 100644 third_party/llama.cpp/patches/0059-Complete-Inkling-and-GLM-metadata-integration.patch create mode 100644 third_party/llama.cpp/patches/0060-skippy-balance-filtered-stages-across-devices.patch diff --git a/.github/workflows/queue-unsloth-layer-packages.yml b/.github/workflows/queue-unsloth-layer-packages.yml index aadfac8817..c33cf38034 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: diff --git a/Cargo.lock b/Cargo.lock index c6b6e722b8..88903d2472 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7438,6 +7438,7 @@ dependencies = [ "serde", "serde_json", "skippy-ffi", + "skippy-runtime", ] [[package]] diff --git a/Justfile b/Justfile index 48af49ecfb..4d3c92c5e0 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 40666e378b..8f29e81d4f 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 135a9a220c..b13d432e21 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 7f058a82c6..166ab20e69 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 15c52163cc..98579c69b3 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(()); 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 2ea9f02c41..bb5f22be31 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, } @@ -58,7 +59,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).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 +143,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 +179,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 +263,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 +301,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 +330,7 @@ fn disabled_family_policy( ) -> FamilyPolicy { FamilyPolicy { activation_wire_dtype, + default_kv_cache_type: None, prefix_cache: FamilyPrefixCachePolicy::Disabled { reason }, } } @@ -621,7 +648,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, 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 56069383e4..718ff8061b 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/resolution.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs index ce071f959d..6993ce43ae 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,19 @@ 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 { + 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 +267,19 @@ 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 { + 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/tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs index 17240816bc..9096b4f833 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,75 @@ 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 family_policy_wires_prefix_cache_by_default_for_supported_models() { let model_file = temp_model_file(); @@ -870,6 +939,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/mesh/capacity.rs b/crates/mesh-llm-host-runtime/src/mesh/capacity.rs new file mode 100644 index 0000000000..1ace29d8a3 --- /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 1616662cc1..b4c5018e02 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 44b86b84bf..f6d8047a40 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 9a4d0681f9..efdf4c1cd2 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/mod.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/mod.rs @@ -8,7 +8,8 @@ pub use mesh_llm_types::mesh::{ MAX_SPLIT_RTT_MS, 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 +77,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 1de6abcb4b..d098639b57 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/node.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/node.rs @@ -129,7 +129,9 @@ pub(crate) fn hardware_snapshot_for_start( role: &NodeRole, max_vram_gb: Option, ) -> NodeHardwareSnapshot { - let mut vram_bytes = hw.vram_bytes; + let local_runtime_capacity_bytes = + super::capacity::capped_capacity_bytes(hw.vram_bytes, max_vram_gb); + let mut vram_bytes = super::capacity::advertised_capacity_bytes(&hw, max_vram_gb); let gpu_name = if matches!(role, NodeRole::Client) { None } else { @@ -160,6 +162,7 @@ pub(crate) fn hardware_snapshot_for_start( NodeHardwareSnapshot { vram_bytes, + local_runtime_capacity_bytes, gpu_name, hostname, is_soc, @@ -280,7 +283,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 +1017,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 +1192,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)), @@ -1725,8 +1733,8 @@ impl Node { // (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/peer_state.rs b/crates/mesh-llm-host-runtime/src/mesh/peer_state.rs index b07d9018ad..8b03009574 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 c9a1bcea6d..b740d8c1b1 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 5e016c6654..eceb3bee52 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 1618b76222..b202dc87df 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 4c01b29127..0d01f54af9 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 e76344e190..2ecf49e3e5 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 d848302dfa..653148318c 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/runtime/local.rs b/crates/mesh-llm-host-runtime/src/runtime/local.rs index daa0c12a5e..6498ca94c0 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local.rs @@ -517,7 +517,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. @@ -618,9 +618,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 @@ -704,9 +709,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 52e16393b1..8286d97a31 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 b45e7a0869..36d3ad242b 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,11 +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, @@ -279,6 +300,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 +311,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 +319,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 +338,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 +355,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 +375,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 2e0a9a95f5..e1b50cf7c6 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,11 @@ 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); +pub(super) const MAX_STAGE_SOURCE_PREPARE_TIMEOUT: Duration = Duration::from_secs(6 * 60 * 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 +181,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 +221,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 +280,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 +332,25 @@ 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) + .clamp( + MIN_STAGE_SOURCE_PREPARE_TIMEOUT, + MAX_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 e132980c13..e1ac153959 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,66 @@ 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 mut package_planning_called = false; + + 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(_) => package_planning_called = true, + } + + assert!(!package_planning_called); +} + #[test] fn split_topology_planner_uses_all_eligible_participants() { let participants = vec![ @@ -67,6 +127,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 +613,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: 321_400_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 <= MAX_STAGE_SOURCE_PREPARE_TIMEOUT); +} + #[test] fn startup_runtime_plan_auto_splits_when_model_exceeds_local_capacity() { assert_eq!( @@ -745,7 +888,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 +1567,105 @@ 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 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, size_tiered, + "family-aware planning and the large-model policy must agree on Q4_0" + ); +} diff --git a/crates/mesh-llm-host-runtime/src/runtime/mod.rs b/crates/mesh-llm-host-runtime/src/runtime/mod.rs index 4b1482cadb..1ad700e6ed 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 3233c69244..8ea8fb828d 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 54e0cbcb40..50adcd24dd 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 0000000000..ee5ace7841 --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/runtime/split_participant_settle.rs @@ -0,0 +1,378 @@ +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); + +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() + 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 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 ade2449ed3..7d2af937ed 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 40a3659de4..e3e540e3bd 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 0000000000..18ea0500cb --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/runtime/startup_retry.rs @@ -0,0 +1,96 @@ +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("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 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 615a49a1f8..fc76ea12ad 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 b5c9b5aaf5..4dda2bbb0d 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 cce4106945..e6d4f47083 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 24e59a2932..1ff45b7d03 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 802caf4a5b..80e89aecca 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 17935797a6..4c4c7e4ec1 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 49978fb754..b7af2dd049 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 76eafd2b55..301c3069bb 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 4666b475c9..32a5e44dd2 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 edd1946feb..5c41965c31 100644 --- a/crates/model-package/src/script.rs +++ b/crates/model-package/src/script.rs @@ -129,7 +129,11 @@ 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("- openai-compatible")); assert!(EMBEDDED_SCRIPT.contains("## Model Overview")); assert!(EMBEDDED_SCRIPT.contains("## Highlights")); @@ -198,6 +202,12 @@ 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""#)); } diff --git a/crates/model-package/src/scripts/split-model-job.sh b/crates/model-package/src/scripts/split-model-job.sh index 03541fc871..576c1d2ab3 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 @@ -522,6 +554,20 @@ 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 "" +) def sha256(path: Path) -> str: digest = hashlib.sha256() @@ -652,7 +698,7 @@ readme = f"""--- library_name: mesh-llm base_model: - {yaml_quote(source_repo)} -pipeline_tag: text-generation +pipeline_tag: {yaml_quote(source_pipeline_tag)} tags: - gguf - mesh-llm @@ -661,7 +707,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. diff --git a/crates/skippy-coordinator/src/topology.rs b/crates/skippy-coordinator/src/topology.rs index 8f311861d9..cd49cffc7a 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 b18a57474d..21ee19ac1b 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,25 @@ pub enum StatePayloadKind { RecurrentOnly, KvRecurrent, } + +#[derive(Args)] +pub struct StageFaParityArgs { + #[arg(long)] + pub model: PathBuf, + #[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 d1fd813ae4..269229389d 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 9ec2ef3a79..d9a1018015 100644 --- a/crates/skippy-correctness/src/runner/mod.rs +++ b/crates/skippy-correctness/src/runner/mod.rs @@ -2,8 +2,10 @@ pub(crate) mod native_mtp; 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/stage_fa_parity.rs b/crates/skippy-correctness/src/runner/stage_fa_parity.rs new file mode 100644 index 0000000000..736f082864 --- /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: "unsloth/inkling-GGUF:UD-Q2_K_XL".to_string(), + 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 99ddd56bb4..e3802ec95d 100644 --- a/crates/skippy-ffi/src/lib.rs +++ b/crates/skippy-ffi/src/lib.rs @@ -5,6 +5,7 @@ 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)] @@ -887,6 +896,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; @@ -1177,6 +1187,20 @@ pub fn skippy_abi_features() -> u64 { unsafe { fns() } } +/// 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 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"))] unsafe extern "C" { pub fn llama_log_set(log_callback: LlamaLogCallback, user_data: *mut c_void); @@ -1795,6 +1819,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 +1870,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 +1881,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 +1916,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-protocol/src/binary/activation.rs b/crates/skippy-protocol/src/binary/activation.rs index 644bbb1176..2e056a9312 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 cb8004b66f..58b1cce798 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, @@ -905,6 +905,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 c89eb8ca51..a95b4ce6eb 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, which changes activation payload length. +// Stage peers must be upgraded together so an older reader rejects the header before it can +// leave the additional sideband bytes in the stream. +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; @@ -183,10 +184,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 +199,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 +213,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 } diff --git a/crates/skippy-quantize/Cargo.toml b/crates/skippy-quantize/Cargo.toml index c75e8b4695..f24f515078 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 6907a8eb6a..4f2e8f6d98 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 }; diff --git a/crates/skippy-quantize/src/gguf_metadata.rs b/crates/skippy-quantize/src/gguf_metadata.rs new file mode 100644 index 0000000000..e960acbee9 --- /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 b871a1f4b4..04fdae7095 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 c845eff4eb..1c559077bc 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; @@ -305,6 +305,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 +403,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 +420,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 +428,88 @@ 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 TensorNameMap::HfToGgufWithMtp { layer_start } = tensor_name_map else { + anyhow::bail!("Inkling MTP conversion requires an MTP-aware tensor name map"); + }; + let depth = inkling_mtp_depth(tensor.name())? + .with_context(|| format!("missing Inkling MTP depth in {}", tensor.name()))?; + let layer = layer_start + .checked_add(depth) + .context("Inkling MTP layer id overflow")?; + 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 +746,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 +770,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 +811,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 +865,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 fcb31c2bd7..a5f63c0528 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 ba09b3e6c3..92b23a6bfd 100644 --- a/crates/skippy-quantize/src/gguf_writer_tests.rs +++ b/crates/skippy-quantize/src/gguf_writer_tests.rs @@ -260,6 +260,90 @@ 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 validates_qwen_dense_native_conversion_fixture() { let root = unique_temp_dir(); @@ -1267,6 +1351,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 e898523f90..96361e2c3c 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 0000000000..64683de1f4 --- /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 347d9f1399..7856e134e2 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, diff --git a/crates/skippy-quantize/src/mtp_attach.rs b/crates/skippy-quantize/src/mtp_attach.rs new file mode 100644 index 0000000000..bb28ac6c97 --- /dev/null +++ b/crates/skippy-quantize/src/mtp_attach.rs @@ -0,0 +1,203 @@ +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 native_mtp_multimodal_feature = + skippy_ffi::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/projector_validate.rs b/crates/skippy-quantize/src/projector_validate.rs new file mode 100644 index 0000000000..9e48045851 --- /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 3fa4b3920b..b0c028ccdc 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.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 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_start + depth, + 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,46 @@ 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" + )); + } } diff --git a/crates/skippy-quantize/src/tokenizer_metadata.rs b/crates/skippy-quantize/src/tokenizer_metadata.rs index a6377f9cc0..06b4566b82 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 = 2; 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 41671fa9e2..3dce8433bb 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/media.rs b/crates/skippy-runtime/src/media.rs index a9a3470a30..93b2119639 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/runtime_events.rs b/crates/skippy-runtime/src/runtime_events.rs index d363781259..424cce4712 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/types.rs b/crates/skippy-runtime/src/types.rs index 0c0057d977..a49ae54417 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/src/binary_transport/binary_messaging.rs b/crates/skippy-server/src/binary_transport/binary_messaging.rs index a9f1dcdf47..96ded79f96 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, @@ -162,7 +162,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 +243,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/connection.rs b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs index f0fc56b4c2..c0ee57922d 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -1,7 +1,10 @@ use super::async_forwarder::AsyncForwarder; use super::reply::drain_deferred_prefill_replies; use super::reply::send_stage_reply; -use super::reply::{configure_prediction_return_stream, reply_window_for_message}; +use super::reply::{ + configure_prediction_return_stream, normalize_downstream_prefix_restore_reply, + reply_window_for_message, +}; use super::summary::BinaryMessageObservation; use super::summary::BinaryRequestSummary; use super::telemetry::UpstreamReplyWriteSpan; @@ -25,7 +28,6 @@ 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; @@ -72,6 +74,36 @@ use std::time::Instant; static BINARY_SESSION_COUNTER: AtomicU64 = AtomicU64::new(1); +/// Tracks runtime session keys created by one binary stage connection so +/// that a mid-request failure (connection error, protocol violation, +/// runtime error) releases the execution lanes those sessions hold. +/// +/// Without this, a generation that errors before the graceful `Stop` +/// message leaks its lane session in [`RuntimeState`]: each retried +/// request uses a fresh session id, so leaked lanes accumulate until +/// every admission fails with "all execution lanes are busy" and only a +/// process restart recovers. +#[derive(Default)] +pub(super) struct ConnectionSessionTracker { + active: std::collections::BTreeSet, +} + +impl ConnectionSessionTracker { + pub(super) fn touch(&mut self, session_key: &str) { + if !self.active.contains(session_key) { + self.active.insert(session_key.to_string()); + } + } + + pub(super) fn stopped(&mut self, session_key: &str) { + self.active.remove(session_key); + } + + pub(super) fn drain(&mut self) -> Vec { + std::mem::take(&mut self.active).into_iter().collect() + } +} + #[allow(clippy::too_many_arguments)] pub(super) fn handle_binary_connection( config: &StageConfig, @@ -81,7 +113,7 @@ pub(super) fn handle_binary_connection( kv: Option<&Arc>, telemetry: &Telemetry, upstream: &mut TcpStream, - mut downstream: Option, + downstream: Option, activation_width: i32, wire_dtype: WireActivationDType, max_inflight: usize, @@ -93,13 +125,98 @@ pub(super) fn handle_binary_connection( prediction_return_sinks: &PredictionReturnSinks, first_message: StageWireMessage, ) -> 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 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, + ); + release_tracked_connection_sessions(config, runtime, telemetry, &mut session_tracker); + result +} + +/// Drops any runtime sessions this connection created but never stopped +/// gracefully, returning their execution lanes to the pool. +fn release_tracked_connection_sessions( + config: &StageConfig, + runtime: &Arc>, + telemetry: &Telemetry, + session_tracker: &mut ConnectionSessionTracker, +) { + let orphaned = session_tracker.drain(); + if orphaned.is_empty() { + return; + } + let Ok(mut runtime) = runtime.lock() else { + return; + }; + for session_key in orphaned { + match runtime.drop_session_timed(&session_key) { + Ok(drop_stats) => { + let mut attrs = crate::telemetry::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) => { + eprintln!( + "failed to reclaim orphaned binary stage session {session_key}: {error:#}" + ); + } + } } +} +#[allow(clippy::too_many_arguments)] +fn handle_binary_connection_messages( + config: &StageConfig, + topology: Option<&StageTopology>, + runtime: &Arc>, + decode_frame_batcher: &DecodeFrameBatcher, + kv: Option<&Arc>, + telemetry: &Telemetry, + upstream: &mut TcpStream, + mut 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, + session_tracker: &mut ConnectionSessionTracker, +) -> Result<()> { let connection_session_id = BINARY_SESSION_COUNTER.fetch_add(1, Ordering::Relaxed); let positional_speculation_supported = !model_requires_recurrent_state(config); let max_deferred_prefill_replies = @@ -146,6 +263,7 @@ 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(); + session_tracker.touch(&session_key); if message.kind == WireMessageKind::VerifyWindow && !positional_speculation_supported { bail!( "stage-state v10 positional speculation requires an attention-only stage; {} contains recurrent state", @@ -278,6 +396,7 @@ pub(super) fn handle_binary_connection( 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(&mut *upstream, stop_stats).context("send stop ACK")?; @@ -442,14 +561,13 @@ pub(super) fn handle_binary_connection( downstream_wire_condition, ) .context("forward prefix cache control")?; - let reply = recv_reply(&mut *downstream).context("prefix cache downstream ACK")?; + 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 = message.kind == WireMessageKind::TryRestorePrefill - && (reply.stats.kv_lookup_misses > 0 - || reply.stats.kv_lookup_errors > 0 - || reply.stats.kv_lookup_hits == 0); + let downstream_missed = + normalize_downstream_prefix_restore_reply(message.kind, &mut reply.stats); control_stats.merge(reply.stats); if downstream_missed { let mut runtime = runtime.lock().expect("runtime lock poisoned"); @@ -1300,3 +1418,32 @@ pub(super) fn handle_binary_connection( } } } + +#[cfg(test)] +mod tests { + use super::ConnectionSessionTracker; + + #[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"); + + // Simulate a mid-request stage error: session-a errored before its + // graceful Stop, session-b completed normally. + tracker.stopped("session-b"); + + assert_eq!(tracker.drain(), vec!["session-a".to_string()]); + // Idempotent: a second drain reclaims nothing. + 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()); + } +} 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 2edd6a919c..17e7f37f50 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/options.rs b/crates/skippy-server/src/binary_transport/options.rs index c24cdb6b31..b3a1bc20ed 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 4ae9728e49..02e0caa092 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 8b767b0244..a015eaef31 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,71 @@ 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<()> { + send_client_ready_hello_if_enabled(stream).context("send downstream client ready hello")?; + 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 +426,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 +437,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); + } } } } @@ -794,8 +868,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 +880,7 @@ mod tests { net::{Shutdown, TcpListener, TcpStream}, os::fd::AsRawFd, thread, - time::Duration, + time::{Duration, Instant}, }; #[test] @@ -848,9 +923,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 +954,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 +965,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 3c0d05d249..030f4ab68c 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 7150ab3ec8..ca04c54412 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 939705b089..8fb146a266 100644 --- a/crates/skippy-server/src/frontend/embedded_execution.rs +++ b/crates/skippy-server/src/frontend/embedded_execution.rs @@ -12,6 +12,7 @@ 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; @@ -36,7 +37,6 @@ 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 DispatchedEmbeddedStage { started: Instant, @@ -296,7 +296,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 +340,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 +354,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 2bb2d2d8a1..6302596dbb 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -35,7 +35,8 @@ 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, + refill_pipeline_ngram_candidates, should_open_upstream_prediction_return, + speculation_after_prefix_restore, }; use openai_frontend::{OpenAiError, OpenAiResult}; use serde_json::json; @@ -61,7 +62,9 @@ impl StageOpenAiBackend { .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() { + if should_open_upstream_prediction_return(request.native_mtp_enabled) + && let Some(prediction_return) = request.prediction_return.as_ref() + { match crate::binary_transport::direct_return::open_downstream_prediction_return_stream( request.config, request_id, @@ -711,7 +714,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 +874,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 +887,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 +926,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 +935,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(); @@ -1451,7 +1467,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 +1488,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), diff --git a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs index 028a81be36..06055ccb15 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,10 @@ pub(super) enum DirectPredictionReturnPath { ReverseFallback, } +pub(super) fn should_open_upstream_prediction_return(native_mtp_enabled: bool) -> bool { + native_mtp_enabled +} + pub(super) fn direct_prediction_return_path( verify_windows_enabled: bool, receiver_registered: bool, @@ -433,7 +455,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 +508,12 @@ mod tests { ); } + #[test] + fn pure_ngram_skips_the_blocking_upstream_return_open() { + assert!(!should_open_upstream_prediction_return(false)); + assert!(should_open_upstream_prediction_return(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 1378d475b8..3d1f569c41 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 4682a1aec2..16e0f54ce0 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/timeouts.rs b/crates/skippy-server/src/frontend/generation/timeouts.rs new file mode 100644 index 0000000000..d7e2880fbd --- /dev/null +++ b/crates/skippy-server/src/frontend/generation/timeouts.rs @@ -0,0 +1,45 @@ +use std::time::Duration; + +const STAGE_REPLY_TIMEOUT_ENV: &str = "SKIPPY_STAGE_REPLY_TIMEOUT_SECS"; +const DEFAULT_STAGE_REPLY_TIMEOUT_SECS: u64 = 30; +const MIN_STAGE_REPLY_TIMEOUT_SECS: u64 = 1; +const MAX_STAGE_REPLY_TIMEOUT_SECS: u64 = 60 * 60; + +pub(in crate::frontend) fn stage_reply_timeout() -> 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 6f0bc3a90a..9ee8a82db5 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/request.rs b/crates/skippy-server/src/frontend/request.rs index 77234ff576..fb7c800c76 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/tests/prefill.rs b/crates/skippy-server/src/frontend/tests/prefill.rs index f6c9f51c5f..7117b73535 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 2402681e8d..bdaedb9386 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 1680a14905..c4eb56811c 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/wire_messages.rs b/crates/skippy-server/src/frontend/wire_messages.rs index 7d82e7a755..bbd3974dd9 100644 --- a/crates/skippy-server/src/frontend/wire_messages.rs +++ b/crates/skippy-server/src/frontend/wire_messages.rs @@ -322,6 +322,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 +354,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/runtime_state.rs b/crates/skippy-server/src/runtime_state.rs index c250a96bf4..c66c52fac3 100644 --- a/crates/skippy-server/src/runtime_state.rs +++ b/crates/skippy-server/src/runtime_state.rs @@ -934,10 +934,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(()) } @@ -1052,6 +1053,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, @@ -1367,9 +1380,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 recurrent_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 35e5cabfae..659c710cc0 100644 --- a/crates/skippy-topology/capabilities/reviewed-family-capabilities.json +++ b/crates/skippy-topology/capabilities/reviewed-family-capabilities.json @@ -2070,5 +2070,28 @@ "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": [] + } } ] diff --git a/crates/skippy-topology/src/family_capability.rs b/crates/skippy-topology/src/family_capability.rs index 05def6beb1..1ad4dfadc2 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", @@ -753,6 +758,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, @@ -1112,6 +1134,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/tests.rs b/crates/skippy-topology/src/tests.rs index ff731c259f..f701948410 100644 --- a/crates/skippy-topology/src/tests.rs +++ b/crates/skippy-topology/src/tests.rs @@ -950,6 +950,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"); diff --git a/docs/LAYER_PACKAGE_REPOS.md b/docs/LAYER_PACKAGE_REPOS.md index 73411f2643..3a956d7e91 100644 --- a/docs/LAYER_PACKAGE_REPOS.md +++ b/docs/LAYER_PACKAGE_REPOS.md @@ -260,6 +260,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 +276,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/FAMILY_CERTIFY.md b/docs/skippy/FAMILY_CERTIFY.md index c645f64351..454bba4912 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 0c2dc326f0..74320d600e 100644 --- a/docs/skippy/FAMILY_STATUS.md +++ b/docs/skippy/FAMILY_STATUS.md @@ -10,6 +10,17 @@ customer-facing source of truth. Last updated: 2026-05-07. +## 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 | Family | Support Level | Recommended Artifact | Stage Plan | Wire | Speculative | Topology Constraint | Other Required Policy | diff --git a/docs/skippy/NEW_MODEL_ONBOARDING.md b/docs/skippy/NEW_MODEL_ONBOARDING.md index f50ccb42c7..c9106ad702 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/llama-parity-candidates.json b/docs/skippy/llama-parity-candidates.json index 22796aefdc..a156eebf58 100644 --- a/docs/skippy/llama-parity-candidates.json +++ b/docs/skippy/llama-parity-candidates.json @@ -523,6 +523,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/scripts/hf-skippy-convert-job.py b/scripts/hf-skippy-convert-job.py new file mode 100644 index 0000000000..6eb17530a4 --- /dev/null +++ b/scripts/hf-skippy-convert-job.py @@ -0,0 +1,211 @@ +# /// 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 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") + 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) + 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 converted_artifact_dir(args: argparse.Namespace) -> Path: + artifact_dir = Path(args.work_dir) / "target" / args.target_prefix + if args.upload_only: + 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 or not any(artifact_dir.glob("*.gguf")): + details = f"; missing {', '.join(missing)}" if missing else "" + raise FileNotFoundError( + f"complete converted artifact not found: {artifact_dir}{details}" + ) + 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 0000000000..e649762cc1 --- /dev/null +++ b/scripts/hf-skippy-mtp-certify-job.py @@ -0,0 +1,168 @@ +# /// 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 os +import shutil +import subprocess +import urllib.parse +import urllib.request +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 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 projector_path(args: argparse.Namespace) -> Path: + if not args.projector_url: + return require_gguf_magic(Path(args.projector)) + scheme = urllib.parse.urlparse(args.projector_url).scheme + if scheme not in ("http", "https"): + raise RuntimeError(f"unsupported projector URL scheme: {scheme!r}") + target = Path(args.projector_local_path) + target.parent.mkdir(parents=True, exist_ok=True) + temporary = target.with_suffix(f"{target.suffix}.part") + print(f"+ download {args.projector_url} -> {target}", flush=True) + with urllib.request.urlopen(args.projector_url) as response, temporary.open("wb") as output: + shutil.copyfileobj(response, output, length=8 * 1024 * 1024) + temporary.replace(target) + 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/third_party/llama.cpp/patches/0048-Add-TML-Inkling-architecture.patch b/third_party/llama.cpp/patches/0048-Add-TML-Inkling-architecture.patch new file mode 100644 index 0000000000..90d212a4d2 --- /dev/null +++ b/third_party/llama.cpp/patches/0048-Add-TML-Inkling-architecture.patch @@ -0,0 +1,5829 @@ +From b3b84213bacdb6bd0be485bd9d4e1e7d7b4da813 Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Wed, 15 Jul 2026 18:14:40 +0000 +Subject: [PATCH 48/60] 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 a309f027..db9004d8 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 91658e6d..89f51871 100644 +--- a/common/chat.cpp ++++ b/common/chat.cpp +@@ -2381,6 +2381,119 @@ static common_chat_params common_chat_params_init_cohere2moe(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) { +@@ -2798,6 +2911,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 45c001b7..88b2d520 100644 +--- a/conversion/__init__.py ++++ b/conversion/__init__.py +@@ -110,6 +110,7 @@ TEXT_MODEL_MAP: dict[str, str] = { + "HunYuanVLForConditionalGeneration": "hunyuan", + "HYV3ForCausalLM": "hunyuan", + "IQuestCoderForCausalLM": "llama", ++ "InklingForConditionalGeneration": "inkling", + "InternLM2ForCausalLM": "internlm", + "InternLM3ForCausalLM": "internlm", + "JAISLMHeadModel": "jais", +@@ -278,6 +279,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 a7cd3fd9..720cfdca 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 00000000..90ef65c3 +--- /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 83936373..086dddd9 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 7ef0ec88..6a5925ce 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 a5a3a58a..b397fb6b 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 eed46cfe..6f7cc39d 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 61237ed0..9e9b1227 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 26af9002..0352e6ba 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 00000000..98458ae8 +--- /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 00000000..bb56b898 +--- /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 e67cc7fd..b07a0122 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 7f4cfd55..a29f9910 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 ab7a3b29..43849b9c 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 e73a7b89..5373b64d 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 d55cc1ec..8b2be232 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 71e3b264..e8f2c060 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 d7dbc8b9..bbf862b3 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 e18ada53..f62c4ebb 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 31cd00f7..013bd838 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 1463169c..1e09177a 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 d3805772..1166178b 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 cdc85b08..da33f642 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 4df94b9a..f7762fae 100644 +--- a/gguf-py/gguf/constants.py ++++ b/gguf-py/gguf/constants.py +@@ -555,6 +555,7 @@ class MODEL_ARCH(IntEnum): + TALKIE = auto() + MELLUM = auto() + NANBEIGE = auto() ++ INKLING = auto() + + + class VISION_PROJECTOR_TYPE(IntEnum): +@@ -778,6 +779,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() +@@ -1159,6 +1168,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] = { +@@ -1380,6 +1390,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", +@@ -4141,6 +4158,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, +@@ -4810,6 +4859,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 8299ac25..21c79dbd 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 +@@ -2496,6 +2498,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 00000000..95874d23 +--- /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 a0945e50..a489e1ad 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" }, +@@ -792,6 +815,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}}, +@@ -833,6 +859,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}}, +@@ -961,6 +994,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; +@@ -1017,6 +1051,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 fb5eb392..74446fda 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 a2467e66..36bf3219 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 ab6e111e..e74a58a5 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 50af97f3..36568582 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 0b3626dc..4d74e352 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 033e172d..eebed9e0 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 b253d915..3b1d9b36 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 3812c594..264b0b5b 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 442387e6..3621b047 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) + "'"); + } +@@ -2113,7 +2115,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) { +@@ -2461,6 +2464,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 36d0480e..4d267514 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 03a6c34f..8ab4e403 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 9164a4dd..74942e81 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 b7c28926..65e43f67 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 00000000..6c9c1aef +--- /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 aeeb7b22..8a14222b 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 158ab28d..794c0f04 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 8ce019b3..d9704672 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -12397,6 +12397,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; +@@ -12976,6 +13073,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]; +@@ -12999,7 +13097,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; + } +@@ -15049,6 +15147,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)); ++ + test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); + test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); + test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); +diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp +index 4dd00efd..fcf025d0 100644 +--- a/tests/test-chat.cpp ++++ b/tests/test-chat.cpp +@@ -2931,6 +2931,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 00000000..94232b32 +--- /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 00000000..4e3aca44 +--- /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 659ea9e9..c7efc95a 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 18a8288b..ae574bc9 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 09204113..539a7c66 100644 +--- a/tools/mtmd/clip-impl.h ++++ b/tools/mtmd/clip-impl.h +@@ -284,6 +284,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" +@@ -346,6 +353,7 @@ + struct clip_ctx; + + enum projector_type { ++ PROJECTOR_TYPE_INKLING, + PROJECTOR_TYPE_MLP, + PROJECTOR_TYPE_MLP_NORM, + PROJECTOR_TYPE_LDP, +@@ -403,6 +411,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 8dc87549..3fd92b77 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 +@@ -355,6 +355,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; +@@ -369,6 +374,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 04614b93..8fe2a5b2 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: +@@ -1294,6 +1298,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: +@@ -1923,7 +1942,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; +@@ -2009,6 +2029,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: + { +@@ -3034,6 +3071,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; + } + +@@ -3431,6 +3473,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: +@@ -3813,6 +3865,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: +@@ -4759,6 +4815,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 00000000..c072f66a +--- /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 caed438e..c671757f 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 ed68951c..c7f0b9a1 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 d8ec72b9..eae32a76 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 36cd463b..01cb9f59 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 115cba51..044d72f6 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 6e61cf3e..119e9fc7 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: + { +@@ -1200,23 +1213,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/0049-Add-staged-execution-support-for-Inkling.patch b/third_party/llama.cpp/patches/0049-Add-staged-execution-support-for-Inkling.patch new file mode 100644 index 0000000000..daea88aaef --- /dev/null +++ b/third_party/llama.cpp/patches/0049-Add-staged-execution-support-for-Inkling.patch @@ -0,0 +1,108 @@ +From f46720057c729128913a708eb6581be7375fc431 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 13:47:13 +1000 +Subject: [PATCH 49/60] 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 6c9c1aef..3d5f8ab5 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 21db2b77..347d6985 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/0050-Advance-native-MTP-depth-per-draft-step.patch b/third_party/llama.cpp/patches/0050-Advance-native-MTP-depth-per-draft-step.patch new file mode 100644 index 0000000000..54b388b927 --- /dev/null +++ b/third_party/llama.cpp/patches/0050-Advance-native-MTP-depth-per-draft-step.patch @@ -0,0 +1,81 @@ +From 40ed78300dc40dd12586b31017dca33980ea8de7 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 14:46:51 +1000 +Subject: [PATCH 50/60] 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 347d6985..a918ce0f 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/0051-Replay-multi-depth-MTP-caches-correctly.patch b/third_party/llama.cpp/patches/0051-Replay-multi-depth-MTP-caches-correctly.patch new file mode 100644 index 0000000000..3c231e4f8b --- /dev/null +++ b/third_party/llama.cpp/patches/0051-Replay-multi-depth-MTP-caches-correctly.patch @@ -0,0 +1,368 @@ +From 68ea61226111c9ab95d959d6a3d355c3dada7228 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 14:58:19 +1000 +Subject: [PATCH 51/60] 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 4df4605b..2d546c20 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 a918ce0f..5d221cda 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) { +@@ -4413,7 +4549,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"); + } +@@ -4678,7 +4814,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/0052-Add-Inkling-multi-depth-MTP-sidecars.patch b/third_party/llama.cpp/patches/0052-Add-Inkling-multi-depth-MTP-sidecars.patch new file mode 100644 index 0000000000..991717f418 --- /dev/null +++ b/third_party/llama.cpp/patches/0052-Add-Inkling-multi-depth-MTP-sidecars.patch @@ -0,0 +1,545 @@ +From 06bac0e35ea1d224100e57af5d5ad1cf04495557 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 15:09:39 +1000 +Subject: [PATCH 52/60] 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 90ef65c3..07984290 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 c4e057ae..4b47903d 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 3621b047..be5cfb88 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -2115,10 +2115,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 3d5f8ab5..d435a1f5 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 5d221cda..61f9c508 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/0053-Carry-multimodal-Inkling-embeddings-into-MTP.patch b/third_party/llama.cpp/patches/0053-Carry-multimodal-Inkling-embeddings-into-MTP.patch new file mode 100644 index 0000000000..bfbdac9c28 --- /dev/null +++ b/third_party/llama.cpp/patches/0053-Carry-multimodal-Inkling-embeddings-into-MTP.patch @@ -0,0 +1,567 @@ +From 5b4896aa14167a09aba5534433c85ca3f49cd089 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 16:04:48 +1000 +Subject: [PATCH 53/60] 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 e9d31f55..cc2382d3 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 0145705b..11f06bf3 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 4b47903d..9d528761 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 42442879..20c121b2 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 36bf3219..27c5129b 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 e74a58a5..6e66f7fe 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 d435a1f5..c054f1d9 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 61f9c508..c76d3fbb 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) { +@@ -4939,6 +5074,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) { +@@ -4950,6 +5091,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); + } +@@ -4961,6 +5103,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/0054-Size-recurrent-memory-for-appended-MTP-layers.patch b/third_party/llama.cpp/patches/0054-Size-recurrent-memory-for-appended-MTP-layers.patch new file mode 100644 index 0000000000..f8ef7df005 --- /dev/null +++ b/third_party/llama.cpp/patches/0054-Size-recurrent-memory-for-appended-MTP-layers.patch @@ -0,0 +1,46 @@ +From d8d24d3fb5c071baec273bfd722532414952535d Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 17 Jul 2026 20:20:45 +1000 +Subject: [PATCH 54/60] 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 ef82eb97..4f1f1df8 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/0055-Restore-recurrent-verify-window-state.patch b/third_party/llama.cpp/patches/0055-Restore-recurrent-verify-window-state.patch new file mode 100644 index 0000000000..51d20c8d53 --- /dev/null +++ b/third_party/llama.cpp/patches/0055-Restore-recurrent-verify-window-state.patch @@ -0,0 +1,317 @@ +From 33aa6b672f3d6ed994ed819636ed08fcab18cfa7 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 24 Jul 2026 03:57:28 +1000 +Subject: [PATCH 55/60] Restore recurrent verify-window state + +Assisted-by: codex +--- + src/skippy.cpp | 241 ++++++++++++++++++++++++++++++++++++++++++++----- + 1 file changed, 218 insertions(+), 23 deletions(-) + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index c76d3fbb..08bb41e9 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -124,6 +124,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 +162,7 @@ struct skippy_session { + llama_token mtp_pending_draft_token = -1; + bool external_mtp_embeddings_available = false; + std::vector external_mtp_embeddings; ++ skippy_verify_checkpoint verify_checkpoint; + }; + + static thread_local skippy_session * g_skippy_external_decode_session = nullptr; +@@ -1590,7 +1602,63 @@ 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) { ++ session->verify_checkpoint = {}; ++ 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(session->ctx, session->seq_id); ++ 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( ++ session->ctx, ++ checkpoint.state.data(), ++ checkpoint.state.size(), ++ session->seq_id); ++ 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); ++ } ++ session->verify_checkpoint = 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, + const void * input_payload, +@@ -2859,6 +2927,7 @@ static void skippy_mtp_clear_session_state(skippy_session * session) { + if (session == nullptr) { + return; + } ++ session->verify_checkpoint = {}; + session->mtp_next_pos = session->n_past; + session->mtp_prefix_valid = session->n_past == 0; + session->mtp_has_pending_h = false; +@@ -4353,6 +4422,97 @@ 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; ++ const size_t prefix_bytes = skippy_activation_payload_bytes(session, token_count, flags); ++ std::vector payload(prefix_bytes); ++ if ((flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0) { ++ std::memcpy(payload.data(), checkpoint.input_payload.data(), prefix_bytes); ++ return payload; ++ } ++ ++ 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); ++ 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 ((flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0) { ++ std::memcpy( ++ payload.data() + prefix_offset, ++ checkpoint.input_payload.data() + source_offset, ++ prefix_hidden_bytes); ++ source_offset += source_hidden_bytes; ++ prefix_offset += prefix_hidden_bytes; ++ } ++ if ((flags & SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0) { ++ std::memcpy( ++ payload.data() + prefix_offset, ++ checkpoint.input_payload.data() + source_offset, ++ prefix_hidden_bytes); ++ } ++ return payload; ++} ++ ++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 = session->verify_checkpoint; ++ const uint64_t token_start = static_cast(checkpoint.token_start); ++ const uint64_t token_end = token_start + checkpoint.token_count; ++ if (!checkpoint.valid || token_count < token_start || token_count >= token_end) { ++ return skippy_success(out_error); ++ } ++ ++ const size_t read = llama_state_seq_set_data( ++ session->ctx, ++ checkpoint.state.data(), ++ checkpoint.state.size(), ++ session->seq_id); ++ if (read != checkpoint.state.size()) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to restore verify-window sequence state"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ ++ session->n_past = checkpoint.token_start; ++ const size_t accepted_count = static_cast(token_count - token_start); ++ 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(); ++ const llama_token * token_ids = checkpoint.token_ids.empty() ? nullptr : checkpoint.token_ids.data(); ++ status = skippy_decode_activation_frame( ++ session, ++ &input_desc, ++ input_payload.data(), ++ token_ids, ++ nullptr, ++ 0, ++ accepted_count, ++ false, ++ out_error); ++ } else if (accepted_count > 0) { ++ status = skippy_decode_tokens( ++ session, ++ checkpoint.token_ids.data(), ++ accepted_count, ++ false, ++ out_error); ++ } ++ if (status == SKIPPY_STATUS_OK) { ++ *out_restored = true; ++ } ++ return status; ++} ++ + extern "C" { + + struct skippy_abi_version skippy_abi_version(void) { +@@ -5551,7 +5711,12 @@ 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) { +@@ -6260,7 +6425,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( +@@ -6826,29 +7003,47 @@ enum skippy_status skippy_trim_session( + 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; ++ 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) { ++ 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; ++ } ++ } else if (dynamic_cast(memory) != nullptr) { ++ // Recurrent state cannot be rolled back by suffix removal; only a ++ // retained verify checkpoint can restore it. ++ if (static_cast(session->n_past) != token_count) { ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, ++ "recurrent trim requires a retained verify checkpoint"); ++ return SKIPPY_STATUS_UNSUPPORTED; ++ } ++ } else { ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "runtime memory type is not supported for trim"); ++ return SKIPPY_STATUS_UNSUPPORTED; + } +- } 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; + } + session->n_past = static_cast(token_count); + if (session->token_history.size() > token_count) { +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0056-skippy-preserve-verify-checkpoint-until-trim.patch b/third_party/llama.cpp/patches/0056-skippy-preserve-verify-checkpoint-until-trim.patch new file mode 100644 index 0000000000..6a00fb2163 --- /dev/null +++ b/third_party/llama.cpp/patches/0056-skippy-preserve-verify-checkpoint-until-trim.patch @@ -0,0 +1,84 @@ +From 8b19b3407034057ef91daf45c49633cce0618ce8 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 24 Jul 2026 14:08:34 +1000 +Subject: [PATCH 56/60] skippy: preserve verify checkpoint until trim + +Assisted-by: goose +--- + src/skippy.cpp | 11 +++++++++-- + 1 file changed, 9 insertions(+), 2 deletions(-) + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 08bb41e9..fd1fd472 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -2927,7 +2927,6 @@ static void skippy_mtp_clear_session_state(skippy_session * session) { + if (session == nullptr) { + return; + } +- session->verify_checkpoint = {}; + session->mtp_next_pos = session->n_past; + session->mtp_prefix_valid = session->n_past == 0; + session->mtp_has_pending_h = false; +@@ -4465,7 +4464,9 @@ static enum skippy_status skippy_restore_verify_prefix( + const skippy_verify_checkpoint & checkpoint = session->verify_checkpoint; + const uint64_t token_start = static_cast(checkpoint.token_start); + const uint64_t token_end = token_start + checkpoint.token_count; +- if (!checkpoint.valid || token_count < token_start || token_count >= token_end) { ++ const uint64_t session_token_count = static_cast(std::max(session->n_past, 0)); ++ if (!checkpoint.valid || session_token_count != token_end || ++ token_count < token_start || token_count >= token_end) { + return skippy_success(out_error); + } + +@@ -5291,6 +5292,7 @@ enum skippy_status skippy_session_set_position( + session->signal_history.resize(static_cast(n_past)); + } + skippy_mtp_clear_session_state(session); ++ session->verify_checkpoint = {}; + return skippy_success(out_error); + } + +@@ -5384,6 +5386,7 @@ enum skippy_status skippy_session_reset( + session->signal_history.clear(); + skippy_clear_chat_sampling(session); + skippy_mtp_clear_session_state(session); ++ session->verify_checkpoint = {}; + session->ctx->synchronize(); + return skippy_success(out_error); + } +@@ -6556,6 +6559,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_checkpoint = {}; + } + + enum skippy_status skippy_export_state( +@@ -6974,6 +6978,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_checkpoint = {}; + session->ctx->synchronize(); + + return skippy_success(out_error); +@@ -7054,6 +7059,7 @@ enum skippy_status skippy_trim_session( + session->signal_history.resize(static_cast(token_count)); + } + skippy_mtp_clear_session_state(session); ++ session->verify_checkpoint = {}; + session->ctx->synchronize(); + + return skippy_success(out_error); +@@ -7206,6 +7212,7 @@ enum skippy_status skippy_session_restore_prefix( + } + session->signal_history.clear(); + skippy_mtp_clear_session_state(session); ++ session->verify_checkpoint = {}; + session->ctx->synchronize(); + return skippy_success(out_error); + } +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0057-skippy-retain-pipelined-verify-checkpoints.patch b/third_party/llama.cpp/patches/0057-skippy-retain-pipelined-verify-checkpoints.patch new file mode 100644 index 0000000000..373f25b795 --- /dev/null +++ b/third_party/llama.cpp/patches/0057-skippy-retain-pipelined-verify-checkpoints.patch @@ -0,0 +1,170 @@ +From cbb9ce751bc15c3895e262f5b73812368526b42c Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 24 Jul 2026 14:54:57 +1000 +Subject: [PATCH 57/60] skippy: retain pipelined verify checkpoints + +Assisted-by: goose +--- + src/skippy.cpp | 61 ++++++++++++++++++++++++++++++-------------------- + 1 file changed, 37 insertions(+), 24 deletions(-) + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index fd1fd472..886efbcc 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -162,7 +162,7 @@ struct skippy_session { + llama_token mtp_pending_draft_token = -1; + bool external_mtp_embeddings_available = false; + std::vector external_mtp_embeddings; +- skippy_verify_checkpoint verify_checkpoint; ++ std::vector verify_checkpoints; + }; + + static thread_local skippy_session * g_skippy_external_decode_session = nullptr; +@@ -1616,7 +1616,6 @@ static enum skippy_status skippy_checkpoint_verify_window( + const skippy_activation_desc * input_desc, + const void * input_payload, + skippy_error ** out_error) { +- session->verify_checkpoint = {}; + llama_memory_t memory = session->ctx->get_memory(); + if (!skippy_memory_needs_verify_checkpoint(memory)) { + return skippy_success(out_error); +@@ -1653,7 +1652,16 @@ static enum skippy_status skippy_checkpoint_verify_window( + const uint8_t * bytes = static_cast(input_payload); + checkpoint.input_payload.assign(bytes, bytes + input_desc->payload_bytes); + } +- session->verify_checkpoint = std::move(checkpoint); ++ // Fail closed: evicting the oldest checkpoint would silently strand a ++ // window that can still be restored. The ceiling is only a growth backstop ++ // (the scheduler bounds in-flight windows to the configured depth). ++ 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); + } + +@@ -4461,35 +4469,40 @@ static enum skippy_status skippy_restore_verify_prefix( + bool * out_restored, + skippy_error ** out_error) { + *out_restored = false; +- const skippy_verify_checkpoint & checkpoint = session->verify_checkpoint; +- const uint64_t token_start = static_cast(checkpoint.token_start); +- const uint64_t token_end = token_start + checkpoint.token_count; +- const uint64_t session_token_count = static_cast(std::max(session->n_past, 0)); +- if (!checkpoint.valid || session_token_count != token_end || +- token_count < token_start || token_count >= token_end) { ++ 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( + session->ctx, +- checkpoint.state.data(), +- checkpoint.state.size(), ++ checkpoint->state.data(), ++ checkpoint->state.size(), + session->seq_id); +- if (read != checkpoint.state.size()) { ++ if (read != checkpoint->state.size()) { + skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to restore verify-window sequence state"); + return SKIPPY_STATUS_RUNTIME_ERROR; + } + +- session->n_past = checkpoint.token_start; ++ 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); + enum skippy_status status = SKIPPY_STATUS_OK; +- if (accepted_count > 0 && checkpoint.has_activation_input) { +- skippy_activation_desc input_desc = checkpoint.input_desc; ++ 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); ++ session, *checkpoint, accepted_count); + input_desc.token_count = static_cast(accepted_count); + input_desc.payload_bytes = input_payload.size(); +- const llama_token * token_ids = checkpoint.token_ids.empty() ? nullptr : checkpoint.token_ids.data(); ++ const llama_token * token_ids = checkpoint->token_ids.empty() ? nullptr : checkpoint->token_ids.data(); + status = skippy_decode_activation_frame( + session, + &input_desc, +@@ -4503,7 +4516,7 @@ static enum skippy_status skippy_restore_verify_prefix( + } else if (accepted_count > 0) { + status = skippy_decode_tokens( + session, +- checkpoint.token_ids.data(), ++ checkpoint->token_ids.data(), + accepted_count, + false, + out_error); +@@ -5292,7 +5305,7 @@ enum skippy_status skippy_session_set_position( + session->signal_history.resize(static_cast(n_past)); + } + skippy_mtp_clear_session_state(session); +- session->verify_checkpoint = {}; ++ session->verify_checkpoints.clear(); + return skippy_success(out_error); + } + +@@ -5386,7 +5399,7 @@ enum skippy_status skippy_session_reset( + session->signal_history.clear(); + skippy_clear_chat_sampling(session); + skippy_mtp_clear_session_state(session); +- session->verify_checkpoint = {}; ++ session->verify_checkpoints.clear(); + session->ctx->synchronize(); + return skippy_success(out_error); + } +@@ -6559,7 +6572,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_checkpoint = {}; ++ session->verify_checkpoints.clear(); + } + + enum skippy_status skippy_export_state( +@@ -6978,7 +6991,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_checkpoint = {}; ++ session->verify_checkpoints.clear(); + session->ctx->synchronize(); + + return skippy_success(out_error); +@@ -7059,7 +7072,7 @@ enum skippy_status skippy_trim_session( + session->signal_history.resize(static_cast(token_count)); + } + skippy_mtp_clear_session_state(session); +- session->verify_checkpoint = {}; ++ session->verify_checkpoints.clear(); + session->ctx->synchronize(); + + return skippy_success(out_error); +@@ -7212,7 +7225,7 @@ enum skippy_status skippy_session_restore_prefix( + } + session->signal_history.clear(); + skippy_mtp_clear_session_state(session); +- session->verify_checkpoint = {}; ++ 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/0058-skippy-filter-Inkling-stage-memory-by-layer-range.patch b/third_party/llama.cpp/patches/0058-skippy-filter-Inkling-stage-memory-by-layer-range.patch new file mode 100644 index 0000000000..2e4f57ae82 --- /dev/null +++ b/third_party/llama.cpp/patches/0058-skippy-filter-Inkling-stage-memory-by-layer-range.patch @@ -0,0 +1,72 @@ +From 7be339fb9da027c7ecd15a7d5eb9e2e56313d312 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Fri, 24 Jul 2026 19:02:35 +1000 +Subject: [PATCH 58/60] skippy: filter Inkling stage memory by layer range + +Assisted-by: goose +--- + src/llama-model.cpp | 11 +++++++++-- + src/skippy.cpp | 4 ++-- + 2 files changed, 11 insertions(+), 4 deletions(-) + +diff --git a/src/llama-model.cpp b/src/llama-model.cpp +index be5cfb88..527b52b8 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -6,6 +6,7 @@ + #include "llama-impl.h" + #include "llama-mmap.h" + #include "llama-cparams.h" ++#include "llama-graph.h" + #include "llama-model-loader.h" + + #include "llama-kv-cache.h" +@@ -2126,8 +2127,14 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, + 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(); }; ++ const skippy_graph_filter stage_filter = skippy_graph_get_filter(); ++ filter_attn = [&, stage_filter](uint32_t il) { ++ return il < hparams.n_layer() && ++ (!stage_filter.enabled || ++ (il >= static_cast(stage_filter.layer_start) && ++ il < static_cast(stage_filter.layer_end))); ++ }; ++ filter_recr = filter_attn; + } + } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { + filter_attn = [&](uint32_t il) { +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 886efbcc..46f8fad6 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -4512,6 +4512,7 @@ static enum skippy_status skippy_restore_verify_prefix( + 0, + accepted_count, + false, ++ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, + out_error); + } else if (accepted_count > 0) { + status = skippy_decode_tokens( +@@ -4519,6 +4520,7 @@ static enum skippy_status skippy_restore_verify_prefix( + checkpoint->token_ids.data(), + accepted_count, + false, ++ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, + out_error); + } + if (status == SKIPPY_STATUS_OK) { +@@ -7060,8 +7062,6 @@ enum skippy_status skippy_trim_session( + 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) { +-- +2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0059-Complete-Inkling-and-GLM-metadata-integration.patch b/third_party/llama.cpp/patches/0059-Complete-Inkling-and-GLM-metadata-integration.patch new file mode 100644 index 0000000000..ff5aa4596c --- /dev/null +++ b/third_party/llama.cpp/patches/0059-Complete-Inkling-and-GLM-metadata-integration.patch @@ -0,0 +1,143 @@ +From 8bd27d799ab7231e2b276645dbc6286a0fdd705e Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Mon, 27 Jul 2026 12:01:26 +1000 +Subject: [PATCH 59/60] 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 3298ed66..16ce255a 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 a489e1ad..836df068 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 74446fda..0033c761 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 527b52b8..8c0589da 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -1138,6 +1138,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 a88aae2c..57c45772 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/0060-skippy-balance-filtered-stages-across-devices.patch b/third_party/llama.cpp/patches/0060-skippy-balance-filtered-stages-across-devices.patch new file mode 100644 index 0000000000..1e0286766f --- /dev/null +++ b/third_party/llama.cpp/patches/0060-skippy-balance-filtered-stages-across-devices.patch @@ -0,0 +1,76 @@ +From 48a67d12017d1c78d563007f5658150b605d8556 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 29 Jul 2026 00:43:18 +1000 +Subject: [PATCH 60/60] 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 ec530d62..1ac8a9f9 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 da58661f..988fcf94 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 8c0589da..dbeefd8a 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -1317,13 +1317,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/tools/relay-fly-legacy/README.md b/tools/relay-fly-legacy/README.md index 8c01621c45..412a1fdfe1 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 50eafac668..842c1d35b0 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. From 533cdd05ad2a857d1028a6ad0b5f3c6b5dfc334e Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:23:48 +1000 Subject: [PATCH 02/26] feat(skippy): add Laguna staged runtime candidate --- .../skippy-topology/src/family_capability.rs | 18 +++++ crates/skippy-topology/src/lib.rs | 7 +- crates/skippy-topology/src/tests.rs | 14 ++++ docs/skippy/llama-parity-candidates.json | 8 +++ ...0046-Support-Laguna-staged-execution.patch | 72 +++++++++++++++++++ 5 files changed, 116 insertions(+), 3 deletions(-) create mode 100644 third_party/llama.cpp/patches/0046-Support-Laguna-staged-execution.patch diff --git a/crates/skippy-topology/src/family_capability.rs b/crates/skippy-topology/src/family_capability.rs index 1ad4dfadc2..045df24f56 100644 --- a/crates/skippy-topology/src/family_capability.rs +++ b/crates/skippy-topology/src/family_capability.rs @@ -150,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", @@ -556,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, @@ -1103,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)); } diff --git a/crates/skippy-topology/src/lib.rs b/crates/skippy-topology/src/lib.rs index 2015b6ff7f..ae2b5b1788 100644 --- a/crates/skippy-topology/src/lib.rs +++ b/crates/skippy-topology/src/lib.rs @@ -12,9 +12,10 @@ 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, + 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, }; diff --git a/crates/skippy-topology/src/tests.rs b/crates/skippy-topology/src/tests.rs index f701948410..88afc329d5 100644 --- a/crates/skippy-topology/src/tests.rs +++ b/crates/skippy-topology/src/tests.rs @@ -1033,6 +1033,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/skippy/llama-parity-candidates.json b/docs/skippy/llama-parity-candidates.json index a156eebf58..fb2145f973 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": "candidate", + "repo": "poolside/Laguna-S-2.1-GGUF", + "include": "laguna-s-2.1-Q4_K_M.gguf", + "notes": "P0 certification target from the official Poolside Q4_K_M GGUF; runtime-slice, package, dtype, state, cache, and mesh evidence remain pending until the Laguna certification run completes" + }, { "llama_model": "mistral3", "family": "mistral", diff --git a/third_party/llama.cpp/patches/0046-Support-Laguna-staged-execution.patch b/third_party/llama.cpp/patches/0046-Support-Laguna-staged-execution.patch new file mode 100644 index 0000000000..3ab8dcff07 --- /dev/null +++ b/third_party/llama.cpp/patches/0046-Support-Laguna-staged-execution.patch @@ -0,0 +1,72 @@ +From d94d179b554ff3de781b122b305a22d1a63f1b2d Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Mon, 27 Jul 2026 17:19:02 +1000 +Subject: [PATCH] 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 fb55ec12..7144a389 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 47f0a549..da0d97bf 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -4112,6 +4112,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) + From 431a94439abccf97319c53cef696fe8fa4ab1361 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Mon, 27 Jul 2026 18:32:22 +1000 Subject: [PATCH 03/26] docs(skippy): record Laguna package certification --- crates/model-package/src/script.rs | 3 ++ .../src/scripts/split-model-job.sh | 46 ++++++++++++++++++- docs/skippy/LLAMA_PARITY.md | 1 + docs/skippy/llama-parity-candidates.json | 4 +- 4 files changed, 50 insertions(+), 4 deletions(-) diff --git a/crates/model-package/src/script.rs b/crates/model-package/src/script.rs index 5c41965c31..8928da8df3 100644 --- a/crates/model-package/src/script.rs +++ b/crates/model-package/src/script.rs @@ -134,6 +134,9 @@ mod tests { 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")); diff --git a/crates/model-package/src/scripts/split-model-job.sh b/crates/model-package/src/scripts/split-model-job.sh index 576c1d2ab3..438612d714 100755 --- a/crates/model-package/src/scripts/split-model-job.sh +++ b/crates/model-package/src/scripts/split-model-job.sh @@ -568,6 +568,7 @@ experimental_warning = ( if experimental else "" ) +api = HfApi(token=os.environ["HF_TOKEN"]) def sha256(path: Path) -> str: digest = hashlib.sha256() @@ -601,6 +602,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"]: @@ -647,6 +680,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), @@ -693,10 +730,16 @@ 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: {yaml_quote(source_pipeline_tag)} tags: @@ -827,7 +870,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/docs/skippy/LLAMA_PARITY.md b/docs/skippy/LLAMA_PARITY.md index 4fdf1d0b75..d42940cb4f 100644 --- a/docs/skippy/LLAMA_PARITY.md +++ b/docs/skippy/LLAMA_PARITY.md @@ -370,6 +370,7 @@ 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@797f31cc813c0d643b7205118a901732f39de8ab` | package validated | untested | untested | package-only validation passed: 48 layers; `0..24` and `24..48` stages materialized at 33,500,313,312 and 34,752,130,336 bytes | | `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/llama-parity-candidates.json b/docs/skippy/llama-parity-candidates.json index fb2145f973..c91d482ad7 100644 --- a/docs/skippy/llama-parity-candidates.json +++ b/docs/skippy/llama-parity-candidates.json @@ -188,10 +188,10 @@ { "llama_model": "laguna", "family": "laguna", - "status": "candidate", + "status": "certified_package_only", "repo": "poolside/Laguna-S-2.1-GGUF", "include": "laguna-s-2.1-Q4_K_M.gguf", - "notes": "P0 certification target from the official Poolside Q4_K_M GGUF; runtime-slice, package, dtype, state, cache, and mesh evidence remain pending until the Laguna certification run completes" + "notes": "package-only certified from poolside/Laguna-S-2.1-GGUF:Q4_K_M at source revision edd093522473dc7313b0738d8b4116b7f8b9745f; HF job 6a6708667ef3c0846496a429 published meshllm/laguna-s-2.1-Q4_K_M-layers, and pinned package revision 797f31cc813c0d643b7205118a901732f39de8ab passed package materialization with 48 layers and two 24-layer stages of 33500313312 and 34752130336 bytes; runtime-slice, activation-wire dtype, state, cache, and mesh evidence remain pending" }, { "llama_model": "mistral3", From b3b043f502cc76564dccc2dc9124cbcd2baa9e6e Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:37:33 +1000 Subject: [PATCH 04/26] fix(skippy): integrate hybrid verify recovery and lane cleanup --- .../mesh-llm-config/src/model_validation.rs | 32 +- .../src/protocol/convert.rs | 4 +- .../src/protocol/tests/announcements.rs | 4 +- crates/skippy-ffi/src/lib.rs | 10 +- crates/skippy-protocol/src/binary/mod.rs | 27 + crates/skippy-protocol/src/binary/types.rs | 16 +- crates/skippy-protocol/src/lib.rs | 12 +- crates/skippy-runtime/src/activation.rs | 3 + crates/skippy-runtime/src/session.rs | 18 + crates/skippy-server/README.md | 11 +- .../src/binary_transport/binary_messaging.rs | 1 + .../binary_messaging/async_forwarder.rs | 75 +++ .../binary_messaging/connection.rs | 75 ++- .../binary_messaging/session_tracker.rs | 102 ++++ .../src/binary_transport/stage_execution.rs | 1 + .../src/frontend/embedded_execution.rs | 58 ++ .../src/frontend/embedded_generation.rs | 41 +- .../src/frontend/generation/server.rs | 2 +- .../src/frontend/linear_proposal.rs | 10 + .../src/frontend/native_mtp/verify_window.rs | 22 +- .../skippy-server/src/frontend/speculative.rs | 38 +- .../src/frontend/wire_messages.rs | 26 + crates/skippy-server/src/runtime_state.rs | 10 + docs/design/TESTING.md | 4 +- docs/design/message_protocol.md | 2 +- docs/skippy/DATA_FLOW.md | 8 +- docs/skippy/PIPELINED_VERIFY_WINDOW.md | 4 +- docs/skippy/WAN_SPLIT_PERF.md | 2 +- .../tests/test_windows_native_runtime_deps.py | 2 +- ...brid-verify-window-state-across-trim.patch | 501 ++++++++++++++++++ ...-accepted-verify-checkpoints-exactly.patch | 160 ++++++ 31 files changed, 1200 insertions(+), 81 deletions(-) create mode 100644 crates/skippy-server/src/binary_transport/binary_messaging/session_tracker.rs create mode 100644 third_party/llama.cpp/patches/0048-Preserve-hybrid-verify-window-state-across-trim.patch create mode 100644 third_party/llama.cpp/patches/0049-skippy-retire-accepted-verify-checkpoints-exactly.patch diff --git a/crates/mesh-llm-config/src/model_validation.rs b/crates/mesh-llm-config/src/model_validation.rs index a564a2d1b6..0505452fb0 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/protocol/convert.rs b/crates/mesh-llm-host-runtime/src/protocol/convert.rs index 72a384ad5a..4232786e6b 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 06dda3a548..691796c9de 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/skippy-ffi/src/lib.rs b/crates/skippy-ffi/src/lib.rs index e3802ec95d..7a499c0f26 100644 --- a/crates/skippy-ffi/src/lib.rs +++ b/crates/skippy-ffi/src/lib.rs @@ -1,6 +1,6 @@ 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 = 34; 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; @@ -856,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; @@ -1553,6 +1554,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, diff --git a/crates/skippy-protocol/src/binary/mod.rs b/crates/skippy-protocol/src/binary/mod.rs index 58b1cce798..f7dfbdce9b 100644 --- a/crates/skippy-protocol/src/binary/mod.rs +++ b/crates/skippy-protocol/src/binary/mod.rs @@ -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; diff --git a/crates/skippy-protocol/src/binary/types.rs b/crates/skippy-protocol/src/binary/types.rs index a95b4ce6eb..4541f7f9f0 100644 --- a/crates/skippy-protocol/src/binary/types.rs +++ b/crates/skippy-protocol/src/binary/types.rs @@ -5,9 +5,9 @@ use super::{ invalid_data, }; -// v11 adds the Inkling MTP embedding sideband, which changes activation payload length. -// Stage peers must be upgraded together so an older reader rejects the header before it can -// leave the additional sideband bytes in the stream. +// 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; @@ -57,6 +57,7 @@ pub enum WireMessageKind { DecodeReadout = 8, DecodeLightCtx = 9, VerifyWindow = 21, + RetireVerifyWindow = 22, StateExport = 13, ConfigureGeneration = 14, ProbePrefill = 15, @@ -96,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) } @@ -141,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")), } } @@ -322,6 +328,7 @@ impl StageStateHeader { kind, WireMessageKind::StateImport | WireMessageKind::StateExport ) || kind.is_session_control() + || kind.is_verify_retirement() || kind.is_generation_control() { return true; @@ -400,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!( @@ -704,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 b2aeec0ba2..da4626698b 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-runtime/src/activation.rs b/crates/skippy-runtime/src/activation.rs index ace1e95714..a50994b769 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/session.rs b/crates/skippy-runtime/src/session.rs index 550309cc54..5bf6891688 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-server/README.md b/crates/skippy-server/README.md index 6bd26ec875..90393a1134 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 diff --git a/crates/skippy-server/src/binary_transport/binary_messaging.rs b/crates/skippy-server/src/binary_transport/binary_messaging.rs index 96ded79f96..559dea069e 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -36,6 +36,7 @@ pub(in crate::binary_transport) mod async_forwarder; mod connection; mod prefill_recording; pub(in crate::binary_transport) mod reply; +mod session_tracker; mod summary; mod telemetry; 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 d334f6c887..a64b50751c 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 c0ee57922d..20135d3834 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -5,6 +5,7 @@ use super::reply::{ configure_prediction_return_stream, normalize_downstream_prefix_restore_reply, reply_window_for_message, }; +use super::session_tracker::{ConnectionSessionTracker, release_tracked_connection_sessions}; use super::summary::BinaryMessageObservation; use super::summary::BinaryRequestSummary; use super::telemetry::UpstreamReplyWriteSpan; @@ -151,50 +152,6 @@ pub(super) fn handle_binary_connection( result } -/// Drops any runtime sessions this connection created but never stopped -/// gracefully, returning their execution lanes to the pool. -fn release_tracked_connection_sessions( - config: &StageConfig, - runtime: &Arc>, - telemetry: &Telemetry, - session_tracker: &mut ConnectionSessionTracker, -) { - let orphaned = session_tracker.drain(); - if orphaned.is_empty() { - return; - } - let Ok(mut runtime) = runtime.lock() else { - return; - }; - for session_key in orphaned { - match runtime.drop_session_timed(&session_key) { - Ok(drop_stats) => { - let mut attrs = crate::telemetry::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) => { - eprintln!( - "failed to reclaim orphaned binary stage session {session_key}: {error:#}" - ); - } - } - } -} - #[allow(clippy::too_many_arguments)] fn handle_binary_connection_messages( config: &StageConfig, @@ -266,7 +223,7 @@ fn handle_binary_connection_messages( session_tracker.touch(&session_key); if message.kind == WireMessageKind::VerifyWindow && !positional_speculation_supported { bail!( - "stage-state v10 positional speculation requires an attention-only stage; {} contains recurrent state", + "stage-state v11 positional speculation requires an attention-only stage; {} contains recurrent state", config.stage_id ); } @@ -403,6 +360,34 @@ fn handle_binary_connection_messages( continue; } + if message.kind.is_verify_retirement() { + if let Some(forwarder) = async_forwarder.as_mut() { + 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")?; + { + let mut runtime = runtime.lock().expect("runtime lock poisoned"); + runtime + .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")?; + } + 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() { 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 0000000000..5f395e2d15 --- /dev/null +++ b/crates/skippy-server/src/binary_transport/binary_messaging/session_tracker.rs @@ -0,0 +1,102 @@ +use super::telemetry::insert_runtime_session_stats; +use crate::{ + runtime_state::RuntimeState, + telemetry::{Telemetry, lifecycle_attrs}, +}; +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, +) { + let orphaned = session_tracker.drain(); + if orphaned.is_empty() { + return; + } + let Ok(mut runtime) = runtime.lock() else { + return; + }; + 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) => { + eprintln!( + "failed to reclaim orphaned binary stage session {session_key}: {error:#}" + ); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::ConnectionSessionTracker; + + #[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()); + } +} diff --git a/crates/skippy-server/src/binary_transport/stage_execution.rs b/crates/skippy-server/src/binary_transport/stage_execution.rs index a015eaef31..e9923b5795 100644 --- a/crates/skippy-server/src/binary_transport/stage_execution.rs +++ b/crates/skippy-server/src/binary_transport/stage_execution.rs @@ -595,6 +595,7 @@ pub(crate) fn run_binary_stage_message( | WireMessageKind::StateExport | WireMessageKind::ConfigureGeneration | WireMessageKind::TrimSession + | WireMessageKind::RetireVerifyWindow | WireMessageKind::ProbePrefill | WireMessageKind::RestorePrefill | WireMessageKind::TryRestorePrefill diff --git a/crates/skippy-server/src/frontend/embedded_execution.rs b/crates/skippy-server/src/frontend/embedded_execution.rs index 8fb146a266..c35bc283bc 100644 --- a/crates/skippy-server/src/frontend/embedded_execution.rs +++ b/crates/skippy-server/src/frontend/embedded_execution.rs @@ -16,6 +16,7 @@ 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; @@ -38,6 +39,13 @@ const DIRECT_RETURN_FALLBACK_POLL: Duration = Duration::from_millis(10); // normal WAN verify traversal while remaining shorter than the HTTP client's // request timeout. +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, stats: StageReplyStats, @@ -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<'_>, diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 6302596dbb..34823496ea 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, @@ -949,7 +950,7 @@ impl StageOpenAiBackend { 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!( @@ -1257,6 +1258,24 @@ impl StageOpenAiBackend { let fully_accepted_window = !native_mtp_verify_decision.rejected && native_mtp_verify_decision.accepted_proposal_tokens == window.proposal_tokens.len(); + let checkpoint_no_longer_needed = verify_checkpoint_no_longer_needed( + native_mtp_verify_decision.commit_count, + window.input_tokens.len(), + ); + if checkpoint_no_longer_needed { + 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(), + }, + )?; + } let pipeline_continues = fully_accepted_window; let accepted_candidate_tokens = native_mtp_verify_decision.accepted_proposal_tokens; @@ -1573,6 +1592,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/generation/server.rs b/crates/skippy-server/src/frontend/generation/server.rs index c802e39d0a..cb8692fc62 100644 --- a/crates/skippy-server/src/frontend/generation/server.rs +++ b/crates/skippy-server/src/frontend/generation/server.rs @@ -313,7 +313,7 @@ pub fn embedded_openai_backend(args: EmbeddedOpenAiArgs) -> Result 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 b6852d3a7c..e9cad05cdf 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,24 @@ 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( + native_mtp_verify_decision.commit_count, + 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/speculative.rs b/crates/skippy-server/src/frontend/speculative.rs index e85d6b0542..6abb8e61e1 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/wire_messages.rs b/crates/skippy-server/src/frontend/wire_messages.rs index bbd3974dd9..4fc48b42bc 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, diff --git a/crates/skippy-server/src/runtime_state.rs b/crates/skippy-server/src/runtime_state.rs index c66c52fac3..05731f3838 100644 --- a/crates/skippy-server/src/runtime_state.rs +++ b/crates/skippy-server/src/runtime_state.rs @@ -512,6 +512,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)?; diff --git a/docs/design/TESTING.md b/docs/design/TESTING.md index 7affe68128..aa0867f5a1 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 6867c337b9..34c5a03b1d 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 58810c6559..9f78d34a17 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/PIPELINED_VERIFY_WINDOW.md b/docs/skippy/PIPELINED_VERIFY_WINDOW.md index 3f7b090b89..f40d78ba53 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/WAN_SPLIT_PERF.md b/docs/skippy/WAN_SPLIT_PERF.md index 9a5631dec5..f40f065909 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/scripts/tests/test_windows_native_runtime_deps.py b/scripts/tests/test_windows_native_runtime_deps.py index 7ee63fbf31..5ad3a94a5a 100644 --- a/scripts/tests/test_windows_native_runtime_deps.py +++ b/scripts/tests/test_windows_native_runtime_deps.py @@ -129,7 +129,7 @@ def test_package_verifier_accepts_closed_windows_dependency_graph(self): "runtime": { "id": artifact.name, "mesh_version": "0.72.1", - "skippy_abi": "0.1.32", + "skippy_abi": "0.1.34", "platform": { "os": "windows", "arch": "x86_64", diff --git a/third_party/llama.cpp/patches/0048-Preserve-hybrid-verify-window-state-across-trim.patch b/third_party/llama.cpp/patches/0048-Preserve-hybrid-verify-window-state-across-trim.patch new file mode 100644 index 0000000000..ff255be73d --- /dev/null +++ b/third_party/llama.cpp/patches/0048-Preserve-hybrid-verify-window-state-across-trim.patch @@ -0,0 +1,501 @@ +From 3de16ed4e7b4dc99758193736d2b4ccc7b311602 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Tue, 28 Jul 2026 19:12:14 +1000 +Subject: [PATCH 48/49] 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 21db2b77d..f42731cb6 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; +@@ -149,6 +161,8 @@ struct skippy_session { + bool mtp_has_pending_draft = false; + llama_pos mtp_pending_draft_pos = 0; + llama_token mtp_pending_draft_token = -1; ++ bool reset_required = false; ++ std::vector verify_checkpoints; + }; + + struct skippy_tensor_meta { +@@ -246,7 +260,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; +@@ -1566,6 +1588,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, +@@ -1763,7 +1855,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) { +@@ -4035,6 +4130,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) { +@@ -4801,6 +5059,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); + } + +@@ -4894,6 +5153,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); + } +@@ -5221,7 +5482,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) { +@@ -5282,7 +5553,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, +@@ -5930,7 +6202,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( +@@ -6049,6 +6333,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( +@@ -6467,6 +6752,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); +@@ -6495,30 +6781,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) { +@@ -6529,6 +6831,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); +@@ -6681,6 +6984,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.54.0 + diff --git a/third_party/llama.cpp/patches/0049-skippy-retire-accepted-verify-checkpoints-exactly.patch b/third_party/llama.cpp/patches/0049-skippy-retire-accepted-verify-checkpoints-exactly.patch new file mode 100644 index 0000000000..febe2f5f19 --- /dev/null +++ b/third_party/llama.cpp/patches/0049-skippy-retire-accepted-verify-checkpoints-exactly.patch @@ -0,0 +1,160 @@ +From 1f12cf3984b66dde5c5340d1d54614856fb123af Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 29 Jul 2026 11:12:07 +1000 +Subject: [PATCH 49/49] 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 e9d31f559..44afc6b74 100644 +--- a/include/skippy.h ++++ b/include/skippy.h +@@ -524,6 +524,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 0145705b2..e89410d03 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 f42731cb6..a22f0dde6 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" + +@@ -6758,6 +6759,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 158ab28df..846d4d954 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.54.0 + From ffa48e9959bc9b4f2977fb0886c3a20aa46ace9b Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:52:53 +1000 Subject: [PATCH 05/26] docs(skippy): record Laguna M5 parity --- docs/skippy/FAMILY_STATUS.md | 9 +++++---- docs/skippy/LLAMA_PARITY.md | 2 +- docs/skippy/llama-parity-candidates.json | 2 +- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/skippy/FAMILY_STATUS.md b/docs/skippy/FAMILY_STATUS.md index 74320d600e..bc6c9a7063 100644 --- a/docs/skippy/FAMILY_STATUS.md +++ b/docs/skippy/FAMILY_STATUS.md @@ -8,7 +8,7 @@ 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-28. ## Context Capacity Contract @@ -99,12 +99,13 @@ correctness smoke or from GGUF metadata alone. ## 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 +Laguna S 2.1 (package-backed two-stage single-step; broader certification pending) ``` ## Exceptions diff --git a/docs/skippy/LLAMA_PARITY.md b/docs/skippy/LLAMA_PARITY.md index d42940cb4f..f162480807 100644 --- a/docs/skippy/LLAMA_PARITY.md +++ b/docs/skippy/LLAMA_PARITY.md @@ -370,7 +370,7 @@ 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@797f31cc813c0d643b7205118a901732f39de8ab` | package validated | untested | untested | package-only validation passed: 48 layers; `0..24` and `24..48` stages materialized at 33,500,313,312 and 34,752,130,336 bytes | +| `laguna` | `meshllm/laguna-s-2.1-Q4_K_M-layers@797f31cc813c0d643b7205118a901732f39de8ab` | package-backed `single-step` passed at split `24` | untested | untested | M5 Max f16 parity passed against pinned full-model revision `edd093522473dc7313b0738d8b4116b7f8b9745f`: baseline token `674` matched staged token `674`; width `3072`, 12,288-byte activation frame, 6,144-byte wire payload; chain, dtype matrix, state/cache, and mesh evidence remain pending | | `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/llama-parity-candidates.json b/docs/skippy/llama-parity-candidates.json index c91d482ad7..113e1ac854 100644 --- a/docs/skippy/llama-parity-candidates.json +++ b/docs/skippy/llama-parity-candidates.json @@ -191,7 +191,7 @@ "status": "certified_package_only", "repo": "poolside/Laguna-S-2.1-GGUF", "include": "laguna-s-2.1-Q4_K_M.gguf", - "notes": "package-only certified from poolside/Laguna-S-2.1-GGUF:Q4_K_M at source revision edd093522473dc7313b0738d8b4116b7f8b9745f; HF job 6a6708667ef3c0846496a429 published meshllm/laguna-s-2.1-Q4_K_M-layers, and pinned package revision 797f31cc813c0d643b7205118a901732f39de8ab passed package materialization with 48 layers and two 24-layer stages of 33500313312 and 34752130336 bytes; runtime-slice, activation-wire dtype, state, cache, and mesh evidence remain pending" + "notes": "package-only certification status retained while broader gates remain pending; source revision edd093522473dc7313b0738d8b4116b7f8b9745f (SHA-256 a34c74e46688122bef83122f4133031bababbefcf57436dde97048c91e2cc6ff) and package revision 797f31cc813c0d643b7205118a901732f39de8ab passed M5 Max package-backed single-step parity at split 24 with f16 activation transport: baseline token 674 matched staged token 674, activation width 3072, payload 12288 bytes, wire payload 6144 bytes; package materialization, production-shaped OpenAI smoke, and two-stage direct prediction return also passed; chain, dtype matrix, state, cache, and mesh evidence remain pending" }, { "llama_model": "mistral3", From 6275d9498caa5cbcea230ad49e8f8650bc74aa26 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:08:55 +1000 Subject: [PATCH 06/26] test(skippy): certify Laguna three-stage parity --- crates/skippy-correctness/src/runner/mod.rs | 1 + .../src/runner/prediction_return.rs | 208 ++++++++++++++++++ .../src/runner/split_chain.rs | 26 ++- docs/skippy/FAMILY_STATUS.md | 2 +- docs/skippy/LLAMA_PARITY.md | 2 +- docs/skippy/llama-parity-candidates.json | 2 +- 6 files changed, 231 insertions(+), 10 deletions(-) create mode 100644 crates/skippy-correctness/src/runner/prediction_return.rs diff --git a/crates/skippy-correctness/src/runner/mod.rs b/crates/skippy-correctness/src/runner/mod.rs index d9a1018015..ed3a8301b8 100644 --- a/crates/skippy-correctness/src/runner/mod.rs +++ b/crates/skippy-correctness/src/runner/mod.rs @@ -1,4 +1,5 @@ pub(crate) mod native_mtp; +mod prediction_return; mod single_step; mod split_chain; mod stage_execution; 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 0000000000..f759f6c1a5 --- /dev/null +++ b/crates/skippy-correctness/src/runner/prediction_return.rs @@ -0,0 +1,208 @@ +use std::{ + env, io, + io::Write, + net::{SocketAddr, TcpListener, TcpStream}, + sync::{ + Arc, + 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, + 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 (sender, receiver) = mpsc::channel(); + let thread = thread::spawn(move || { + let result = accept_prediction_return(listener, &thread_shutdown, &sender); + if let Err(error) = result { + let _ = sender.send(Err(format!("{error:#}"))); + } + }); + Ok(Self { + bind_addr, + receiver, + shutdown, + 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 Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} + +fn accept_prediction_return( + listener: TcpListener, + shutdown: &AtomicBool, + 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")?; + 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(); + } +} diff --git a/crates/skippy-correctness/src/runner/split_chain.rs b/crates/skippy-correctness/src/runner/split_chain.rs index c10667b9fd..52f4b9cbd8 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/docs/skippy/FAMILY_STATUS.md b/docs/skippy/FAMILY_STATUS.md index bc6c9a7063..d6fdf5f738 100644 --- a/docs/skippy/FAMILY_STATUS.md +++ b/docs/skippy/FAMILY_STATUS.md @@ -105,7 +105,7 @@ smoke, reviewed topology records, and family-specific policy notes are updated. ```text Gemma text -Laguna S 2.1 (package-backed two-stage single-step; broader certification pending) +Laguna S 2.1 (package-backed M5 two-stage single-step and three-stage chain parity; broader certification pending) ``` ## Exceptions diff --git a/docs/skippy/LLAMA_PARITY.md b/docs/skippy/LLAMA_PARITY.md index f162480807..827305b3d6 100644 --- a/docs/skippy/LLAMA_PARITY.md +++ b/docs/skippy/LLAMA_PARITY.md @@ -370,7 +370,7 @@ 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@797f31cc813c0d643b7205118a901732f39de8ab` | package-backed `single-step` passed at split `24` | untested | untested | M5 Max f16 parity passed against pinned full-model revision `edd093522473dc7313b0738d8b4116b7f8b9745f`: baseline token `674` matched staged token `674`; width `3072`, 12,288-byte activation frame, 6,144-byte wire payload; chain, dtype matrix, state/cache, and mesh evidence remain pending | +| `laguna` | `meshllm/laguna-s-2.1-Q4_K_M-layers@797f31cc813c0d643b7205118a901732f39de8ab` | package-backed `single-step` passed at split `24`; three-stage `chain` passed at splits `16,32` | untested | untested | M5 Max f16 parity passed against pinned full-model revision `edd093522473dc7313b0738d8b4116b7f8b9745f`: baseline token `674` matched both two-stage and three-stage predictions; width `3072`, 12,288-byte activation frame, 6,144-byte wire payload; dtype matrix, state/cache, and M4+M5 mesh evidence remain pending | | `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/llama-parity-candidates.json b/docs/skippy/llama-parity-candidates.json index 113e1ac854..4fad6b0c18 100644 --- a/docs/skippy/llama-parity-candidates.json +++ b/docs/skippy/llama-parity-candidates.json @@ -191,7 +191,7 @@ "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 797f31cc813c0d643b7205118a901732f39de8ab passed M5 Max package-backed single-step parity at split 24 with f16 activation transport: baseline token 674 matched staged token 674, activation width 3072, payload 12288 bytes, wire payload 6144 bytes; package materialization, production-shaped OpenAI smoke, and two-stage direct prediction return also passed; chain, dtype matrix, state, cache, and mesh evidence remain pending" + "notes": "package-only certification status retained while broader gates remain pending; source revision edd093522473dc7313b0738d8b4116b7f8b9745f (SHA-256 a34c74e46688122bef83122f4133031bababbefcf57436dde97048c91e2cc6ff) and package revision 797f31cc813c0d643b7205118a901732f39de8ab 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; package materialization, production-shaped OpenAI smoke, and direct prediction return also passed; dtype matrix, state, cache, and M4+M5 mesh evidence remain pending" }, { "llama_model": "mistral3", From 4ee94f73615f3ffb205ec32d2d8238f424778162 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:48:49 +1000 Subject: [PATCH 07/26] docs(skippy): record Laguna distributed serving smoke --- docs/skippy/FAMILY_STATUS.md | 2 +- docs/skippy/LLAMA_PARITY.md | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/docs/skippy/FAMILY_STATUS.md b/docs/skippy/FAMILY_STATUS.md index d6fdf5f738..d6ccf966ac 100644 --- a/docs/skippy/FAMILY_STATUS.md +++ b/docs/skippy/FAMILY_STATUS.md @@ -105,7 +105,7 @@ smoke, reviewed topology records, and family-specific policy notes are updated. ```text Gemma text -Laguna S 2.1 (package-backed M5 two-stage single-step and three-stage chain parity; broader certification pending) +Laguna S 2.1 (package-backed M5 two-stage single-step and three-stage chain parity, plus M5 Metal + Vast CUDA distributed OpenAI serving at split 24; broader certification pending) ``` ## Exceptions diff --git a/docs/skippy/LLAMA_PARITY.md b/docs/skippy/LLAMA_PARITY.md index 827305b3d6..a9a7c59659 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,7 +371,7 @@ 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@797f31cc813c0d643b7205118a901732f39de8ab` | package-backed `single-step` passed at split `24`; three-stage `chain` passed at splits `16,32` | untested | untested | M5 Max f16 parity passed against pinned full-model revision `edd093522473dc7313b0738d8b4116b7f8b9745f`: baseline token `674` matched both two-stage and three-stage predictions; width `3072`, 12,288-byte activation frame, 6,144-byte wire payload; dtype matrix, state/cache, and M4+M5 mesh evidence remain pending | +| `laguna` | `meshllm/laguna-s-2.1-Q4_K_M-layers@797f31cc813c0d643b7205118a901732f39de8ab` | package-backed `single-step` passed at split `24`; three-stage `chain` passed at splits `16,32`; distributed OpenAI serving passed at split `24` | untested | untested | M5 Max f16 parity passed against pinned full-model revision `edd093522473dc7313b0738d8b4116b7f8b9745f`: baseline token `674` matched both two-stage and three-stage predictions. A package-backed M5 Metal `0..24` + Australian Vast.ai RTX PRO 6000 Blackwell CUDA `24..48` run at commit `226fb584` returned HTTP 200 for a 43-token prompt plus one generated token (`19`); the decode activation used the expected 12,288-byte frame and 6,144-byte f16 wire payload. Dtype matrix, state/cache, and mesh-node orchestration remain pending. | | `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 | From f0cd1da5725f7f4d57f58976de1a6ac93c4aa833 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:09:50 +1000 Subject: [PATCH 08/26] docs(skippy): record real Laguna mesh evidence --- docs/skippy/FAMILY_STATUS.md | 2 +- docs/skippy/LLAMA_PARITY.md | 2 +- docs/skippy/llama-parity-candidates.json | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/skippy/FAMILY_STATUS.md b/docs/skippy/FAMILY_STATUS.md index d6ccf966ac..402707463c 100644 --- a/docs/skippy/FAMILY_STATUS.md +++ b/docs/skippy/FAMILY_STATUS.md @@ -105,7 +105,7 @@ smoke, reviewed topology records, and family-specific policy notes are updated. ```text Gemma text -Laguna S 2.1 (package-backed M5 two-stage single-step and three-stage chain parity, plus M5 Metal + Vast CUDA distributed OpenAI serving at split 24; broader certification pending) +Laguna S 2.1 (package-backed M5 two-stage single-step and three-stage chain parity, plus ordinary M5 Metal 0..36 + Vast CUDA 36..48 Mesh serving at configured context 262144 with a 44,460-token prompt; broader certification pending) ``` ## Exceptions diff --git a/docs/skippy/LLAMA_PARITY.md b/docs/skippy/LLAMA_PARITY.md index a9a7c59659..e5a77050bd 100644 --- a/docs/skippy/LLAMA_PARITY.md +++ b/docs/skippy/LLAMA_PARITY.md @@ -371,7 +371,7 @@ Rows with distributed evidence call out the second backend explicitly. | `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@797f31cc813c0d643b7205118a901732f39de8ab` | package-backed `single-step` passed at split `24`; three-stage `chain` passed at splits `16,32`; distributed OpenAI serving passed at split `24` | untested | untested | M5 Max f16 parity passed against pinned full-model revision `edd093522473dc7313b0738d8b4116b7f8b9745f`: baseline token `674` matched both two-stage and three-stage predictions. A package-backed M5 Metal `0..24` + Australian Vast.ai RTX PRO 6000 Blackwell CUDA `24..48` run at commit `226fb584` returned HTTP 200 for a 43-token prompt plus one generated token (`19`); the decode activation used the expected 12,288-byte frame and 6,144-byte f16 wire payload. Dtype matrix, state/cache, and mesh-node orchestration remain pending. | +| `laguna` | `meshllm/laguna-s-2.1-Q4_K_M-layers@797f31cc813c0d643b7205118a901732f39de8ab` | package-backed `single-step` passed at split `24`; three-stage `chain` passed at splits `16,32`; ordinary two-node Mesh OpenAI serving passed at `0..36 / 36..48` with configured context `262144` | untested | untested | M5 Max f16 parity passed against pinned full-model revision `edd093522473dc7313b0738d8b4116b7f8b9745f`: baseline token `674` matched both two-stage and three-stage predictions. The real Mesh run used M5 Max/Metal plus an Australian Vast RTX 6000 Ada/CUDA worker, normal planning and lifecycle, and a direct Iroh path with observed RTT samples from 24 ms to 320 ms. `doctor split` reported ready with no blockers. A 44,460-token prompt plus deterministic 65-token completion passed with exact output agreement. Suffix N-gram accepted 32 of 64 speculative tokens on the repeat/copy workload and raised decode from 5.72 to 9.61 tok/s; a 56-token low-overlap control produced no proposals. Exact-prefix reuse did not occur (`cached_tokens=0`), leaving streaming TTFT near 190 seconds. Equal `24/24` placement at full context exhausted the 48 GB CUDA worker; the working capacity-bounded topology was `36/12`. Dtype matrix, state/cache, tool-loop, and larger-context soak evidence remain pending. | | `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/llama-parity-candidates.json b/docs/skippy/llama-parity-candidates.json index 4fad6b0c18..89c37cd881 100644 --- a/docs/skippy/llama-parity-candidates.json +++ b/docs/skippy/llama-parity-candidates.json @@ -191,7 +191,7 @@ "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 797f31cc813c0d643b7205118a901732f39de8ab 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; package materialization, production-shaped OpenAI smoke, and direct prediction return also passed; dtype matrix, state, cache, and M4+M5 mesh evidence remain pending" + "notes": "package-only certification status retained while broader gates remain pending; source revision edd093522473dc7313b0738d8b4116b7f8b9745f (SHA-256 a34c74e46688122bef83122f4133031bababbefcf57436dde97048c91e2cc6ff) and package revision 797f31cc813c0d643b7205118a901732f39de8ab 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 Vast RTX 6000 Ada CUDA private Mesh run passed normal 0..36/36..48 placement, direct Iroh transport, doctor readiness, configured context 262144, a 44460-token prompt, exact 65-token completion agreement, and suffix N-gram verification with 32/64 speculative tokens accepted; the 24/24 full-context CUDA stage exhausted 48 GB, exact-prefix reuse reported zero cached tokens, and dtype matrix, state/cache, tool-loop, and larger-context soak evidence remain pending" }, { "llama_model": "mistral3", From 92d1b90bdd647c88dd720a3c6c9c4cfdbe00c640 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:33:44 +1000 Subject: [PATCH 09/26] fix(skippy): make stage memory truly layer-local --- .../src/inference/skippy/deployment.rs | 2 +- .../src/inference/skippy/family_policy.rs | 102 +++++++++++++++++- .../src/inference/skippy/stage/mod.rs | 8 +- ...staged-runtime-memory-to-layer-range.patch | 72 +++++++++++++ 4 files changed, 179 insertions(+), 5 deletions(-) create mode 100644 third_party/llama.cpp/patches/0047-Filter-staged-runtime-memory-to-layer-range.patch 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 572d36c14a..e11dca3249 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 bb5f22be31..f13e59d9f2 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 @@ -51,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, @@ -66,7 +84,8 @@ impl FamilyPolicy { // 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).unwrap_or(0); + 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 @@ -343,7 +362,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(), @@ -351,10 +379,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 { @@ -430,6 +463,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); @@ -795,4 +866,29 @@ 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, + 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/stage/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs index db2da2f8ff..8200e6ecfc 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/third_party/llama.cpp/patches/0047-Filter-staged-runtime-memory-to-layer-range.patch b/third_party/llama.cpp/patches/0047-Filter-staged-runtime-memory-to-layer-range.patch new file mode 100644 index 0000000000..296b5c2d39 --- /dev/null +++ b/third_party/llama.cpp/patches/0047-Filter-staged-runtime-memory-to-layer-range.patch @@ -0,0 +1,72 @@ +From 3a5f43b7fe0e90ac28f0a2856f87c291f5990a29 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Tue, 28 Jul 2026 14:23:09 +1000 +Subject: [PATCH] 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 10633a51..7ad59b29 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -2042,6 +2042,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; + +@@ -2102,7 +2118,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 +@@ -2127,6 +2143,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( +@@ -2204,6 +2223,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) From 93f3b18ece282f9311cddce589035e13222d7af7 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:35:27 +1000 Subject: [PATCH 10/26] fix(packaging): preserve source revisions --- crates/mesh-llm-commands/src/model_package.rs | 27 +++++++++- crates/model-package/src/script.rs | 7 +++ .../src/scripts/split-model-job.sh | 21 ++++---- .../src/runner/prediction_return.rs | 49 +++++++++++++++++-- 4 files changed, 89 insertions(+), 15 deletions(-) diff --git a/crates/mesh-llm-commands/src/model_package.rs b/crates/mesh-llm-commands/src/model_package.rs index 98579c69b3..8414ac9118 100644 --- a/crates/mesh-llm-commands/src/model_package.rs +++ b/crates/mesh-llm-commands/src/model_package.rs @@ -315,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); @@ -674,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/model-package/src/script.rs b/crates/model-package/src/script.rs index 8928da8df3..fd6a2912be 100644 --- a/crates/model-package/src/script.rs +++ b/crates/model-package/src/script.rs @@ -214,4 +214,11 @@ mod tests { 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 438612d714..76605a2106 100755 --- a/crates/model-package/src/scripts/split-model-job.sh +++ b/crates/model-package/src/scripts/split-model-job.sh @@ -461,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 @@ -471,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", @@ -498,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", diff --git a/crates/skippy-correctness/src/runner/prediction_return.rs b/crates/skippy-correctness/src/runner/prediction_return.rs index f759f6c1a5..c58f82236b 100644 --- a/crates/skippy-correctness/src/runner/prediction_return.rs +++ b/crates/skippy-correctness/src/runner/prediction_return.rs @@ -1,9 +1,9 @@ use std::{ env, io, io::Write, - net::{SocketAddr, TcpListener, TcpStream}, + net::{Shutdown, SocketAddr, TcpListener, TcpStream}, sync::{ - Arc, + Arc, Mutex, atomic::{AtomicBool, Ordering}, mpsc, }, @@ -24,6 +24,7 @@ pub(super) struct PredictionReturnListener { bind_addr: SocketAddr, receiver: mpsc::Receiver>, shutdown: Arc, + connection: Arc>>, thread: Option>, } @@ -39,9 +40,12 @@ impl PredictionReturnListener { .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, &sender); + let result = + accept_prediction_return(listener, &thread_shutdown, &thread_connection, &sender); if let Err(error) = result { let _ = sender.send(Err(format!("{error:#}"))); } @@ -50,6 +54,7 @@ impl PredictionReturnListener { bind_addr, receiver, shutdown, + connection, thread: Some(thread), }) } @@ -75,6 +80,11 @@ impl PredictionReturnListener { 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(); } @@ -84,6 +94,7 @@ impl Drop for PredictionReturnListener { fn accept_prediction_return( listener: TcpListener, shutdown: &AtomicBool, + connection: &Mutex>, sender: &mpsc::Sender>, ) -> Result<()> { let mut stream = loop { @@ -102,6 +113,18 @@ fn accept_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(); @@ -205,4 +228,24 @@ mod tests { 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)); + } } From cd3d8a6c6cfd393b3c92ac0790ea480cdfa2a79a Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Tue, 28 Jul 2026 21:47:31 +1000 Subject: [PATCH 11/26] feat(skippy): honor package verification depth --- .../skippy/resolver/native_mtp_tests.rs | 14 +++-- .../inference/skippy/resolver/speculative.rs | 2 +- crates/skippy-model-package/src/package.rs | 3 + crates/skippy-model-package/src/preflight.rs | 61 ++++++++++++++++++- crates/skippy-runtime/src/package.rs | 4 ++ .../reviewed-family-capabilities.json | 20 ++++++ crates/skippy-topology/src/tests.rs | 10 +++ docs/LAYER_PACKAGE_REPOS.md | 5 +- docs/skippy/FAMILY_STATUS.md | 2 +- docs/skippy/LLAMA_PARITY.md | 2 +- docs/skippy/SUFFIX_NGRAM_PROPOSER.md | 14 ++++- docs/skippy/llama-parity-candidates.json | 2 +- docs/specs/layer-package-repos.md | 7 ++- 13 files changed, 130 insertions(+), 16 deletions(-) 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 df1b58fe67..172acf54de 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/speculative.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs index 0320d219f5..dbacb3ee90 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/skippy-model-package/src/package.rs b/crates/skippy-model-package/src/package.rs index 49f954a1ff..5b67c5cd38 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 +275,8 @@ struct PackageWindowPolicy { initial_window: u32, min_window: u32, max_window: u32, + #[serde(default)] + pipeline_depth: Option, } #[derive(Debug, Deserialize)] @@ -1014,6 +1018,16 @@ fn validate_window_policy( "use positive window sizes", ); } + if window.pipeline_depth == Some(0) { + report.error( + "invalid_window_policy_pipeline_depth", + format!( + "speculative strategy {name} window_policy.pipeline_depth must be greater than zero" + ), + Some("model-package.json".to_string()), + "set pipeline_depth to a positive in-flight verification-window capacity", + ); + } if window.min_window > window.max_window { report.error( "invalid_window_policy_bounds", @@ -1135,6 +1149,7 @@ 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, } } @@ -1735,7 +1750,8 @@ mod tests { "default": "fixed", "initial_window": 1, "min_window": 1, - "max_window": 1 + "max_window": 1, + "pipeline_depth": 2 } } } @@ -1770,6 +1786,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 +2121,48 @@ 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(); + } + fn assert_issue(report: &PackagePreflightReport, code: &str) { assert!( report.issues.iter().any(|issue| issue.code == code), diff --git a/crates/skippy-runtime/src/package.rs b/crates/skippy-runtime/src/package.rs index 69a13be16b..efbcc6eea2 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-topology/capabilities/reviewed-family-capabilities.json b/crates/skippy-topology/capabilities/reviewed-family-capabilities.json index 659c710cc0..840640c648 100644 --- a/crates/skippy-topology/capabilities/reviewed-family-capabilities.json +++ b/crates/skippy-topology/capabilities/reviewed-family-capabilities.json @@ -2093,5 +2093,25 @@ "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/tests.rs b/crates/skippy-topology/src/tests.rs index 88afc329d5..a044d44c88 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, diff --git a/docs/LAYER_PACKAGE_REPOS.md b/docs/LAYER_PACKAGE_REPOS.md index 3a956d7e91..30f5980dce 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 diff --git a/docs/skippy/FAMILY_STATUS.md b/docs/skippy/FAMILY_STATUS.md index 402707463c..c7235b07b5 100644 --- a/docs/skippy/FAMILY_STATUS.md +++ b/docs/skippy/FAMILY_STATUS.md @@ -105,7 +105,7 @@ smoke, reviewed topology records, and family-specific policy notes are updated. ```text Gemma text -Laguna S 2.1 (package-backed M5 two-stage single-step and three-stage chain parity, plus ordinary M5 Metal 0..36 + Vast CUDA 36..48 Mesh serving at configured context 262144 with a 44,460-token prompt; broader certification pending) +Laguna S 2.1 Q4_K_M (pinned package-backed M5 two-stage single-step and three-stage chain parity, plus ordinary M5 Metal 0..36 + Australian Vast CUDA 36..48 Mesh serving with a 44,460-token prompt; suffix N-gram depth 2 is published as the package default, with a fresh live no-override confirmation still pending; broader family certification remains pending) ``` ## Exceptions diff --git a/docs/skippy/LLAMA_PARITY.md b/docs/skippy/LLAMA_PARITY.md index e5a77050bd..4804ac475c 100644 --- a/docs/skippy/LLAMA_PARITY.md +++ b/docs/skippy/LLAMA_PARITY.md @@ -371,7 +371,7 @@ Rows with distributed evidence call out the second backend explicitly. | `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@797f31cc813c0d643b7205118a901732f39de8ab` | package-backed `single-step` passed at split `24`; three-stage `chain` passed at splits `16,32`; ordinary two-node Mesh OpenAI serving passed at `0..36 / 36..48` with configured context `262144` | untested | untested | M5 Max f16 parity passed against pinned full-model revision `edd093522473dc7313b0738d8b4116b7f8b9745f`: baseline token `674` matched both two-stage and three-stage predictions. The real Mesh run used M5 Max/Metal plus an Australian Vast RTX 6000 Ada/CUDA worker, normal planning and lifecycle, and a direct Iroh path with observed RTT samples from 24 ms to 320 ms. `doctor split` reported ready with no blockers. A 44,460-token prompt plus deterministic 65-token completion passed with exact output agreement. Suffix N-gram accepted 32 of 64 speculative tokens on the repeat/copy workload and raised decode from 5.72 to 9.61 tok/s; a 56-token low-overlap control produced no proposals. Exact-prefix reuse did not occur (`cached_tokens=0`), leaving streaming TTFT near 190 seconds. Equal `24/24` placement at full context exhausted the 48 GB CUDA worker; the working capacity-bounded topology was `36/12`. Dtype matrix, state/cache, tool-loop, and larger-context soak evidence remain pending. | +| `laguna` | `meshllm/laguna-s-2.1-Q4_K_M-layers@0c467ad441ee94cb5a76f626294d963c4048507d` | package-backed `single-step` passed at split `24`; three-stage `chain` passed at splits `16,32`; ordinary two-node Mesh OpenAI serving passed at `0..36 / 36..48` with configured context `262144` | untested | untested | M5 Max f16 parity passed against pinned full-model revision `edd093522473dc7313b0738d8b4116b7f8b9745f`: baseline token `674` matched both two-stage and three-stage predictions. The real Mesh run used M5 Max/Metal plus an Australian Vast RTX 6000 Ada/CUDA worker, normal planning and lifecycle, and direct Iroh transport. A final three-request recovery probe used a 44,460-token prompt plus deterministic 65-token completion and returned byte-identical content on every request; decode measured 19.74, 18.24, and 15.42 tok/s. Each request proposed 64 suffix N-gram tokens and accepted 32 at verify pipeline depth `2`; the exact-prefix hit reused 44,416 prompt tokens. A 63-request stability harness returned 63 HTTP 200 responses with no unsupported-trim, all-lanes-busy, reset, decode, slot, or proactive-eviction fatal errors. Structured tool use remains uncertified: the formal harness passed 6/13 checks because textual `` output was not consistently projected into OpenAI `tool_calls`. Package manifest SHA-256 `0250cfb54ceeb94a9c71e48df447f780e32fc625553844d6403770f315be0237` now publishes suffix N-gram (`min=5`, `max=32`, proposal cap `48`) and verify pipeline depth `2` as the default. Resolver coverage proves a no-config launch consumes that package policy; a fresh live M5 + Vast no-override confirmation remains pending. Equal `24/24` placement at full context exhausted the 48 GB CUDA worker; the working capacity-bounded topology was `36/12`. Q8 wire, state handoff, structured tool use, and full native-context saturation 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/SUFFIX_NGRAM_PROPOSER.md b/docs/skippy/SUFFIX_NGRAM_PROPOSER.md index ba1bbd7fc4..396d84d964 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/llama-parity-candidates.json b/docs/skippy/llama-parity-candidates.json index 89c37cd881..6b1613ca54 100644 --- a/docs/skippy/llama-parity-candidates.json +++ b/docs/skippy/llama-parity-candidates.json @@ -191,7 +191,7 @@ "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 797f31cc813c0d643b7205118a901732f39de8ab 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 Vast RTX 6000 Ada CUDA private Mesh run passed normal 0..36/36..48 placement, direct Iroh transport, doctor readiness, configured context 262144, a 44460-token prompt, exact 65-token completion agreement, and suffix N-gram verification with 32/64 speculative tokens accepted; the 24/24 full-context CUDA stage exhausted 48 GB, exact-prefix reuse reported zero cached tokens, and dtype matrix, state/cache, tool-loop, and larger-context soak evidence remain pending" + "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", diff --git a/docs/specs/layer-package-repos.md b/docs/specs/layer-package-repos.md index 4f4b0beaf1..03e6a92cd3 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. | From 57368b2dfeb53b0bd55b15308d09447c7bc045f7 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:16:33 +1000 Subject: [PATCH 12/26] test(skippy): assert Laguna cache policy --- .../mesh-llm-host-runtime/src/inference/skippy/family_policy.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 f13e59d9f2..750114d8e1 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 @@ -706,7 +706,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 { From 115e77af1e6a91be661076195e05933ecb08143d Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:41:38 +1000 Subject: [PATCH 13/26] llama: linearize Laguna Inkling and recovery patches --- ...048-Support-Laguna-staged-execution.patch} | 10 +- ...taged-runtime-memory-to-layer-range.patch} | 16 +- ...> 0050-Add-TML-Inkling-architecture.patch} | 130 +++---- ...taged-execution-support-for-Inkling.patch} | 8 +- ...nce-native-MTP-depth-per-draft-step.patch} | 6 +- ...ay-multi-depth-MTP-caches-correctly.patch} | 12 +- ...dd-Inkling-multi-depth-MTP-sidecars.patch} | 16 +- ...timodal-Inkling-embeddings-into-MTP.patch} | 26 +- ...estore-recurrent-verify-window-state.patch | 317 ------------------ ...rent-memory-for-appended-MTP-layers.patch} | 6 +- ...reserve-verify-checkpoint-until-trim.patch | 84 ----- ...nkling-and-GLM-metadata-integration.patch} | 16 +- ...-retain-pipelined-verify-checkpoints.patch | 170 ---------- ...ance-filtered-stages-across-devices.patch} | 12 +- ...-Inkling-stage-memory-by-layer-range.patch | 72 ---- ...rid-verify-window-state-across-trim.patch} | 44 +-- ...accepted-verify-checkpoints-exactly.patch} | 18 +- ...t-trims-after-failed-verify-recovery.patch | 27 ++ 18 files changed, 187 insertions(+), 803 deletions(-) rename third_party/llama.cpp/patches/{0046-Support-Laguna-staged-execution.patch => 0048-Support-Laguna-staged-execution.patch} (91%) rename third_party/llama.cpp/patches/{0047-Filter-staged-runtime-memory-to-layer-range.patch => 0049-Filter-staged-runtime-memory-to-layer-range.patch} (87%) rename third_party/llama.cpp/patches/{0048-Add-TML-Inkling-architecture.patch => 0050-Add-TML-Inkling-architecture.patch} (99%) rename third_party/llama.cpp/patches/{0049-Add-staged-execution-support-for-Inkling.patch => 0051-Add-staged-execution-support-for-Inkling.patch} (96%) rename third_party/llama.cpp/patches/{0050-Advance-native-MTP-depth-per-draft-step.patch => 0052-Advance-native-MTP-depth-per-draft-step.patch} (94%) rename third_party/llama.cpp/patches/{0051-Replay-multi-depth-MTP-caches-correctly.patch => 0053-Replay-multi-depth-MTP-caches-correctly.patch} (97%) rename third_party/llama.cpp/patches/{0052-Add-Inkling-multi-depth-MTP-sidecars.patch => 0054-Add-Inkling-multi-depth-MTP-sidecars.patch} (98%) rename third_party/llama.cpp/patches/{0053-Carry-multimodal-Inkling-embeddings-into-MTP.patch => 0055-Carry-multimodal-Inkling-embeddings-into-MTP.patch} (97%) delete mode 100644 third_party/llama.cpp/patches/0055-Restore-recurrent-verify-window-state.patch rename third_party/llama.cpp/patches/{0054-Size-recurrent-memory-for-appended-MTP-layers.patch => 0056-Size-recurrent-memory-for-appended-MTP-layers.patch} (91%) delete mode 100644 third_party/llama.cpp/patches/0056-skippy-preserve-verify-checkpoint-until-trim.patch rename third_party/llama.cpp/patches/{0059-Complete-Inkling-and-GLM-metadata-integration.patch => 0057-Complete-Inkling-and-GLM-metadata-integration.patch} (95%) delete mode 100644 third_party/llama.cpp/patches/0057-skippy-retain-pipelined-verify-checkpoints.patch rename third_party/llama.cpp/patches/{0060-skippy-balance-filtered-stages-across-devices.patch => 0058-skippy-balance-filtered-stages-across-devices.patch} (92%) delete mode 100644 third_party/llama.cpp/patches/0058-skippy-filter-Inkling-stage-memory-by-layer-range.patch rename third_party/llama.cpp/patches/{0048-Preserve-hybrid-verify-window-state-across-trim.patch => 0059-Preserve-hybrid-verify-window-state-across-trim.patch} (93%) rename third_party/llama.cpp/patches/{0049-skippy-retire-accepted-verify-checkpoints-exactly.patch => 0060-skippy-retire-accepted-verify-checkpoints-exactly.patch} (92%) create mode 100644 third_party/llama.cpp/patches/0061-Reject-trims-after-failed-verify-recovery.patch diff --git a/third_party/llama.cpp/patches/0046-Support-Laguna-staged-execution.patch b/third_party/llama.cpp/patches/0048-Support-Laguna-staged-execution.patch similarity index 91% rename from third_party/llama.cpp/patches/0046-Support-Laguna-staged-execution.patch rename to third_party/llama.cpp/patches/0048-Support-Laguna-staged-execution.patch index 3ab8dcff07..94dc4f1a0e 100644 --- a/third_party/llama.cpp/patches/0046-Support-Laguna-staged-execution.patch +++ b/third_party/llama.cpp/patches/0048-Support-Laguna-staged-execution.patch @@ -1,7 +1,7 @@ -From d94d179b554ff3de781b122b305a22d1a63f1b2d Mon Sep 17 00:00:00 2001 +From 3140a1e0aa272ccaf89d280fa14f293e9234ed08 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Mon, 27 Jul 2026 17:19:02 +1000 -Subject: [PATCH] Support Laguna staged execution +Subject: [PATCH 48/61] Support Laguna staged execution --- src/models/laguna.cpp | 19 ++++++++++++++++--- @@ -9,7 +9,7 @@ Subject: [PATCH] Support Laguna staged execution 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/src/models/laguna.cpp b/src/models/laguna.cpp -index fb55ec12..7144a389 100644 +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 @@ -56,10 +56,10 @@ index fb55ec12..7144a389 100644 cb(cur, "result_norm", -1); res->t_embd = cur; diff --git a/src/skippy.cpp b/src/skippy.cpp -index 47f0a549..da0d97bf 100644 +index 21db2b77d..3965e0252 100644 --- a/src/skippy.cpp +++ b/src/skippy.cpp -@@ -4112,6 +4112,7 @@ static enum skippy_status skippy_finish_model_open( +@@ -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 && diff --git a/third_party/llama.cpp/patches/0047-Filter-staged-runtime-memory-to-layer-range.patch b/third_party/llama.cpp/patches/0049-Filter-staged-runtime-memory-to-layer-range.patch similarity index 87% rename from third_party/llama.cpp/patches/0047-Filter-staged-runtime-memory-to-layer-range.patch rename to third_party/llama.cpp/patches/0049-Filter-staged-runtime-memory-to-layer-range.patch index 296b5c2d39..ae62f3f393 100644 --- a/third_party/llama.cpp/patches/0047-Filter-staged-runtime-memory-to-layer-range.patch +++ b/third_party/llama.cpp/patches/0049-Filter-staged-runtime-memory-to-layer-range.patch @@ -1,23 +1,22 @@ -From 3a5f43b7fe0e90ac28f0a2856f87c291f5990a29 Mon Sep 17 00:00:00 2001 +From 426444c8c8ee62b6322ac6794ada658df0f077f4 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Tue, 28 Jul 2026 14:23:09 +1000 -Subject: [PATCH] Filter staged runtime memory to layer range +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 10633a51..7ad59b29 100644 +index 442387e6a..b23a7e8d8 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp -@@ -2042,6 +2042,22 @@ ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int i +@@ -2047,6 +2047,22 @@ ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int i return layers[il].rope_short; } @@ -40,7 +39,7 @@ index 10633a51..7ad59b29 100644 llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const llama_cparams & cparams) const { llama_memory_i * res; -@@ -2102,7 +2118,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, +@@ -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, @@ -49,7 +48,7 @@ index 10633a51..7ad59b29 100644 } 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 -@@ -2127,6 +2143,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, +@@ -2132,6 +2148,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, }; } @@ -59,7 +58,7 @@ index 10633a51..7ad59b29 100644 if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { // Use hybrid-iswa for hybrid models with SWA res = new llama_memory_hybrid_iswa( -@@ -2204,6 +2223,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, +@@ -2209,6 +2228,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } } @@ -70,3 +69,4 @@ index 10633a51..7ad59b29 100644 -- 2.50.1 (Apple Git-155) + diff --git a/third_party/llama.cpp/patches/0048-Add-TML-Inkling-architecture.patch b/third_party/llama.cpp/patches/0050-Add-TML-Inkling-architecture.patch similarity index 99% rename from third_party/llama.cpp/patches/0048-Add-TML-Inkling-architecture.patch rename to third_party/llama.cpp/patches/0050-Add-TML-Inkling-architecture.patch index 90d212a4d2..71379ec73f 100644 --- a/third_party/llama.cpp/patches/0048-Add-TML-Inkling-architecture.patch +++ b/third_party/llama.cpp/patches/0050-Add-TML-Inkling-architecture.patch @@ -1,7 +1,7 @@ -From b3b84213bacdb6bd0be485bd9d4e1e7d7b4da813 Mon Sep 17 00:00:00 2001 +From ef9aea72c4da2cb4d92851ef8b4dd456fb48d070 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 15 Jul 2026 18:14:40 +0000 -Subject: [PATCH 48/60] Add TML Inkling architecture +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 @@ -86,7 +86,7 @@ and audio support, and backend op tests at production shapes. create mode 100644 tools/mtmd/models/inkling.cpp diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp -index a309f027..db9004d8 100644 +index a309f0276..db9004d84 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( @@ -121,7 +121,7 @@ index a309f027..db9004d8 100644 // Build ID parsers if keys are provided common_peg_parser id_parser = eps(); diff --git a/common/chat.cpp b/common/chat.cpp -index 91658e6d..89f51871 100644 +index 91658e6df..89f518716 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -2381,6 +2381,119 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t @@ -260,7 +260,7 @@ index 91658e6d..89f51871 100644 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 45c001b7..88b2d520 100644 +index 45c001b78..88b2d5201 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -110,6 +110,7 @@ TEXT_MODEL_MAP: dict[str, str] = { @@ -280,7 +280,7 @@ index 45c001b7..88b2d520 100644 "JanusForConditionalGeneration": "januspro", "KimiK25ForConditionalGeneration": "kimivl", diff --git a/conversion/base.py b/conversion/base.py -index a7cd3fd9..720cfdca 100644 +index a7cd3fd90..720cfdca3 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -268,8 +268,17 @@ class ModelBase: @@ -339,7 +339,7 @@ index a7cd3fd9..720cfdca 100644 self.gguf_writer.write_kv_data_to_file() diff --git a/conversion/inkling.py b/conversion/inkling.py new file mode 100644 -index 00000000..90ef65c3 +index 000000000..90ef65c3b --- /dev/null +++ b/conversion/inkling.py @@ -0,0 +1,350 @@ @@ -694,7 +694,7 @@ index 00000000..90ef65c3 + + raise ValueError(f"unexpected Inkling vision tensor {name!r}") diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h -index 83936373..086dddd9 100644 +index 839363731..086dddd96 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -11,7 +11,7 @@ extern "C" { @@ -707,7 +707,7 @@ index 83936373..086dddd9 100644 #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h -index 7ef0ec88..6a5925ce 100644 +index 7ef0ec88f..6a5925ce4 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -435,8 +435,9 @@ extern "C" { @@ -759,7 +759,7 @@ index 7ef0ec88..6a5925ce 100644 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 a5a3a58a..b397fb6b 100644 +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( @@ -790,7 +790,7 @@ index a5a3a58a..b397fb6b 100644 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 eed46cfe..6f7cc39d 100644 +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 @@ -818,7 +818,7 @@ index eed46cfe..6f7cc39d 100644 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 61237ed0..9e9b1227 100644 +index 61237ed06..9e9b1227b 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -8340,10 +8340,11 @@ template @@ -917,7 +917,7 @@ index 61237ed0..9e9b1227 100644 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 26af9002..0352e6ba 100644 +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) { @@ -967,7 +967,7 @@ index 26af9002..0352e6ba 100644 } diff --git a/ggml/src/ggml-cuda/fattn-banded.cu b/ggml/src/ggml-cuda/fattn-banded.cu new file mode 100644 -index 00000000..98458ae8 +index 000000000..98458ae86 --- /dev/null +++ b/ggml/src/ggml-cuda/fattn-banded.cu @@ -0,0 +1,247 @@ @@ -1220,7 +1220,7 @@ index 00000000..98458ae8 +} diff --git a/ggml/src/ggml-cuda/fattn-banded.cuh b/ggml/src/ggml-cuda/fattn-banded.cuh new file mode 100644 -index 00000000..bb56b898 +index 000000000..bb56b8986 --- /dev/null +++ b/ggml/src/ggml-cuda/fattn-banded.cuh @@ -0,0 +1,7 @@ @@ -1232,7 +1232,7 @@ index 00000000..bb56b898 + +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 e67cc7fd..b07a0122 100644 +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 { @@ -1281,7 +1281,7 @@ index e67cc7fd..b07a0122 100644 // 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 7f4cfd55..a29f9910 100644 +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( @@ -1522,7 +1522,7 @@ index 7f4cfd55..a29f9910 100644 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 ab7a3b29..43849b9c 100644 +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 @@ -1550,7 +1550,7 @@ index ab7a3b29..43849b9c 100644 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 e73a7b89..5373b64d 100644 +index e73a7b890..5373b64d9 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -25,6 +25,7 @@ @@ -1734,7 +1734,7 @@ index e73a7b89..5373b64d 100644 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 d55cc1ec..8b2be232 100644 +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( @@ -1760,7 +1760,7 @@ index d55cc1ec..8b2be232 100644 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 71e3b264..e8f2c060 100644 +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_ @@ -1816,7 +1816,7 @@ index 71e3b264..e8f2c060 100644 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 d7dbc8b9..bbf862b3 100644 +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( @@ -1838,7 +1838,7 @@ index d7dbc8b9..bbf862b3 100644 if constexpr (has_fusion) { diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu -index e18ada53..f62c4ebb 100644 +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( @@ -1860,7 +1860,7 @@ index e18ada53..f62c4ebb 100644 // 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 31cd00f7..013bd838 100644 +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 @@ -1873,7 +1873,7 @@ index 31cd00f7..013bd838 100644 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 1463169c..1e09177a 100644 +index 1463169cf..1e09177a4 100644 --- a/ggml/src/ggml-cuda/ssm-conv.cu +++ b/ggml/src/ggml-cuda/ssm-conv.cu @@ -5,8 +5,8 @@ @@ -1940,7 +1940,7 @@ index 1463169c..1e09177a 100644 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 d3805772..1166178b 100644 +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 @@ -1952,7 +1952,7 @@ index d3805772..1166178b 100644 if (rpc_get) { diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c -index cdc85b08..da33f642 100644 +index cdc85b083..da33f642a 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1067,6 +1067,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { @@ -2060,7 +2060,7 @@ index cdc85b08..da33f642 100644 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 4df94b9a..f7762fae 100644 +index 4df94b9a3..f7762fae1 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -555,6 +555,7 @@ class MODEL_ARCH(IntEnum): @@ -2156,7 +2156,7 @@ index 4df94b9a..f7762fae 100644 GEMMA3NV = "gemma3nv" GEMMA3NA = "gemma3na" diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py -index 8299ac25..21c79dbd 100644 +index 8299ac25b..21c79dbd8 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -63,6 +63,7 @@ class TensorNameMap: @@ -2231,7 +2231,7 @@ index 8299ac25..21c79dbd 100644 "model.layers.{bid}.residual_layernorm", diff --git a/models/templates/Inkling.jinja b/models/templates/Inkling.jinja new file mode 100644 -index 00000000..95874d23 +index 000000000..95874d23d --- /dev/null +++ b/models/templates/Inkling.jinja @@ -0,0 +1,514 @@ @@ -2750,7 +2750,7 @@ index 00000000..95874d23 +{%- endif -%} +{#- Unsloth translation to jinja from TML's parser #} diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp -index a0945e50..a489e1ad 100644 +index a0945e501..a489e1ad5 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -144,6 +144,7 @@ static const std::map LLM_ARCH_NAMES = { @@ -2841,7 +2841,7 @@ index a0945e50..a489e1ad 100644 default: return true; diff --git a/src/llama-arch.h b/src/llama-arch.h -index fb5eb392..74446fda 100644 +index fb5eb3920..74446fdaf 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -149,6 +149,7 @@ enum llm_arch { @@ -2897,7 +2897,7 @@ index fb5eb392..74446fda 100644 LLM_TENSOR_VISEXP_ATTN_OUT, LLM_TENSOR_VISEXP_FFN_GATE, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp -index a2467e66..36bf3219 100644 +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( @@ -2955,7 +2955,7 @@ index a2467e66..36bf3219 100644 res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); diff --git a/src/llama-graph.h b/src/llama-graph.h -index ab6e111e..e74a58a5 100644 +index ab6e111e1..e74a58a56 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1083,7 +1083,8 @@ struct llm_graph_context { @@ -2969,7 +2969,7 @@ index ab6e111e..e74a58a5 100644 // 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 50af97f3..36568582 100644 +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 { @@ -2985,7 +2985,7 @@ index 50af97f3..36568582 100644 // for RWKV models return token_shift_count * n_embd; diff --git a/src/llama-hparams.h b/src/llama-hparams.h -index 0b3626dc..4d74e352 100644 +index 0b3626dc8..4d74e3527 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -80,6 +80,17 @@ struct llama_hparams { @@ -3007,7 +3007,7 @@ index 0b3626dc..4d74e352 100644 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 033e172d..eebed9e0 100644 +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 { @@ -3136,7 +3136,7 @@ index 033e172d..eebed9e0 100644 kv->set_input_k_rot(dst); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h -index b253d915..3b1d9b36 100644 +index b253d9158..3b1d9b36f 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -197,6 +197,9 @@ public: @@ -3188,7 +3188,7 @@ index b253d915..3b1d9b36 100644 // // update context diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp -index 3812c594..264b0b5b 100644 +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) { @@ -3200,7 +3200,7 @@ index 3812c594..264b0b5b 100644 default: return true; diff --git a/src/llama-model.cpp b/src/llama-model.cpp -index 442387e6..3621b047 100644 +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 @@ -3212,7 +3212,7 @@ index 442387e6..3621b047 100644 default: throw std::runtime_error(std::string("unsupported model architecture: '") + llm_arch_name(arch) + "'"); } -@@ -2113,7 +2115,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, +@@ -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; @@ -3222,7 +3222,7 @@ index 442387e6..3621b047 100644 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) { -@@ -2461,6 +2464,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { +@@ -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: @@ -3231,7 +3231,7 @@ index 442387e6..3621b047 100644 // 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 36d0480e..4d267514 100644 +index 36d0480e5..4d267514b 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -532,6 +532,15 @@ struct llama_layer { @@ -3251,7 +3251,7 @@ index 36d0480e..4d267514 100644 struct llama_device { diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index 03a6c34f..8ab4e403 100644 +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 @@ -3272,7 +3272,7 @@ index 03a6c34f..8ab4e403 100644 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 9164a4dd..74942e81 100644 +index 9164a4dd8..74942e812 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -433,6 +433,12 @@ struct llm_tokenizer_bpe : llm_tokenizer { @@ -3300,7 +3300,7 @@ index 9164a4dd..74942e81 100644 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 b7c28926..65e43f67 100644 +index b7c289263..65e43f671 100644 --- a/src/llama-vocab.h +++ b/src/llama-vocab.h @@ -65,6 +65,7 @@ enum llama_vocab_pre_type { @@ -3313,7 +3313,7 @@ index b7c28926..65e43f67 100644 struct LLM_KV; diff --git a/src/models/inkling.cpp b/src/models/inkling.cpp new file mode 100644 -index 00000000..6c9c1aef +index 000000000..6c9c1aef8 --- /dev/null +++ b/src/models/inkling.cpp @@ -0,0 +1,695 @@ @@ -4013,7 +4013,7 @@ index 00000000..6c9c1aef + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h -index aeeb7b22..8a14222b 100644 +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 { @@ -4037,7 +4037,7 @@ index aeeb7b22..8a14222b 100644 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 158ab28d..794c0f04 100644 +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) @@ -4056,7 +4056,7 @@ index 158ab28d..794c0f04 100644 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 8ce019b3..d9704672 100644 +index 8ce019b3b..d97046725 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -12397,6 +12397,103 @@ struct test_glm_dsa_compact_k_gather : public test_case { @@ -4204,7 +4204,7 @@ index 8ce019b3..d9704672 100644 test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp -index 4dd00efd..fcf025d0 100644 +index 4dd00efdd..fcf025d07 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -2931,6 +2931,85 @@ static void test_template_output_peg_parsers(bool detailed_debug) { @@ -4295,7 +4295,7 @@ index 4dd00efd..fcf025d0 100644 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 00000000..94232b32 +index 000000000..94232b32d --- /dev/null +++ b/tests/test-flash-attn-bias.cpp @@ -0,0 +1,525 @@ @@ -4826,7 +4826,7 @@ index 00000000..94232b32 +} diff --git a/tests/test-flash-attn-generic-hash.cpp b/tests/test-flash-attn-generic-hash.cpp new file mode 100644 -index 00000000..4e3aca44 +index 000000000..4e3aca44f --- /dev/null +++ b/tests/test-flash-attn-generic-hash.cpp @@ -0,0 +1,142 @@ @@ -4973,7 +4973,7 @@ index 00000000..4e3aca44 + return 0; +} diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp -index 659ea9e9..c7efc95a 100644 +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) { @@ -4987,7 +4987,7 @@ index 659ea9e9..c7efc95a 100644 // FIXME some models are segfaulting with WebGPU: #ifdef GGML_USE_WEBGPU diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt -index 18a8288b..ae574bc9 100644 +index 18a8288ba..ae574bc98 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -38,6 +38,7 @@ add_library(mtmd @@ -4999,7 +4999,7 @@ index 18a8288b..ae574bc9 100644 models/kimik25.cpp models/nemotron-v2-vl.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h -index 09204113..539a7c66 100644 +index 092041138..539a7c661 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -284,6 +284,13 @@ @@ -5033,7 +5033,7 @@ index 09204113..539a7c66 100644 { PROJECTOR_TYPE_LDP, "ldp" }, { PROJECTOR_TYPE_LDPV2, "ldpv2"}, diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h -index 8dc87549..3fd92b77 100644 +index 8dc875497..3fd92b77f 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -33,7 +33,7 @@ enum resize_algo { @@ -5071,7 +5071,7 @@ index 8dc87549..3fd92b77 100644 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 04614b93..8fe2a5b2 100644 +index 04614b93b..8fe2a5b2d 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 @@ -5194,7 +5194,7 @@ index 04614b93..8fe2a5b2 100644 case PROJECTOR_TYPE_LDPV2: diff --git a/tools/mtmd/models/inkling.cpp b/tools/mtmd/models/inkling.cpp new file mode 100644 -index 00000000..c072f66a +index 000000000..c072f66a2 --- /dev/null +++ b/tools/mtmd/models/inkling.cpp @@ -0,0 +1,106 @@ @@ -5305,7 +5305,7 @@ index 00000000..c072f66a + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h -index caed438e..c671757f 100644 +index caed438ec..c671757f1 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -12,6 +12,17 @@ struct clip_graph_siglip : clip_graph { @@ -5327,7 +5327,7 @@ index caed438e..c671757f 100644 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 ed68951c..c7f0b9a1 100644 +index ed68951c0..c7f0b9a10 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -283,6 +283,7 @@ struct filter_params { @@ -5423,7 +5423,7 @@ index ed68951c..c7f0b9a1 100644 // mtmd_audio_preprocessor_whisper // diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h -index d8ec72b9..eae32a76 100644 +index d8ec72b9d..eae32a762 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 { @@ -5444,7 +5444,7 @@ index d8ec72b9..eae32a76 100644 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 36cd463b..01cb9f59 100644 +index 36cd463b2..01cb9f592 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -68,6 +68,9 @@ struct img_tool { @@ -5714,7 +5714,7 @@ index 36cd463b..01cb9f59 100644 // mtmd_image_preprocessor_dyn_size // diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h -index 115cba51..044d72f6 100644 +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 { @@ -5752,7 +5752,7 @@ index 115cba51..044d72f6 100644 // 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 6e61cf3e..119e9fc7 100644 +index 6e61cf3e5..119e9fc79 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -393,6 +393,12 @@ struct mtmd_context { diff --git a/third_party/llama.cpp/patches/0049-Add-staged-execution-support-for-Inkling.patch b/third_party/llama.cpp/patches/0051-Add-staged-execution-support-for-Inkling.patch similarity index 96% rename from third_party/llama.cpp/patches/0049-Add-staged-execution-support-for-Inkling.patch rename to third_party/llama.cpp/patches/0051-Add-staged-execution-support-for-Inkling.patch index daea88aaef..029727fc7f 100644 --- a/third_party/llama.cpp/patches/0049-Add-staged-execution-support-for-Inkling.patch +++ b/third_party/llama.cpp/patches/0051-Add-staged-execution-support-for-Inkling.patch @@ -1,7 +1,7 @@ -From f46720057c729128913a708eb6581be7375fc431 Mon Sep 17 00:00:00 2001 +From de940c16a6bf7999e9a9b164a3716d6d9f6c7468 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 17 Jul 2026 13:47:13 +1000 -Subject: [PATCH 49/60] Add staged execution support for Inkling +Subject: [PATCH 51/61] Add staged execution support for Inkling --- src/models/inkling.cpp | 30 +++++++++++++++++++++++------- @@ -9,7 +9,7 @@ Subject: [PATCH 49/60] Add staged execution support for Inkling 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/src/models/inkling.cpp b/src/models/inkling.cpp -index 6c9c1aef..3d5f8ab5 100644 +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 @@ -92,7 +92,7 @@ index 6c9c1aef..3d5f8ab5 100644 if (inp_out_ids) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); diff --git a/src/skippy.cpp b/src/skippy.cpp -index 21db2b77..347d6985 100644 +index 3965e0252..f6ccad78c 100644 --- a/src/skippy.cpp +++ b/src/skippy.cpp @@ -4187,6 +4187,7 @@ static enum skippy_status skippy_finish_model_open( diff --git a/third_party/llama.cpp/patches/0050-Advance-native-MTP-depth-per-draft-step.patch b/third_party/llama.cpp/patches/0052-Advance-native-MTP-depth-per-draft-step.patch similarity index 94% rename from third_party/llama.cpp/patches/0050-Advance-native-MTP-depth-per-draft-step.patch rename to third_party/llama.cpp/patches/0052-Advance-native-MTP-depth-per-draft-step.patch index 54b388b927..4efacbbb70 100644 --- a/third_party/llama.cpp/patches/0050-Advance-native-MTP-depth-per-draft-step.patch +++ b/third_party/llama.cpp/patches/0052-Advance-native-MTP-depth-per-draft-step.patch @@ -1,14 +1,14 @@ -From 40ed78300dc40dd12586b31017dca33980ea8de7 Mon Sep 17 00:00:00 2001 +From 1d78b6bfc85213cc7bcacfb60c691f7424c4bd64 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 17 Jul 2026 14:46:51 +1000 -Subject: [PATCH 50/60] Advance native MTP depth per draft step +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 347d6985..a918ce0f 100644 +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) { diff --git a/third_party/llama.cpp/patches/0051-Replay-multi-depth-MTP-caches-correctly.patch b/third_party/llama.cpp/patches/0053-Replay-multi-depth-MTP-caches-correctly.patch similarity index 97% rename from third_party/llama.cpp/patches/0051-Replay-multi-depth-MTP-caches-correctly.patch rename to third_party/llama.cpp/patches/0053-Replay-multi-depth-MTP-caches-correctly.patch index 3c231e4f8b..3515bb5b4a 100644 --- a/third_party/llama.cpp/patches/0051-Replay-multi-depth-MTP-caches-correctly.patch +++ b/third_party/llama.cpp/patches/0053-Replay-multi-depth-MTP-caches-correctly.patch @@ -1,7 +1,7 @@ -From 68ea61226111c9ab95d959d6a3d355c3dada7228 Mon Sep 17 00:00:00 2001 +From 8e80340a6f85a49e53320c79e3d353b7285b84db Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 17 Jul 2026 14:58:19 +1000 -Subject: [PATCH 51/60] Replay multi-depth MTP caches correctly +Subject: [PATCH 53/61] Replay multi-depth MTP caches correctly --- src/models/step35.cpp | 12 ++- @@ -9,7 +9,7 @@ Subject: [PATCH 51/60] Replay multi-depth MTP caches correctly 2 files changed, 189 insertions(+), 49 deletions(-) diff --git a/src/models/step35.cpp b/src/models/step35.cpp -index 4df4605b..2d546c20 100644 +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 @@ -34,7 +34,7 @@ index 4df4605b..2d546c20 100644 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 a918ce0f..5d221cda 100644 +index 1484f4f3b..52670efe8 100644 --- a/src/skippy.cpp +++ b/src/skippy.cpp @@ -1822,6 +1822,14 @@ struct skippy_mtp_depth_scope { @@ -345,7 +345,7 @@ index a918ce0f..5d221cda 100644 } if (out_mtp_draft != nullptr && token_count > 0) { -@@ -4413,7 +4549,7 @@ static enum skippy_status skippy_finish_model_open( +@@ -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); @@ -354,7 +354,7 @@ index a918ce0f..5d221cda 100644 } else { fprintf(stderr, "skippy: native MTP sidecar unavailable for this final stage; continuing without drafts\n"); } -@@ -4678,7 +4814,7 @@ enum skippy_status skippy_model_attach_mtp_draft_model( +@@ -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); diff --git a/third_party/llama.cpp/patches/0052-Add-Inkling-multi-depth-MTP-sidecars.patch b/third_party/llama.cpp/patches/0054-Add-Inkling-multi-depth-MTP-sidecars.patch similarity index 98% rename from third_party/llama.cpp/patches/0052-Add-Inkling-multi-depth-MTP-sidecars.patch rename to third_party/llama.cpp/patches/0054-Add-Inkling-multi-depth-MTP-sidecars.patch index 991717f418..d8f60501a6 100644 --- a/third_party/llama.cpp/patches/0052-Add-Inkling-multi-depth-MTP-sidecars.patch +++ b/third_party/llama.cpp/patches/0054-Add-Inkling-multi-depth-MTP-sidecars.patch @@ -1,7 +1,7 @@ -From 06bac0e35ea1d224100e57af5d5ad1cf04495557 Mon Sep 17 00:00:00 2001 +From a7a426fc1ca1c0fd28933bcd2bf49191e5041580 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 17 Jul 2026 15:09:39 +1000 -Subject: [PATCH 52/60] Add Inkling multi-depth MTP sidecars +Subject: [PATCH 54/61] Add Inkling multi-depth MTP sidecars --- conversion/inkling.py | 78 ++++++++++++++++- @@ -12,7 +12,7 @@ Subject: [PATCH 52/60] Add Inkling multi-depth MTP sidecars 5 files changed, 247 insertions(+), 54 deletions(-) diff --git a/conversion/inkling.py b/conversion/inkling.py -index 90ef65c3..07984290 100644 +index 90ef65c3b..079842903 100644 --- a/conversion/inkling.py +++ b/conversion/inkling.py @@ -1,5 +1,7 @@ @@ -152,7 +152,7 @@ index 90ef65c3..07984290 100644 @ModelBase.register("InklingForConditionalGeneration") class InklingMmprojModel(MmprojModel): diff --git a/src/llama-context.cpp b/src/llama-context.cpp -index c4e057ae..4b47903d 100644 +index c4e057aee..4b47903d3 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -158,6 +158,12 @@ llama_context::llama_context( @@ -169,10 +169,10 @@ index c4e057ae..4b47903d 100644 // 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 3621b047..be5cfb88 100644 +index f95591860..4ffe8df0f 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp -@@ -2115,10 +2115,20 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, +@@ -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; @@ -195,7 +195,7 @@ index 3621b047..be5cfb88 100644 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 3d5f8ab5..d435a1f5 100644 +index 3d5f8ab56..d435a1f5d 100644 --- a/src/models/inkling.cpp +++ b/src/models/inkling.cpp @@ -13,6 +13,9 @@ @@ -492,7 +492,7 @@ index 3d5f8ab5..d435a1f5 100644 res->t_embd = cur; diff --git a/src/skippy.cpp b/src/skippy.cpp -index 5d221cda..61f9c508 100644 +index 52670efe8..913c4d5ee 100644 --- a/src/skippy.cpp +++ b/src/skippy.cpp @@ -1823,7 +1823,7 @@ struct skippy_mtp_depth_scope { diff --git a/third_party/llama.cpp/patches/0053-Carry-multimodal-Inkling-embeddings-into-MTP.patch b/third_party/llama.cpp/patches/0055-Carry-multimodal-Inkling-embeddings-into-MTP.patch similarity index 97% rename from third_party/llama.cpp/patches/0053-Carry-multimodal-Inkling-embeddings-into-MTP.patch rename to third_party/llama.cpp/patches/0055-Carry-multimodal-Inkling-embeddings-into-MTP.patch index bfbdac9c28..842265636d 100644 --- a/third_party/llama.cpp/patches/0053-Carry-multimodal-Inkling-embeddings-into-MTP.patch +++ b/third_party/llama.cpp/patches/0055-Carry-multimodal-Inkling-embeddings-into-MTP.patch @@ -1,7 +1,7 @@ -From 5b4896aa14167a09aba5534433c85ca3f49cd089 Mon Sep 17 00:00:00 2001 +From 7412e7ae58b83de1864433828e4a786693e58b60 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 17 Jul 2026 16:04:48 +1000 -Subject: [PATCH 53/60] Carry multimodal Inkling embeddings into MTP +Subject: [PATCH 55/61] Carry multimodal Inkling embeddings into MTP --- include/skippy.h | 3 +- @@ -15,7 +15,7 @@ Subject: [PATCH 53/60] Carry multimodal Inkling embeddings into MTP 8 files changed, 218 insertions(+), 30 deletions(-) diff --git a/include/skippy.h b/include/skippy.h -index e9d31f55..cc2382d3 100644 +index e9d31f559..cc2382d37 100644 --- a/include/skippy.h +++ b/include/skippy.h @@ -43,7 +43,8 @@ enum skippy_activation_layout { @@ -29,7 +29,7 @@ index e9d31f55..cc2382d3 100644 #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 0145705b..11f06bf3 100644 +index 0145705b2..11f06bf3e 100644 --- a/include/skippy/common.h +++ b/include/skippy/common.h @@ -55,6 +55,7 @@ enum skippy_feature { @@ -41,7 +41,7 @@ index 0145705b..11f06bf3 100644 enum skippy_status { diff --git a/src/llama-context.cpp b/src/llama-context.cpp -index 4b47903d..9d528761 100644 +index 4b47903d3..9d528761b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4113,7 +4113,10 @@ int32_t llama_encode( @@ -57,7 +57,7 @@ index 4b47903d..9d528761 100644 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 42442879..20c121b2 100644 +index 42442879e..20c121b24 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -20,6 +20,8 @@ class llama_batch_allocr; @@ -70,7 +70,7 @@ index 42442879..20c121b2 100644 struct llama_memory_i; struct llama_memory_context_i; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp -index 36bf3219..27c5129b 100644 +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; @@ -134,7 +134,7 @@ index 36bf3219..27c5129b 100644 return res; diff --git a/src/llama-graph.h b/src/llama-graph.h -index e74a58a5..6e66f7fe 100644 +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 { @@ -161,7 +161,7 @@ index e74a58a5..6e66f7fe 100644 // 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 d435a1f5..c054f1d9 100644 +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 @@ -197,7 +197,7 @@ index d435a1f5..c054f1d9 100644 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 61f9c508..c76d3fbb 100644 +index 913c4d5ee..cc50cc501 100644 --- a/src/skippy.cpp +++ b/src/skippy.cpp @@ -149,8 +149,12 @@ struct skippy_session { @@ -531,7 +531,7 @@ index 61f9c508..c76d3fbb 100644 } void skippy_error_free(struct skippy_error * error) { -@@ -4939,6 +5074,12 @@ enum skippy_status skippy_session_begin_external_decode( +@@ -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; } @@ -544,7 +544,7 @@ index 61f9c508..c76d3fbb 100644 const skippy_runtime_config & config = session->stage_model->config; if (config.filter_tensors_on_load) { -@@ -4950,6 +5091,7 @@ enum skippy_status skippy_session_begin_external_decode( +@@ -4951,6 +5092,7 @@ enum skippy_status skippy_session_begin_external_decode( filter.include_output = config.include_output; skippy_graph_set_filter(filter); } @@ -552,7 +552,7 @@ index 61f9c508..c76d3fbb 100644 return skippy_success(out_error); } -@@ -4961,6 +5103,9 @@ enum skippy_status skippy_session_end_external_decode( +@@ -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; } diff --git a/third_party/llama.cpp/patches/0055-Restore-recurrent-verify-window-state.patch b/third_party/llama.cpp/patches/0055-Restore-recurrent-verify-window-state.patch deleted file mode 100644 index 51d20c8d53..0000000000 --- a/third_party/llama.cpp/patches/0055-Restore-recurrent-verify-window-state.patch +++ /dev/null @@ -1,317 +0,0 @@ -From 33aa6b672f3d6ed994ed819636ed08fcab18cfa7 Mon Sep 17 00:00:00 2001 -From: Mesh-LLM CI -Date: Fri, 24 Jul 2026 03:57:28 +1000 -Subject: [PATCH 55/60] Restore recurrent verify-window state - -Assisted-by: codex ---- - src/skippy.cpp | 241 ++++++++++++++++++++++++++++++++++++++++++++----- - 1 file changed, 218 insertions(+), 23 deletions(-) - -diff --git a/src/skippy.cpp b/src/skippy.cpp -index c76d3fbb..08bb41e9 100644 ---- a/src/skippy.cpp -+++ b/src/skippy.cpp -@@ -124,6 +124,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 +162,7 @@ struct skippy_session { - llama_token mtp_pending_draft_token = -1; - bool external_mtp_embeddings_available = false; - std::vector external_mtp_embeddings; -+ skippy_verify_checkpoint verify_checkpoint; - }; - - static thread_local skippy_session * g_skippy_external_decode_session = nullptr; -@@ -1590,7 +1602,63 @@ 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) { -+ session->verify_checkpoint = {}; -+ 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(session->ctx, session->seq_id); -+ 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( -+ session->ctx, -+ checkpoint.state.data(), -+ checkpoint.state.size(), -+ session->seq_id); -+ 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); -+ } -+ session->verify_checkpoint = 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, - const void * input_payload, -@@ -2859,6 +2927,7 @@ static void skippy_mtp_clear_session_state(skippy_session * session) { - if (session == nullptr) { - return; - } -+ session->verify_checkpoint = {}; - session->mtp_next_pos = session->n_past; - session->mtp_prefix_valid = session->n_past == 0; - session->mtp_has_pending_h = false; -@@ -4353,6 +4422,97 @@ 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; -+ const size_t prefix_bytes = skippy_activation_payload_bytes(session, token_count, flags); -+ std::vector payload(prefix_bytes); -+ if ((flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0) { -+ std::memcpy(payload.data(), checkpoint.input_payload.data(), prefix_bytes); -+ return payload; -+ } -+ -+ 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); -+ 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 ((flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0) { -+ std::memcpy( -+ payload.data() + prefix_offset, -+ checkpoint.input_payload.data() + source_offset, -+ prefix_hidden_bytes); -+ source_offset += source_hidden_bytes; -+ prefix_offset += prefix_hidden_bytes; -+ } -+ if ((flags & SKIPPY_ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0) { -+ std::memcpy( -+ payload.data() + prefix_offset, -+ checkpoint.input_payload.data() + source_offset, -+ prefix_hidden_bytes); -+ } -+ return payload; -+} -+ -+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 = session->verify_checkpoint; -+ const uint64_t token_start = static_cast(checkpoint.token_start); -+ const uint64_t token_end = token_start + checkpoint.token_count; -+ if (!checkpoint.valid || token_count < token_start || token_count >= token_end) { -+ return skippy_success(out_error); -+ } -+ -+ const size_t read = llama_state_seq_set_data( -+ session->ctx, -+ checkpoint.state.data(), -+ checkpoint.state.size(), -+ session->seq_id); -+ if (read != checkpoint.state.size()) { -+ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to restore verify-window sequence state"); -+ return SKIPPY_STATUS_RUNTIME_ERROR; -+ } -+ -+ session->n_past = checkpoint.token_start; -+ const size_t accepted_count = static_cast(token_count - token_start); -+ 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(); -+ const llama_token * token_ids = checkpoint.token_ids.empty() ? nullptr : checkpoint.token_ids.data(); -+ status = skippy_decode_activation_frame( -+ session, -+ &input_desc, -+ input_payload.data(), -+ token_ids, -+ nullptr, -+ 0, -+ accepted_count, -+ false, -+ out_error); -+ } else if (accepted_count > 0) { -+ status = skippy_decode_tokens( -+ session, -+ checkpoint.token_ids.data(), -+ accepted_count, -+ false, -+ out_error); -+ } -+ if (status == SKIPPY_STATUS_OK) { -+ *out_restored = true; -+ } -+ return status; -+} -+ - extern "C" { - - struct skippy_abi_version skippy_abi_version(void) { -@@ -5551,7 +5711,12 @@ 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) { -@@ -6260,7 +6425,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( -@@ -6826,29 +7003,47 @@ enum skippy_status skippy_trim_session( - 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; -+ 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) { -+ 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; -+ } -+ } else if (dynamic_cast(memory) != nullptr) { -+ // Recurrent state cannot be rolled back by suffix removal; only a -+ // retained verify checkpoint can restore it. -+ if (static_cast(session->n_past) != token_count) { -+ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, -+ "recurrent trim requires a retained verify checkpoint"); -+ return SKIPPY_STATUS_UNSUPPORTED; -+ } -+ } else { -+ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "runtime memory type is not supported for trim"); -+ return SKIPPY_STATUS_UNSUPPORTED; - } -- } 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; - } - session->n_past = static_cast(token_count); - if (session->token_history.size() > token_count) { --- -2.50.1 (Apple Git-155) - diff --git a/third_party/llama.cpp/patches/0054-Size-recurrent-memory-for-appended-MTP-layers.patch b/third_party/llama.cpp/patches/0056-Size-recurrent-memory-for-appended-MTP-layers.patch similarity index 91% rename from third_party/llama.cpp/patches/0054-Size-recurrent-memory-for-appended-MTP-layers.patch rename to third_party/llama.cpp/patches/0056-Size-recurrent-memory-for-appended-MTP-layers.patch index f8ef7df005..afbc226acd 100644 --- a/third_party/llama.cpp/patches/0054-Size-recurrent-memory-for-appended-MTP-layers.patch +++ b/third_party/llama.cpp/patches/0056-Size-recurrent-memory-for-appended-MTP-layers.patch @@ -1,14 +1,14 @@ -From d8d24d3fb5c071baec273bfd722532414952535d Mon Sep 17 00:00:00 2001 +From 326edb5713500457c0a82fbb450d84fcf3549cb0 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Fri, 17 Jul 2026 20:20:45 +1000 -Subject: [PATCH 54/60] Size recurrent memory for appended MTP layers +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 ef82eb97..4f1f1df8 100644 +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( diff --git a/third_party/llama.cpp/patches/0056-skippy-preserve-verify-checkpoint-until-trim.patch b/third_party/llama.cpp/patches/0056-skippy-preserve-verify-checkpoint-until-trim.patch deleted file mode 100644 index 6a00fb2163..0000000000 --- a/third_party/llama.cpp/patches/0056-skippy-preserve-verify-checkpoint-until-trim.patch +++ /dev/null @@ -1,84 +0,0 @@ -From 8b19b3407034057ef91daf45c49633cce0618ce8 Mon Sep 17 00:00:00 2001 -From: Mesh-LLM CI -Date: Fri, 24 Jul 2026 14:08:34 +1000 -Subject: [PATCH 56/60] skippy: preserve verify checkpoint until trim - -Assisted-by: goose ---- - src/skippy.cpp | 11 +++++++++-- - 1 file changed, 9 insertions(+), 2 deletions(-) - -diff --git a/src/skippy.cpp b/src/skippy.cpp -index 08bb41e9..fd1fd472 100644 ---- a/src/skippy.cpp -+++ b/src/skippy.cpp -@@ -2927,7 +2927,6 @@ static void skippy_mtp_clear_session_state(skippy_session * session) { - if (session == nullptr) { - return; - } -- session->verify_checkpoint = {}; - session->mtp_next_pos = session->n_past; - session->mtp_prefix_valid = session->n_past == 0; - session->mtp_has_pending_h = false; -@@ -4465,7 +4464,9 @@ static enum skippy_status skippy_restore_verify_prefix( - const skippy_verify_checkpoint & checkpoint = session->verify_checkpoint; - const uint64_t token_start = static_cast(checkpoint.token_start); - const uint64_t token_end = token_start + checkpoint.token_count; -- if (!checkpoint.valid || token_count < token_start || token_count >= token_end) { -+ const uint64_t session_token_count = static_cast(std::max(session->n_past, 0)); -+ if (!checkpoint.valid || session_token_count != token_end || -+ token_count < token_start || token_count >= token_end) { - return skippy_success(out_error); - } - -@@ -5291,6 +5292,7 @@ enum skippy_status skippy_session_set_position( - session->signal_history.resize(static_cast(n_past)); - } - skippy_mtp_clear_session_state(session); -+ session->verify_checkpoint = {}; - return skippy_success(out_error); - } - -@@ -5384,6 +5386,7 @@ enum skippy_status skippy_session_reset( - session->signal_history.clear(); - skippy_clear_chat_sampling(session); - skippy_mtp_clear_session_state(session); -+ session->verify_checkpoint = {}; - session->ctx->synchronize(); - return skippy_success(out_error); - } -@@ -6556,6 +6559,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_checkpoint = {}; - } - - enum skippy_status skippy_export_state( -@@ -6974,6 +6978,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_checkpoint = {}; - session->ctx->synchronize(); - - return skippy_success(out_error); -@@ -7054,6 +7059,7 @@ enum skippy_status skippy_trim_session( - session->signal_history.resize(static_cast(token_count)); - } - skippy_mtp_clear_session_state(session); -+ session->verify_checkpoint = {}; - session->ctx->synchronize(); - - return skippy_success(out_error); -@@ -7206,6 +7212,7 @@ enum skippy_status skippy_session_restore_prefix( - } - session->signal_history.clear(); - skippy_mtp_clear_session_state(session); -+ session->verify_checkpoint = {}; - session->ctx->synchronize(); - return skippy_success(out_error); - } --- -2.50.1 (Apple Git-155) - diff --git a/third_party/llama.cpp/patches/0059-Complete-Inkling-and-GLM-metadata-integration.patch b/third_party/llama.cpp/patches/0057-Complete-Inkling-and-GLM-metadata-integration.patch similarity index 95% rename from third_party/llama.cpp/patches/0059-Complete-Inkling-and-GLM-metadata-integration.patch rename to third_party/llama.cpp/patches/0057-Complete-Inkling-and-GLM-metadata-integration.patch index ff5aa4596c..95ee485200 100644 --- a/third_party/llama.cpp/patches/0059-Complete-Inkling-and-GLM-metadata-integration.patch +++ b/third_party/llama.cpp/patches/0057-Complete-Inkling-and-GLM-metadata-integration.patch @@ -1,7 +1,7 @@ -From 8bd27d799ab7231e2b276645dbc6286a0fdd705e Mon Sep 17 00:00:00 2001 +From b01316df2de211495a1754050a1643509ae03b7a Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Mon, 27 Jul 2026 12:01:26 +1000 -Subject: [PATCH 59/60] Complete Inkling and GLM metadata integration +Subject: [PATCH 57/61] Complete Inkling and GLM metadata integration --- common/stage-chat.cpp | 6 ++++++ @@ -12,7 +12,7 @@ Subject: [PATCH 59/60] Complete Inkling and GLM metadata integration 5 files changed, 50 insertions(+), 7 deletions(-) diff --git a/common/stage-chat.cpp b/common/stage-chat.cpp -index 3298ed66..16ce255a 100644 +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( @@ -36,7 +36,7 @@ index 3298ed66..16ce255a 100644 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 a489e1ad..836df068 100644 +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 = { @@ -49,7 +49,7 @@ index a489e1ad..836df068 100644 { 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 74446fda..0033c761 100644 +index 74446fdaf..0033c7612 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -263,6 +263,8 @@ enum llm_kv { @@ -62,10 +62,10 @@ index 74446fda..0033c761 100644 LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, LLM_KV_ATTENTION_INDEXER_TYPES, diff --git a/src/llama-model.cpp b/src/llama-model.cpp -index 527b52b8..8c0589da 100644 +index 4ffe8df0f..70c105e6a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp -@@ -1138,6 +1138,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { +@@ -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); @@ -74,7 +74,7 @@ index 527b52b8..8c0589da 100644 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 a88aae2c..57c45772 100644 +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) { diff --git a/third_party/llama.cpp/patches/0057-skippy-retain-pipelined-verify-checkpoints.patch b/third_party/llama.cpp/patches/0057-skippy-retain-pipelined-verify-checkpoints.patch deleted file mode 100644 index 373f25b795..0000000000 --- a/third_party/llama.cpp/patches/0057-skippy-retain-pipelined-verify-checkpoints.patch +++ /dev/null @@ -1,170 +0,0 @@ -From cbb9ce751bc15c3895e262f5b73812368526b42c Mon Sep 17 00:00:00 2001 -From: Mesh-LLM CI -Date: Fri, 24 Jul 2026 14:54:57 +1000 -Subject: [PATCH 57/60] skippy: retain pipelined verify checkpoints - -Assisted-by: goose ---- - src/skippy.cpp | 61 ++++++++++++++++++++++++++++++-------------------- - 1 file changed, 37 insertions(+), 24 deletions(-) - -diff --git a/src/skippy.cpp b/src/skippy.cpp -index fd1fd472..886efbcc 100644 ---- a/src/skippy.cpp -+++ b/src/skippy.cpp -@@ -162,7 +162,7 @@ struct skippy_session { - llama_token mtp_pending_draft_token = -1; - bool external_mtp_embeddings_available = false; - std::vector external_mtp_embeddings; -- skippy_verify_checkpoint verify_checkpoint; -+ std::vector verify_checkpoints; - }; - - static thread_local skippy_session * g_skippy_external_decode_session = nullptr; -@@ -1616,7 +1616,6 @@ static enum skippy_status skippy_checkpoint_verify_window( - const skippy_activation_desc * input_desc, - const void * input_payload, - skippy_error ** out_error) { -- session->verify_checkpoint = {}; - llama_memory_t memory = session->ctx->get_memory(); - if (!skippy_memory_needs_verify_checkpoint(memory)) { - return skippy_success(out_error); -@@ -1653,7 +1652,16 @@ static enum skippy_status skippy_checkpoint_verify_window( - const uint8_t * bytes = static_cast(input_payload); - checkpoint.input_payload.assign(bytes, bytes + input_desc->payload_bytes); - } -- session->verify_checkpoint = std::move(checkpoint); -+ // Fail closed: evicting the oldest checkpoint would silently strand a -+ // window that can still be restored. The ceiling is only a growth backstop -+ // (the scheduler bounds in-flight windows to the configured depth). -+ 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); - } - -@@ -4461,35 +4469,40 @@ static enum skippy_status skippy_restore_verify_prefix( - bool * out_restored, - skippy_error ** out_error) { - *out_restored = false; -- const skippy_verify_checkpoint & checkpoint = session->verify_checkpoint; -- const uint64_t token_start = static_cast(checkpoint.token_start); -- const uint64_t token_end = token_start + checkpoint.token_count; -- const uint64_t session_token_count = static_cast(std::max(session->n_past, 0)); -- if (!checkpoint.valid || session_token_count != token_end || -- token_count < token_start || token_count >= token_end) { -+ 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( - session->ctx, -- checkpoint.state.data(), -- checkpoint.state.size(), -+ checkpoint->state.data(), -+ checkpoint->state.size(), - session->seq_id); -- if (read != checkpoint.state.size()) { -+ if (read != checkpoint->state.size()) { - skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to restore verify-window sequence state"); - return SKIPPY_STATUS_RUNTIME_ERROR; - } - -- session->n_past = checkpoint.token_start; -+ 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); - enum skippy_status status = SKIPPY_STATUS_OK; -- if (accepted_count > 0 && checkpoint.has_activation_input) { -- skippy_activation_desc input_desc = checkpoint.input_desc; -+ 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); -+ session, *checkpoint, accepted_count); - input_desc.token_count = static_cast(accepted_count); - input_desc.payload_bytes = input_payload.size(); -- const llama_token * token_ids = checkpoint.token_ids.empty() ? nullptr : checkpoint.token_ids.data(); -+ const llama_token * token_ids = checkpoint->token_ids.empty() ? nullptr : checkpoint->token_ids.data(); - status = skippy_decode_activation_frame( - session, - &input_desc, -@@ -4503,7 +4516,7 @@ static enum skippy_status skippy_restore_verify_prefix( - } else if (accepted_count > 0) { - status = skippy_decode_tokens( - session, -- checkpoint.token_ids.data(), -+ checkpoint->token_ids.data(), - accepted_count, - false, - out_error); -@@ -5292,7 +5305,7 @@ enum skippy_status skippy_session_set_position( - session->signal_history.resize(static_cast(n_past)); - } - skippy_mtp_clear_session_state(session); -- session->verify_checkpoint = {}; -+ session->verify_checkpoints.clear(); - return skippy_success(out_error); - } - -@@ -5386,7 +5399,7 @@ enum skippy_status skippy_session_reset( - session->signal_history.clear(); - skippy_clear_chat_sampling(session); - skippy_mtp_clear_session_state(session); -- session->verify_checkpoint = {}; -+ session->verify_checkpoints.clear(); - session->ctx->synchronize(); - return skippy_success(out_error); - } -@@ -6559,7 +6572,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_checkpoint = {}; -+ session->verify_checkpoints.clear(); - } - - enum skippy_status skippy_export_state( -@@ -6978,7 +6991,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_checkpoint = {}; -+ session->verify_checkpoints.clear(); - session->ctx->synchronize(); - - return skippy_success(out_error); -@@ -7059,7 +7072,7 @@ enum skippy_status skippy_trim_session( - session->signal_history.resize(static_cast(token_count)); - } - skippy_mtp_clear_session_state(session); -- session->verify_checkpoint = {}; -+ session->verify_checkpoints.clear(); - session->ctx->synchronize(); - - return skippy_success(out_error); -@@ -7212,7 +7225,7 @@ enum skippy_status skippy_session_restore_prefix( - } - session->signal_history.clear(); - skippy_mtp_clear_session_state(session); -- session->verify_checkpoint = {}; -+ 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-balance-filtered-stages-across-devices.patch b/third_party/llama.cpp/patches/0058-skippy-balance-filtered-stages-across-devices.patch similarity index 92% rename from third_party/llama.cpp/patches/0060-skippy-balance-filtered-stages-across-devices.patch rename to third_party/llama.cpp/patches/0058-skippy-balance-filtered-stages-across-devices.patch index 1e0286766f..45de955287 100644 --- a/third_party/llama.cpp/patches/0060-skippy-balance-filtered-stages-across-devices.patch +++ b/third_party/llama.cpp/patches/0058-skippy-balance-filtered-stages-across-devices.patch @@ -1,7 +1,7 @@ -From 48a67d12017d1c78d563007f5658150b605d8556 Mon Sep 17 00:00:00 2001 +From 77251842c0f5c2f288ac0a2f9039a41478ca6370 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Wed, 29 Jul 2026 00:43:18 +1000 -Subject: [PATCH 60/60] skippy: balance filtered stages across devices +Subject: [PATCH 58/61] skippy: balance filtered stages across devices --- src/llama-model-loader.cpp | 4 ++++ @@ -10,7 +10,7 @@ Subject: [PATCH 60/60] skippy: balance filtered stages across devices 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp -index ec530d62..1ac8a9f9 100644 +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() { @@ -25,7 +25,7 @@ index ec530d62..1ac8a9f9 100644 return g_skippy_last_tensor_filtered; } diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h -index da58661f..988fcf94 100644 +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 { @@ -37,10 +37,10 @@ index da58661f..988fcf94 100644 enum llama_fver { diff --git a/src/llama-model.cpp b/src/llama-model.cpp -index 8c0589da..dbeefd8a 100644 +index 70c105e6a..0e45dc460 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp -@@ -1317,13 +1317,29 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { +@@ -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); diff --git a/third_party/llama.cpp/patches/0058-skippy-filter-Inkling-stage-memory-by-layer-range.patch b/third_party/llama.cpp/patches/0058-skippy-filter-Inkling-stage-memory-by-layer-range.patch deleted file mode 100644 index 2e4f57ae82..0000000000 --- a/third_party/llama.cpp/patches/0058-skippy-filter-Inkling-stage-memory-by-layer-range.patch +++ /dev/null @@ -1,72 +0,0 @@ -From 7be339fb9da027c7ecd15a7d5eb9e2e56313d312 Mon Sep 17 00:00:00 2001 -From: Mesh-LLM CI -Date: Fri, 24 Jul 2026 19:02:35 +1000 -Subject: [PATCH 58/60] skippy: filter Inkling stage memory by layer range - -Assisted-by: goose ---- - src/llama-model.cpp | 11 +++++++++-- - src/skippy.cpp | 4 ++-- - 2 files changed, 11 insertions(+), 4 deletions(-) - -diff --git a/src/llama-model.cpp b/src/llama-model.cpp -index be5cfb88..527b52b8 100644 ---- a/src/llama-model.cpp -+++ b/src/llama-model.cpp -@@ -6,6 +6,7 @@ - #include "llama-impl.h" - #include "llama-mmap.h" - #include "llama-cparams.h" -+#include "llama-graph.h" - #include "llama-model-loader.h" - - #include "llama-kv-cache.h" -@@ -2126,8 +2127,14 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, - 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(); }; -+ const skippy_graph_filter stage_filter = skippy_graph_get_filter(); -+ filter_attn = [&, stage_filter](uint32_t il) { -+ return il < hparams.n_layer() && -+ (!stage_filter.enabled || -+ (il >= static_cast(stage_filter.layer_start) && -+ il < static_cast(stage_filter.layer_end))); -+ }; -+ filter_recr = filter_attn; - } - } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { - filter_attn = [&](uint32_t il) { -diff --git a/src/skippy.cpp b/src/skippy.cpp -index 886efbcc..46f8fad6 100644 ---- a/src/skippy.cpp -+++ b/src/skippy.cpp -@@ -4512,6 +4512,7 @@ static enum skippy_status skippy_restore_verify_prefix( - 0, - accepted_count, - false, -+ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, - out_error); - } else if (accepted_count > 0) { - status = skippy_decode_tokens( -@@ -4519,6 +4520,7 @@ static enum skippy_status skippy_restore_verify_prefix( - checkpoint->token_ids.data(), - accepted_count, - false, -+ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, - out_error); - } - if (status == SKIPPY_STATUS_OK) { -@@ -7060,8 +7062,6 @@ enum skippy_status skippy_trim_session( - 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) { --- -2.50.1 (Apple Git-155) - diff --git a/third_party/llama.cpp/patches/0048-Preserve-hybrid-verify-window-state-across-trim.patch b/third_party/llama.cpp/patches/0059-Preserve-hybrid-verify-window-state-across-trim.patch similarity index 93% rename from third_party/llama.cpp/patches/0048-Preserve-hybrid-verify-window-state-across-trim.patch rename to third_party/llama.cpp/patches/0059-Preserve-hybrid-verify-window-state-across-trim.patch index ff255be73d..351b1ce7e4 100644 --- a/third_party/llama.cpp/patches/0048-Preserve-hybrid-verify-window-state-across-trim.patch +++ b/third_party/llama.cpp/patches/0059-Preserve-hybrid-verify-window-state-across-trim.patch @@ -1,7 +1,7 @@ -From 3de16ed4e7b4dc99758193736d2b4ccc7b311602 Mon Sep 17 00:00:00 2001 +From 57280d642a5806d9464577bfbbc276d5afcab79b Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Tue, 28 Jul 2026 19:12:14 +1000 -Subject: [PATCH 48/49] Preserve hybrid verify-window state across trim +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. @@ -13,7 +13,7 @@ Assisted-by: codex 1 file changed, 332 insertions(+), 28 deletions(-) diff --git a/src/skippy.cpp b/src/skippy.cpp -index 21db2b77d..f42731cb6 100644 +index cc50cc501..9da45701e 100644 --- a/src/skippy.cpp +++ b/src/skippy.cpp @@ -9,6 +9,7 @@ @@ -42,16 +42,16 @@ index 21db2b77d..f42731cb6 100644 struct skippy_session { skippy_model * stage_model = nullptr; llama_context * ctx = nullptr; -@@ -149,6 +161,8 @@ struct skippy_session { - bool mtp_has_pending_draft = false; - llama_pos mtp_pending_draft_pos = 0; +@@ -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; }; - struct skippy_tensor_meta { -@@ -246,7 +260,15 @@ static enum skippy_status skippy_success(skippy_error ** out_error) { + 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; } @@ -68,7 +68,7 @@ index 21db2b77d..f42731cb6 100644 static int32_t skippy_stage_layer_count(const llama_model * model) { if (model == nullptr) { return 0; -@@ -1566,6 +1588,76 @@ static bool skippy_has_activation_payload(const skippy_activation_desc * desc, c +@@ -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; } @@ -145,7 +145,7 @@ index 21db2b77d..f42731cb6 100644 static enum skippy_status skippy_validate_frame_input_sequences( skippy_session * session, const skippy_activation_desc * input_desc, -@@ -1763,7 +1855,10 @@ static enum skippy_status skippy_decode_batch( +@@ -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; } @@ -157,7 +157,7 @@ index 21db2b77d..f42731cb6 100644 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) { -@@ -4035,6 +4130,169 @@ static enum skippy_status skippy_verify_activation_frame( +@@ -4353,6 +4448,169 @@ static enum skippy_status skippy_verify_activation_frame( return status; } @@ -327,7 +327,7 @@ index 21db2b77d..f42731cb6 100644 extern "C" { struct skippy_abi_version skippy_abi_version(void) { -@@ -4801,6 +5059,7 @@ enum skippy_status skippy_session_set_position( +@@ -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); @@ -335,7 +335,7 @@ index 21db2b77d..f42731cb6 100644 return skippy_success(out_error); } -@@ -4894,6 +5153,8 @@ enum skippy_status skippy_session_reset( +@@ -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); @@ -344,7 +344,7 @@ index 21db2b77d..f42731cb6 100644 session->ctx->synchronize(); return skippy_success(out_error); } -@@ -5221,7 +5482,17 @@ enum skippy_status skippy_verify_tokens( +@@ -5552,7 +5813,17 @@ enum skippy_status skippy_verify_tokens( return SKIPPY_STATUS_INVALID_ARGUMENT; } @@ -363,7 +363,7 @@ index 21db2b77d..f42731cb6 100644 if (status == SKIPPY_STATUS_OK) { const int32_t n_tokens = static_cast(token_count); for (int32_t i = 0; i < n_tokens; ++i) { -@@ -5282,7 +5553,8 @@ static enum skippy_status skippy_prefill_chunk_frame_impl( +@@ -5613,7 +5884,8 @@ static enum skippy_status skippy_prefill_chunk_frame_impl( return status; } @@ -373,7 +373,7 @@ index 21db2b77d..f42731cb6 100644 status = skippy_decode_activation_frame( session, input_desc, -@@ -5930,7 +6202,19 @@ enum skippy_status skippy_verify_tokens_frame_sampled( +@@ -6261,7 +6533,19 @@ enum skippy_status skippy_verify_tokens_frame_sampled( } } @@ -394,7 +394,7 @@ index 21db2b77d..f42731cb6 100644 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( -@@ -6049,6 +6333,7 @@ static void skippy_update_session_state_after_import( +@@ -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); @@ -402,7 +402,7 @@ index 21db2b77d..f42731cb6 100644 } enum skippy_status skippy_export_state( -@@ -6467,6 +6752,7 @@ enum skippy_status skippy_import_kv_page( +@@ -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); @@ -410,7 +410,7 @@ index 21db2b77d..f42731cb6 100644 session->ctx->synchronize(); return skippy_success(out_error); -@@ -6495,30 +6781,46 @@ enum skippy_status skippy_trim_session( +@@ -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; } @@ -480,7 +480,7 @@ index 21db2b77d..f42731cb6 100644 } session->n_past = static_cast(token_count); if (session->token_history.size() > token_count) { -@@ -6529,6 +6831,7 @@ enum skippy_status skippy_trim_session( +@@ -6860,6 +7162,7 @@ enum skippy_status skippy_trim_session( session->signal_history.resize(static_cast(token_count)); } skippy_mtp_clear_session_state(session); @@ -488,7 +488,7 @@ index 21db2b77d..f42731cb6 100644 session->ctx->synchronize(); return skippy_success(out_error); -@@ -6681,6 +6984,7 @@ enum skippy_status skippy_session_restore_prefix( +@@ -7012,6 +7315,7 @@ enum skippy_status skippy_session_restore_prefix( } session->signal_history.clear(); skippy_mtp_clear_session_state(session); @@ -497,5 +497,5 @@ index 21db2b77d..f42731cb6 100644 return skippy_success(out_error); } -- -2.54.0 +2.50.1 (Apple Git-155) diff --git a/third_party/llama.cpp/patches/0049-skippy-retire-accepted-verify-checkpoints-exactly.patch b/third_party/llama.cpp/patches/0060-skippy-retire-accepted-verify-checkpoints-exactly.patch similarity index 92% rename from third_party/llama.cpp/patches/0049-skippy-retire-accepted-verify-checkpoints-exactly.patch rename to third_party/llama.cpp/patches/0060-skippy-retire-accepted-verify-checkpoints-exactly.patch index febe2f5f19..be34a36621 100644 --- a/third_party/llama.cpp/patches/0049-skippy-retire-accepted-verify-checkpoints-exactly.patch +++ b/third_party/llama.cpp/patches/0060-skippy-retire-accepted-verify-checkpoints-exactly.patch @@ -1,7 +1,7 @@ -From 1f12cf3984b66dde5c5340d1d54614856fb123af Mon Sep 17 00:00:00 2001 +From 60f5917b3718cc74b2298571b4fac924e9ecf081 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI Date: Wed, 29 Jul 2026 11:12:07 +1000 -Subject: [PATCH 49/49] skippy: retire accepted verify checkpoints exactly +Subject: [PATCH 60/61] skippy: retire accepted verify checkpoints exactly --- include/skippy.h | 6 +++++ @@ -15,10 +15,10 @@ Subject: [PATCH 49/49] skippy: retire accepted verify checkpoints exactly create mode 100644 tests/test-skippy-verify-checkpoint-retirement.cpp diff --git a/include/skippy.h b/include/skippy.h -index e9d31f559..44afc6b74 100644 +index cc2382d37..f60e7c2ee 100644 --- a/include/skippy.h +++ b/include/skippy.h -@@ -524,6 +524,12 @@ LLAMA_API enum skippy_status skippy_trim_session( +@@ -525,6 +525,12 @@ LLAMA_API enum skippy_status skippy_trim_session( uint64_t token_count, struct skippy_error ** out_error); @@ -32,7 +32,7 @@ index e9d31f559..44afc6b74 100644 struct skippy_session * session, int32_t layer_start, diff --git a/include/skippy/common.h b/include/skippy/common.h -index 0145705b2..e89410d03 100644 +index 11f06bf3e..d1e45b44d 100644 --- a/include/skippy/common.h +++ b/include/skippy/common.h @@ -26,7 +26,7 @@ extern "C" { @@ -75,7 +75,7 @@ index 000000000..8793000c1 + return true; +} diff --git a/src/skippy.cpp b/src/skippy.cpp -index f42731cb6..a22f0dde6 100644 +index 9da45701e..cb20422f9 100644 --- a/src/skippy.cpp +++ b/src/skippy.cpp @@ -1,4 +1,5 @@ @@ -84,7 +84,7 @@ index f42731cb6..a22f0dde6 100644 #include "skippy/devices.h" #include "skippy-signals.h" -@@ -6758,6 +6759,26 @@ enum skippy_status skippy_import_kv_page( +@@ -7089,6 +7090,26 @@ enum skippy_status skippy_import_kv_page( return skippy_success(out_error); } @@ -112,7 +112,7 @@ index f42731cb6..a22f0dde6 100644 struct skippy_session * session, uint64_t token_count, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt -index 158ab28df..846d4d954 100644 +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) @@ -156,5 +156,5 @@ index 000000000..d09acccb4 + return 0; +} -- -2.54.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 0000000000..ec4cfcedd3 --- /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) + From 8ebda853712f121009d42f036b2f85e837131e58 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:50:55 +1000 Subject: [PATCH 14/26] fix(skippy): reconcile combined family recovery state --- .../src/inference/skippy/family_policy.rs | 4 +- .../binary_messaging/connection.rs | 59 ------------------- crates/skippy-topology/src/lib.rs | 5 +- 3 files changed, 4 insertions(+), 64 deletions(-) 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 750114d8e1..fd822bce81 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 @@ -84,8 +84,7 @@ impl FamilyPolicy { // 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); + 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 @@ -877,6 +876,7 @@ mod tests { 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, 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 20135d3834..9caf509de0 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -75,36 +75,6 @@ use std::time::Instant; static BINARY_SESSION_COUNTER: AtomicU64 = AtomicU64::new(1); -/// Tracks runtime session keys created by one binary stage connection so -/// that a mid-request failure (connection error, protocol violation, -/// runtime error) releases the execution lanes those sessions hold. -/// -/// Without this, a generation that errors before the graceful `Stop` -/// message leaks its lane session in [`RuntimeState`]: each retried -/// request uses a fresh session id, so leaked lanes accumulate until -/// every admission fails with "all execution lanes are busy" and only a -/// process restart recovers. -#[derive(Default)] -pub(super) struct ConnectionSessionTracker { - active: std::collections::BTreeSet, -} - -impl ConnectionSessionTracker { - pub(super) fn touch(&mut self, session_key: &str) { - if !self.active.contains(session_key) { - self.active.insert(session_key.to_string()); - } - } - - pub(super) fn stopped(&mut self, session_key: &str) { - self.active.remove(session_key); - } - - pub(super) fn drain(&mut self) -> Vec { - std::mem::take(&mut self.active).into_iter().collect() - } -} - #[allow(clippy::too_many_arguments)] pub(super) fn handle_binary_connection( config: &StageConfig, @@ -1403,32 +1373,3 @@ fn handle_binary_connection_messages( } } } - -#[cfg(test)] -mod tests { - use super::ConnectionSessionTracker; - - #[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"); - - // Simulate a mid-request stage error: session-a errored before its - // graceful Stop, session-b completed normally. - tracker.stopped("session-b"); - - assert_eq!(tracker.drain(), vec!["session-a".to_string()]); - // Idempotent: a second drain reclaims nothing. - 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()); - } -} diff --git a/crates/skippy-topology/src/lib.rs b/crates/skippy-topology/src/lib.rs index ae2b5b1788..d06b39c2b5 100644 --- a/crates/skippy-topology/src/lib.rs +++ b/crates/skippy-topology/src/lib.rs @@ -15,9 +15,8 @@ pub use family_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, + 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, From a47cb0ccbb69e91aad7ccf09bb50bc21069304a5 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 31 Jul 2026 00:25:51 +1000 Subject: [PATCH 15/26] Harden combined Skippy family support --- Cargo.lock | 1 + .../inference/skippy/resolver/resolution.rs | 12 + .../src/inference/skippy/resolver/tests.rs | 27 + .../src/runtime/local_split/loading.rs | 6 +- .../src/runtime/local_split/tests.rs | 4 +- crates/skippy-ffi/src/lib.rs | 21 +- crates/skippy-model-package/Cargo.toml | 1 + crates/skippy-model-package/src/preflight.rs | 62 +- crates/skippy-quantize/src/mtp_attach.rs | 5 +- .../src/binary_transport/binary_messaging.rs | 3 + .../binary_messaging/connection.rs | 670 ++++-------------- .../binary_messaging/control_messages.rs | 407 +++++++++++ .../binary_messaging/message_receive.rs | 34 + .../binary_messaging/session_lifecycle.rs | 57 ++ .../binary_messaging/telemetry.rs | 241 +++++++ .../src/binary_transport/stage_execution.rs | 10 +- scripts/hf-skippy-mtp-certify-job.py | 75 +- .../tests/test_hf_skippy_mtp_certify_job.py | 55 ++ .../tests/test_windows_native_runtime_deps.py | 2 +- ...-Inkling-MTP-and-KV-contiguity-state.patch | 269 +++++++ 20 files changed, 1413 insertions(+), 549 deletions(-) create mode 100644 crates/skippy-server/src/binary_transport/binary_messaging/control_messages.rs create mode 100644 crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs create mode 100644 crates/skippy-server/src/binary_transport/binary_messaging/session_lifecycle.rs create mode 100644 scripts/tests/test_hf_skippy_mtp_certify_job.py create mode 100644 third_party/llama.cpp/patches/0062-Harden-Inkling-MTP-and-KV-contiguity-state.patch diff --git a/Cargo.lock b/Cargo.lock index 88903d2472..310c7eb5d1 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", ] 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 6993ce43ae..2111840795 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 @@ -246,6 +246,12 @@ fn resolve_cache_type_k( 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( @@ -276,6 +282,12 @@ fn resolve_cache_type_v( 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( 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 9096b4f833..a557b51dea 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 @@ -783,6 +783,33 @@ cache_type_v = "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(); 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 e1b50cf7c6..f3dd625c5e 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 @@ -21,7 +21,6 @@ use std::path::{Path, PathBuf}; use std::time::Duration; pub(super) const MIN_STAGE_SOURCE_PREPARE_TIMEOUT: Duration = Duration::from_secs(30 * 60); -pub(super) const MAX_STAGE_SOURCE_PREPARE_TIMEOUT: Duration = Duration::from_secs(6 * 60 * 60); const STAGE_SOURCE_PREPARE_ALLOWANCE: Duration = Duration::from_secs(10 * 60); const STAGE_SOURCE_MIN_BYTES_PER_SEC: u64 = 16 * 1024 * 1024; @@ -345,10 +344,7 @@ pub(super) fn stage_source_prepare_timeout( 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) - .clamp( - MIN_STAGE_SOURCE_PREPARE_TIMEOUT, - MAX_STAGE_SOURCE_PREPARE_TIMEOUT, - ) + .max(MIN_STAGE_SOURCE_PREPARE_TIMEOUT) } pub(super) fn split_runtime_stage_load_request( 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 e1ac153959..878768f347 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 @@ -616,7 +616,7 @@ fn split_startup_error_messages_include_specific_blocker_tokens() { #[test] fn stage_source_prepare_timeout_scales_with_assigned_package_bytes() { let package = skippy::SkippyPackageIdentity { - source_model_bytes: 321_400_000_000, + source_model_bytes: 975_000_000_000, layer_count: 66, ..package(66) }; @@ -642,7 +642,7 @@ fn stage_source_prepare_timeout_scales_with_assigned_package_bytes() { assert!(small_timeout > MIN_STAGE_SOURCE_PREPARE_TIMEOUT); assert!(large_timeout > small_timeout); - assert!(large_timeout <= MAX_STAGE_SOURCE_PREPARE_TIMEOUT); + assert!(large_timeout > Duration::from_secs(6 * 60 * 60)); } #[test] diff --git a/crates/skippy-ffi/src/lib.rs b/crates/skippy-ffi/src/lib.rs index 7a499c0f26..a76464b9a1 100644 --- a/crates/skippy-ffi/src/lib.rs +++ b/crates/skippy-ffi/src/lib.rs @@ -1,6 +1,6 @@ pub const ABI_VERSION_MAJOR: u32 = 0; pub const ABI_VERSION_MINOR: u32 = 1; -pub const ABI_VERSION_PATCH: u32 = 34; +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; @@ -1183,9 +1183,14 @@ 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. @@ -1194,6 +1199,14 @@ 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 { diff --git a/crates/skippy-model-package/Cargo.toml b/crates/skippy-model-package/Cargo.toml index fe5c4eb59d..feae008418 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/preflight.rs b/crates/skippy-model-package/src/preflight.rs index 122d87ef7a..709fad0d26 100644 --- a/crates/skippy-model-package/src/preflight.rs +++ b/crates/skippy-model-package/src/preflight.rs @@ -3,6 +3,7 @@ use std::fs; use std::path::Path; use serde::{Deserialize, Serialize}; +use skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH; mod artifact_io; @@ -1018,14 +1019,21 @@ fn validate_window_policy( "use positive window sizes", ); } - if window.pipeline_depth == Some(0) { + 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 greater than zero" + "speculative strategy {name} window_policy.pipeline_depth must be between 1 and {MAX_VERIFY_WINDOW_PIPELINE_DEPTH}" ), Some("model-package.json".to_string()), - "set pipeline_depth to a positive in-flight verification-window capacity", + 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 { @@ -2163,6 +2171,54 @@ mod tests { 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-quantize/src/mtp_attach.rs b/crates/skippy-quantize/src/mtp_attach.rs index bb28ac6c97..698a6ec6a6 100644 --- a/crates/skippy-quantize/src/mtp_attach.rs +++ b/crates/skippy-quantize/src/mtp_attach.rs @@ -47,8 +47,9 @@ pub(crate) fn run_validate_mtp_attach(args: ValidateMtpAttachArgs) -> Result<()> skippy_ffi::native_runtime_loaded(), "validate-mtp-attach requires a statically linked standalone build or a loaded native runtime" ); - let native_mtp_multimodal_feature = - skippy_ffi::abi_features() & skippy_ffi::FEATURE_INKLING_MTP_MM != 0; + 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" diff --git a/crates/skippy-server/src/binary_transport/binary_messaging.rs b/crates/skippy-server/src/binary_transport/binary_messaging.rs index 559dea069e..987cffbdfc 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging.rs @@ -34,8 +34,11 @@ 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; 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 9caf509de0..eb012622d9 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -1,15 +1,18 @@ use super::async_forwarder::AsyncForwarder; -use super::reply::drain_deferred_prefill_replies; -use super::reply::send_stage_reply; -use super::reply::{ - configure_prediction_return_stream, normalize_downstream_prefix_restore_reply, - reply_window_for_message, +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::session_lifecycle::align_session_to_message; use super::session_tracker::{ConnectionSessionTracker, 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, }; @@ -20,14 +23,12 @@ 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::stage_execution::binary_message_attrs; use crate::binary_transport::stage_execution::binary_message_session_id; @@ -35,7 +36,6 @@ 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; @@ -48,33 +48,23 @@ use crate::kv_integration::{KvStageIntegration, model_requires_recurrent_state}; 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, @@ -144,7 +134,7 @@ fn handle_binary_connection_messages( first_message: StageWireMessage, session_tracker: &mut ConnectionSessionTracker, ) -> Result<()> { - let connection_session_id = BINARY_SESSION_COUNTER.fetch_add(1, Ordering::Relaxed); + let connection_session_id = next_connection_session_id(); let positional_speculation_supported = !model_requires_recurrent_state(config); let max_deferred_prefill_replies = reply_credit_limit.unwrap_or_else(|| max_inflight.saturating_sub(1)); @@ -169,20 +159,15 @@ fn handle_binary_connection_messages( 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); @@ -197,347 +182,114 @@ fn handle_binary_connection_messages( 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, - ); - } + 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, - ); - 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(&mut *upstream, stop_stats).context("send stop ACK")?; + handle_stop( + config, + runtime, + kv, + telemetry, + upstream, + downstream.as_mut(), + wire_dtype, + downstream_wire_condition, + &message, + &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() { - if let Some(forwarder) = async_forwarder.as_mut() { - 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")?; - { - let mut runtime = runtime.lock().expect("runtime lock poisoned"); - runtime - .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")?; - } + 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() { - 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( + handle_session_control( + runtime, + upstream, downstream.as_mut(), + wire_dtype, + downstream_wire_condition, + &message, + &session_key, &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")?; + &mut pending_reply_stats, + async_forwarder.as_mut(), + )?; 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( + 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 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")?; + &mut pending_reply_stats, + async_forwarder.as_mut(), + prediction_return_sinks, + &mut prediction_return_streams, + )?; 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_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, - &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 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_id, + &mut pending_prefill_replies, + &mut pending_reply_stats, + async_forwarder.as_mut(), + prediction_return_sinks, + &mut prediction_return_streams, + )?; continue; } @@ -568,37 +320,17 @@ fn handle_binary_connection_messages( } 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, @@ -1222,154 +954,44 @@ fn handle_binary_connection_messages( 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", + 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, - ); - 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, - message_start_unix_nanos, - message_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 0000000000..ad69ee5e2e --- /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 0000000000..8450c2c9a1 --- /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/session_lifecycle.rs b/crates/skippy-server/src/binary_transport/binary_messaging/session_lifecycle.rs new file mode 100644 index 0000000000..75025fc471 --- /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/telemetry.rs b/crates/skippy-server/src/binary_transport/binary_messaging/telemetry.rs index a3cba4a44a..bb61159991 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/stage_execution.rs b/crates/skippy-server/src/binary_transport/stage_execution.rs index e9923b5795..9afcfad2ba 100644 --- a/crates/skippy-server/src/binary_transport/stage_execution.rs +++ b/crates/skippy-server/src/binary_transport/stage_execution.rs @@ -117,7 +117,15 @@ pub(in crate::binary_transport) fn take_ready_downstream( } fn complete_downstream_ready(stream: &mut TcpStream, timeout: Duration) -> Result<()> { - send_client_ready_hello_if_enabled(stream).context("send downstream client ready hello")?; + 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")?; diff --git a/scripts/hf-skippy-mtp-certify-job.py b/scripts/hf-skippy-mtp-certify-job.py index e649762cc1..d9f0b3d3ab 100644 --- a/scripts/hf-skippy-mtp-certify-job.py +++ b/scripts/hf-skippy-mtp-certify-job.py @@ -7,14 +7,22 @@ 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) @@ -73,19 +81,72 @@ def require_gguf_magic(path: Path) -> Path: 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)) - scheme = urllib.parse.urlparse(args.projector_url).scheme - if scheme not in ("http", "https"): - raise RuntimeError(f"unsupported projector URL scheme: {scheme!r}") + 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") - print(f"+ download {args.projector_url} -> {target}", flush=True) - with urllib.request.urlopen(args.projector_url) as response, temporary.open("wb") as output: - shutil.copyfileobj(response, output, length=8 * 1024 * 1024) - temporary.replace(target) + 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) 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 0000000000..8461ce02b1 --- /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 5ad3a94a5a..637695945f 100644 --- a/scripts/tests/test_windows_native_runtime_deps.py +++ b/scripts/tests/test_windows_native_runtime_deps.py @@ -129,7 +129,7 @@ def test_package_verifier_accepts_closed_windows_dependency_graph(self): "runtime": { "id": artifact.name, "mesh_version": "0.72.1", - "skippy_abi": "0.1.34", + "skippy_abi": "0.1.35", "platform": { "os": "windows", "arch": "x86_64", 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 0000000000..a308bd296c --- /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) + From 21f052883fe9e4bcd9cb11f91ae8abc84d656632 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 31 Jul 2026 01:10:07 +1000 Subject: [PATCH 16/26] fix(skippy): harden combined family runtime --- .../queue-unsloth-layer-packages.yml | 2 + .../src/runtime/local_split.rs | 10 -- .../src/runtime/local_split/tests.rs | 15 +- .../src/runtime/startup_retry.rs | 11 ++ crates/skippy-correctness/src/cli.rs | 2 + .../src/runner/stage_fa_parity.rs | 2 +- crates/skippy-quantize/src/backend.rs | 2 +- crates/skippy-quantize/src/gguf_writer.rs | 24 ++- .../skippy-quantize/src/gguf_writer_tests.rs | 34 +++++ crates/skippy-quantize/src/tensor_map.rs | 3 +- .../skippy-quantize/src/tokenizer_metadata.rs | 2 +- .../binary_messaging/session_tracker.rs | 11 +- .../src/frontend/embedded_generation.rs | 12 +- .../frontend/embedded_generation/lifecycle.rs | 14 +- .../src/frontend/native_mtp/verify_window.rs | 6 +- crates/skippy-server/src/runtime_state.rs | 20 +-- scripts/hf-skippy-convert-job.py | 2 + ...kling-MTP-sidecar-and-verify-cleanup.patch | 140 ++++++++++++++++++ 18 files changed, 262 insertions(+), 50 deletions(-) create mode 100644 third_party/llama.cpp/patches/0063-Harden-Inkling-MTP-sidecar-and-verify-cleanup.patch diff --git a/.github/workflows/queue-unsloth-layer-packages.yml b/.github/workflows/queue-unsloth-layer-packages.yml index c33cf38034..38b6b922ab 100644 --- a/.github/workflows/queue-unsloth-layer-packages.yml +++ b/.github/workflows/queue-unsloth-layer-packages.yml @@ -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/crates/mesh-llm-host-runtime/src/runtime/local_split.rs b/crates/mesh-llm-host-runtime/src/runtime/local_split.rs index 36d3ad242b..cba5ad109c 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_split.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_split.rs @@ -205,16 +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" ); - 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 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 878768f347..9bf59e9495 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 @@ -82,7 +82,6 @@ fn canonical_coordinator_is_identical_with_divergent_observer_signals() { 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 mut package_planning_called = false; let gate = canonical_coordinator_gate(local.node_id, vec![local, coordinator]) .expect("canonical coordinator gate"); @@ -90,10 +89,10 @@ fn noncanonical_gate_returns_standby_without_invoking_package_planning() { CanonicalCoordinatorGate::Standby { coordinator: selected, } => assert_eq!(selected, coordinator.node_id), - CanonicalCoordinatorGate::Coordinator(_) => package_planning_called = true, + CanonicalCoordinatorGate::Coordinator(_) => { + panic!("local node must not be elected coordinator") + } } - - assert!(!package_planning_called); } #[test] @@ -1643,6 +1642,10 @@ fn real_inkling_metadata_plans_family_kv_not_size_tiered() { 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); @@ -1665,7 +1668,7 @@ fn real_inkling_metadata_plans_family_kv_not_size_tiered() { "inkling must resolve a q4_0 family K/V default" ); assert_eq!( - planned, size_tiered, - "family-aware planning and the large-model policy must agree on Q4_0" + 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/startup_retry.rs b/crates/mesh-llm-host-runtime/src/runtime/startup_retry.rs index 18ea0500cb..ec547fc17c 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/startup_retry.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/startup_retry.rs @@ -8,6 +8,7 @@ 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")) } @@ -54,6 +55,16 @@ mod tests { )); } + #[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( diff --git a/crates/skippy-correctness/src/cli.rs b/crates/skippy-correctness/src/cli.rs index 21ee19ac1b..63291938f7 100644 --- a/crates/skippy-correctness/src/cli.rs +++ b/crates/skippy-correctness/src/cli.rs @@ -330,6 +330,8 @@ pub enum StatePayloadKind { 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)] diff --git a/crates/skippy-correctness/src/runner/stage_fa_parity.rs b/crates/skippy-correctness/src/runner/stage_fa_parity.rs index 736f082864..8205a49487 100644 --- a/crates/skippy-correctness/src/runner/stage_fa_parity.rs +++ b/crates/skippy-correctness/src/runner/stage_fa_parity.rs @@ -89,7 +89,7 @@ fn decode_boundary( filter_tensors_on_load: true, }; let selection = select_layer_package_parts(&PackageStageRequest { - model_id: "unsloth/inkling-GGUF:UD-Q2_K_XL".to_string(), + 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(), diff --git a/crates/skippy-quantize/src/backend.rs b/crates/skippy-quantize/src/backend.rs index 4f2e8f6d98..e4ac486dc5 100644 --- a/crates/skippy-quantize/src/backend.rs +++ b/crates/skippy-quantize/src/backend.rs @@ -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/gguf_writer.rs b/crates/skippy-quantize/src/gguf_writer.rs index 1c559077bc..cddf9d38ff 100644 --- a/crates/skippy-quantize/src/gguf_writer.rs +++ b/crates/skippy-quantize/src/gguf_writer.rs @@ -457,14 +457,24 @@ fn inkling_w13_tensor_sources( tensor_name_map: TensorNameMap, output_type: Option, ) -> Result> { - let TensorNameMap::HfToGgufWithMtp { layer_start } = tensor_name_map else { - anyhow::bail!("Inkling MTP conversion requires an MTP-aware tensor name map"); + 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()))? }; - let depth = inkling_mtp_depth(tensor.name())? - .with_context(|| format!("missing Inkling MTP depth in {}", tensor.name()))?; - let layer = layer_start - .checked_add(depth) - .context("Inkling MTP layer id overflow")?; ensure!( tensor.shape().len() == 2, "Inkling MTP fused w13 tensor {} must be rank 2, got {:?}", diff --git a/crates/skippy-quantize/src/gguf_writer_tests.rs b/crates/skippy-quantize/src/gguf_writer_tests.rs index 92b23a6bfd..3182bc0f83 100644 --- a/crates/skippy-quantize/src/gguf_writer_tests.rs +++ b/crates/skippy-quantize/src/gguf_writer_tests.rs @@ -344,6 +344,40 @@ fn writes_inkling_mtp_streaming_transforms() { 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 parsed = parse_test_gguf(&fs::read(&output).unwrap()); + assert_eq!(parsed.tensor_count, 2); + assert_eq!(parsed.tensor("blk.3.ffn_gate.weight").dims, vec![2, 2]); + assert_eq!(parsed.tensor("blk.3.ffn_up.weight").dims, vec![2, 2]); + fs::remove_dir_all(root).unwrap(); +} + #[test] fn validates_qwen_dense_native_conversion_fixture() { let root = unique_temp_dir(); diff --git a/crates/skippy-quantize/src/tensor_map.rs b/crates/skippy-quantize/src/tensor_map.rs index b0c028ccdc..4afae0fc1a 100644 --- a/crates/skippy-quantize/src/tensor_map.rs +++ b/crates/skippy-quantize/src/tensor_map.rs @@ -90,7 +90,8 @@ pub(crate) fn is_shared_mtp_context_tensor(name: &str) -> bool { } pub(crate) fn is_inkling_fused_w13(name: &str) -> bool { - name.starts_with("model.mtp.layers.") && name.ends_with(".mlp.w13_dn.weight") + (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> { diff --git a/crates/skippy-quantize/src/tokenizer_metadata.rs b/crates/skippy-quantize/src/tokenizer_metadata.rs index 06b4566b82..05c5877c53 100644 --- a/crates/skippy-quantize/src/tokenizer_metadata.rs +++ b/crates/skippy-quantize/src/tokenizer_metadata.rs @@ -8,7 +8,7 @@ use serde_json::Value; use crate::gguf_writer::GgufKv; const TOKEN_TYPE_NORMAL: i32 = 1; -const TOKEN_TYPE_UNUSED: i32 = 2; +const TOKEN_TYPE_UNUSED: i32 = 5; const TOKEN_TYPE_CONTROL: i32 = 3; pub(crate) fn push_tokenizer_metadata( 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 index 5f395e2d15..dfec1fa597 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/session_tracker.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/session_tracker.rs @@ -44,8 +44,15 @@ pub(super) fn release_tracked_connection_sessions( if orphaned.is_empty() { return; } - let Ok(mut runtime) = runtime.lock() else { - return; + let mut runtime = match runtime.lock() { + Ok(runtime) => runtime, + Err(_) => { + eprintln!( + "failed to reclaim {} orphaned binary stage session(s): runtime lock poisoned", + orphaned.len() + ); + return; + } }; for session_key in orphaned { match runtime.drop_session_timed(&session_key) { diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 34823496ea..51e9c62fed 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -63,8 +63,14 @@ impl StageOpenAiBackend { .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 should_open_upstream_prediction_return(request.native_mtp_enabled) - && let Some(prediction_return) = request.prediction_return.as_ref() + let standalone_ngram_return = !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_return, + ) && let Some(prediction_return) = request.prediction_return.as_ref() { match crate::binary_transport::direct_return::open_downstream_prediction_return_stream( request.config, @@ -945,7 +951,7 @@ 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, )? { diff --git a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs index 06055ccb15..a468348fa0 100644 --- a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs +++ b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs @@ -120,8 +120,11 @@ pub(super) enum DirectPredictionReturnPath { ReverseFallback, } -pub(super) fn should_open_upstream_prediction_return(native_mtp_enabled: bool) -> bool { - native_mtp_enabled +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 direct_prediction_return_path( @@ -509,9 +512,10 @@ mod tests { } #[test] - fn pure_ngram_skips_the_blocking_upstream_return_open() { - assert!(!should_open_upstream_prediction_return(false)); - assert!(should_open_upstream_prediction_return(true)); + 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] 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 e9cad05cdf..7049364e80 100644 --- a/crates/skippy-server/src/frontend/native_mtp/verify_window.rs +++ b/crates/skippy-server/src/frontend/native_mtp/verify_window.rs @@ -209,10 +209,8 @@ 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( - native_mtp_verify_decision.commit_count, - consumed_positions, - ); + let checkpoint_no_longer_needed = + verify_checkpoint_no_longer_needed(committed_positions, consumed_positions); if checkpoint_no_longer_needed { self.retire_verify_window( request, diff --git a/crates/skippy-server/src/runtime_state.rs b/crates/skippy-server/src/runtime_state.rs index 05731f3838..0264567394 100644 --- a/crates/skippy-server/src/runtime_state.rs +++ b/crates/skippy-server/src/runtime_state.rs @@ -894,10 +894,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(()) } @@ -925,10 +926,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(()) } @@ -1395,7 +1397,7 @@ mod tests { }; #[test] - fn recurrent_prefix_restore_moves_tracked_position_backwards() { + 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); diff --git a/scripts/hf-skippy-convert-job.py b/scripts/hf-skippy-convert-job.py index 6eb17530a4..c40dab0180 100644 --- a/scripts/hf-skippy-convert-job.py +++ b/scripts/hf-skippy-convert-job.py @@ -128,6 +128,8 @@ def convert(args: argparse.Namespace, root: Path) -> Path: "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") 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 0000000000..76af7c7305 --- /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 + From 635d240a8012c182fb290f7724f53e2e23e981f4 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 31 Jul 2026 02:04:05 +1000 Subject: [PATCH 17/26] fix(skippy): retire final committed verify span --- .../skippy-quantize/src/gguf_writer_tests.rs | 21 +++++++++-- .../src/frontend/embedded_generation.rs | 35 +++++++++---------- 2 files changed, 35 insertions(+), 21 deletions(-) diff --git a/crates/skippy-quantize/src/gguf_writer_tests.rs b/crates/skippy-quantize/src/gguf_writer_tests.rs index 3182bc0f83..b076c2eff2 100644 --- a/crates/skippy-quantize/src/gguf_writer_tests.rs +++ b/crates/skippy-quantize/src/gguf_writer_tests.rs @@ -371,10 +371,25 @@ fn writes_inkling_trunk_fused_w13_streaming_transforms() { ) .unwrap(); - let parsed = parse_test_gguf(&fs::read(&output).unwrap()); + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); assert_eq!(parsed.tensor_count, 2); - assert_eq!(parsed.tensor("blk.3.ffn_gate.weight").dims, vec![2, 2]); - assert_eq!(parsed.tensor("blk.3.ffn_up.weight").dims, vec![2, 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(); } diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 51e9c62fed..e12ec0879e 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -1264,24 +1264,6 @@ impl StageOpenAiBackend { let fully_accepted_window = !native_mtp_verify_decision.rejected && native_mtp_verify_decision.accepted_proposal_tokens == window.proposal_tokens.len(); - let checkpoint_no_longer_needed = verify_checkpoint_no_longer_needed( - native_mtp_verify_decision.commit_count, - window.input_tokens.len(), - ); - if checkpoint_no_longer_needed { - 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(), - }, - )?; - } let pipeline_continues = fully_accepted_window; let accepted_candidate_tokens = native_mtp_verify_decision.accepted_proposal_tokens; @@ -1381,6 +1363,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; From 9319897560e038bccdce9a395320ce2d72441431 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 31 Jul 2026 02:24:09 +1000 Subject: [PATCH 18/26] fix(skippy): retire failed lane before replacement --- .../src/frontend/embedded_generation/lifecycle.rs | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs index a468348fa0..ca776fa95a 100644 --- a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs +++ b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs @@ -365,9 +365,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(()); } @@ -397,7 +403,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(()) From 779d27b6b73e293f17b6dd10158da3a63d5e3c07 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 31 Jul 2026 04:02:33 +1000 Subject: [PATCH 19/26] fix(skippy): retire partial exact replay trials --- .../skippy-server/src/frontend/prefix_cache.rs | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/crates/skippy-server/src/frontend/prefix_cache.rs b/crates/skippy-server/src/frontend/prefix_cache.rs index be25393b0e..bc982bc0e2 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)); + } } From 51da4a90b5db4e667e3ec8c407f546c633b07d50 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 31 Jul 2026 04:02:36 +1000 Subject: [PATCH 20/26] fix(skippy): close combined review gaps --- .../src/runtime/split_participant_settle.rs | 18 ++++++++++++++++- crates/skippy-quantize/src/tensor_map.rs | 20 ++++++++++++++----- crates/skippy-server/README.md | 2 +- 3 files changed, 33 insertions(+), 7 deletions(-) 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 index ee5ace7841..925687053b 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/split_participant_settle.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/split_participant_settle.rs @@ -17,6 +17,10 @@ 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)] @@ -82,7 +86,7 @@ impl<'a> SplitMembershipWait<'a> { node, model_name, model_ref, - deadline: tokio::time::Instant::now() + timeout, + deadline: tokio::time::Instant::now() + membership_settle_timeout(timeout), barrier: SplitMembershipSettleBarrier::default(), last_logged_signature: Vec::new(), } @@ -323,6 +327,18 @@ mod tests { 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(); diff --git a/crates/skippy-quantize/src/tensor_map.rs b/crates/skippy-quantize/src/tensor_map.rs index 4afae0fc1a..26013add19 100644 --- a/crates/skippy-quantize/src/tensor_map.rs +++ b/crates/skippy-quantize/src/tensor_map.rs @@ -121,17 +121,16 @@ fn normalize_inkling_mtp_source_name(name: &str, layer_start: u32) -> Result() .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_start + depth, - suffix - ))) + Ok(Some(format!("model.layers.{layer}.{suffix}"))) } fn map_mtp_source_tensor(name: &str) -> Result> { @@ -532,4 +531,15 @@ mod tests { "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-server/README.md b/crates/skippy-server/README.md index 90393a1134..d3b4eb0dc7 100644 --- a/crates/skippy-server/README.md +++ b/crates/skippy-server/README.md @@ -100,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. From bec59731bf68a45f5cfe026154122a3b3b14ff8a Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 31 Jul 2026 04:02:46 +1000 Subject: [PATCH 21/26] docs(skippy): promote pinned Laguna Q4 package --- docs/skippy/FAMILY_STATUS.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/skippy/FAMILY_STATUS.md b/docs/skippy/FAMILY_STATUS.md index c7235b07b5..b0d1fc54b3 100644 --- a/docs/skippy/FAMILY_STATUS.md +++ b/docs/skippy/FAMILY_STATUS.md @@ -8,7 +8,7 @@ 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-07-28. +Last updated: 2026-07-31. ## Context Capacity Contract @@ -55,6 +55,7 @@ correctness smoke or from GGUF metadata alone. | 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. | @@ -105,7 +106,6 @@ smoke, reviewed topology records, and family-specific policy notes are updated. ```text Gemma text -Laguna S 2.1 Q4_K_M (pinned package-backed M5 two-stage single-step and three-stage chain parity, plus ordinary M5 Metal 0..36 + Australian Vast CUDA 36..48 Mesh serving with a 44,460-token prompt; suffix N-gram depth 2 is published as the package default, with a fresh live no-override confirmation still pending; broader family certification remains pending) ``` ## Exceptions From a3a60e2e5cedeb3c4e21bcda526996eb961f1de5 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 31 Jul 2026 04:04:49 +1000 Subject: [PATCH 22/26] refactor(skippy): isolate prediction return startup --- .../src/frontend/embedded_generation.rs | 33 ++----------------- .../frontend/embedded_generation/lifecycle.rs | 31 +++++++++++++++++ 2 files changed, 34 insertions(+), 30 deletions(-) diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index e12ec0879e..5f5dfe9bbd 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -35,9 +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, should_open_upstream_prediction_return, - speculation_after_prefix_restore, + 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; @@ -62,33 +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; - let standalone_ngram_return = !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_return, - ) && 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 = (|| { diff --git a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs index ca776fa95a..9396a3a120 100644 --- a/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs +++ b/crates/skippy-server/src/frontend/embedded_generation/lifecycle.rs @@ -127,6 +127,37 @@ pub(super) fn should_open_upstream_prediction_return( 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, From b2658577ab1feebe38b748866751fa2eeb056bb5 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:58:51 +1000 Subject: [PATCH 23/26] fix: address consolidated model review --- .../src/api/split_readiness.rs | 4 +- crates/mesh-llm-host-runtime/src/mesh/mod.rs | 4 +- crates/mesh-llm-host-runtime/src/mesh/node.rs | 235 +--------- .../src/mesh/node/startup.rs | 228 ++++++++++ crates/skippy-model-package/src/preflight.rs | 429 +----------------- .../src/preflight/artifacts.rs | 427 +++++++++++++++++ crates/skippy-quantize/src/direct_convert.rs | 3 +- crates/skippy-quantize/src/gguf_writer.rs | 52 +++ .../skippy-quantize/src/gguf_writer_tests.rs | 58 +++ crates/skippy-quantize/src/main.rs | 5 +- crates/skippy-quantize/src/native_convert.rs | 54 ++- scripts/hf-skippy-convert-job.py | 42 +- scripts/tests/test_hf_skippy_convert_job.py | 56 +++ 13 files changed, 932 insertions(+), 665 deletions(-) create mode 100644 crates/mesh-llm-host-runtime/src/mesh/node/startup.rs create mode 100644 crates/skippy-model-package/src/preflight/artifacts.rs create mode 100644 scripts/tests/test_hf_skippy_convert_job.py 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 1ec3327e6a..0b33823710 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/mesh/mod.rs b/crates/mesh-llm-host-runtime/src/mesh/mod.rs index efdf4c1cd2..67bf98877a 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/mod.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/mod.rs @@ -5,8 +5,10 @@ //! 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, max_split_rtt_ms, split_allow_relay_paths, diff --git a/crates/mesh-llm-host-runtime/src/mesh/node.rs b/crates/mesh-llm-host-runtime/src/mesh/node.rs index d098639b57..f871aeae85 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,221 +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 local_runtime_capacity_bytes = - super::capacity::capped_capacity_bytes(hw.vram_bytes, max_vram_gb); - let mut vram_bytes = 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, - } -} - -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, @@ -1729,7 +1510,7 @@ 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 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 0000000000..83ebf4d2dd --- /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/skippy-model-package/src/preflight.rs b/crates/skippy-model-package/src/preflight.rs index 709fad0d26..3da3fb1c36 100644 --- a/crates/skippy-model-package/src/preflight.rs +++ b/crates/skippy-model-package/src/preflight.rs @@ -6,12 +6,18 @@ use serde::{Deserialize, Serialize}; use skippy_protocol::MAX_VERIFY_WINDOW_PIPELINE_DEPTH; mod artifact_io; +mod artifacts; use crate::generation_manifest::{ PackageGeneration, PackageGenerationExperimentalPolicy, PackageGenerationPolicy, PackageGenerationThresholds, }; use artifact_io::{file_sha256, safe_relative_path, sha256_bytes}; +#[cfg(test)] +use artifacts::validate_artifact_sha; +use artifacts::{ + build_stage_reports, collect_artifacts, validate_artifacts, validate_layer_coverage, +}; #[derive(Debug, Clone, Copy, Default)] pub(crate) struct PackagePreflightOptions { @@ -1161,429 +1167,6 @@ fn preflight_window_policy(window: &PackageWindowPolicy) -> PreflightWindowPolic } } -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, 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 0000000000..f3f6234561 --- /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-quantize/src/direct_convert.rs b/crates/skippy-quantize/src/direct_convert.rs index a07d8619cb..28cd045078 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_writer.rs b/crates/skippy-quantize/src/gguf_writer.rs index cddf9d38ff..1c0ffc5baa 100644 --- a/crates/skippy-quantize/src/gguf_writer.rs +++ b/crates/skippy-quantize/src/gguf_writer.rs @@ -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, diff --git a/crates/skippy-quantize/src/gguf_writer_tests.rs b/crates/skippy-quantize/src/gguf_writer_tests.rs index b076c2eff2..8a89f6b4a5 100644 --- a/crates/skippy-quantize/src/gguf_writer_tests.rs +++ b/crates/skippy-quantize/src/gguf_writer_tests.rs @@ -1237,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(); diff --git a/crates/skippy-quantize/src/main.rs b/crates/skippy-quantize/src/main.rs index 7856e134e2..43874e830b 100644 --- a/crates/skippy-quantize/src/main.rs +++ b/crates/skippy-quantize/src/main.rs @@ -1000,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/native_convert.rs b/crates/skippy-quantize/src/native_convert.rs index 454abd3348..3b98312592 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/scripts/hf-skippy-convert-job.py b/scripts/hf-skippy-convert-job.py index c40dab0180..784787e3ba 100644 --- a/scripts/hf-skippy-convert-job.py +++ b/scripts/hf-skippy-convert-job.py @@ -8,6 +8,7 @@ from __future__ import annotations import argparse +import json import os import shutil import subprocess @@ -134,6 +135,7 @@ def convert(args: argparse.Namespace, root: Path) -> Path: 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 @@ -149,16 +151,42 @@ def upload(args: argparse.Namespace, artifact_dir: Path) -> None: ) +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: - 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 or not any(artifact_dir.glob("*.gguf")): - details = f"; missing {', '.join(missing)}" if missing else "" - raise FileNotFoundError( - f"complete converted artifact not found: {artifact_dir}{details}" - ) + validate_converted_artifact(artifact_dir) return artifact_dir 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 0000000000..958483a5bc --- /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() From 64ab9f32ef1e342edc53997f29352bca3485cc40 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 31 Jul 2026 16:58:49 +1000 Subject: [PATCH 24/26] fix(skippy): surface orphan cleanup failures --- .../binary_messaging/connection.rs | 10 ++- .../binary_messaging/session_tracker.rs | 74 +++++++++++++++---- 2 files changed, 65 insertions(+), 19 deletions(-) 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 eb012622d9..0a9579aab5 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -7,7 +7,10 @@ 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::session_lifecycle::align_session_to_message; -use super::session_tracker::{ConnectionSessionTracker, release_tracked_connection_sessions}; +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; @@ -108,8 +111,9 @@ pub(super) fn handle_binary_connection( first_message, &mut session_tracker, ); - release_tracked_connection_sessions(config, runtime, telemetry, &mut session_tracker); - result + 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)] 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 index dfec1fa597..77030200e9 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/session_tracker.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/session_tracker.rs @@ -3,6 +3,7 @@ use crate::{ runtime_state::RuntimeState, telemetry::{Telemetry, lifecycle_attrs}, }; +use anyhow::{Context, Result, bail}; use serde_json::json; use skippy_protocol::StageConfig; use std::{ @@ -39,21 +40,18 @@ pub(super) fn release_tracked_connection_sessions( runtime: &Arc>, telemetry: &Telemetry, session_tracker: &mut ConnectionSessionTracker, -) { +) -> Result<()> { let orphaned = session_tracker.drain(); if orphaned.is_empty() { - return; + return Ok(()); } - let mut runtime = match runtime.lock() { - Ok(runtime) => runtime, - Err(_) => { - eprintln!( - "failed to reclaim {} orphaned binary stage session(s): runtime lock poisoned", - orphaned.len() - ); - return; - } - }; + 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) => { @@ -75,17 +73,39 @@ pub(super) fn release_tracked_connection_sessions( telemetry.emit("stage.binary_session_orphan_reclaimed", attrs); } Err(error) => { - eprintln!( - "failed to reclaim orphaned binary stage session {session_key}: {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; + use super::{ConnectionSessionTracker, combine_connection_and_cleanup_results}; + use anyhow::anyhow; #[test] fn tracker_drains_sessions_that_never_saw_a_stop() { @@ -106,4 +126,26 @@ mod tests { 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")); + } } From 6c54cd012fd59c4a246d54c9dca2cfa3b8aa01d9 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Sat, 1 Aug 2026 11:39:39 +1000 Subject: [PATCH 25/26] skippy: complete Inkling tool path and operator notes --- .../binary_messaging/connection.rs | 9 +- .../src/frontend/generation/server.rs | 13 +-- .../src/frontend/tool_emulation.rs | 54 +++++++++-- .../skippy-server/src/kv_integration/mod.rs | 2 - docs/SKIPPY_SPLITS.md | 77 +++++++++++++++ docs/skippy/FAMILY_STATUS.md | 15 ++- docs/skippy/LLAMA_PARITY.md | 3 +- ...cept-Inkling-arguments-tool-call-key.patch | 94 +++++++++++++++++++ 8 files changed, 238 insertions(+), 29 deletions(-) create mode 100644 third_party/llama.cpp/patches/0064-Accept-Inkling-arguments-tool-call-key.patch 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 0a9579aab5..3c809768e5 100644 --- a/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs +++ b/crates/skippy-server/src/binary_transport/binary_messaging/connection.rs @@ -47,7 +47,7 @@ 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; @@ -139,7 +139,6 @@ fn handle_binary_connection_messages( session_tracker: &mut ConnectionSessionTracker, ) -> Result<()> { let connection_session_id = next_connection_session_id(); - let positional_speculation_supported = !model_requires_recurrent_state(config); let max_deferred_prefill_replies = reply_credit_limit.unwrap_or_else(|| max_inflight.saturating_sub(1)); let mut pending_prefill_replies = 0usize; @@ -180,12 +179,6 @@ fn handle_binary_connection_messages( let session_id = binary_message_session_id(connection_session_id, &message); let session_key = session_id.to_string(); session_tracker.touch(&session_key); - if message.kind == WireMessageKind::VerifyWindow && !positional_speculation_supported { - bail!( - "stage-state v11 positional speculation requires an attention-only stage; {} contains recurrent state", - config.stage_id - ); - } emit_binary_message_received( telemetry, config, diff --git a/crates/skippy-server/src/frontend/generation/server.rs b/crates/skippy-server/src/frontend/generation/server.rs index cb8692fc62..58539c0a06 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`). 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/kv_integration/mod.rs b/crates/skippy-server/src/kv_integration/mod.rs index 0c364ad559..15d7158d9c 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/docs/SKIPPY_SPLITS.md b/docs/SKIPPY_SPLITS.md index 13bb69c177..228624be82 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/skippy/FAMILY_STATUS.md b/docs/skippy/FAMILY_STATUS.md index b0d1fc54b3..90b95daec9 100644 --- a/docs/skippy/FAMILY_STATUS.md +++ b/docs/skippy/FAMILY_STATUS.md @@ -104,9 +104,18 @@ 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 4804ac475c..4c300ba6eb 100644 --- a/docs/skippy/LLAMA_PARITY.md +++ b/docs/skippy/LLAMA_PARITY.md @@ -371,7 +371,8 @@ Rows with distributed evidence call out the second backend explicitly. | `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` passed at split `24`; three-stage `chain` passed at splits `16,32`; ordinary two-node Mesh OpenAI serving passed at `0..36 / 36..48` with configured context `262144` | untested | untested | M5 Max f16 parity passed against pinned full-model revision `edd093522473dc7313b0738d8b4116b7f8b9745f`: baseline token `674` matched both two-stage and three-stage predictions. The real Mesh run used M5 Max/Metal plus an Australian Vast RTX 6000 Ada/CUDA worker, normal planning and lifecycle, and direct Iroh transport. A final three-request recovery probe used a 44,460-token prompt plus deterministic 65-token completion and returned byte-identical content on every request; decode measured 19.74, 18.24, and 15.42 tok/s. Each request proposed 64 suffix N-gram tokens and accepted 32 at verify pipeline depth `2`; the exact-prefix hit reused 44,416 prompt tokens. A 63-request stability harness returned 63 HTTP 200 responses with no unsupported-trim, all-lanes-busy, reset, decode, slot, or proactive-eviction fatal errors. Structured tool use remains uncertified: the formal harness passed 6/13 checks because textual `` output was not consistently projected into OpenAI `tool_calls`. Package manifest SHA-256 `0250cfb54ceeb94a9c71e48df447f780e32fc625553844d6403770f315be0237` now publishes suffix N-gram (`min=5`, `max=32`, proposal cap `48`) and verify pipeline depth `2` as the default. Resolver coverage proves a no-config launch consumes that package policy; a fresh live M5 + Vast no-override confirmation remains pending. Equal `24/24` placement at full context exhausted the 48 GB CUDA worker; the working capacity-bounded topology was `36/12`. Q8 wire, state handoff, structured tool use, and full native-context saturation remain unproven. | +| `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/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 0000000000..0831afcdb3 --- /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) + From 102aa6fb2e7779a00140124ac906a001011783ea Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Mon, 3 Aug 2026 08:45:21 +1000 Subject: [PATCH 26/26] llama: refresh Inkling patch for updated upstream --- .../0050-Add-TML-Inkling-architecture.patch | 113 +++++++++--------- 1 file changed, 56 insertions(+), 57 deletions(-) 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 index 71379ec73f..40d1adc6d3 100644 --- a/third_party/llama.cpp/patches/0050-Add-TML-Inkling-architecture.patch +++ b/third_party/llama.cpp/patches/0050-Add-TML-Inkling-architecture.patch @@ -1,4 +1,4 @@ -From ef9aea72c4da2cb4d92851ef8b4dd456fb48d070 Mon Sep 17 00:00:00 2001 +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 @@ -86,7 +86,7 @@ and audio support, and backend op tests at production shapes. create mode 100644 tools/mtmd/models/inkling.cpp diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp -index a309f0276..db9004d84 100644 +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( @@ -121,10 +121,10 @@ index a309f0276..db9004d84 100644 // Build ID parsers if keys are provided common_peg_parser id_parser = eps(); diff --git a/common/chat.cpp b/common/chat.cpp -index 91658e6df..89f518716 100644 +index 71871b3df..ae1569ab0 100644 --- a/common/chat.cpp +++ b/common/chat.cpp -@@ -2381,6 +2381,119 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t +@@ -2641,6 +2641,119 @@ static common_chat_params common_chat_params_init_minimax_m3(const common_chat_t return data; } @@ -244,7 +244,7 @@ index 91658e6df..89f518716 100644 namespace workaround { static void map_developer_role_to_system(json & messages) { -@@ -2798,6 +2911,14 @@ std::optional common_chat_try_specialized_template( +@@ -3058,6 +3171,14 @@ std::optional common_chat_try_specialized_template( return common_chat_params_init_cohere2moe(tmpl, params); } @@ -260,10 +260,10 @@ index 91658e6df..89f518716 100644 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 45c001b78..88b2d5201 100644 +index 1a47b851a..abdf76cdb 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py -@@ -110,6 +110,7 @@ TEXT_MODEL_MAP: dict[str, str] = { +@@ -111,6 +111,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "HunYuanVLForConditionalGeneration": "hunyuan", "HYV3ForCausalLM": "hunyuan", "IQuestCoderForCausalLM": "llama", @@ -271,7 +271,7 @@ index 45c001b78..88b2d5201 100644 "InternLM2ForCausalLM": "internlm", "InternLM3ForCausalLM": "internlm", "JAISLMHeadModel": "jais", -@@ -278,6 +279,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = { +@@ -279,6 +280,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = { "GraniteSpeechPlusForConditionalGeneration": "granite", "HunYuanVLForConditionalGeneration": "hunyuan", "Idefics3ForConditionalGeneration": "smolvlm", @@ -1952,7 +1952,7 @@ index d38057721..1166178b7 100644 if (rpc_get) { diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c -index cdc85b083..da33f642a 100644 +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] = { @@ -2060,10 +2060,10 @@ index cdc85b083..da33f642a 100644 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 4df94b9a3..f7762fae1 100644 +index 124ea28b0..643063e66 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py -@@ -555,6 +555,7 @@ class MODEL_ARCH(IntEnum): +@@ -556,6 +556,7 @@ class MODEL_ARCH(IntEnum): TALKIE = auto() MELLUM = auto() NANBEIGE = auto() @@ -2071,7 +2071,7 @@ index 4df94b9a3..f7762fae1 100644 class VISION_PROJECTOR_TYPE(IntEnum): -@@ -778,6 +779,14 @@ class MODEL_TENSOR(IntEnum): +@@ -779,6 +780,14 @@ class MODEL_TENSOR(IntEnum): SHORTCONV_CONV = auto() SHORTCONV_INPROJ = auto() SHORTCONV_OUTPROJ = auto() @@ -2086,7 +2086,7 @@ index 4df94b9a3..f7762fae1 100644 VISEXP_ATTN_QKV = auto() VISEXP_ATTN_OUT = auto() VISEXP_GATE = auto() -@@ -1159,6 +1168,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { +@@ -1168,6 +1177,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.TALKIE: "talkie", MODEL_ARCH.MELLUM: "mellum", MODEL_ARCH.NANBEIGE: "nanbeige", @@ -2094,7 +2094,7 @@ index 4df94b9a3..f7762fae1 100644 } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { -@@ -1380,6 +1390,13 @@ TENSOR_NAMES: dict[MODEL_TENSOR, 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", @@ -2108,7 +2108,7 @@ index 4df94b9a3..f7762fae1 100644 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", -@@ -4141,6 +4158,38 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { +@@ -4161,6 +4178,38 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_UP_EXP, MODEL_TENSOR.FFN_EXP_PROBS_B, ], @@ -2147,7 +2147,7 @@ index 4df94b9a3..f7762fae1 100644 MODEL_ARCH.SMALLTHINKER: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, -@@ -4810,6 +4859,7 @@ class GGUFValueType(IntEnum): +@@ -4834,6 +4883,7 @@ class GGUFValueType(IntEnum): class VisionProjectorType: @@ -2156,7 +2156,7 @@ index 4df94b9a3..f7762fae1 100644 GEMMA3NV = "gemma3nv" GEMMA3NA = "gemma3na" diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py -index 8299ac25b..21c79dbd8 100644 +index 1e991b873..5a0175a56 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -63,6 +63,7 @@ class TensorNameMap: @@ -2175,7 +2175,7 @@ index 8299ac25b..21c79dbd8 100644 ), MODEL_TENSOR.DENSE_2_OUT: ( "dense_2_out", # embeddinggemma -@@ -2496,6 +2498,53 @@ class TensorNameMap: +@@ -2548,6 +2550,53 @@ class TensorNameMap: # architecture-specific block mappings arch_block_mappings_cfg: dict[MODEL_ARCH, dict[MODEL_TENSOR, tuple[str, ...]]] = { @@ -2750,7 +2750,7 @@ index 000000000..95874d23d +{%- endif -%} +{#- Unsloth translation to jinja from TML's parser #} diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp -index a0945e501..a489e1ad5 100644 +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 = { @@ -2800,7 +2800,7 @@ index a0945e501..a489e1ad5 100644 { 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" }, -@@ -792,6 +815,9 @@ static const std::map LLM_TENSOR_INFOS = { +@@ -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}}, @@ -2810,7 +2810,7 @@ index a0945e501..a489e1ad5 100644 {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}}, -@@ -833,6 +859,13 @@ static const std::map LLM_TENSOR_INFOS = { +@@ -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}}, @@ -2824,7 +2824,7 @@ index a0945e501..a489e1ad5 100644 {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}}, -@@ -961,6 +994,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { +@@ -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: @@ -2832,7 +2832,7 @@ index a0945e501..a489e1ad5 100644 return true; default: return false; -@@ -1017,6 +1051,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { +@@ -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: @@ -2841,7 +2841,7 @@ index a0945e501..a489e1ad5 100644 default: return true; diff --git a/src/llama-arch.h b/src/llama-arch.h -index fb5eb3920..74446fdaf 100644 +index cbc97085e..6a283eb6b 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -149,6 +149,7 @@ enum llm_arch { @@ -3231,7 +3231,7 @@ index b23a7e8d8..f95591860 100644 // 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 36d0480e5..4d267514b 100644 +index d6a40fa30..8d00d5fa1 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -532,6 +532,15 @@ struct llama_layer { @@ -4056,10 +4056,10 @@ index 158ab28df..794c0f047 100644 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 8ce019b3b..d97046725 100644 +index e9e926e87..9031e71e9 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp -@@ -12397,6 +12397,103 @@ struct test_glm_dsa_compact_k_gather : public test_case { +@@ -12407,6 +12407,103 @@ struct test_glm_dsa_compact_k_gather : public test_case { } }; @@ -4163,7 +4163,7 @@ index 8ce019b3b..d97046725 100644 // GGML_OP_CROSS_ENTROPY_LOSS struct test_cross_entropy_loss : public test_case { const ggml_type type; -@@ -12976,6 +13073,7 @@ struct test_generic_op : public test_case { +@@ -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: @@ -4171,7 +4171,7 @@ index 8ce019b3b..d97046725 100644 { // Scale error with kv length to account for accumulating floating point error const int64_t kv = sources[1].ne[1]; -@@ -12999,7 +13097,7 @@ struct test_generic_op : public test_case { +@@ -13009,7 +13107,7 @@ struct test_generic_op : public test_case { } // FLASH_ATTN_EXT: src[3] is the KQ mask @@ -4180,7 +4180,7 @@ index 8ce019b3b..d97046725 100644 init_tensor_kq_mask(t); continue; } -@@ -15049,6 +15147,22 @@ static std::vector> make_test_cases_eval() { +@@ -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)); @@ -4200,14 +4200,14 @@ index 8ce019b3b..d97046725 100644 + 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)); + - test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); - test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); - test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); + // 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 4dd00efdd..fcf025d07 100644 +index 01b07953a..d0f985049 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp -@@ -2931,6 +2931,85 @@ static void test_template_output_peg_parsers(bool detailed_debug) { +@@ -2996,6 +2996,85 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .run(); } @@ -4987,7 +4987,7 @@ index 659ea9e9c..c7efc95a4 100644 // FIXME some models are segfaulting with WebGPU: #ifdef GGML_USE_WEBGPU diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt -index 18a8288ba..ae574bc98 100644 +index 15040e4af..461c4f429 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -38,6 +38,7 @@ add_library(mtmd @@ -4999,10 +4999,10 @@ index 18a8288ba..ae574bc98 100644 models/kimik25.cpp models/nemotron-v2-vl.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h -index 092041138..539a7c661 100644 +index 589fc724e..fc011f714 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h -@@ -284,6 +284,13 @@ +@@ -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" @@ -5016,7 +5016,7 @@ index 092041138..539a7c661 100644 // mobilenetv5 (gemma3n) definitions #define TN_MNV5_STEM_CONV "v.conv_stem.conv.weight" #define TN_MNV5_STEM_BIAS "v.conv_stem.conv.bias" -@@ -346,6 +353,7 @@ +@@ -353,6 +360,7 @@ struct clip_ctx; enum projector_type { @@ -5024,7 +5024,7 @@ index 092041138..539a7c661 100644 PROJECTOR_TYPE_MLP, PROJECTOR_TYPE_MLP_NORM, PROJECTOR_TYPE_LDP, -@@ -403,6 +411,7 @@ enum projector_type { +@@ -411,6 +419,7 @@ enum projector_type { }; static std::map PROJECTOR_TYPE_NAMES = { @@ -5033,7 +5033,7 @@ index 092041138..539a7c661 100644 { PROJECTOR_TYPE_LDP, "ldp" }, { PROJECTOR_TYPE_LDPV2, "ldpv2"}, diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h -index 8dc875497..3fd92b77f 100644 +index 146eabce2..bfe021922 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -33,7 +33,7 @@ enum resize_algo { @@ -5045,7 +5045,7 @@ index 8dc875497..3fd92b77f 100644 }; // Padding style for img_tool::resize -@@ -355,6 +355,11 @@ struct qf_block { +@@ -363,6 +363,11 @@ struct qf_block { std::vector qf_proj_layers; }; @@ -5057,7 +5057,7 @@ index 8dc875497..3fd92b77f 100644 struct clip_model { clip_modality modality = CLIP_MODALITY_VISION; projector_type proj_type = PROJECTOR_TYPE_MLP; -@@ -369,6 +374,12 @@ struct clip_model { +@@ -377,6 +382,12 @@ struct clip_model { ggml_tensor * norm_embd_w = nullptr; ggml_tensor * norm_embd_b = nullptr; @@ -5071,7 +5071,7 @@ index 8dc875497..3fd92b77f 100644 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 04614b93b..8fe2a5b2d 100644 +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 @@ -5085,7 +5085,7 @@ index 04614b93b..8fe2a5b2d 100644 case PROJECTOR_TYPE_GEMMA3: case PROJECTOR_TYPE_IDEFICS3: case PROJECTOR_TYPE_LFM2: -@@ -1294,6 +1298,21 @@ struct clip_model_loader { +@@ -1298,6 +1302,21 @@ struct clip_model_loader { // model-specific params switch (model.proj_type) { @@ -5107,7 +5107,7 @@ index 04614b93b..8fe2a5b2d 100644 case PROJECTOR_TYPE_MLP: case PROJECTOR_TYPE_MLP_NORM: case PROJECTOR_TYPE_LDP: -@@ -1923,7 +1942,8 @@ struct clip_model_loader { +@@ -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 = ( @@ -5117,7 +5117,7 @@ index 04614b93b..8fe2a5b2d 100644 // layers const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; -@@ -2009,6 +2029,23 @@ struct clip_model_loader { +@@ -2057,6 +2077,23 @@ struct clip_model_loader { switch (model.proj_type) { @@ -5141,7 +5141,7 @@ index 04614b93b..8fe2a5b2d 100644 case PROJECTOR_TYPE_MLP: case PROJECTOR_TYPE_MLP_NORM: { -@@ -3034,6 +3071,11 @@ struct clip_model_loader { +@@ -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); @@ -5153,7 +5153,7 @@ index 04614b93b..8fe2a5b2d 100644 return batch; } -@@ -3431,6 +3473,16 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { +@@ -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) { @@ -5170,7 +5170,7 @@ index 04614b93b..8fe2a5b2d 100644 case PROJECTOR_TYPE_MLP: case PROJECTOR_TYPE_MLP_NORM: case PROJECTOR_TYPE_JANUS_PRO: -@@ -3813,6 +3865,10 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 +@@ -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) { @@ -5181,7 +5181,7 @@ index 04614b93b..8fe2a5b2d 100644 case PROJECTOR_TYPE_MINICPMV: { // inspired from siglip: -@@ -4759,6 +4815,10 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 +@@ -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) { @@ -5305,7 +5305,7 @@ index 000000000..c072f66a2 + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h -index caed438ec..c671757f1 100644 +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 { @@ -5327,7 +5327,7 @@ index caed438ec..c671757f1 100644 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 ed68951c0..c7f0b9a10 100644 +index fea03557d..13e31ff54 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -283,6 +283,7 @@ struct filter_params { @@ -5423,7 +5423,7 @@ index ed68951c0..c7f0b9a10 100644 // mtmd_audio_preprocessor_whisper // diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h -index d8ec72b9d..eae32a762 100644 +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 { @@ -5752,7 +5752,7 @@ index 115cba51e..044d72f61 100644 // 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 6e61cf3e5..119e9fc79 100644 +index 93ca8cbcf..a70f326a9 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -393,6 +393,12 @@ struct mtmd_context { @@ -5782,7 +5782,7 @@ index 6e61cf3e5..119e9fc79 100644 case PROJECTOR_TYPE_QWEN2A: case PROJECTOR_TYPE_QWEN25O: { -@@ -1200,23 +1213,35 @@ struct mtmd_tokenizer { +@@ -1204,23 +1217,35 @@ struct mtmd_tokenizer { } size_t n_tokens = 0; @@ -5826,4 +5826,3 @@ index 6e61cf3e5..119e9fc79 100644 image_tokens->nx = n_tokens; -- 2.50.1 (Apple Git-155) -