Preserve image-aware sparse QSA for Qwen vision

tools/issue107-indexer-reference.py

@@ -0,0 +1,53 @@
+"""Real-geometry image-aware QSA preparation, checked independently with NumPy."""
+import hashlib
+import json
+from pathlib import Path
+import subprocess
+import sys
+
+checkout, output = map(Path, sys.argv[1:3])
+assert subprocess.check_output(["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
+sys.path.insert(0, str(checkout))
+import mlx.core as mx
+import numpy as np
+from mtplx.models.qwen4_exp import TextArgs, QSAIndexer, _build_mrope_axes
+from mtplx.attention_context import vision_rope
+
+output.mkdir(parents=True, exist_ok=True)
+args = TextArgs(hidden_size=640, head_dim=256, indexer_head_dim=128, indexer_n_heads=4,
+    indexer_kv_heads=1, indexer_compress_ratio=4, indexer_budget=2048,
+    rope_parameters=dict(mrope_interleaved=True, mrope_section=[11,11,10], partial_rotary_factor=.25,
+                         rope_theta=10000000, rope_type="default"))
+idx = QSAIndexer(args)
+axes = _build_mrope_axes([11,11,10],True)
+table = np.stack([np.arange(2060, dtype=np.int32)//(axis+1) for axis in range(3)])
+def pattern(shape, seed):
+    return (((np.arange(np.prod(shape), dtype=np.int64)*17+seed)%97-48)/32).astype(np.float32).reshape(shape)
+raw_q = pattern((1,16,4,128),3)
+raw_k = pattern((1,2072,128),11)
+
+for dtype_name, dtype in [("f32",mx.float32),("bf16",mx.bfloat16)]:
+    idx.q_layernorm.weight = mx.ones((128,),dtype=dtype)
+    idx.k_layernorm.weight = mx.ones((128,),dtype=dtype)
+    q, k = mx.array(raw_q).astype(dtype), mx.array(raw_k).astype(dtype)
+    for mode in ["table", "equal", "delta", "empty"]:
+        delta = -7 if mode in ("table","delta","empty") else 0
+        table_value = table if mode=="table" else np.broadcast_to(np.arange(2060,dtype=np.int32),(3,2060)) if mode=="equal" else np.zeros((3,0),dtype=np.int32) if mode=="empty" else None
+        with vision_rope(None if table_value is None else mx.array(table_value),delta):
+            query = idx._prepare_queries(q,2056)
+            pooled = idx._pool_keys_eager(k,0,518)
+        for name,value in [("query",query),("pooled",pooled)]:
+            data=np.asarray(value.astype(mx.float32))
+            (output/f"{dtype_name}-{mode}-{name}.f32").write_bytes(data.tobytes())
+            print(json.dumps(dict(case="indexer",dtype=dtype_name,mode=mode,part=name,shape=list(data.shape),sha256=hashlib.sha256(data.tobytes()).hexdigest())),flush=True)
+        if dtype_name=="f32":
+            def independent(values, positions):
+                values = values / np.sqrt(np.mean(values*values,axis=-1,keepdims=True)+args.rms_norm_eps)
+                pos3=np.stack([table_value[:,p] if table_value is not None and p<table_value.shape[1] else np.full(3,p+delta) for p in positions],axis=1)
+                angles=pos3[axes].T.astype(np.float32)*np.asarray(idx._inv_freq)[None,:]
+                angles=np.concatenate([angles,angles],axis=-1)[None,:,None,:]
+                first=values[...,:64];rotated=np.concatenate([-first[...,32:],first[...,:32]],axis=-1)
+                return np.concatenate([first*np.cos(angles)+rotated*np.sin(angles),values[...,64:]],axis=-1)
+            np.testing.assert_allclose(np.asarray(query),independent(raw_q,np.arange(2056,2072)),rtol=3e-6,atol=3e-6)
+            np.testing.assert_allclose(np.asarray(pooled),independent(raw_k.reshape(1,518,4,128).mean(axis=2)[:,:,None],np.arange(0,2072,4))[:,:,0],rtol=3e-6,atol=3e-6)
+print(json.dumps(dict(case="independent_numpy",passed=True)),flush=True)

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)