Preserve image-aware sparse QSA for Qwen vision

tools/issue107-model-reference.py

@@ -0,0 +1,48 @@
+"""Pinned MTPLX oracle for issue 107's real long two-image product request."""
+import argparse
+import hashlib
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+checkout, root, image, output = map(Path, sys.argv[1:5])
+assert subprocess.check_output(["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
+sys.path.insert(0, str(checkout))
+from mtplx.server.openai import _server_runtime_env_overrides
+os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(root), generation_mode="mtp"), None))
+os.environ["MTPLX_SUSTAINED_PREFILL"] = "1"
+os.environ["MTPLX_PREFILL_CHUNK_SIZE"] = "2048"
+import mlx.core as mx
+import numpy as np
+from mlx_lm.utils import load_model
+from mtplx.models import qwen4_exp as qwen
+from mtplx.runtime import _load_tokenizer_resilient, MTPLXRuntime
+from mtplx.mtp_patch import MTPContract
+from mtplx.sampling import SamplerConfig
+from mtplx.generation import generate_mtpk
+from mtplx.server.openai import ChatMessage, _encode_messages_uncached, _materialize_vision_splice
+
+output.mkdir(parents=True, exist_ok=True)
+print(json.dumps(dict(case="environment", overrides={k:v for k,v in os.environ.items() if k.startswith("MTPLX_")})), flush=True)
+tokenizer = _load_tokenizer_resilient(root, json.loads((root/"config.json").read_text()))
+model, _ = load_model(root, lazy=False, strict=True, get_model_classes=lambda **_: (qwen.Model,qwen.ModelArgs))
+model.post_weight_load(root)
+assert model.attach_mtp(root)
+runtime = MTPLXRuntime(model, tokenizer, root, True, MTPContract())
+text = "Describe these images" + "<|vision_start|><|image_pad|><|vision_end|>" * 2
+ids = _encode_messages_uncached(tokenizer, [ChatMessage(role="user", content=text)], enable_thinking=True,
+    reasoning_effort="low", preserve_reasoning_history=True, tools=None)
+ids, splice = _materialize_vision_splice(argparse.Namespace(args=argparse.Namespace(model=str(root))), [image.read_bytes()]*2, ids)
+assert len(ids) > 2048
+(output/"input.json").write_text(json.dumps(ids))
+out = generate_mtpk(runtime, ids, max_tokens=1, sampler=SamplerConfig(temperature=0., top_p=.95, top_k=20),
+    speculative_depth=3, seed=1, stop_token_ids=set(tokenizer.eos_token_ids), mtp_history_policy="committed",
+    verify_strategy="batched", vision_splice=splice, capture_final_state=True,
+    prefill_callback=lambda event: print(json.dumps(dict(case="prefill", **event)), flush=True))
+assert out.final_state.safe_to_commit
+logits = np.asarray(out.final_state.final_logits.astype(mx.float32)).reshape(-1)
+(output/"final.f32").write_bytes(logits.tobytes())
+print(json.dumps(dict(case="long_vision", prompt_tokens=len(ids), tokens=out.tokens, text=out.text,
+    final_logits_sha256=hashlib.sha256(logits.tobytes()).hexdigest())), flush=True)