legion-ref/mesh-llm
1
0
Fork 0
mirror of https://github.com/Mesh-LLM/mesh-llm.git synced 2026-08-08 22:23:19 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

skippy: land Laguna Q4 and experimental Inkling split serving (#1118)

Browse source 
* feat(skippy): add experimental Inkling text split serving

* feat(skippy): add Laguna staged runtime candidate

* docs(skippy): record Laguna package certification

* fix(skippy): integrate hybrid verify recovery and lane cleanup

* docs(skippy): record Laguna M5 parity

* test(skippy): certify Laguna three-stage parity

* docs(skippy): record Laguna distributed serving smoke

* docs(skippy): record real Laguna mesh evidence

* fix(skippy): make stage memory truly layer-local

* fix(packaging): preserve source revisions

* feat(skippy): honor package verification depth

* test(skippy): assert Laguna cache policy

* llama: linearize Laguna Inkling and recovery patches

* fix(skippy): reconcile combined family recovery state

* Harden combined Skippy family support

* fix(skippy): harden combined family runtime

* fix(skippy): retire final committed verify span

* fix(skippy): retire failed lane before replacement

* fix(skippy): retire partial exact replay trials

* fix(skippy): close combined review gaps

* docs(skippy): promote pinned Laguna Q4 package

* refactor(skippy): isolate prediction return startup

* fix: address consolidated model review

* fix(skippy): surface orphan cleanup failures

* skippy: complete Inkling tool path and operator notes

* llama: refresh Inkling patch for updated upstream

---------

Co-authored-by: Michael Neale <14976+michaelneale@users.noreply.github.com>
This commit is contained in:
Michael Neale 2026-08-03 17:03:05 +10:00 • committed by GitHub
parent d1c7948a9f
commit 6adbd5b0f6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
171 changed files with 17623 additions and 1931 deletions
Download patch file Download diff file Expand all files Collapse all files

4
.github/workflows/queue-unsloth-layer-packages.yml vendored
View file

@ -31,7 +31,7 @@ on:
required: false
default: "main"
quant_preference:
description: "Comma-separated 4-bit quant preference"
description: "Comma-separated quant preference"
required: false
default: "UD-Q4_K_XL,UD-Q4_K_M,Q4_K_XL,Q4_K_M"
split_candidate_vram_gib:
@ -64,6 +64,8 @@ jobs:
steps:
- uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0
with:
persist-credentials: false
- uses: dtolnay/rust-toolchain@4cda84d5c5c54efe2404f9d843567869ab1699d4 # stable 2026-07-16

2
Cargo.lock generated
View file

@ -7395,6 +7395,7 @@ dependencies = [
"serde_json",
"sha2 0.10.9",
"skippy-ffi",
"skippy-protocol",
"skippy-runtime",
"tokio",
]
@ -7438,6 +7439,7 @@ dependencies = [
"serde",
"serde_json",
"skippy-ffi",
"skippy-runtime",
]
[[package]]

8
Justfile
View file

@ -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="":

1
crates/llama-quant-ffi/src/lib.rs
View file

@ -38,6 +38,7 @@ pub enum LlamaFileType {
MostlyMxfp4Moe = 38,
MostlyNvfp4 = 39,
MostlyQ1_0 = 40,
MostlyQ2_0 = 41,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]

3
crates/mesh-llm-cli/src/models.rs
View file

@ -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,

28
crates/mesh-llm-cli/src/parser.rs
View file

@ -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:?}"),
}
}
}

172
crates/mesh-llm-commands/src/model_package.rs
View file

@ -3,7 +3,7 @@ use tokio_stream::StreamExt;
use ::model_package::jobs::HfJobsClient;
use ::model_package::permissions;
use ::model_package::prepare::{self, DiscoveredQuant, PrepareParams};
use ::model_package::prepare::{self, DiscoveredQuant, PrepareJob, PrepareParams};
use ::model_package::script;
use serde_json::json;
@ -16,6 +16,7 @@ pub struct ModelPrepareArgs<'a> {
pub flavor: &'a str,
pub timeout: &'a str,
pub mesh_llm_ref: &'a str,
pub experimental: bool,
pub dry_run: bool,
pub confirm: bool,
pub follow: bool,
@ -37,6 +38,7 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> {
flavor,
timeout,
mesh_llm_ref,
experimental,
dry_run,
confirm,
follow,
@ -95,7 +97,13 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> {
// If no quant specified, list available quants and exit.
// This path doesn't need HF_TOKEN — works for public repos.
if source_quant.is_none() {
return run_list_quants(&hf_client, source_repo, json).await;
return run_list_quants(
&hf_client,
source_repo,
source_model_ref.revision.as_deref(),
json,
)
.await;
}
let submitting = confirm && !dry_run;
@ -116,12 +124,14 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> {
eprintln!("🔍 Resolving source...");
let params = PrepareParams {
source_repo: source_repo.to_string(),
source_revision: source_model_ref.revision.clone(),
quant: source_quant.map(|s| s.to_string()),
target: target.map(|s| s.to_string()),
model_id: model_id.map(|s| s.to_string()),
flavor: flavor.to_string(),
timeout_seconds,
mesh_llm_ref: mesh_llm_ref.to_string(),
experimental,
hf_token: jobs_client
.as_ref()
.map(|client| client.token().to_string()),
@ -129,58 +139,7 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> {
let job = prepare::resolve(&hf_client, params, &perms).await?;
// Print resolved info.
let shard_info = model_ref::split_gguf_shard_info(&job.source_file);
let shard_str = if let Some(shard) = shard_info {
format!(" ({} shards)", shard.total)
} else {
String::new()
};
eprintln!(" Repo: {}", job.source_repo);
eprintln!(" File: {}{}", job.source_file, shard_str);
eprintln!();
eprintln!(
"🔑 Permissions: {} ({})",
perms.username,
if perms.is_meshllm_member {
"meshllm org member"
} else {
"not in meshllm org"
}
);
eprintln!(" Target: {}", job.target_repo);
eprintln!(
" Catalog: meshllm/catalog ({})",
if job.catalog_create_pr {
"will open PR"
} else {
"direct commit"
}
);
eprintln!();
eprintln!(
"📋 Job: {}, timeout {}, mesh-llm@{}",
job.spec.flavor,
format_timeout(job.spec.timeout_seconds),
job.spec
.environment
.get("MESH_LLM_REF")
.map(|s| s.as_str())
.unwrap_or("main")
);
eprintln!(
" Hardware: {} {} ({})",
job.job_plan.pretty_name,
hardware_label(job.job_plan.cpu.as_deref(), job.job_plan.ram.as_deref()),
job.job_plan.selection_reason
);
eprintln!(
" Pricing: ${:.6}/{}, max {}",
job.job_plan.unit_cost_usd,
job.job_plan.unit_label,
format_cost(job.job_plan.max_cost_usd)
);
print_prepare_job(&job, &perms);
if !submitting {
let redacted = redacted_spec(&job.spec);
@ -191,9 +150,12 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> {
"dryRun": true,
"confirmRequired": true,
"sourceRepo": job.source_repo,
"sourceRevision": job.source_revision,
"sourceFile": job.source_file,
"projectors": job.projectors,
"targetRepo": job.target_repo,
"modelId": job.model_id,
"experimental": job.experimental,
"jobPlan": job.job_plan,
"spec": redacted,
}))?
@ -232,9 +194,12 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> {
"jobUrl": job_url,
"namespace": job.namespace,
"sourceRepo": job.source_repo,
"sourceRevision": job.source_revision,
"sourceFile": job.source_file,
"projectors": job.projectors,
"targetRepo": job.target_repo,
"modelId": job.model_id,
"experimental": job.experimental,
"jobPlan": job.job_plan,
}))?
);
@ -251,19 +216,89 @@ pub async fn dispatch_model_package(args: ModelPrepareArgs<'_>) -> Result<()> {
Ok(())
}
fn print_prepare_job(job: &PrepareJob, perms: &permissions::PermissionCheck) {
let shard_info = model_ref::split_gguf_shard_info(&job.source_file);
let shard_str = if let Some(shard) = shard_info {
format!(" ({} shards)", shard.total)
} else {
String::new()
};
eprintln!(" Repo: {}", job.source_repo);
eprintln!(" Commit: {}", job.source_revision);
eprintln!(" File: {}{}", job.source_file, shard_str);
for projector in &job.projectors {
eprintln!(" MMProj: {}", projector.path);
}
eprintln!();
eprintln!(
"🔑 Permissions: {} ({})",
perms.username,
if perms.is_meshllm_member {
"meshllm org member"
} else {
"not in meshllm org"
}
);
eprintln!(" Target: {}", job.target_repo);
eprintln!(
" Release: {}",
if job.experimental {
"experimental (public, not cataloged until HF PR merge)"
} else {
"stable"
}
);
eprintln!(
" Catalog: meshllm/catalog ({})",
if job.catalog_create_pr {
"will open PR"
} else {
"direct commit"
}
);
eprintln!();
eprintln!(
"📋 Job: {}, timeout {}, mesh-llm@{}",
job.spec.flavor,
format_timeout(job.spec.timeout_seconds),
job.spec
.environment
.get("MESH_LLM_REF")
.map(|s| s.as_str())
.unwrap_or("main")
);
eprintln!(
" Hardware: {} {} ({})",
job.job_plan.pretty_name,
hardware_label(job.job_plan.cpu.as_deref(), job.job_plan.ram.as_deref()),
job.job_plan.selection_reason
);
eprintln!(
" Pricing: ${:.6}/{}, max {}",
job.job_plan.unit_cost_usd,
job.job_plan.unit_label,
format_cost(job.job_plan.max_cost_usd)
);
}
async fn run_list_quants(
client: &hf_hub::HFClient,
source_repo: &str,
source_revision: Option<&str>,
json_output: bool,
) -> Result<()> {
let quants = prepare::list_quants(client, source_repo).await?;
let inventory = prepare::list_inventory(client, source_repo, source_revision).await?;
let quants = inventory.quants;
if json_output {
println!(
"{}",
serde_json::to_string_pretty(&json!({
"sourceRepo": source_repo,
"sourceRevision": source_revision,
"quants": quants,
"projectors": inventory.projectors,
}))?
);
return Ok(());
@ -280,13 +315,20 @@ async fn run_list_quants(
eprintln!();
eprintln!("Specify one as a model ref, e.g.:");
eprintln!(
" mesh-llm models package {}:{}",
source_repo, quants[0].name
" mesh-llm models package {}",
source_quant_ref(source_repo, source_revision, &quants[0].name)
);
Ok(())
}
fn source_quant_ref(source_repo: &str, source_revision: Option<&str>, quant: &str) -> String {
source_revision.map_or_else(
|| format!("{source_repo}:{quant}"),
|revision| format!("{source_repo}@{revision}:{quant}"),
)
}
fn print_quant_table(quants: &[DiscoveredQuant]) {
// Find the longest name for alignment.
let max_name = quants.iter().map(|q| q.name.len()).max().unwrap_or(0);
@ -639,4 +681,20 @@ mod tests {
fn parse_timeout_mixed() {
assert_eq!(parse_timeout("1h30m45s").unwrap(), 5445);
}
#[test]
fn source_quant_ref_preserves_revision() {
assert_eq!(
source_quant_ref("poolside/Laguna-S-2.1-GGUF", Some("abc123"), "Q4_K_M"),
"poolside/Laguna-S-2.1-GGUF@abc123:Q4_K_M"
);
}
#[test]
fn source_quant_ref_omits_absent_revision() {
assert_eq!(
source_quant_ref("poolside/Laguna-S-2.1-GGUF", None, "Q4_K_M"),
"poolside/Laguna-S-2.1-GGUF:Q4_K_M"
);
}
}

32
crates/mesh-llm-config/src/model_validation.rs
View file

@ -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(

4
crates/mesh-llm-host-runtime/src/api/split_readiness.rs
View file

@ -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())))
);
}

2
crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs
View file

@ -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
}

137
crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs
View file

