Save inference parity implementation and evaluation harness

src/engine/metal/qwen_mtplx/policy.rs

@@ -0,0 +1,275 @@
+//! Qwen server defaults and the native application's execution settings.
+//! Overrides are model-owned values; never mutate the process environment.
+use super::session_cache::{BoundaryPolicy, DecodeOptions, TurnOptions};
+use super::{ar, prefill, sampling::SamplerConfig, speculative};
+use crate::settings::{EngineSettings, QWEN_DEFAULT_PREFILL_CHUNK};
+use std::collections::BTreeMap;
+
+pub(super) type Overrides = BTreeMap<&'static str, &'static str>;
+
+pub(super) fn truthy(value: Option<&str>) -> bool {
+    value.is_some_and(|s| {
+        matches!(
+            s.trim().to_ascii_lowercase().as_str(),
+            "1" | "true" | "yes" | "on"
+        )
+    })
+}
+
+fn falsey(value: Option<&str>) -> bool {
+    value.is_some_and(|s| {
+        matches!(
+            s.trim().to_ascii_lowercase().as_str(),
+            "0" | "false" | "no" | "off"
+        )
+    })
+}
+
+fn strict_bool(name: &str, raw: Option<String>) -> Result<bool, String> {
+    match raw
+        .as_deref()
+        .unwrap_or_default()
+        .trim()
+        .to_ascii_lowercase()
+        .as_str()
+    {
+        "" | "0" | "false" | "no" | "off" | "disable" | "disabled" => Ok(false),
+        "1" | "true" | "yes" | "on" | "enable" | "enabled" => Ok(true),
+        _ => Err(format!("{name}={raw:?} is not a boolean")),
+    }
+}
+
+pub(super) fn value(overrides: &Overrides, name: &str) -> Option<String> {
+    overrides
+        .get(name)
+        .map(|s| (*s).to_owned())
+        .or_else(|| std::env::var(name).ok())
+}
+
+/// _server_runtime_env_overrides for this family, serial AR/native batched MTP.
+pub(super) fn server_overrides(mtp: bool, get: impl Fn(&str) -> Option<String>) -> Overrides {
+    let mut result = Overrides::new();
+    if mtp {
+        result.insert("MTPLX_SKIP_VERIFY_SNAPSHOT", "0");
+    }
+    for key in [
+        "MTPLX_AR_PIPELINE",
+        "MTPLX_FAMILY_CAPTURE_COMMIT",
+        "MTPLX_FUSED_HC_V3",
+        "MTPLX_FUSED_GDN_INPROJ",
+        "MTPLX_FUSED_GATE_UP",
+        "MTPLX_FUSED_GDN_CONVNORM",
+        "MTPLX_FUSED_GDN_STEP",
+        "MTPLX_FUSED_CONVNORM_VERIFY",
+        "MTPLX_QSA_GATHER",
+    ] {
+        if get(key).is_none() {
+            result.insert(key, "1");
+        }
+    }
+    let compiled = ["MTPLX_COMPILED_GDN", "MTPLX_QWEN4EXP_COMPILE"];
+    if !compiled
+        .iter()
+        .any(|key| get(key).is_some_and(|s| !truthy(Some(&s))))
+    {
+        for key in compiled {
+            if get(key).is_none() {
+                result.insert(key, "1");
+            }
+        }
+    }
+    if get("MTPLX_NAX_VERIFY").is_none() {
+        result.insert("MTPLX_NAX_VERIFY", "0");
+    }
+    result
+}
+
+pub(super) fn integer(raw: Option<&str>, default: usize, minimum: usize) -> Result<usize, String> {
+    let number = raw
+        .and_then(|s| s.trim().replace('_', "").parse::<i128>().ok())
+        .unwrap_or(default as i128)
+        .max(minimum as i128);
+    usize::try_from(number).map_err(|_| "MTPLX setting exceeds the supported integer range".into())
+}
+
+/// Application chunk size is the native request-local override, not an env
+/// write. The bounded/sustained route keeps only the final prefill logits.
+pub(super) fn turn_options(
+    settings: &EngineSettings,
+    get: impl Fn(&str) -> Option<String>,
+) -> Result<(TurnOptions<'static>, Option<BoundaryPolicy>), String> {
+    if !settings.model.is_qwen38() || settings.context_tokens <= 0 {
+        return Err("canonical Qwen requires Qwen settings and a positive context".into());
+    }
+    let flag = |key| truthy(get(key).as_deref());
+    let enabled = |key| !falsey(get(key).as_deref());
+    let number = |key, default, minimum| integer(get(key).as_deref(), default, minimum);
+    let mtp = settings.speculative.glm_mtp;
+    let chunk = if settings.execution.prefill_chunk == 0 {
+        QWEN_DEFAULT_PREFILL_CHUNK
+    } else {
+        settings.execution.prefill_chunk
+    };
+    if chunk > i32::MAX as u32 {
+        return Err("Qwen prefill chunk exceeds int32".into());
+    }
+    let decode = if mtp {
+        DecodeOptions::Mtp(speculative::Options {
+            prompt_ids: &[],
+            depth: 3,
+            max_tokens: usize::MAX,
+            seed: 0,
+            stops: &[],
+            committed_history: true,
+            family_capture: strict_bool(
+                "MTPLX_FAMILY_CAPTURE_COMMIT",
+                get("MTPLX_FAMILY_CAPTURE_COMMIT"),
+            )?,
+            lazy_history: flag("MTPLX_LAZY_MTP_HISTORY_APPEND"),
+            defer_repair_eval: strict_bool(
+                "MTPLX_DEFER_REPAIR_EVAL",
+                get("MTPLX_DEFER_REPAIR_EVAL"),
+            )?,
+            capture_final_state: true,
+        })
+    } else {
+        DecodeOptions::Ar(ar::Options {
+            max_tokens: usize::MAX,
+            seed: 0,
+            stop_tokens: &[],
+            pipeline: flag("MTPLX_AR_PIPELINE"),
+            sync_eval: !flag("MTPLX_ASYNC_AR")
+                || get("MTPLX_EVAL_AUDIT").is_some_and(|s| !s.is_empty()),
+            return_hidden: false,
+            capture_final_state: true,
+        })
+    };
+    let boundaries = if enabled("MTPLX_GDN_BOUNDARY_CAPTURE") {
+        Some(BoundaryPolicy {
+            tail_interval: number("MTPLX_GDN_BOUNDARY_TAIL_INTERVAL", 256, 0)?,
+            max_count: number("MTPLX_GDN_BOUNDARY_MAX", 8, 2)?,
+        })
+    } else {
+        None
+    };
+    Ok((
+        TurnOptions {
+            prefill: prefill::Options {
+                chunk_size: Some(chunk as usize),
+                final_logits_only: true,
+                external_cache_only: flag("MTPLX_PREFILL_OMLX_EXTERNAL")
+                    || (flag("MTPLX_PREFILL_STOCK_CACHE_ONLY")
+                        && flag("MTPLX_ALLOW_UNSAFE_PREFILL_STOCK_CACHE_ONLY")),
+                external_emit_logits: enabled("MTPLX_PREFILL_EXTERNAL_EMIT_LOGITS"),
+                return_hidden: mtp,
+            },
+            decode,
+            sampler: SamplerConfig::default(),
+            draft_sampler: SamplerConfig::default(),
+            lazy_kv: enabled("MTPLX_SESSION_LAZY_SNAPSHOT"),
+            store_min_suffix: if enabled("MTPLX_SESSION_STORE_ON_PREFILL") {
+                Some(number(
+                    "MTPLX_SESSION_STORE_ON_PREFILL_MIN_SUFFIX",
+                    1024,
+                    1,
+                )?)
+            } else {
+                None
+            },
+            fused_max: number("MTPLX_SMALL_SUFFIX_FUSED_MAX", 512, 0)?,
+        },
+        boundaries,
+    ))
+}
+
+#[test]
+#[ignore = "requires DS4_QWEN38_POLICY_REFERENCE from pinned MTPLX; CPU-only"]
+fn qwen_server_policy_matches_original_functions() {
+    let path = std::env::var_os("DS4_QWEN38_POLICY_REFERENCE").unwrap();
+    let contents = std::fs::read_to_string(path).unwrap();
+    let mut count = 0;
+    for case in contents
+        .lines()
+        .filter_map(|line| serde_json::from_str::<serde_json::Value>(line).ok())
+    {
+        if case["event"] != "policy_reference" {
+            continue;
+        }
+        let mtp = case["mtp"].as_bool().unwrap();
+        let lookup = |key: &str| case["environment"][key].as_str().map(str::to_owned);
+        let overrides = server_overrides(mtp, lookup);
+        assert_eq!(
+            serde_json::to_value(&overrides).unwrap(),
+            case["overrides"],
+            "overrides case {count}"
+        );
+        let get = |key: &str| {
+            overrides
+                .get(key)
+                .map(|s| (*s).to_owned())
+                .or_else(|| lookup(key))
+        };
+        let mut settings = crate::settings::effective_settings(
+            crate::model::ModelChoice::Qwen38FlashNext,
+            &crate::settings::GenerationPreferences::default(),
+            &crate::settings::RuntimePreferences::default(),
+            std::path::Path::new("/unused-policy-models"),
+        )
+        .unwrap()
+        .engine;
+        settings.speculative.glm_mtp = mtp;
+        let resolved = turn_options(&settings, get);
+        if case["error"].is_string() {
+            assert!(
+                resolved.is_err(),
+                "invalid configuration accepted in case {count}"
+            );
+            count += 1;
+            continue;
+        }
+        let (actual, boundary) = resolved.unwrap();
+        let expected = &case["options"];
+        assert_eq!(actual.prefill.chunk_size.unwrap(), expected["chunk"]);
+        assert_eq!(actual.prefill.final_logits_only, expected["final_only"]);
+        assert_eq!(
+            actual.prefill.external_cache_only,
+            expected["external_cache_only"]
+        );
+        assert_eq!(
+            actual.prefill.external_emit_logits,
+            expected["external_emit_logits"]
+        );
+        assert_eq!(actual.lazy_kv, expected["lazy_kv"]);
+        assert_eq!(actual.store_min_suffix.is_some(), expected["store"]);
+        if let Some(min) = actual.store_min_suffix {
+            assert_eq!(min, expected["store_min"]);
+        }
+        assert_eq!(actual.fused_max, expected["fused_max"]);
+        assert_eq!(boundary.is_some(), expected["boundaries"]);
+        if let Some(b) = boundary {
+            assert_eq!(b.max_count, expected["boundary_max"]);
+            assert_eq!(b.tail_interval, expected["boundary_tail"]);
+        }
+        match actual.decode {
+            DecodeOptions::Ar(o) => {
+                assert!(!mtp);
+                assert_eq!(o.pipeline, expected["pipeline"]);
+                assert_eq!(o.sync_eval, expected["sync_eval"]);
+            }
+            DecodeOptions::Mtp(o) => {
+                assert!(mtp && o.committed_history && o.capture_final_state);
+                assert_eq!(o.depth, 3);
+                assert_eq!(o.family_capture, expected["capture"]);
+                assert_eq!(o.lazy_history, expected["lazy_history"]);
+                assert_eq!(o.defer_repair_eval, expected["defer_repair"]);
+            }
+        }
+        settings.execution.prefill_chunk = 77;
+        assert_eq!(
+            turn_options(&settings, get).unwrap().0.prefill.chunk_size,
+            Some(77)
+        );
+        count += 1;
+    }
+    assert_eq!(count, 268);
+}