Pipeline MTP-anchored n-gram verify windows

crates/mesh-llm-cli/src/lib.rs

@@ -12,6 +12,7 @@ pub use mesh_llm_events::LogFormat;
 pub use parser::{
     AuthCommand, BinaryFlavor, Cli, Command, ConfigCommand, DiscoveryScope, DoctorCommand,
     GpuCommand, MeshDiscoveryMode, MeshGuardrailCliMode, NormalizedRuntimeArgs, PluginCommand,
-    RuntimeSurface, SkillAgentArg, SkillCommand, TrustCommand, TrustPolicy,
-    legacy_runtime_surface_warning, normalize_runtime_surface_args, validate_discovery_mode_args,
+    RuntimeSurface, SkillAgentArg, SkillCommand, SpeculativeNgramProposerCli, TrustCommand,
+    TrustPolicy, legacy_runtime_surface_warning, normalize_runtime_surface_args,
+    validate_discovery_mode_args,
 };

crates/mesh-llm-cli/src/parser.rs

@@ -395,6 +395,21 @@ pub enum MeshGuardrailCliMode {
     Enforce,
 }
 
+#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)]
+pub enum SpeculativeNgramProposerCli {
+    Simple,
+    Cache,
+}
+
+impl SpeculativeNgramProposerCli {
+    pub fn as_str(self) -> &'static str {
+        match self {
+            Self::Simple => "simple",
+            Self::Cache => "cache",
+        }
+    }
+}
+
 impl MeshGuardrailCliMode {
     pub const fn as_str(self) -> &'static str {
         match self {
@@ -547,6 +562,70 @@ pub struct Cli {
     #[arg(long, hide = true)]
     pub no_draft: bool,
 
+    /// Override the package speculative decoding strategy for this invocation.
+    #[arg(long, hide = true)]
+    pub speculative_strategy: Option<String>,
+
+    /// Override the N-gram proposer kind for this invocation.
+    #[arg(long, value_enum, hide = true)]
+    pub speculative_ngram_proposer: Option<SpeculativeNgramProposerCli>,
+
+    /// Minimum matching N-gram length for a direct N-gram proposer.
+    #[arg(long, hide = true)]
+    pub speculative_ngram_min: Option<u32>,
+
+    /// Maximum matching N-gram length for a direct N-gram proposer.
+    #[arg(long, hide = true)]
+    pub speculative_ngram_max: Option<u32>,
+
+    /// Cap N-gram tokens proposed in one verify window.
+    #[arg(long, hide = true)]
+    pub speculative_ngram_max_proposal_tokens: Option<u32>,
+
+    /// Initial N-gram extension length for a composite MTP strategy.
+    #[arg(long, hide = true)]
+    pub speculative_extension_initial_tokens: Option<u32>,
+
+    /// Maximum N-gram extension length for a composite MTP strategy.
+    #[arg(long, hide = true)]
+    pub speculative_extension_max_tokens: Option<u32>,
+
+    /// Consecutive weak extensions before the composite strategy backs off.
+    #[arg(long, hide = true)]
+    pub speculative_extension_tail_backoff_proposals: Option<u32>,
+
+    /// Native MTP rejection cooldown in generated tokens.
+    #[arg(long, hide = true)]
+    pub speculative_native_mtp_reject_cooldown_tokens: Option<u32>,
+
+    /// Suppress native MTP drafts while its rejection cooldown is active.
+    #[arg(long, hide = true)]
+    pub speculative_native_mtp_suppress_cooldown_drafts: bool,
+
+    /// Keep native MTP drafts during its rejection cooldown.
+    #[arg(
+        long,
+        hide = true,
+        conflicts_with = "speculative_native_mtp_suppress_cooldown_drafts"
+    )]
+    pub speculative_native_mtp_allow_cooldown_drafts: bool,
+
+    /// Maximum native MTP drafts suppressed by a cooldown.
+    #[arg(long, hide = true)]
+    pub speculative_native_mtp_suppress_cooldown_draft_limit: Option<u32>,
+
+    /// Minimum tokens to include in a pipelined verify window.
+    #[arg(long, hide = true)]
+    pub speculative_verify_window_min_tokens: Option<u32>,
+
+    /// Maximum tokens to include in a pipelined verify window.
+    #[arg(long, hide = true)]
+    pub speculative_verify_window_max_tokens: Option<u32>,
+
+    /// Number of in-flight pipelined verify windows.
+    #[arg(long, hide = true)]
+    pub speculative_verify_window_pipeline_depth: Option<u32>,
+
     /// Force tensor split even if the model fits on one node.
     #[arg(long, hide = true)]
     pub split: bool,
@@ -1508,6 +1587,39 @@ mod tests {
         );
     }
 