@ -9,6 +9,7 @@ use crate::models::gguf::{GgufCompactMeta, scan_gguf_compact_meta};
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct FamilyPolicy {
pub(crate) activation_wire_dtype: StageWireDType,
pub(crate) default_kv_cache_type: Option<&'static str>,
pub(crate) prefix_cache: FamilyPrefixCachePolicy,
}
@ -50,6 +51,24 @@ impl FamilyPolicy {
pub(crate) fn stage_kv_cache_config_for_stage(
&self,
config: &StageConfig,
) -> Option<StageKvCacheConfig> {
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<StageKvCacheConfig> {
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<StageKvCacheConfig> {
match self.prefix_cache {
FamilyPrefixCachePolicy::Disabled { .. } => None,
@ -58,7 +77,14 @@ impl FamilyPolicy {
min_tokens,
max_entries,
} => {
let max_bytes = derive_stage_cache_max_bytes(config)?;
// Layer-package configs can be resolved before their GGUF
// parts are materialized, so there may be no scannable model
// metadata here yet. Keep the certified family cache enabled
// in that case: the resident cache still enforces its
// ctx-derived token budget, while zero means no additional
// byte cap. Disabling the cache entirely made every packaged
// model silently miss the family default.
let max_bytes = derive_stage_cache_max_bytes(config, package_meta).unwrap_or(0);
// The family policy's `max_entries` is a generous
// upper bound on cache cardinality. The real ceiling
// is the unified KV cell pool size: each resident
@ -135,6 +161,18 @@ pub(crate) fn family_policy_for_stage_config(config: &StageConfig) -> FamilyPoli
.unwrap_or_else(|| family_policy_for_model_id(&config.model_id))
}
/// Family policy derived from already-scanned GGUF metadata.
///
/// Split topology planning uses this so it applies the same family K/V
/// defaults that stage loading will apply, instead of re-deriving a
/// size-tiered guess that can badly under-estimate the KV budget.
pub(crate) fn family_policy_for_compact_meta(
meta: &GgufCompactMeta,
model_id: Option<&str>,
) -> FamilyPolicy {
family_policy_for_gguf_meta(meta, model_id)
}
pub(crate) fn family_policy_for_model_path(
path: impl AsRef<Path>,
model_id: Option<&str>,
@ -159,10 +197,14 @@ fn family_policy_for_gguf_meta(meta: &GgufCompactMeta, model_id: Option<&str>) -
}
fn family_policy_for_capability(capability: &FamilyCapabilityRecord) -> FamilyPolicy {
family_policy_for_normalized_family_id(
let mut policy = family_policy_for_normalized_family_id(
capability.family_id.as_str(),
wire_dtype_from_capability(capability.default_wire_dtype),
)
);
if capability.family_id == "inkling" {
policy.default_kv_cache_type = Some("q4_0");
}
policy
}
fn family_policy_for_model_id(model_id: &str) -> FamilyPolicy {
@ -239,6 +281,7 @@ fn family_policy_for_normalized_family_id(
fn resident_kv_policy(activation_wire_dtype: StageWireDType) -> FamilyPolicy {
FamilyPolicy {
activation_wire_dtype,
default_kv_cache_type: None,
prefix_cache: FamilyPrefixCachePolicy::Auto {
payload: FamilyPrefixCachePayload::ResidentKv,
min_tokens: 256,
@ -276,6 +319,7 @@ fn resident_kv_policy(activation_wire_dtype: StageWireDType) -> FamilyPolicy {
fn kv_recurrent_policy(activation_wire_dtype: StageWireDType) -> FamilyPolicy {
FamilyPolicy {
activation_wire_dtype,
default_kv_cache_type: None,
prefix_cache: FamilyPrefixCachePolicy::Auto {
payload: FamilyPrefixCachePayload::KvRecurrent,
min_tokens: 256,
@ -304,6 +348,7 @@ fn disabled_family_policy(
) -> FamilyPolicy {
FamilyPolicy {
activation_wire_dtype,
default_kv_cache_type: None,
prefix_cache: FamilyPrefixCachePolicy::Disabled { reason },
}
}
@ -316,7 +361,16 @@ fn wire_dtype_from_capability(dtype: WireDType) -> StageWireDType {
}
}
fn derive_stage_cache_max_bytes(config: &StageConfig) -> Option<u64> {
fn derive_stage_cache_max_bytes(
config: &StageConfig,
package_meta: Option<&GgufCompactMeta>,
) -> Option<u64> {
if let Some(max_bytes) =
package_meta.and_then(|meta| estimate_stage_cache_max_bytes(config, meta))
{
return Some(max_bytes);
}
[
config.materialized_path.as_deref(),
config.source_model_path.as_deref(),
@ -324,10 +378,15 @@ fn derive_stage_cache_max_bytes(config: &StageConfig) -> Option<u64> {
]
.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<GgufCompactMeta> {
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<u64> {
let stage_layers = config.layer_end.checked_sub(config.layer_start)?;
if stage_layers == 0 {
@ -403,6 +462,8 @@ fn ggml_block_bytes(elements: u64, block_size: u64, type_size: u64) -> Option<u6
#[cfg(test)]
mod tests {
use std::fs;
use super::*;
use skippy_protocol::{FlashAttentionType, LoadMode};
use skippy_topology::{STAGE_RUNTIME_LLAMA_FAMILY_EXPECTATIONS, reviewed_capability_records};
@ -468,6 +529,42 @@ mod tests {
}
}
fn push_gguf_string(bytes: &mut Vec<u8>, 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<u8>, 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<u8>, key: &str, value: &str) {
push_gguf_string(bytes, key);
bytes.extend_from_slice(&8u32.to_le_bytes());
push_gguf_string(bytes, value);
}
fn write_package_metadata(package_dir: &Path) {
let mut bytes = Vec::new();
bytes.extend_from_slice(b"GGUF");
bytes.extend_from_slice(&2u32.to_le_bytes());
bytes.extend_from_slice(&0i64.to_le_bytes());
bytes.extend_from_slice(&8i64.to_le_bytes());
push_gguf_string_kv(&mut bytes, "general.architecture", "llama");
push_gguf_u32(&mut bytes, "llama.context_length", 8192);
push_gguf_u32(&mut bytes, "llama.embedding_length", 4096);
push_gguf_u32(&mut bytes, "llama.block_count", 32);
push_gguf_u32(&mut bytes, "llama.attention.head_count", 32);
push_gguf_u32(&mut bytes, "llama.attention.head_count_kv", 8);
push_gguf_u32(&mut bytes, "llama.attention.key_length", 128);
push_gguf_u32(&mut bytes, "llama.attention.value_length", 128);
let shared_dir = package_dir.join("shared");
fs::create_dir_all(&shared_dir).expect("create package shared directory");
fs::write(shared_dir.join("metadata.gguf"), bytes).expect("write package metadata");
}
#[test]
fn qwen_policy_comes_from_gguf_architecture() {
let policy = family_policy_for_gguf_meta(&meta("qwen3"), None);
@ -608,7 +705,7 @@ mod tests {
| "minicpm3" | "plamo" | "plamo3" | "plm" | "refact" | "smallthinker"
| "smollm3" | "arcee" | "chatglm" | "codeshell" | "deci" | "xverse" | "apertus"
| "bitnet" | "command_r" | "starcoder" | "ernie4_5" | "ernie4_5_moe" | "qwen"
| "jais" | "jais2" | "nemotron" | "llama4" | "mistral4" | "seed_oss" => {
| "jais" | "jais2" | "nemotron" | "llama4" | "mistral4" | "seed_oss" | "laguna" => {
assert_eq!(
policy.prefix_cache,
FamilyPrefixCachePolicy::Auto {
@ -621,7 +718,7 @@ mod tests {
}
"qwen3next" | "falcon_h1" | "jamba" | "lfm2" | "mamba" | "mamba2" | "rwkv6"
| "rwkv7" | "granite_hybrid" | "qwen35" | "qwen35moe" | "plamo2" | "nemotron_h"
| "nemotron_h_moe" | "lfm2moe" | "kimi_linear" => assert_eq!(
| "nemotron_h_moe" | "lfm2moe" | "kimi_linear" | "inkling" => assert_eq!(
policy.prefix_cache,
FamilyPrefixCachePolicy::Auto {
payload: FamilyPrefixCachePayload::KvRecurrent,
@ -768,4 +865,30 @@ mod tests {
assert!(estimate_stage_cache_max_bytes(&config, &kv_meta()).is_none());
}
#[test]
fn package_metadata_enables_cache_for_remote_package_paths() {
let package_dir = tempfile::tempdir().expect("package directory");
write_package_metadata(package_dir.path());
let mut config = stage_config();
config.materialized_path = None;
config.source_model_path = Some("/source/not-downloaded/model.gguf".to_string());
config.model_path = Some("hf://mesh-llm/laguna-layers".to_string());
let policy = FamilyPolicy {
activation_wire_dtype: StageWireDType::F16,
default_kv_cache_type: None,
prefix_cache: FamilyPrefixCachePolicy::Auto {
payload: FamilyPrefixCachePayload::ResidentKv,
min_tokens: 256,
max_entries: 16,
},
};
let cache = policy
.stage_kv_cache_config_for_package(&config, package_dir.path())
.expect("package metadata should provide the cache byte budget");
assert_eq!(cache.payload, StageKvCachePayload::ResidentKv);
assert_eq!(cache.max_bytes, 3_211_264);
}
}

4
crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs
View file

@ -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::{

14
crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs
View file

@ -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)

48
crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs
View file

@ -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<ResolvedModelFitConfig> {
let kv = resolve_kv_defaults(context, kv_policy);
let throughput = resolve_throughput_defaults(context);
@ -188,8 +188,8 @@ fn resolve_model_fit_config(
.and_then(|defaults| defaults.ubatch),
BUILTIN_UBATCH,
);
let cache_type_k = resolve_cache_type_k(context, &kv, kv_policy);
let cache_type_v = resolve_cache_type_v(context, &kv, kv_policy);
let cache_type_k = resolve_cache_type_k(context, &kv, kv_policy, family_policy);
let cache_type_v = resolve_cache_type_v(context, &kv, kv_policy, family_policy);
let kv_offload = resolve_kv_offload(context, &kv);
let flash_attention = context
.model_fit
@ -237,11 +237,25 @@ fn resolve_cache_type_k(
context: &ResolverContext<'_>,
kv: &KvDefaults,
kv_policy: KvCachePolicy,
family_policy: &super::super::family_policy::FamilyPolicy,
) -> String {
if let Some(explicit) = context
.model_fit
.and_then(|fit| non_auto_string(fit.cache_type_k.as_deref()))
{
return explicit.to_string();
}
if let Some(family_default) = family_policy.default_kv_cache_type {
if let Some(explicit) = context
.global_model_fit
.and_then(|fit| non_auto_string(fit.cache_type_k.as_deref()))
{
return explicit.to_string();
}
return family_default.to_string();
}
resolve_field_string(
context
.model_fit
.and_then(|fit| non_auto_string(fit.cache_type_k.as_deref())),
None,
kv.model_macro
.as_ref()
.and_then(|defaults| defaults.cache_type_k.as_deref()),
@ -259,11 +273,25 @@ fn resolve_cache_type_v(
context: &ResolverContext<'_>,
kv: &KvDefaults,
kv_policy: KvCachePolicy,
family_policy: &super::super::family_policy::FamilyPolicy,
) -> String {
if let Some(explicit) = context
.model_fit
.and_then(|fit| non_auto_string(fit.cache_type_v.as_deref()))
{
return explicit.to_string();
}
if let Some(family_default) = family_policy.default_kv_cache_type {
if let Some(explicit) = context
.global_model_fit
.and_then(|fit| non_auto_string(fit.cache_type_v.as_deref()))
{
return explicit.to_string();
}
return family_default.to_string();
}
resolve_field_string(
context
.model_fit
.and_then(|fit| non_auto_string(fit.cache_type_v.as_deref())),
None,
kv.model_macro
.as_ref()
.and_then(|defaults| defaults.cache_type_v.as_deref()),

2
crates/mesh-llm-host-runtime/src/inference/skippy/resolver/speculative.rs
View file

@ -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,
}
}

133
crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs
View file

@ -714,6 +714,102 @@ fn family_policy_beats_builtin_wire_dtype_when_config_is_unset() {
assert_eq!(resolved.skippy.activation_wire_dtype, StageWireDType::F32);
}
#[test]
fn inkling_family_defaults_to_f32_wire_and_q4_kv() {
let resolved = resolve_skippy_config(SkippyConfigResolveRequest {
mesh_config: &MeshConfig::default(),
model_id: "meshllm/inkling-UD-Q2_K_XL-layers",
model_path: Path::new("/models/inkling.gguf"),
model_bytes: 316 * 1024 * 1024 * 1024,
allocatable_memory_bytes: None,
request_defaults: None,
package_generation: None,
})
.unwrap();
assert_eq!(resolved.skippy.activation_wire_dtype, StageWireDType::F32);
assert_eq!(resolved.model_fit.cache_type_k, "q4_0");
assert_eq!(resolved.model_fit.cache_type_v, "q4_0");
}
#[test]
fn inkling_family_kv_default_beats_generic_saver_macro() {
let mesh_config = parse_config(
r#"
[[models]]
model = "meshllm/inkling-UD-Q2_K_XL-layers"
[models.model_fit]
kv_cache_policy = "saver"
"#,
);
let resolved = resolve_skippy_config(SkippyConfigResolveRequest {
mesh_config: &mesh_config,
model_id: "meshllm/inkling-UD-Q2_K_XL-layers",
model_path: Path::new("/models/inkling.gguf"),
model_bytes: 316 * 1024 * 1024 * 1024,
allocatable_memory_bytes: None,
request_defaults: None,
package_generation: None,
})
.unwrap();
assert_eq!(resolved.model_fit.cache_type_k, "q4_0");
assert_eq!(resolved.model_fit.cache_type_v, "q4_0");
}
#[test]
fn explicit_inkling_kv_override_beats_family_default() {
let mesh_config = parse_config(
r#"
[[models]]
model = "meshllm/inkling-UD-Q2_K_XL-layers"
cache_type_k = "q8_0"
cache_type_v = "q8_0"
"#,
);
let resolved = resolve_skippy_config(SkippyConfigResolveRequest {
mesh_config: &mesh_config,
model_id: "meshllm/inkling-UD-Q2_K_XL-layers",
model_path: Path::new("/models/inkling.gguf"),
model_bytes: 316 * 1024 * 1024 * 1024,
allocatable_memory_bytes: None,
request_defaults: None,
package_generation: None,
})
.unwrap();
assert_eq!(resolved.model_fit.cache_type_k, "q8_0");
assert_eq!(resolved.model_fit.cache_type_v, "q8_0");
}
#[test]
fn explicit_global_inkling_kv_override_beats_family_default() {
let mesh_config = parse_config(
r#"
[defaults.model_fit]
cache_type_k = "q8_0"
cache_type_v = "q8_0"
[[models]]
model = "meshllm/inkling-UD-Q2_K_XL-layers"
"#,
);
let resolved = resolve_skippy_config(SkippyConfigResolveRequest {
mesh_config: &mesh_config,
model_id: "meshllm/inkling-UD-Q2_K_XL-layers",
model_path: Path::new("/models/inkling.gguf"),
model_bytes: 316 * 1024 * 1024 * 1024,
allocatable_memory_bytes: None,
request_defaults: None,
package_generation: None,
})
.unwrap();
assert_eq!(resolved.model_fit.cache_type_k, "q8_0");
assert_eq!(resolved.model_fit.cache_type_v, "q8_0");
}
#[test]
fn family_policy_wires_prefix_cache_by_default_for_supported_models() {
let model_file = temp_model_file();
@ -870,6 +966,43 @@ fn layer_package_translation_does_not_treat_hf_ref_as_direct_gguf() {
assert_eq!(options.config.load_mode, LoadMode::LayerPackage);
assert_eq!(options.config.model_path.as_deref(), Some(package_ref));
let kv_cache = options
.config
.kv_cache
.expect("packaged supported family should retain its cache policy");
assert_eq!(kv_cache.mode, StageKvCacheMode::LookupRecord);
assert_eq!(kv_cache.payload, StageKvCachePayload::ResidentKv);
assert_eq!(kv_cache.max_bytes, 0);
}
#[test]
fn inkling_layer_package_retains_recurrent_cache_policy_before_materialization() {
let config = MeshConfig::default();
let package_ref = "hf://meshllm/inkling-UD-Q2_K_XL-layers";
let resolved = resolve_skippy_config(SkippyConfigResolveRequest {
mesh_config: &config,
model_id: "meshllm/inkling-UD-Q2_K_XL-layers",
model_path: Path::new(package_ref),
model_bytes: 316 * 1024 * 1024 * 1024,
allocatable_memory_bytes: None,
request_defaults: None,
package_generation: None,
})
.expect("Inkling package config should resolve");
let stage = resolved
.to_stage_config(Some(fake_hf_package_identity(66)), LoadMode::LayerPackage)
.expect("Inkling package stage config should build");
let kv_cache = stage
.kv_cache
.expect("Inkling package should retain its recurrent cache policy");
assert_eq!(kv_cache.mode, StageKvCacheMode::LookupRecord);
assert_eq!(kv_cache.payload, StageKvCachePayload::KvRecurrent);
assert!(kv_cache.max_entries > 0);
assert!(kv_cache.max_entries <= 16);
assert_eq!(kv_cache.max_bytes, 0);
assert_eq!(kv_cache.min_tokens, 256);
}
#[test]

8
crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs
View file

@ -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)
}

145
crates/mesh-llm-host-runtime/src/mesh/capacity.rs Normal file
View file

@ -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<f64>) -> 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<f64>) -> 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<u64>, 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);
}
}

37
crates/mesh-llm-host-runtime/src/mesh/connections.rs
View file

@ -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<String>,
pub(crate) hostname: Option<String>,
pub(crate) is_soc: Option<bool>,
@ -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<tokio::sync::Semaphore> {
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 {

22
crates/mesh-llm-host-runtime/src/mesh/direct_path.rs
View file

@ -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"
);
}

8
crates/mesh-llm-host-runtime/src/mesh/mod.rs
View file

@ -5,10 +5,13 @@
//! mixed-version compatibility. Skippy activation transport remains on the
//! latency-sensitive `skippy-stage/2` ALPN.
#[cfg(test)]
pub(crate) use mesh_llm_types::mesh::MAX_SPLIT_RTT_MS;
pub use mesh_llm_types::mesh::{
MAX_SPLIT_RTT_MS, ModelDemand, ModelRuntimeDescriptor, ModelSourceKind, ServedModelDescriptor,
ModelDemand, ModelRuntimeDescriptor, ModelSourceKind, ServedModelDescriptor,
ServedModelIdentity, ServedModelMetadata, infer_available_model_descriptors,
infer_local_served_model_descriptor, infer_served_model_descriptors,
infer_local_served_model_descriptor, infer_served_model_descriptors, max_split_rtt_ms,
split_allow_relay_paths,
};
use anyhow::{Context, Result};
@ -76,6 +79,7 @@ pub(crate) fn elapsed_ms_u64(duration: std::time::Duration) -> u64 {
}
mod artifact_transfer_io;
mod capacity;
mod connection_reservation;
mod connections;
mod direct_path;

241
crates/mesh-llm-host-runtime/src/mesh/node.rs
View file

@ -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<f64>) -> u64 {
let mut detected = crate::system::hardware::survey().vram_bytes;
if let Some(cap) = max_vram_gb {
let cap_bytes = (cap * 1e9) as u64;
if cap_bytes < detected {
detected = cap_bytes;
}
}
detected
}
mod startup;
pub use startup::detect_vram_bytes_capped;
use startup::{
bind_mesh_endpoint, init_owner_runtime, startup_secret_key, wait_for_endpoint_online,
};
pub(crate) use startup::{default_plugin_event_source, hardware_snapshot_for_start};
/// Lightweight routing table for passive nodes (clients + standby GPU).
/// Contains just enough info to route requests to the right host.
@ -33,218 +29,6 @@ pub struct RouteEntry {
pub vram_gb: f64,
}
pub(crate) async fn startup_secret_key(role: &NodeRole) -> Result<SecretKey> {
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<iroh::endpoint::RelayMode> {
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<Endpoint> {
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<f64>,
) -> NodeHardwareSnapshot {
let mut vram_bytes = hw.vram_bytes;
let gpu_name = if matches!(role, NodeRole::Client) {
None
} else {
hw.gpu_name
};
let hostname = hw.hostname;
let is_soc = Some(hw.is_soc);
let gpu_vram = (!hw.gpu_vram.is_empty()).then(|| {
hw.gpu_vram
.iter()
.map(|b| b.to_string())
.collect::<Vec<_>>()
.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::<Vec<_>>()
.join(","),
)
};
log_detected_vram(&mut vram_bytes, max_vram_gb);
NodeHardwareSnapshot {
vram_bytes,
gpu_name,
hostname,
is_soc,
gpu_vram,
gpu_reserved_bytes,
}
}
pub(crate) fn detected_vram_log(vram_bytes: u64, max_vram_gb: Option<f64>) -> 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<f64>) {
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<u64>,
) {
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<String>,
) -> Result<OwnerRuntimeInit> {
let trust_store = owner_config
.map(|config| config.trust_store.clone())
.unwrap_or_default();
let trust_policy = owner_config
.map(|config| config.trust_policy)
.unwrap_or_default();
let owner_attestation = match owner_config.and_then(|config| config.keypair.as_ref()) {
Some(keypair) => Some(load_or_refresh_owner_attestation(
keypair,
endpoint_id,
owner_config.and_then(|config| config.node_label.clone()),
hostname,
)?),
None => None,
};
Ok(OwnerRuntimeInit {
trust_store,
trust_policy,
owner_attestation,
})
}
pub(crate) fn default_plugin_event_source(endpoint_id: EndpointId, source_peer_id: &mut String) {
if source_peer_id.is_empty() {
*source_peer_id = endpoint_id_hex(endpoint_id);
}
}
#[derive(Clone)]
pub struct Node {
pub(crate) endpoint: Endpoint,
@ -280,7 +64,10 @@ pub struct Node {
pub(crate) join_targets: Arc<Mutex<Vec<EndpointAddr>>>,
pub(crate) first_joined_mesh_ts: Arc<Mutex<Option<u64>>>,
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<usize>,
pub peer_change_rx: watch::Receiver<usize>,
pub(crate) inflight_requests: Arc<std::sync::atomic::AtomicUsize>,
@ -1011,6 +798,7 @@ impl Node {
std::sync::atomic::AtomicBool::new(false),
)),
vram_bytes: hardware.vram_bytes,
local_runtime_capacity_bytes: hardware.local_runtime_capacity_bytes,
peer_change_tx,
peer_change_rx,
inflight_requests: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
@ -1185,6 +973,7 @@ impl Node {
std::sync::atomic::AtomicBool::new(false),
)),
vram_bytes: 0,
local_runtime_capacity_bytes: 0,
peer_change_tx,
peer_change_rx,
inflight_requests: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
@ -1721,12 +1510,12 @@ impl Node {
};
if let Some(peer) = updated_peer {
tracing::info!("Peer {} RTT: {}ms", id.fmt_short(), rtt_ms);
// If RTT dropped from above the split threshold (80ms) to below it
// If RTT dropped from above the configured split threshold to below it
// (e.g. relay → direct), trigger a re-election so the peer can now
// be included in split mode.
let became_split_eligible = old_rtt
.map(|old| old > MAX_SPLIT_RTT_MS && rtt_ms <= MAX_SPLIT_RTT_MS)
.unwrap_or(rtt_ms <= MAX_SPLIT_RTT_MS);
.map(|old| old > max_split_rtt_ms() && rtt_ms <= max_split_rtt_ms())
.unwrap_or(rtt_ms <= max_split_rtt_ms());
if became_split_eligible {
emit_mesh_info(format!(
"📡 Peer {} RTT improved ({}ms → {}ms) — re-electing for split",

228
crates/mesh-llm-host-runtime/src/mesh/node/startup.rs Normal file
View file

@ -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<f64>) -> 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<SecretKey> {
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<iroh::endpoint::RelayMode> {
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<Endpoint> {
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<f64>,
) -> 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::<Vec<_>>()
.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::<Vec<_>>()
.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<f64>) -> 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<f64>) {
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<u64>,
) {
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<String>,
) -> Result<OwnerRuntimeInit> {
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);
}
}

6
crates/mesh-llm-host-runtime/src/mesh/peer_state.rs
View file

@ -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<PeerInfo> {
self.state
.lock()

10
crates/mesh-llm-host-runtime/src/mesh/stage_artifacts.rs
View file

@ -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<crate::inference::skippy::StageControlResponse> {
// 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)),
}

21
crates/mesh-llm-host-runtime/src/mesh/stage_transport.rs
View file

@ -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<SplitStagePathRejection> {
pub(crate) fn stage_path_rejection(self) -> Option<SplitStagePathRejection> {
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())

1
crates/mesh-llm-host-runtime/src/mesh/tests/connections.rs
View file

@ -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)),

50
crates/mesh-llm-host-runtime/src/mesh/tests/direct_path.rs
View file

@ -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()

60
crates/mesh-llm-host-runtime/src/models/capabilities.rs
View file

@ -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<PathBuf>,
) -> 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<PathBuf>,
) -> 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,
},
);

2
crates/mesh-llm-host-runtime/src/models/mod.rs
View file

@ -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;

4
crates/mesh-llm-host-runtime/src/protocol/convert.rs
View file

@ -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,

4
crates/mesh-llm-host-runtime/src/protocol/tests/announcements.rs
View file

@ -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()

24
crates/mesh-llm-host-runtime/src/runtime/local.rs
View file

@ -531,7 +531,7 @@ pub(super) async fn start_runtime_local_model(
let my_vram = spec
.capacity_budget_bytes
.or_else(|| spec.pinned_gpu.map(|gpu| gpu.allocatable_vram_bytes()))
.unwrap_or_else(|| spec.node.vram_bytes());
.unwrap_or_else(|| spec.node.local_runtime_capacity_bytes());
// For split/layer-package models, compute the local share of model weights
// and the layer fraction so the context planner budgets correctly.
@ -632,9 +632,14 @@ async fn start_runtime_skippy_model(
let capabilities = models::runtime_verified_model_capabilities(
&model_name,
spec.model_path,
models::RuntimeMediaCapabilityEvidence {
vision_projector_loaded: resolved.hardware.projector_path.is_some(),
},
models::runtime_media_capability_evidence(
resolved
.hardware
.projector_path
.as_deref()
.map(PathBuf::from),
)
.await,
);
let embedded_openai = resolved.to_embedded_openai_args(0, false)?;
let mut options = resolved
@ -718,9 +723,14 @@ async fn start_runtime_layer_package_model(
let capabilities = models::runtime_verified_model_capabilities(
&model_name,
spec.model_path,
models::RuntimeMediaCapabilityEvidence {
vision_projector_loaded: resolved.hardware.projector_path.is_some(),
},
models::runtime_media_capability_evidence(
resolved
.hardware
.projector_path
.as_deref()
.map(PathBuf::from),
)
.await,
);
let activation_width = skippy_stage_activation_width(package.activation_width, &model_name)?;
let run_id = format!("mesh-skippy-{}", now_unix_nanos());

183
crates/mesh-llm-host-runtime/src/runtime/local_package.rs
View file

@ -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<u64> {
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<u64>,
timeout: Duration,
) -> Result<SplitParticipantSnapshot> {
let deadline = tokio::time::Instant::now() + timeout;
let mut best: Vec<SplitParticipant> = Vec::new();
let mut best_excluded: Vec<SplitParticipantExclusion> = 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<SplitParticipant>,
best_excluded: &mut Vec<SplitParticipantExclusion>,
) {
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<SplitParticipant>,
excluded: Vec<SplitParticipantExclusion>,
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<SplitParticipant>) {
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,

125
crates/mesh-llm-host-runtime/src/runtime/local_split.rs
View file

@ -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<T> {
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<SplitRuntimeStart> {
let coordinator_start = elect_split_start_coordinator(&spec, model_ref).await?;
let (canonical_coordinator, settled_membership) = match coordinator_start {
CanonicalCoordinatorGate::Coordinator(membership) => (spec.node.id(), membership),
CanonicalCoordinatorGate::Standby { coordinator } => {
return Ok(SplitRuntimeStart::Standby { coordinator });
}
};
let run_id = format!("mesh-split-{}", now_unix_nanos());
let topology_id = format!("topology-{run_id}");
let split_setup =
prepare_split_runtime_start(&spec, model_ref, &topology_id, Duration::from_secs(30))
.await?;
let split_setup = prepare_split_runtime_start(
&spec,
model_ref,
&topology_id,
&settled_membership,
canonical_coordinator,
Duration::from_secs(30),
)
.await?;
let SplitRuntimeStartPreparation {
package,
participant_snapshot,
@ -166,6 +186,12 @@ pub(super) async fn start_runtime_split_model(
let stage0 = stages
.first()
.context("split topology did not produce stage 0")?;
anyhow::ensure!(
stage0.node_id == canonical_coordinator,
"split topology stage 0 {} does not match canonical coordinator {}",
stage0.node_id.fmt_short(),
canonical_coordinator.fmt_short()
);
tracing::info!(
model_ref,
topology_id,
@ -179,21 +205,6 @@ pub(super) async fn start_runtime_split_model(
excluded = ?split_participant_exclusion_labels(&participant_snapshot.excluded),
"split topology planned; elected coordinator from stage 0"
);
if let Some(standby) =
split_runtime_standby_start(spec.node, model_ref, &topology_id, &run_id, stage0)
{
return Ok(standby);
}
tracing::info!(
model_ref,
topology_id,
run_id,
local_node = %spec.node.id().fmt_short(),
context_length = planned_topology.context_length,
parallel_lanes = planned_topology.slots,
"split topology election selected local node as coordinator"
);
let ctx_size = planned_topology.context_length;
let slots = planned_topology.slots;
let projector_path = spec
@ -279,6 +290,8 @@ async fn prepare_split_runtime_start(
spec: &LocalRuntimeModelStartSpec<'_>,
model_ref: &str,
topology_id: &str,
settled_membership: &[SplitParticipant],
canonical_coordinator: iroh::EndpointId,
timeout: Duration,
) -> Result<SplitRuntimeStartPreparation> {
let package = resolve_split_runtime_package(spec.model_path, model_ref).await?;
@ -288,6 +301,7 @@ async fn prepare_split_runtime_start(
model_ref,
&package,
spec.pinned_gpu.map(|gpu| gpu.allocatable_vram_bytes()),
settled_membership,
timeout,
)
.await?;
@ -295,6 +309,7 @@ async fn prepare_split_runtime_start(
let kv_bytes_per_token = split_runtime_kv_bytes_per_token(
&package,
&compact_meta,
model_ref,
spec.cache_type_k_override,
spec.cache_type_v_override,
)?;
@ -313,6 +328,13 @@ async fn prepare_split_runtime_start(
&participant_snapshot.participants,
)
.await?;
anyhow::ensure!(
locked_stages
.first()
.is_some_and(|stage| stage.node_id == canonical_coordinator),
"split topology lock stage 0 must be canonical coordinator {}",
canonical_coordinator.fmt_short()
);
plan_locked_runtime_slice_topology_with_resources(
topology_id,
model_ref,
@ -323,13 +345,14 @@ async fn prepare_split_runtime_start(
&locked_stages,
)?
} else {
plan_runtime_slice_topology_with_resources(
plan_runtime_slice_topology_with_resources_and_stage0(
topology_id,
model_ref,
&package,
&participant_snapshot.participants,
&participant_snapshot.excluded,
resources,
Some(canonical_coordinator),
)?
};
Ok(SplitRuntimeStartPreparation {
@ -342,27 +365,47 @@ async fn prepare_split_runtime_start(
})
}
fn split_runtime_standby_start(
node: &mesh::Node,
async fn elect_split_start_coordinator(
spec: &LocalRuntimeModelStartSpec<'_>,
model_ref: &str,
topology_id: &str,
run_id: &str,
stage0: &RuntimeSliceStagePlan,
) -> Option<SplitRuntimeStart> {
if stage0.node_id == node.id() {
return None;
) -> Result<CanonicalCoordinatorGate<Vec<SplitParticipant>>> {
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<SplitParticipant>,
) -> Result<CanonicalCoordinatorGate<Vec<SplitParticipant>>> {
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<iroh::EndpointId> {
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};

35
crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs
View file

@ -20,6 +20,10 @@ use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::Duration;
pub(super) const MIN_STAGE_SOURCE_PREPARE_TIMEOUT: Duration = Duration::from_secs(30 * 60);
const STAGE_SOURCE_PREPARE_ALLOWANCE: Duration = Duration::from_secs(10 * 60);
const STAGE_SOURCE_MIN_BYTES_PER_SEC: u64 = 16 * 1024 * 1024;
pub(super) struct SplitGenerationLoadSpec<'a> {
pub(super) node: &'a mesh::Node,
pub(super) mesh_config: &'a plugin::MeshConfig,
@ -176,7 +180,14 @@ pub(super) async fn load_split_runtime_generation_inner(
stage_index: downstream.stage_index,
endpoint: downstream_endpoint,
});
let vision_projector_loaded = runtime_options.config.projector_path.is_some();
let media_capability_evidence = models::runtime_media_capability_evidence(
runtime_options
.config
.projector_path
.as_deref()
.map(std::path::PathBuf::from),
)
.await;
let node_for_hook = spec.node.clone();
let model_ref = spec.model_ref.to_string();
let reporter_model_ref = model_ref.clone();
@ -209,9 +220,7 @@ pub(super) async fn load_split_runtime_generation_inner(
let capabilities = models::runtime_verified_model_capabilities(
spec.model_ref,
spec.model_path,
models::RuntimeMediaCapabilityEvidence {
vision_projector_loaded,
},
media_capability_evidence,
);
spec.node
@ -270,7 +279,7 @@ pub(super) async fn load_downstream_split_runtime_stages(
spec.node,
stage.node_id,
&load,
Duration::from_secs(30 * 60),
stage_source_prepare_timeout(spec.package, stage),
)
.await
.with_context(|| {
@ -322,6 +331,22 @@ pub(super) async fn load_downstream_split_runtime_stages(
.context("split topology missing downstream stage")
}
pub(super) fn stage_source_prepare_timeout(
package: &skippy::SkippyPackageIdentity,
stage: &RuntimeSliceStagePlan,
) -> Duration {
let package_layers = u64::from(package.layer_count.max(1));
let stage_layers = u64::from(stage.layer_end.saturating_sub(stage.layer_start).max(1));
let estimated_stage_bytes = package
.source_model_bytes
.saturating_mul(stage_layers)
.div_ceil(package_layers);
let transfer_secs = estimated_stage_bytes.div_ceil(STAGE_SOURCE_MIN_BYTES_PER_SEC);
Duration::from_secs(transfer_secs)
.saturating_add(STAGE_SOURCE_PREPARE_ALLOWANCE)
.max(MIN_STAGE_SOURCE_PREPARE_TIMEOUT)
}
pub(super) fn split_runtime_stage_load_request(
spec: &SplitGenerationLoadSpec<'_>,
settings: &SplitGenerationLoadSettings<'_>,

250
crates/mesh-llm-host-runtime/src/runtime/local_split/tests.rs
View file

@ -36,6 +36,65 @@ fn runtime_local_targets_keep_duplicate_same_model_ports() {
);
}
#[test]
fn canonical_coordinator_is_identical_with_divergent_observer_signals() {
let capacities = [
(1, 24_000_000_000),
(2, 48_000_000_000),
(3, 48_000_000_000),
];
let observer_a = capacities
.into_iter()
.map(|(seed, capacity)| {
SplitParticipant::new(make_id(seed), capacity, None).with_package_signals(
SplitParticipantPackageSignal {
cached_slice_bytes: u64::from(seed) * 10_000,
missing_artifact_bytes: u64::from(4 - seed) * 20_000,
availability_score: u32::from(seed),
},
Some(u32::from(seed) * 40),
true,
)
})
.collect::<Vec<_>>();
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::<Vec<_>>();
let expected = [make_id(2), make_id(3)].into_iter().min().unwrap();
assert_eq!(canonical_split_coordinator(&observer_a), Some(expected));
assert_eq!(canonical_split_coordinator(&observer_b), Some(expected));
}
#[test]
fn noncanonical_gate_returns_standby_without_invoking_package_planning() {
let local = SplitParticipant::new(make_id(1), 24_000_000_000, None);
let coordinator = SplitParticipant::new(make_id(2), 48_000_000_000, None);
let gate = canonical_coordinator_gate(local.node_id, vec![local, coordinator])
.expect("canonical coordinator gate");
match gate {
CanonicalCoordinatorGate::Standby {
coordinator: selected,
} => assert_eq!(selected, coordinator.node_id),
CanonicalCoordinatorGate::Coordinator(_) => {
panic!("local node must not be elected coordinator")
}
}
}
#[test]
fn split_topology_planner_uses_all_eligible_participants() {
let participants = vec![
@ -67,6 +126,57 @@ fn split_topology_planner_uses_all_eligible_participants() {
assert_eq!(stages.last().unwrap().layer_end, 40);
}
#[test]
fn resource_planner_keeps_canonical_coordinator_at_stage_zero() {
let canonical = SplitParticipant::new(make_id(1), 48_000_000_000, None).with_package_signals(
SplitParticipantPackageSignal {
cached_slice_bytes: 0,
missing_artifact_bytes: 40_000_000,
availability_score: 0,
},
Some(200),
true,
);
let fast_a = SplitParticipant::new(make_id(2), 32_000_000_000, None).with_package_signals(
SplitParticipantPackageSignal {
cached_slice_bytes: 40_000_000,
missing_artifact_bytes: 0,
availability_score: 40,
},
Some(1),
true,
);
let fast_b = SplitParticipant::new(make_id(3), 32_000_000_000, None).with_package_signals(
SplitParticipantPackageSignal {
cached_slice_bytes: 40_000_000,
missing_artifact_bytes: 0,
availability_score: 40,
},
Some(1),
true,
);
let participants = [canonical, fast_a, fast_b];
let package = package(40);
let planned = plan_runtime_slice_topology_with_resources_and_stage0(
"topology-test",
"test-model",
&package,
&participants,
&[],
SplitTopologyResourceInputs {
native_context_length: 65_536,
kv_bytes_per_token: 64 * 1024,
ctx_size_override: Some(65_536),
parallel_override: Some(1),
},
Some(canonical.node_id),
)
.expect("canonical stage-zero topology");
assert_eq!(planned.stages.first().unwrap().node_id, canonical.node_id);
}
#[test]
fn split_topology_planner_prefers_cached_participant_in_runtime_path() {
let cold = SplitParticipant::new(make_id(1), 24_000_000_000, None).with_package_signals(
@ -502,6 +612,38 @@ fn split_startup_error_messages_include_specific_blocker_tokens() {
assert!(timeout.contains("30s"));
}
#[test]
fn stage_source_prepare_timeout_scales_with_assigned_package_bytes() {
let package = skippy::SkippyPackageIdentity {
source_model_bytes: 975_000_000_000,
layer_count: 66,
..package(66)
};
let small_stage = RuntimeSliceStagePlan {
stage_id: "stage-2".to_string(),
stage_index: 2,
node_id: make_id(2),
layer_start: 59,
layer_end: 66,
parameter_bytes: 0,
};
let large_stage = RuntimeSliceStagePlan {
stage_id: "stage-0".to_string(),
stage_index: 0,
node_id: make_id(1),
layer_start: 0,
layer_end: 39,
parameter_bytes: 0,
};
let small_timeout = stage_source_prepare_timeout(&package, &small_stage);
let large_timeout = stage_source_prepare_timeout(&package, &large_stage);
assert!(small_timeout > MIN_STAGE_SOURCE_PREPARE_TIMEOUT);
assert!(large_timeout > small_timeout);
assert!(large_timeout > Duration::from_secs(6 * 60 * 60));
}
#[test]
fn startup_runtime_plan_auto_splits_when_model_exceeds_local_capacity() {
assert_eq!(
@ -745,7 +887,7 @@ fn split_participant_signature_includes_vram_for_stability() {
}
#[test]
fn split_participant_signature_includes_package_signals_for_stability() {
fn split_participant_signature_includes_package_signals_for_claim_identity() {
let node_id = make_id(9);
let first = vec![SplitParticipant::new(node_id, 24_000_000_000, None)];
let second = vec![
@ -1424,3 +1566,109 @@ fn split_topology_minimum_rejects_single_stage_split_candidate() {
]));
assert!(!split_stages_meet_minimum(&[stage(1, 0, 0, 40)]));
}
#[test]
fn split_planning_uses_family_kv_defaults_for_inkling() {
let mut meta = crate::models::gguf::GgufCompactMeta {
architecture: "inkling".to_string(),
context_length: 65_536,
embedding_size: 4096,
head_count: 32,
kv_head_count: 8,
layer_count: 66,
key_length: 128,
value_length: 128,
..Default::default()
};
meta.kv_head_counts = vec![8; 66];
// Inkling's reviewed family default keeps both planning and stage loading
// on quantized Q4_0 K/V rather than silently expanding to F16.
let mut identity = package(66);
identity.source_model_bytes = 318 * 1024 * 1024 * 1024;
let planned =
split_runtime_kv_bytes_per_token(&identity, &meta, "tml/inkling-q2", None, None).unwrap();
let expected_q4 = crate::models::gguf::GgufKvCacheQuant::from_llama_args("q4_0", "q4_0")
.unwrap()
.kv_cache_bytes_per_token(&meta)
.unwrap();
assert_eq!(planned, expected_q4);
// Explicit user overrides still win over the family default.
let overridden = split_runtime_kv_bytes_per_token(
&identity,
&meta,
"tml/inkling-q2",
Some("f16"),
Some("f16"),
)
.unwrap();
assert!(overridden > planned);
}
/// Validates the finding-#1 fix against a real Inkling layer package.
///
/// Set `INKLING_METADATA_GGUF` to a package's `shared/metadata.gguf` to run it;
/// skipped otherwise so CI stays hermetic.
#[test]
fn real_inkling_metadata_plans_family_kv_not_size_tiered() {
let Ok(path) = std::env::var("INKLING_METADATA_GGUF") else {
eprintln!("skip: INKLING_METADATA_GGUF not set");
return;
};
let meta = crate::models::gguf::scan_gguf_compact_meta(std::path::Path::new(&path))
.expect("scan real inkling metadata gguf");
eprintln!(
"REAL META arch={} layers={} embed={} kv_heads={} k_len={} v_len={} ctx={}",
meta.architecture,
meta.layer_count,
meta.embedding_size,
meta.kv_head_count,
meta.key_length,
meta.value_length,
meta.context_length
);
let model_ref = "unsloth/inkling-GGUF:UD-Q2_K_XL";
let policy = crate::inference::skippy::family_policy_for_compact_meta(&meta, Some(model_ref));
eprintln!(
"FAMILY default_kv_cache_type={:?}",
policy.default_kv_cache_type
);
let mut identity = package(meta.layer_count);
identity.source_model_bytes = 318 * 1024 * 1024 * 1024;
let planned =
split_runtime_kv_bytes_per_token(&identity, &meta, model_ref, None, None).unwrap();
let expected_q4 = crate::models::gguf::GgufKvCacheQuant::from_llama_args("q4_0", "q4_0")
.unwrap()
.kv_cache_bytes_per_token(&meta)
.unwrap();
let size_tiered = {
let p =
crate::inference::skippy::KvCachePolicy::for_model_size(identity.source_model_bytes);
split_kv_cache_quant(&p, None, None)
.kv_cache_bytes_per_token(&meta)
.unwrap()
};
let ctx = u64::from(meta.context_length.max(1));
eprintln!(
"KV/token planned={planned} size_tiered={size_tiered} ratio={:.2}x | @ctx{ctx}: planned={:.1}GiB size_tiered={:.1}GiB under_budget={:.1}GiB",
planned as f64 / size_tiered.max(1) as f64,
(planned * ctx) as f64 / (1024.0 * 1024.0 * 1024.0),
(size_tiered * ctx) as f64 / (1024.0 * 1024.0 * 1024.0),
((planned - size_tiered.min(planned)) * ctx) as f64 / (1024.0 * 1024.0 * 1024.0),
);
assert_eq!(
policy.default_kv_cache_type,
Some("q4_0"),
"inkling must resolve a q4_0 family K/V default"
);
assert_eq!(
planned, expected_q4,
"family-aware planning must use the Inkling Q4_0 K/V default"
);
}

2
crates/mesh-llm-host-runtime/src/runtime/mod.rs
View file

@ -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)]

2
crates/mesh-llm-host-runtime/src/runtime/model_lifecycle/load.rs
View file

@ -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;

2
crates/mesh-llm-host-runtime/src/runtime/model_lifecycle/reconciliation.rs
View file

@ -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()

394
crates/mesh-llm-host-runtime/src/runtime/split_participant_settle.rs Normal file
View file

@ -0,0 +1,394 @@
use super::local_package::{
SPLIT_DEFAULT_MIN_PARTICIPANTS, SplitParticipant, SplitParticipantSnapshot,
collect_split_participant_membership, collect_split_participants,
ensure_split_participant_timeout_has_quorum,
};
use super::split_planning::{split_participant_exclusion_labels, split_participant_labels};
use crate::inference::skippy;
use crate::mesh;
use anyhow::Result;
use std::time::Duration;
const SPLIT_PARTICIPANT_POLL_INTERVAL: Duration = Duration::from_millis(500);
/// An automatic split cannot know how many nodes the operator intends to start.
/// Wait for additions/capacity changes to stop instead of claiming the first
/// two-node quorum that happens to become visible.
const SPLIT_MEMBERSHIP_SETTLE_DWELL: Duration = Duration::from_secs(8);
/// Even an immediately stable two-node quorum gets a bounded discovery window.
const SPLIT_FIRST_QUORUM_OBSERVATION: Duration = Duration::from_secs(8);
fn membership_settle_timeout(requested: Duration) -> Duration {
requested.max(SPLIT_FIRST_QUORUM_OBSERVATION.max(SPLIT_MEMBERSHIP_SETTLE_DWELL))
}
type SplitMembershipSignature = Vec<(String, u64)>;
#[derive(Debug, Default)]
struct SplitMembershipSettleBarrier {
signature: SplitMembershipSignature,
first_quorum_observed: Option<tokio::time::Instant>,
stable_since: Option<tokio::time::Instant>,
}
impl SplitMembershipSettleBarrier {
fn observe(&mut self, participants: &[SplitParticipant], now: tokio::time::Instant) -> bool {
let signature = split_membership_signature(participants);
if signature != self.signature {
self.signature = signature;
self.stable_since = Some(now);
}
if participants.len() >= SPLIT_DEFAULT_MIN_PARTICIPANTS
&& self.first_quorum_observed.is_none()
{
self.first_quorum_observed = Some(now);
}
self.is_ready(participants.len(), now)
}
fn is_ready(&self, participant_count: usize, now: tokio::time::Instant) -> bool {
if participant_count < SPLIT_DEFAULT_MIN_PARTICIPANTS {
return false;
}
let Some(first_quorum_observed) = self.first_quorum_observed else {
return false;
};
let Some(stable_since) = self.stable_since else {
return false;
};
now.saturating_duration_since(first_quorum_observed) >= SPLIT_FIRST_QUORUM_OBSERVATION
&& now.saturating_duration_since(stable_since) >= SPLIT_MEMBERSHIP_SETTLE_DWELL
}
fn stable_for(&self, now: tokio::time::Instant) -> Duration {
self.stable_since
.map(|stable_since| now.saturating_duration_since(stable_since))
.unwrap_or_default()
}
}
struct SplitMembershipWait<'a> {
node: &'a mesh::Node,
model_name: &'a str,
model_ref: &'a str,
deadline: tokio::time::Instant,
barrier: SplitMembershipSettleBarrier,
last_logged_signature: SplitMembershipSignature,
}
impl<'a> SplitMembershipWait<'a> {
fn new(
node: &'a mesh::Node,
model_name: &'a str,
model_ref: &'a str,
timeout: Duration,
) -> Self {
Self {
node,
model_name,
model_ref,
deadline: tokio::time::Instant::now() + membership_settle_timeout(timeout),
barrier: SplitMembershipSettleBarrier::default(),
last_logged_signature: Vec::new(),
}
}
async fn run(mut self) -> Result<SplitParticipantSnapshot> {
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<SplitParticipantSnapshot> {
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<u64>,
expected_node_ids: Vec<String>,
deadline: tokio::time::Instant,
}
impl<'a> SplitEligibilityWait<'a> {
async fn run(self) -> Result<SplitParticipantSnapshot> {
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<SplitParticipantSnapshot> {
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<SplitParticipantSnapshot> {
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<u64>,
expected_membership: &[SplitParticipant],
timeout: Duration,
) -> Result<SplitParticipantSnapshot> {
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<String> {
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::<Vec<_>>();
assert!(!barrier.observe(&participants, start + Duration::from_secs(seconds)));
}
let participants = (1..=6).map(participant).collect::<Vec<_>>();
assert!(!barrier.observe(&participants, start + Duration::from_secs(14)));
assert!(barrier.observe(&participants, start + Duration::from_secs(15)));
}
#[test]
fn two_node_cohort_settles_after_bounded_discovery_window() {
let start = tokio::time::Instant::now();
let participants = vec![participant(1), participant(2)];
let mut barrier = SplitMembershipSettleBarrier::default();
assert!(!barrier.observe(&participants, start));
assert!(!barrier.observe(&participants, start + Duration::from_secs(7)));
assert!(barrier.observe(&participants, start + Duration::from_secs(8)));
}
#[test]
fn short_caller_timeout_still_allows_the_settle_barrier() {
assert_eq!(
membership_settle_timeout(Duration::from_secs(1)),
Duration::from_secs(8)
);
assert_eq!(
membership_settle_timeout(Duration::from_secs(30)),
Duration::from_secs(30)
);
}
#[test]
fn volatile_package_and_rtt_signals_do_not_reset_membership_dwell() {
let start = tokio::time::Instant::now();
let mut participants = vec![participant(1), participant(2)];
let mut barrier = SplitMembershipSettleBarrier::default();
assert!(!barrier.observe(&participants, start));
participants[0] = participants[0].with_package_signals(
SplitParticipantPackageSignal {
cached_slice_bytes: 10_000,
missing_artifact_bytes: 90_000,
availability_score: 3,
},
Some(80),
true,
);
participants[1] = participants[1].with_package_signals(
SplitParticipantPackageSignal {
cached_slice_bytes: 100_000,
missing_artifact_bytes: 0,
availability_score: 30,
},
Some(4),
true,
);
assert!(barrier.observe(&participants, start + Duration::from_secs(8)));
assert_eq!(
split_membership_signature(&participants),
vec![
(make_id(1).to_string(), 1_000_000_000),
(make_id(2).to_string(), 2_000_000_000),
]
);
assert_ne!(split_participant_signature(&participants), Vec::new());
}
#[test]
fn capacity_change_restarts_membership_dwell() {
let start = tokio::time::Instant::now();
let participants = vec![participant(1), participant(2)];
let mut barrier = SplitMembershipSettleBarrier::default();
assert!(!barrier.observe(&participants, start));
let changed = vec![
participant(1),
SplitParticipant::new(make_id(2), 9_000, None),
];
assert!(!barrier.observe(&changed, start + Duration::from_secs(7)));
assert!(!barrier.observe(&changed, start + Duration::from_secs(14)));
assert!(barrier.observe(&changed, start + Duration::from_secs(15)));
}
}

96
crates/mesh-llm-host-runtime/src/runtime/split_planning.rs
View file

@ -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<SplitTopologyPlan> {
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<SplitTopologyPlan> {
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<PlannedRuntimeSliceTopology> {
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<iroh::EndpointId>,
) -> Result<PlannedRuntimeSliceTopology> {
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<iroh::EndpointId>,
) -> Result<SplitTopologyPlan> {
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))

9
crates/mesh-llm-host-runtime/src/runtime/startup_handles.rs
View file

@ -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);

107
crates/mesh-llm-host-runtime/src/runtime/startup_retry.rs Normal file
View file

@ -0,0 +1,107 @@
pub(super) fn is_retryable_split_start_failure(message: &str) -> bool {
split_participants_are_still_converging(message)
|| split_control_transport_failed(message)
|| split_stage_source_preparation_timed_out(message)
}
fn split_participants_are_still_converging(message: &str) -> bool {
message.contains("at least two participating nodes")
|| message.contains("at least two stage participants")
|| message.contains("split_capacity_shortfall")
|| message.contains("canonical coordinator")
|| (message.contains("split topology lock stage")
&& message.contains("matched 0 eligible nodes"))
}
fn split_control_transport_failed(message: &str) -> bool {
let is_control_operation = message.contains("load split stage")
|| message.contains("prepare split stage")
|| message.contains("stage_control_unreachable");
let is_transport_failure = message.contains("connection lost")
|| message.contains("stream finished early")
|| message.contains("timeout waiting for stage control response");
is_control_operation && is_transport_failure
}
fn split_stage_source_preparation_timed_out(message: &str) -> bool {
message.contains("stage_source_prepare_timeout")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn participant_shortage_is_retryable() {
assert!(is_retryable_split_start_failure(
"split runtime needs at least two participating nodes for model"
));
assert!(is_retryable_split_start_failure(
"split runtime needs at least two stage participants"
));
}
#[test]
fn transient_capacity_shortfall_is_retryable() {
assert!(is_retryable_split_start_failure(
"split_capacity_shortfall: unable to plan split topology: max_placeable_layers_at_evaluated_shape=21/66"
));
}
#[test]
fn missing_locked_participant_is_retryable() {
assert!(is_retryable_split_start_failure(
"split topology lock stage 3 selector \"worker\" matched 0 eligible nodes; available: local"
));
}
#[test]
fn canonical_coordinator_mismatch_is_retryable() {
assert!(is_retryable_split_start_failure(
"split topology stage 0 node-a does not match canonical coordinator node-b"
));
assert!(is_retryable_split_start_failure(
"split topology lock stage 0 must be canonical coordinator node-b"
));
}
#[test]
fn ambiguous_or_invalid_lock_is_not_retryable() {
assert!(!is_retryable_split_start_failure(
"split topology lock stage 3 selector \"worker\" matched 2 eligible nodes"
));
assert!(!is_retryable_split_start_failure(
"split topology lock manifest abc does not match resolved package manifest def"
));
}
#[test]
fn stage_control_transport_failure_is_retryable() {
assert!(is_retryable_split_start_failure(
"load split stage stage-1: connection lost: closed"
));
assert!(is_retryable_split_start_failure(
"prepare split stage stage-2: stage_control_unreachable: stream finished early"
));
assert!(is_retryable_split_start_failure(
"load split stage stage-3: timeout waiting for stage control response"
));
}
#[test]
fn stage_source_preparation_timeout_is_retryable() {
assert!(is_retryable_split_start_failure(
"prepare split stage stage-1: stage_source_prepare_timeout: timed out waiting for stage source availability after 30m"
));
}
#[test]
fn runtime_or_unrelated_transport_failure_is_not_retryable() {
assert!(!is_retryable_split_start_failure(
"load skippy stage 0 runtime: unable to allocate CUDA0 buffer"
));
assert!(!is_retryable_split_start_failure(
"artifact transfer connection lost: closed"
));
}
}

4
crates/mesh-llm-host-runtime/src/sdk.rs
View file

@ -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;

27
crates/mesh-llm-types/src/mesh/mod.rs
View file

@ -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<u32> = 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<bool> = std::sync::OnceLock::new();
*VALUE.get_or_init(|| {
std::env::var("MESH_SPLIT_ALLOW_RELAY")
.map(|raw| raw.trim() == "1")
.unwrap_or(false)
})
}

1
crates/mesh-llm/src/commands/mod.rs
View file

@ -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,

2
crates/mesh-llm/src/commands/models/mod.rs
View file

@ -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,

143
crates/model-artifact/src/gguf.rs
View file

@ -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<Option<Vec<u32>>> {
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<Option<f32>> {
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<Option<bool>> {
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<u32>,
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<bool>,
pub has_audio_encoder: Option<bool>,
}
impl GgufCompactMeta {
pub fn effective_kv_head_count(&self) -> Option<u32> {
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<GgufCompactMeta> {
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<GgufCompactMeta> {
Some(meta)
}
/// Scan the modality flags stored in a multimodal projector GGUF.
pub fn scan_gguf_projector_meta(path: &Path) -> Option<GgufProjectorMeta> {
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<u8>, 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<u8>, 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<u8>, 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::<Vec<_>>();
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();

17
crates/model-artifact/src/gguf/kv_cache.rs
View file

@ -153,9 +153,24 @@ fn cache_bytes_per_token(
vector_length: u32,
cache_type: GgufKvCacheType,
) -> Option<u64> {
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)?

153
crates/model-package/src/bin/queue-unsloth-layer-packages.rs
View file

@ -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<DiscoveredProjector>,
target_repo: String,
model_layer_repos: Vec<String>,
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::<Vec<_>>()
.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<JobSpec> {
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::<Vec<_>>()
.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));
}

2
crates/model-package/src/jobs.rs
View file

@ -41,6 +41,8 @@ pub struct JobVolume {
pub mount_path: String,
#[serde(rename = "readOnly", skip_serializing_if = "Option::is_none")]
pub read_only: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub revision: Option<String>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]

151
crates/model-package/src/prepare.rs
View file

@ -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<String>,
pub quant: Option<String>,
pub target: Option<String>,
pub model_id: Option<String>,
pub flavor: String,
pub timeout_seconds: u64,
pub mesh_llm_ref: String,
pub experimental: bool,
pub hf_token: Option<String>,
}
/// 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<DiscoveredProjector>,
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<DiscoveredQuant>,
pub projectors: Vec<DiscoveredProjector>,
}
/// List all available GGUF quant variants in a HF model repo.
pub async fn list_quants(client: &HFClient, repo: &str) -> Result<Vec<DiscoveredQuant>> {
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<RepoGgufInventory> {
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<Vec<Discovered
}
}
Ok(discover_quants_from_gguf_files(gguf_files))
Ok(discover_inventory_from_gguf_files(gguf_files))
}
/// Separate model GGUF distributions from multimodal projector sidecars.
pub fn discover_inventory_from_gguf_files(gguf_files: Vec<(String, u64)>) -> 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::<Vec<_>>();
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, &params.source_repo).await?;
let (owner, name) = parse_repo(&params.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, &params.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(),
&params.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::<Vec<_>>()
.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(

22
crates/model-package/src/script.rs
View file

@ -129,7 +129,14 @@ mod tests {
#[test]
fn embedded_script_writes_rich_model_card() {
assert!(EMBEDDED_SCRIPT.contains("pipeline_tag: text-generation"));
assert!(EMBEDDED_SCRIPT.contains("SOURCE_PIPELINE_TAG"));
assert!(EMBEDDED_SCRIPT.contains("PACKAGE_EXPERIMENTAL"));
assert!(EMBEDDED_SCRIPT.contains("pipeline_tag: {yaml_quote(source_pipeline_tag)}"));
assert!(EMBEDDED_SCRIPT.contains("- experimental"));
assert!(EMBEDDED_SCRIPT.contains("Experimental package:"));
assert!(EMBEDDED_SCRIPT.contains("resolve_upstream_license"));
assert!(EMBEDDED_SCRIPT.contains("license_frontmatter"));
assert!(EMBEDDED_SCRIPT.contains("could not resolve upstream license metadata"));
assert!(EMBEDDED_SCRIPT.contains("- openai-compatible"));
assert!(EMBEDDED_SCRIPT.contains("## Model Overview"));
assert!(EMBEDDED_SCRIPT.contains("## Highlights"));
@ -198,7 +205,20 @@ mod tests {
assert!(EMBEDDED_SCRIPT.contains(r#"MOUNTED_SOURCE_PATH="/source/${SOURCE_FILE}""#));
assert!(EMBEDDED_SCRIPT.contains(r#"WRITE_PACKAGE_INPUT="$MOUNTED_SOURCE_PATH""#));
assert!(EMBEDDED_SCRIPT.contains(r#"--source-file "$SOURCE_FILE""#));
assert!(EMBEDDED_SCRIPT.contains("SOURCE_PROJECTOR_FILES"));
assert!(
EMBEDDED_SCRIPT
.contains(r#"WRITE_PACKAGE_PROJECTOR_ARGS+=(--projector "$PROJECTOR_PATH")"#)
);
assert!(EMBEDDED_SCRIPT.contains(r#""${WRITE_PACKAGE_PROJECTOR_ARGS[@]}""#));
assert!(EMBEDDED_SCRIPT.contains(r#"time "$SLICER" write-package "$WRITE_PACKAGE_INPUT""#));
assert!(!EMBEDDED_SCRIPT.contains(r#"time $SLICER write-package "$SOURCE_PATH""#));
}
#[test]
fn embedded_script_preserves_catalog_source_revision() {
assert!(EMBEDDED_SCRIPT.contains(r#""source_revision": source_revision"#));
assert!(EMBEDDED_SCRIPT.contains(r#"variants[variant_name]["source"] = source_entry"#));
assert!(EMBEDDED_SCRIPT.contains(r#"existing_variant["source"] = source_entry"#));
}
}

119
crates/model-package/src/scripts/split-model-job.sh
View file

@ -6,8 +6,11 @@ set -euo pipefail
#
# Environment variables (set by mesh-llm model-package job spec):
# SOURCE_REPO, SOURCE_FILE, SOURCE_QUANT, TARGET_REPO, MODEL_ID, SOURCE_REVISION
# SOURCE_PROJECTOR_FILES — optional newline-delimited repo-relative mmproj GGUFs
# SOURCE_PIPELINE_TAG — source model pipeline tag for the published model card
# MESH_LLM_REF — git ref to build from (default: main)
# CATALOG_CREATE_PR — "true" to open a PR for catalog updates (non-org members)
# PACKAGE_EXPERIMENTAL — "true" to label the public package as not runtime-certified
# HF_TOKEN — injected as a secret by HF Jobs
#
# Volumes:
@ -287,6 +290,34 @@ else
WRITE_PACKAGE_IDENTITY_ARGS=()
echo " Source mount: not available; falling back to Hugging Face cache download"
fi
WRITE_PACKAGE_PROJECTOR_ARGS=()
while IFS= read -r PROJECTOR_FILE; do
if [ -z "$PROJECTOR_FILE" ]; then
continue
fi
MOUNTED_PROJECTOR_PATH="/source/${PROJECTOR_FILE}"
if [ -f "$MOUNTED_PROJECTOR_PATH" ]; then
PROJECTOR_PATH="$MOUNTED_PROJECTOR_PATH"
else
echo " Projector mount missing; downloading ${PROJECTOR_FILE} at ${SOURCE_REVISION}"
PROJECTOR_PATH="$("$VENV_DIR/bin/python3" - "$PROJECTOR_FILE" <<'PYTHON'
from huggingface_hub import hf_hub_download
import os
import sys
print(hf_hub_download(
repo_id=os.environ["SOURCE_REPO"],
filename=sys.argv[1],
revision=os.environ["SOURCE_REVISION"],
cache_dir=os.environ["HF_HUB_CACHE"],
token=os.environ.get("HF_TOKEN"),
))
PYTHON
)"
fi
echo " Projector: $PROJECTOR_PATH"
WRITE_PACKAGE_PROJECTOR_ARGS+=(--projector "$PROJECTOR_PATH")
done <<< "${SOURCE_PROJECTOR_FILES:-}"
echo " Hugging Face cache: $HF_HUB_CACHE"
echo " Package workspace: $PACKAGE_DIR"
echo " Temporary workspace: $TMPDIR"
@ -309,6 +340,7 @@ set +e
time "$SLICER" write-package "$WRITE_PACKAGE_INPUT" \
--out-dir "$PACKAGE_DIR" \
--after-artifact-command "$ARTIFACT_UPLOAD_HOOK" \
"${WRITE_PACKAGE_PROJECTOR_ARGS[@]}" \
"${WRITE_PACKAGE_IDENTITY_ARGS[@]}"
WRITE_PACKAGE_STATUS=$?
set -e
@ -429,6 +461,13 @@ package_entry = {
"type": "layer-package",
"repo": target_repo,
"layer_count": layer_count,
"source_revision": source_revision,
}
source_entry = {
"repo": source_repo,
"file": source_file,
"revision": source_revision,
}
# Handle both dict-style and list-style variants
@ -439,14 +478,11 @@ if isinstance(variants, dict):
packages = variants[variant_name].get("packages", [])
packages = [p for p in packages if p.get("repo") != target_repo]
packages.append(package_entry)
variants[variant_name]["source"] = source_entry
variants[variant_name]["packages"] = packages
else:
variants[variant_name] = {
"source": {
"repo": source_repo,
"file": source_file,
"revision": source_revision,
},
"source": source_entry,
"curated": {
"name": variant_name,
"size": f"{layer_count} layers",
@ -466,14 +502,11 @@ else:
packages = existing_variant.get("packages", [])
packages = [p for p in packages if p.get("repo") != target_repo]
packages.append(package_entry)
existing_variant["source"] = source_entry
existing_variant["packages"] = packages
else:
variants.append({
"source": {
"repo": source_repo,
"file": source_file,
"revision": source_revision,
},
"source": source_entry,
"curated": {
"name": variant_name,
"size": f"{layer_count} layers",
@ -522,6 +555,21 @@ source_revision = os.environ.get("SOURCE_REVISION", "main")
target_repo = os.environ["TARGET_REPO"]
model_id = os.environ.get("MODEL_ID", manifest.get("model_id", target_repo))
mesh_llm_ref = os.environ.get("MESH_LLM_REF", "main")
source_pipeline_tag = os.environ.get("SOURCE_PIPELINE_TAG", "text-generation").strip()
if not source_pipeline_tag:
source_pipeline_tag = "text-generation"
experimental = os.environ.get("PACKAGE_EXPERIMENTAL", "false").lower() == "true"
experimental_tag = "- experimental\n" if experimental else ""
experimental_warning = (
"> [!WARNING]\n"
"> **Experimental package:** artifact integrity may be validated, but runtime, "
"split-correctness, and multimodal certification are still pending. This package "
"is not discoverable through `meshllm/catalog@main` until its Hugging Face catalog "
"PR is reviewed and merged.\n\n"
if experimental
else ""
)
api = HfApi(token=os.environ["HF_TOKEN"])
def sha256(path: Path) -> str:
digest = hashlib.sha256()
@ -555,6 +603,38 @@ def code(value) -> str:
def yaml_quote(value: str) -> str:
return json.dumps(value)
def card_value(info, key: str):
card_data = getattr(info, "card_data", None)
if card_data is None:
return None
if isinstance(card_data, dict):
return card_data.get(key)
return getattr(card_data, key, None)
def first_model_id(value):
if isinstance(value, list):
value = value[0] if value else None
if isinstance(value, dict):
return value.get("id") or value.get("modelId")
return value if isinstance(value, str) and value else None
def resolve_upstream_license():
try:
source_info = api.model_info(source_repo, revision=source_revision)
source_license = card_value(source_info, "license")
if source_license:
return str(source_license), source_repo
base_repo = first_model_id(card_value(source_info, "base_model"))
if base_repo:
base_info = api.model_info(base_repo)
base_license = card_value(base_info, "license")
if base_license:
return str(base_license), base_repo
except Exception as error:
print(f" WARNING: could not resolve upstream license metadata: {error}")
return None, None
def infer_model_family(name: str) -> str:
lowered = name.lower()
for family in ["Qwen3", "Qwen2.5", "DeepSeek", "Kimi", "Gemma", "GLM", "Llama"]:
@ -601,6 +681,10 @@ activation_width = manifest.get("activation_width") or "not recorded"
skippy_abi = manifest.get("skippy_abi_version") or "not recorded"
source_sha = source_model.get("sha256") or "not recorded"
canonical_ref = source_model.get("canonical_ref") or f"{source_repo}@{source_revision}/{source_file}"
upstream_license, license_source_repo = resolve_upstream_license()
license_frontmatter = (
f"license: {yaml_quote(upstream_license)}\n" if upstream_license else ""
)
file_rows = [
("Manifest", "model-package.json", "Package schema, source identity, checksums", manifest_hash),
@ -647,12 +731,18 @@ rows = [
("Source file", code(source_file)),
("Package repo", link(target_repo, f"https://huggingface.co/{target_repo}")),
]
if upstream_license and license_source_repo:
rows.append((
"License",
f"{code(upstream_license)} from "
f"{link(license_source_repo, f'https://huggingface.co/{license_source_repo}')}",
))
readme = f"""---
library_name: mesh-llm
base_model:
{license_frontmatter}base_model:
- {yaml_quote(source_repo)}
pipeline_tag: text-generation
pipeline_tag: {yaml_quote(source_pipeline_tag)}
tags:
- gguf
- mesh-llm
@ -661,7 +751,7 @@ tags:
- distributed-inference
- local-inference
- openai-compatible
---
{experimental_tag}---
<div align="center">
<a href="https://www.meshllm.cloud">
@ -681,7 +771,7 @@ tags:
</p>
</div>
GGUF layer package for running **{display_name}** across a local Mesh LLM cluster.
{experimental_warning}GGUF layer package for running **{display_name}** across a local Mesh LLM cluster.
This package is derived from [{source_repo}](https://huggingface.co/{source_repo}) and keeps the original GGUF distribution split into per-layer artifacts for distributed inference.
@ -781,7 +871,6 @@ skippy-model-package write-package "{source_path}" --out-dir "{package_dir}"
Path("/tmp/README.md").write_text(readme)
api = HfApi(token=os.environ["HF_TOKEN"])
api.upload_file(
path_or_fileobj="/tmp/README.md",
path_in_repo="README.md",

30
crates/skippy-coordinator/src/topology.rs
View file

@ -113,6 +113,20 @@ pub enum TopologyPlanError {
}
pub fn plan_topology(input: &TopologyPlanningInput) -> Result<TopologyPlan, TopologyPlanError> {
plan_topology_with_required_stage0(input, None)
}
pub fn plan_topology_with_stage0(
input: &TopologyPlanningInput,
stage0_node_id: &str,
) -> Result<TopologyPlan, TopologyPlanError> {
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<TopologyPlan, TopologyPlanError> {
validate_input(input)?;
let minimum_context = minimum_valid_context(input.native_context_length);
@ -137,6 +151,9 @@ pub fn plan_topology(input: &TopologyPlanningInput) -> Result<TopologyPlan, Topo
else {
return;
};
if !candidate_has_required_stage0(&candidate, required_stage0_node_id) {
return;
}
if best_for_count
.as_ref()
.is_none_or(|current| candidate_better_for_same_shape(&candidate, current))
@ -417,6 +434,19 @@ fn latency_aware_planning(_input: &TopologyPlanningInput, nodes: &[UsableNode])
.any(|node| node.stage_transfer_latency_ms.is_some())
}
fn candidate_has_required_stage0(
candidate: &CandidatePlan,
required_stage0_node_id: Option<&str>,
) -> 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

25
crates/skippy-correctness/src/cli.rs
View file

@ -19,6 +19,7 @@ pub enum CommandKind {
StateHandoff(StateHandoffArgs),
NativeMtpOpenAiAb(Box<NativeMtpOpenAiAbArgs>),
GlmDsaStage0Trace(Box<GlmDsaStage0TraceArgs>),
StageFaParity(StageFaParityArgs),
}
#[derive(Args, Clone)]
@ -324,3 +325,27 @@ pub enum StatePayloadKind {
RecurrentOnly,
KvRecurrent,
}
#[derive(Args)]
pub struct StageFaParityArgs {
#[arg(long)]
pub model: PathBuf,
#[arg(long, default_value = "unsloth/inkling-GGUF:UD-Q2_K_XL")]
pub model_id: String,
#[arg(long, default_value_t = 0)]
pub layer_start: u32,
#[arg(long)]
pub layer_end: u32,
#[arg(long, default_value_t = 2048)]
pub ctx_size: u32,
#[arg(long, default_value_t = 99)]
pub n_gpu_layers: i32,
#[arg(long, default_value = "Hello")]
pub prompt: String,
#[arg(long, default_value_t = 5e-3)]
pub max_abs: f32,
#[arg(long)]
pub enabled_output: Option<PathBuf>,
#[arg(long)]
pub disabled_output: Option<PathBuf>,
}

3
crates/skippy-correctness/src/main.rs
View file

@ -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),
}
}

3
crates/skippy-correctness/src/runner/mod.rs
View file

@ -1,9 +1,12 @@
pub(crate) mod native_mtp;
mod prediction_return;
mod single_step;
mod split_chain;
mod stage_execution;
mod stage_fa_parity;
mod state_handoff;
pub use single_step::single_step;
pub use split_chain::{chain, dtype_matrix, split_scan};
pub use stage_fa_parity::stage_fa_parity;
pub use state_handoff::state_handoff;

251
crates/skippy-correctness/src/runner/prediction_return.rs Normal file
View file

@ -0,0 +1,251 @@
use std::{
env, io,
io::Write,
net::{Shutdown, SocketAddr, TcpListener, TcpStream},
sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
mpsc,
},
thread::{self, JoinHandle},
time::Duration,
};
use anyhow::{Context, Result, anyhow, bail};
use skippy_protocol::binary::{
READY_MAGIC, StageReply, WireMessageKind, read_stage_message, recv_ready, recv_reply,
send_ready,
};
const CLIENT_READY_HELLO_ENV: &str = "SKIPPY_STAGE_CLIENT_READY_HELLO";
const CLIENT_READY_HELLO_PEEK_TIMEOUT: Duration = Duration::from_millis(500);
pub(super) struct PredictionReturnListener {
bind_addr: SocketAddr,
receiver: mpsc::Receiver<Result<StageReply, String>>,
shutdown: Arc<AtomicBool>,
connection: Arc<Mutex<Option<TcpStream>>>,
thread: Option<JoinHandle<()>>,
}
impl PredictionReturnListener {
pub(super) fn start() -> Result<Self> {
let listener = TcpListener::bind("127.0.0.1:0")
.context("bind correctness prediction return listener")?;
let bind_addr = listener
.local_addr()
.context("read correctness prediction return listener address")?;
listener
.set_nonblocking(true)
.context("set correctness prediction return listener nonblocking")?;
let shutdown = Arc::new(AtomicBool::new(false));
let thread_shutdown = shutdown.clone();
let connection = Arc::new(Mutex::new(None));
let thread_connection = connection.clone();
let (sender, receiver) = mpsc::channel();
let thread = thread::spawn(move || {
let result =
accept_prediction_return(listener, &thread_shutdown, &thread_connection, &sender);
if let Err(error) = result {
let _ = sender.send(Err(format!("{error:#}")));
}
});
Ok(Self {
bind_addr,
receiver,
shutdown,
connection,
thread: Some(thread),
})
}
pub(super) fn endpoint(&self) -> String {
format!("tcp://{}", self.bind_addr)
}
pub(super) fn receive(&self, timeout: Duration) -> Result<StageReply> {
match self.receiver.recv_timeout(timeout) {
Ok(Ok(reply)) => Ok(reply),
Ok(Err(error)) => Err(anyhow!(error)),
Err(mpsc::RecvTimeoutError::Timeout) => {
bail!("timed out waiting for direct prediction return")
}
Err(mpsc::RecvTimeoutError::Disconnected) => {
bail!("direct prediction return listener disconnected")
}
}
}
}
impl Drop for PredictionReturnListener {
fn drop(&mut self) {
self.shutdown.store(true, Ordering::SeqCst);
if let Ok(connection) = self.connection.lock()
&& let Some(stream) = connection.as_ref()
{
let _ = stream.shutdown(Shutdown::Both);
}
if let Some(thread) = self.thread.take() {
let _ = thread.join();
}
}
}
fn accept_prediction_return(
listener: TcpListener,
shutdown: &AtomicBool,
connection: &Mutex<Option<TcpStream>>,
sender: &mpsc::Sender<Result<StageReply, String>>,
) -> Result<()> {
let mut stream = loop {
match listener.accept() {
Ok((stream, _)) => break stream,
Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
if shutdown.load(Ordering::SeqCst) {
return Ok(());
}
thread::sleep(Duration::from_millis(10));
}
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Err(error) => return Err(error).context("accept direct prediction return"),
}
};
stream
.set_nonblocking(false)
.context("set direct prediction return stream blocking")?;
let shutdown_stream = stream
.try_clone()
.context("clone direct prediction return stream for shutdown")?;
{
let mut connection = connection
.lock()
.map_err(|_| anyhow!("direct prediction return connection lock poisoned"))?;
if shutdown.load(Ordering::SeqCst) {
return Ok(());
}
*connection = Some(shutdown_stream);
}
consume_optional_client_ready_hello(&mut stream)?;
send_ready(&mut stream).context("send direct prediction return ready")?;
stream.flush().ok();
let open =
read_stage_message(&mut stream, 0).context("read direct prediction return open message")?;
if open.kind != WireMessageKind::PredictionReturnOpen {
bail!("expected direct prediction return open message");
}
loop {
match recv_reply(&mut stream) {
Ok(reply) => {
if sender.send(Ok(reply)).is_err() {
return Ok(());
}
}
Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(()),
Err(error) => return Err(error).context("read direct prediction return reply"),
}
}
}
fn consume_optional_client_ready_hello(stream: &mut TcpStream) -> Result<()> {
if !client_ready_hello_enabled() {
return Ok(());
}
let previous_timeout = stream
.read_timeout()
.context("read direct prediction return timeout")?;
stream
.set_read_timeout(Some(CLIENT_READY_HELLO_PEEK_TIMEOUT))
.context("set direct prediction return hello timeout")?;
let mut bytes = [0_u8; 4];
let peek_result = stream.peek(&mut bytes);
stream
.set_read_timeout(previous_timeout)
.context("restore direct prediction return timeout")?;
match peek_result {
Ok(4) if i32::from_le_bytes(bytes) == READY_MAGIC => {
recv_ready(stream).context("consume direct prediction return client ready hello")?;
}
Ok(_) => {}
Err(error)
if matches!(
error.kind(),
io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut
) => {}
Err(error) => {
return Err(error).context("peek direct prediction return client ready hello");
}
}
Ok(())
}
fn client_ready_hello_enabled() -> bool {
env::var(CLIENT_READY_HELLO_ENV)
.map(|value| matches!(value.as_str(), "1" | "true" | "TRUE" | "yes" | "on"))
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
use skippy_protocol::binary::{
StageStateHeader, StageWireMessage, WireActivationDType, send_reply_predicted,
write_stage_message,
};
#[test]
fn receives_prediction_over_direct_return_endpoint() {
let listener = PredictionReturnListener::start().unwrap();
let endpoint = listener.endpoint();
let address = endpoint.strip_prefix("tcp://").unwrap().to_string();
let client = thread::spawn(move || {
let mut stream = TcpStream::connect(address).unwrap();
if client_ready_hello_enabled() {
send_ready(&mut stream).unwrap();
}
recv_ready(&mut stream).unwrap();
let kind = WireMessageKind::PredictionReturnOpen;
let open = StageWireMessage {
kind,
pos_start: 0,
token_count: 0,
state: StageStateHeader::new(kind, WireActivationDType::F32),
request_id: 11,
session_id: 13,
sampling: None,
chat_sampling_metadata: None,
tokens: Vec::new(),
positions: Vec::new(),
activation: Vec::new(),
raw_bytes: Vec::new(),
};
write_stage_message(&mut stream, &open, WireActivationDType::F32).unwrap();
send_reply_predicted(&mut stream, 674).unwrap();
});
let reply = listener.receive(Duration::from_secs(1)).unwrap();
assert_eq!(reply.predicted, 674);
client.join().unwrap();
}
#[test]
fn drop_stops_when_connected_peer_stalls() {
let listener = PredictionReturnListener::start().unwrap();
let address = listener
.endpoint()
.strip_prefix("tcp://")
.unwrap()
.to_string();
let mut stream = TcpStream::connect(address).unwrap();
if client_ready_hello_enabled() {
send_ready(&mut stream).unwrap();
}
recv_ready(&mut stream).unwrap();
let started = std::time::Instant::now();
drop(listener);
assert!(started.elapsed() < Duration::from_secs(1));
}
}

26
crates/skippy-correctness/src/runner/split_chain.rs
View file

@ -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<BinaryChainResult> {
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<BinaryChainResult> {
"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<BinaryChainResult> {
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<BinaryChainResult> {
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<BinaryChainResult> {
})?;
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),

134
crates/skippy-correctness/src/runner/stage_fa_parity.rs Normal file
View file

@ -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<skippy_runtime::ActivationFrame> {
let config = RuntimeConfig {
stage_index: 0,
layer_start: args.layer_start,
layer_end: args.layer_end,
ctx_size: args.ctx_size,
lane_count: 1,
n_batch: None,
n_ubatch: None,
n_threads: None,
n_threads_batch: None,
n_gpu_layers: args.n_gpu_layers,
mmap: None,
mlock: false,
selected_backend_device: None,
cache_type_k: GGML_TYPE_F16,
cache_type_v: GGML_TYPE_F16,
flash_attn_type,
load_mode: RuntimeLoadMode::LayerPackage,
projector_path: None,
include_embeddings: true,
include_output: false,
filter_tensors_on_load: true,
};
let selection = select_layer_package_parts(&PackageStageRequest {
model_id: args.model_id.clone(),
topology_id: "stage-fa-parity".to_string(),
package_ref: args.model.display().to_string(),
stage_id: "stage-0".to_string(),
layer_start: args.layer_start,
layer_end: args.layer_end,
include_embeddings: true,
include_output: false,
})
.context("select package parts")?;
let model = StageModel::open_from_parts(&selection.absolute_paths, &config)
.context("open stage model")?;
let tokens = model
.tokenize(&args.prompt, true)
.context("tokenize prompt")?;
if tokens.is_empty() {
bail!("prompt produced no token");
}
let mut session = model.create_session().context("create stage session")?;
if tokens.len() == 1 {
let (_, frame) = session
.decode_step_frame(tokens[0], None, 0)
.context("decode stage boundary")?;
Ok(frame)
} else {
session
.prefill_chunk_frame(&tokens, None, 0)
.context("prefill stage boundary")
}
}
fn payload_f32(payload: &[u8]) -> Result<Vec<f32>> {
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())
}

103
crates/skippy-ffi/src/lib.rs
View file

@ -1,10 +1,11 @@
pub const ABI_VERSION_MAJOR: u32 = 0;
pub const ABI_VERSION_MINOR: u32 = 1;
pub const ABI_VERSION_PATCH: u32 = 33;
pub const ABI_VERSION_PATCH: u32 = 35;
pub const FEATURE_BACKEND_DEVICES: u64 = 1 << 23;
pub const FEATURE_RUNTIME_EVENTS: u64 = 1 << 24;
pub const FEATURE_NATIVE_MTP_N1: u64 = 1 << 25;
pub const FEATURE_NGRAM_CACHE_DRAFT: u64 = 1 << 26;
pub const FEATURE_INKLING_MTP_MM: u64 = 1 << 27;
#[cfg(feature = "dynamic-runtime")]
mod dynamic_library;
@ -18,19 +19,20 @@ pub struct AbiVersion {
}
/// Whether a native runtime reporting `version` can back this binary's ABI
/// bindings. Required symbol signatures may change between patches (for
/// example `skippy_apply_chat_template_json` gained an argument in 0.1.28),
/// so older runtimes must be rejected at load time.
/// bindings. Required symbol signatures and by-value struct layouts may change
/// between patches, so the loader requires an exact ABI match.
pub const fn runtime_abi_supported(version: AbiVersion) -> bool {
version.major == ABI_VERSION_MAJOR
&& version.minor == ABI_VERSION_MINOR
&& version.patch >= ABI_VERSION_PATCH
&& version.patch == ABI_VERSION_PATCH
}
use std::ffi::{c_char, c_int, c_void};
pub type LlamaLogCallback =
Option<unsafe extern "C" fn(level: c_int, text: *const c_char, user_data: *mut c_void)>;
pub type MtmdProgressCallback =
Option<unsafe extern "C" fn(progress: f32, user_data: *mut c_void) -> bool>;
pub type SkippyRuntimeEventCallback =
Option<unsafe extern "C" fn(event: *const SkippyRuntimeEventV1, user_data: *mut c_void)>;
@ -349,6 +351,9 @@ pub struct MtmdContextParams {
pub image_max_tokens: c_int,
pub cb_eval: *mut c_void,
pub cb_eval_user_data: *mut c_void,
pub batch_max_tokens: c_int,
pub progress_callback: MtmdProgressCallback,
pub progress_callback_user_data: *mut c_void,
}
#[repr(C)]
@ -377,6 +382,8 @@ pub struct ActivationDesc {
pub flags: u64,
}
pub const ACTIVATION_FLAG_INKLING_MTP_EMBD: u64 = 1 << 2;
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct LogitBias {
@ -422,6 +429,7 @@ pub enum LlamaFileType {
MostlyMxfp4Moe = 38,
MostlyNvfp4 = 39,
MostlyQ1_0 = 40,
MostlyQ2_0 = 41,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@ -461,7 +469,8 @@ pub enum GgmlType {
Mxfp4 = 39,
Nvfp4 = 40,
Q1_0 = 41,
Count = 42,
Q2_0 = 42,
Count = 43,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@ -847,6 +856,7 @@ mod dynamic {
skippy_session_last_token_signal(session: *mut Session, out_signal: *mut TokenSignal, out_error: *mut *mut Error) -> Status;
skippy_session_signal_window(session: *mut Session, window_tokens: u32, out_window: *mut GenerationSignalWindow, out_error: *mut *mut Error) -> Status;
skippy_trim_session(session: *mut Session, token_count: u64, out_error: *mut *mut Error) -> Status;
skippy_retire_verify_checkpoint(session: *mut Session, token_start: u64, token_count: u64, out_error: *mut *mut Error) -> Status;
skippy_export_state(session: *mut Session, layer_start: i32, layer_end: i32, output: *mut c_void, output_capacity: usize, out_bytes: *mut usize, out_error: *mut *mut Error) -> Status;
skippy_import_state(session: *mut Session, layer_start: i32, layer_end: i32, input: *const c_void, input_bytes: usize, out_error: *mut *mut Error) -> Status;
skippy_export_full_state(session: *mut Session, layer_start: i32, layer_end: i32, output: *mut c_void, output_capacity: usize, out_bytes: *mut usize, out_error: *mut *mut Error) -> Status;
@ -887,6 +897,7 @@ mod dynamic {
mtmd_decode_use_mrope(ctx: *const MtmdContext) -> bool;
mtmd_input_chunk_get_type(chunk: *const Opaque) -> MtmdInputChunkType;
mtmd_input_chunk_get_n_tokens(chunk: *const Opaque) -> usize;
mtmd_input_chunk_get_tokens_text(chunk: *const Opaque, out_count: *mut usize) -> *const i32;
mtmd_input_chunk_get_tokens_image(chunk: *const Opaque) -> *const Opaque;
mtmd_helper_image_get_decoder_pos(image: *const Opaque, pos_0: i32, out_pos: *mut MtmdDecoderPos);
mtmd_helper_eval_chunks(ctx: *mut MtmdContext, lctx: *mut Opaque, chunks: *const MtmdInputChunks, n_past: i32, seq_id: i32, n_batch: i32, logits_last: bool, new_n_past: *mut i32) -> c_int;
@ -1172,9 +1183,36 @@ pub use dynamic::*;
/// Returns the skippy ABI feature bitmask.
/// Requires the native runtime to be loaded first (checked by caller).
pub fn skippy_abi_features() -> u64 {
let fns = dynamic::skippy_abi_features_optional()
.expect("skippy_abi_features not available in loaded runtime");
unsafe { fns() }
try_abi_features().expect("skippy_abi_features not available in loaded runtime")
}
/// Returns the Skippy ABI feature bitmask when the loaded dynamic runtime
/// exports feature probing.
#[cfg(feature = "dynamic-runtime")]
pub fn try_abi_features() -> Option<u64> {
dynamic::skippy_abi_features_optional().map(|features| unsafe { features() })
}
/// Returns the active Skippy ABI feature bitmask through a safe Rust wrapper.
#[cfg(feature = "dynamic-runtime")]
pub fn abi_features() -> u64 {
skippy_abi_features()
}
/// Returns the statically linked Skippy ABI feature bitmask.
#[cfg(not(feature = "dynamic-runtime"))]
pub fn try_abi_features() -> Option<u64> {
// SAFETY: the statically linked ABI exposes this nullary query with no
// caller-owned pointers or lifetime requirements.
Some(unsafe { skippy_abi_features() })
}
/// Returns the statically linked Skippy ABI feature bitmask.
#[cfg(not(feature = "dynamic-runtime"))]
pub fn abi_features() -> u64 {
// SAFETY: the statically linked ABI exposes this nullary query with no
// caller-owned pointers or lifetime requirements.
unsafe { skippy_abi_features() }
}
#[cfg(not(feature = "dynamic-runtime"))]
@ -1529,6 +1567,13 @@ unsafe extern "C" {
out_error: *mut *mut Error,
) -> Status;
pub fn skippy_retire_verify_checkpoint(
session: *mut Session,
token_start: u64,
token_count: u64,
out_error: *mut *mut Error,
) -> Status;
pub fn skippy_export_state(
session: *mut Session,
layer_start: i32,
@ -1795,6 +1840,11 @@ unsafe extern "C" {
pub fn mtmd_input_chunk_get_n_tokens(chunk: *const Opaque) -> usize;
pub fn mtmd_input_chunk_get_tokens_text(
chunk: *const Opaque,
out_count: *mut usize,
) -> *const i32;
pub fn mtmd_input_chunk_get_tokens_image(chunk: *const Opaque) -> *const Opaque;
pub fn mtmd_helper_image_get_decoder_pos(
@ -1841,6 +1891,7 @@ pub type Opaque = c_void;
#[cfg(test)]
mod tests {
use super::*;
use std::mem::{offset_of, size_of};
const fn version(major: u32, minor: u32, patch: u32) -> AbiVersion {
AbiVersion {
@ -1851,21 +1902,21 @@ mod tests {
}
#[test]
fn accepts_current_and_newer_patch_runtimes() {
fn accepts_current_patch_runtime() {
assert!(runtime_abi_supported(version(
ABI_VERSION_MAJOR,
ABI_VERSION_MINOR,
ABI_VERSION_PATCH,
)));
assert!(runtime_abi_supported(version(
}
#[test]
fn rejects_other_patch_runtimes() {
assert!(!runtime_abi_supported(version(
ABI_VERSION_MAJOR,
ABI_VERSION_MINOR,
ABI_VERSION_PATCH + 1,
)));
}
#[test]
fn rejects_older_patch_runtimes() {
assert!(!runtime_abi_supported(version(
ABI_VERSION_MAJOR,
ABI_VERSION_MINOR,
@ -1886,4 +1937,26 @@ mod tests {
ABI_VERSION_PATCH,
)));
}
#[test]
#[cfg(target_pointer_width = "64")]
fn mtmd_context_params_matches_native_layout() {
assert_eq!(size_of::<MtmdContextParams>(), 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());
}
}

1
crates/skippy-model-package/Cargo.toml
View file

@ -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" }

3
crates/skippy-model-package/src/package.rs
View file

@ -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<u32>,
}
#[derive(Debug, Deserialize, Serialize)]
@ -663,6 +665,7 @@ pub(crate) fn package_generation(tensors: &[TensorInfo]) -> Option<PackageGenera
initial_window: 1,
min_window: 1,
max_window: 1,
pipeline_depth: None,
}),
proposer: Some(strategy_id.clone()),
primary: None,

546
crates/skippy-model-package/src/preflight.rs
View file

@ -3,14 +3,21 @@ use std::fs;
use std::path::Path;
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 {
@ -182,6 +189,8 @@ pub(crate) struct PreflightWindowPolicy {
pub initial_window: u32,
pub min_window: u32,
pub max_window: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub pipeline_depth: Option<u32>,
}
#[derive(Debug, Deserialize)]
@ -273,6 +282,8 @@ struct PackageWindowPolicy {
initial_window: u32,
min_window: u32,
max_window: u32,
#[serde(default)]
pipeline_depth: Option<u32>,
}
#[derive(Debug, Deserialize)]
@ -1014,6 +1025,23 @@ fn validate_window_policy(
"use positive window sizes",
);
}
if window.pipeline_depth.is_some_and(|depth| {
depth == 0
|| usize::try_from(depth)
.map(|depth| depth > MAX_VERIFY_WINDOW_PIPELINE_DEPTH)
.unwrap_or(true)
}) {
report.error(
"invalid_window_policy_pipeline_depth",
format!(
"speculative strategy {name} window_policy.pipeline_depth must be between 1 and {MAX_VERIFY_WINDOW_PIPELINE_DEPTH}"
),
Some("model-package.json".to_string()),
format!(
"set pipeline_depth to an in-flight verification-window capacity no greater than {MAX_VERIFY_WINDOW_PIPELINE_DEPTH}"
),
);
}
if window.min_window > window.max_window {
report.error(
"invalid_window_policy_bounds",
@ -1135,432 +1163,10 @@ fn preflight_window_policy(window: &PackageWindowPolicy) -> PreflightWindowPolic
initial_window: window.initial_window,
min_window: window.min_window,
max_window: window.max_window,
pipeline_depth: window.pipeline_depth,
}
}
fn collect_artifacts(manifest: &PackageManifest) -> Vec<ArtifactSpec> {
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<u32>,
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::<u32, usize>::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<PreflightIssue>,
) -> 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<PreflightIssue>) {
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<PreflightIssue>,
) -> 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<u64>,
size_matches_manifest: Option<bool>,
sha256_matches_manifest: Option<bool>,
) -> 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<usize>,
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<String, StageArtifact>,
) -> 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::<Vec<_>>();
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<String, StageArtifact> {
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<PreflightIssue>,
code: impl Into<String>,
@ -1735,7 +1341,8 @@ mod tests {
"default": "fixed",
"initial_window": 1,
"min_window": 1,
"max_window": 1
"max_window": 1,
"pipeline_depth": 2
}
}
}
@ -1770,6 +1377,7 @@ mod tests {
assert_eq!(window_policy.initial_window, 1);
assert_eq!(window_policy.min_window, 1);
assert_eq!(window_policy.max_window, 1);
assert_eq!(window_policy.pipeline_depth, Some(2));
fs::remove_dir_all(dir).unwrap();
}
@ -2104,6 +1712,96 @@ mod tests {
fs::remove_dir_all(dir).unwrap();
}
#[test]
fn preflight_rejects_zero_verify_window_pipeline_depth() {
let dir = unique_test_dir("zero-window-pipeline-depth");
let package = write_package_fixture(&dir, true);
write_generation_to_manifest(
&package,
serde_json::json!({
"speculative_decoding": {
"default": "ngram-suffix",
"proposers": {
"suffix": {
"type": "ngram-suffix",
"ngram_min": 5,
"ngram_max": 32,
"max_proposal_tokens": 48,
"history_scope": "request"
}
},
"strategies": {
"ngram-suffix": {
"type": "ngram-suffix",
"proposer": "suffix",
"window_policy": {
"default": "fixed",
"initial_window": 32,
"min_window": 1,
"max_window": 32,
"pipeline_depth": 0
}
}
}
}
}),
);
let report = preflight_package(&package, &PackagePreflightOptions::default());
assert!(!report.valid);
assert_issue(&report, "invalid_window_policy_pipeline_depth");
fs::remove_dir_all(dir).unwrap();
}
#[test]
fn preflight_enforces_verify_window_pipeline_depth_maximum() {
for (depth, expected_valid) in [
(MAX_VERIFY_WINDOW_PIPELINE_DEPTH, true),
(MAX_VERIFY_WINDOW_PIPELINE_DEPTH + 1, false),
] {
let dir = unique_test_dir(&format!("window-pipeline-depth-{depth}"));
let package = write_package_fixture(&dir, true);
write_generation_to_manifest(
&package,
serde_json::json!({
"speculative_decoding": {
"default": "ngram-suffix",
"proposers": {
"suffix": {
"type": "ngram-suffix",
"ngram_min": 5,
"ngram_max": 32,
"max_proposal_tokens": 48,
"history_scope": "request"
}
},
"strategies": {
"ngram-suffix": {
"type": "ngram-suffix",
"proposer": "suffix",
"window_policy": {
"default": "fixed",
"initial_window": 32,
"min_window": 1,
"max_window": 32,
"pipeline_depth": depth
}
}
}
}
}),
);
let report = preflight_package(&package, &PackagePreflightOptions::default());
assert_eq!(report.valid, expected_valid, "pipeline depth {depth}");
if !expected_valid {
assert_issue(&report, "invalid_window_policy_pipeline_depth");
}
fs::remove_dir_all(dir).unwrap();
}
}
fn assert_issue(report: &PackagePreflightReport, code: &str) {
assert!(
report.issues.iter().any(|issue| issue.code == code),

427
crates/skippy-model-package/src/preflight/artifacts.rs Normal file
View file

@ -0,0 +1,427 @@
use super::*;
pub(super) fn collect_artifacts(manifest: &PackageManifest) -> Vec<ArtifactSpec> {
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<u32>,
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::<u32, usize>::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<PreflightIssue>,
) -> 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<PreflightIssue>) {
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<PreflightIssue>,
) -> 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<u64>,
size_matches_manifest: Option<bool>,
sha256_matches_manifest: Option<bool>,
) -> 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<usize>,
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<String, StageArtifact>,
) -> 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::<Vec<_>>();
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<String, StageArtifact> {
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()
}

5
crates/skippy-protocol/src/binary/activation.rs
View file

@ -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

73
crates/skippy-protocol/src/binary/mod.rs
View file

@ -15,8 +15,8 @@ pub use codec::{
send_reply_predicted_with_tokens_window_and_stats, write_stage_message,
};
pub use types::{
ACTIVATION_FLAG_GEMMA3N_ALTUP, ACTIVATION_FLAG_RWKV7_V_FIRST, LLAMA_TOKEN_NULL,
MAX_STAGE_ACTIVATION_BYTES, MAX_STAGE_CHAT_SAMPLING_METADATA_BYTES,
ACTIVATION_FLAG_GEMMA3N_ALTUP, ACTIVATION_FLAG_INKLING_MTP_EMBD, ACTIVATION_FLAG_RWKV7_V_FIRST,
LLAMA_TOKEN_NULL, MAX_STAGE_ACTIVATION_BYTES, MAX_STAGE_CHAT_SAMPLING_METADATA_BYTES,
MAX_STAGE_DECODED_ACTIVATION_BYTES, MAX_STAGE_LOGIT_BIAS, MAX_STAGE_PREDICTED_TOKENS,
MAX_STAGE_SIDEBAND_VALUES, MAX_STAGE_STATE_IMPORT_BYTES, READY_MAGIC,
STAGE_LOGIT_BIAS_WIRE_BYTES, STAGE_SAMPLING_CONFIG_BASE_BYTES, STAGE_STATE_HEADER_BYTES,
@ -626,6 +626,33 @@ mod tests {
const { assert!(STAGE_WIRE_FIXED_HEADER_BYTES <= 80) };
}
#[test]
fn verify_retirement_round_trips_exact_identity() {
let kind = WireMessageKind::RetireVerifyWindow;
let message = StageWireMessage {
kind,
pos_start: 128,
token_count: 8,
state: StageStateHeader::new(kind, WireActivationDType::F32),
request_id: 23,
session_id: 29,
sampling: None,
chat_sampling_metadata: None,
tokens: Vec::new(),
positions: Vec::new(),
activation: Vec::new(),
raw_bytes: Vec::new(),
};
let mut bytes = Vec::new();
write_stage_message(&mut bytes, &message, WireActivationDType::F32).unwrap();
let decoded = read_stage_message(Cursor::new(bytes), 2048).unwrap();
assert_eq!(decoded.kind, kind);
assert_eq!(decoded.pos_start, 128);
assert_eq!(decoded.token_count, 8);
assert!(decoded.state.matches_kind(kind));
}
#[test]
fn session_control_messages_are_fixed_header_only() {
let kind = WireMessageKind::TrimSession;
@ -905,6 +932,48 @@ mod tests {
);
}
#[test]
fn inkling_mtp_embedding_sideband_activation_round_trips() {
let mut state =
StageStateHeader::new(WireMessageKind::PrefillEmbd, WireActivationDType::F32);
state.source_stage_index = 0;
state.flags |= state_flags::INKLING_MTP_EMBD_SIDEBAND;
let mut activation = Vec::new();
for value in [1.0_f32, 2.0, 3.0, 4.0] {
activation.extend_from_slice(&value.to_le_bytes());
}
let message = StageWireMessage {
kind: WireMessageKind::PrefillEmbd,
pos_start: 0,
token_count: 1,
state,
request_id: 7,
session_id: 9,
sampling: None,
chat_sampling_metadata: None,
tokens: Vec::new(),
positions: Vec::new(),
activation,
raw_bytes: Vec::new(),
};
let mut bytes = Vec::new();
write_stage_message(&mut bytes, &message, WireActivationDType::F32).unwrap();
let decoded = read_stage_message(Cursor::new(bytes), 2).unwrap();
assert_eq!(decoded.activation.len(), 16);
assert_eq!(
activation_frame_flags_from_state_flags(decoded.state.flags),
ACTIVATION_FLAG_INKLING_MTP_EMBD
);
assert_eq!(
activation_state_flags_from_frame_flags(ACTIVATION_FLAG_INKLING_MTP_EMBD),
state_flags::INKLING_MTP_EMBD_SIDEBAND
);
assert_eq!(
decoded.activation_f32_payload(2).unwrap(),
message.activation
);
}
#[test]
fn f32_activation_payload_can_be_moved_without_clone() {
let state = StageStateHeader::new(WireMessageKind::DecodeEmbd, WireActivationDType::F32);

25
crates/skippy-protocol/src/binary/types.rs
View file

@ -5,9 +5,10 @@ use super::{
invalid_data,
};
// v10 makes the coordinator the sole owner of verify-window acceptance and removes the
// redundant tail-stage acceptance/correction fields. Stage peers must be upgraded together.
pub const STAGE_STATE_VERSION: i32 = 10;
// v11 adds the Inkling MTP embedding sideband and makes the coordinator the sole owner of
// verify-window acceptance, removing redundant tail-stage acceptance/correction fields. Stage
// peers must be upgraded together so older readers reject the changed payload contract.
pub const STAGE_STATE_VERSION: i32 = 11;
pub const MAX_STAGE_LOGIT_BIAS: usize = 256;
pub const MAX_STAGE_PREDICTED_TOKENS: usize = 262_144;
pub const MAX_STAGE_SIDEBAND_VALUES: usize = 1_048_576;
@ -56,6 +57,7 @@ pub enum WireMessageKind {
DecodeReadout = 8,
DecodeLightCtx = 9,
VerifyWindow = 21,
RetireVerifyWindow = 22,
StateExport = 13,
ConfigureGeneration = 14,
ProbePrefill = 15,
@ -95,6 +97,10 @@ impl WireMessageKind {
matches!(self, Self::TrimSession)
}
pub fn is_verify_retirement(self) -> bool {
matches!(self, Self::RetireVerifyWindow)
}
pub fn is_generation_control(self) -> bool {
matches!(self, Self::ConfigureGeneration)
}
@ -140,6 +146,7 @@ impl TryFrom<i32> for WireMessageKind {
19 => Ok(Self::TrimSession),
20 => Ok(Self::PredictionReturnOpen),
21 => Ok(Self::VerifyWindow),
22 => Ok(Self::RetireVerifyWindow),
_ => Err(invalid_data("unknown stage message kind")),
}
}
@ -183,10 +190,12 @@ pub mod state_flags {
pub const CHAT_SAMPLING_METADATA: i32 = 1 << 5;
pub const RWKV7_V_FIRST_SIDEBAND: i32 = 1 << 6;
pub const GEMMA3N_ALTUP_SIDEBAND: i32 = 1 << 7;
pub const INKLING_MTP_EMBD_SIDEBAND: i32 = 1 << 8;
}
pub const ACTIVATION_FLAG_RWKV7_V_FIRST: u64 = 1 << 0;
pub const ACTIVATION_FLAG_GEMMA3N_ALTUP: u64 = 1 << 1;
pub const ACTIVATION_FLAG_INKLING_MTP_EMBD: u64 = 1 << 2;
pub fn activation_frame_flags_from_state_flags(flags: i32) -> u64 {
let mut frame_flags = 0;
@ -196,6 +205,9 @@ pub fn activation_frame_flags_from_state_flags(flags: i32) -> u64 {
if (flags & state_flags::GEMMA3N_ALTUP_SIDEBAND) != 0 {
frame_flags |= ACTIVATION_FLAG_GEMMA3N_ALTUP;
}
if (flags & state_flags::INKLING_MTP_EMBD_SIDEBAND) != 0 {
frame_flags |= ACTIVATION_FLAG_INKLING_MTP_EMBD;
}
frame_flags
}
@ -207,6 +219,9 @@ pub fn activation_state_flags_from_frame_flags(flags: u64) -> i32 {
if (flags & ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0 {
state |= state_flags::GEMMA3N_ALTUP_SIDEBAND;
}
if (flags & ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0 {
state |= state_flags::INKLING_MTP_EMBD_SIDEBAND;
}
state
}
@ -313,6 +328,7 @@ impl StageStateHeader {
kind,
WireMessageKind::StateImport | WireMessageKind::StateExport
) || kind.is_session_control()
|| kind.is_verify_retirement()
|| kind.is_generation_control()
{
return true;
@ -391,7 +407,7 @@ pub struct StageWireMessage {
impl StageWireMessage {
/// The committed session position that must exist before this message runs.
///
/// Stage-state v10 makes this absolute position authoritative: a worker whose
/// Stage-state v11 makes this absolute position authoritative: a worker whose
/// speculative KV is ahead must rewind locally before executing the message.
pub fn authoritative_session_position(&self) -> Option<u64> {
if !matches!(
@ -695,6 +711,7 @@ fn expected_phase(kind: WireMessageKind) -> WireStagePhase {
WireMessageKind::StateImport | WireMessageKind::StateExport
)
|| kind.is_session_control()
|| kind.is_verify_retirement()
|| kind.is_generation_control()
|| kind.is_prefix_cache_control()
{

12
crates/skippy-protocol/src/lib.rs
View file

@ -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}")
);
}

6
crates/skippy-quantize/Cargo.toml
View file

@ -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"

4
crates/skippy-quantize/src/backend.rs
View file

@ -129,7 +129,7 @@ fn skippy_abi_capabilities(skippy_runtime_libraries: &[PathBuf]) -> SkippyAbiCap
let load_error = load_skippy_runtime_for_probe(skippy_runtime_libraries);
let runtime_loaded = skippy_ffi::native_runtime_loaded();
let feature_mask = if runtime_loaded {
std::panic::catch_unwind(skippy_ffi::skippy_abi_features).ok()
std::panic::catch_unwind(skippy_ffi::abi_features).ok()
} else {
None
};
@ -170,7 +170,7 @@ fn skippy_abi_reason(
return "no Skippy native runtime library was loaded for ABI probing".to_string();
}
if feature_mask.is_none() {
return "loaded Skippy runtime does not expose skippy_abi_features".to_string();
return "loaded Skippy runtime does not expose abi_features".to_string();
}
if gguf_slice_write {
return "loaded Skippy ABI exposes GGUF slice writing and the linked llama symbols can be used for GGUF quantization, but not HF checkpoint conversion".to_string();

3
crates/skippy-quantize/src/direct_convert.rs
View file

@ -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,

207
crates/skippy-quantize/src/gguf_metadata.rs Normal file
View file

@ -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<bool> },
ArrayF32 { key: String, value: Vec<f32> },
ArrayI32 { key: String, value: Vec<i32> },
ArrayString { key: String, value: Vec<String> },
ArrayU32 { key: String, value: Vec<u32> },
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<bool>) -> Self {
Self::ArrayBool {
key: key.to_string(),
value,
}
}
pub(crate) fn array_f32(key: &str, value: Vec<f32>) -> Self {
Self::ArrayF32 {
key: key.to_string(),
value,
}
}
pub(crate) fn array_i32(key: &str, value: Vec<i32>) -> Self {
Self::ArrayI32 {
key: key.to_string(),
value,
}
}
pub(crate) fn array_string(key: &str, value: Vec<String>) -> Self {
Self::ArrayString {
key: key.to_string(),
value,
}
}
pub(crate) fn array_u32(key: &str, value: Vec<u32>) -> 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<W: Write>(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<W: Write>(
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<W: Write>(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<W: Write>(writer: &mut W, value: &str) -> Result<()> {
writer.write_all(&(value.len() as u64).to_le_bytes())?;
writer.write_all(value.as_bytes())?;
Ok(())
}

76
crates/skippy-quantize/src/gguf_template.rs
View file

@ -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<Option<u32>> {
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<Vec<GgufKv>> {
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()

399
crates/skippy-quantize/src/gguf_writer.rs
View file

@ -7,9 +7,17 @@ use anyhow::{Context, Result, ensure};
use serde::Serialize;
use crate::float_convert::{FloatDType, convert_float_chunk, target_dtype_for_tensor};
pub(crate) use crate::gguf_metadata::GgufKv;
use crate::gguf_metadata::write_kv;
#[cfg(test)]
use crate::gguf_metadata::{
GGUF_TYPE_ARRAY, GGUF_TYPE_BOOL, GGUF_TYPE_FLOAT32, GGUF_TYPE_INT32, GGUF_TYPE_STRING,
GGUF_TYPE_UINT16, GGUF_TYPE_UINT32, GGUF_TYPE_UINT64,
};
use crate::hf_checkpoint::{SafetensorFile, SafetensorTensorInfo, open_safetensor_files};
use crate::tensor_map::{
TensorNameMap, hf_layer_id, is_mtp_source_tensor, is_shared_mtp_context_tensor,
TensorNameMap, hf_layer_id, inkling_mtp_depth, is_inkling_fused_w13, is_mtp_source_tensor,
is_shared_mtp_context_tensor,
};
use crate::types::ConvertOutputType;
@ -23,14 +31,6 @@ use glm_dsa::{
const GGUF_MAGIC: &[u8; 4] = b"GGUF";
const GGUF_VERSION: u32 = 3;
const GGUF_ALIGNMENT: u64 = 32;
const GGUF_TYPE_BOOL: u32 = 7;
const GGUF_TYPE_UINT32: u32 = 4;
const GGUF_TYPE_INT32: u32 = 5;
const GGUF_TYPE_FLOAT32: u32 = 6;
const GGUF_TYPE_STRING: u32 = 8;
const GGUF_TYPE_ARRAY: u32 = 9;
const GGUF_TYPE_UINT16: u32 = 2;
const GGUF_TYPE_UINT64: u32 = 10;
const GGML_TYPE_F32: u32 = 0;
const GGML_TYPE_F16: u32 = 1;
const GGML_TYPE_BF16: u32 = 30;
@ -96,6 +96,58 @@ pub(crate) fn validate_raw_safetensors_gguf(
})
}
pub(crate) fn recommended_raw_safetensors_gguf_split_count(
source: &Path,
mut options: RawGgufWriteOptions,
max_tensor_bytes: u64,
) -> Result<u32> {
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::<u64>()
<= 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<SafetensorFile>,
tensors: Vec<TensorSource>,
@ -305,6 +357,15 @@ fn collect_tensor_sources(
if !tensor_selection.includes(tensor.name())? {
continue;
}
if is_inkling_fused_w13(tensor.name()) {
tensors.extend(inkling_w13_tensor_sources(
file_index,
tensor,
tensor_name_map,
output_type,
)?);
continue;
}
if matches!(
tensor_name_map,
TensorNameMap::HfToGguf | TensorNameMap::HfToGgufWithMtp { .. }
@ -394,9 +455,11 @@ impl TensorSource {
let source_dtype = FloatDType::from_safetensor(tensor.dtype()).with_context(|| {
format!("unsupported dtype {} for {}", tensor.dtype(), tensor.name())
})?;
let target_dtype = target_dtype_for_tensor(source_dtype, output_type, tensor.shape())?;
let name = tensor_name_map.map_tensor_name(tensor.name())?;
let target_dtype =
target_dtype_for_mapped_tensor(source_dtype, output_type, tensor.shape(), &name)?;
let element_count = tensor_element_count(tensor)?;
let dims = mapped_tensor_dims(tensor.shape(), &name)?;
Ok(Self {
segments: vec![TensorSegment {
file_index,
@ -409,7 +472,7 @@ impl TensorSource {
transform: TensorTransform::Identity,
}],
name,
dims: tensor.shape().iter().rev().copied().collect(),
dims,
ggml_type: ggml_type_for_dtype(target_dtype),
byte_len: tensor_byte_len(element_count, target_dtype)?,
gguf_offset: 0,
@ -417,6 +480,98 @@ impl TensorSource {
}
}
fn target_dtype_for_mapped_tensor(
source_dtype: FloatDType,
output_type: Option<ConvertOutputType>,
shape: &[u64],
mapped_name: &str,
) -> Result<FloatDType> {
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<Vec<u64>> {
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<ConvertOutputType>,
) -> Result<Vec<TensorSource>> {
let layer = if let Some(depth) = inkling_mtp_depth(tensor.name())? {
let TensorNameMap::HfToGgufWithMtp { layer_start } = tensor_name_map else {
anyhow::bail!("Inkling MTP conversion requires an MTP-aware tensor name map");
};
layer_start
.checked_add(depth)
.context("Inkling MTP layer id overflow")?
} else {
ensure!(
matches!(
tensor_name_map,
TensorNameMap::HfToGguf | TensorNameMap::HfToGgufWithMtp { .. }
),
"Inkling fused w13 conversion requires an HF tensor name map"
);
hf_layer_id(tensor.name())?
.with_context(|| format!("missing Inkling layer id in {}", tensor.name()))?
};
ensure!(
tensor.shape().len() == 2,
"Inkling MTP fused w13 tensor {} must be rank 2, got {:?}",
tensor.name(),
tensor.shape()
);
ensure!(
tensor.shape()[0].is_multiple_of(2),
"Inkling MTP fused w13 tensor {} must have an even row count",
tensor.name()
);
let source_dtype = FloatDType::from_safetensor(tensor.dtype())
.with_context(|| format!("unsupported dtype {} for {}", tensor.dtype(), tensor.name()))?;
let output_shape = [tensor.shape()[0] / 2, tensor.shape()[1]];
let target_dtype = target_dtype_for_tensor(source_dtype, output_type, &output_shape)?;
let element_count = output_shape[0]
.checked_mul(output_shape[1])
.context("Inkling MTP w13 output element count overflow")?;
let target_byte_len = tensor_byte_len(element_count, target_dtype)?;
let dims = output_shape.iter().rev().copied().collect::<Vec<_>>();
Ok([("ffn_gate", 0_u64), ("ffn_up", 1_u64)]
.into_iter()
.map(|(projection, parity)| TensorSource {
segments: vec![TensorSegment {
file_index,
source_name: tensor.name().to_string(),
source_dtype,
target_dtype,
element_count,
source_byte_len: tensor.byte_len(),
target_byte_len,
transform: TensorTransform::AlternatingRows {
parity,
row_elements: tensor.shape()[1],
},
}],
name: format!("blk.{layer}.{projection}.weight"),
dims: dims.clone(),
ggml_type: ggml_type_for_dtype(target_dtype),
byte_len: target_byte_len,
gguf_offset: 0,
})
.collect())
}
struct TensorSegment {
file_index: usize,
source_name: String,
@ -653,92 +808,6 @@ fn raw_metadata(source: &Path, tensor_count: usize) -> Vec<GgufKv> {
]
}
#[derive(Debug, Clone)]
pub(crate) enum GgufKv {
ArrayF32 { key: String, value: Vec<f32> },
ArrayI32 { key: String, value: Vec<i32> },
ArrayString { key: String, value: Vec<String> },
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<f32>) -> Self {
Self::ArrayF32 {
key: key.to_string(),
value,
}
}
pub(crate) fn array_i32(key: &str, value: Vec<i32>) -> Self {
Self::ArrayI32 {
key: key.to_string(),
value,
}
}
pub(crate) fn array_string(key: &str, value: Vec<String>) -> 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<W: Write>(
writer: &mut W,
metadata: &[GgufKv],
@ -763,82 +832,6 @@ fn write_header_and_tensor_table<W: Write>(
Ok(())
}
fn write_kv<W: Write>(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<W: Write>(
writer: &mut W,
key: &str,
element_type: u32,
len: usize,
) -> Result<()> {
ensure!(!key.is_empty(), "GGUF metadata key cannot be empty");
ensure!(
len > 0,
"GGUF array metadata {key:?} cannot be empty because llama.cpp rejects empty arrays"
);
write_string(writer, key)?;
write_u32(writer, GGUF_TYPE_ARRAY)?;
write_u32(writer, element_type)?;
write_u64(writer, len as u64)
}
fn stream_tensor_data(
writer: &mut File,
files: &[SafetensorFile],
@ -880,6 +873,13 @@ fn stream_segment(
segment: &TensorSegment,
buffer_size: usize,
) -> Result<u64> {
if let TensorTransform::AlternatingRows {
parity,
row_elements,
} = segment.transform
{
return stream_alternating_rows(writer, file, segment, buffer_size, parity, row_elements);
}
if let Some(written) = stream_transformed_segment(writer, file, segment, buffer_size)? {
return Ok(written);
}
@ -927,6 +927,59 @@ fn stream_segment(
Ok(output_bytes)
}
/// Deinterleave alternating rows of a fused SwiGLU tensor (Inkling MTP fused
/// w13): parity 0 keeps even rows (gate), parity 1 keeps odd rows (up).
fn stream_alternating_rows(
writer: &mut File,
file: &SafetensorFile,
segment: &TensorSegment,
buffer_size: usize,
parity: u64,
row_elements: u64,
) -> Result<u64> {
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)

11
crates/skippy-quantize/src/gguf_writer/glm_dsa.rs
View file

@ -368,9 +368,11 @@ fn metadata_u32(metadata: &[GgufKv], key: &str) -> Option<u32> {
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,

198
crates/skippy-quantize/src/gguf_writer_tests.rs
View file

@ -260,6 +260,139 @@ fn writes_qwen_style_mtp_only_tensors_with_shared_context() {
fs::remove_dir_all(root).unwrap();
}
#[test]
fn writes_inkling_mtp_streaming_transforms() {
let root = unique_temp_dir();
fs::create_dir_all(&root).unwrap();
let w13 = (1_u32..=8)
.flat_map(|value| (value as f32).to_le_bytes())
.collect::<Vec<_>>();
let bf16_values = [0x80, 0x3f, 0x00, 0x40, 0x40, 0x40, 0x80, 0x40];
write_safetensor(
&root.join("model.safetensors"),
&[
("model.llm.embed.weight", "F32", &[1], &[1, 0, 0, 0]),
("model.llm.embed_norm.weight", "F32", &[1], &[2, 0, 0, 0]),
("model.llm.norm.weight", "F32", &[1], &[3, 0, 0, 0]),
("model.llm.unembed.weight", "F32", &[1], &[4, 0, 0, 0]),
(
"model.mtp.layers.0.embed_norm.weight",
"F32",
&[1],
&[5, 0, 0, 0],
),
(
"model.mtp.layers.0.transformer_block.attn.rel_logits_proj.proj",
"BF16",
&[2, 2],
&bf16_values,
),
(
"model.mtp.layers.0.transformer_block.attn.k_sconv.weight",
"BF16",
&[2, 1, 2],
&bf16_values,
),
(
"model.mtp.layers.0.transformer_block.mlp.w13_dn.weight",
"F32",
&[4, 2],
&w13,
),
],
);
let output = root.join("inkling-mtp.gguf");
write_raw_safetensors_gguf(
&root,
&output,
RawGgufWriteOptions {
buffer_size: 9,
metadata: None,
tensor_name_map: TensorNameMap::HfToGgufWithMtp { layer_start: 66 },
split: None,
output_type: Some(ConvertOutputType::Bf16),
tensor_selection: TensorSelection::MtpOnly { layer_start: 66 },
},
)
.unwrap();
let bytes = fs::read(&output).unwrap();
let parsed = parse_test_gguf(&bytes);
assert_eq!(parsed.tensor_count, 9);
let shortconv = parsed.tensor("blk.66.shortconv_k.weight");
assert_eq!(shortconv.dims, vec![2, 2]);
assert_eq!(shortconv.ggml_type, GGML_TYPE_F32);
let rel_proj = parsed.tensor("blk.66.attn_rel_proj.weight");
assert_eq!(rel_proj.ggml_type, GGML_TYPE_F32);
let gate = parsed.tensor("blk.66.ffn_gate.weight");
let up = parsed.tensor("blk.66.ffn_up.weight");
assert_eq!(gate.dims, vec![2, 2]);
assert_eq!(up.dims, vec![2, 2]);
assert_eq!(gate.ggml_type, GGML_TYPE_BF16);
assert_eq!(up.ggml_type, GGML_TYPE_BF16);
let gate_expected = [0x80, 0x3f, 0x00, 0x40, 0xa0, 0x40, 0xc0, 0x40];
let up_expected = [0x40, 0x40, 0x80, 0x40, 0xe0, 0x40, 0x00, 0x41];
assert_eq!(
&bytes[gate.absolute_offset..gate.absolute_offset + gate_expected.len()],
gate_expected
);
assert_eq!(
&bytes[up.absolute_offset..up.absolute_offset + up_expected.len()],
up_expected
);
fs::remove_dir_all(root).unwrap();
}
#[test]
fn writes_inkling_trunk_fused_w13_streaming_transforms() {
let root = unique_temp_dir();
fs::create_dir_all(&root).unwrap();
let w13 = (1_u32..=8)
.flat_map(|value| (value as f32).to_le_bytes())
.collect::<Vec<_>>();
write_safetensor(
&root.join("model.safetensors"),
&[("model.layers.3.mlp.w13_dn.weight", "F32", &[4, 2], &w13)],
);
let output = root.join("inkling-trunk.gguf");
write_raw_safetensors_gguf(
&root,
&output,
RawGgufWriteOptions {
buffer_size: 9,
metadata: None,
tensor_name_map: TensorNameMap::HfToGguf,
split: None,
output_type: Some(ConvertOutputType::Bf16),
tensor_selection: TensorSelection::All,
},
)
.unwrap();
let bytes = fs::read(&output).unwrap();
let parsed = parse_test_gguf(&bytes);
assert_eq!(parsed.tensor_count, 2);
let gate = parsed.tensor("blk.3.ffn_gate.weight");
let up = parsed.tensor("blk.3.ffn_up.weight");
assert_eq!(gate.dims, vec![2, 2]);
assert_eq!(up.dims, vec![2, 2]);
assert_eq!(gate.ggml_type, GGML_TYPE_BF16);
assert_eq!(up.ggml_type, GGML_TYPE_BF16);
let gate_expected = [0x80, 0x3f, 0x00, 0x40, 0xa0, 0x40, 0xc0, 0x40];
let up_expected = [0x40, 0x40, 0x80, 0x40, 0xe0, 0x40, 0x00, 0x41];
assert_eq!(
&bytes[gate.absolute_offset..gate.absolute_offset + gate_expected.len()],
gate_expected
);
assert_eq!(
&bytes[up.absolute_offset..up.absolute_offset + up_expected.len()],
up_expected
);
fs::remove_dir_all(root).unwrap();
}
#[test]
fn validates_qwen_dense_native_conversion_fixture() {
let root = unique_temp_dir();
@ -1104,6 +1237,64 @@ fn native_splits_are_byte_balanced_not_tensor_count_balanced() {
fs::remove_dir_all(root).unwrap();
}
#[test]
fn recommends_enough_byte_balanced_splits_for_the_size_limit() {
let root = unique_temp_dir();
fs::create_dir_all(&root).unwrap();
write_safetensor(
&root.join("model.safetensors"),
&[
("a.weight", "F32", &[4], &[1; 16]),
("b.weight", "F32", &[4], &[2; 16]),
("c.weight", "F32", &[4], &[3; 16]),
],
);
let split_count = recommended_raw_safetensors_gguf_split_count(
&root,
RawGgufWriteOptions {
buffer_size: 4,
metadata: None,
tensor_name_map: TensorNameMap::Raw,
split: None,
output_type: None,
tensor_selection: TensorSelection::All,
},
20,
)
.unwrap();
assert_eq!(split_count, 3);
fs::remove_dir_all(root).unwrap();
}
#[test]
fn rejects_a_size_limit_smaller_than_one_selected_tensor() {
let root = unique_temp_dir();
fs::create_dir_all(&root).unwrap();
write_safetensor(
&root.join("model.safetensors"),
&[("a.weight", "F32", &[4], &[1; 16])],
);
let error = recommended_raw_safetensors_gguf_split_count(
&root,
RawGgufWriteOptions {
buffer_size: 4,
metadata: None,
tensor_name_map: TensorNameMap::Raw,
split: None,
output_type: None,
tensor_selection: TensorSelection::All,
},
15,
)
.unwrap_err();
assert!(error.to_string().contains("largest selected tensor"));
fs::remove_dir_all(root).unwrap();
}
#[test]
fn keeps_rank_one_f32_tensor_as_f32_for_bf16_output() {
let root = unique_temp_dir();
@ -1267,6 +1458,13 @@ fn read_string_array_or_skip(cursor: &mut std::io::Cursor<&[u8]>) -> Option<Vec<
fn skip_array_items(cursor: &mut std::io::Cursor<&[u8]>, 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);
}

35
crates/skippy-quantize/src/hf_checkpoint.rs
View file

@ -334,6 +334,11 @@ fn discover_safetensors(source: &Path) -> Result<Vec<PathBuf>> {
);
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();

243
crates/skippy-quantize/src/inkling_metadata.rs Normal file
View file

@ -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<Option<u32>> {
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<Vec<GgufKv>> {
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<Vec<bool>> {
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<Vec<bool>> {
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<u32> {
nested_u32(config, path).with_context(|| format!("config missing {}", path.join(".")))
}
fn nested_u32(config: &Value, path: &[&str]) -> Option<u32> {
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<bool> {
path.iter()
.try_fold(config, |value, key| value.get(key))?
.as_bool()
}
fn required_u32(config: &Value, key: &str) -> Result<u32> {
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<f32> {
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")
}

15
crates/skippy-quantize/src/main.rs
View file

@ -12,19 +12,23 @@ mod command_reports;
mod direct_convert;
mod direct_quantize;
mod float_convert;
mod gguf_metadata;
mod gguf_template;
mod gguf_writer;
mod hf_checkpoint;
mod imatrix;
mod inkling_metadata;
mod llama_load;
mod locking;
mod manifest;
mod memory_budget;
mod mtp_attach;
mod native_convert;
mod native_quantize;
mod output;
mod plan_convert;
mod preflight;
mod projector_validate;
mod quantize;
mod records;
mod residency;
@ -56,6 +60,7 @@ use memory_budget::{
MemoryBudgetPlanInput, MemoryPolicy, MemorySize, effective_stream_buffer_bytes,
native_convert_stream_working_set_bytes, print_memory_budget_plan,
};
use mtp_attach::{ValidateMtpAttachArgs, run_validate_mtp_attach};
use native_convert::{build_native_convert_command, run_native_convert};
use native_quantize::{build_native_quantize_command, run_native_quantize};
use output::{
@ -64,6 +69,7 @@ use output::{
};
use plan_convert::{PlanConvertArgs, run_plan_convert};
use preflight::run_job_preflight;
use projector_validate::{ValidateProjectorArgs, run_validate_projector};
use records::{WindowRunRecordInput, unix_timestamp_ms, write_window_record};
use residency::remove_dir_if_exists;
use splits::{
@ -109,6 +115,8 @@ enum Command {
RunQuantWindow(RunQuantWindowArgs),
VerifyJob(VerifyJobArgs),
ValidateLlamaLoad(ValidateLlamaLoadArgs),
ValidateMtpAttach(ValidateMtpAttachArgs),
ValidateProjector(ValidateProjectorArgs),
ValidateTensorTypes(ValidateTensorTypesArgs),
ValidateSplits(ValidateSplitsArgs),
}
@ -510,6 +518,8 @@ fn main() -> Result<()> {
args.json,
),
Command::ValidateLlamaLoad(args) => run_validate_llama_load(args),
Command::ValidateMtpAttach(args) => run_validate_mtp_attach(args),
Command::ValidateProjector(args) => run_validate_projector(args),
Command::ValidateTensorTypes(args) => validate_tensor_types_command(&args.file, args.json),
Command::ValidateSplits(args) => validate_splits_command(
&args.root,
@ -990,9 +1000,10 @@ fn init_convert(args: InitConvertArgs) -> Result<()> {
}
fn convert_job(args: ConvertJobArgs) -> Result<()> {
let manifest = convert_manifest_from_args(&args.init)?;
let manifest_path = args.init.manifest.clone();
let runner = prepare_convert_runner(args.run.runner)?;
let mut manifest = convert_manifest_from_args(&args.init)?;
native_convert::apply_native_convert_split_max_size(&runner, &mut manifest)?;
let manifest_path = args.init.manifest.clone();
if args.run.preflight_only {
return run_job_preflight(
&manifest_path,

204
crates/skippy-quantize/src/mtp_attach.rs Normal file
View file

@ -0,0 +1,204 @@
use std::path::{Path, PathBuf};
use anyhow::{Context, Result, ensure};
use clap::Parser;
use serde::Serialize;
use skippy_runtime::{
FlashAttentionType, GGML_TYPE_F16, ModelInfo, RuntimeConfig, RuntimeLoadMode, StageModel,
};
use crate::output::{print_json_pretty, print_success};
#[derive(Debug, Parser)]
pub(crate) struct ValidateMtpAttachArgs {
#[arg(long = "model", required = true)]
model_parts: Vec<PathBuf>,
#[arg(long)]
mtp_draft: PathBuf,
#[arg(long)]
layer_count: u32,
#[arg(long)]
mtp_layer_count: Option<u32>,
#[arg(long)]
projector: Option<PathBuf>,
#[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<PathBuf>,
mtp_draft: PathBuf,
projector: Option<PathBuf>,
layer_count: u32,
mtp_layer_count: u32,
ctx_size: u32,
native_mtp_multimodal_feature: bool,
session_created: bool,
}
pub(crate) fn run_validate_mtp_attach(args: ValidateMtpAttachArgs) -> Result<()> {
validate_paths(&args)?;
ensure!(
skippy_ffi::native_runtime_loaded(),
"validate-mtp-attach requires a statically linked standalone build or a loaded native runtime"
);
let abi_features = skippy_ffi::try_abi_features()
.context("loaded native runtime does not expose Skippy ABI feature probing")?;
let native_mtp_multimodal_feature = abi_features & skippy_ffi::FEATURE_INKLING_MTP_MM != 0;
ensure!(
native_mtp_multimodal_feature,
"native runtime does not advertise Inkling multimodal MTP support"
);
let target_config = runtime_config(
args.layer_count,
args.ctx_size,
args.n_gpu_layers,
args.projector.as_deref(),
);
let mut target = if args.model_parts.len() == 1 {
StageModel::open(&args.model_parts[0], &target_config)
} else {
StageModel::open_from_parts(&args.model_parts, &target_config)
}
.context("open target model for MTP attach validation")?;
let mtp_layer_count = match args.mtp_layer_count {
Some(layer_count) => layer_count,
None => infer_layer_count(&args.mtp_draft)?,
};
ensure!(
mtp_layer_count > 0,
"--mtp-layer-count must be greater than zero"
);
let draft_config = runtime_config(mtp_layer_count, args.ctx_size, args.n_gpu_layers, None);
target
.attach_mtp_draft_model(&args.mtp_draft, &draft_config)
.context("attach MTP draft model")?;
let _session = target
.create_session()
.context("create target session after MTP attach")?;
let report = MtpAttachReport {
model_parts: args.model_parts,
mtp_draft: args.mtp_draft,
projector: args.projector,
layer_count: args.layer_count,
mtp_layer_count,
ctx_size: args.ctx_size,
native_mtp_multimodal_feature,
session_created: true,
};
if args.json {
print_json_pretty(&report)?;
} else {
print_success(format!(
"MTP attach valid: parts={} layers={} mtp_layers={} projector={} session_created=true",
report.model_parts.len(),
report.layer_count,
report.mtp_layer_count,
report
.projector
.as_deref()
.map_or_else(|| "none".to_string(), |path| path.display().to_string())
));
}
Ok(())
}
fn validate_paths(args: &ValidateMtpAttachArgs) -> Result<()> {
ensure!(
args.layer_count > 0,
"--layer-count must be greater than zero"
);
ensure!(args.ctx_size > 0, "--ctx-size must be greater than zero");
for path in &args.model_parts {
ensure!(
path.is_file(),
"model part does not exist: {}",
path.display()
);
}
ensure!(
args.mtp_draft.is_file(),
"MTP draft does not exist: {}",
args.mtp_draft.display()
);
if let Some(projector) = &args.projector {
ensure!(
projector.is_file(),
"projector does not exist: {}",
projector.display()
);
}
Ok(())
}
fn infer_layer_count(path: &Path) -> Result<u32> {
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);
}
}

54
crates/skippy-quantize/src/native_convert.rs
View file

@ -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::<MemorySize>()
.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,

68
crates/skippy-quantize/src/projector_validate.rs Normal file
View file

@ -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(())
}

141
crates/skippy-quantize/src/tensor_map.rs
View file

@ -18,23 +18,31 @@ impl TensorNameMap {
}
fn map_hf_to_gguf(name: &str, mtp_layer_start: Option<u32>) -> Result<String> {
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<u32>) -> Result<String> {
}
pub(crate) fn is_mtp_source_tensor(name: &str) -> bool {
is_qwen_mtp_source_tensor(name)
is_inkling_mtp_source_tensor(name)
|| is_qwen_mtp_source_tensor(name)
|| map_mtp_source_tensor(name).is_ok_and(|mapped| mapped.is_some())
}
@ -73,9 +82,57 @@ pub(crate) fn is_shared_mtp_context_tensor(name: &str) -> bool {
| "model.norm.weight"
| "norm.weight"
| "lm_head.weight"
| "model.llm.embed.weight"
| "model.llm.embed_norm.weight"
| "model.llm.norm.weight"
| "model.llm.unembed.weight"
)
}
pub(crate) fn is_inkling_fused_w13(name: &str) -> bool {
(name.starts_with("model.layers.") || name.starts_with("model.mtp.layers."))
&& name.ends_with(".mlp.w13_dn.weight")
}
pub(crate) fn inkling_mtp_depth(name: &str) -> Result<Option<u32>> {
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::<u32>()
.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<Option<String>> {
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::<u32>()
.map_err(|err| anyhow!("malformed Inkling MTP depth in {name}: {err}"))?;
let layer = layer_start
.checked_add(depth)
.ok_or_else(|| anyhow!("Inkling MTP layer id overflow for {name}"))?;
let suffix = match suffix {
"embed_norm.weight" => "enorm.weight",
"hidden_norm.weight" => "hnorm.weight",
"input_proj.weight" => "eh_proj.weight",
value => value.strip_prefix("transformer_block.").unwrap_or(value),
};
Ok(Some(format!("model.layers.{layer}.{suffix}")))
}
fn map_mtp_source_tensor(name: &str) -> Result<Option<String>> {
let mapped = match name {
"pre_projection" | "pre_projection.weight" => "nextn.pre_projection.weight".to_string(),
@ -182,6 +239,29 @@ impl<'a> HfLayerTensor<'a> {
"self_attn.kv_a_proj_with_mqa.weight" => Ok(format!("blk.{bid}.attn_kv_a_mqa.weight")),
"self_attn.kv_b_proj.weight" => Ok(format!("blk.{bid}.attn_kv_b.weight")),
"self_attn.kv_a_layernorm.weight" => Ok(format!("blk.{bid}.attn_kv_a_norm.weight")),
"attn_norm.weight" => Ok(format!("blk.{bid}.attn_norm.weight")),
"attn.wq_du.weight" => Ok(format!("blk.{bid}.attn_q.weight")),
"attn.wk_dv.weight" => Ok(format!("blk.{bid}.attn_k.weight")),
"attn.wv_dv.weight" => Ok(format!("blk.{bid}.attn_v.weight")),
"attn.wr_du.weight" => Ok(format!("blk.{bid}.attn_r.weight")),
"attn.wo_ud.weight" => Ok(format!("blk.{bid}.attn_output.weight")),
"attn.q_norm.weight" => Ok(format!("blk.{bid}.attn_q_norm.weight")),
"attn.k_norm.weight" => Ok(format!("blk.{bid}.attn_k_norm.weight")),
"attn.rel_logits_proj.proj" | "attn.rel_logits_proj.weight" => {
Ok(format!("blk.{bid}.attn_rel_proj.weight"))
}
"attn.k_sconv.weight" => Ok(format!("blk.{bid}.shortconv_k.weight")),
"attn.v_sconv.weight" => Ok(format!("blk.{bid}.shortconv_v.weight")),
"attn_sconv.weight" => Ok(format!("blk.{bid}.shortconv_attn.weight")),
"mlp_sconv.weight" => Ok(format!("blk.{bid}.shortconv_mlp.weight")),
"mlp_norm.weight" => Ok(format!("blk.{bid}.ffn_norm.weight")),
"mlp.w2_md.weight" => Ok(format!("blk.{bid}.ffn_down.weight")),
"mlp.global_scale" | "mlp.global_scale.weight" => {
Ok(format!("blk.{bid}.ffn_gscale.weight"))
}
"mlp.w13_dn.weight" => {
bail!("Inkling fused w13 tensor requires streaming deinterleave")
}
"self_attn.indexer.k_norm.weight" => Ok(format!("blk.{bid}.indexer.k_norm.weight")),
"self_attn.indexer.k_norm.bias" => Ok(format!("blk.{bid}.indexer.k_norm.bias")),
"self_attn.indexer.weights_proj.weight" => Ok(format!("blk.{bid}.indexer.proj.weight")),
@ -409,4 +489,57 @@ mod tests {
"blk.33.attn_q.weight"
);
}
#[test]
fn recognizes_and_maps_inkling_mtp_tensors() {
let map = TensorNameMap::HfToGgufWithMtp { layer_start: 66 };
for (source, expected) in [
(
"model.mtp.layers.0.embed_norm.weight",
"blk.66.nextn.enorm.weight",
),
(
"model.mtp.layers.1.hidden_norm.weight",
"blk.67.nextn.hnorm.weight",
),
(
"model.mtp.layers.2.input_proj.weight",
"blk.68.nextn.eh_proj.weight",
),
(
"model.mtp.layers.3.transformer_block.attn.wq_du.weight",
"blk.69.attn_q.weight",
),
(
"model.mtp.layers.4.transformer_block.attn.rel_logits_proj.proj",
"blk.70.attn_rel_proj.weight",
),
(
"model.mtp.layers.7.transformer_block.mlp.global_scale",
"blk.73.ffn_gscale.weight",
),
] {
assert!(is_mtp_source_tensor(source));
assert_eq!(map.map_tensor_name(source).unwrap(), expected);
}
assert_eq!(
map.map_tensor_name("model.llm.embed_norm.weight").unwrap(),
"token_embd_norm.weight"
);
assert!(is_shared_mtp_context_tensor("model.llm.unembed.weight"));
assert!(is_inkling_fused_w13(
"model.mtp.layers.0.transformer_block.mlp.w13_dn.weight"
));
}
#[test]
fn rejects_inkling_mtp_layer_overflow() {
let error = TensorNameMap::HfToGgufWithMtp {
layer_start: u32::MAX,
}
.map_tensor_name("model.mtp.layers.1.embed_norm.weight")
.unwrap_err();
assert!(error.to_string().contains("MTP layer id overflow"));
}
}

86
crates/skippy-quantize/src/tokenizer_metadata.rs
View file

@ -8,6 +8,7 @@ use serde_json::Value;
use crate::gguf_writer::GgufKv;
const TOKEN_TYPE_NORMAL: i32 = 1;
const TOKEN_TYPE_UNUSED: i32 = 5;
const TOKEN_TYPE_CONTROL: i32 = 3;
pub(crate) fn push_tokenizer_metadata(
@ -64,7 +65,7 @@ struct BpeVocabMetadata {
added_tokens: BTreeMap<String, u32>,
}
fn read_byte_level_bpe(tokenizer: &Value, _config: &Value) -> Result<BpeVocabMetadata> {
fn read_byte_level_bpe(tokenizer: &Value, config: &Value) -> Result<BpeVocabMetadata> {
let model = tokenizer
.get("model")
.and_then(Value::as_object)
@ -86,7 +87,16 @@ fn read_byte_level_bpe(tokenizer: &Value, _config: &Value) -> Result<BpeVocabMet
.and_then(Value::as_object)
.context("tokenizer.json model missing object field vocab")?;
let added_tokens = collect_added_tokens(tokenizer);
let vocab_size = tokenizer_vocab_size(raw_vocab, &added_tokens)?;
let tokenizer_vocab_size = tokenizer_vocab_size(raw_vocab, &added_tokens)?;
let inkling_unpadded_vocab = inkling_text_u32(config, "unpadded_vocab_size")
.and_then(|value| usize::try_from(value).ok());
let vocab_size = inkling_text_u32(config, "vocab_size")
.and_then(|value| usize::try_from(value).ok())
.unwrap_or(tokenizer_vocab_size);
ensure!(
vocab_size >= 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<BpeVocabMet
}
}
if let Some(unpadded_vocab) = inkling_unpadded_vocab {
ensure!(
unpadded_vocab <= vocab_size,
"Inkling unpadded_vocab_size exceeds vocab_size"
);
for index in unpadded_vocab..vocab_size {
ensure!(
tokens[index].is_empty(),
"real token found at/above Inkling unpadded_vocab_size: {index}"
);
tokens[index] = format!("[PAD{index}]");
token_types[index] = TOKEN_TYPE_UNUSED;
}
}
let missing = tokens.iter().position(String::is_empty);
ensure!(
missing.is_none(),
@ -238,6 +263,9 @@ fn tokenizer_pre(config: &Value) -> 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<u32> {
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<GgufKv>,
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();

10
crates/skippy-quantize/src/types.rs
View file

@ -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<String> {
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()));

3
crates/skippy-runtime/src/activation.rs
View file

@ -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],

15
crates/skippy-runtime/src/media.rs
View file

@ -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,
});

4
crates/skippy-runtime/src/package.rs
View file

@ -104,6 +104,7 @@ pub struct PackageWindowPolicyInfo {
pub initial_window: u32,
pub min_window: u32,
pub max_window: u32,
pub pipeline_depth: Option<u32>,
}
#[derive(Debug, Clone)]
@ -288,6 +289,8 @@ struct PackageWindowPolicy {
initial_window: u32,
min_window: u32,
max_window: u32,
#[serde(default)]
pipeline_depth: Option<u32>,
}
#[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,

2
crates/skippy-runtime/src/runtime_events.rs
View file

@ -355,7 +355,7 @@ fn abi_features_bitmask() -> Option<u64> {
}
#[cfg(not(feature = "dynamic-native-runtime"))]
{
Some(unsafe { skippy_ffi::skippy_abi_features() })
Some(skippy_ffi::abi_features())
}
}

18
crates/skippy-runtime/src/session.rs
View file

@ -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<Vec<i32>> {
if token_ids.is_empty() {
return Ok(Vec::new());

1
crates/skippy-runtime/src/types.rs
View file

@ -233,6 +233,7 @@ pub struct MediaPrefill {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MediaPrefillChunkFrame {
pub token_count: usize,
pub tokens: Vec<i32>,
pub positions: Vec<i32>,
pub output: ActivationFrame,
}

13
crates/skippy-server/README.md
View file

@ -15,7 +15,7 @@ mesh/openai-frontend; diagnostic and benchmark clients may connect directly to
the first stage.
The full request/reply path is tip-to-tip: token IDs enter at the driver-facing
tip, and activations flow through the stage chain. Stage protocol generation 3
tip, and activations flow through the stage chain. Stage protocol generation 4
is a compatibility-breaking contract: prediction-bearing replies return
directly from the final/readout tip to the driver-facing stage instead of being
relayed back through intermediate stages. Middle-out is the prefill optimization
@ -80,10 +80,11 @@ unload or replan.
## Notes
- `serve-binary` is the tuned binary stage-to-stage path.
- `serve-binary` participates in the breaking generation-3 stage protocol.
Stage compatibility requires `stage-generation-3`; direct prediction return is
part of that generation's contract, so older chained-reply peers are rejected
during split planning instead of being mixed into a generation-3 topology.
- `serve-binary` participates in the breaking generation-4 stage protocol.
Stage compatibility requires `stage-generation-4`; direct prediction return and
exact verify-checkpoint retirement are part of that generation's contract, so
older peers are rejected during split planning instead of being mixed into a
generation-4 topology.
- `serve-binary` accepts upstream protocol connections concurrently. Model
execution remains serialized by the per-process runtime lock, but readiness,
abandoned, or broken connections do not monopolize the listener and block the
@ -99,7 +100,7 @@ unload or replan.
`/v1/completions` using the shared `openai-frontend` crate for a local
final/single-stage config with no downstream peer. Split serving uses
embedded stage-0 OpenAI serving from `serve-binary --openai-bind-addr` because
generation-3 prediction returns flow directly from the final stage to stage 0.
generation-4 prediction returns flow directly from the final stage to stage 0.
The older standalone `serve-openai --first-stage-addr` adapter is no longer
supported. `--model-id` is the exact served model id to advertise
and accept, for example `org/repo:Q4_K_M`; it is not parsed as stage topology.

12
crates/skippy-server/src/binary_transport/binary_messaging.rs
View file

@ -11,8 +11,8 @@ use std::{
};
use super::stage_execution::{
consume_optional_client_ready_hello, prepare_binary_stage_connection,
take_warm_or_connect_downstream, warm_downstream_preconnect_enabled,
consume_optional_client_ready_hello, prepare_binary_stage_connection, take_ready_downstream,
warm_downstream_preconnect_enabled,
};
use super::{
decode_batcher::DecodeFrameBatcher,
@ -34,8 +34,12 @@ use skippy_protocol::binary::{WireMessageKind, read_stage_message, send_ready};
pub(in crate::binary_transport) mod async_forwarder;
mod connection;
mod control_messages;
mod message_receive;
mod prefill_recording;
pub(in crate::binary_transport) mod reply;
mod session_lifecycle;
mod session_tracker;
mod summary;
mod telemetry;
@ -162,7 +166,7 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc<AtomicBool>) -> R
speculative: openai_options.speculative.clone(),
native_mtp_enabled: native_mtp_enabled
&& openai_options.speculative.native_mtp.enabled,
native_mtp_draft_model_path: None,
native_mtp_draft_model_path: openai_options.native_mtp_draft_model_path,
native_mtp_max_tokens: openai_options.native_mtp_max_tokens,
native_mtp_min_tokens: openai_options.native_mtp_min_tokens,
activation_width,
@ -243,7 +247,7 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc<AtomicBool>) -> 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,

75
crates/skippy-server/src/binary_transport/binary_messaging/async_forwarder.rs
View file

@ -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() {

758
crates/skippy-server/src/binary_transport/binary_messaging/connection.rs
View file

@ -1,11 +1,21 @@
use super::async_forwarder::AsyncForwarder;
use super::reply::drain_deferred_prefill_replies;
use super::control_messages::{
handle_generation_control, handle_prefix_cache_control, handle_session_control, handle_stop,
handle_verify_retirement,
};
use super::message_receive::{next_connection_session_id, receive_next_message};
use super::reply::reply_window_for_message;
use super::reply::send_stage_reply;
use super::reply::{configure_prediction_return_stream, reply_window_for_message};
use super::session_lifecycle::align_session_to_message;
use super::session_tracker::{
ConnectionSessionTracker, combine_connection_and_cleanup_results,
release_tracked_connection_sessions,
};
use super::summary::BinaryMessageObservation;
use super::summary::BinaryRequestSummary;
use super::telemetry::UpstreamReplyWriteSpan;
use super::telemetry::{
BinaryMessageTiming, emit_binary_message_received, emit_binary_message_timing,
emit_upstream_reply_write_span, insert_runtime_session_stats, record_prefill_edge_transport,
record_verify_window_timing,
};
@ -16,23 +26,19 @@ use crate::binary_transport::binary_kv::accumulate_prefill_tokens;
use crate::binary_transport::binary_kv::add_binary_record_stats;
use crate::binary_transport::binary_kv::emit_binary_proactive_eviction;
use crate::binary_transport::binary_kv::maybe_lookup_binary_prefill;
use crate::binary_transport::binary_kv::maybe_prefix_cache_control;
use crate::binary_transport::binary_kv::maybe_record_binary_full_prefill;
use crate::binary_transport::direct_return;
use crate::binary_transport::direct_return::PredictionReturnSinks;
use crate::binary_transport::forwarded_stage_message_timed;
use crate::binary_transport::kv_eviction::binary_proactive_eviction_plan;
use crate::binary_transport::kv_eviction::evict_binary_resident_prefix_for_decode;
use crate::binary_transport::restore_prefill_decode::handle_binary_restore_prefill_decode_control;
use crate::binary_transport::run_binary_stage_message;
use crate::binary_transport::send_client_ready_hello_if_enabled;
use crate::binary_transport::stage_execution::binary_message_attrs;
use crate::binary_transport::stage_execution::binary_message_session_id;
use crate::binary_transport::stage_execution::decode_record_tokens_sideband;
use crate::binary_transport::stage_execution::elapsed_ms;
use crate::binary_transport::stage_execution::empty_activation_frame;
use crate::binary_transport::stage_execution::input_activation_frame;
use crate::binary_transport::stage_execution::insert_optional_unix_nanos;
use crate::binary_transport::stage_execution::is_decode_frame_batch_candidate;
use crate::binary_transport::stage_execution::nanos_delta_ms;
use crate::binary_transport::stage_execution::runtime_sampling_config;
@ -41,39 +47,77 @@ use crate::binary_transport::stage_execution::stage_mask;
use crate::binary_transport::stage_execution::token_sideband_or_fill;
use crate::binary_transport::stage_output_activation_capacity;
use crate::binary_transport::write_stage_message_conditioned;
use crate::kv_integration::{KvStageIntegration, model_requires_recurrent_state};
use crate::kv_integration::KvStageIntegration;
use crate::runtime_state::RuntimeState;
use crate::telemetry::Telemetry;
use crate::telemetry::now_unix_nanos;
use anyhow::Context;
use anyhow::Result;
use anyhow::bail;
use anyhow::{Context, Result, bail};
use serde_json::json;
use skippy_protocol::StageConfig;
use skippy_protocol::StageTopology;
use skippy_protocol::binary::StageReply;
use skippy_protocol::binary::StageReplyStats;
use skippy_protocol::binary::StageWireMessage;
use skippy_protocol::binary::WireActivationDType;
use skippy_protocol::binary::WireMessageKind;
use skippy_protocol::binary::WireReplyKind;
use skippy_protocol::binary::read_stage_message;
use skippy_protocol::binary::recv_reply;
use skippy_protocol::binary::send_reply_ack;
use skippy_protocol::binary::send_reply_ack_with_stats;
use skippy_protocol::{StageConfig, StageTopology};
use std::collections::BTreeMap;
use std::io;
use std::net::TcpStream;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::sync::{Arc, Mutex};
use std::time::Instant;
static BINARY_SESSION_COUNTER: AtomicU64 = AtomicU64::new(1);
#[allow(clippy::too_many_arguments)]
pub(super) fn handle_binary_connection(
config: &StageConfig,
topology: Option<&StageTopology>,
runtime: &Arc<Mutex<RuntimeState>>,
decode_frame_batcher: &DecodeFrameBatcher,
kv: Option<&Arc<KvStageIntegration>>,
telemetry: &Telemetry,
upstream: &mut TcpStream,
downstream: Option<TcpStream>,
activation_width: i32,
wire_dtype: WireActivationDType,
max_inflight: usize,
reply_credit_limit: Option<usize>,
async_prefill_forward: bool,
downstream_wire_condition: WireCondition,
downstream_connect_timeout_secs: u64,
native_mtp_enabled: bool,
prediction_return_sinks: &PredictionReturnSinks,
first_message: StageWireMessage,
) -> Result<()> {
let mut session_tracker = ConnectionSessionTracker::default();
let result = handle_binary_connection_messages(
config,
topology,
runtime,
decode_frame_batcher,
kv,
telemetry,
upstream,
downstream,
activation_width,
wire_dtype,
max_inflight,
reply_credit_limit,
async_prefill_forward,
downstream_wire_condition,
downstream_connect_timeout_secs,
native_mtp_enabled,
prediction_return_sinks,
first_message,
&mut session_tracker,
);
let cleanup_result =
release_tracked_connection_sessions(config, runtime, telemetry, &mut session_tracker);
combine_connection_and_cleanup_results(result, cleanup_result)
}
#[allow(clippy::too_many_arguments)]
fn handle_binary_connection_messages(
config: &StageConfig,
topology: Option<&StageTopology>,
runtime: &Arc<Mutex<RuntimeState>>,
@ -92,16 +136,9 @@ pub(super) fn handle_binary_connection(
native_mtp_enabled: bool,
prediction_return_sinks: &PredictionReturnSinks,
first_message: StageWireMessage,
session_tracker: &mut ConnectionSessionTracker,
) -> Result<()> {
if let Some(downstream) = downstream.as_mut() {
send_client_ready_hello_if_enabled(&mut *downstream)
.context("send downstream client ready hello")?;
skippy_protocol::binary::recv_ready(&mut *downstream)
.context("downstream binary stage did not become ready")?;
}
let connection_session_id = BINARY_SESSION_COUNTER.fetch_add(1, Ordering::Relaxed);
let positional_speculation_supported = !model_requires_recurrent_state(config);
let connection_session_id = next_connection_session_id();
let max_deferred_prefill_replies =
reply_credit_limit.unwrap_or_else(|| max_inflight.saturating_sub(1));
let mut pending_prefill_replies = 0usize;
@ -125,20 +162,15 @@ pub(super) fn handle_binary_connection(
loop {
let recv_start_unix_nanos = now_unix_nanos() as u64;
let recv_started = Instant::now();
let mut message = if let Some(message) = next_message.take() {
message
} else {
match read_stage_message(&mut *upstream, activation_width) {
Ok(message) => message,
Err(error)
if error.kind() == io::ErrorKind::UnexpectedEof
&& pending_prefill_replies == 0
&& request_summary.message_count == 0 =>
{
return Ok(());
}
Err(error) => return Err(error).context("read binary stage message"),
}
let Some(mut message) = receive_next_message(
upstream,
activation_width,
next_message.take(),
pending_prefill_replies,
request_summary.message_count,
)?
else {
return Ok(());
};
let recv_end_unix_nanos = now_unix_nanos() as u64;
let recv_read_ms = elapsed_ms(recv_started);
@ -146,325 +178,115 @@ pub(super) fn handle_binary_connection(
let message_started = Instant::now();
let session_id = binary_message_session_id(connection_session_id, &message);
let session_key = session_id.to_string();
if message.kind == WireMessageKind::VerifyWindow && !positional_speculation_supported {
bail!(
"stage-state v10 positional speculation requires an attention-only stage; {} contains recurrent state",
config.stage_id
);
}
if telemetry.is_debug_enabled() {
let mut recv_attrs = binary_message_attrs(config, session_id, &message);
recv_attrs.insert(
"llama_stage.recv_start_unix_nanos".to_string(),
json!(recv_start_unix_nanos),
);
recv_attrs.insert(
"llama_stage.recv_end_unix_nanos".to_string(),
json!(recv_end_unix_nanos),
);
recv_attrs.insert("llama_stage.recv_read_ms".to_string(), json!(recv_read_ms));
recv_attrs.insert(
"skippy.upstream_message_wait_ms".to_string(),
json!(recv_read_ms),
);
recv_attrs.insert(
"llama_stage.source_stage_index".to_string(),
json!(message.state.source_stage_index),
);
recv_attrs.insert(
"llama_stage.configured_upstream_stage_index".to_string(),
json!(config.upstream.as_ref().map(|peer| peer.stage_index)),
);
recv_attrs.insert(
"llama_stage.message_wire_bytes".to_string(),
json!(message.estimated_wire_bytes()),
);
recv_attrs.insert(
"skippy.activation_bytes".to_string(),
json!(message.activation.len()),
);
telemetry.emit_debug_span(
"stage.binary_recv",
recv_attrs,
recv_start_unix_nanos,
recv_end_unix_nanos,
);
}
session_tracker.touch(&session_key);
emit_binary_message_received(
telemetry,
config,
session_id,
&message,
recv_start_unix_nanos,
recv_end_unix_nanos,
recv_read_ms,
);
if message.kind == WireMessageKind::Stop {
if pending_prefill_replies != 0 {
bail!("cannot stop with {pending_prefill_replies} deferred prefill replies");
}
let mut stop_stats = std::mem::take(&mut pending_reply_stats);
request_summary.emit(telemetry, config, session_id);
request_summary = BinaryRequestSummary::default();
if let Some(downstream) = downstream.as_mut() {
if let Some(forwarder) = async_forwarder.as_mut() {
forwarder
.flush()
.context("flush async forwards before stop")?;
}
write_stage_message_conditioned(
&mut *downstream,
&message,
wire_dtype,
downstream_wire_condition,
)
.context("forward binary stop")?;
let reply = recv_reply(&mut *downstream).context("stop downstream ACK")?;
if reply.kind != WireReplyKind::Ack {
bail!("stop expected downstream ACK");
}
stop_stats.merge(reply.stats);
}
let reset_start_unix_nanos = now_unix_nanos() as u64;
let reset_timer = Instant::now();
let lock_timer = Instant::now();
let mut runtime = runtime.lock().expect("runtime lock poisoned");
let runtime_lock_wait_ms = elapsed_ms(lock_timer);
let accumulated = std::mem::take(&mut accumulated_prefill_tokens);
for (prefill_session_key, tokens) in accumulated {
let record = maybe_record_binary_full_prefill(
config,
&mut runtime,
kv,
telemetry,
&prefill_session_key,
&message,
&tokens,
);
if record.recorded_pages > 0 {
stop_stats.kv_recorded_pages += record.recorded_pages as i64;
stop_stats.kv_record_stage_mask |= stage_mask(config.stage_index);
}
}
let drop_stats = runtime
.drop_session_timed(&session_key)
.context("reset binary stage session")?;
drop(runtime);
let reset_end_unix_nanos = now_unix_nanos() as u64;
let mut reset_attrs = binary_message_attrs(config, session_id, &message);
reset_attrs.insert(
"llama_stage.runtime_lock_wait_ms".to_string(),
json!(runtime_lock_wait_ms),
);
reset_attrs.insert(
"llama_stage.session_reset_ms".to_string(),
json!(drop_stats.reset_ms),
);
reset_attrs.insert(
"llama_stage.session_reset".to_string(),
json!(drop_stats.reset_session),
);
reset_attrs.insert(
"llama_stage.lane_discarded".to_string(),
json!(drop_stats.lane_discarded),
);
if let Some(reason) = drop_stats.lane_discard_reason.as_deref() {
reset_attrs.insert("llama_stage.lane_discard_reason".to_string(), json!(reason));
}
reset_attrs.insert(
"llama_stage.elapsed_ms".to_string(),
json!(elapsed_ms(reset_timer)),
);
insert_runtime_session_stats(
&mut reset_attrs,
"llama_stage.runtime_sessions_after",
&drop_stats.stats_after,
);
telemetry.emit_debug_span(
"stage.binary_session_stop",
reset_attrs,
reset_start_unix_nanos,
reset_end_unix_nanos,
);
prediction_return_streams.remove(&(message.request_id, message.session_id));
prediction_return_sinks.remove(message.request_id, message.session_id);
send_reply_ack_with_stats(&mut *upstream, stop_stats).context("send stop ACK")?;
continue;
}
if message.kind.is_session_control() {
let mut control_stats = std::mem::take(&mut pending_reply_stats);
if let Some(forwarder) = async_forwarder.as_mut() {
forwarder
.flush()
.context("flush async forwards before session control")?;
}
drain_deferred_prefill_replies(
downstream.as_mut(),
&mut pending_prefill_replies,
&mut control_stats,
)
.context("drain deferred replies before session control")?;
{
let mut runtime = runtime.lock().expect("runtime lock poisoned");
match message.kind {
WireMessageKind::TrimSession => runtime
.trim_session(&session_key, message.token_count.max(0) as u64)
.context("trim binary stage session")?,
_ => unreachable!("session control checked above"),
}
}
if let Some(downstream) = downstream.as_mut() {
write_stage_message_conditioned(
&mut *downstream,
&message,
wire_dtype,
downstream_wire_condition,
)
.context("forward session control")?;
let reply =
recv_reply(&mut *downstream).context("session control downstream ACK")?;
if reply.kind != WireReplyKind::Ack {
bail!("session control expected downstream ACK");
}
control_stats.merge(reply.stats);
}
send_reply_ack_with_stats(&mut *upstream, control_stats)
.context("session control ack")?;
continue;
}
if message.kind.is_generation_control() {
let mut generation_stats = std::mem::take(&mut pending_reply_stats);
if let Some(forwarder) = async_forwarder.as_mut() {
forwarder
.flush()
.context("flush async forwards before generation config")?;
}
drain_deferred_prefill_replies(
downstream.as_mut(),
&mut pending_prefill_replies,
&mut generation_stats,
)
.context("drain deferred replies before generation config")?;
if let Some(downstream) = downstream.as_mut() {
write_stage_message_conditioned(
&mut *downstream,
&message,
wire_dtype,
downstream_wire_condition,
)
.context("forward generation config")?;
let reply =
recv_reply(&mut *downstream).context("generation config downstream ACK")?;
if reply.kind != WireReplyKind::Ack {
bail!("generation config expected downstream ACK");
}
generation_stats.merge(reply.stats);
} else {
if let Some(metadata) = message.chat_sampling_metadata.as_deref() {
let sampling = runtime_sampling_config(message.sampling.as_ref());
let mut runtime = runtime.lock().expect("runtime lock poisoned");
runtime
.configure_chat_sampling(
&session_key,
metadata,
message.state.prompt_token_count.max(0) as u64,
sampling.as_ref(),
)
.context("configure binary stage generation")?;
}
configure_prediction_return_stream(
config,
topology,
message.request_id,
message.session_id,
wire_dtype,
downstream_connect_timeout_secs,
prediction_return_sinks,
&mut prediction_return_streams,
);
}
send_reply_ack_with_stats(&mut *upstream, generation_stats)
.context("generation config ack")?;
continue;
}
if message.kind.is_prefix_cache_control() {
let control_started = Instant::now();
let mut control_stats = std::mem::take(&mut pending_reply_stats);
if let Some(forwarder) = async_forwarder.as_mut() {
forwarder
.flush()
.context("flush async forwards before prefix cache control")?;
}
drain_deferred_prefill_replies(
downstream.as_mut(),
&mut pending_prefill_replies,
&mut control_stats,
)
.context("drain deferred replies before prefix cache control")?;
if message.kind == WireMessageKind::TryRestorePrefillDecode {
handle_binary_restore_prefill_decode_control(
config,
topology,
runtime,
kv,
telemetry,
&session_key,
session_id,
message,
downstream.as_mut(),
wire_dtype,
downstream_wire_condition,
activation_width,
control_started,
control_stats,
prediction_return_sinks,
&mut prediction_return_streams,
downstream_connect_timeout_secs,
native_mtp_enabled,
)
.context("handle restore-prefill-decode control")?;
continue;
}
let token_ids = token_sideband_or_fill(&message)
.context("read prefix cache control token sideband")?;
let local = maybe_prefix_cache_control(
handle_stop(
config,
runtime,
kv,
telemetry,
&session_key,
upstream,
downstream.as_mut(),
wire_dtype,
downstream_wire_condition,
&message,
&token_ids,
);
control_stats.merge(local.stats);
if local.hit
&& let Some(downstream) = downstream.as_mut()
{
write_stage_message_conditioned(
&mut *downstream,
&message,
wire_dtype,
downstream_wire_condition,
)
.context("forward prefix cache control")?;
let reply = recv_reply(&mut *downstream).context("prefix cache downstream ACK")?;
if reply.kind != WireReplyKind::Ack {
bail!("prefix cache control expected downstream ACK");
}
let downstream_missed = message.kind == WireMessageKind::TryRestorePrefill
&& (reply.stats.kv_lookup_misses > 0
|| reply.stats.kv_lookup_errors > 0
|| reply.stats.kv_lookup_hits == 0);
control_stats.merge(reply.stats);
if downstream_missed {
let mut runtime = runtime.lock().expect("runtime lock poisoned");
let _ = runtime.drop_session_timed(&session_key);
}
}
let mut attrs = binary_message_attrs(config, session_id, &message);
attrs.insert("skippy.kv.control_hit".to_string(), json!(local.hit));
attrs.insert(
"llama_stage.elapsed_ms".to_string(),
json!(elapsed_ms(control_started)),
);
telemetry.emit_debug("stage.binary_prefix_cache_control", attrs);
send_reply_ack_with_stats(&mut *upstream, control_stats)
.context("prefix cache control ack")?;
&session_key,
session_id,
pending_prefill_replies,
&mut pending_reply_stats,
&mut request_summary,
&mut accumulated_prefill_tokens,
async_forwarder.as_mut(),
session_tracker,
&mut prediction_return_streams,
prediction_return_sinks,
)?;
continue;
}
if message.kind.is_verify_retirement() {
handle_verify_retirement(
runtime,
downstream.as_mut(),
wire_dtype,
downstream_wire_condition,
&message,
&session_key,
async_forwarder.as_mut(),
)?;
continue;
}
if message.kind.is_session_control() {
handle_session_control(
runtime,
upstream,
downstream.as_mut(),
wire_dtype,
downstream_wire_condition,
&message,
&session_key,
&mut pending_prefill_replies,
&mut pending_reply_stats,
async_forwarder.as_mut(),
)?;
continue;
}
if message.kind.is_generation_control() {
handle_generation_control(
config,
topology,
runtime,
upstream,
downstream.as_mut(),
wire_dtype,
downstream_wire_condition,
downstream_connect_timeout_secs,
&message,
&session_key,
&mut pending_prefill_replies,
&mut pending_reply_stats,
async_forwarder.as_mut(),
prediction_return_sinks,
&mut prediction_return_streams,
)?;
continue;
}
if message.kind.is_prefix_cache_control() {
handle_prefix_cache_control(
config,
topology,
runtime,
kv,
telemetry,
upstream,
downstream.as_mut(),
wire_dtype,
downstream_wire_condition,
downstream_connect_timeout_secs,
activation_width,
native_mtp_enabled,
message,
&session_key,
session_id,
&mut pending_prefill_replies,
&mut pending_reply_stats,
async_forwarder.as_mut(),
prediction_return_sinks,
&mut prediction_return_streams,
)?;
continue;
}
@ -495,37 +317,17 @@ pub(super) fn handle_binary_connection(
}
let token_ids = token_sideband_or_fill(&message)?;
let mut session_auto_align_count = 0usize;
let mut session_auto_align_ms = 0.0;
let mut session_auto_align_trimmed_tokens = 0u64;
if let Some(target_token_count) = message.authoritative_session_position() {
let align_started = Instant::now();
let align = {
let mut runtime = runtime.lock().expect("runtime lock poisoned");
runtime
.align_session_to_token_count_if_ahead(&session_key, target_token_count)
.context("auto-align binary stage session")?
};
if let Some(align) = align {
let align_ms = elapsed_ms(align_started);
session_auto_align_count = 1;
session_auto_align_ms = align_ms;
session_auto_align_trimmed_tokens = align
.before_token_count
.saturating_sub(align.after_token_count);
let mut attrs = binary_message_attrs(config, session_id, &message);
attrs.insert(
"llama_stage.session_auto_align_before_tokens".to_string(),
json!(align.before_token_count),
);
attrs.insert(
"llama_stage.session_auto_align_after_tokens".to_string(),
json!(align.after_token_count),
);
attrs.insert("llama_stage.elapsed_ms".to_string(), json!(align_ms));
telemetry.emit_debug("stage.binary_session_auto_align", attrs);
}
}
let auto_align = align_session_to_message(
config,
runtime,
telemetry,
&session_key,
session_id,
&message,
)?;
let session_auto_align_count = auto_align.count;
let session_auto_align_ms = auto_align.elapsed_ms;
let session_auto_align_trimmed_tokens = auto_align.trimmed_tokens;
if message.kind.is_prefill() {
accumulate_prefill_tokens(
&mut accumulated_prefill_tokens,
@ -1149,154 +951,44 @@ pub(super) fn handle_binary_connection(
upstream_message_wait_ms: recv_read_ms,
});
if telemetry.is_debug_enabled() {
let mut timing_attrs = binary_message_attrs(config, session_id, &message);
timing_attrs.insert(
"llama_stage.message_start_unix_nanos".to_string(),
json!(message_start_unix_nanos),
);
timing_attrs.insert(
"llama_stage.message_end_unix_nanos".to_string(),
json!(message_end_unix_nanos),
);
timing_attrs.insert(
"llama_stage.compute_start_unix_nanos".to_string(),
json!(compute_start_unix_nanos),
);
timing_attrs.insert(
"llama_stage.compute_end_unix_nanos".to_string(),
json!(compute_end_unix_nanos),
);
timing_attrs.insert("llama_stage.compute_ms".to_string(), json!(compute_ms));
timing_attrs.insert("llama_stage.recv_read_ms".to_string(), json!(recv_read_ms));
timing_attrs.insert(
"skippy.upstream_message_wait_ms".to_string(),
json!(recv_read_ms),
);
timing_attrs.insert(
"llama_stage.input_activation_decode_ms".to_string(),
json!(input_activation_decode_ms),
);
timing_attrs.insert(
"llama_stage.runtime_lock_wait_ms".to_string(),
json!(runtime_lock_wait_ms),
);
timing_attrs.insert(
"llama_stage.runtime_lock_hold_ms".to_string(),
json!(runtime_lock_hold_ms),
);
timing_attrs.insert(
"llama_stage.runtime_lock_acquires".to_string(),
json!(runtime_lock_acquires),
);
if let Some(stats) = runtime_sessions_before.as_ref() {
insert_runtime_session_stats(
&mut timing_attrs,
"llama_stage.runtime_sessions_before",
stats,
);
}
if let Some(stats) = runtime_sessions_after.as_ref() {
insert_runtime_session_stats(
&mut timing_attrs,
"llama_stage.runtime_sessions_after",
stats,
);
}
timing_attrs.insert(
"llama_stage.forward_write_ms".to_string(),
json!(forward_write_ms),
);
timing_attrs.insert(
"llama_stage.activation_encode_ms".to_string(),
json!(forward_activation_encode_ms),
);
timing_attrs.insert(
"llama_stage.downstream_wait_ms".to_string(),
json!(downstream_wait_ms),
);
timing_attrs.insert("skippy.compute_ms".to_string(), json!(compute_ms));
timing_attrs.insert(
"skippy.forward_write_ms".to_string(),
json!(forward_write_ms),
);
timing_attrs.insert(
"skippy.downstream_wait_ms".to_string(),
json!(downstream_wait_ms),
);
timing_attrs.insert(
"skippy.upstream_reply_ms".to_string(),
json!(upstream_reply_ms),
);
timing_attrs.insert("llama_stage.forward_mode".to_string(), json!(forward_mode));
insert_optional_unix_nanos(
&mut timing_attrs,
"llama_stage.forward_write_start_unix_nanos",
forward_write_start_unix_nanos,
);
insert_optional_unix_nanos(
&mut timing_attrs,
"llama_stage.forward_write_end_unix_nanos",
forward_write_end_unix_nanos,
);
insert_optional_unix_nanos(
&mut timing_attrs,
"llama_stage.downstream_wait_start_unix_nanos",
downstream_wait_start_unix_nanos,
);
insert_optional_unix_nanos(
&mut timing_attrs,
"llama_stage.downstream_wait_end_unix_nanos",
downstream_wait_end_unix_nanos,
);
insert_optional_unix_nanos(
&mut timing_attrs,
"llama_stage.upstream_reply_start_unix_nanos",
upstream_reply_start_unix_nanos,
);
insert_optional_unix_nanos(
&mut timing_attrs,
"llama_stage.upstream_reply_end_unix_nanos",
upstream_reply_end_unix_nanos,
);
timing_attrs.insert(
"skippy.message_elapsed_ms".to_string(),
json!(message_elapsed_ms),
);
timing_attrs.insert(
"skippy.input_activation_bytes".to_string(),
json!(input_activation_bytes),
);
timing_attrs.insert(
"skippy.output_activation_bytes".to_string(),
json!(output.payload.len()),
);
timing_attrs.insert(
"skippy.prefill_credit_limit".to_string(),
json!(max_deferred_prefill_replies),
);
timing_attrs.insert(
"skippy.prefill_pending_replies_before".to_string(),
json!(pending_prefill_replies_before),
);
timing_attrs.insert(
"skippy.prefill_pending_replies_after".to_string(),
json!(pending_prefill_replies),
);
timing_attrs.insert(
"skippy.prefill_credit_wait_count".to_string(),
json!(credit_wait_count),
);
timing_attrs.insert(
"skippy.prefill_deferred_replies_drained".to_string(),
json!(deferred_prefill_replies_drained),
);
telemetry.emit_debug_span(
"stage.binary_message_timing",
timing_attrs,
emit_binary_message_timing(
telemetry,
config,
session_id,
&message,
BinaryMessageTiming {
message_start_unix_nanos,
message_end_unix_nanos,
);
}
compute_start_unix_nanos,
compute_end_unix_nanos,
forward_write_start_unix_nanos,
forward_write_end_unix_nanos,
downstream_wait_start_unix_nanos,
downstream_wait_end_unix_nanos,
upstream_reply_start_unix_nanos,
upstream_reply_end_unix_nanos,
compute_ms,
recv_read_ms,
input_activation_decode_ms,
runtime_lock_wait_ms,
runtime_lock_hold_ms,
runtime_lock_acquires,
runtime_sessions_before: runtime_sessions_before.as_ref(),
runtime_sessions_after: runtime_sessions_after.as_ref(),
forward_write_ms,
forward_activation_encode_ms,
downstream_wait_ms,
upstream_reply_ms,
forward_mode,
message_elapsed_ms,
input_activation_bytes,
output_activation_bytes: output.payload.len(),
max_deferred_prefill_replies,
pending_prefill_replies_before,
pending_prefill_replies_after: pending_prefill_replies,
credit_wait_count,
deferred_prefill_replies_drained,
},
);
}
}

407
crates/skippy-server/src/binary_transport/binary_messaging/control_messages.rs Normal file
View file

@ -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<Mutex<RuntimeState>>,
kv: Option<&Arc<KvStageIntegration>>,
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<String, Vec<i32>>,
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<Mutex<RuntimeState>>,
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<Mutex<RuntimeState>>,
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<Mutex<RuntimeState>>,
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<Mutex<RuntimeState>>,
kv: Option<&Arc<KvStageIntegration>>,
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")
}

34
crates/skippy-server/src/binary_transport/binary_messaging/message_receive.rs Normal file
View file

@ -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<StageWireMessage>,
pending_prefill_replies: usize,
observed_message_count: usize,
) -> Result<Option<StageWireMessage>> {
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"),
}
}

Some files were not shown because too many files have changed in this diff Show more