"""d3 text and video requests on MLX (Apple GPU): upstream prompt and answer code; mlx-vlm's Qwen3.5 vision tower and language model; all questions of a request in ONE batched forward (right-padded rows, key-padding mask on the noncausal full-attention layers, causal Gated DeltaNet); FP32 readout on the host. m = D3MLXVideo("", "", max_pixels=100_352) m.system_one(state=..., questions={...}, videos=[frames]) # or videos=None for a text request""" import sys, time from pathlib import Path import mlx.core as mx import numpy as np import torch _KEY_MASK = None _LAYER_TYPES = None def _patch(layer_types): """Two fixes to mlx-vlm's Qwen3.5 for d3. (1) Layer kinds come from the checkpoint's `layer_types`: mlx-vlm assumes full attention every `full_attention_interval` layers, but d3-edge has it at layers 2, 5, 8-11, so half its layers were built as the wrong kind and ran with random weights. (2) d3's full-attention layers are noncausal (`noncausal_full_attention`): every real token is visible, only padding keys are masked. Imported late: importing mlx_vlm's qwen3_5 registers its own Qwen3-VL processor, which upstream d3's request preparation must not get.""" global _LAYER_TYPES _LAYER_TYPES = list(layer_types) import mlx_vlm.models.qwen3_5.language as qlang if not getattr(qlang, "_d3_patched", False): init = qlang.Qwen3_5DecoderLayer.__init__ def layer_init(self, args, layer_idx): if _LAYER_TYPES is not None: kind = _LAYER_TYPES[layer_idx] # the interval the stock __init__ consults; make it agree with this layer's kind args.full_attention_interval = 1 if kind == "full_attention" else 10 ** 6 init(self, args, layer_idx) causal = qlang._create_qwen3_5_attention_mask def noncausal(h, cache): return _KEY_MASK if _KEY_MASK is not None else causal(h, cache) qlang.Qwen3_5DecoderLayer.__init__ = layer_init qlang._create_qwen3_5_attention_mask = noncausal qlang._d3_patched = True class D3MLXVideo: def __init__(self, root, mlx_dir, max_pixels=None): root = Path(root) sys.path.insert(0, str(root)) import d3_runtime as up if max_pixels: up.VIDEO_MAX_PIXELS = max_pixels from safetensors.torch import load_file self.up = up self.prep = up.D3.from_pretrained(str(root), device="meta", verify="none") # prompts and positions only self.processor, self.temperature, self.backbone = self.prep.processor, self.prep.temperature, self.prep.backbone self.readout = load_file(str(root / "readout.safetensors"))["weight"].float().numpy() self.video_token = self.backbone.config.video_token_id import json _patch(json.loads((root / "config.json").read_text())["text_config"]["layer_types"]) from mlx_vlm.utils import load_model self.model = load_model(Path(mlx_dir), strict=True) self.lm = self.model.language_model.model self.model_name = root.name + "-mlx" self.timing = {} def system_one(self, *, state, questions, videos=None): global _KEY_MASK up, proc = self.up, self.processor t0 = time.perf_counter() prepared = self.prep.prepare(state, questions, videos=videos) vis = None if videos: media = prepared.media grids = media["video_grid_thw"] pv = mx.array(media["pixel_values_videos"].float().numpy()) vis, _ = self.model.vision_tower(pv, mx.array(grids.numpy())) mx.eval(vis) # MLX is lazy: time the vision tower here t1 = time.perf_counter() rows, keys = [], [] for key in prepared.runnable: if videos: expanded = proc.get_text_with_replacements([prepared.texts[key]], [], prepared.video_texts)[0][0] enc = proc.tokenizer([expanded], return_token_type_ids=False, return_tensors="pt") ids = enc["input_ids"] types = torch.as_tensor(proc.create_mm_token_type_ids(ids)) with torch.no_grad(): pos = self.backbone.compute_3d_position_ids( input_ids=ids, inputs_embeds=None, image_grid_thw=None, video_grid_thw=grids, attention_mask=enc["attention_mask"], past_key_values=None, mm_token_type_ids=types)[:, 0] rows.append((ids[0].numpy(), pos.numpy())) else: # text: upstream's token ids, positions 0..n-1 on all three axes row = np.array(prepared.sequences[key], np.int32) rows.append((row, np.tile(np.arange(len(row), dtype=np.int32)[None], (3, 1)))) keys.append(key) B, L = len(rows), max(len(r[0]) for r in rows) ids = np.zeros((B, L), np.int32); pos = np.zeros((3, B, L), np.int32); valid = np.zeros((B, L), bool) lengths = [] for b, (r, p) in enumerate(rows): n = len(r); lengths.append(n) ids[b, :n] = r; ids[b, n:] = r[-1] pos[:, b, :n] = p; pos[:, b, n:] = p[:, -1:] valid[b, :n] = True emb = self.lm.embed_tokens(mx.array(ids)) if vis is not None: is_video = mx.array(ids == self.video_token) vis = vis.astype(emb.dtype) per_row = int(is_video[0].sum().item()) # every row carries the same video tokens, in the same order idx = mx.cumsum(is_video.astype(mx.int32), axis=1) - 1 emb = mx.where(is_video[..., None], vis[mx.clip(idx, 0, per_row - 1)], emb) _KEY_MASK = mx.array(valid)[:, None, None, :] try: h = self.lm(None, inputs_embeds=emb, position_ids=mx.array(pos)) finally: _KEY_MASK = None last = mx.stack([h[b, lengths[b] - 1] for b in range(B)]).astype(mx.float32) hs = np.array(last) t2 = time.perf_counter() answers, used = {}, 0 for b, key in enumerate(keys): question = prepared.questions[key] logits = self.readout[:len(question.keys)].astype(np.float64) @ hs[b].astype(np.float64) / self.temperature p = np.exp(logits - logits.max()) answers[key] = up.product_answer(question, (p / p.sum()).tolist()) used += lengths[b] self.timing = {"prepare+vision": round(t1 - t0, 3), "text": round(t2 - t1, 3)} return {"model": self.model_name, "answers": answers, "usage": {"input_tokens": used, "output_tokens": 0}}