+    #[test]
+    fn serve_parses_speculative_decode_overrides() {
+        let normalized = normalize_runtime_surface_args([
+            "mesh-llm",
+            "serve",
+            "--speculative-strategy",
+            "mtp-cache",
+            "--speculative-ngram-proposer",
+            "cache",
+            "--speculative-ngram-min",
+            "2",
+            "--speculative-ngram-max",
+            "6",
+            "--speculative-extension-max-tokens",
+            "8",
+            "--speculative-native-mtp-allow-cooldown-drafts",
+            "--speculative-verify-window-pipeline-depth",
+            "3",
+        ]);
+        let cli = Cli::try_parse_from(normalized.normalized).expect("clap parse");
+
+        assert_eq!(cli.speculative_strategy.as_deref(), Some("mtp-cache"));
+        assert_eq!(
+            cli.speculative_ngram_proposer,
+            Some(SpeculativeNgramProposerCli::Cache)
+        );
+        assert_eq!(cli.speculative_ngram_min, Some(2));
+        assert_eq!(cli.speculative_ngram_max, Some(6));
+        assert_eq!(cli.speculative_extension_max_tokens, Some(8));
+        assert!(cli.speculative_native_mtp_allow_cooldown_drafts);
+        assert_eq!(cli.speculative_verify_window_pipeline_depth, Some(3));
+    }
+
     #[test]
     fn legacy_runtime_surface_warning_for_top_level_serve_flags() {
         let normalized =

crates/mesh-llm-config/src/lib.rs

@@ -38,8 +38,8 @@ pub use validate::{
 mod tests {
     use super::{
         ConfigStore, GpuAssignment, LocalServingNodeConfig, MeshConfig, ModelRuntimeKind,
-        built_in_config_schema, canonicalize_built_in_config_identifier, parse_config_toml,
-        validate_config,
+        SpeculativeConfig, built_in_config_schema, canonicalize_built_in_config_identifier,
+        parse_config_toml, validate_config,
     };
     use std::collections::{BTreeMap, BTreeSet};
     use std::fs;
@@ -55,6 +55,31 @@ mod tests {
         assert!(config.models.is_empty());
     }
 
+    #[test]
+    fn speculative_config_precedence_keeps_lower_layer_fields() {
+        let defaults = SpeculativeConfig {
+            strategy: Some("mtp-cache".to_string()),
+            verify_window_pipeline_depth: Some(2),
+            ..Default::default()
+        };
+        let model = SpeculativeConfig {
+            ngram_max_proposal_tokens: Some(6),
+            ..Default::default()
+        };
+        let overrides = SpeculativeConfig {
+            strategy: Some("mtp".to_string()),
+            verify_window_pipeline_depth: Some(3),
+            ..Default::default()
+        };
+
+        let resolved =
+            SpeculativeConfig::with_precedence(Some(&overrides), Some(&model), Some(&defaults));
+
+        assert_eq!(resolved.strategy.as_deref(), Some("mtp"));
+        assert_eq!(resolved.ngram_max_proposal_tokens, Some(6));
+        assert_eq!(resolved.verify_window_pipeline_depth, Some(3));
+    }
+
     #[test]
     fn plugin_startup_config_round_trips_from_toml() {
         let config: MeshConfig = toml::from_str(
@@ -731,7 +756,6 @@ gpu_id = "pci:0000:65:00.0"
             "models",
             "plugins",
             "settings",
-            "strategy",
         ];
 
         let mut total = 0usize;

crates/mesh-llm-config/src/model.rs

@@ -559,10 +559,80 @@ pub struct SpeculativeConfig {
     pub draft_cache_type_v: Option<String>,
     pub ngram_min: Option<u32>,
     pub ngram_max: Option<u32>,
+    pub ngram_proposer: Option<String>,
+    pub ngram_max_proposal_tokens: Option<u32>,
+    pub extension_initial_tokens: Option<u32>,
+    pub extension_max_tokens: Option<u32>,
+    pub extension_tail_backoff_proposals: Option<u32>,
+    pub native_mtp_reject_cooldown_tokens: Option<u32>,
+    pub native_mtp_suppress_cooldown_drafts: Option<bool>,
+    pub native_mtp_suppress_cooldown_draft_limit: Option<u32>,
+    pub verify_window_min_tokens: Option<u32>,
+    pub verify_window_max_tokens: Option<u32>,
+    pub verify_window_pipeline_depth: Option<u32>,
     pub spec_default: Option<BoolOrAuto>,
     pub(crate) legacy_draft_model_path_used: bool,
 }
 
+impl SpeculativeConfig {
+    /// Resolves the three supported policy layers without discarding fields
+    /// that are not overridden by a more specific layer.
+    pub fn with_precedence(
+        overrides: Option<&Self>,
+        model: Option<&Self>,
+        defaults: Option<&Self>,
+    ) -> Self {
+        macro_rules! pick {
+            ($field:ident) => {
+                overrides
+                    .and_then(|config| config.$field.clone())
+                    .or_else(|| model.and_then(|config| config.$field.clone()))
+                    .or_else(|| defaults.and_then(|config| config.$field.clone()))
+            };
+        }
+
+        Self {
+            strategy: pick!(strategy),
+            mode: pick!(mode),
+            draft_model: pick!(draft_model),
+            draft_hf_repo: pick!(draft_hf_repo),
+            draft_hf_file: pick!(draft_hf_file),
+            draft_selection_policy: pick!(draft_selection_policy),
+            pairing_fault: pick!(pairing_fault),
+            draft_max_tokens: pick!(draft_max_tokens),
+            draft_min_tokens: pick!(draft_min_tokens),
+            draft_acceptance_threshold: pick!(draft_acceptance_threshold),
+            draft_split_probability: pick!(draft_split_probability),
+            draft_gpu_layers: pick!(draft_gpu_layers),
+            draft_device: pick!(draft_device),
+            draft_threads: pick!(draft_threads),
+            draft_cache_type_k: pick!(draft_cache_type_k),
+            draft_cache_type_v: pick!(draft_cache_type_v),
+            ngram_min: pick!(ngram_min),
+            ngram_max: pick!(ngram_max),
+            ngram_proposer: pick!(ngram_proposer),
+            ngram_max_proposal_tokens: pick!(ngram_max_proposal_tokens),
+            extension_initial_tokens: pick!(extension_initial_tokens),
+            extension_max_tokens: pick!(extension_max_tokens),
+            extension_tail_backoff_proposals: pick!(extension_tail_backoff_proposals),
+            native_mtp_reject_cooldown_tokens: pick!(native_mtp_reject_cooldown_tokens),
+            native_mtp_suppress_cooldown_drafts: pick!(native_mtp_suppress_cooldown_drafts),
+            native_mtp_suppress_cooldown_draft_limit: pick!(
+                native_mtp_suppress_cooldown_draft_limit
+            ),
+            verify_window_min_tokens: pick!(verify_window_min_tokens),
+            verify_window_max_tokens: pick!(verify_window_max_tokens),
+            verify_window_pipeline_depth: pick!(verify_window_pipeline_depth),
+            spec_default: pick!(spec_default),
+            legacy_draft_model_path_used: overrides
+                .filter(|config| config.draft_model.is_some())
+                .or_else(|| model.filter(|config| config.draft_model.is_some()))
+                .or_else(|| defaults.filter(|config| config.draft_model.is_some()))
+                .is_some_and(|config| config.legacy_draft_model_path_used),
+        }
+    }
+}
+
 /// Raw deserialization helper that accepts both `draft_model` and the legacy
 /// `draft_model_path` key. The public `SpeculativeConfig` is constructed from
 /// this after detecting which key was used.
@@ -608,6 +678,28 @@ struct SpeculativeConfigRaw {
     #[serde(default)]
     ngram_max: Option<u32>,
     #[serde(default)]
+    ngram_proposer: Option<String>,
+    #[serde(default)]
+    ngram_max_proposal_tokens: Option<u32>,
+    #[serde(default)]
+    extension_initial_tokens: Option<u32>,
+    #[serde(default)]
+    extension_max_tokens: Option<u32>,
+    #[serde(default)]
+    extension_tail_backoff_proposals: Option<u32>,
+    #[serde(default)]
+    native_mtp_reject_cooldown_tokens: Option<u32>,
+    #[serde(default)]
+    native_mtp_suppress_cooldown_drafts: Option<bool>,
+    #[serde(default)]
+    native_mtp_suppress_cooldown_draft_limit: Option<u32>,
+    #[serde(default)]
+    verify_window_min_tokens: Option<u32>,
+    #[serde(default)]
+    verify_window_max_tokens: Option<u32>,
+    #[serde(default)]
+    verify_window_pipeline_depth: Option<u32>,
+    #[serde(default)]
     spec_default: Option<BoolOrAuto>,
 }
 
@@ -643,6 +735,17 @@ impl<'de> Deserialize<'de> for SpeculativeConfig {
             draft_cache_type_v: raw.draft_cache_type_v,
             ngram_min: raw.ngram_min,
             ngram_max: raw.ngram_max,
+            ngram_proposer: raw.ngram_proposer,
+            ngram_max_proposal_tokens: raw.ngram_max_proposal_tokens,
+            extension_initial_tokens: raw.extension_initial_tokens,
+            extension_max_tokens: raw.extension_max_tokens,
+            extension_tail_backoff_proposals: raw.extension_tail_backoff_proposals,
+            native_mtp_reject_cooldown_tokens: raw.native_mtp_reject_cooldown_tokens,
+            native_mtp_suppress_cooldown_drafts: raw.native_mtp_suppress_cooldown_drafts,
+            native_mtp_suppress_cooldown_draft_limit: raw.native_mtp_suppress_cooldown_draft_limit,
+            verify_window_min_tokens: raw.verify_window_min_tokens,
+            verify_window_max_tokens: raw.verify_window_max_tokens,
+            verify_window_pipeline_depth: raw.verify_window_pipeline_depth,
             spec_default: raw.spec_default,
             legacy_draft_model_path_used: legacy_used,
         })
@@ -656,7 +759,7 @@ impl Serialize for SpeculativeConfig {
     {
         use serde::ser::SerializeMap;
 
-        let mut map = serializer.serialize_map(Some(21))?;
+        let mut map = serializer.serialize_map(Some(32))?;
         map.serialize_entry("strategy", &self.strategy)?;
         map.serialize_entry("mode", &self.mode)?;
         if self.legacy_draft_model_path_used {
@@ -684,6 +787,32 @@ impl Serialize for SpeculativeConfig {
         map.serialize_entry("draft_cache_type_v", &self.draft_cache_type_v)?;
         map.serialize_entry("ngram_min", &self.ngram_min)?;
         map.serialize_entry("ngram_max", &self.ngram_max)?;
+        map.serialize_entry("ngram_proposer", &self.ngram_proposer)?;
+        map.serialize_entry("ngram_max_proposal_tokens", &self.ngram_max_proposal_tokens)?;
+        map.serialize_entry("extension_initial_tokens", &self.extension_initial_tokens)?;
+        map.serialize_entry("extension_max_tokens", &self.extension_max_tokens)?;
+        map.serialize_entry(
+            "extension_tail_backoff_proposals",
+            &self.extension_tail_backoff_proposals,
+        )?;
+        map.serialize_entry(
+            "native_mtp_reject_cooldown_tokens",
+            &self.native_mtp_reject_cooldown_tokens,
+        )?;
+        map.serialize_entry(
+            "native_mtp_suppress_cooldown_drafts",
+            &self.native_mtp_suppress_cooldown_drafts,
+        )?;
+        map.serialize_entry(
+            "native_mtp_suppress_cooldown_draft_limit",
+            &self.native_mtp_suppress_cooldown_draft_limit,
+        )?;
+        map.serialize_entry("verify_window_min_tokens", &self.verify_window_min_tokens)?;
+        map.serialize_entry("verify_window_max_tokens", &self.verify_window_max_tokens)?;
+        map.serialize_entry(
+            "verify_window_pipeline_depth",
+            &self.verify_window_pipeline_depth,
+        )?;
         map.serialize_entry("spec_default", &self.spec_default)?;
         map.end()
     }